[mlir][vector] Add pattern to distribute vector reduction to GPU shuffles

[ThomasRaoux]

Summary

Add a pattern to do ad hoc lowering of vector.reduction to a sequence of warp shuffles. This allow distributing reduction on a warp for GPU targets.
Also add an execution test for warp reduction.

co-authored with @springerm

[csigg]

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

746

These %e look wrong.

769

Could you use rewriter.notifyMatchFailure() to make debugging easier? Here and below.

784–785

Would it be cleaner to initialize those vectors during construction instead of push_back?

785

.getType()?

790

Would getResults().back() be easier to parse?

791

I would have expected this to be 'laneId' of some sort. Does this happen latter in the lowering?

807

Could we use 'bfly' reduction instead of broadcasting separately?

[ThomasRaoux]

Address review comments

[ThomasRaoux]

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

746

good catch, fixed it.

784–785

indeed.

790

The new op has extra results so this is not the last element of the newWarpOp results.

791

The op will return Value vector<1xf32> this extract just converts it back to a scalar. The laneId indexing is implicitly done due to the distribution of the vector onto simt lanes.

807

That’s a good point. One potential problem I see with the

807

sorry outdated reply, changed it to butterfly reduction.

[springerm]

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

806

Better to wrap in updateRootInplace

[mehdi_amini]

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

28 ↗

(On Diff #436659)

Seems like a layering violation?

[springerm]

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

28 ↗

(On Diff #436659)

Note this is VectorTransforms not VectorDialect. But this code could also live in GPUTransforms (adding a dependency on VectorDialect to GPUTransforms instead).

[ThomasRaoux]

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

28 ↗

(On Diff #436659)

Addressed in https://reviews.llvm.org/D127867

[ThomasRaoux]

address review comment

mlir/include/mlir/Dialect/Vector/Transforms/VectorDistribution.h

 /// Move scalar operations with no dependency on the warp op outside of the
 /// region.
 void moveScalarUniformCode(WarpExecuteOnLane0Op op);
 
 /// Collect patterns to propagate warp distribution.
 void populatePropagateWarpVectorDistributionPatterns(
     RewritePatternSet &pattern);
 
-/// Collect patterns to distribute vector reduction ops using GPU warp shuffle
-/// ops.
-void populateReductionToGPUWarpShufflePatterns(RewritePatternSet &pattern);
+/// Lambda signature to compute a reduction of a distributed value for the given
+/// reduction kind and size.
+using DistributedReductionFn =
+    std::function<Value(Location, OpBuilder &, Value, CombiningKind, uint32_t)>;
+
+/// Collect patterns to distribute vector reduction ops using given lamdba to
+/// distribute reduction op.
+void populateDistributeReduction(RewritePatternSet &pattern,
+                                 DistributedReductionFn distributedReductionFn);
 
 } // namespace vector
 } // namespace mlir
 #endif // MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORDISTRIBUTION_H_

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

 //===- VectorDistribute.cpp - patterns to do vector distribution ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/Transforms/SideEffectUtils.h"
 
 using namespace mlir;
for (const auto &res : llvm::enumerate(resultIdx)) {
           .replaceAllUsesWith(newForOp.getResult(res.index()));
       newForOp->setOperand(res.index() + 3, warpOp.getResult(res.value()));
     }
     return success();
   }
 };
 
 /// A pattern that extracts vector.reduction ops from a WarpExecuteOnLane0Op.
-/// The vector is reduced in parallel. Currently limited to vector<32x...>
-/// values. Every lane reduces two scalars, 5 times in a row.
-/// E.g.:
+/// The vector is reduced in parallel. Currently limited to vector size matching
+/// the warpOp size. E.g.:
 /// ```
-/// %r = vector_ext.warp_execute_on_lane_0(%laneid) -> (f32) {
+/// %r = vector_ext.warp_execute_on_lane_0(%laneid)[32] -> (f32) {
 ///   %0 = "some_def"() : () -> (vector<32xf32>)
 ///   %1 = vector.reduction "add", %0 : vector<32xf32> into f32
 ///   vector_ext.yield %1 : f32
 /// }
 /// ```
 /// is lowered to:
 /// ```
