From 18b8b753ba707973626e737a9a6c03e6e3958c1c Mon Sep 17 00:00:00 2001 From: Jeff Niu Date: Thu, 19 Dec 2024 13:18:42 -0800 Subject: [PATCH] [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. --- lib/Tools/LinearLayout.cpp | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/lib/Tools/LinearLayout.cpp b/lib/Tools/LinearLayout.cpp index 28c2c16568..d4407abbbd 100644 --- 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> inDimsA; - SmallVector> 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 {