Improve dpp implementation

Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>

Author: Muzammiluddin-Syed-ECE

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

@@ -367,112 +367,112 @@ struct VectorSubgroupReduceToShuffles final
  bool matchClustered = false;
 };
 
-std::optional<Value> createSubgroupDPPReduction(OpBuilder &b, Location loc,
- Value input,
- gpu::AllReduceOperation mode,
- const ClusterInfo &ci,
- amdgpu::Chipset chipset) {
- Value result = input;
+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;
- Value lane0 =
- b.create<arith::ConstantOp>(loc, b.getI32Type(), b.getI32IntegerAttr(0));
- Value lane32 =
- b.create<arith::ConstantOp>(loc, b.getI32Type(), b.getI32IntegerAttr(32));
-
- auto dppReduceAcrossLanes = [&](int numLanes,
- Value res) -> std::optional<Value> {
- Value dppResult, laneVal;
-
- switch (numLanes) {
- case 2:
- // Perform reduction between all lanes N <-> N+1.
- dppResult = b.create<amdgpu::DPPOp>(
- loc, res.getType(), res, res, amdgpu::DPPPerm::quad_perm,
- b.getI32ArrayAttr({1, 0, 3, 2}), allRows, allBanks, boundCtrl);
- break;
- case 4:
- // Perform reduction between all lanes N <-> N+2.
- dppResult = b.create<amdgpu::DPPOp>(
- loc, res.getType(), res, res, amdgpu::DPPPerm::quad_perm,
- b.getI32ArrayAttr({2, 3, 0, 1}), allRows, allBanks, boundCtrl);
- break;
- case 8:
- // Perform reduction between all lanes N <-> 7-N,
- // e.g lane[0] <-> lane[7], lane[1] <-> lane[6]..., lane[3] <-> lane[4].
- dppResult = b.create<amdgpu::DPPOp>(
- loc, res.getType(), res, res, amdgpu::DPPPerm::row_half_mirror,
- b.getUnitAttr(), allRows, allBanks, boundCtrl);
- break;
- case 16:
- // Perform reduction between all lanes N <-> 15-N,
- // e.g lane[0] <-> lane[15], lane[1] <-> lane[14]..., lane[7] <-> lane[8].
- dppResult = b.create<amdgpu::DPPOp>(
- loc, result.getType(), res, res, amdgpu::DPPPerm::row_mirror,
- b.getUnitAttr(), allRows, allBanks, boundCtrl);
- break;
- case 32:
- if (chipset.majorVersion <= 9) {
- // Broadcast last value from each row to next row.
- // Use row mask to avoid polluting rows 1 and 3.
- dppResult = b.create<amdgpu::DPPOp>(loc, res.getType(), res, res,
- amdgpu::DPPPerm::row_bcast_15,
- b.getUnitAttr(), 0xa, allBanks,
- /*bound_ctrl*/ false);
- } else if (chipset.majorVersion <= 12) {
- // Use a permute lane to cross rows (row 1 <-> row 0, row 3 <-> row 2).
- dppResult = b.create<ROCDL::PermlaneX16Op>(loc, res.getType(), res, res,
- -1, -1, /*fi=*/true,
- /*bound_ctrl=*/false);
- if (ci.subgroupSize == 32) {
- dppResult =
- b.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
- }
- } else {
- return std::nullopt;
- }
- break;
- case 64:
- if (chipset.majorVersion <= 9) {
- // Broadcast 31st lane value to rows 2 and 3.
- // Use row mask to avoid polluting rows 0 and 1.
- dppResult = b.create<amdgpu::DPPOp>(loc, res.getType(), res, res,
- amdgpu::DPPPerm::row_bcast_31,
- b.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.
- dppResult =
- b.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane32);
- laneVal = b.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
- return vector::makeArithReduction(
- b, loc, gpu::convertReductionKind(mode), dppResult, laneVal);
- } else {
- return std::nullopt;
+ 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));
+ dpp =
+ rewriter.create<ROCDL::ReadlaneOp>(loc, res.getType(), res, lane0);
  }
- break;
- default:
- // Should never reach here given previous validation of ClusterInfo.
- llvm_unreachable("ERROR: Unexpected cluster size.");
- return std::nullopt;
+ } else {
+ return rewriter.notifyMatchFailure(
+ op, "Subgroup reduce lowering to DPP not currently supported for "
+ "this device.");
  }
- return vector::makeArithReduction(b, loc, gpu::convertReductionKind(mode),
- res, dppResult);
- };
-
- for (unsigned cs = 2; cs <= ci.clusterSize; cs <<= 1) {
- if (auto dpp = dppReduceAcrossLanes(cs, result)) {
- result = *dpp;
- continue;
+ }
+ 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.");
  }
- return std::nullopt;
+ res = vector::makeArithReduction(rewriter, loc,
+ gpu::convertReductionKind(mode), res, dpp);
  }
-
- assert(result.getType() == input.getType());
- return result;
+ assert(res.getType() == input.getType());
+ return res;
 }
 
 /// Collect a set of patterns to lower `gpu.subgroup_reduce` into `amdgpu.dpp`
@@ -500,22 +500,21 @@ struct ScalarSubgroupReduceToDPP final
  return failure();
 
  if (ci->clusterStride != 1)
- return failure();
+ return rewriter.notifyMatchFailure(
+ op, "Supgroup 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");
 
- Location loc = op.getLoc();
- std::optional<Value> dpp = createSubgroupDPPReduction(
- rewriter, loc, op.getValue(), op.getOp(), *ci, chipset);
- if (!dpp)
- return rewriter.notifyMatchFailure(
- op, "Subgroup reduce lowering to DPP not currently supported for "
- "this device.");
+ FailureOr<Value> dpp = createSubgroupDPPReduction(
+ rewriter, op, op.getValue(), op.getOp(), *ci, chipset);
+ if (failed(dpp))
+ return failure();
 
- rewriter.replaceOp(op, *dpp);
+ rewriter.replaceOp(op, dpp.value());
  return success();
  }