diff --git a/library/aes.c b/library/aes.c
index 6d090a1..d2d1c0c 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -677,7 +677,7 @@
     uint32_t *RK, X0, X1, X2, X3, Y0, Y1, Y2, Y3;
 
 #if defined(POLARSSL_AESNI_C) && defined(POLARSSL_HAVE_X86_64)
-    if( aesni_supported() )
+    if( aesni_supports( POLARSSL_AESNI_AES ) )
         return( aesni_crypt_ecb( ctx, mode, input, output ) );
 #endif
 
diff --git a/library/aesni.c b/library/aesni.c
index 7628a03..9b41c36 100644
--- a/library/aesni.c
+++ b/library/aesni.c
@@ -37,24 +37,24 @@
 #if defined(POLARSSL_HAVE_X86_64)
 
 /*
- * AES-NI support detection routine, [AES-WP] figure 23
+ * AES-NI support detection routine
  */
-int aesni_supported( void )
+int aesni_supports( unsigned int what )
 {
-    static int supported = -1;
-    unsigned int c;
+    static int done = 0;
+    static unsigned int c = 0;
 
-    if( supported == -1 )
+    if( ! done )
     {
         asm( "movl  $1, %%eax   \n"
              "cpuid             \n"
              : "=c" (c)
              :
              : "eax", "ebx", "edx" );
-        supported = ( ( c & 0x02000000 ) != 0 );
+        done = 1;
     }
 
-    return( supported );
+    return( ( c & what ) != 0 );
 }
 
 /*