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();
   }