Add tests for unit masked loads/stores.

Change-Id: I6d89d53f24e67073a40567c5852bc0ea97192e63
diff --git a/kelvin_test_utils/BUILD b/kelvin_test_utils/BUILD
index 06b6f25..82eabb6 100644
--- a/kelvin_test_utils/BUILD
+++ b/kelvin_test_utils/BUILD
@@ -25,6 +25,15 @@
 )
 
 py_library(
+    name = "rvv_type_util",
+    srcs = ["rvv_type_util.py"],
+    deps = [
+        requirement("numpy"),
+    ],
+    visibility = ["//visibility:public"],
+)
+
+py_library(
     name = "spi_master",
     srcs = ["spi_master.py"],
     deps = [
diff --git a/kelvin_test_utils/rvv_type_util.py b/kelvin_test_utils/rvv_type_util.py
new file mode 100644
index 0000000..6d47598
--- /dev/null
+++ b/kelvin_test_utils/rvv_type_util.py
@@ -0,0 +1,26 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+# Maps
+DTYPE_TO_SEW = {
+    np.uint8: 0b000,
+    np.uint16: 0b001,
+    np.uint32: 0b010,
+}
+
+def construct_vtype(ma, ta, sew, lmul):
+    """Constructs the vtype register."""
+    return (ma << 7) | (ta << 6) | (sew << 3) | lmul
\ No newline at end of file
diff --git a/tests/cocotb/BUILD b/tests/cocotb/BUILD
index 3cadbae..c319045 100644
--- a/tests/cocotb/BUILD
+++ b/tests/cocotb/BUILD
@@ -173,6 +173,7 @@
 # BEGIN_TESTCASES_FOR_rvv_load_store_test
 RVV_LOAD_STORE_TESTCASES = [
     "load_store_bits",
+    "load_unit_masked",
     "load8_index8",
     "load8_index8_seg",
     "load8_index16",
@@ -198,6 +199,7 @@
     "load_store32_unit_m2",
     "load8_segment2_stride6_m1",
     "load16_segment2_stride6_m1",
+    "store_unit_masked",
     "store8_index8",
     "store16_index8",
     "store16_index16",
@@ -388,6 +390,7 @@
         "seed": "42",
         "test_module": ["rvv_load_store_test.py"],
         "deps": [
+            "//kelvin_test_utils:rvv_type_util",
             "//kelvin_test_utils:sim_test_fixture",
             "@bazel_tools//tools/python/runfiles",
             requirement("tqdm"),
diff --git a/tests/cocotb/rvv/load_store/BUILD b/tests/cocotb/rvv/load_store/BUILD
index 1de8c1b..4ab7ca0 100644
--- a/tests/cocotb/rvv/load_store/BUILD
+++ b/tests/cocotb/rvv/load_store/BUILD
@@ -23,6 +23,9 @@
         "load_store_bits": {
             "srcs": ["load_store_bits.cc"],
         },
+        "load_unit_masked": {
+            "srcs": ["load_unit_masked.cc"],
+        },
         "load8_index8": {
             "srcs": ["load8_index8.cc"],
         },
@@ -128,6 +131,9 @@
         "store32_seg_unit": {
             "srcs": ["store32_seg_unit.cc"],
         },
+        "store_unit_masked": {
+            "srcs": ["store_unit_masked.cc"],
+        },
     },
 )
 
@@ -135,6 +141,7 @@
     name = "rvv_load_store_tests",
     srcs = [
         ":load_store_bits.elf",
+        ":load_unit_masked.elf",
         ":load8_index8.elf",
         ":load8_index8_seg.elf",
         ":load8_index16.elf",
@@ -170,5 +177,6 @@
         ":store8_seg_unit",
         ":store16_seg_unit",
         ":store32_seg_unit",
+        ":store_unit_masked.elf",
     ],
 )
\ No newline at end of file
diff --git a/tests/cocotb/rvv/load_store/load_unit_masked.cc b/tests/cocotb/rvv/load_store/load_unit_masked.cc
new file mode 100644
index 0000000..f6d2749
--- /dev/null
+++ b/tests/cocotb/rvv/load_store/load_unit_masked.cc
@@ -0,0 +1,83 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <riscv_vector.h>
+#include <stdint.h>
+
+size_t vl __attribute__((section(".data"))) = 16;
+size_t vtype __attribute__((section(".data"))) = 0;
+size_t load_filler __attribute__((section(".data"))) = 0;
+uint8_t mask_data[16] __attribute__((section(".data")));
+uint8_t load_data[128] __attribute__((section(".data")));
+// Allow load_addr to be overwritten to exercise AXI memory
+uint8_t* load_addr __attribute__((section(".data"))) = load_data;
+uint8_t store_data[128] __attribute__((section(".data")));
+
+extern "C" {
+__attribute__((used, retain)) void test_unit_load8() {
+  asm("vsetvl zero, %[vl], %[vtype];"
+      "vlm.v v0, %[mask_data];"
+      "vmv.v.x v8, %[load_filler];"
+      "vle8.v v8, %[load_addr], v0.t;"
+      "vse8.v v8, %[store_data];"
+      : [store_data] "=m"(store_data)
+      : [vl] "r"(vl),
+        [vtype] "r"(vtype),
+        [load_filler] "r"(load_filler),
+        [mask_data] "m"(mask_data),
+        [load_addr] "m"(*reinterpret_cast<const uint8_t(*)[128]>(load_addr))
+      : "v0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15",
+        "vl", "vtype");
+}
+
+__attribute__((used, retain)) void test_unit_load16() {
+  asm("vsetvl zero, %[vl], %[vtype];"
+      "vlm.v v0, %[mask_data];"
+      "vmv.v.x v8, %[load_filler];"
+      "vle16.v v8, %[load_addr], v0.t;"
+      "vse16.v v8, %[store_data];"
+      : [store_data] "=m"(store_data)
+      : [vl] "r"(vl),
+        [vtype] "r"(vtype),
+        [load_filler] "r"(load_filler),
+        [mask_data] "m"(mask_data),
+        [load_addr] "m"(*reinterpret_cast<const uint16_t(*)[64]>(load_addr))
+      : "v0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15",
+        "vl", "vtype");
+}
+
+__attribute__((used, retain)) void test_unit_load32() {
+  asm("vsetvl zero, %[vl], %[vtype];"
+      "vlm.v v0, %[mask_data];"
+      "vmv.v.x v8, %[load_filler];"
+      "vle32.v v8, %[load_addr], v0.t;"
+      "vse32.v v8, %[store_data];"
+      : [store_data] "=m"(store_data)
+      : [vl] "r"(vl),
+        [vtype] "r"(vtype),
+        [load_filler] "r"(load_filler),
+        [mask_data] "m"(mask_data),
+        [load_addr] "m"(*reinterpret_cast<const uint32_t(*)[32]>(load_addr))
+      : "v0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15",
+        "vl", "vtype");
+}
+}
+
+void (*impl)() __attribute__((section(".data"))) = &test_unit_load8;
+
+int main(int argc, char** argv) {
+  impl();
+
+  return 0;
+}
diff --git a/tests/cocotb/rvv/load_store/store_unit_masked.cc b/tests/cocotb/rvv/load_store/store_unit_masked.cc
new file mode 100644
index 0000000..1043ed5
--- /dev/null
+++ b/tests/cocotb/rvv/load_store/store_unit_masked.cc
@@ -0,0 +1,76 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <riscv_vector.h>
+#include <stdint.h>
+
+size_t vl __attribute__((section(".data"))) = 16;
+size_t vtype __attribute__((section(".data"))) = 0;
+uint8_t mask_data[16] __attribute__((section(".data")));
+uint8_t load_data[128] __attribute__((section(".data")));
+uint8_t store_data[128] __attribute__((section(".data")));
+// Allow store_addr to be overwritten to exercise AXI memory
+uint8_t* store_addr __attribute__((section(".data"))) = store_data;
+
+extern "C" {
+__attribute__((used, retain)) void test_unit_store8() {
+  asm("vsetvl zero, %[vl], %[vtype];"
+      "vlm.v v0, %[mask_data];"
+      "vle8.v v8, %[load_data];"
+      "vse8.v v8, %[store_addr], v0.t;"
+      : [store_addr] "=m"(*reinterpret_cast<uint8_t(*)[128]>(store_addr))
+      : [vl] "r"(vl),
+        [vtype] "r"(vtype),
+        [mask_data] "m"(mask_data),
+        [load_data] "m"(load_data)
+      : "v0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15",
+        "vl", "vtype");
+}
+
+__attribute__((used, retain)) void test_unit_store16() {
+  asm("vsetvl zero, %[vl], %[vtype];"
+      "vlm.v v0, %[mask_data];"
+      "vle16.v v8, %[load_data];"
+      "vse16.v v8, %[store_addr], v0.t;"
+      : [store_addr] "=m"(*reinterpret_cast<uint16_t(*)[64]>(store_addr))
+      : [vl] "r"(vl),
+        [vtype] "r"(vtype),
+        [mask_data] "m"(mask_data),
+        [load_data] "m"(load_data)
+      : "v0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15",
+        "vl", "vtype");
+}
+
+__attribute__((used, retain)) void test_unit_store32() {
+  asm("vsetvl zero, %[vl], %[vtype];"
+      "vlm.v v0, %[mask_data];"
+      "vle32.v v8, %[load_data];"
+      "vse32.v v8, %[store_addr], v0.t;"
+      : [store_addr] "=m"(*reinterpret_cast<uint32_t(*)[32]>(store_addr))
+      : [vl] "r"(vl),
+        [vtype] "r"(vtype),
+        [mask_data] "m"(mask_data),
+        [load_data] "m"(load_data)
+      : "v0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15",
+        "vl", "vtype");
+}
+}
+
+void (*impl)() __attribute__((section(".data"))) = &test_unit_store8;
+
+int main(int argc, char** argv) {
+  impl();
+
+  return 0;
+}
diff --git a/tests/cocotb/rvv_load_store_test.py b/tests/cocotb/rvv_load_store_test.py
index 9d12c63..b6f3b86 100644
--- a/tests/cocotb/rvv_load_store_test.py
+++ b/tests/cocotb/rvv_load_store_test.py
@@ -13,10 +13,12 @@
 # limitations under the License.
 
 import cocotb
+import itertools
 import numpy as np
 import tqdm
 
 from bazel_tools.tools.python.runfiles import runfiles
+from kelvin_test_utils.rvv_type_util import construct_vtype, DTYPE_TO_SEW
 from kelvin_test_utils.sim_test_fixture import Fixture
 
 
@@ -373,6 +375,193 @@
         })
         assert (actual_output == expected_output).all(), debug_msg
 
