[Mlir-commits] [mlir] [MLIR][XeGPU] Add simple rank-based sg layout creation (PR #172867)

[Artem Kroviakov]

================
@@ -538,6 +538,55 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
   return false;
 }
 
+FailureOr<std::pair<SmallVector<int>, SmallVector<int>>>
+chooseLayout(llvm::ArrayRef<int64_t> wgShape, int64_t sgCount) {
+  const size_t rank = wgShape.size();
+
+  // Step 1. Factorize sgCount into prime factors.
+  SmallVector<int> layout;
+  int64_t temp = sgCount;
+  for (int64_t i = 2; i * i <= temp; ++i) {
+    while (temp % i == 0) {
+      layout.push_back(i);
+      temp /= i;
+    }
+  }
+  if (temp > 1)
+    layout.push_back(temp);
+
+  if (layout.size() < rank)
+    return failure();
+
+  // Step 2. Fuse two smallest factors until we have `rank` factors.
+  while (layout.size() > rank) {
+    std::sort(layout.begin(), layout.end());
+    int64_t a = layout[0];
+    int64_t b = layout[1];
+    layout.erase(layout.begin());
+    layout[0] = a * b;
+  }
+
+  SmallVector<int> data;
+  for (auto [i, dim] : llvm::enumerate(layout)) {
+    if (wgShape[i] % dim != 0)
+      return failure();
+    data.push_back(wgShape[i] / dim);
+  }
----------------
akroviakov wrote:

The only goal of the current approach is _layout stability_, even at the cost of missed optimizations for some cases or "wrong/unusual" layouts for corner cases. This is what I remember from @Jianhui-Li motivation.

Making this utility aware of inst_data turns it into more of a layout _solver_ and enforcing inst_data alignment risks layout changes that are explicitly undesired for this utility. 

https://github.com/llvm/llvm-project/pull/172867