Add masked SRAM write support

Author: AnZhong24

python/assassyn/ir/memory/base.md

@@ -33,6 +33,7 @@ Base class for memory modules that provides common functionality for SRAM and DR
 - `re: Value` - Read enable signal (combinational input)  
 - `addr: Value` - Address signal (combinational input)
 - `wdata: Value` - Write data signal (combinational input)
+- `wmask: Value | None` - Optional per-bit write mask for partial writes
 - `addr_width: int` - Width of the address in bits (derived as log2(depth))
 - `_payload: Array` - Array holding the memory contents (private, not for direct access, owned by the memory instance)
 
@@ -48,7 +49,7 @@ Initialize memory base class with validation and setup.
 **Returns:** None
 
 **Explanation:**
-This constructor validates all input parameters and sets up the memory module infrastructure. It enforces that depth must be a power of 2 to enable efficient address decoding using log2 operations. The constructor creates a `RegArray` instance with the specified width and depth to serve as the memory payload, using the instance name for proper identification in generated code. All signal attributes are initialized to None and will be assigned during the `build()` method of concrete implementations.
+This constructor validates all input parameters and sets up the memory module infrastructure. It enforces that depth must be a power of 2 to enable efficient address decoding using log2 operations. The constructor creates a `RegArray` instance with the specified width and depth to serve as the memory payload, using the instance name for proper identification in generated code. All signal attributes are initialized to None and will be assigned during the `build()` method of concrete implementations. `wmask` defaults to `None`, so existing memories keep the original full-word write behavior unless masked writes are explicitly requested.
 
 **Address Width Derivation Logic:** The address width is calculated as `addr_width = log2(depth)` because:
 1. **Power-of-2 Constraint**: Depth must be a power of 2 to enable efficient address decoding

python/assassyn/ir/memory/base.py

@@ -24,6 +24,7 @@ class MemoryBase(Downstream):
     re: Value       # Read enable signal
     addr: Value     # Address signal
     wdata: Value    # Write data signal
+    wmask: Value | None  # Optional write mask signal
 
     # Derived Values
     addr_width: int # Width of the address in bits
@@ -70,3 +71,4 @@ def __init__(self, width: int, depth: int, init_file: str | None):
         self.re = None
         self.addr = None
         self.wdata = None
+        self.wmask = None

python/assassyn/ir/memory/sram.md

@@ -44,7 +44,7 @@ Initialize SRAM module with read data buffer.
 **Explanation:**
 This constructor calls the parent `MemoryBase.__init__()` method to inherit base memory functionality, then creates an additional `dout` register buffer. Both arrays record the SRAM instance as their owner: `_payload` is created in the base class with `owner=self`, and `dout` uses the same override. Downstream passes rely on `Array.is_payload(SRAM)` to distinguish the payload buffer from auxiliary registers. Using `Bits` type ensures compatibility with array read operations that return raw bit values.
 
-### `def build(self, we, re, addr, wdata)`
+### `def build(self, we, re, addr, wdata, wmask=None)`
 
 Build the SRAM module with combinational logic for synchronous memory operations.
 
@@ -53,6 +53,7 @@ Build the SRAM module with combinational logic for synchronous memory operations
 - `re: Value` - Read enable signal
 - `addr: Value` - Address signal  
 - `wdata: Value` - Write data signal
+- `wmask: Value | None` - Optional per-bit write mask. `None` preserves full-word writes.
 
 **Returns:** None
 
@@ -61,7 +62,7 @@ This method implements the core SRAM functionality using the `@combinational` de
 
 1. **Signal Assignment:** All input signals are stored as instance attributes for memory operations
 2. **Mutual Exclusion:** Uses `assume(~(we & re))` from [intrinsic.py](../expr/intrinsic.md) to enforce that read and write operations cannot be enabled simultaneously
-3. **Write Operation:** When `we` is enabled, writes `wdata` to `_payload[addr]` using conditional execution
+3. **Write Operation:** When `we` is enabled, either writes `wdata` directly to `_payload[addr]` or merges it with the previous word using `wmask`
 4. **Read Operation:** When `re` is enabled, reads `_payload[addr]` and stores the result in `dout[0]` for downstream modules to access
 
 **SRAM Read Data Timing:** The relationship between read enable timing and `dout` buffer update:
@@ -73,6 +74,7 @@ This method implements the core SRAM functionality using the `@combinational` de
 **Technical Details:**
 - Uses `Condition` blocks for conditional execution of read/write operations
 - Enforces mutual exclusion between read and write operations using `assume` intrinsic
