[Layouts] Fix invertAndCompose when there are identity dimensions (…

…#5468) I found a counterexample where `invertAndCompose` was giving an incorrect result. Thanks for @lezcano for pointing out a possible solution: if input dim 0 and 4 are identity dimensions, then they are added back to the reduced inverse as dimensions 3 and 4. Them reshaping the result will assign the wrong bases to the wrong dimensions. Use a transpose instead to order the input and output dimensions correctly instead.
intel · Dec 28, 2024 · 18b8b75 · 18b8b75
1 parent df900e4
commit 18b8b75
Showing 1 changed file with 4 additions and 11 deletions.
diff --git a/lib/Tools/LinearLayout.cpp b/lib/Tools/LinearLayout.cpp
@@ -917,17 +917,10 @@ LinearLayout LinearLayout::invertAndCompose(const LinearLayout &outer) const {
     ret *= LinearLayout::identity1D(A.getInDimSize(dim), dim, dim);
   }
 
-  // Reshape the result
-  SmallVector<std::pair<StringAttr, int32_t>> inDimsA;
-  SmallVector<std::pair<StringAttr, int32_t>> inDimsB;
-  for (auto dim : A.getInDimNames()) {
-    inDimsA.push_back({dim, A.getInDimSize(dim)});
-  }
-  for (auto dim : B.getInDimNames()) {
-    inDimsB.push_back({dim, B.getInDimSize(dim)});
-  }
-  ret = ret.reshapeIns(inDimsB).reshapeOuts(inDimsA);
-  return ret;
+  // Reorder the dimensions in the result to match the order expected by the
+  // current and outer layouts.
+  return ret.transposeIns(llvm::to_vector(B.getInDimNames()))
+      .transposeOuts(llvm::to_vector(A.getInDimNames()));
 }
 
 LinearLayout LinearLayout::invert() const {