[MLIR][XeGPU] Add unroll patterns for XeGPU (1/N) (#137010)

Similar to vector ops, XeGPU ops need to be unrolled into smaller shapes
such that they can be dispatched into a hardware instruction. This PR
marks the initial phase of a series dedicated to incorporating unroll
patterns for XeGPU operations. In this installment, we introduce
patterns for the following operations:
1. createNd
2. updateNd
3. prefetchNd
4. loadNd
5. storeNd
6. dpas

Author: chencha3

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

@@ -303,7 +303,6 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout"> {
       return LayoutAttr::get(getContext(), getSgLayout(), getSgData(), nullptr,
                              getLaneLayout(), getLaneData(), getOrder());
     }
-
   }];
 
   let assemblyFormat = "`<` struct(params) `>`";

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

@@ -14,11 +14,67 @@ 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 {
+  /// Callback function that indicates whether vector unrolling should be
+  /// attempted on the operation.
+  using FilterConstraintFnType = std::function<LogicalResult(Operation *op)>;
+  FilterConstraintFnType filterConstraint = nullptr;
+  UnrollOptions &setFilterConstraint(FilterConstraintFnType constraint) {
+    filterConstraint = std::move(constraint);
+    return *this;
+  }
+
+  /// Function that computes the target shape for unrolling. It returns an
+  /// optional vector of integers representing the shape. If it returns
+  /// `std::nullopt`, unrolling is aborted for the given operation.
+  using NativeShapeFnType =
+      std::function<std::optional<SmallVector<int64_t>>(Operation *op)>;
+  NativeShapeFnType nativeShape = nullptr;
+  UnrollOptions &setNativeShapeFn(NativeShapeFnType fn) {
+    nativeShape = std::move(fn);
+    return *this;
+  }
+
+  /// Function that converts a ShapedType (TensorDescType or VectorType)
+  /// into the unrolled type based on the tileShape. It returns a vector of
+  /// types representing the unrolled types for simplicity.
+  using UnrolledTypeFnType = std::function<SmallVector<Type>(
+      ShapedType type, ArrayRef<int64_t> tileShape)>;
+  UnrolledTypeFnType getUnrolledTypes = nullptr;
+  UnrollOptions &setUnrolledTypesFn(UnrolledTypeFnType fn) {
+    getUnrolledTypes = 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 patterns to unroll xegpu operations to a smaller shapes.
+/// Users can control whether an operation to be unrolled or not, as well as
+/// its target shape via `options` structure. (via setting filterConstraint
+/// and nativeShape respectively, both of them are function refs taking `op` as
+/// input).
+/// An `op` is unrolled to the `targetShape` as follows, for each of its
+/// operands:
+///   1. the unrolled type `unrolledType` and number of unrolled instances
+///   `numUnrolledInstances` are computed from the `targetShape`.
+///   2. pack each operand. ExtractStridedSlice are created to break-up the
+///   vector operands. And BuiltinUnrealizedCastop are created to break-up
+///    the TensorDesc operands.
+///   3. the original op is cloned `numUnrolledInstances` times, once for each
+///   result.
+///   4. unpack the results. InsertStridedSlice are inserted for VectorType
+///   result, and BuiltinUnrealizedCastOp are inserted for TensorDescType result
+///   to re-assemble the slices into the original shape.
+void populateXeGPUUnrollPatterns(RewritePatternSet &patterns,
+                                 const UnrollOptions &options);
+
 } // namespace xegpu
 } // namespace mlir
 

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/XeGPU/IR/XeGPU.h"
 #include "mlir/IR/Builders.h"

mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_mlir_dialect_library(MLIRXeGPUTransforms
   XeGPUFoldAliasOps.cpp
   XeGPUSubgroupDistribute.cpp
+  XeGPUUnroll.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/XeGPU