From 3d72ccac607a9d161b6e57a6ca72467e49bc0c01 Mon Sep 17 00:00:00 2001
From: David Harris <David_Harris@hmc.edu>
Date: Sun, 10 Mar 2024 22:37:50 -0700
Subject: [PATCH] AES simplification

---
 src/ieu/aes_instructions/aes32dsi.sv  | 33 ++++++++-------------------
 src/ieu/aes_instructions/aes32dsmi.sv |  4 +---
 src/ieu/aes_instructions/aes32esi.sv  |  7 +-----
 src/ieu/aes_instructions/aes32esmi.sv |  9 ++------
 src/ieu/aes_instructions/aes64ks2.sv  |  2 +-
 5 files changed, 14 insertions(+), 41 deletions(-)

diff --git a/src/ieu/aes_instructions/aes32dsi.sv b/src/ieu/aes_instructions/aes32dsi.sv
index 29c68b67f..dd399d75d 100644
--- a/src/ieu/aes_instructions/aes32dsi.sv
+++ b/src/ieu/aes_instructions/aes32dsi.sv
@@ -33,28 +33,13 @@ module aes32dsi(
 );
 
    logic [4:0] 			 shamt;
-   logic [31:0] 		    SboxIn32;
-   logic [7:0] 			 SboxIn;
-   logic [7:0] 			 SboxOut;
-   logic [31:0] 		    so;
-   logic [31:0] 		    sorotate;   
-   
-   // shamt = bs * 8
-   assign shamt = {bs, 3'b0};
-   
-   // Shift rs2 right by shamt and take the lower byte
-   assign SboxIn32 = (rs2 >> shamt);
-   assign SboxIn = SboxIn32[7:0];
-   
-   // Apply inverse sbox to si
-   aesinvsbox inv_sbox(SboxIn, SboxOut);
-   
-   // Pad output of inverse substitution box
-   assign so = {24'h0, SboxOut};
-   
-   // Rotate the substitution box output left by shamt (bs * 8)
-   assign sorotate = (so << shamt) | (so >> (32 - shamt)); 
-   
-   // Set result to "X(rs1)[31..0] ^ rol32(so, unsigned(shamt));"
-   assign DataOut = rs1 ^ sorotate;   
+   logic [7:0] 			 SboxIn, SboxOut;
+   logic [31:0] 		    so, sorotate;
+  
+   assign shamt = {bs, 3'b0};                               // shamt = bs * 8 (convert bytes to bits)
+   assign SboxIn = rs2[shamt +: 8];                         // Shift rs2 right by shamt and take the lower byte
+   aesinvsbox inv_sbox(SboxIn, SboxOut);                    // Apply inverse sbox 
+   assign so = {24'h0, SboxOut};                            // Pad output of inverse substitution box
+   assign sorotate = (so << shamt) | (so >> (32 - shamt));  // Rotate the substitution box output left by shamt (bs * 8)
+   assign DataOut = rs1 ^ sorotate;                         // Set result to "X(rs1)[31..0] ^ rol32(so, unsigned(shamt));"
 endmodule
diff --git a/src/ieu/aes_instructions/aes32dsmi.sv b/src/ieu/aes_instructions/aes32dsmi.sv
index a249e7e93..4f8b5d9c0 100644
--- a/src/ieu/aes_instructions/aes32dsmi.sv
+++ b/src/ieu/aes_instructions/aes32dsmi.sv
@@ -33,7 +33,6 @@ module aes32dsmi(
 );
 
    logic [4:0] 			  shamt;
-   logic [31:0] 		     SboxIn32;
    logic [7:0] 			  SboxIn;
    logic [7:0] 			  SboxOut;
    logic [31:0] 		     so;
@@ -44,8 +43,7 @@ module aes32dsmi(
    assign shamt = {bs, 3'b0};
    
    // Shift rs2 right by shamt and take the lower byte
-   assign SboxIn32 = (rs2 >> shamt);
-   assign SboxIn = SboxIn32[7:0];
+   assign SboxIn = rs2[shamt +: 8];                         // Shift rs2 right by shamt and take the lower byte
    
    // Apply inverse sbox to si
    aesinvsbox inv_sbox(SboxIn, SboxOut);
diff --git a/src/ieu/aes_instructions/aes32esi.sv b/src/ieu/aes_instructions/aes32esi.sv
index c8a6dd0b4..d18fbcfa1 100644
--- a/src/ieu/aes_instructions/aes32esi.sv
+++ b/src/ieu/aes_instructions/aes32esi.sv
@@ -33,7 +33,6 @@ module aes32esi(
 );                
                 
    logic [4:0] 			 shamt;
-   logic [31:0] 		    SboxIn32;
    logic [7:0] 			 SboxIn;
    logic [7:0] 			 SboxOut;
    logic [31:0] 		    so;
@@ -42,11 +41,7 @@ module aes32esi(
    // Shift bs by 3 to get shamt
    assign shamt = {bs, 3'b0};
    
-   // Shift rs2 right by shamt to get sbox input
-   assign SboxIn32 = (rs2 >> shamt);
-   
-   // Take the bottom byte as an input to the substitution box
-   assign SboxIn = SboxIn32[7:0];
+   assign SboxIn = rs2[shamt +: 8];                         // Shift rs2 right by shamt and take the lower byte
    
    // Substitute
    aessbox subbox(SboxIn, SboxOut);
diff --git a/src/ieu/aes_instructions/aes32esmi.sv b/src/ieu/aes_instructions/aes32esmi.sv
index b00495f48..558824e7c 100644
--- a/src/ieu/aes_instructions/aes32esmi.sv
+++ b/src/ieu/aes_instructions/aes32esmi.sv
@@ -33,7 +33,6 @@ module aes32esmi(
 );                
                 
    logic [4:0] 			  shamt;
-   logic [31:0] 		     SboxIn32;
    logic [7:0] 			  SboxIn;
    logic [7:0] 			  SboxOut;
    logic [31:0] 		     so;
@@ -43,12 +42,8 @@ module aes32esmi(
    // Shift bs by 3 to get shamt
    assign shamt = {bs, 3'b0};
    
-   // Shift rs2 right by shamt to get sbox input
-   assign SboxIn32 = (rs2 >> shamt);
-   
-   // Take the bottom byte as an input to the substitution box
-   assign SboxIn = SboxIn32[7:0];
-   
+   assign SboxIn = rs2[shamt +: 8];                         // Shift rs2 right by shamt and take the lower byte
+  
    // Substitute
    aessbox sbox(SboxIn, SboxOut);
    
diff --git a/src/ieu/aes_instructions/aes64ks2.sv b/src/ieu/aes_instructions/aes64ks2.sv
index 0d7a84b32..64484e06e 100644
--- a/src/ieu/aes_instructions/aes64ks2.sv
+++ b/src/ieu/aes_instructions/aes64ks2.sv
@@ -34,6 +34,6 @@ module aes64ks2(
    logic [31:0] 		    w0, w1;
    
    assign w0 = rs1[63:32] ^ rs2[31:0];
-   assign w1 = rs1[63:32] ^ rs2[31:0] ^ rs2[63:32];   
+   assign w1 = w0 ^ rs2[63:32];   
    assign rd = {w1, w0};   
 endmodule