Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use AESNI for AES-CTR if available #371

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 155 additions & 2 deletions Source/C++/Crypto/Ap4AesBlockCipher.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1933,7 +1933,7 @@ AP4_AesCtrBlockCipher::Process(const AP4_UI08* input,
} else {
AP4_SetMemory(counter, 0, AP4_AES_BLOCK_SIZE);
}

// process all blocks
while (input_size) {
AP4_UI08 block[AP4_AES_BLOCK_SIZE];
Expand Down Expand Up @@ -1963,6 +1963,150 @@ AP4_AesCtrBlockCipher::Process(const AP4_UI08* input,
return AP4_SUCCESS;
}

#if AP4_AES_BLOCK_SIZE == 16 && AP4_AES_KEY_LENGTH == 16 && defined(__AES__) && defined(__SSE4_2__)
#define AP4_ENABLE_AESNI

#include <immintrin.h>

#define cpuid(func,ax,bx,cx,dx)\
__asm__ __volatile__ ("cpuid":\
"=a" (ax), "=b" (bx), "=c" (cx), "=d" (dx) : "a" (func));

static bool g_SupportAesNI = false;
__attribute__((constructor)) static void detect_aesni()
{
unsigned int a,b,c,d;
cpuid(1, a,b,c,d);
g_SupportAesNI = c & 0x2000000;
}

#endif

#ifdef AP4_ENABLE_AESNI
template <AP4_Size NB>
static void aesni_process_NB_blocks(AP4_UI08* Out, AP4_UI08 const* In, __m128i& CB, AP4_UI08 const* Keys)
{
__m128i CBs[NB];
const __m128i CB_ = CB;
#pragma unroll
for (AP4_Size i = 0; i < NB; ++i) {
CBs[i] = _mm_add_epi64(CB_, _mm_set_epi32(0,i+1,0,0));
}
CB = CBs[NB-1];
const __m128i Bswap64 = _mm_set_epi8(8,9,10,11,12,13,14,15,7,6,5,4,3,2,1,0);
#pragma unroll
for (AP4_Size i = 0; i < NB; ++i) {
CBs[i] = _mm_shuffle_epi8(CBs[i], Bswap64);
}

__m128i EncrCBs[NB];
const __m128i Key0 = _mm_loadu_si128((const __m128i*)Keys);
#pragma unroll
for (AP4_Size i = 0; i < NB; ++i) {
EncrCBs[i] = _mm_xor_si128(CBs[i], Key0);
}
#pragma unroll
for (AP4_Size R = 1; R < 10; R++) {
const __m128i Key = _mm_loadu_si128((const __m128i*)&Keys[R*AP4_AES_KEY_LENGTH]);
#pragma unroll
for (AP4_Size i = 0; i < NB; ++i) {
EncrCBs[i] = _mm_aesenc_si128(EncrCBs[i], Key);
}
}

const __m128i KeyLast = _mm_loadu_si128((const __m128i*)&Keys[10*AP4_AES_KEY_LENGTH]);
#pragma unroll
for (AP4_Size i = 0; i < NB; ++i) {
EncrCBs[i] = _mm_aesenclast_si128(EncrCBs[i],
_mm_xor_si128(KeyLast, _mm_loadu_si128((const __m128i*)&In[i*AP4_AES_BLOCK_SIZE])));
}

#pragma unroll
for (AP4_Size i = 0; i < NB; ++i) {
_mm_storeu_si128((__m128i*)&Out[i*AP4_AES_BLOCK_SIZE], EncrCBs[i]);
}
}

static void process(AP4_UI08* output, AP4_UI08 const* input, AP4_Size input_size, const AP4_UI08* iv, const AP4_UI08* Keys)
{
union {
AP4_UI08 B[AP4_AES_BLOCK_SIZE];
__m128i V;
} Counter;

if (iv) {
memcpy(&Counter.B[0], iv, AP4_AES_BLOCK_SIZE);
Counter.V = _mm_shuffle_epi8(Counter.V, _mm_set_epi8(8,9,10,11,12,13,14,15,7,6,5,4,3,2,1,0));
Counter.V = _mm_sub_epi64(Counter.V, _mm_set_epi32(0,1,0,0));
} else {
memset(&Counter.B[0], 0, AP4_AES_BLOCK_SIZE);
}

// First, process blocks eight by eight (Intel recommandation)
const AP4_Size Size8B = 8*AP4_AES_BLOCK_SIZE;
const AP4_Size End8B = (input_size/Size8B)*Size8B;
for (AP4_Size i = 0; i < End8B; i += Size8B) {
aesni_process_NB_blocks<8>(&output[i], &input[i], Counter.V, Keys);
}
// Process the remaining blocks!
const AP4_Size RemBlocks = (input_size-End8B)/AP4_AES_BLOCK_SIZE;
AP4_Size CurIdx = End8B;
switch (RemBlocks) {
#define FINAL_BLOCKS(N)\
case N:\
aesni_process_NB_blocks<N>(&output[CurIdx], &input[CurIdx], Counter.V, Keys);\
CurIdx += N*AP4_AES_BLOCK_SIZE;\
break;

FINAL_BLOCKS(7)
FINAL_BLOCKS(6)
FINAL_BLOCKS(5)
FINAL_BLOCKS(4)
FINAL_BLOCKS(3)
FINAL_BLOCKS(2)
FINAL_BLOCKS(1)
#undef FINAL_BLOCKS
}
const AP4_Size Rem = input_size-CurIdx;
assert(Rem < 16 && "too many remaining bytes!");
if (Rem > 0) {
// Last block
AP4_UI08 LastBlock[AP4_AES_BLOCK_SIZE];
memcpy(&LastBlock[0], &input[CurIdx], Rem);
aesni_process_NB_blocks<1>(&LastBlock[0], &LastBlock[0], Counter.V, Keys);
memcpy(&output[CurIdx], &LastBlock[0], Rem);
}
}

class AP4_AesNICtrBlockCipher : public AP4_AesBlockCipher
{
public:
AP4_AesNICtrBlockCipher(CipherDirection direction,
unsigned int counter_size,
aes_ctx* context) :
AP4_AesBlockCipher(direction, CTR, context)
{
assert(counter_size == 8 && "counter size must be 8 bytes!");
}

// AP4_BlockCipher methods
virtual AP4_Result Process(const AP4_UI08* input,
AP4_Size input_size,
AP4_UI08* output,
const AP4_UI08* iv);
};

AP4_Result AP4_AesNICtrBlockCipher::Process(const AP4_UI08* input,
AP4_Size input_size,
AP4_UI08* output,
const AP4_UI08* iv)
{
assert(m_Context->n_rnd == 10 && "this only works for AES128!");
process(output, input, input_size, iv, (const AP4_UI08*) &m_Context->k_sch[0]);
return AP4_SUCCESS;
}
#endif // AP4_ENABLE_AESNI

/*----------------------------------------------------------------------
| AP4_AesBlockCipher::Create
+---------------------------------------------------------------------*/
Expand Down Expand Up @@ -1994,7 +2138,16 @@ AP4_AesBlockCipher::Create(const AP4_UI08* key,
if (ctr_params) {
counter_size = ctr_params->counter_size;
}
cipher = new AP4_AesCtrBlockCipher(direction, counter_size, context);
#ifdef AP4_ENABLE_AESNI
if (g_SupportAesNI && (counter_size == 8)) {
cipher = new AP4_AesNICtrBlockCipher(direction, counter_size, context);
}
else {
#endif
cipher = new AP4_AesCtrBlockCipher(direction, counter_size, context);
#ifdef AP4_ENABLE_AESNI
}
#endif
break;
}

Expand Down