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

[Adam Siemieniuk]

================
@@ -487,10 +531,63 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
   } else if (layoutKind == LayoutKind::Lane) {
     return !(anchorLayout.getEffectiveLaneLayoutAsInt().empty() ||
              anchorLayout.getEffectiveLaneDataAsInt().empty());
+  } else if (layoutKind == LayoutKind::Subgroup) {
+    return !(anchorLayout.getEffectiveSgLayoutAsInt().empty() ||
+             anchorLayout.getEffectiveSgDataAsInt().empty());
   }
   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);
----------------
adam-smnk wrote:

nit: maybe more graceful error? Or are there preconditions that really ensures that ratio is correct?

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