[Mlir-commits] [mlir] [mlir][XeGPU][Transform] Add gather/scatter coalescing analysis. (PR #201684)

[Jianhui Li]

================
@@ -446,6 +450,195 @@ struct TestXeGPULayoutInterface
   }
 };
 
+struct TestXeGPUCoalesceGatherScatter
+    : public PassWrapper<TestXeGPUCoalesceGatherScatter, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUCoalesceGatherScatter)
+
+  StringRef getArgument() const final {
+    return "test-xegpu-coalesce-gather-scatter";
+  }
+
+  StringRef getDescription() const final {
+    return "Test the XeGPU contiguity analysis and its coalescing consumer.";
+  }
+
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<arith::ArithDialect>();
+    registry.insert<vector::VectorDialect>();
+    registry.insert<xegpu::XeGPUDialect>();
+  }
+
+  TestXeGPUCoalesceGatherScatter() = default;
+  TestXeGPUCoalesceGatherScatter(const TestXeGPUCoalesceGatherScatter &pass)
+      : PassWrapper(pass) {}
+
+  Option<unsigned> maxChunkSize{
+      *this, "max-chunk-size",
+      llvm::cl::desc("Upper bound on the produced lane_data FCD."),
+      llvm::cl::init(8)};
+
+  Option<bool> analyzeOnly{
+      *this, "analyze-only",
+      llvm::cl::desc("Only run the analysis (stamp contiguous_chunk "
+                     "attributes); do not apply."),
+      llvm::cl::init(false)};
+
+  void runOnOperation() override {
+    xegpu::runCoalesceGatherScatterAnalysis(getOperation());
+    if (analyzeOnly)
+      return;
+    getOperation()->walk([&](Operation *op) {
+      if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+        applyContiguousChunk(load, maxChunkSize);
+      else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+        applyContiguousChunk(store, maxChunkSize);
+    });
+  }
+
+private:
+  /// Largest power-of-two `<= bound` that divides `numLanes`.
+  static int64_t largestPow2Divisor(int64_t numLanes, int64_t bound) {
+    if (bound < 2 || numLanes < 2)
+      return 1;
+    int64_t f = std::min<int64_t>(bound, numLanes);
+    // Round down to power of 2.
+    if (!llvm::isPowerOf2_64(f))
+      f = static_cast<int64_t>(llvm::bit_floor(static_cast<uint64_t>(f)));
+    while (f >= 2) {
+      if (numLanes % f == 0)
+        return f;
+      f /= 2;
+    }
+    return 1;
+  }
+
+  /// Look up the subgroup size from the enclosing gpu.module's xevm.target.
+  /// Falls back to 16 when no target chip is found or the chip is unknown,
+  /// matching the typical Intel Xe2 default.
+  static unsigned lookupSubgroupSize(Operation *op) {
+    const auto *uArch =
+        xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+    return uArch ? static_cast<unsigned>(uArch->getSubgroupSize()) : 16u;
----------------
Jianhui-Li wrote:

why 16u in case of no uArch?  In layout propagation, in case of no uArch info, we exit the pass. 

The abstraction/utility here has little value (just 2 lines) and encode different behavior, why not just copy paste the code pattern?  

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