+- Supports optional masked writes with `merged = (old_data & ~wmask) | (wdata & wmask)`
 - Provides immediate data access without request/response cycles
 - Read data is buffered in `dout` register for downstream module consumption
 - Follows the combinational downstream module pattern for same-cycle signal processing

python/assassyn/ir/memory/sram.py

@@ -34,25 +34,32 @@ def __init__(self, width: int, depth: int, init_file: str | None):
         )
 
     @combinational
-    def build(self, we, re, addr, wdata):  # pylint: disable=too-many-arguments
+    def build(self, we, re, addr, wdata, wmask=None):  # pylint: disable=too-many-arguments
         '''The constructor for the SRAM module.
 
         Args:
             we: Value: The write enable signal.
             re: Value: The read enable signal.
             addr: Value: The address signal.
             wdata: Value: The write data signal.
+            wmask: Optional per-bit write mask. `None` keeps full-word writes.
         '''
         self.we = we
         self.re = re
         self.addr = addr
         self.wdata = wdata
+        self.wmask = wmask
 
         # Enforce that we and re cannot be both enabled
         assume(~(we & re))
 
         with Condition(we):
-            self._payload[addr] = wdata
+            if wmask is None:
+                self._payload[addr] = wdata
+            else:
+                old_data = self._payload[addr]
+                merged = (old_data & (~self.wmask)) | (wdata & self.wmask)
+                self._payload[addr] = merged
         with Condition(re):
             self.dout[0] = self._payload[addr]
 

python/ci-tests/README.md

@@ -15,6 +15,7 @@
 | `test_select, test_select1hot`                               | gramme : `select`                      |
 | `test_testbench`                                             | Usage of `with Cycle(1):`              |
 | `test_explict_pop, test_peek`                                | gramme in `Port`                       |
+| `test_sram, test_sram_masked`                                | SRAM full-word access and masked write |
 |                                                              |                                        |
 | `test_fifo1, test_bind, `<br>`test_eager_bind, test_imbalance, `<br>`test_fifo_valid, test_wait_until` | sth about **Pure Sequential Logic**    |
 | `test_comb_expose, test_toposort`<br />`test_downstream, `   | sth about **Pure Combinational Logic** |
@@ -63,4 +64,7 @@
 14. `test_explict_pop`
     + An alternative method for reading port data.
 15. `test_peek`
-    + Similar to the operation of viewing the top of a queue in a `Queue`. It corresponds to the `front()` operation in the STL of C++ queues. Essentially, it is looking at the top element of the queue without removing it.
+    + Similar to the operation of viewing the top of a queue in a `Queue`. It corresponds to the `front()` operation in the STL of C++ queues. Essentially, it is looking at the top element of the queue without removing it.
+16. `test_sram, test_sram_masked`
+    + `test_sram` covers the original full-word SRAM read/write behavior.
+    + `test_sram_masked` adds regression coverage for masked writes, including zero mask, each byte lane, both halfword lanes, a non-byte-aligned bit mask, and cross-address isolation.

python/ci-tests/test_sram_masked.py

