Add tests for unit/unit-segmented loads/stores.

Check combinations of SEW and EEW.

Change-Id: I7de787fb11ce4d554c338e069ed8c0f068bb0612
diff --git a/kelvin_test_utils/rvv_type_util.py b/kelvin_test_utils/rvv_type_util.py
index 6d47598..0a57750 100644
--- a/kelvin_test_utils/rvv_type_util.py
+++ b/kelvin_test_utils/rvv_type_util.py
@@ -21,6 +21,45 @@
     np.uint32: 0b010,
 }
 
+SEWS = [
+    0b000,
+    0b001,
+    0b010,
+]
+
+SEW_TO_LMULS_AND_VLMAXS = {
+    0b000: [
+        (0b110, 4),
+        (0b111, 8),
+        (0b000, 16),
+        (0b001, 32),
+        (0b010, 64),
+        (0b011, 128),
+    ],
+    0b001: [
+        (0b111, 4),
+        (0b000, 8),
+        (0b001, 16),
+        (0b010, 32),
+        (0b011, 64),
+    ],
+    0b010: [
+        (0b000, 4),
+        (0b001, 8),
+        (0b010, 16),
+        (0b011, 32),
+    ],
+}
+
+LMUL_TO_EMUL = {
+    0b110: 1,
+    0b111: 1,
+    0b000: 1,
+    0b001: 2,
+    0b010: 4,
+    0b011: 8,
+}
+
 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 e2cabff..69bce27 100644
--- a/tests/cocotb/BUILD
+++ b/tests/cocotb/BUILD
@@ -175,6 +175,7 @@
 RVV_LOAD_STORE_TESTCASES = [
     "load_store_bits",
     "load_unit_masked",
+    "load_unit_all_vtypes_test",
     "load8_index8",
     "load8_index8_seg",
     "load8_index16",
@@ -204,6 +205,7 @@
     "load8_segment2_stride6_m1",
     "load16_segment2_stride6_m1",
     "store_unit_masked",
+    "store_unit_all_vtypes_test",
     "store8_index8",
     "store8_index8_seg",
     "store16_index8",
diff --git a/tests/cocotb/rvv/load_store/BUILD b/tests/cocotb/rvv/load_store/BUILD
index c62c274..32ede03 100644
--- a/tests/cocotb/rvv/load_store/BUILD
+++ b/tests/cocotb/rvv/load_store/BUILD
@@ -26,6 +26,9 @@
         "load_unit_masked": {
             "srcs": ["load_unit_masked.cc"],
         },
+        "load_unit_vtype": {
+            "srcs": ["load_unit_vtype.cc"],
+        },
         "load8_index8": {
             "srcs": ["load8_index8.cc"],
         },
@@ -146,6 +149,9 @@
         "store_unit_masked": {
             "srcs": ["store_unit_masked.cc"],
         },
+        "store_unit_vtype": {
+            "srcs": ["store_unit_vtype.cc"],
+        },
     },
 )
 
@@ -154,6 +160,7 @@
     srcs = [
         ":load_store_bits.elf",
         ":load_unit_masked.elf",
+        ":load_unit_vtype.elf",
         ":load8_index8.elf",
         ":load8_index8_seg.elf",
         ":load8_index16.elf",
@@ -194,5 +201,6 @@
         ":store16_seg_unit",
         ":store32_seg_unit",
         ":store_unit_masked.elf",
+        ":store_unit_vtype.elf",
     ],
 )
