ARM AES support

osu-crypto · Oct 25, 2024 · 280d6a3 · 280d6a3
1 parent e59724e
commit 280d6a3
Show file tree

Hide file tree

Showing 9 changed files with 476 additions and 320 deletions.
diff --git a/CMakePresets.json b/CMakePresets.json
@@ -98,6 +98,7 @@
       "cacheVariables": {
         "CMAKE_BUILD_TYPE": "RelWithDebInfo",
         "LIBOTE_STD_VER": "20",
+        "ENABLE_PORTABLE_AES": true,
         "ENABLE_GMP": false,
         "ENABLE_ALL_OT": true,
         "ENABLE_RELIC": false,

diff --git a/cryptoTools b/cryptoTools
diff --git a/frontend/benchmark.h b/frontend/benchmark.h
@@ -606,4 +606,88 @@ namespace osuCrypto
         std::cout << "ENABLE_Silent_VOLE = false" << std::endl;
 #endif
     }
+
+
+    void AESBenchmark(const oc::CLP& cmd)
+    {
+        u64 n = roundUpTo(cmd.getOr("n", 1ull << cmd.getOr("nn", 20)), 8);
+        u64 t =cmd.getOr("t", 10);
+        using AES_ = AES;// details::AES<details::AESTypes::Portable>;
+
+        auto unroll8 = [](AES_& aes, block* __restrict s)
+            {
+                block b[8];
+                b[0] = AES_::firstFn(s[0], aes.mRoundKey[0]);
+                b[1] = AES_::firstFn(s[1], aes.mRoundKey[0]);
+                b[2] = AES_::firstFn(s[2], aes.mRoundKey[0]);
+                b[3] = AES_::firstFn(s[3], aes.mRoundKey[0]);
+                b[4] = AES_::firstFn(s[4], aes.mRoundKey[0]);
+                b[5] = AES_::firstFn(s[5], aes.mRoundKey[0]);
+                b[6] = AES_::firstFn(s[6], aes.mRoundKey[0]);
+                b[7] = AES_::firstFn(s[7], aes.mRoundKey[0]);
+
+                for (u64 i = 1; i < 9; ++i)
+                {
+                    b[0] = AES_::roundFn(b[0], aes.mRoundKey[i]);
+                    b[1] = AES_::roundFn(b[1], aes.mRoundKey[i]);
+                    b[2] = AES_::roundFn(b[2], aes.mRoundKey[i]);
+                    b[3] = AES_::roundFn(b[3], aes.mRoundKey[i]);
+                    b[4] = AES_::roundFn(b[4], aes.mRoundKey[i]);
+                    b[5] = AES_::roundFn(b[5], aes.mRoundKey[i]);
+                    b[6] = AES_::roundFn(b[6], aes.mRoundKey[i]);
+                    b[7] = AES_::roundFn(b[7], aes.mRoundKey[i]);
+                }
+
+
+                b[0] = AES_::penultimateFn(b[0], aes.mRoundKey[9]);
+                b[1] = AES_::penultimateFn(b[1], aes.mRoundKey[9]);
+                b[2] = AES_::penultimateFn(b[2], aes.mRoundKey[9]);
+                b[3] = AES_::penultimateFn(b[3], aes.mRoundKey[9]);
+                b[4] = AES_::penultimateFn(b[4], aes.mRoundKey[9]);
+                b[5] = AES_::penultimateFn(b[5], aes.mRoundKey[9]);
+                b[6] = AES_::penultimateFn(b[6], aes.mRoundKey[9]);
+                b[7] = AES_::penultimateFn(b[7], aes.mRoundKey[9]);
+                s[0] = AES_::finalFn(b[0], aes.mRoundKey[10]);
+                s[1] = AES_::finalFn(b[1], aes.mRoundKey[10]);
+                s[2] = AES_::finalFn(b[2], aes.mRoundKey[10]);
+                s[3] = AES_::finalFn(b[3], aes.mRoundKey[10]);
+                s[4] = AES_::finalFn(b[4], aes.mRoundKey[10]);
+                s[5] = AES_::finalFn(b[5], aes.mRoundKey[10]);
+                s[6] = AES_::finalFn(b[6], aes.mRoundKey[10]);
+                s[7] = AES_::finalFn(b[7], aes.mRoundKey[10]);
+
+            };
+
+        oc::AlignedUnVector<block> x(n);
+        auto n8 = n / 8;
+        AES_ aes(block(42352345, 3245345234676534));
+        Timer timer;
+        timer.setTimePoint("begin");
+        for (u64 tt = 0; tt < t; ++tt)
+        {
+            for (u64 i = 0; i < n; i += 8)
+            {
+                unroll8(aes, x.data() + i);
+            }
+            timer.setTimePoint("unroll");
+        }
+
+        for (u64 tt = 0; tt < t; ++tt)
+        {
+            for (u64 i = 0; i < n; i += 8)
+            {
+                aes.ecbEncBlocks<8>(x.data() + i, x.data() + i);
+            }
+            timer.setTimePoint("aes <>");
+        }
+
+        for (u64 tt = 0; tt < t; ++tt)
+        {
+                aes.ecbEncBlocks(x, x);
+            timer.setTimePoint("aes ");
+        }
+
+        std::cout << timer << std::endl;
+
+    }
 }
