Fix riscv vector load segment, indexed tests.

PiperOrigin-RevId: 874288325
Change-Id: I06d4cc9fba1d593c0b5a5784a10ab3c7888ccfd3

riscv/riscv_vector_memory_instructions.cc

diff --git a/riscv/riscv_vector_memory_instructions.cc b/riscv/riscv_vector_memory_instructions.cc
index 29a1d2b..5144796 100644
--- a/riscv/riscv_vector_memory_instructions.cc
+++ b/riscv/riscv_vector_memory_instructions.cc
@@ -594,7 +594,7 @@
   }
   // Index lmul is scaled from the lmul by the relative size of the index
   // element to the SEW (selected element width).
-  int index_emul = (element_width * lmul8) / element_width;
+  int index_emul = (index_width * lmul8) / element_width;
   // Validate that index_emul has a legal value.
   if ((index_emul > 64) || (index_emul == 0)) {
     // TODO: signal vector error.
@@ -644,7 +644,8 @@
     }
     for (int field = 0; field < num_fields; field++) {
       masks[field * num_segments + i] = mask_value;
-      addresses[field * num_segments + i] = base + offset + field;
+      addresses[field * num_segments + i] =
+          base + offset + field * element_width;
     }
   }
   auto* context = new VectorLoadContext(data_db, mask_db, element_width, start,

riscv/test/riscv_vector_memory_instructions_test.cc

diff --git a/riscv/test/riscv_vector_memory_instructions_test.cc b/riscv/test/riscv_vector_memory_instructions_test.cc
index ed39c73..180d39c 100644
--- a/riscv/test/riscv_vector_memory_instructions_test.cc
+++ b/riscv/test/riscv_vector_memory_instructions_test.cc
@@ -770,8 +770,12 @@
               bool mask = (kA5Mask[mask_index] >> mask_offset) & 0x1;
               if (mask && (count < num_values)) {
                 // First compute the expected value, then compare it.
-                int address = 4096 + IndexValue<IndexType>(count) +
-                              field * sizeof(ValueType);
+                int values_per_reg = kVectorLengthInBytes / sizeof(IndexType);
+                auto index_span = vreg_[kVs2 + count / values_per_reg]
+                                      ->data_buffer()
+                                      ->Get<IndexType>();
+                IndexType index_val = index_span[count % values_per_reg];
+                int address = 4096 + index_val + field * sizeof(ValueType);
                 ValueType value = ComputeValue<ValueType>(address);
 
                 EXPECT_EQ(value, span[i])
@@ -1681,63 +1685,63 @@
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei8_16) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint8_t, uint16_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei8_32) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint8_t, uint32_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei8_64) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint8_t, uint64_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei16_8) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint16_t, uint8_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei16_16) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint16_t, uint16_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei16_32) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint16_t, uint32_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei16_64) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint16_t, uint64_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei32_8) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint32_t, uint8_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei32_16) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint32_t, uint16_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei32_32) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint32_t, uint32_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei32_64) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint32_t, uint64_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei64_8) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint64_t, uint8_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei64_16) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint64_t, uint16_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei64_32) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint64_t, uint32_t>();
 }
 
 TEST_F(RV32VInstructionsTest, Vluxsegei64_64) {
-  VectorLoadIndexedSegmentHelper<uint8_t, uint8_t>();
+  VectorLoadIndexedSegmentHelper<uint64_t, uint64_t>();
 }
 
 // Test Vector store strided.