Address comments

Author: victor-eds

third_party/intel/lib/TritonIntelGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -452,37 +452,41 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     StringAttr kWarp = str_attr("warp");
     StringAttr kBlock = str_attr("block");
 
-    LinearLayout comp = srcLayout.invertAndCompose(dstLayout);
+    LinearLayout comp = dstLayout.invertAndCompose(srcLayout);
     std::optional<LinearLayout> conversion = comp.divideRight(
         LinearLayout::identity1D(comp.getInDimSize(kWarp), kWarp, kWarp) *
         LinearLayout::identity1D(comp.getInDimSize(kBlock), kBlock, kBlock));
     assert(conversion && "Expecting valid conversion");
     // Expected conversion is:
     // - register=1 -> (0, 1)
-    //    register=2 -> (0, 2)
-    //    register=4 -> (0, 4)
-    //    register=8 -> (0, 8)
-    //    register=N -> (N, 0)
-    //    ...
-    //  - lane=1 -> (1, 0)
-    //    lane=2 -> (2, 0)
-    //    lane=4 -> (4, 0)
-    //    lane=8 -> (8, 0)
-    // where out dims are: [register (size 2*N), lane (size 16)]
-    std::vector<std::vector<int32_t>> registerBases{
-        {0, 1}, {0, 2}, {0, 4}, {0, 8}};
-    {
-      // Populate register bases for N > 8.
-      std::vector<int32_t> base(2);
-      for (int32_t i = 16, n = conversion->getInDimSize(kRegister); i < n;
-           i *= 2) {
-        base.front() = i;
-        registerBases.push_back(base);
+    // ...
+    // - register=i -> (0, 2**(i-1))
+    // ...
+    // - register=N -> (0, 2**(N-1))
+    // - lane=1 -> (0, 1)
+    // ...
+    // - lane=j -> (2**(j-1), 0)
+    // ...
+    //   lane=M -> (2**(M-1), 0)
+    // where out dims are: [register (size 2**(N-1)), lane (size 2**(M-1))]
+    //
+    // With N = M.
+    const auto buildBasis = [&](int32_t size, std::size_t index) {
+      std::vector<std::vector<int32_t>> basis;
+      std::vector<int32_t> curr(2);
+      for (int32_t i = 1; i < size; i *= 2) {
+        curr[index] = i;
+        basis.push_back(curr);
       }
-    }
+      return basis;
+    };
+
+    constexpr std::size_t laneIndex = 0;
+    constexpr std::size_t registerIndex = 1;
+    int32_t size = conversion->getInDimSize(kLane);
     std::array<std::pair<StringAttr, std::vector<std::vector<int32_t>>>, 2>
-        bases{{{kRegister, std::move(registerBases)},
-               {kLane, {{1, 0}, {2, 0}, {4, 0}, {8, 0}}}}};
+        bases{{{kRegister, buildBasis(size, registerIndex)},
+               {kLane, buildBasis(size, laneIndex)}}};
     std::array<StringAttr, 2> outDimNames{kRegister, kLane};
     return conversion == LinearLayout(bases, outDimNames);
   }