@@ -0,0 +1,161 @@
+"""CI test for SRAM masked-write support."""
+
+import re
+
+from assassyn.frontend import *
+from assassyn.test import run_test
+
+
+ADDR0 = 0x12
+ADDR1 = 0x34
+
+OPS = [
+    {"kind": "write", "addr": ADDR0, "wdata": 0x11223344, "wmask": 0xFFFFFFFF},
+    {"kind": "read", "addr": ADDR0, "expected": 0x11223344},
+    {"kind": "write", "addr": ADDR0, "wdata": 0xFFFFFFFF, "wmask": 0x00000000},
+    {"kind": "read", "addr": ADDR0, "expected": 0x11223344},
+    {"kind": "write", "addr": ADDR0, "wdata": 0x000000AA, "wmask": 0x000000FF},
+    {"kind": "read", "addr": ADDR0, "expected": 0x112233AA},
+    {"kind": "write", "addr": ADDR0, "wdata": 0x0000BB00, "wmask": 0x0000FF00},
+    {"kind": "read", "addr": ADDR0, "expected": 0x1122BBAA},
+    {"kind": "write", "addr": ADDR0, "wdata": 0x00CC0000, "wmask": 0x00FF0000},
+    {"kind": "read", "addr": ADDR0, "expected": 0x11CCBBAA},
+    {"kind": "write", "addr": ADDR0, "wdata": 0xDD000000, "wmask": 0xFF000000},
+    {"kind": "read", "addr": ADDR0, "expected": 0xDDCCBBAA},
+    {"kind": "write", "addr": ADDR0, "wdata": 0x0000EEFF, "wmask": 0x0000FFFF},
+    {"kind": "read", "addr": ADDR0, "expected": 0xDDCCEEFF},
+    {"kind": "write", "addr": ADDR0, "wdata": 0xA1B20000, "wmask": 0xFFFF0000},
+    {"kind": "read", "addr": ADDR0, "expected": 0xA1B2EEFF},
+    {"kind": "write", "addr": ADDR1, "wdata": 0x55667788, "wmask": 0xFFFFFFFF},
+    {"kind": "read", "addr": ADDR1, "expected": 0x55667788},
+    {"kind": "write", "addr": ADDR1, "wdata": 0x00990000, "wmask": 0x00FF0000},
+    {"kind": "read", "addr": ADDR1, "expected": 0x55997788},
+    {"kind": "write", "addr": ADDR1, "wdata": 0xAA5500CC, "wmask": 0x0F0F00F0},
+    {"kind": "read", "addr": ADDR1, "expected": 0x5A9577C8},
+    {"kind": "read", "addr": ADDR0, "expected": 0xA1B2EEFF},
+]
+
+
+class ReadObserver(Module):
+
+    def __init__(self):
+        super().__init__(
+            ports={
+                "step": Port(Bits(8)),
+                "addr": Port(Bits(9)),
+            }
+        )
+
+    @module.combinational
+    def build(self, rdata: RegArray):
+        step, addr = self.pop_all_ports(True)
+        log(
+            "masked_read step={} addr=0x{:03x} data=0x{:08x}",
+            step,
+            addr,
+            rdata[0].bitcast(Bits(32)),
+        )
+
+
+class Launcher(Module):
+
+    def __init__(self, target):
+        super().__init__(ports={})
+        self.target = target
+
+    @module.combinational
+    def build(self):
+        self.target.async_called()
+
+
+class MaskedDriver(Module):
+
+    def __init__(self, observer):
+        super().__init__(ports={})
+        self.name = "Driver"
+        self.observer = observer
+
+    @module.combinational
+    def build(self):
+        phase_bits = max(1, (len(OPS) + 1).bit_length())
+        phase = RegArray(UInt(phase_bits), 1, initializer=[0])
+        state = phase[0]
+        next_state = state + UInt(phase_bits)(1)
+        (phase & self)[0] <= next_state
+
+        we = Bits(1)(0)
+        re = Bits(1)(0)
+        addr = Bits(9)(0)
+        wdata = Bits(32)(0)
+        wmask = Bits(32)(0)
+
+        for idx, op in enumerate(OPS):
+            is_step = state == UInt(phase_bits)(idx)
+            addr_bits = Bits(9)(op["addr"])
+            addr = is_step.select(addr_bits, addr)
+
+            if op["kind"] == "write":
+                we = is_step.select(Bits(1)(1), we)
+                wdata = is_step.select(Bits(32)(op["wdata"]), wdata)
+                wmask = is_step.select(Bits(32)(op["wmask"]), wmask)
+            else:
+                re = is_step.select(Bits(1)(1), re)
+                with Condition(is_step):
+                    self.observer.async_called(step=Bits(8)(idx), addr=addr_bits)
+
+        sram = SRAM(32, 512, None)
+        sram.build(we, re, addr, wdata, wmask)
+
+        with Condition(state == UInt(phase_bits)(len(OPS))):
+            finish()
+
+        return sram
+
+
+READ_RE = re.compile(
+    r"masked_read step=(\d+) addr=0x([0-9a-fA-F]+) data=0x([0-9a-fA-F]+)"
+)
+
+
+def check(raw):
+    expected_reads = [
+        (idx, op["addr"], op["expected"])
+        for idx, op in enumerate(OPS)
+        if op["kind"] == "read"
+    ]
+
+    actual_reads = []
+    for line in raw.splitlines():
+        if "[readobserver" not in line.lower():
+            continue
+        match = READ_RE.search(line)
+        assert match is not None, f"Unexpected ReadObserver log line: {line}"
+        actual_reads.append(
+            (
+                int(match.group(1)),
+                int(match.group(2), 16),
+                int(match.group(3), 16),
+            )
+        )
+
+    assert actual_reads == expected_reads, (
+        f"Masked SRAM reads mismatch.\n"
+        f"expected={expected_reads}\n"
+        f"actual={actual_reads}"
+    )
+
+
+def test_sram_masked_write():
+    def top():
+        observer = ReadObserver()
+        driver = MaskedDriver(observer)
+        sram = driver.build()
+        launcher = Launcher(driver)
+        launcher.build()
+        observer.build(sram.dout)
+
+    run_test("sram_masked", top, check, sim_threshold=200, idle_threshold=200)
+
+
+if __name__ == "__main__":
+    test_sram_masked_write()