+@cocotb.test()
+async def load_unit_masked(dut):
+    """Test masked unit stores."""
+
+    fixture = await Fixture.Create(dut)
+    r = runfiles.Create()
+
+    await fixture.load_elf_and_lookup_symbols(
+        r.Rlocation(
+            'kelvin_hw/tests/cocotb/rvv/load_store/load_unit_masked.elf'),
+        [ "impl", "vtype", "load_data", "load_addr", "vl", "load_filler",
+          "mask_data", "store_data", "test_unit_load8", "test_unit_load16",
+          "test_unit_load32"],
+    )
+
+    cases = [
+        (np.uint8, 0b110, 4), # SEW8, mf4
+        (np.uint8, 0b110, 3), # SEW8, mf4
+        (np.uint8, 0b111, 8), # SEW8, mf2
+        (np.uint8, 0b111, 7), # SEW8, mf2
+        (np.uint8, 0b000, 16), # SEW8, m1
+        (np.uint8, 0b000, 15), # SEW8, m1
+        (np.uint8, 0b001, 32), # SEW8, m2
+        (np.uint8, 0b001, 31), # SEW8, m2
+        (np.uint8, 0b010, 64), # SEW8, m4
+        (np.uint8, 0b010, 63), # SEW8, m4
+        (np.uint8, 0b011, 128), # SEW8, m8
+        (np.uint8, 0b011, 127), # SEW8, m8
+        (np.uint16, 0b111, 4), # SEW16, mf2
+        (np.uint16, 0b111, 3), # SEW16, mf2
+        (np.uint16, 0b000, 8), # SEW16, m1
+        (np.uint16, 0b000, 7), # SEW16, m1
+        (np.uint16, 0b001, 16), # SEW16, m2
+        (np.uint16, 0b001, 15), # SEW16, m2
+        (np.uint16, 0b010, 32), # SEW16, m4
+        (np.uint16, 0b010, 31), # SEW16, m4
+        (np.uint16, 0b011, 64), # SEW16, m8
+        (np.uint16, 0b011, 63), # SEW16, m8
+        (np.uint32, 0b000, 4), # SEW32, m1
+        (np.uint32, 0b000, 3), # SEW32, m1
+        (np.uint32, 0b001, 8), # SEW32, m2
+        (np.uint32, 0b001, 7), # SEW32, m2
+        (np.uint32, 0b010, 16), # SEW32, m4
+        (np.uint32, 0b010, 15), # SEW32, m4
+        (np.uint32, 0b011, 32), # SEW32, m8
+        (np.uint32, 0b011, 31), # SEW32, m8
+    ]
+
+    dtype_to_function = {
+        np.uint8: "test_unit_load8",
+        np.uint16: "test_unit_load16",
+        np.uint32: "test_unit_load32",
+    }
+
+    all_cases = itertools.product([False, True], cases)
+    total_loops = 2 * len(cases)
+
+    rng = np.random.default_rng()
+    for use_axi, (dtype, lmul, vl) in tqdm.tqdm(all_cases, total=total_loops):
+        vtype = construct_vtype(1, 1, DTYPE_TO_SEW[dtype], lmul)
+        mask_bytes = (vl + 7) // 8
+        mask_data = rng.integers(0, 256, mask_bytes, dtype=np.uint8)
+        min_value = np.iinfo(dtype).min
+        max_value = np.iinfo(dtype).max + 1
+        load_data = rng.integers(min_value, max_value, vl, dtype=dtype)
+        load_filler = np.iinfo(dtype).max
+
+        if use_axi:
+            await fixture.write_word(
+                'load_addr', fixture.core_mini_axi.memory_base_addr)
+
+        await fixture.write_ptr('impl', dtype_to_function[dtype])
+        await fixture.write_word('vl', vl)
+        await fixture.write_word('load_filler', load_filler)
+        await fixture.write_word('vtype', vtype)
+        if use_axi:
+            load_data_size = vl * np.dtype(dtype).itemsize
+            fixture.core_mini_axi.memory[0:load_data_size] = \
+                load_data.view(np.uint8)
+        else:
+            await fixture.write('load_data', load_data)
+        await fixture.write('mask_data', mask_data)
+
+        await fixture.run_to_halt()
+
+        actual_output = (await fixture.read(
+            'store_data', vl * np.dtype(dtype).itemsize)).view(dtype)
+
+        mask_bits = np.concat([
+            list(reversed(np.unpackbits(x))) for x in mask_data])
+        for i in range(vl):
+            if mask_bits[i]:
+                assert load_data[i] == actual_output[i]
+            else:
+                assert load_filler == actual_output[i]
+
+@cocotb.test()
+async def store_unit_masked(dut):
+    """Test masked unit stores."""
+    fixture = await Fixture.Create(dut)
+    r = runfiles.Create()
+
+    await fixture.load_elf_and_lookup_symbols(
+        r.Rlocation(
+            'kelvin_hw/tests/cocotb/rvv/load_store/store_unit_masked.elf'),
+        [ "impl", "vtype", "load_data", "vl", "mask_data", "store_data",
+          "store_addr", "test_unit_store8", "test_unit_store16",
+          "test_unit_store32" ],
+    )
+
+    cases = [
+        (np.uint8, 0b110, 4), # SEW8, mf4
+        (np.uint8, 0b110, 3), # SEW8, mf4
+        (np.uint8, 0b111, 8), # SEW8, mf2
+        (np.uint8, 0b111, 7), # SEW8, mf2
+        (np.uint8, 0b000, 16), # SEW8, m1
+        (np.uint8, 0b000, 15), # SEW8, m1
+        (np.uint8, 0b001, 32), # SEW8, m2
+        (np.uint8, 0b001, 31), # SEW8, m2
+        (np.uint8, 0b010, 64), # SEW8, m4
+        (np.uint8, 0b010, 63), # SEW8, m4
+        (np.uint8, 0b011, 128), # SEW8, m8
+        (np.uint8, 0b011, 127), # SEW8, m8
+        (np.uint16, 0b111, 4), # SEW16, mf2
+        (np.uint16, 0b111, 3), # SEW16, mf2
+        (np.uint16, 0b000, 8), # SEW16, m1
+        (np.uint16, 0b000, 7), # SEW16, m1
+        (np.uint16, 0b001, 16), # SEW16, m2
+        (np.uint16, 0b001, 15), # SEW16, m2
+        (np.uint16, 0b010, 32), # SEW16, m4
+        (np.uint16, 0b010, 31), # SEW16, m4
+        (np.uint16, 0b011, 64), # SEW16, m8
+        (np.uint16, 0b011, 63), # SEW16, m8
+        (np.uint32, 0b000, 4), # SEW32, m1
+        (np.uint32, 0b000, 3), # SEW32, m1
+        (np.uint32, 0b001, 8), # SEW32, m2
+        (np.uint32, 0b001, 7), # SEW32, m2
+        (np.uint32, 0b010, 16), # SEW32, m4
+        (np.uint32, 0b010, 15), # SEW32, m4
+        (np.uint32, 0b011, 32), # SEW32, m8
+        (np.uint32, 0b011, 31), # SEW32, m8
+    ]
+
+    dtype_to_function = {
+        np.uint8: "test_unit_store8",
+        np.uint16: "test_unit_store16",
+        np.uint32: "test_unit_store32",
+    }
+
+    all_cases = itertools.product([False, True], cases)
+    total_loops = 2 * len(cases)
+
+    rng = np.random.default_rng()
+    for use_axi, (dtype, lmul, vl) in tqdm.tqdm(all_cases, total=total_loops):
+        vtype = construct_vtype(1, 1, DTYPE_TO_SEW[dtype], lmul)
+        mask_bytes = (vl + 7) // 8
+        mask_data = rng.integers(0, 256, mask_bytes, dtype=np.uint8)
+        min_value = np.iinfo(dtype).min
+        max_value = np.iinfo(dtype).max + 1
+        load_data = rng.integers(min_value, max_value, vl, dtype=dtype)
+
+        if use_axi:
+            await fixture.write_word(
+                'store_addr', fixture.core_mini_axi.memory_base_addr)
+
+        await fixture.write_ptr('impl', dtype_to_function[dtype])
+        await fixture.write_word('vl', vl)
+        await fixture.write_word('vtype', vtype)
+        await fixture.write('load_data', load_data)
+        await fixture.write('mask_data', mask_data)
+
+        await fixture.run_to_halt()
+
+        if use_axi:
+            output_size = vl * np.dtype(dtype).itemsize
+            actual_output = \
+                fixture.core_mini_axi.memory[0:output_size].view(dtype)
+        else:
+            actual_output = (await fixture.read(
+                'store_data', vl * np.dtype(dtype).itemsize)).view(dtype)
+
+        mask_bits = np.concat([
+            list(reversed(np.unpackbits(x))) for x in mask_data])
+        for i in range(vl):
+            if mask_bits[i]:
+                assert load_data[i] == actual_output[i]
+
 
 @cocotb.test()
 async def load8_stride2_m1(dut):