[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

@@ -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 {