Update SHA-256 VM chip to support SHA-512 and SHA-384

Author: Avaneesh-axiom

extensions/sha2/circuit/Cargo.toml

@@ -1,9 +1,9 @@
 [package]
-name = "openvm-sha256-circuit"
+name = "openvm-sha2-circuit"
 version.workspace = true
 authors.workspace = true
 edition.workspace = true
-description = "OpenVM circuit extension for sha256"
+description = "OpenVM circuit extension for SHA-2"
 
 [dependencies]
 openvm-stark-backend = { workspace = true }
@@ -13,16 +13,16 @@ openvm-circuit-primitives-derive = { workspace = true }
 openvm-circuit-derive = { workspace = true }
 openvm-circuit = { workspace = true }
 openvm-instructions = { workspace = true }
-openvm-sha256-transpiler = { workspace = true }
+openvm-sha2-transpiler = { workspace = true }
 openvm-rv32im-circuit = { workspace = true }
-openvm-sha256-air = { workspace = true }
+openvm-sha2-air = { workspace = true }
 
 derive-new.workspace = true
 derive_more = { workspace = true, features = ["from"] }
 rand.workspace = true
 serde.workspace = true
 sha2 = { version = "0.10", default-features = false }
-strum = { workspace = true }
+ndarray = { workspace = true, default-features = false }
 
 [dev-dependencies]
 openvm-stark-sdk = { workspace = true }
@@ -37,3 +37,6 @@ mimalloc = ["openvm-circuit/mimalloc"]
 jemalloc = ["openvm-circuit/jemalloc"]
 jemalloc-prof = ["openvm-circuit/jemalloc-prof"]
 nightly-features = ["openvm-circuit/nightly-features"]
+
+[package.metadata.cargo-shear]
+ignored = ["ndarray"]

extensions/sha2/circuit/README.md

@@ -1,28 +1,43 @@
-# SHA256 VM Extension
+# SHA-2 VM Extension
 
-This crate contains the circuit for the SHA256 VM extension.
+This crate contains circuits for the SHA-2 family of hash functions.
+We support SHA-256, SHA-512, and SHA-384.
 
-## SHA-256 Algorithm Summary
+## SHA-2 Algorithms Summary
 
-See the [FIPS standard](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf), in particular, section 6.2 for reference.
+The SHA-256, SHA-512, and SHA-384 algorithms are similar in structure.
+We will first describe the SHA-256 algorithm, and then describe the differences between the three algorithms.
+
+See the [FIPS standard](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf) for reference. In particular, sections 6.2, 6.4, and 6.5.
 
 In short the SHA-256 algorithm works as follows.
 1. Pad the message to 512 bits and split it into 512-bit 'blocks'.
-2. Initialize a hash state consisting of eight 32-bit words.
+2. Initialize a hash state consisting of eight 32-bit words to a specific constant value.
 3. For each block, 
-    1. split the message into 16 32-bit words and produce 48 more 'message schedule' words based on them.
-    2. apply 64 'rounds' to update the hash state based on the message schedule.
-    3. add the previous block's final hash state to the current hash state (modulo `2^32`).
+    1. split the message into 16 32-bit words and produce 48 more words based on them. The 16 message words together with the 48 additional words are called the 'message schedule'.
+    2. apply a scrambling function 64 times to the hash state to update it based on the message schedule. We call each update a 'round'.
+    3. add the previous block's final hash state to the current hash state (modulo $2^{32}$).
 4. The output is the final hash state
 
+The differences with the SHA-512 algorithm are that:
+- SHA-512 uses 64-bit words, 1024-bit blocks, performs 80 rounds, and produces a 512-bit output. 
+- all the arithmetic is done modulo $2^{64}$.
+- the initial hash state is different.
+
+The SHA-384 algorithm is a truncation of the SHA-512 output to 384 bits, and the only difference is that the initial hash state is different.
+
 ## Design Overview
 
