[Mlir-commits] [mlir] [MLIR][XeGPU] Add unroll patterns for XeGPU (1/N) (PR #137010)

[Chao Chen]

================
@@ -14,11 +14,53 @@ class RewritePatternSet;
 
 namespace xegpu {
 
+/// Options to control the XeGPU unrolling. Its main purpose is to
+/// provide a way to customize the native shape of the operation.
+struct UnrollOptions {
+  using FilterConstraintFnType = std::function<LogicalResult(Operation *op)>;
+  /// Callback function that indicates whether vector unrolling should be
+  /// attempted on the operation.
+  FilterConstraintFnType filterConstraint = nullptr;
+  UnrollOptions &setFilterConstraint(FilterConstraintFnType constraint) {
+    filterConstraint = std::move(constraint);
+    return *this;
+  }
+
+  using NativeShapeFnType =
+      std::function<std::optional<SmallVector<int64_t>>(Operation *op)>;
+  /// Function that returns the shape to unroll to for a given operation.
+  /// The unrolling is aborted if the function returns `std::nullopt`.
+  NativeShapeFnType nativeShape = nullptr;
+  UnrollOptions &setNativeShapeFn(NativeShapeFnType fn) {
+    nativeShape = std::move(fn);
+    return *this;
+  }
+};
+
 /// Appends patterns for folding aliasing ops into XeGPU ops into `patterns`.
 void populateXeGPUFoldAliasOpsPatterns(RewritePatternSet &patterns);
 /// Appends patterns for XeGPU SIMT distribution into `patterns`.
 void populateXeGPUSubgroupDistributePatterns(RewritePatternSet &patterns);
 
+/// Collect a set of pattern to unroll xegpu operations to a smaller shapes.
+/// Users can control whether an operation to be unrolled or not, as well as
+/// the its target shape via `options` structure. (via setting filterConstraint
----------------
chencha3 wrote:

Fixed, thanks

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