[Mlir-commits] [mlir] [mlir][gpu][nvvm] Lower subgroup query ops to NVVM (PR #209091)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Jul 12 23:32:36 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Johnny Lin (johnny19436)
<details>
<summary>Changes</summary>
`convert-gpu-to-nvvm` currently lacks lowering coverage for `gpu.subgroup_id`, `gpu.num_subgroups`, and `gpu.subgroup_size`, while these subgroup query ops are already part of the GPU dialect.
This PR adds NVVM lowering for the three ops by mapping them to NVVM special registers (`warpid`, `nwarpid`, and `warpsize`), propagates `upper_bound` to range attributes where representable, and keeps index-bitwidth handling consistent with existing GPUToNVVM lowering patterns.
It also adds/extends `GPUToNVVM` conversion tests to cover:
- positive lowering of all three subgroup query ops,
- `upper_bound`/range behavior, and
- index-bitwidth=32 coverage.
Tested with:
- `ninja -C build mlir-opt`
- `build/bin/llvm-lit -sv mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir`
---
Full diff: https://github.com/llvm/llvm-project/pull/209091.diff
3 Files Affected:
- (modified) mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp (+114-13)
- (modified) mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir (+6-3)
- (modified) mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir (+21-6)
``````````diff
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 80420c26537c3..52271ddd333f0 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -35,9 +35,12 @@
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/MathExtras.h"
+
#include "../GPUCommon/GPUOpsLowering.h"
#include "../GPUCommon/IndexIntrinsicsOpLowering.h"
#include "../GPUCommon/OpToFuncCallLowering.h"
+#include <limits>
#include <optional>
namespace mlir {
@@ -101,6 +104,31 @@ static constexpr llvm::StringLiteral kNVVMNamedBarrierIdPrefix =
static constexpr int32_t kNVVMFirstNamedBarrierId = 1;
static constexpr int32_t kNVVMLastNamedBarrierId = 15;
static constexpr int32_t kNVVMWarpSize = 32;
+static constexpr uint32_t kMaxDim = std::numeric_limits<uint32_t>::max();
+static constexpr uint32_t kMaxSubgroupSize = 128;
+
+// Truncate or extend the result depending on the index bitwidth specified by
+// the LLVMTypeConverter options.
+static Value truncOrExtToIndexBitwidth(ConversionPatternRewriter &rewriter,
+ Location loc, Value value,
+ const LLVMTypeConverter &converter) {
+ auto valueType = cast<IntegerType>(value.getType());
+ unsigned valueBitwidth = valueType.getWidth();
+ unsigned indexBitwidth = converter.getIndexTypeBitwidth();
+ if (indexBitwidth > valueBitwidth) {
+ return LLVM::SExtOp::create(rewriter, loc,
+ IntegerType::get(rewriter.getContext(),
+ indexBitwidth),
+ value);
+ }
+ if (indexBitwidth < valueBitwidth) {
+ return LLVM::TruncOp::create(rewriter, loc,
+ IntegerType::get(rewriter.getContext(),
+ indexBitwidth),
+ value);
+ }
+ return value;
+}
static FailureOr<StringAttr>
createNVVMNamedBarrierIdGlobal(gpu::InitializeNamedBarrierOp op,
@@ -256,7 +284,6 @@ struct GPULaneIdOpToNVVM : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
matchAndRewrite(gpu::LaneIdOp op, gpu::LaneIdOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op->getLoc();
- MLIRContext *context = rewriter.getContext();
LLVM::ConstantRangeAttr bounds = nullptr;
if (std::optional<APInt> upperBound = op.getUpperBound())
bounds = rewriter.getAttr<LLVM::ConstantRangeAttr>(
@@ -266,21 +293,93 @@ struct GPULaneIdOpToNVVM : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
/*bitWidth=*/32, /*lower=*/0, /*upper=*/kWarpSize);
Value newOp =
NVVM::LaneIdOp::create(rewriter, loc, rewriter.getI32Type(), bounds);
- // Truncate or extend the result depending on the index bitwidth specified
- // by the LLVMTypeConverter options.
- const unsigned indexBitwidth = getTypeConverter()->getIndexTypeBitwidth();
- if (indexBitwidth > 32) {
- newOp = LLVM::SExtOp::create(
- rewriter, loc, IntegerType::get(context, indexBitwidth), newOp);
- } else if (indexBitwidth < 32) {
- newOp = LLVM::TruncOp::create(
- rewriter, loc, IntegerType::get(context, indexBitwidth), newOp);
- }
+ newOp = truncOrExtToIndexBitwidth(rewriter, loc, newOp, *getTypeConverter());
rewriter.replaceOp(op, {newOp});
return success();
}
};
+struct GPUSubgroupSizeOpToNVVM
+ : ConvertOpToLLVMPattern<gpu::SubgroupSizeOp> {
+ using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(gpu::SubgroupSizeOp op, gpu::SubgroupSizeOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ LLVM::ConstantRangeAttr bounds = nullptr;
+ if (std::optional<APInt> upperBound = op.getUpperBound()) {
+ uint32_t subgroupUpperBound =
+ static_cast<uint32_t>(upperBound->getLimitedValue(kMaxDim));
+ bounds = rewriter.getAttr<LLVM::ConstantRangeAttr>(
+ /*bitWidth=*/32, /*lower=*/1,
+ /*upper=*/llvm::SaturatingAdd(subgroupUpperBound, uint32_t{1}));
+ } else {
+ bounds = rewriter.getAttr<LLVM::ConstantRangeAttr>(
+ /*bitWidth=*/32, /*lower=*/1, /*upper=*/kMaxSubgroupSize + 1);
+ }
+
+ Value subgroupSize =
+ NVVM::WarpSizeOp::create(rewriter, op.getLoc(), rewriter.getI32Type(),
+ bounds);
+ subgroupSize = truncOrExtToIndexBitwidth(rewriter, op.getLoc(), subgroupSize,
+ *getTypeConverter());
+ rewriter.replaceOp(op, subgroupSize);
+ return success();
+ }
+};
+
+struct GPUSubgroupIdOpToNVVM : ConvertOpToLLVMPattern<gpu::SubgroupIdOp> {
+ using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(gpu::SubgroupIdOp op, gpu::SubgroupIdOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ // TODO: Derive tighter bounds from known launch dimensions when available.
+ uint32_t upperBound = kMaxDim;
+ if (std::optional<APInt> specifiedBound = op.getUpperBound())
+ upperBound = static_cast<uint32_t>(specifiedBound->getLimitedValue(kMaxDim));
+ LLVM::ConstantRangeAttr bounds = rewriter.getAttr<LLVM::ConstantRangeAttr>(
+ /*bitWidth=*/32, /*lower=*/0, /*upper=*/upperBound);
+
+ Value subgroupId =
+ NVVM::WarpIdOp::create(rewriter, op.getLoc(), rewriter.getI32Type(),
+ bounds);
+ subgroupId = truncOrExtToIndexBitwidth(rewriter, op.getLoc(), subgroupId,
+ *getTypeConverter());
+ rewriter.replaceOp(op, subgroupId);
+ return success();
+ }
+};
+
+struct GPUNumSubgroupsOpToNVVM
+ : ConvertOpToLLVMPattern<gpu::NumSubgroupsOp> {
+ using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+ LogicalResult matchAndRewrite(gpu::NumSubgroupsOp op,
+ gpu::NumSubgroupsOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const
+ override {
+ LLVM::ConstantRangeAttr bounds = nullptr;
+ if (std::optional<APInt> upperBound = op.getUpperBound()) {
+ uint32_t numSubgroupsUpperBound =
+ static_cast<uint32_t>(upperBound->getLimitedValue(kMaxDim));
+ bounds = rewriter.getAttr<LLVM::ConstantRangeAttr>(
+ /*bitWidth=*/32, /*lower=*/1,
+ /*upper=*/llvm::SaturatingAdd(numSubgroupsUpperBound, uint32_t{1}));
+ }
+ // Without an explicit upper bound, the target-independent default for
+ // gpu.num_subgroups cannot be encoded as a finite i32 upper-exclusive range.
+
+ Value numSubgroups =
+ NVVM::WarpDimOp::create(rewriter, op.getLoc(), rewriter.getI32Type(),
+ bounds);
+ numSubgroups = truncOrExtToIndexBitwidth(rewriter, op.getLoc(),
+ numSubgroups, *getTypeConverter());
+ rewriter.replaceOp(op, numSubgroups);
+ return success();
+ }
+};
+
struct GPUBallotOpToNVVM : public ConvertOpToLLVMPattern<gpu::BallotOp> {
using ConvertOpToLLVMPattern<gpu::BallotOp>::ConvertOpToLLVMPattern;
@@ -666,8 +765,10 @@ void mlir::populateGpuToNVVMConversionPatterns(
patterns.add<gpu::index_lowering::OpLowering<
gpu::GridDimOp, NVVM::GridDimXOp, NVVM::GridDimYOp, NVVM::GridDimZOp>>(
converter, IndexKind::Grid, IntrType::Dim, benefit);
- patterns.add<GPULaneIdOpToNVVM, GPUBallotOpToNVVM, GPUShuffleOpLowering,
- GPUReturnOpLowering>(converter, benefit);
+ patterns.add<GPULaneIdOpToNVVM, GPUSubgroupIdOpToNVVM,
+ GPUNumSubgroupsOpToNVVM, GPUSubgroupSizeOpToNVVM,
+ GPUBallotOpToNVVM, GPUShuffleOpLowering, GPUReturnOpLowering>(
+ converter, benefit);
patterns.add<GPUDynamicSharedMemoryOpLowering>(
converter, NVVM::kSharedMemoryAlignmentBit, benefit);
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir
index 1bc7acfe46a4a..62bb55af5b026 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir
@@ -7,7 +7,7 @@ gpu.module @test_module_0 {
func.func @gpu_index_ops()
-> (index, index, index, index, index, index,
index, index, index, index, index, index,
- index) {
+ index, index, index, index) {
%tIdX = gpu.thread_id x
%tIdY = gpu.thread_id y
%tIdZ = gpu.thread_id z
@@ -26,13 +26,16 @@ gpu.module @test_module_0 {
// CHECK-NOT: = llvm.sext %{{.*}} : i32 to i64
%laneId = gpu.lane_id
+ %subgroupSize = gpu.subgroup_size : index
+ %subgroupId = gpu.subgroup_id : index
+ %numSubgroups = gpu.num_subgroups : index
func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
%bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
- %laneId
+ %laneId, %subgroupSize, %subgroupId, %numSubgroups
: index, index, index, index, index, index,
index, index, index, index, index, index,
- index
+ index, index, index, index
}
}
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index b96069ac41a44..c18bb0e363d64 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -9,7 +9,7 @@ gpu.module @test_module_0 {
func.func @gpu_index_ops()
-> (index, index, index, index, index, index,
index, index, index, index, index, index,
- index) {
+ index, index, index, index) {
// CHECK: = nvvm.read.ptx.sreg.tid.x : i32
// CHECK: = llvm.sext %{{.*}} : i32 to i64
@@ -55,13 +55,22 @@ gpu.module @test_module_0 {
// CHECK: = nvvm.read.ptx.sreg.laneid range <i32, 0, 32> : i32
// CHECK: = llvm.sext %{{.*}} : i32 to i64
%laneId = gpu.lane_id
+ // CHECK: = nvvm.read.ptx.sreg.warpsize range <i32, 1, 129> : i32
+ // CHECK: = llvm.sext %{{.*}} : i32 to i64
+ %subgroupSize = gpu.subgroup_size : index
+ // CHECK: = nvvm.read.ptx.sreg.warpid range <i32, 0, -1> : i32
+ // CHECK: = llvm.sext %{{.*}} : i32 to i64
+ %subgroupId = gpu.subgroup_id : index
+ // CHECK: = nvvm.read.ptx.sreg.nwarpid : i32
+ // CHECK: = llvm.sext %{{.*}} : i32 to i64
+ %numSubgroups = gpu.num_subgroups : index
func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
%bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
- %laneId
+ %laneId, %subgroupSize, %subgroupId, %numSubgroups
: index, index, index, index, index, index,
index, index, index, index, index, index,
- index
+ index, index, index, index
}
}
@@ -998,15 +1007,21 @@ gpu.module @test_module_48 {
gpu.module @test_module_49 {
// CHECK-LABEL: func @explicit_id_bounds()
- func.func @explicit_id_bounds() -> (index, index, index) {
+ func.func @explicit_id_bounds() -> (index, index, index, index, index, index) {
// CHECK: = nvvm.read.ptx.sreg.tid.x range <i32, 0, 32> : i32
%0 = gpu.thread_id x upper_bound 32
// CHECK: = nvvm.read.ptx.sreg.ntid.x range <i32, 1, 33> : i32
%1 = gpu.block_dim x upper_bound 32
// CHECK: = nvvm.read.ptx.sreg.laneid range <i32, 0, 16> : i32
%2 = gpu.lane_id upper_bound 16
-
- return %0, %1, %2 : index, index, index
+ // CHECK: = nvvm.read.ptx.sreg.warpsize range <i32, 1, 33> : i32
+ %3 = gpu.subgroup_size upper_bound 32 : index
+ // CHECK: = nvvm.read.ptx.sreg.warpid range <i32, 0, 16> : i32
+ %4 = gpu.subgroup_id upper_bound 16 : index
+ // CHECK: = nvvm.read.ptx.sreg.nwarpid range <i32, 1, 17> : i32
+ %5 = gpu.num_subgroups upper_bound 16 : index
+
+ return %0, %1, %2, %3, %4, %5 : index, index, index, index, index, index
}
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/209091
More information about the Mlir-commits
mailing list