From 55c021fe82cc723adcdef43b98fab3889f0d2248 Mon Sep 17 00:00:00 2001
From: MerryMage <MerryMage@users.noreply.github.com>
Date: Fri, 19 Jun 2020 12:11:35 +0100
Subject: [PATCH] emit_x64_aes: AESNI implementations of all opcodes

---
 src/backend/x64/emit_x64_aes.cpp | 48 ++++++++++++++++++++++++++------
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/src/backend/x64/emit_x64_aes.cpp b/src/backend/x64/emit_x64_aes.cpp
index 17c63d7d..431a919c 100644
--- a/src/backend/x64/emit_x64_aes.cpp
+++ b/src/backend/x64/emit_x64_aes.cpp
@@ -17,8 +17,7 @@ namespace AES = Common::Crypto::AES;
 
 using AESFn = void(AES::State&, const AES::State&);
 
-static void EmitAESFunction(RegAlloc::ArgumentInfo args, EmitContext& ctx, BlockOfCode& code,
-                            IR::Inst* inst, AESFn fn) {
+static void EmitAESFunction(RegAlloc::ArgumentInfo args, EmitContext& ctx, BlockOfCode& code, IR::Inst* inst, AESFn fn) {
     constexpr u32 stack_space = static_cast<u32>(sizeof(AES::State)) * 2;
     const Xbyak::Xmm input = ctx.reg_alloc.UseXmm(args[0]);
     const Xbyak::Xmm result = ctx.reg_alloc.ScratchXmm();
@@ -26,16 +25,13 @@ static void EmitAESFunction(RegAlloc::ArgumentInfo args, EmitContext& ctx, Block
 
     ctx.reg_alloc.HostCall(nullptr);
     code.sub(rsp, stack_space + ABI_SHADOW_SPACE);
+
     code.lea(code.ABI_PARAM1, ptr[rsp + ABI_SHADOW_SPACE]);
     code.lea(code.ABI_PARAM2, ptr[rsp + ABI_SHADOW_SPACE + sizeof(AES::State)]);
-
     code.movaps(xword[code.ABI_PARAM2], input);
-
     code.CallFunction(fn);
-
     code.movaps(result, xword[rsp + ABI_SHADOW_SPACE]);
 
-    // Free memory
     code.add(rsp, stack_space + ABI_SHADOW_SPACE);
 
     ctx.reg_alloc.DefineValue(inst, result);
@@ -44,12 +40,34 @@ static void EmitAESFunction(RegAlloc::ArgumentInfo args, EmitContext& ctx, Block
 void EmitX64::EmitAESDecryptSingleRound(EmitContext& ctx, IR::Inst* inst) {
     auto args = ctx.reg_alloc.GetArgumentInfo(inst);
 
+    if (code.HasAESNI()) {
+        const Xbyak::Xmm data = ctx.reg_alloc.UseScratchXmm(args[0]);
+        const Xbyak::Xmm zero = ctx.reg_alloc.ScratchXmm();
+
+        code.pxor(zero, zero);
+        code.aesdeclast(data, zero);
+
+        ctx.reg_alloc.DefineValue(inst, data);
+        return;
+    }
+
     EmitAESFunction(args, ctx, code, inst, AES::DecryptSingleRound);
 }
 
 void EmitX64::EmitAESEncryptSingleRound(EmitContext& ctx, IR::Inst* inst) {
     auto args = ctx.reg_alloc.GetArgumentInfo(inst);
 
+    if (code.HasAESNI()) {
+        const Xbyak::Xmm data = ctx.reg_alloc.UseScratchXmm(args[0]);
+        const Xbyak::Xmm zero = ctx.reg_alloc.ScratchXmm();
+
+        code.pxor(zero, zero);
+        code.aesenclast(data, zero);
+
+        ctx.reg_alloc.DefineValue(inst, data);
+        return;
+    }
+
     EmitAESFunction(args, ctx, code, inst, AES::EncryptSingleRound);
 }
 
@@ -62,13 +80,27 @@ void EmitX64::EmitAESInverseMixColumns(EmitContext& ctx, IR::Inst* inst) {
         code.aesimc(data, data);
 
         ctx.reg_alloc.DefineValue(inst, data);
-    } else {
-        EmitAESFunction(args, ctx, code, inst, AES::InverseMixColumns);
+        return;
     }
+
+    EmitAESFunction(args, ctx, code, inst, AES::InverseMixColumns);
 }
 
 void EmitX64::EmitAESMixColumns(EmitContext& ctx, IR::Inst* inst) {
     auto args = ctx.reg_alloc.GetArgumentInfo(inst);
+
+    if (code.HasAESNI()) {
+        const Xbyak::Xmm data = ctx.reg_alloc.UseScratchXmm(args[0]);
+        const Xbyak::Xmm zero = ctx.reg_alloc.ScratchXmm();
+
+        code.pxor(zero, zero);
+        code.aesdeclast(data, zero);
+        code.aesenc(data, zero);
+
+        ctx.reg_alloc.DefineValue(inst, data);
+        return;
+    }
+
     EmitAESFunction(args, ctx, code, inst, AES::MixColumns);
 }