\ No newline at end of file
diff --git a/tests/cocotb/rvv/load_store/load_unit_vtype.cc b/tests/cocotb/rvv/load_store/load_unit_vtype.cc
new file mode 100644
index 0000000..6426645
--- /dev/null
+++ b/tests/cocotb/rvv/load_store/load_unit_vtype.cc
@@ -0,0 +1,93 @@
+// 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>
+
+// Tests loads after setting vtype. Loads/stores should be agnostic to vtype.
+
+size_t vl __attribute__((section(".data"))) = 16;
+size_t vtype __attribute__((section(".data"))) = 0;
+uint8_t load_data[256] __attribute__((section(".data")));
+uint8_t store_data[256] __attribute__((section(".data")));
+
+extern "C" {
+
+#define CREATE_LOAD_FN(name, data_bits) \
+__attribute__((used, retain)) void name() { \
+  size_t store_vl = 8*__riscv_vlenb(); \
+  asm("vsetvl zero, %[vl], %[vtype];" \
+      "vle" #data_bits ".v v8, %[load_data];" \
+      "vsetvli zero, %[store_vl], e8, m8, ta, ma;" \
+      "vse8.v v8, %[store_data];" \
+      : [store_data] "=m"(store_data) \
+      : [vl] "r"(vl), \
+        [store_vl] "r"(store_vl), \
+        [vtype] "r"(vtype), \
+        [load_data] "m"(load_data) \
+      : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", \
+        "vl", "vtype"); \
+}
+
+#define CREATE_SEGMENT_LOAD_FN(name, data_bits, segment) \
+__attribute__((used, retain)) void name() { \
+  size_t store_vl = 8*__riscv_vlenb(); \
+  asm("vsetvl zero, %[vl], %[vtype];" \
+      "vlseg" #segment "e" #data_bits ".v v8, %[load_data];" \
+      "vsetvli zero, %[store_vl], e8, m8, ta, ma;" \
+      "vse8.v v8, %[store_data];" \
+      : [store_data] "=m"(store_data) \
+      : [vl] "r"(vl), \
+        [store_vl] "r"(store_vl), \
+        [vtype] "r"(vtype), \
+        [load_data] "m"(load_data) \
+      : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", \
+        "vl", "vtype"); \
+}
+
+CREATE_LOAD_FN(test_vle8, 8)
+CREATE_SEGMENT_LOAD_FN(test_vlseg2e8, 8, 2)
+CREATE_SEGMENT_LOAD_FN(test_vlseg3e8, 8, 3)
+CREATE_SEGMENT_LOAD_FN(test_vlseg4e8, 8, 4)
+CREATE_SEGMENT_LOAD_FN(test_vlseg5e8, 8, 5)
+CREATE_SEGMENT_LOAD_FN(test_vlseg6e8, 8, 6)
+CREATE_SEGMENT_LOAD_FN(test_vlseg7e8, 8, 7)
+CREATE_SEGMENT_LOAD_FN(test_vlseg8e8, 8, 8)
+
+CREATE_LOAD_FN(test_vle16, 16)
+CREATE_SEGMENT_LOAD_FN(test_vlseg2e16, 16, 2)
+CREATE_SEGMENT_LOAD_FN(test_vlseg3e16, 16, 3)
+CREATE_SEGMENT_LOAD_FN(test_vlseg4e16, 16, 4)
+CREATE_SEGMENT_LOAD_FN(test_vlseg5e16, 16, 5)
+CREATE_SEGMENT_LOAD_FN(test_vlseg6e16, 16, 6)
+CREATE_SEGMENT_LOAD_FN(test_vlseg7e16, 16, 7)
+CREATE_SEGMENT_LOAD_FN(test_vlseg8e16, 16, 8)
+
+CREATE_LOAD_FN(test_vle32, 32)
+CREATE_SEGMENT_LOAD_FN(test_vlseg2e32, 32, 2)
+CREATE_SEGMENT_LOAD_FN(test_vlseg3e32, 32, 3)
+CREATE_SEGMENT_LOAD_FN(test_vlseg4e32, 32, 4)
+CREATE_SEGMENT_LOAD_FN(test_vlseg5e32, 32, 5)
+CREATE_SEGMENT_LOAD_FN(test_vlseg6e32, 32, 6)
+CREATE_SEGMENT_LOAD_FN(test_vlseg7e32, 32, 7)
+CREATE_SEGMENT_LOAD_FN(test_vlseg8e32, 32, 8)
+
+}
+
+void (*impl)() __attribute__((section(".data"))) = &test_vle8;
+
+int main(int argc, char** argv) {
+  impl();
+  return 0;
+}
diff --git a/tests/cocotb/rvv/load_store/store_unit_vtype.cc b/tests/cocotb/rvv/load_store/store_unit_vtype.cc
new file mode 100644
index 0000000..a50230f
--- /dev/null
+++ b/tests/cocotb/rvv/load_store/store_unit_vtype.cc
@@ -0,0 +1,93 @@
+// 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>
+
+// Tests unit stores after setting vtype. Unit stores should be vtype-agnostic.
+
+size_t vl __attribute__((section(".data"))) = 16;
+size_t vtype __attribute__((section(".data"))) = 0;
+uint8_t load_data[256] __attribute__((section(".data")));
+uint8_t store_data[256] __attribute__((section(".data")));
+
+extern "C" {
+
+#define CREATE_STORE_FN(name, data_bits) \
+__attribute__((used, retain)) void name() { \
+  size_t load_vl = 8*__riscv_vlenb(); \
+  asm("vsetvli zero, %[load_vl], e8, m8, ta, ma;" \
+      "vle8.v v8, %[load_data];" \
+      "vsetvl zero, %[vl], %[vtype];" \
+      "vse" #data_bits ".v v8, %[store_data];" \
+      : [store_data] "=m"(store_data) \
+      : [vl] "r"(vl), \
+        [load_vl] "r"(load_vl), \
+        [vtype] "r"(vtype), \
+        [load_data] "m"(load_data) \
+      : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", \
+        "vl", "vtype"); \
+}
+
+#define CREATE_SEGMENT_STORE_FN(name, data_bits, segment) \
+__attribute__((used, retain)) void name() { \
+  size_t load_vl = 8*__riscv_vlenb(); \
+  asm("vsetvli zero, %[load_vl], e8, m8, ta, ma;" \
+      "vle8.v v8, %[load_data];" \
+      "vsetvl zero, %[vl], %[vtype];" \
+      "vsseg" #segment "e" #data_bits ".v v8, %[store_data];" \
+      : [store_data] "=m"(store_data) \
+      : [vl] "r"(vl), \
+        [load_vl] "r"(load_vl), \
+        [vtype] "r"(vtype), \
+        [load_data] "m"(load_data) \
+      : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", \
+        "vl", "vtype"); \
+}
+
+CREATE_STORE_FN(test_vse8, 8)
+CREATE_SEGMENT_STORE_FN(test_vsseg2e8, 8, 2)
+CREATE_SEGMENT_STORE_FN(test_vsseg3e8, 8, 3)
+CREATE_SEGMENT_STORE_FN(test_vsseg4e8, 8, 4)
+CREATE_SEGMENT_STORE_FN(test_vsseg5e8, 8, 5)
+CREATE_SEGMENT_STORE_FN(test_vsseg6e8, 8, 6)
+CREATE_SEGMENT_STORE_FN(test_vsseg7e8, 8, 7)
+CREATE_SEGMENT_STORE_FN(test_vsseg8e8, 8, 8)
+
+CREATE_STORE_FN(test_vse16, 16)
+CREATE_SEGMENT_STORE_FN(test_vsseg2e16, 16, 2)
+CREATE_SEGMENT_STORE_FN(test_vsseg3e16, 16, 3)
+CREATE_SEGMENT_STORE_FN(test_vsseg4e16, 16, 4)
+CREATE_SEGMENT_STORE_FN(test_vsseg5e16, 16, 5)
+CREATE_SEGMENT_STORE_FN(test_vsseg6e16, 16, 6)
+CREATE_SEGMENT_STORE_FN(test_vsseg7e16, 16, 7)
+CREATE_SEGMENT_STORE_FN(test_vsseg8e16, 16, 8)
+
+CREATE_STORE_FN(test_vse32, 32)
+CREATE_SEGMENT_STORE_FN(test_vsseg2e32, 32, 2)
+CREATE_SEGMENT_STORE_FN(test_vsseg3e32, 32, 3)
+CREATE_SEGMENT_STORE_FN(test_vsseg4e32, 32, 4)
+CREATE_SEGMENT_STORE_FN(test_vsseg5e32, 32, 5)
+CREATE_SEGMENT_STORE_FN(test_vsseg6e32, 32, 6)
+CREATE_SEGMENT_STORE_FN(test_vsseg7e32, 32, 7)
+CREATE_SEGMENT_STORE_FN(test_vsseg8e32, 32, 8)
+
+}
+
+void (*impl)() __attribute__((section(".data"))) = &test_vse8;
+
+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 900d336..bfb22d1 100644
--- a/tests/cocotb/rvv_load_store_test.py
+++ b/tests/cocotb/rvv_load_store_test.py
@@ -18,7 +18,7 @@
 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.rvv_type_util import construct_vtype, DTYPE_TO_SEW, SEWS, SEW_TO_LMULS_AND_VLMAXS, LMUL_TO_EMUL
 from kelvin_test_utils.sim_test_fixture import Fixture
 
 
@@ -2500,3 +2500,160 @@
     routputs = (await fixture.core_mini_axi.read(target_out_addr, vl)).view(
         np.uint8)
     assert (input_data == routputs).all()
+
+@cocotb.test()
+async def load_unit_all_vtypes_test(dut):
+    """Testbench to test RVV Unit/segmented loads, with all vtypes."""
+    fixture = await Fixture.Create(dut)
+    r = runfiles.Create()
+    functions = [
+        ("test_vle8",      np.uint8, 1),
+        ("test_vlseg2e8",  np.uint8, 2),
+        ("test_vlseg3e8",  np.uint8, 3),
+        ("test_vlseg4e8",  np.uint8, 4),
+        ("test_vlseg5e8",  np.uint8, 5),
+        ("test_vlseg6e8",  np.uint8, 6),
+        ("test_vlseg7e8",  np.uint8, 7),
+        ("test_vlseg8e8",  np.uint8, 8),
+        ("test_vle16",     np.uint16, 1),
+        ("test_vlseg2e16", np.uint16, 2),
+        ("test_vlseg3e16", np.uint16, 3),
+        ("test_vlseg4e16", np.uint16, 4),
+        ("test_vlseg5e16", np.uint16, 5),
+        ("test_vlseg6e16", np.uint16, 6),
+        ("test_vlseg7e16", np.uint16, 7),
+        ("test_vlseg8e16", np.uint16, 8),
+        ("test_vle32",     np.uint32, 1),
+        ("test_vlseg2e32", np.uint32, 2),
+        ("test_vlseg3e32", np.uint32, 3),
+        ("test_vlseg4e32", np.uint32, 4),
+        ("test_vlseg5e32", np.uint32, 5),
+        ("test_vlseg6e32", np.uint32, 6),
+        ("test_vlseg7e32", np.uint32, 7),
+        ("test_vlseg8e32", np.uint32, 8),
+    ]
+
+    await fixture.load_elf_and_lookup_symbols(
+        r.Rlocation('kelvin_hw/tests/cocotb/rvv/load_store/load_unit_vtype.elf'),
+        ['vl', 'vtype', 'load_data', 'store_data', 'impl'] +
+            list(f[0] for f in functions),
+    )
+
+    vlenb = 16
+    with tqdm.tqdm(functions) as t:
+      for (function, dtype, segments) in t:
+        for sew in SEWS:
+          for lmul, vlmax in SEW_TO_LMULS_AND_VLMAXS[DTYPE_TO_SEW[dtype]]:
+            if (LMUL_TO_EMUL[lmul] * segments) > 8:
+              continue
+
+            # TODO(derekjchow): Remove this when bug is fixed
+            if sew != DTYPE_TO_SEW[dtype]:
+              continue
+
+            t.set_postfix({
+                'function': function,
+                'sew': sew,
+                'lmul': lmul,
+            })
+            vtype = construct_vtype(1, 1, sew, lmul)
+            await fixture.write_ptr('impl', function)
+            await fixture.write_word('vtype', vtype)
+            await fixture.write_word('vl', vlmax)
+            load_data = np.random.randint(
+                0, 255, 256, dtype=np.uint8).view(dtype)
+            await fixture.write('load_data', load_data)
+            await fixture.write('store_data',
+                                np.zeros(256, dtype=np.uint8))
+
+            await fixture.run_to_halt()
+
+            store_data = (await fixture.read('store_data', 256))
+
+            regsize = vlenb * LMUL_TO_EMUL[lmul]
+            for segment in range(segments):
+              segment_reg = store_data[segment*regsize:(segment+1)*regsize]
+              store_result = segment_reg.view(dtype)[0:vlmax]
+              expected_result = load_data[segment:segments*vlmax:segments]
+              assert (store_result == expected_result).all()
+
+
+@cocotb.test()
+async def store_unit_all_vtypes_test(dut):
+    """Testbench to test RVV Unit/segmented stores, with all vtypes."""
+    fixture = await Fixture.Create(dut)
+    r = runfiles.Create()
+    functions = [
+        ("test_vse8",      np.uint8, 1),
+        ("test_vsseg2e8",  np.uint8, 2),
+        ("test_vsseg3e8",  np.uint8, 3),
+        ("test_vsseg4e8",  np.uint8, 4),
+        ("test_vsseg5e8",  np.uint8, 5),
+        ("test_vsseg6e8",  np.uint8, 6),
+        ("test_vsseg7e8",  np.uint8, 7),
+        ("test_vsseg8e8",  np.uint8, 8),
+        ("test_vse16",     np.uint16, 1),
+        ("test_vsseg2e16", np.uint16, 2),
+        ("test_vsseg3e16", np.uint16, 3),
+        ("test_vsseg4e16", np.uint16, 4),
+        ("test_vsseg5e16", np.uint16, 5),
+        ("test_vsseg6e16", np.uint16, 6),
+        ("test_vsseg7e16", np.uint16, 7),
+        ("test_vsseg8e16", np.uint16, 8),
+        ("test_vse32",     np.uint32, 1),
+        ("test_vsseg2e32", np.uint32, 2),
+        ("test_vsseg3e32", np.uint32, 3),
+        ("test_vsseg4e32", np.uint32, 4),
+        ("test_vsseg5e32", np.uint32, 5),
+        ("test_vsseg6e32", np.uint32, 6),
+        ("test_vsseg7e32", np.uint32, 7),
+        ("test_vsseg8e32", np.uint32, 8),
+    ]
+
+    await fixture.load_elf_and_lookup_symbols(
+        r.Rlocation('kelvin_hw/tests/cocotb/rvv/load_store/store_unit_vtype.elf'),
+        ['vl', 'vtype', 'load_data', 'store_data', 'impl'] +
+            list(f[0] for f in functions),
+    )
+
+    vlenb = 16
+    with tqdm.tqdm(functions) as t:
+      for (function, dtype, segments) in t:
+        for sew in SEWS:
+          for lmul, vlmax in SEW_TO_LMULS_AND_VLMAXS[DTYPE_TO_SEW[dtype]]:
+            if (LMUL_TO_EMUL[lmul] * segments) > 8:
+              continue
+
+            # TODO(derekjchow): Remove this when bug is fixed
+            if sew != DTYPE_TO_SEW[dtype]:
+              continue
+
+            t.set_postfix({
+                'function': function,
+                'sew': sew,
+                'lmul': lmul,
+            })
+
+            vtype = construct_vtype(1, 1, sew, lmul)
+            await fixture.write_ptr('impl', function)
+            await fixture.write_word('vtype', vtype)
+            await fixture.write_word('vl', vlmax)
+            load_data = np.random.randint(
+                0, 255, 256, dtype=np.uint8).view(dtype)
+            await fixture.write('load_data', load_data)
+            await fixture.write('store_data',
+                                np.zeros(256, dtype=np.uint8))
+
+            await fixture.run_to_halt()
+
+            store_data = (await fixture.read(
+                'store_data',
+                segments * vlmax * np.dtype(dtype).itemsize)).view(dtype)
+
+            # Load and transpose stored segments
+            regsize = vlenb * LMUL_TO_EMUL[lmul] // np.dtype(dtype).itemsize
+            segment_regs = [
+                load_data[x*regsize:(x*regsize)+vlmax] for x in range(segments)
+            ]
+            expected_result = np.transpose(np.stack(segment_regs)).flatten()
+            assert (expected_result == store_data).all()