[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.

Author: Mogball

lib/Tools/LinearLayout.cpp

@@ -943,17 +943,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 {

unittest/Tools/LinearLayoutTest.cpp

@@ -513,6 +513,59 @@ TEST_F(LinearLayoutTest, InvertAndCompose_BroadcastedDims2) {
   EXPECT_EQ(c.compose(b), a.transposeOuts(llvm::to_vector(b.getOutDimNames())));
 }
 
+TEST_F(LinearLayoutTest, InvertAndCompose_IdentityInDim) {
+  SmallVector<StringAttr> outDims = {S("dim0"), S("dim1"), S("dim2"),
+                                     S("dim3"), S("dim4"), S("dim5"),
+                                     S("dim6"), S("dim7"), S("dim8")};
+
+  LinearLayout src({{S("register"),
+                     {
+                         {0, 0, 0, 0, 0, 0, 0, 0, 1},
+                         {0, 0, 0, 0, 0, 0, 0, 1, 0},
+                     }},
+                    {S("lane"),
+                     {
+                         {0, 0, 0, 0, 0, 0, 1, 0, 0},
+                         {0, 0, 0, 0, 0, 1, 0, 0, 0},
+                         {0, 0, 0, 0, 1, 0, 0, 0, 0},
+                         {0, 0, 0, 1, 0, 0, 0, 0, 0},
+                         {0, 0, 1, 0, 0, 0, 0, 0, 0},
+                     }},
+                    {S("warp"),
+                     {
+                         {0, 1, 0, 0, 0, 0, 0, 0, 0},
+                         {1, 0, 0, 0, 0, 0, 0, 0, 0},
+                     }},
+                    {S("block"), {}}},
+                   outDims);
+  LinearLayout dst({{S("register"),
+                     {
+                         {0, 0, 0, 0, 0, 0, 0, 0, 1},
+                         {0, 0, 0, 0, 0, 0, 0, 1, 0},
+                     }},
+                    {S("lane"),
+                     {
+                         {1, 0, 0, 0, 0, 0, 0, 0, 0},
+                         {0, 1, 0, 0, 0, 0, 0, 0, 0},
+                         {0, 0, 1, 0, 0, 0, 0, 0, 0},
+                         {0, 0, 0, 1, 0, 0, 0, 0, 0},
+                         {0, 0, 0, 0, 1, 0, 0, 0, 0},
+                     }},
+                    {S("warp"),
+                     {
+                         {0, 0, 0, 0, 0, 1, 0, 0, 0},
+                         {0, 0, 0, 0, 0, 0, 1, 0, 0},
+                     }},
+                    {S("block"), {}}},
+                   outDims);
+
+  LinearLayout cvt = dst.invertAndCompose(src);
+  SmallVector<std::pair<StringAttr, int32_t>> k = {
+      {S("register"), 3}, {S("lane"), 0}, {S("warp"), 2}, {S("block"), 0}};
+
+  EXPECT_EQ(dst.apply(k), src.apply(cvt.apply(k)));
+}
+
 TEST_F(LinearLayoutTest, NumConsecutiveInOut) {
   EXPECT_EQ(
       1,