-This chip produces an AIR that consists of 17 rows for each block (512 bits) in the message, and no more rows.
-The first 16 rows of each block are called 'round rows', and each of them represents four rounds of the SHA-256 algorithm.
-Each row constrains updates to the working variables on each round, and it also constrains the message schedule words based on previous rounds.
-The final row is called a 'digest row' and it produces a final hash for the block, computed as the sum of the working variables and the previous block's final hash.
+We reuse the same AIR code to produce circuits for all three algorithms.
+To achieve this, we parameterize the AIR by constants (such as the word size, number of rounds, and block size) that are specific to each algorithm.
+
+This chip produces an AIR that consists of $R+1$ rows for each block of the message, and no more rows
+(for SHA-256, $R = 16$ and for SHA-512 and SHA-384, $R = 20$).
+The first $R$ rows of each block are called 'round rows', and each of them constrains four rounds of the hash algorithm.
+Each row constrains updates to the working variables on each round, and also constrains the message schedule words based on previous rounds.
+The final row of each block is called a 'digest row' and it produces a final hash for the block, computed as the sum of the working variables and the previous block's final hash.
 
-Note that this chip only supports messages of length less than `2^29` bytes.
+Note that this chip only supports messages of length less than $2^{29}$ bytes.
 
 ### Storing working variables
 
