From d7d7caf363051bd477e64d3bd9ef14a939fcb822 Mon Sep 17 00:00:00 2001
From: "Wang, Quintin" <quintin.wang@intel.com>
Date: Fri, 27 Dec 2024 08:58:53 +0000
Subject: [PATCH] Cherry-pick upstream PR
 https://github.com/triton-lang/triton/pull/5468

---
 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 0ab563908a..cd61abe5e2 100644
--- a/lib/Tools/LinearLayout.cpp
+++ b/lib/Tools/LinearLayout.cpp
@@ -942,17 +942,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()));
 }
 
 llvm::MapVector<StringAttr, int32_t>