From 8af25a45e6732be07661973ed1e4866fab4715d6 Mon Sep 17 00:00:00 2001
From: David Harris <David_Harris@hmc.edu>
Date: Mon, 11 Mar 2024 01:36:46 -0700
Subject: [PATCH] AES32 sharing logic

---
 src/ieu/aes_instructions/aes32d.sv | 15 ++++-----------
 src/ieu/aes_instructions/aes32e.sv | 15 ++++-----------
 src/ieu/kmu/zknde32.sv             | 23 +++++++++++++++++------
 3 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/src/ieu/aes_instructions/aes32d.sv b/src/ieu/aes_instructions/aes32d.sv
index 606328af4..95b75fc80 100644
--- a/src/ieu/aes_instructions/aes32d.sv
+++ b/src/ieu/aes_instructions/aes32d.sv
@@ -26,23 +26,16 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////
 
 module aes32d(
-   input  logic [1:0]  bs,
-   input  logic [31:0] rs1,
-   input  logic [31:0] rs2,
+   input  logic [7:0]  SboxIn,
    input  logic        finalround,
    output logic [31:0] result
 );
 
-   logic [4:0] 			  shamt;
-   logic [7:0] 			  SboxIn, SboxOut;
-   logic [31:0] 		     so, mixed, rotin, rotout;
+   logic [7:0] 			  SboxOut;
+   logic [31:0] 		     so, mixed;
    
-   assign shamt = {bs, 3'b0};                     // shamt = bs * 8 (convert bytes to bits)
-   assign SboxIn = rs2[shamt +: 8];               // select byte bs of rs2
    aesinvsbox inv_sbox(SboxIn, SboxOut);          // Apply inverse sbox to si
    assign so = {24'h0, SboxOut};                  // Pad output of inverse substitution box
    aesinvmixcolumns mix(so, mixed);               // Run so through the mixword AES function
-   mux2 #(32) rmux(mixed, so, finalround, rotin); // on final round, skip mixcolumns
-   rotate #(32) rot(rotin, shamt, rotout);        // Rotate left by shamt (bs * 8)
-   assign result = rs1 ^ rotout;                  // xor with running value
+   mux2 #(32) rmux(mixed, so, finalround, result); // on final round, skip mixcolumns
 endmodule
diff --git a/src/ieu/aes_instructions/aes32e.sv b/src/ieu/aes_instructions/aes32e.sv
index 8ee5cb858..969f8a25f 100644
--- a/src/ieu/aes_instructions/aes32e.sv
+++ b/src/ieu/aes_instructions/aes32e.sv
@@ -26,23 +26,16 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////
 
 module aes32e(
-   input  logic [1:0]  bs,
-   input  logic [31:0] rs1,
-   input  logic [31:0] rs2,
+   input  logic [7:0]  SboxIn,
    input  logic        finalround,
    output logic [31:0] result
 );                
                 
-   logic [4:0] 			  shamt;
-   logic [7:0] 			  SboxIn, SboxOut;
-   logic [31:0] 		     so, mixed, rotin, rotout;
+   logic [7:0] 			  SboxOut;
+   logic [31:0] 		     so, mixed;
    
-   assign shamt = {bs, 3'b0};                     // shamt = bs * 8 (convert bytes to bits)
-   assign SboxIn = rs2[shamt +: 8];               // select byte bs of rs2
    aessbox sbox(SboxIn, SboxOut);                 // Substitute
    assign so = {24'h0, SboxOut};                  // Pad sbox output
    aesmixcolumns mwd(so, mixed);                  // Mix Word using aesmixword component
-   mux2 #(32) rmux(mixed, so, finalround, rotin); // on final round, skip mixcolumns
-   rotate #(32) mrot(rotin, shamt, rotout);       // Rotate the mixcolumns output left by shamt (bs * 8)
-   assign result = rs1 ^ rotout;                  // xor with running value
+   mux2 #(32) rmux(mixed, so, finalround, result); // on final round, skip mixcolumns
 endmodule
diff --git a/src/ieu/kmu/zknde32.sv b/src/ieu/kmu/zknde32.sv
index 32eed01ac..24b045f31 100644
--- a/src/ieu/kmu/zknde32.sv
+++ b/src/ieu/kmu/zknde32.sv
@@ -34,16 +34,27 @@ module zknde32 import cvw::*; #(parameter cvw_t P) (
    output logic [31:0] ZKNDEResult
 );
 
-     logic [31:0]        ZKNEResult, ZKNDResult;             
+    logic [4:0] 	shamt;
+    logic [7:0]     SboxIn;
+    logic [31:0]    ZKNEResult, ZKNDResult, rotin, rotout;             
 
-    if (P.ZKND_SUPPORTED) aes32d aes32d(.bs(Funct7[6:5]), .rs1(A), .rs2(B), .finalround(ZKNSelect[2]), .result(ZKNDResult));
-    if (P.ZKNE_SUPPORTED) aes32e aes32e(.bs(Funct7[6:5]), .rs1(A), .rs2(B), .finalround(ZKNSelect[2]), .result(ZKNEResult));
+    // Initial shamt and Sbox input selection steps shared between encrypt and decrypt
+    assign shamt = {Funct7[6:5], 3'b0};                     // shamt = bs * 8 (convert bytes to bits)
+    assign SboxIn = B[shamt +: 8];               // select byte bs of rs2
+
+    // Handle logic specific to encrypt or decrypt
+    if (P.ZKND_SUPPORTED) aes32d aes32d(.SboxIn, .finalround(ZKNSelect[2]), .result(ZKNDResult));
+    if (P.ZKNE_SUPPORTED) aes32e aes32e(.SboxIn, .finalround(ZKNSelect[2]), .result(ZKNEResult));
 
     // Mux result if both decrypt and encrypt are supported; otherwise, choose the only result
     if (P.ZKND_SUPPORTED & P.ZKNE_SUPPORTED) 
-        mux2 #(32) zknmux(ZKNDResult, ZKNEResult, ZKNSelect[0], ZKNDEResult); 
+        mux2 #(32) zknmux(ZKNDResult, ZKNEResult, ZKNSelect[0], rotin); 
     else if (P.ZKND_SUPPORTED)
-        assign ZKNDEResult = ZKNDResult;
+        assign rotin = ZKNDResult;
     else 
-        assign ZKNDEResult = ZKNEResult;
+        assign rotin = ZKNEResult;
+
+    // final rotate and XOR steps shared between encrypt and decrypt
+    rotate #(32) mrot(rotin, shamt, rotout);       // Rotate the mixcolumns output left by shamt (bs * 8)
+    assign ZKNDEResult = A ^ rotout;               // xor with running value
 endmodule