-/// %0 = vector_ext.warp_execute_on_lane_0(%laneid) -> (vector<1xf32>) {
+/// %0 = vector_ext.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {
 ///   %1 = "some_def"() : () -> (vector<32xf32>)
 ///   vector_ext.yield %1 : vector<32xf32>
 /// }
 /// %a = vector.extract %0[0] : vector<1xf32>
-/// %r0, %s0 = gpu.shuffle xor %e, %c1, %c32 : f32
-/// %a0 = arith.addf %a, %r0 : f32
-/// %r1, %s1 = gpu.shuffle xor %a0, %c2, %c32 : f32
-/// %a1 = arith.addf %a0, %r1 : f32
-/// ...
-/// %r4, %s4 = gpu.shuffle xor %a3, %c16, %c32 : f32
-/// %r = arith.addf %a3, %r4 : f32
+/// %r = ("warp.reduction %a")
 /// ```
-struct ReductionToGPUWarpShuffle
-    : public OpRewritePattern<WarpExecuteOnLane0Op> {
-  using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
+struct WarpOpReduction : public OpRewritePattern<WarpExecuteOnLane0Op> {
+  WarpOpReduction(MLIRContext *context,
+                  DistributedReductionFn distributedReductionFn,
+                  PatternBenefit benefit = 1)
+      : OpRewritePattern<WarpExecuteOnLane0Op>(context, benefit),
+        distributedReductionFn(distributedReductionFn) {}
 
   LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
                                 PatternRewriter &rewriter) const override {
     OpOperand *yieldOperand = getWarpResult(
         warpOp, [](Operation *op) { return isa<vector::ReductionOp>(op); });
     if (!yieldOperand)
       return failure();
 
     auto reductionOp =
         cast<vector::ReductionOp>(yieldOperand->get().getDefiningOp());
     auto vectorType = reductionOp.getVector().getType().cast<VectorType>();
     // Only rank 1 vectors supported.
     if (vectorType.getRank() != 1)
       return rewriter.notifyMatchFailure(
           warpOp, "Only rank 1 reductions can be distributed.");
     // Only warp_size-sized vectors supported.
     if (static_cast<uint64_t>(vectorType.getShape()[0]) != warpOp.getWarpSize())
       return rewriter.notifyMatchFailure(
           warpOp, "Reduction vector dimension must match was size.");
     // Only f32 and i32 element types are supported.
     if (!reductionOp.getType().isF32() &&
         !reductionOp.getType().isSignlessInteger(32))
       return rewriter.notifyMatchFailure(
           warpOp,
           "Reduction distribution currently only supports 32bits types.");
 
     Location yieldLoc = yieldOperand->getOwner()->getLoc();
 
     // Return vector that will be reduced from the WarpExecuteOnLane0Op.
     unsigned operandIndex = yieldOperand->getOperandNumber();
     SmallVector<Value> yieldValues = {reductionOp.getVector()};
     SmallVector<Type> retTypes = {VectorType::get({1}, reductionOp.getType())};
     unsigned numResults = warpOp.getNumResults();
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
         rewriter, warpOp, yieldValues, retTypes);
     rewriter.setInsertionPointAfter(newWarpOp);
 
     // Every lane has one scalar value. These should be reduced.
     Value laneValVec = newWarpOp.getResult(numResults);
     Value laneVal = rewriter.create<vector::ExtractOp>(yieldLoc, laneValVec, 0);
-
-    // Parallel reduction using butterfly shuffles.
-    for (uint64_t i = 1; i < newWarpOp.getWarpSize(); i <<= 1) {
-      Value shuffled =
-          rewriter
-              .create<gpu::ShuffleOp>(reductionOp.getLoc(), laneVal, i,
-                                      /*width=*/newWarpOp.getWarpSize(),
-                                      /*mode=*/gpu::ShuffleMode::XOR)
-              .result();
-      laneVal = makeArithReduction(rewriter, reductionOp.getLoc(),
-                                   reductionOp.getKind(), laneVal, shuffled);
-    }
-
+    laneVal =
+        distributedReductionFn(reductionOp.getLoc(), rewriter, laneVal,
+                               reductionOp.getKind(), newWarpOp.getWarpSize());
     newWarpOp.getResult(operandIndex).replaceAllUsesWith(laneVal);
     return success();
   }
