[Mlir-commits] [mlir] [mlir][xegpu] Add support for `vector.reduction` and `vector.multi_reduction` subgroup to work-item distribution. (PR #180308)

[Artem Kroviakov]

================
@@ -362,6 +395,137 @@ struct SgToWiPrefetchNd : public OpConversionPattern<xegpu::PrefetchNdOp> {
   }
 };
 
+/// This pattern distributes a subgroup-level vector.reduction op to
+/// workitem-level. This require shuffling the data across the workitems (using
+/// gpu::ShuffleOp) and reducing in stages until all workitems have the final
+/// result.
+struct SgToWiVectorReduction : public OpConversionPattern<vector::ReductionOp> {
+  using OpConversionPattern<vector::ReductionOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::ReductionOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto layout = xegpu::getDistributeLayoutAttr(op.getVector());
+
+    // If no layout, nothing to do.
+    if (!layout || !layout.isForSubgroup())
+      return failure();
+
+    VectorType vectorType = op.getSourceVectorType();
+
+    // Only rank 1 vectors supported.
+    if (vectorType.getRank() != 1)
----------------
akroviakov wrote:

> I can't think of any meaningful use of rank0 as source

I agree, but it would be a [valid IR](https://github.com/llvm/llvm-project/blob/79ec0ce5d8b72d26bb25c48083b693b15b53da55/mlir/test/Dialect/Vector/canonicalize.mlir#L3079-L3085), even though such a reduction folds during canonicalization.

> So lets keep this check for now and add it later.

Sure.

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