Fix index size

Samsung · Jan 20, 2025 · 4ec592e · 4ec592e
1 parent 36eda15
commit 4ec592e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/compiler/luci/pass/src/helpers/ArrayIndex.cpp b/compiler/luci/pass/src/helpers/ArrayIndex.cpp
@@ -46,7 +46,7 @@ uint32_t Array4DIndex::operator()(uint32_t i0, uint32_t i1, uint32_t i2, uint32_
   return i0 * _strides[0] + i1 * _strides[1] + i2 * _strides[2] + i3 * _strides[3];
 }
 
-uint32_t Array4DIndex::size(void) const { return _strides[3]; }
+uint32_t Array4DIndex::size(void) const { return _dim[0] * _dim[1] * _dim[2] * _dim[3]; }
 
 uint32_t Array4DIndex::stride(uint32_t axis) const { return _strides[axis]; }
 

diff --git a/compiler/luci/pass/src/helpers/ArrayIndex.test.cpp b/compiler/luci/pass/src/helpers/ArrayIndex.test.cpp
@@ -24,11 +24,15 @@ TEST(LuciPassHelpersArrayIndex, array_index_4d)
 
   EXPECT_EQ(idx(0, 0, 0, 0), 0);
 
+  // stride
   EXPECT_EQ(idx(1, 0, 0, 0), idx.stride(0));
   EXPECT_EQ(idx(0, 1, 0, 0), idx.stride(1));
   EXPECT_EQ(idx(0, 0, 1, 0), idx.stride(2));
   EXPECT_EQ(idx(0, 0, 0, 1), idx.stride(3));
 
+  // size
+  EXPECT_EQ(idx.size(), 5 * 4 * 3 * 2);
+
   EXPECT_EQ(idx(4, 3, 2, 1), 4 * 4 * 3 * 2 + 3 * 3 * 2 + 2 * 2 + 1);
 }