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

[Jianhui Li]

================
@@ -446,6 +446,48 @@ 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 coalesce-gather-scatter analysis + apply APIs.";
+  }
+
+  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 xegpu.coalesce_hint "
+                     "attributes); do not apply."),
+      llvm::cl::init(false)};
+
+  void runOnOperation() override {
+    xegpu::CoalesceGatherScatterAnalysisOptions options;
+    options.maxChunkSize = maxChunkSize;
+    xegpu::runCoalesceGatherScatterAnalysis(getOperation(), options);
+    if (!analyzeOnly)
----------------
Jianhui-Li wrote:

I see.  I am confused by the co-location. I don't see value of have the helpers there, since the contiguity analysis should not involve these layout related function at all. Please move them to here. 

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