[Mlir-commits] [mlir] [MLIR][XeGPU] Add sg layout propagation (PR #170879)

[Igor Zamyatin]

================
@@ -538,6 +538,56 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
   return false;
 }
 
+FailureOr<std::pair<SmallVector<int>, SmallVector<int>>>
+chooseLayout(llvm::ArrayRef<int64_t> big, llvm::ArrayRef<int> small,
+             const int64_t count) {
+  const int64_t n = big.size();
+  assert(n == small.size());
+  for (int dim = 0; dim < n; ++dim) {
+    if (big[dim] % small[dim])
+      return failure();
+  }
+  // Fill the large shape with smaller ones
+  SmallVector<int> tiles(n);
+  for (int64_t dim = 0; dim < n; ++dim) {
+    tiles[dim] = big[dim] / small[dim];
+    assert(!(tiles[dim] % 2) || tiles[dim] == 1);
+  }
+  // The baseline layout is based on the smallest data
+  int64_t totalTiles = llvm::product_of(tiles);
+  SmallVector<int> bestLayout = tiles;
+  SmallVector<int> bestData{small};
+  // Fit the layout to the given count
+  while (totalTiles > count) {
+    // Stack tiles along the longest layout dim
+    int64_t maxDim = 0;
+    for (int64_t i = 1; i < n; ++i)
+      if (bestLayout[i] > bestLayout[maxDim])
+        maxDim = i;
+    // If the longest dim is 1, cannot divide further
+    if (bestLayout[maxDim] == 1)
+      return failure();
+    // Merge data tiles
+    bestLayout[maxDim] /= 2;
+    bestData[maxDim] *= 2;
+    totalTiles = llvm::product_of(bestLayout);
+  }
+  assert(totalTiles == count);
+  return std::make_pair(bestLayout, bestData);
+}
+
+FailureOr<int64_t> getNumSg(Operation *op, const int sgSize) {
----------------
Garra1980 wrote:

How about some comment here?

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