diff --git a/frontend/main.cpp b/frontend/main.cpp
@@ -117,6 +117,8 @@ int main(int argc, char** argv)
 			ExConvCodeOldBench(cmd);
 		else if (cmd.isSet("tungsten"))
 			TungstenCodeBench(cmd);
+		else if (cmd.isSet("aes"))
+			AESBenchmark(cmd);
 
 		return 0;
 	}

diff --git a/libOTe/NChooseOne/Kkrt/KkrtNcoOtReceiver.h b/libOTe/NChooseOne/Kkrt/KkrtNcoOtReceiver.h
@@ -15,6 +15,7 @@
 #include <vector>
 #include <cryptoTools/Crypto/AES.h>
 #include <cryptoTools/Common/Timer.h>
+#include <cryptoTools/Crypto/MultiKeyAES.h>
 
 #ifdef GetMessage
 #undef GetMessage

diff --git a/libOTe/NChooseOne/Kkrt/KkrtNcoOtSender.h b/libOTe/NChooseOne/Kkrt/KkrtNcoOtSender.h
@@ -17,6 +17,7 @@
 
 #include <cryptoTools/Common/Timer.h>
 #include <cryptoTools/Network/Channel.h>
+#include <cryptoTools/Crypto/MultiKeyAES.h>
 
 #include <array>
 #include <vector>

diff --git a/libOTe/TwoChooseOne/Kos/KosOtExtReceiver.h b/libOTe/TwoChooseOne/Kos/KosOtExtReceiver.h
@@ -14,6 +14,7 @@
 #include <cryptoTools/Crypto/PRNG.h>
 #include <cryptoTools/Common/Timer.h>
 #include "libOTe/Tools/Coproto.h"
+#include <cryptoTools/Crypto/MultiKeyAES.h>
 
 namespace osuCrypto
 {

diff --git a/libOTe/TwoChooseOne/Kos/KosOtExtSender.h b/libOTe/TwoChooseOne/Kos/KosOtExtSender.h
@@ -14,6 +14,7 @@
 #include <cryptoTools/Common/Timer.h>
 #include <cryptoTools/Crypto/PRNG.h>
 #include <array>
+#include <cryptoTools/Crypto/MultiKeyAES.h>
 
 namespace osuCrypto {
+15 −1		cmake/cryptoToolsBuildOptions.cmake
+6 −1		cryptoTools/CMakeLists.txt
+2 −2		cryptoTools/Circuit/BetaCircuit.cpp
+1 −1		cryptoTools/Circuit/MxCircuitLibrary.cpp
+19 −0		cryptoTools/Common/block.h
+8 −3		cryptoTools/Common/config.h.in
+209 −320		cryptoTools/Crypto/AES.cpp
+667 −682		cryptoTools/Crypto/AES.h
+203 −0		cryptoTools/Crypto/MultiKeyAES.h
+154 −124		tests_cryptoTools/AES_Tests.cpp