@@ -50,7 +65,7 @@ Since we can reliably constrain values from four rounds ago, we can build up `in
 
 The last block of every message should have the `is_last_block` flag set to `1`.
 Note that `is_last_block` is not constrained to be true for the last block of every message, instead it *defines* what the last block of a message is.
-For instance, if we produce an air with 10 blocks and only the last block has `is_last_block = 1` then the constraints will interpret it as a single message of length 10 blocks.
+For instance, if we produce a trace with 10 blocks and only the last block has `is_last_block = 1` then the constraints will interpret it as a single message of length 10 blocks.
 If, however, we set `is_last_block` to true for the 6th block, the trace will be interpreted as hashing two messages, each of length 5 blocks.
 
 Note that we do constrain, however, that the very last block of the trace has `is_last_block = 1`.
@@ -63,11 +78,11 @@ We use this trick in several places in this chip.
 
 ### Block index counter variables
 
-There are two "block index" counter variables in each row of the air named `global_block_idx` and `local_block_idx`.
-Both of these variables take on the same value on all 17 rows in a block. 
+There are two "block index" counter variables in each row named `global_block_idx` and `local_block_idx`.
+Both of these variables take on the same value on all $R+1$ rows in a block. 
 
 The `global_block_idx` is the index of the block in the entire trace.
-The very first 17 rows in the trace will have `global_block_idx = 1` and the counter will increment by 1 between blocks.  
+The very first block in the trace will have `global_block_idx = 1` on each row and the counter will increment by 1 between blocks.  
 The padding rows will all have `global_block_idx = 0`.
 The `global_block_idx` is used in interaction constraints to constrain the value of `hash` between blocks.
 
@@ -79,15 +94,16 @@ The `local_block_idx` is used to calculate the length of the message processed s
 
 ### VM air vs SubAir
 
-The SHA-256 VM extension chip uses the `Sha256Air` SubAir to help constrain the SHA-256 hash.
-The VM extension air constrains the correctness of the SHA message padding, while the SubAir adds all other constraints related to the hash algorithm.
-The VM extension air also constrains memory reads and writes.
+The SHA-2 VM extension chip uses the `Sha2Air` SubAir to help constrain the appropriate SHA-2 hash algorithm.
+The SubAir is also parameterized by the specific SHA-2 variant's constants.
+The VM extension AIR constrains the correctness of the message padding, while the SubAir adds all other constraints related to the hash algorithm.
+The VM extension AIR also constrains memory reads and writes.
 
 ### A gotcha about padding rows
 
 There are two senses of the word padding used in the context of this chip and this can be confusing.
-First, we use padding to refer to the extra bits added to the message that is input to the SHA-256 algorithm in order to make the input's length a multiple of 512 bits.
-So, we may use the term 'padding rows' to refer to round rows that correspond to the padded bits of a message (as in `Sha256VmAir::eval_padding_row`).
+First, we use padding to refer to the extra bits added to the message that is input to the hash algorithm in order to make the input's length a multiple of the block size.
+So, we may use the term 'padding rows' to refer to round rows that correspond to the padded bits of a message (as in `Sha2VmAir::eval_padding_row`).
 Second, the dummy rows that are added to the trace to make the trace height a power of 2 are also called padding rows (see the `is_padding_row` flag).
 In the SubAir, padding row probably means dummy row.
-In the VM air, it probably refers to SHA-256 padding.
+In the VM air, it probably refers to the message padding.

extensions/sha2/circuit/src/extension.rs

@@ -13,15 +13,15 @@ use openvm_rv32im_circuit::{
     Rv32I, Rv32IExecutor, Rv32IPeriphery, Rv32Io, Rv32IoExecutor, Rv32IoPeriphery, Rv32M,
     Rv32MExecutor, Rv32MPeriphery,
 };
-use openvm_sha256_transpiler::Rv32Sha256Opcode;
+use openvm_sha2_air::{Sha256Config, Sha384Config, Sha512Config};
+use openvm_sha2_transpiler::Rv32Sha2Opcode;
 use openvm_stark_backend::p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
-use strum::IntoEnumIterator;
 
 use crate::*;
 
 #[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)]
-pub struct Sha256Rv32Config {
+pub struct Sha2Rv32Config {
     #[system]
     pub system: SystemConfig,
     #[extension]
@@ -31,38 +31,40 @@ pub struct Sha256Rv32Config {
     #[extension]
     pub io: Rv32Io,
     #[extension]
-    pub sha256: Sha256,
+    pub sha2: Sha2,
 }
 
-impl Default for Sha256Rv32Config {
+impl Default for Sha2Rv32Config {
     fn default() -> Self {
         Self {
             system: SystemConfig::default().with_continuations(),
             rv32i: Rv32I,
             rv32m: Rv32M::default(),
             io: Rv32Io,
-            sha256: Sha256,
+            sha2: Sha2,
         }
     }
 }
 
 #[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
-pub struct Sha256;
+pub struct Sha2;
 
 #[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)]
-pub enum Sha256Executor<F: PrimeField32> {
-    Sha256(Sha256VmChip<F>),
+pub enum Sha2Executor<F: PrimeField32> {
+    Sha256(Sha2VmChip<F, Sha256Config>),
+    Sha512(Sha2VmChip<F, Sha512Config>),
+    Sha384(Sha2VmChip<F, Sha384Config>),
 }
 
 #[derive(From, ChipUsageGetter, Chip, AnyEnum)]
-pub enum Sha256Periphery<F: PrimeField32> {
+pub enum Sha2Periphery<F: PrimeField32> {
     BitwiseOperationLookup(SharedBitwiseOperationLookupChip<8>),
     Phantom(PhantomChip<F>),
 }
 
-impl<F: PrimeField32> VmExtension<F> for Sha256 {
-    type Executor = Sha256Executor<F>;
-    type Periphery = Sha256Periphery<F>;
+impl<F: PrimeField32> VmExtension<F> for Sha2 {
+    type Executor = Sha2Executor<F>;
+    type Periphery = Sha2Periphery<F>;
 
     fn build(
         &self,
@@ -81,18 +83,32 @@ impl<F: PrimeField32> VmExtension<F> for Sha256 {
             chip
         };
 
-        let sha256_chip = Sha256VmChip::new(
+        let sha256_chip = Sha2VmChip::<F, Sha256Config>::new(
+            builder.system_port(),
+            builder.system_config().memory_config.pointer_max_bits,
+            bitwise_lu_chip.clone(),
+            builder.new_bus_idx(),
+            builder.system_base().offline_memory(),
+        );
+        inventory.add_executor(sha256_chip, vec![Rv32Sha2Opcode::SHA256.global_opcode()])?;
+
+        let sha512_chip = Sha2VmChip::<F, Sha512Config>::new(
+            builder.system_port(),
+            builder.system_config().memory_config.pointer_max_bits,
+            bitwise_lu_chip.clone(),
+            builder.new_bus_idx(),
+            builder.system_base().offline_memory(),
+        );
+        inventory.add_executor(sha512_chip, vec![Rv32Sha2Opcode::SHA512.global_opcode()])?;
+
+        let sha384_chip = Sha2VmChip::<F, Sha384Config>::new(
             builder.system_port(),
             builder.system_config().memory_config.pointer_max_bits,
             bitwise_lu_chip,
             builder.new_bus_idx(),
-            Rv32Sha256Opcode::CLASS_OFFSET,
             builder.system_base().offline_memory(),
         );
-        inventory.add_executor(
-            sha256_chip,
-            Rv32Sha256Opcode::iter().map(|x| x.global_opcode()),
-        )?;
+        inventory.add_executor(sha384_chip, vec![Rv32Sha2Opcode::SHA384.global_opcode()])?;
 
         Ok(inventory)
     }

extensions/sha2/circuit/src/lib.rs

@@ -1,5 +1,5 @@
-mod sha256_chip;
-pub use sha256_chip::*;
+mod sha2_chip;
+pub use sha2_chip::*;
 
 mod extension;
 pub use extension::*;