Address comments

nod-ai · Jan 31, 2025 · 6c23ff0 · 6c23ff0
1 parent e5ce2a4
commit 6c23ff0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Utils/AMDAIEDmaUtils.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Utils/AMDAIEDmaUtils.cpp
@@ -584,20 +584,21 @@ bool DmaDimConfig::isValidAccessPattern(SmallVector<int64_t> sizes,
   assert(sizes.size() == strides.size() &&
          "`sizes` and `strides` should have the same size");
   // No need to check the unit dimensions.
-  for (int i = 0; i < sizes.size(); i++) {
-    if (sizes[i] == 1) {
-      sizes.erase(sizes.begin() + i);
-      strides.erase(strides.begin() + i);
-      i--;
+  SmallVector<int64_t> nonUnitSizes;
+  SmallVector<int64_t> nonUnitStrides;
+  for (size_t i = 0; i < sizes.size(); ++i) {
+    if (sizes[i] != 1) {
+      nonUnitSizes.push_back(sizes[i]);
+      nonUnitStrides.push_back(strides[i]);
     }
   }
-  SmallVector<int64_t> maxSizes = getMaxSizes(sizes.size());
-  assert(maxSizes.size() >= sizes.size() &&
+  SmallVector<int64_t> maxSizes = getMaxSizes(nonUnitSizes.size());
+  assert(maxSizes.size() >= nonUnitSizes.size() &&
          "Max number of dimensions exceeded");
-  size_t frontToDrop = maxSizes.size() - sizes.size();
-  if (anyOutOfRange(sizes, maxSizes, frontToDrop)) return false;
-  SmallVector<int64_t> maxStrides = getMaxStrides(sizes.size());
-  if (anyOutOfRange(strides, maxStrides, frontToDrop)) return false;
+  size_t frontToDrop = maxSizes.size() - nonUnitSizes.size();
+  if (anyOutOfRange(nonUnitSizes, maxSizes, frontToDrop)) return false;
+  SmallVector<int64_t> maxStrides = getMaxStrides(nonUnitSizes.size());
+  if (anyOutOfRange(nonUnitStrides, maxStrides, frontToDrop)) return false;
   return true;
 }
 

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/AMDAIEDmaUtilsTest.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/AMDAIEDmaUtilsTest.cpp
@@ -592,6 +592,8 @@ TEST_F(FoldTest, UnitDimsMerge) {
   EXPECT_TRUE(checkFoldUnitDims({2, 0, 1, 0}, {1, 32, 1, 8},
                                 {1024, 32, 1024, 1}, {96, 0}, {32, 8},
                                 {32, 1}));
+  EXPECT_TRUE(checkFoldUnitDims({0, 0, 1, 0}, {2, 32, 1, 8}, {0, 32, 1024, 1},
+                                {0, 32, 0}, {2, 32, 8}, {0, 32, 1}));
   EXPECT_TRUE(
       checkFoldUnitDims({2, 2, 15}, {1, 1, 10}, {4, 6, 10}, {17}, {10}, {10}));
   EXPECT_TRUE(checkFoldUnitDims({3, 1, 15}, {1, 1, 10}, {4, 6, 10}, {1, 15},
@@ -607,6 +609,8 @@ TEST_F(FoldTest, UnitDimsFoldAndMerge) {
                                 {1}, {1}, {96}));
   EXPECT_TRUE(checkFoldUnitDims({1, 0, 1, 0}, {1, 1, 1, 8}, {1024, 32, 1024, 1},
                                 {2048}, {8}, {1}));
+  EXPECT_TRUE(checkFoldUnitDims({0, 0, 1, 0}, {1, 32, 1, 8}, {0, 32, 1024, 1},
+                                {32, 0}, {32, 8}, {32, 1}));
 }
 
 TEST_F(FoldTest, FoldRepetitionCount) {