feat: update memory chiplet to be element-addressable

air/src/constraints/chiplets/memory/tests.rs

@@ -10,7 +10,7 @@ use crate::{
     chiplets::memory,
     trace::{
         chiplets::{
-            memory::{Selectors, MEMORY_COPY_READ, MEMORY_INIT_READ, MEMORY_WRITE},
+            memory::{Selectors, MEMORY_COPY_READ, MEMORY_INIT_READ, MEMORY_WRITE_SELECTOR},
             MEMORY_TRACE_OFFSET,
         },
         TRACE_WIDTH,
@@ -30,7 +30,7 @@ fn test_memory_write() {
 
     // Write to a new context.
     let result = get_constraint_evaluation(
-        MEMORY_WRITE,
+        MEMORY_WRITE_SELECTOR,
         MemoryTestDeltaType::Context,
         &old_values,
         &new_values,
@@ -39,7 +39,7 @@ fn test_memory_write() {
 
     // Write to a new address in the same context.
     let result = get_constraint_evaluation(
-        MEMORY_WRITE,
+        MEMORY_WRITE_SELECTOR,
         MemoryTestDeltaType::Address,
         &old_values,
         &new_values,
@@ -48,7 +48,7 @@ fn test_memory_write() {
 
     // Write to the same context and address at a new clock cycle.
     let result = get_constraint_evaluation(
-        MEMORY_WRITE,
+        MEMORY_WRITE_SELECTOR,
         MemoryTestDeltaType::Clock,
         &old_values,
         &new_values,

air/src/trace/chiplets/memory.rs

@@ -7,6 +7,7 @@ use super::{create_range, Felt, Range, ONE, ZERO};
 /// Number of columns needed to record an execution trace of the memory chiplet.
 pub const TRACE_WIDTH: usize = 15;
 
+// TODO(plafer): get rid of all "selector" constants
 /// Number of selector columns in the trace.
 pub const NUM_SELECTORS: usize = 2;
 
@@ -16,16 +17,27 @@ pub const NUM_SELECTORS: usize = 2;
 /// read / write) is to be applied at a specific row of the memory execution trace.
 pub type Selectors = [Felt; NUM_SELECTORS];
 
-// --- OPERATION SELECTORS ------------------------------------------------------------------------
-
 /// Specifies an operation that initializes new memory and then reads it.
 pub const MEMORY_INIT_READ: Selectors = [ONE, ZERO];
 
 /// Specifies an operation that copies existing memory and then reads it.
 pub const MEMORY_COPY_READ: Selectors = [ONE, ONE];
 
 /// Specifies a memory write operation.
-pub const MEMORY_WRITE: Selectors = [ZERO, ZERO];
+pub const MEMORY_WRITE_SELECTOR: Selectors = [ZERO, ZERO];
+
+// --- OPERATION SELECTORS ------------------------------------------------------------------------
+
+/// Specifies the value of the `READ_WRITE` column when the operation is a write.
+pub const MEMORY_WRITE: Felt = ZERO;
+/// Specifies the value of the `READ_WRITE` column when the operation is a read.
+pub const MEMORY_READ: Felt = ONE;
+/// Specifies the value of the `ELEMENT_OR_WORD` column when the operation is over an element.
+pub const MEMORY_ACCESS_ELEMENT: Felt = ZERO;
+/// Specifies the value of the `ELEMENT_OR_WORD` column when the operation is over a word.
+pub const MEMORY_ACCESS_WORD: Felt = ONE;
+
+// TODO(plafer): figure out the new labels
 
 /// Unique label computed as 1 plus the full chiplet selector with the bits reversed.
 /// mem_read selector=[1, 1, 0, 1], rev(selector)=[1, 0, 1, 1], +1=[1, 1, 0, 0]
@@ -37,15 +49,24 @@ pub const MEMORY_WRITE_LABEL: u8 = 0b0100;
 
 // --- COLUMN ACCESSOR INDICES WITHIN THE CHIPLET -------------------------------------------------
 
+// TODO(plafer): replace all uses with `WORD_SIZE`
 /// The number of elements accessible in one read or write memory access.
 pub const NUM_ELEMENTS: usize = 4;
 
+/// Column to hold the whether the operation is a read or write.
+pub const READ_WRITE_COL_IDX: usize = 0;
+/// Column to hold the whether the operation was over an element or a word.
+pub const ELEMENT_OR_WORD_COL_IDX: usize = READ_WRITE_COL_IDX + 1;
 /// Column to hold the context ID of the current memory context.
-pub const CTX_COL_IDX: usize = NUM_SELECTORS;
+pub const CTX_COL_IDX: usize = ELEMENT_OR_WORD_COL_IDX + 1;
 /// Column to hold the memory address.
 pub const ADDR_COL_IDX: usize = CTX_COL_IDX + 1;
+/// Column to hold the first bit of the index of the address in the batch.
+pub const IDX0_COL_IDX: usize = ADDR_COL_IDX + 1;
+/// Column to hold the second bit of the index of the address in the batch.
+pub const IDX1_COL_IDX: usize = IDX0_COL_IDX + 1;
 /// Column for the clock cycle in which the memory operation occurred.
-pub const CLK_COL_IDX: usize = ADDR_COL_IDX + 1;
+pub const CLK_COL_IDX: usize = IDX1_COL_IDX + 1;
 /// Columns to hold the values stored at a given memory context, address, and clock cycle after
 /// the memory operation. When reading from a new address, these are initialized to zero.
 pub const V_COL_RANGE: Range<usize> = create_range(CLK_COL_IDX + 1, NUM_ELEMENTS);
@@ -58,3 +79,6 @@ pub const D1_COL_IDX: usize = D0_COL_IDX + 1;
 /// Column for the inverse of the delta between two consecutive context IDs, addresses, or clock
 /// cycles, used to enforce that changes are correctly constrained.
 pub const D_INV_COL_IDX: usize = D1_COL_IDX + 1;
+/// Column to hold the flag indicating whether the current memory operation is in the same batch and
+/// same context as the previous operation.
+pub const FLAG_SAME_BATCH_AND_CONTEXT: usize = D_INV_COL_IDX + 1;

processor/src/chiplets/memory/mod.rs

@@ -2,10 +2,14 @@ use alloc::{collections::BTreeMap, vec::Vec};
 
 use miden_air::{
     trace::chiplets::memory::{
-        ADDR_COL_IDX, CLK_COL_IDX, CTX_COL_IDX, D0_COL_IDX, D1_COL_IDX, D_INV_COL_IDX, V_COL_RANGE,
+        ADDR_COL_IDX, CLK_COL_IDX, CTX_COL_IDX, D0_COL_IDX, D1_COL_IDX, D_INV_COL_IDX,
+        ELEMENT_OR_WORD_COL_IDX, FLAG_SAME_BATCH_AND_CONTEXT, IDX0_COL_IDX, IDX1_COL_IDX,
+        MEMORY_ACCESS_ELEMENT, MEMORY_ACCESS_WORD, MEMORY_READ, MEMORY_WRITE, READ_WRITE_COL_IDX,
+        V_COL_RANGE,
     },
     RowIndex,
 };
+use vm_core::{WORD_SIZE, ZERO};
 
 use super::{
     utils::{split_element_u32_into_u16, split_u32_into_u16},
@@ -14,7 +18,7 @@ use super::{
 use crate::{system::ContextId, ExecutionError};
 
 mod segment;
-use segment::MemorySegmentTrace;
+use segment::{MemoryOperation, MemorySegmentTrace};
 
 #[cfg(test)]
 mod tests;
@@ -28,38 +32,39 @@ const INIT_MEM_VALUE: Word = EMPTY_WORD;
 // RANDOM ACCESS MEMORY
 // ================================================================================================
 
-// TODO(plafer): update
+// TODO(plafer): finish documenting (especially for the values of d0/d1 for batch, and f_scb)
 /// Memory controller for the VM.
 ///
 /// This component is responsible for tracking current memory state of the VM, as well as for
 /// building an execution trace of all memory accesses.
 ///
 /// The memory is comprised of one or more segments, each segment accessible from a specific
-/// execution context. The root (kernel) context has context ID 0, and all additional contexts
-/// have increasing IDs. Within each segment, the memory is word-addressable. That is, four field
-/// elements are located at each memory address, and we can read and write elements to/from memory
-/// in batches of four.
+/// execution context. The root (kernel) context has context ID 0, and all additional contexts have
+/// increasing IDs. Within each segment, the memory is element-addressable, even though the trace
+/// tracks words (that we refer to as "batches") for optimization purposes. That is, a single field
+/// element is located at each memory address, and we can read and write elements to/from memory
+/// either individually or in batches of four.
 ///
-/// Memory for a a given address is always initialized to zeros. That is, reading from an address
-/// before writing to it will return four ZERO elements.
+/// Memory for a a given address is always initialized to zero. That is, reading from an address
+/// before writing to it will return ZERO.
 ///
 /// ## Execution trace
 /// The layout of the memory access trace is shown below.
 ///
-///   s0   s1   ctx  addr   clk   v0   v1   v2   v3   d0   d1   d_inv
-/// ├────┴────┴────┴──────┴─────┴────┴────┴────┴────┴────┴────┴───────┤
+///   rw   ew   ctx  batch   idx0   idx1  clk   v0   v1   v2   v3   d0   d1   d_inv   f_scb
+/// ├────┴────┴────┴───────┴──────┴──────┴────┴────┴────┴────┴────┴────┴────┴───────┴───────┤
 ///
 /// In the above, the meaning of the columns is as follows:
-/// - `s0` is a selector column used to identify whether the memory access is a read or a write. A
-///   value of ZERO indicates a write, and ONE indicates a read.
-/// - `s1` is a selector column used to identify whether the memory access is a read of an existing
-///   memory value or not (i.e., this context/addr combination already existed and is being read). A
-///   value of ONE indicates a read of existing memory, meaning the previous value must be copied.
+/// - `rw` is a selector column used to identify whether the memory operation is a read or a write.
+/// - `ew` is a selector column used to identify whether the memory operation is over an element or
+///   a word.
 /// - `ctx` contains execution context ID. Values in this column must increase monotonically but
 ///   there can be gaps between two consecutive context IDs of up to 2^32. Also, two consecutive
 ///   values can be the same.
-/// - `addr` contains memory address. Values in this column must increase monotonically for a given
-///   context but there can be gaps between two consecutive values of up to 2^32. Also, two
+/// - `batch` contains the the index of the batch of addresses, computed as the address of the first
+///   element in the batch divided by 4. For example, the value of `batch` for the batch of
+///   addresses 40, 41, 42, and 43 is 10. Values in this column must increase monotonically for a
+///   given context but there can be gaps between two consecutive values of up to 2^32. Also, two
 ///   consecutive values can be the same.
 /// - `clk` contains clock cycle at which a memory operation happened. Values in this column must
 ///   increase monotonically for a given context and memory address but there can be gaps between
@@ -75,6 +80,8 @@ const INIT_MEM_VALUE: Word = EMPTY_WORD;
 ///     `old_clk` - 1).
 /// - `d_inv` contains the inverse of the delta between two consecutive context IDs, addresses, or
 ///   clock cycles computed as described above.
+/// - `f_scb` is a flag indicating whether the context and the batch are the same as in the previous
+///   row.
 ///
 /// For the first row of the trace, values in `d0`, `d1`, and `d_inv` are set to zeros.
 #[derive(Debug, Default)]
@@ -154,9 +161,7 @@ impl Memory {
             .try_into()
             .map_err(|_| ExecutionError::MemoryAddressOutOfBounds(addr.as_int()))?;
         self.num_trace_rows += 1;
-        let _ = self.trace.entry(ctx).or_default().read(ctx, addr, Felt::from(clk));
-        // TODO(plafer): `MemoryTraceSegment` needs to read Felts
-        todo!()
+        self.trace.entry(ctx).or_default().read(ctx, addr, Felt::from(clk))
     }
 
     /// Returns a word located in memory at the specified context/address.
@@ -178,9 +183,16 @@ impl Memory {
             .as_int()
             .try_into()
             .map_err(|_| ExecutionError::MemoryAddressOutOfBounds(addr.as_int()))?;
-        // TODO(plafer): check for alignment
+        if addr % WORD_SIZE as u32 != 0 {
+            return Err(ExecutionError::MemoryUnalignedWordAccess {
+                addr,
+                ctx,
+                clk: Felt::from(clk),
+            });
+        }
+
         self.num_trace_rows += 1;
-        self.trace.entry(ctx).or_default().read(ctx, addr, Felt::from(clk))
+        self.trace.entry(ctx).or_default().read_word(ctx, addr, Felt::from(clk))
     }
 
     /// Writes the provided field element at the specified context/address.
@@ -200,13 +212,7 @@ impl Memory {
             .try_into()
             .map_err(|_| ExecutionError::MemoryAddressOutOfBounds(addr.as_int()))?;
         self.num_trace_rows += 1;
-        // TODO(plafer): fix
-        self.trace.entry(ctx).or_default().write(
-            ctx,
-            addr,
-            Felt::from(clk),
-            [value, value, value, value],
-        )
+        self.trace.entry(ctx).or_default().write(ctx, addr, Felt::from(clk), value)
     }
 
     /// Writes the provided word at the specified context/address.
@@ -226,9 +232,16 @@ impl Memory {
             .as_int()
             .try_into()
             .map_err(|_| ExecutionError::MemoryAddressOutOfBounds(addr.as_int()))?;
-        // TODO(plafer): check for alignment
+        if addr % WORD_SIZE as u32 != 0 {
+            return Err(ExecutionError::MemoryUnalignedWordAccess {
+                addr,
+                ctx,
+                clk: Felt::from(clk),
+            });
+        }
+
         self.num_trace_rows += 1;
-        self.trace.entry(ctx).or_default().write(ctx, addr, Felt::from(clk), value)
+        self.trace.entry(ctx).or_default().write_word(ctx, addr, Felt::from(clk), value)
     }
 
     // EXECUTION TRACE GENERATION
@@ -301,13 +314,33 @@ impl Memory {
                 let felt_addr = Felt::from(addr);
                 for memory_access in addr_trace {
                     let clk = memory_access.clk();
-                    let value = memory_access.value();
+                    let value = memory_access.word();
 
-                    let selectors = memory_access.op_selectors();
-                    trace.set(row, 0, selectors[0]);
-                    trace.set(row, 1, selectors[1]);
+                    match memory_access.operation() {
+                        MemoryOperation::Read => trace.set(row, READ_WRITE_COL_IDX, MEMORY_READ),
+                        MemoryOperation::Write => trace.set(row, READ_WRITE_COL_IDX, MEMORY_WRITE),
+                    }
+                    let (idx1, idx0) = match memory_access.access_type() {
+                        segment::MemoryAccessType::Element { addr_idx_in_word } => {
+                            trace.set(row, ELEMENT_OR_WORD_COL_IDX, MEMORY_ACCESS_ELEMENT);
+
+                            match addr_idx_in_word {
+                                0 => (ZERO, ZERO),
+                                1 => (ZERO, ONE),
+                                2 => (ONE, ZERO),
+                                3 => (ONE, ONE),
+                                _ => panic!("invalid address index in word: {addr_idx_in_word}"),
+                            }
+                        },
+                        segment::MemoryAccessType::Word => {
+                            trace.set(row, ELEMENT_OR_WORD_COL_IDX, MEMORY_ACCESS_WORD);
+                            (ZERO, ZERO)
+                        },
+                    };
                     trace.set(row, CTX_COL_IDX, ctx);
                     trace.set(row, ADDR_COL_IDX, felt_addr);
+                    trace.set(row, IDX0_COL_IDX, idx0);
+                    trace.set(row, IDX1_COL_IDX, idx1);
                     trace.set(row, CLK_COL_IDX, clk);
                     for (idx, col) in V_COL_RANGE.enumerate() {
                         trace.set(row, col, value[idx]);
@@ -328,6 +361,12 @@ impl Memory {
                     // TODO: switch to batch inversion to improve efficiency.
                     trace.set(row, D_INV_COL_IDX, delta.inv());
 
+                    if prev_ctx == ctx && prev_addr == felt_addr {
+                        trace.set(row, FLAG_SAME_BATCH_AND_CONTEXT, ONE);
+                    } else {
+                        trace.set(row, FLAG_SAME_BATCH_AND_CONTEXT, ZERO);
+                    };
+
                     // update values for the next iteration of the loop
                     prev_ctx = ctx;
                     prev_addr = felt_addr;