From 7d87c4f6c5f631d62aaaed761935dfe91856c0e4 Mon Sep 17 00:00:00 2001
From: David Harris <David_Harris@hmc.edu>
Date: Mon, 11 Mar 2024 00:53:39 -0700
Subject: [PATCH] AES64 simplification

---
 src/ieu/bmu/bitmanipalu.sv |  3 ++-
 src/ieu/kmu/zkn64.sv       | 33 ++++++++++++++++++++++++++++++---
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/ieu/bmu/bitmanipalu.sv b/src/ieu/bmu/bitmanipalu.sv
index 3d07dc498..bf6914b88 100644
--- a/src/ieu/bmu/bitmanipalu.sv
+++ b/src/ieu/bmu/bitmanipalu.sv
@@ -117,7 +117,8 @@ module bitmanipalu import cvw::*; #(parameter cvw_t P) (
     if (P.ZKND_SUPPORTED) aes32d aes32d(.bs(Funct7[6:5]), .rs1(ABMU), .rs2(BBMU), .finalround(ZBBSelect[2]), .result(ZKNDResult));
     if (P.ZKNE_SUPPORTED) aes32e aes32e(.bs(Funct7[6:5]), .rs1(ABMU), .rs2(BBMU), .finalround(ZBBSelect[2]), .result(ZKNEResult));
   end else 
-    if (P.ZKND_SUPPORTED | P.ZKNE_SUPPORTED) zkn64 ZKND64(.A(ABMU), .B(BBMU), .Funct7, .round(Rs2E[3:0]), .ZKNSelect(ZBBSelect[3:0]), .ZKNDResult, .ZKNEResult); // *** simplify to only one output
+    if (P.ZKND_SUPPORTED | P.ZKNE_SUPPORTED) 
+      zkn64 #(P) ZKND64(.A(ABMU), .B(BBMU), .Funct7, .round(Rs2E[3:0]), .ZKNSelect(ZBBSelect[3:0]), .ZKNDResult, .ZKNEResult); // *** simplify to only one output
 
 /*
   // ZKND Unit
diff --git a/src/ieu/kmu/zkn64.sv b/src/ieu/kmu/zkn64.sv
index 1960bbaa8..0cc2238f6 100644
--- a/src/ieu/kmu/zkn64.sv
+++ b/src/ieu/kmu/zkn64.sv
@@ -26,7 +26,7 @@
 // and limitations under the License.
 ////////////////////////////////////////////////////////////////////////////////////////////////
 
-module zkn64 (
+module zkn64 import cvw::*; #(parameter cvw_t P) (
    input  logic [63:0] A, B,
    input  logic [6:0] 	    Funct7,
    input  logic [3:0] 	    round,
@@ -34,6 +34,33 @@ module zkn64 (
    output logic [63:0] ZKNDResult, ZKNEResult
 );
    
-    zknd64 #(64) ZKND64(.A, .B, .Funct7, .round, .ZKNDSelect(ZKNSelect[3:0]), .ZKNDResult); // *** strip out parameter unneded
-    zkne64 #(64) ZKNE64(.A, .B, .Funct7, .round, .ZKNESelect(ZKNSelect[2:0]), .ZKNEResult);
+//    zknd64 #(64) ZKND64(.A, .B, .Funct7, .round, .ZKNDSelect(ZKNSelect[3:0]), .ZKNDResult); // *** strip out parameter unneded
+//    zkne64 #(64) ZKNE64(.A, .B, .Funct7, .round, .ZKNESelect(ZKNSelect[2:0]), .ZKNEResult);
+
+
+    logic [63:0] 	     aes64dRes, aes64eRes, aes64imRes, aes64ks1iRes, aes64ks2Res;
+   
+    if (P.ZKND_SUPPORTED) // ZKND supports aes64ds, aes64dsm, aes64im
+        aes64d    aes64d(.rs1(A), .rs2(B), .finalround(ZKNSelect[2]), .aes64im(ZKNSelect[3]), .result(aes64dRes)); // decode AES
+    if (P.ZKNE_SUPPORTED) // ZKNE supports aes64es, aes64esm
+        aes64e    aes64e(.rs1(A), .rs2(B), .finalround(ZKNSelect[2]), .result(aes64eRes));
+
+    // Both ZKND and ZKNE support aes64ks1i and aes64ks2 instructions
+    aes64ks1i aes64ks1i(.round, .rs1(A), .result(aes64ks1iRes));
+    aes64ks2  aes64ks2(.rs2(B), .rs1(A), .result(aes64ks2Res));
+   
+    mux3 #(64) zkndmux(aes64dRes, aes64ks1iRes, aes64ks2Res, ZKNSelect[1:0], ZKNDResult);
+    mux3 #(64) zknemux(aes64eRes, aes64ks1iRes, aes64ks2Res, ZKNSelect[1:0], ZKNEResult);
+
+/*
+     logic [63:0] 	     aes64eRes, aes64ks1iRes, aes64ks2Res;
+   
+   // RV64
+   aes64ks1i aes64ks1i(.round, .rs1(A), .result(aes64ks1iRes));
+   aes64ks2  aes64ks2(.rs2(B), .rs1(A), .result(aes64ks2Res)); 
+   
+   // 010 is a placeholder to match the select of ZKND's AES64KS1I since they share some instruction
+   mux4 #(WIDTH) zknemux(aes64eRes, aes64dRes, aes64ks1iRes, aes64ks2Res, ZKNESelect[1:0], ZKNEResult);   
+*/
+
 endmodule