[mlir][AMDGPU] Implement gpu.subgroup_reduce with DPP intrinsics on AMD GPUs (#133204)

When performing cross-lane reductions using subgroup_reduce ops across contiguous lanes on AMD GPUs, lower to Data Parallel Primitives (DPP) ops when possible. This reduces latency on applicable devices. See related [Issue](iree-org/iree#20007) To do: - Improve lowering to subgroup_reduce in compatible matvecs (these get directly lowered to gpu.shuffles in an earlier pass) --------- Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>

Author: Muzammiluddin-Syed-ECE

mlir/include/mlir/Dialect/GPU/Transforms/Passes.h

@@ -13,6 +13,7 @@
 #ifndef MLIR_DIALECT_GPU_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_GPU_TRANSFORMS_PASSES_H_
 
+#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Utils/GPUUtils.h"
 #include "mlir/IR/PatternMatch.h"
@@ -68,6 +69,20 @@ void populateGpuLowerClusteredSubgroupReduceToShufflePatterns(
  RewritePatternSet &patterns, unsigned subgroupSize,
  unsigned shuffleBitwidth = 32, PatternBenefit benefit = 1);
 
+/// Collect a set of patterns to lower `gpu.subgroup_reduce` into `amdgpu.dpp`
+/// ops over scalar types. Assumes that the subgroup has
+/// `subgroupSize` lanes. Applicable only to AMD GPUs.
+void populateGpuLowerSubgroupReduceToDPPPatterns(RewritePatternSet &patterns,
+ unsigned subgroupSize,
+ amdgpu::Chipset chipset,
+ PatternBenefit benefit = 1);
+
+/// Disjoint counterpart of `populateGpuLowerSubgroupReduceToDPPPatterns`
+/// that only matches `gpu.subgroup_reduce` ops with a `cluster_size`.
+void populateGpuLowerClusteredSubgroupReduceToDPPPatterns(
+ RewritePatternSet &patterns, unsigned subgroupSize, amdgpu::Chipset chipset,
+ PatternBenefit benefit = 1);
+
 /// Collect all patterns to rewrite ops within the GPU dialect.
 inline void populateGpuRewritePatterns(RewritePatternSet &patterns) {
  populateGpuAllReducePatterns(patterns);

mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp

@@ -10,15 +10,19 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
+#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Utils/GPUUtils.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include <cassert>
@@ -362,6 +366,163 @@ struct VectorSubgroupReduceToShuffles final
  unsigned shuffleBitwidth = 0;
  bool matchClustered = false;
 };
+
+static FailureOr<Value>
+createSubgroupDPPReduction(PatternRewriter &rewriter, gpu::SubgroupReduceOp op,
+ Value input, gpu::AllReduceOperation mode,
+ const ClusterInfo &ci, amdgpu::Chipset chipset) {
+ Location loc = op.getLoc();
+ Value dpp;
+ Value res = input;
+ constexpr int allRows = 0xf;
+ constexpr int allBanks = 0xf;
+ const bool boundCtrl = true;
+ if (ci.clusterSize >= 2) {
+ // Perform reduction between all lanes N <-> N+1.
+ dpp = rewriter.create<amdgpu::DPPOp>(
+ loc, res.getType(), res, res, amdgpu::DPPPerm::quad_perm,
+ rewriter.getI32ArrayAttr({1, 0, 3, 2}), allRows, allBanks, boundCtrl);
+ res = vector::makeArithReduction(rewriter, loc,
+ gpu::convertReductionKind(mode), res, dpp);
+ }
+
+ if (ci.clusterSize >= 4) {
+ // Perform reduction between all lanes N <-> N+2.
+ dpp = rewriter.create<amdgpu::DPPOp>(
+ loc, res.getType(), res, res, amdgpu::DPPPerm::quad_perm,
+ rewriter.getI32ArrayAttr({2, 3, 0, 1}), allRows, allBanks, boundCtrl);
+ res = vector::makeArithReduction(rewriter, loc,
+ gpu::convertReductionKind(mode), res, dpp);
+ }
+ if (ci.clusterSize >= 8) {
+ // Perform reduction between all lanes N <-> 7-N,
+ // e.g lane[0] <-> lane[7], lane[1] <-> lane[6]..., lane[3] <-> lane[4].
+ dpp = rewriter.create<amdgpu::DPPOp>(
+ loc, res.getType(), res, res, amdgpu::DPPPerm::row_half_mirror,
+ rewriter.getUnitAttr(), allRows, allBanks, boundCtrl);
+ res = vector::makeArithReduction(rewriter, loc,
+ gpu::convertReductionKind(mode), res, dpp);
+ }
+ if (ci.clusterSize >= 16) {
+ // Perform reduction between all lanes N <-> 15-N,
+ // e.g lane[0] <-> lane[15], lane[1] <-> lane[14]..., lane[7] <-> lane[8].
+ dpp = rewriter.create<amdgpu::DPPOp>(
+ loc, res.getType(), res, res, amdgpu::DPPPerm::row_mirror,
+ rewriter.getUnitAttr(), allRows, allBanks, boundCtrl);
+ res = vector::makeArithReduction(rewriter, loc,
+ gpu::convertReductionKind(mode), res, dpp);
+ }
+ if (ci.clusterSize >= 32) {
+ if (chipset.majorVersion <= 9) {
+ // Broadcast last value from each row to next row.
+ // Use row mask to avoid polluting rows 1 and 3.
+ dpp = rewriter.create<amdgpu::DPPOp>(
+ loc, res.getType(), res, res, amdgpu::DPPPerm::row_bcast_15,
+ rewriter.getUnitAttr(), 0xa, allBanks,
+ /*bound_ctrl*/ false);
+ res = vector::makeArithReduction(
+ rewriter, loc, gpu::convertReductionKind(mode), res, dpp);
+ } else if (chipset.majorVersion <= 12) {
+ // Use a permute lane to cross rows (row 1 <-> row 0, row 3 <-> row 2).
+ Value uint32Max = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-1));
+ dpp = rewriter.create<ROCDL::PermlaneX16Op>(loc, res.getType(), res, res,
+ uint32Max, uint32Max,
+ /*fi=*/true,
+ /*bound_ctrl=*/false);
+ res = vector::makeArithReduction(
+ rewriter, loc, gpu::convertReductionKind(mode), res, dpp);
+ if (ci.subgroupSize == 32) {
+ Value lane0 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
+ res =
+ rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
+ }
+ } else {
+ return rewriter.notifyMatchFailure(
+ op, "Subgroup reduce lowering to DPP not currently supported for "
+ "this device.");
+ }
+ }
+ if (ci.clusterSize >= 64) {
+ if (chipset.majorVersion <= 9) {
+ // Broadcast 31st lane value to rows 2 and 3.
+ // Use row mask to avoid polluting rows 0 and 1.
+ dpp = rewriter.create<amdgpu::DPPOp>(
+ loc, res.getType(), res, res, amdgpu::DPPPerm::row_bcast_31,
+ rewriter.getUnitAttr(), 0xc, allBanks,
+ /*bound_ctrl*/ false);
+
+ } else if (chipset.majorVersion <= 12) {
+ // Assume reduction across 32 lanes has been done.
+ // Perform final reduction manually by summing values in lane 0 and
+ // lane 32.
+ Value lane0 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
+ Value lane32 = rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(32));
+ dpp = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane32);
+ res = rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
+ } else {
+ return rewriter.notifyMatchFailure(
+ op, "Subgroup reduce lowering to DPP not currently supported for "
+ "this device.");
+ }
+ res = vector::makeArithReduction(rewriter, loc,
+ gpu::convertReductionKind(mode), res, dpp);
+ }
+ assert(res.getType() == input.getType());
+ return res;
+}
+
+/// Collect a set of patterns to lower `gpu.subgroup_reduce` into `amdgpu.dpp`
+/// ops over scalar types. Assumes that the subgroup has
+/// `subgroupSize` lanes. Applicable only to AMD GPUs.
+struct ScalarSubgroupReduceToDPP final
+ : OpRewritePattern<gpu::SubgroupReduceOp> {
+ ScalarSubgroupReduceToDPP(MLIRContext *ctx, unsigned subgroupSize,
+ bool matchClustered, amdgpu::Chipset chipset,
+ PatternBenefit benefit)
+ : OpRewritePattern(ctx, benefit), subgroupSize(subgroupSize),
+ matchClustered(matchClustered), chipset(chipset) {}
+
+ LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.getClusterSize().has_value() != matchClustered) {
+ return rewriter.notifyMatchFailure(
+ op, llvm::formatv("op is {0}clustered but pattern is configured to "
+ "only match {1}clustered ops",
+ matchClustered ? "non-" : "",
+ matchClustered ? "" : "non-"));
+ }
+ auto ci = getAndValidateClusterInfo(op, subgroupSize);
+ if (failed(ci))
+ return failure();
+
+ if (ci->clusterStride != 1)
+ return rewriter.notifyMatchFailure(
+ op, "Subgroup reductions using DPP are currently only available for "
+ "clusters of contiguous lanes.");
+
+ Type valueTy = op.getType();
+ if (!valueTy.isIntOrFloat())
+ return rewriter.notifyMatchFailure(
+ op, "Value type is not a compatible scalar.");
+
+ FailureOr<Value> dpp = createSubgroupDPPReduction(
+ rewriter, op, op.getValue(), op.getOp(), *ci, chipset);
+ if (failed(dpp))
+ return failure();
+
+ rewriter.replaceOp(op, dpp.value());
+ return success();
+ }
+
+private:
+ unsigned subgroupSize = 0;
+ bool matchClustered = false;
+ amdgpu::Chipset chipset;
+};
 } // namespace
 
 void mlir::populateGpuBreakDownSubgroupReducePatterns(
@@ -372,6 +533,22 @@ void mlir::populateGpuBreakDownSubgroupReducePatterns(
  patterns.add<ScalarizeSingleElementReduce>(patterns.getContext(), benefit);
 }
 
+void mlir::populateGpuLowerSubgroupReduceToDPPPatterns(
+ RewritePatternSet &patterns, unsigned subgroupSize, amdgpu::Chipset chipset,
+ PatternBenefit benefit) {
+ patterns.add<ScalarSubgroupReduceToDPP>(patterns.getContext(), subgroupSize,
+ /*matchClustered=*/false, chipset,
+ benefit);
+}
+
+void mlir::populateGpuLowerClusteredSubgroupReduceToDPPPatterns(
+ RewritePatternSet &patterns, unsigned subgroupSize, amdgpu::Chipset chipset,
+ PatternBenefit benefit) {
+ patterns.add<ScalarSubgroupReduceToDPP>(patterns.getContext(), subgroupSize,
+ /*matchClustered=*/true, chipset,
+ benefit);
+}
+
 void mlir::populateGpuLowerSubgroupReduceToShufflePatterns(
  RewritePatternSet &patterns, unsigned subgroupSize,
  unsigned shuffleBitwidth, PatternBenefit benefit) {