+
+private:
+  DistributedReductionFn distributedReductionFn;
 };
 
 } // namespace
 
 void mlir::vector::populateWarpExecuteOnLane0OpToScfForPattern(
     RewritePatternSet &patterns,
     const WarpExecuteOnLane0LoweringOptions &options) {
   patterns.add<WarpOpToScfForPattern>(patterns.getContext(), options);
 }
 
 void mlir::vector::populateDistributeTransferWriteOpPatterns(
     RewritePatternSet &patterns, DistributionMapFn distributionMapFn) {
   patterns.add<WarpOpTransferWrite>(patterns.getContext(), distributionMapFn);
 }
 
 void mlir::vector::populatePropagateWarpVectorDistributionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<WarpOpElementwise, WarpOpTransferRead, WarpOpDeadResult,
                WarpOpBroadcast, WarpOpForwardOperand, WarpOpScfForOp>(
       patterns.getContext());
 }
 
-void mlir::vector::populateReductionToGPUWarpShufflePatterns(
-    RewritePatternSet &patterns) {
-  patterns.add<ReductionToGPUWarpShuffle>(patterns.getContext());
+void mlir::vector::populateDistributeReduction(
+    RewritePatternSet &patterns,
+    DistributedReductionFn distributedReductionFn) {
+  patterns.add<WarpOpReduction>(patterns.getContext(), distributedReductionFn);
 }
 
 void mlir::vector::moveScalarUniformCode(WarpExecuteOnLane0Op warpOp) {
   Block *body = warpOp.getBody();
 
   // Keep track of the ops we want to hoist.
   llvm::SmallSetVector<Operation *, 8> opsToMove;
 

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
   // The symbol table inserts at the end of the module, but globals are a bit
   // nicer if they are at the beginning.
   global->moveBefore(&moduleOp.front());
 
   builder.restoreInsertionPoint(ip);
   return builder.create<memref::GetGlobalOp>(loc, memrefType, symbolName);
 }
 
+static Value warpReduction(Location loc, OpBuilder &builder, Value input,
+                           CombiningKind kind, uint32_t size) {
+  Value laneVal = input;
+  // Parallel reduction using butterfly shuffles.
+  for (uint64_t i = 1; i < size; i <<= 1) {
+    Value shuffled = builder
+                         .create<gpu::ShuffleOp>(loc, laneVal, i,
+                                                 /*width=*/size,
+                                                 /*mode=*/gpu::ShuffleMode::XOR)
+                         .result();
+    laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);
+  }
+  return laneVal;
+}
+
 struct TestVectorDistribution
     : public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution)
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<scf::SCFDialect, memref::MemRefDialect, gpu::GPUDialect,
                     AffineDialect>();
   }
if (distributeTransferWriteOps) {
       };
       RewritePatternSet patterns(ctx);
       populateDistributeTransferWriteOpPatterns(patterns, distributionFn);
       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
     }
     if (propagateDistribution) {
       RewritePatternSet patterns(ctx);
       vector::populatePropagateWarpVectorDistributionPatterns(patterns);
-      vector::populateReductionToGPUWarpShufflePatterns(patterns);
+      vector::populateDistributeReduction(patterns, warpReduction);
       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
     }
     WarpExecuteOnLane0LoweringOptions options;
     options.warpAllocationFn = allocateGlobalSharedMemory;
     options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,
                                       WarpExecuteOnLane0Op warpOp) {
       builder.create<gpu::BarrierOp>(loc);
     };

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     deps = [
         ":AffineDialect",
         ":ArithmeticDialect",
         ":ArithmeticUtils",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":DialectUtils",
         ":FuncDialect",
-        ":GPUDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":Support",
         ":TensorDialect",
         ":Transforms",