add UnrollOption

Author: chencha3

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

@@ -296,6 +296,16 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout"> {
         return attr.size();
       return 0;
     }
+
+    LayoutAttr dropSgLayoutAndData() {
+      return LayoutAttr::get(getContext(), nullptr, nullptr, getInstData(),
+                             getLaneLayout(), getLaneData(), getOrder());
+    }
+
+    LayoutAttr dropInstData() {
+      return LayoutAttr::get(getContext(), getSgLayout(), getSgData(), nullptr,
+                             getLaneLayout(), getLaneData(), getOrder());
+    }
   }];
 
   let assemblyFormat = "`<` struct(params) `>`";

mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h

@@ -12,17 +12,36 @@
 namespace mlir {
 class RewritePatternSet;
 
-namespace vector {
-struct UnrollVectorOptions;
-} // namespace vector
 
 namespace xegpu {
+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 of the vector 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);
 
 void populateXeGPUUnrollPatterns(RewritePatternSet &patterns,
-                                 const vector::UnrollVectorOptions &options);
+                                 const UnrollOptions &options);
 
 } // namespace xegpu
 } // namespace mlir

mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp

@@ -9,7 +9,6 @@
 #include "mlir/Dialect/XeGPU/Transforms/Passes.h"
 
 #include "mlir/Dialect/Utils/IndexingUtils.h"
-#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Dialect/XeGPU/IR/XeGPU.h"
 #include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -34,7 +33,7 @@ namespace {
 template <typename SourceOp>
 struct UnrollPattern : public OpRewritePattern<SourceOp> {
   UnrollPattern(MLIRContext *context,
-                const vector::UnrollVectorOptions &options,
+                const xegpu::UnrollOptions &options,
                 PatternBenefit benefit = 1)
       : OpRewritePattern<SourceOp>(context, benefit), options(options) {}
 
@@ -64,10 +63,7 @@ struct UnrollPattern : public OpRewritePattern<SourceOp> {
     auto layout = dyn_cast_if_present<xegpu::LayoutAttr>(attr);
     if (!layout || layout.getLaneLayout() == nullptr)
       return xegpu::LayoutAttr();
-    return xegpu::LayoutAttr::get(
-        layout.getContext(), nullptr /* sg_layout */, nullptr /* sg_data */,
-        nullptr /* inst_data */, layout.getLaneLayout(), layout.getLaneData(),
-        layout.getOrder());
+    return layout.dropInstData();
   };
 
   SmallVector<Type> convertType(ShapedType type,
@@ -167,7 +163,7 @@ struct UnrollPattern : public OpRewritePattern<SourceOp> {
   const char *const unpackAttrName = "__xetile_blocking_unpack__";
   const char *const blockAttrName = "__xetile_blocking_inner_block__";
 
-  vector::UnrollVectorOptions options;
+  xegpu::UnrollOptions options;
 };
 
 struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
@@ -479,7 +475,7 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
 
 void mlir::xegpu::populateXeGPUUnrollPatterns(
     RewritePatternSet &patterns,
-    const mlir::vector::UnrollVectorOptions &options) {
+    const xegpu::UnrollOptions &options) {
   patterns.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
                UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp>(
       patterns.getContext(), options);

mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir

@@ -95,5 +95,4 @@ gpu.module @test {
     %c = xegpu.dpas %a, %b : vector<32x32xf16>, vector<32x32xf16> -> vector<32x32xf32>
     gpu.return %c : vector<32x32xf32>
   }
-
-}
+}

mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp

@@ -43,7 +43,7 @@ struct TestXeGPUUnrollingPatterns
       : PassWrapper(pass) {}
 
   void runOnOperation() override {
-    vector::UnrollVectorOptions options;
+    xegpu::UnrollOptions options;
     options.setNativeShapeFn(
         [&](Operation *op) -> std::optional<SmallVector<int64_t>> {
           if (isa<xegpu::CreateNdDescOp, xegpu::UpdateNdOffsetOp,