[Mlir-commits] [mlir] [MLIR][ROCDL] Add conversion for gpu.subgroup_id to ROCDL (PR #136405)
Alan Li
llvmlistbot at llvm.org
Tue Apr 22 18:23:54 PDT 2025
https://github.com/lialan updated https://github.com/llvm/llvm-project/pull/136405
>From d75c2106ae48c3e86a2a572d2efe715c54cbdcfb Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Sat, 19 Apr 2025 00:51:15 -0400
Subject: [PATCH 1/7] [MLIR][ROCDL] Add conversion for gpu.subgroup_id to ROCDL
Creates `rocdl.wave_id` op with llvm conversion to:
`__builtin_amdgcn_s_get_waveid_in_workgroup`
---
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 8 ++++
.../GPUToROCDL/LowerGpuOpsToROCDLOps.cpp | 48 ++++++++++++++-----
.../Conversion/GPUToROCDL/gpu-to-rocdl.mlir | 10 ++--
mlir/test/Target/LLVMIR/rocdl.mlir | 6 +++
4 files changed, 58 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 186a4f53f93cb..09d22da0d4c72 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -204,6 +204,14 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
}];
}
+// the intrinsic function name is too long so we use a shorter name for rocdl.
+def ROCDL_WaveIdOp : LLVM_IntrOpBase<ROCDL_Dialect, "wave_id",
+ "amdgcn_s_get_waveid_in_workgroup", [], [], [Pure], 1>,
+ Arguments<(ins)> {
+ let results = (outs LLVM_Type:$res);
+ let assemblyFormat = "attr-dict `:` type($res)";
+}
+
//===----------------------------------------------------------------------===//
// Thread index and Block index
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index e6dd6f135884e..315bc7157cd83 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -80,6 +80,24 @@ static constexpr StringLiteral amdgcnDataLayout =
"64-S32-A5-G1-ni:7:8:9";
namespace {
+
+// Truncate or extend the result depending on the index bitwidth specified
+// by the LLVMTypeConverter options.
+static Value truncOrExtToLLVMType(ConversionPatternRewriter &rewriter,
+ Location loc, Value value,
+ const LLVMTypeConverter *converter) {
+ auto intWidth = cast<IntegerType>(value.getType()).getWidth();
+ auto indexBitwidth = converter->getIndexTypeBitwidth();
+ if (indexBitwidth > intWidth) {
+ return rewriter.create<LLVM::SExtOp>(
+ loc, IntegerType::get(rewriter.getContext(), indexBitwidth), value);
+ } else if (indexBitwidth < intWidth) {
+ return rewriter.create<LLVM::TruncOp>(
+ loc, IntegerType::get(rewriter.getContext(), indexBitwidth), value);
+ }
+ return value;
+}
+
struct GPULaneIdOpToROCDL : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
using ConvertOpToLLVMPattern<gpu::LaneIdOp>::ConvertOpToLLVMPattern;
@@ -98,16 +116,7 @@ struct GPULaneIdOpToROCDL : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
rewriter.create<ROCDL::MbcntLoOp>(loc, intTy, ValueRange{minus1, zero});
Value laneId = rewriter.create<ROCDL::MbcntHiOp>(
loc, intTy, ValueRange{minus1, mbcntLo});
- // Truncate or extend the result depending on the index bitwidth specified
- // by the LLVMTypeConverter options.
- const unsigned indexBitwidth = getTypeConverter()->getIndexTypeBitwidth();
- if (indexBitwidth > 32) {
- laneId = rewriter.create<LLVM::SExtOp>(
- loc, IntegerType::get(context, indexBitwidth), laneId);
- } else if (indexBitwidth < 32) {
- laneId = rewriter.create<LLVM::TruncOp>(
- loc, IntegerType::get(context, indexBitwidth), laneId);
- }
+ laneId = truncOrExtToLLVMType(rewriter, loc, laneId, getTypeConverter());
rewriter.replaceOp(op, {laneId});
return success();
}
@@ -190,6 +199,21 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
}
};
+struct GPUSubgroupIdOpToROCDL : ConvertOpToLLVMPattern<gpu::SubgroupIdOp> {
+ using ConvertOpToLLVMPattern<gpu::SubgroupIdOp>::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(gpu::SubgroupIdOp op, gpu::SubgroupIdOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto int32Type = IntegerType::get(rewriter.getContext(), 32);
+ Value waveIdOp = rewriter.create<ROCDL::WaveIdOp>(op.getLoc(), int32Type);
+ waveIdOp = truncOrExtToLLVMType(rewriter, op.getLoc(), waveIdOp,
+ getTypeConverter());
+ rewriter.replaceOp(op, {waveIdOp});
+ return success();
+ }
+};
+
/// Import the GPU Ops to ROCDL Patterns.
#include "GPUToROCDL.cpp.inc"
@@ -405,7 +429,9 @@ void mlir::populateGpuToROCDLConversionPatterns(
// TODO: Add alignment for workgroup memory
patterns.add<GPUDynamicSharedMemoryOpLowering>(converter);
- patterns.add<GPUShuffleOpLowering, GPULaneIdOpToROCDL>(converter);
+ patterns
+ .add<GPUShuffleOpLowering, GPULaneIdOpToROCDL, GPUSubgroupIdOpToROCDL>(
+ converter);
populateMathToROCDLConversionPatterns(converter, patterns);
}
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
index 071cae9d5789f..a06b77dcff038 100644
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
@@ -11,7 +11,7 @@ gpu.module @test_module {
func.func @gpu_index_ops()
-> (index, index, index, index, index, index,
index, index, index, index, index, index,
- index) {
+ index, index) {
// CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i64
// CHECK: rocdl.workitem.id.x : i32
@@ -59,12 +59,16 @@ gpu.module @test_module {
// CHECK: = llvm.sext %{{.*}} : i32 to i64
%laneId = gpu.lane_id
+ // CHECK: = rocdl.wave_id : i32
+ // CHECK: = llvm.sext %{{.*}} : i32 to i64
+ %waveId = gpu.subgroup_id : index
+
func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
%bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
- %laneId
+ %laneId, %waveId
: index, index, index, index, index, index,
index, index, index, index, index, index,
- index
+ index, index
}
}
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 3db1f7b2b6427..f5767dd1fc95a 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -88,6 +88,12 @@ llvm.func @rocdl.lane_id() -> i32 {
llvm.return %3 : i32
}
+llvm.func @rocdl.wave_id() -> i32 {
+ // CHECK: call i32 @llvm.amdgcn.s.get.waveid.in.workgroup()
+ %0 = rocdl.wave_id : i32
+ llvm.return %0 : i32
+}
+
llvm.func @rocdl.swizzle(%src : i32) -> i32 {
// CHECK-LABEL: rocdl.swizzle
// CHECK: call i32 @llvm.amdgcn.ds.swizzle
>From 2fc9d61012255e266eb43c32fc1c5329d512a6e0 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Sat, 19 Apr 2025 23:04:32 -0400
Subject: [PATCH 2/7] update comments; also add a guard to guard off gfx9.
---
.../Conversion/GPUToROCDL/GPUToROCDLPass.h | 4 +-
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 5 ++-
.../GPUToROCDL/LowerGpuOpsToROCDLOps.cpp | 43 ++++++++++++-------
3 files changed, 33 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h b/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h
index 1a917932a9a84..2d8aaf8371627 100644
--- a/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h
+++ b/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h
@@ -10,6 +10,7 @@
#include "mlir/Conversion/GPUToROCDL/Runtimes.h"
#include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
+#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
#include <memory>
namespace mlir {
@@ -32,7 +33,8 @@ class GPUModuleOp;
/// The resulting pattern set should be run over a gpu.module op
void populateGpuToROCDLConversionPatterns(const LLVMTypeConverter &converter,
RewritePatternSet &patterns,
- gpu::amd::Runtime runtime);
+ gpu::amd::Runtime runtime,
+ mlir::amdgpu::Chipset chipset);
/// Configure target to convert from the GPU dialect to ROCDL.
void configureGpuToROCDLConversionLegality(ConversionTarget &target);
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 09d22da0d4c72..f5da5e1fcfa19 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -204,9 +204,10 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
}];
}
-// the intrinsic function name is too long so we use a shorter name for rocdl.
+// The LLVM intrinsic function name is rather mouthful,
+// so here we opt to use a shorter rocdl name.
def ROCDL_WaveIdOp : LLVM_IntrOpBase<ROCDL_Dialect, "wave_id",
- "amdgcn_s_get_waveid_in_workgroup", [], [], [Pure], 1>,
+ "amdgcn_s_get_waveid_in_workgroup", [], [], [], 1>,
Arguments<(ins)> {
let results = (outs LLVM_Type:$res);
let assemblyFormat = "attr-dict `:` type($res)";
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 315bc7157cd83..be3ecae005ff3 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -85,15 +85,15 @@ namespace {
// by the LLVMTypeConverter options.
static Value truncOrExtToLLVMType(ConversionPatternRewriter &rewriter,
Location loc, Value value,
- const LLVMTypeConverter *converter) {
+ const LLVMTypeConverter &converter) {
auto intWidth = cast<IntegerType>(value.getType()).getWidth();
- auto indexBitwidth = converter->getIndexTypeBitwidth();
+ auto indexBitwidth = converter.getIndexTypeBitwidth();
+ auto indexBitwidthType =
+ IntegerType::get(rewriter.getContext(), converter.getIndexTypeBitwidth());
if (indexBitwidth > intWidth) {
- return rewriter.create<LLVM::SExtOp>(
- loc, IntegerType::get(rewriter.getContext(), indexBitwidth), value);
+ return rewriter.create<LLVM::SExtOp>(loc, indexBitwidthType, value);
} else if (indexBitwidth < intWidth) {
- return rewriter.create<LLVM::TruncOp>(
- loc, IntegerType::get(rewriter.getContext(), indexBitwidth), value);
+ return rewriter.create<LLVM::TruncOp>(loc, indexBitwidthType, value);
}
return value;
}
@@ -116,7 +116,7 @@ struct GPULaneIdOpToROCDL : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
rewriter.create<ROCDL::MbcntLoOp>(loc, intTy, ValueRange{minus1, zero});
Value laneId = rewriter.create<ROCDL::MbcntHiOp>(
loc, intTy, ValueRange{minus1, mbcntLo});
- laneId = truncOrExtToLLVMType(rewriter, loc, laneId, getTypeConverter());
+ laneId = truncOrExtToLLVMType(rewriter, loc, laneId, *getTypeConverter());
rewriter.replaceOp(op, {laneId});
return success();
}
@@ -199,16 +199,27 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
}
};
-struct GPUSubgroupIdOpToROCDL : ConvertOpToLLVMPattern<gpu::SubgroupIdOp> {
- using ConvertOpToLLVMPattern<gpu::SubgroupIdOp>::ConvertOpToLLVMPattern;
+struct GPUSubgroupIdOpToROCDL final
+ : ConvertOpToLLVMPattern<gpu::SubgroupIdOp> {
+ using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+ GPUSubgroupIdOpToROCDL(MLIRContext *ctx, mlir::amdgpu::Chipset chipset)
+ : ConvertOpToLLVMPattern(ctx), chipset(chipset) {}
+
+ mlir::amdgpu::Chipset chipset;
LogicalResult
matchAndRewrite(gpu::SubgroupIdOp op, gpu::SubgroupIdOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
+ if (chipset.majorVersion < 10) {
+ return rewriter.notifyMatchFailure(
+ op, "SubgroupIdOp is not yet supported on this architecture");
+ }
+
auto int32Type = IntegerType::get(rewriter.getContext(), 32);
Value waveIdOp = rewriter.create<ROCDL::WaveIdOp>(op.getLoc(), int32Type);
waveIdOp = truncOrExtToLLVMType(rewriter, op.getLoc(), waveIdOp,
- getTypeConverter());
+ *getTypeConverter());
rewriter.replaceOp(op, {waveIdOp});
return success();
}
@@ -343,7 +354,8 @@ struct LowerGpuOpsToROCDLOpsPass final
populateAMDGPUToROCDLConversionPatterns(converter, llvmPatterns,
*maybeChipset);
- populateGpuToROCDLConversionPatterns(converter, llvmPatterns, runtime);
+ populateGpuToROCDLConversionPatterns(converter, llvmPatterns, runtime,
+ *maybeChipset);
configureGpuToROCDLConversionLegality(target);
if (failed(applyPartialConversion(m, target, std::move(llvmPatterns))))
signalPassFailure();
@@ -391,7 +403,7 @@ void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) {
void mlir::populateGpuToROCDLConversionPatterns(
const LLVMTypeConverter &converter, RewritePatternSet &patterns,
- mlir::gpu::amd::Runtime runtime) {
+ mlir::gpu::amd::Runtime runtime, mlir::amdgpu::Chipset chipset) {
using gpu::index_lowering::IndexKind;
using gpu::index_lowering::IntrType;
using mlir::gpu::amd::Runtime;
@@ -429,10 +441,9 @@ void mlir::populateGpuToROCDLConversionPatterns(
// TODO: Add alignment for workgroup memory
patterns.add<GPUDynamicSharedMemoryOpLowering>(converter);
- patterns
- .add<GPUShuffleOpLowering, GPULaneIdOpToROCDL, GPUSubgroupIdOpToROCDL>(
- converter);
-
+ patterns.add<GPUShuffleOpLowering, GPULaneIdOpToROCDL>(converter);
+ patterns.add(
+ std::make_unique<GPUSubgroupIdOpToROCDL>(patterns.getContext(), chipset));
populateMathToROCDLConversionPatterns(converter, patterns);
}
>From 2208dbcfba5751d52cf811f0273e970cc63dcef8 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Sun, 20 Apr 2025 14:31:05 -0400
Subject: [PATCH 3/7] updates
---
.../Conversion/GPUToROCDL/GPUToROCDLPass.h | 6 +---
.../GPUToROCDL/LowerGpuOpsToROCDLOps.cpp | 32 +++++--------------
.../GPUToROCDL/gpu-to-rocdl-chipset.mlir | 13 ++++++++
.../Conversion/GPUToROCDL/gpu-to-rocdl.mlir | 10 ++----
4 files changed, 25 insertions(+), 36 deletions(-)
create mode 100644 mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir
diff --git a/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h b/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h
index 2d8aaf8371627..f53291c04158f 100644
--- a/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h
+++ b/mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h
@@ -43,11 +43,7 @@ void configureGpuToROCDLConversionLegality(ConversionTarget &target);
/// index bitwidth used for the lowering of the device side index computations
/// is configurable.
std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
-createLowerGpuOpsToROCDLOpsPass(
- const std::string &chipset = "gfx900",
- unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout,
- bool useBarePtrCallConv = false,
- gpu::amd::Runtime runtime = gpu::amd::Runtime::Unknown);
+createLowerGpuOpsToROCDLOpsPass();
} // namespace mlir
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index be3ecae005ff3..3d844333ef3b4 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -203,10 +203,11 @@ struct GPUSubgroupIdOpToROCDL final
: ConvertOpToLLVMPattern<gpu::SubgroupIdOp> {
using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
- GPUSubgroupIdOpToROCDL(MLIRContext *ctx, mlir::amdgpu::Chipset chipset)
- : ConvertOpToLLVMPattern(ctx), chipset(chipset) {}
+ GPUSubgroupIdOpToROCDL(const LLVMTypeConverter &converter,
+ const mlir::amdgpu::Chipset &chipset)
+ : ConvertOpToLLVMPattern(converter), chipset(chipset) {}
- mlir::amdgpu::Chipset chipset;
+ const mlir::amdgpu::Chipset chipset;
LogicalResult
matchAndRewrite(gpu::SubgroupIdOp op, gpu::SubgroupIdOp::Adaptor adaptor,
@@ -235,19 +236,7 @@ struct GPUSubgroupIdOpToROCDL final
// code.
struct LowerGpuOpsToROCDLOpsPass final
: public impl::ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> {
- LowerGpuOpsToROCDLOpsPass() = default;
- LowerGpuOpsToROCDLOpsPass(const std::string &chipset, unsigned indexBitwidth,
- bool useBarePtrCallConv,
- gpu::amd::Runtime runtime) {
- if (this->chipset.getNumOccurrences() == 0)
- this->chipset = chipset;
- if (this->indexBitwidth.getNumOccurrences() == 0)
- this->indexBitwidth = indexBitwidth;
- if (this->useBarePtrCallConv.getNumOccurrences() == 0)
- this->useBarePtrCallConv = useBarePtrCallConv;
- if (this->runtime.getNumOccurrences() == 0)
- this->runtime = runtime;
- }
+ using Base::Base;
void getDependentDialects(DialectRegistry ®istry) const override {
Base::getDependentDialects(registry);
@@ -442,16 +431,11 @@ void mlir::populateGpuToROCDLConversionPatterns(
patterns.add<GPUDynamicSharedMemoryOpLowering>(converter);
patterns.add<GPUShuffleOpLowering, GPULaneIdOpToROCDL>(converter);
- patterns.add(
- std::make_unique<GPUSubgroupIdOpToROCDL>(patterns.getContext(), chipset));
+ patterns.add<GPUSubgroupIdOpToROCDL>(converter, chipset);
populateMathToROCDLConversionPatterns(converter, patterns);
}
std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
-mlir::createLowerGpuOpsToROCDLOpsPass(const std::string &chipset,
- unsigned indexBitwidth,
- bool useBarePtrCallConv,
- gpu::amd::Runtime runtime) {
- return std::make_unique<LowerGpuOpsToROCDLOpsPass>(
- chipset, indexBitwidth, useBarePtrCallConv, runtime);
+mlir::createLowerGpuOpsToROCDLOpsPass() {
+ return std::make_unique<LowerGpuOpsToROCDLOpsPass>();
}
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir
new file mode 100644
index 0000000000000..4a5570ae6e0e1
--- /dev/null
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx1030' -split-input-file | FileCheck %s
+
+gpu.module @test_module {
+ // CHECK-LABEL: func @gpu_subgroup_id()
+ func.func @gpu_subgroup_id() -> (index) {
+ // expected-error at +1 {{failed to legalize operation 'gpu.subgroup_id' that was explicitly marked illegal}}
+ // CHECK: = rocdl.wave_id : i32
+ // CHECK: = llvm.sext %{{.*}} : i32 to i64
+ %waveId = gpu.subgroup_id : index
+ func.return %waveId : index
+ }
+}
+
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
index a06b77dcff038..071cae9d5789f 100644
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
@@ -11,7 +11,7 @@ gpu.module @test_module {
func.func @gpu_index_ops()
-> (index, index, index, index, index, index,
index, index, index, index, index, index,
- index, index) {
+ index) {
// CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i64
// CHECK: rocdl.workitem.id.x : i32
@@ -59,16 +59,12 @@ gpu.module @test_module {
// CHECK: = llvm.sext %{{.*}} : i32 to i64
%laneId = gpu.lane_id
- // CHECK: = rocdl.wave_id : i32
- // CHECK: = llvm.sext %{{.*}} : i32 to i64
- %waveId = gpu.subgroup_id : index
-
func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
%bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
- %laneId, %waveId
+ %laneId
: index, index, index, index, index, index,
index, index, index, index, index, index,
- index, index
+ index
}
}
>From 42036bf7eab2fcb3667eea4f98bacaad115954df Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Sun, 20 Apr 2025 18:44:36 -0400
Subject: [PATCH 4/7] updates
---
mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp | 8 +++++---
1 file changed, 5 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 3d844333ef3b4..7a578cf2e7617 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -86,13 +86,15 @@ namespace {
static Value truncOrExtToLLVMType(ConversionPatternRewriter &rewriter,
Location loc, Value value,
const LLVMTypeConverter &converter) {
- auto intWidth = cast<IntegerType>(value.getType()).getWidth();
- auto indexBitwidth = converter.getIndexTypeBitwidth();
+ int64_t intWidth = cast<IntegerType>(value.getType()).getWidth();
+ int64_t indexBitwidth = converter.getIndexTypeBitwidth();
auto indexBitwidthType =
IntegerType::get(rewriter.getContext(), converter.getIndexTypeBitwidth());
+ // TODO: use <=> in C++20
if (indexBitwidth > intWidth) {
return rewriter.create<LLVM::SExtOp>(loc, indexBitwidthType, value);
- } else if (indexBitwidth < intWidth) {
+ }
+ if (indexBitwidth < intWidth) {
return rewriter.create<LLVM::TruncOp>(loc, indexBitwidthType, value);
}
return value;
>From 73176787c0153b84f5863417fe1aa47603091788 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 22 Apr 2025 17:46:10 -0400
Subject: [PATCH 5/7] Redo subgroup id
---
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 5 +--
.../GPUToROCDL/LowerGpuOpsToROCDLOps.cpp | 34 ++++++++++++++-----
.../GPUToROCDL/gpu-to-rocdl-chipset.mlir | 13 -------
.../Conversion/GPUToROCDL/gpu-to-rocdl.mlir | 22 ++++++++++++
4 files changed, 48 insertions(+), 26 deletions(-)
delete mode 100644 mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index f5da5e1fcfa19..6a9b4f1690d57 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -204,10 +204,7 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
}];
}
-// The LLVM intrinsic function name is rather mouthful,
-// so here we opt to use a shorter rocdl name.
-def ROCDL_WaveIdOp : LLVM_IntrOpBase<ROCDL_Dialect, "wave_id",
- "amdgcn_s_get_waveid_in_workgroup", [], [], [], 1>,
+def ROCDL_WaveIdOp : ROCDL_IntrOp<"s.get.waveid.in.workgroup", [], [], [Pure], 1>,
Arguments<(ins)> {
let results = (outs LLVM_Type:$res);
let assemblyFormat = "attr-dict `:` type($res)";
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 7a578cf2e7617..9f1b137970aa7 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -214,16 +214,32 @@ struct GPUSubgroupIdOpToROCDL final
LogicalResult
matchAndRewrite(gpu::SubgroupIdOp op, gpu::SubgroupIdOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- if (chipset.majorVersion < 10) {
- return rewriter.notifyMatchFailure(
- op, "SubgroupIdOp is not yet supported on this architecture");
- }
-
auto int32Type = IntegerType::get(rewriter.getContext(), 32);
- Value waveIdOp = rewriter.create<ROCDL::WaveIdOp>(op.getLoc(), int32Type);
- waveIdOp = truncOrExtToLLVMType(rewriter, op.getLoc(), waveIdOp,
- *getTypeConverter());
- rewriter.replaceOp(op, {waveIdOp});
+ auto loc = op.getLoc();
+ LLVM::IntegerOverflowFlags flags =
+ LLVM::IntegerOverflowFlags::nsw | LLVM::IntegerOverflowFlags::nuw;
+ // w_id.x + w_dim.x * (w_id.y + w_dim.y * w_id.z)) / subgroup_size
+ Value workitemIdX = rewriter.create<ROCDL::ThreadIdXOp>(loc, int32Type);
+ Value workitemIdY = rewriter.create<ROCDL::ThreadIdYOp>(loc, int32Type);
+ Value workitemIdZ = rewriter.create<ROCDL::ThreadIdZOp>(loc, int32Type);
+ Value workitemDimX = rewriter.create<ROCDL::BlockDimXOp>(loc, int32Type);
+ Value workitemDimY = rewriter.create<ROCDL::BlockDimYOp>(loc, int32Type);
+ Value dimYxIdZ = rewriter.create<LLVM::MulOp>(loc, int32Type, workitemDimY,
+ workitemIdZ, flags);
+ Value dimYxIdZPlusIdY = rewriter.create<LLVM::AddOp>(
+ loc, int32Type, dimYxIdZ, workitemIdY, flags);
+ Value dimYxIdZPlusIdYTimesDimX = rewriter.create<LLVM::MulOp>(
+ loc, int32Type, workitemDimX, dimYxIdZPlusIdY, flags);
+ Value workitemIdXPlusDimYxIdZPlusIdYTimesDimX =
+ rewriter.create<LLVM::AddOp>(loc, int32Type, workitemIdX,
+ dimYxIdZPlusIdYTimesDimX, flags);
+ Value subgroupSize = rewriter.create<LLVM::ConstantOp>(
+ loc, IntegerType::get(rewriter.getContext(), 32), 64);
+ Value waveIdOp = rewriter.create<LLVM::SDivOp>(
+ loc, workitemIdXPlusDimYxIdZPlusIdYTimesDimX, subgroupSize);
+
+ rewriter.replaceOp(op, {truncOrExtToLLVMType(rewriter, loc, waveIdOp,
+ *getTypeConverter())});
return success();
}
};
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir
deleted file mode 100644
index 4a5570ae6e0e1..0000000000000
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-chipset.mlir
+++ /dev/null
@@ -1,13 +0,0 @@
-// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx1030' -split-input-file | FileCheck %s
-
-gpu.module @test_module {
- // CHECK-LABEL: func @gpu_subgroup_id()
- func.func @gpu_subgroup_id() -> (index) {
- // expected-error at +1 {{failed to legalize operation 'gpu.subgroup_id' that was explicitly marked illegal}}
- // CHECK: = rocdl.wave_id : i32
- // CHECK: = llvm.sext %{{.*}} : i32 to i64
- %waveId = gpu.subgroup_id : index
- func.return %waveId : index
- }
-}
-
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
index 071cae9d5789f..41a78648d7047 100644
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
@@ -740,3 +740,25 @@ gpu.module @test_module {
gpu.module @test_custom_data_layout attributes {llvm.data_layout = "e"} {
}
+
+// -----
+
+gpu.module @test_module {
+ // CHECK-LABEL: func @gpu_subgroup_id()
+ func.func @gpu_subgroup_id() -> (index) {
+ // CHECK: %[[widx:.*]] = rocdl.workitem.id.x : i32
+ // CHECK: %[[widy:.*]] = rocdl.workitem.id.y : i32
+ // CHECK: %[[widz:.*]] = rocdl.workitem.id.z : i32
+ // CHECK: %[[dimx:.*]] = rocdl.workgroup.dim.x : i32
+ // CHECK: %[[dimy:.*]] = rocdl.workgroup.dim.y : i32
+ // CHECK: %[[int5:.*]] = llvm.mul %[[dimy]], %[[widz]] overflow<nsw, nuw> : i32
+ // CHECK: %[[int6:.*]] = llvm.add %[[int5]], %[[widy]] overflow<nsw, nuw> : i32
+ // CHECK: %[[int7:.*]] = llvm.mul %[[dimx]], %[[int6]] overflow<nsw, nuw> : i32
+ // CHECK: %[[int8:.*]] = llvm.add %[[widx]], %[[int7]] overflow<nsw, nuw> : i32
+ // CHECK: %[[ssize:.*]] = llvm.mlir.constant(64 : i32) : i32
+ // CHECK: = llvm.sdiv %[[int8]], %[[ssize]] : i32
+ // CHECK: = llvm.sext %10 : i32 to i64
+ %subgroupId = gpu.subgroup_id : index
+ func.return %subgroupId : index
+ }
+}
>From 90f6ddefb97078265fa1f108958cfcc32463c2c2 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 22 Apr 2025 18:14:19 -0400
Subject: [PATCH 6/7] Remove rocdl op because we cannot use it.
---
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 6 ------
1 file changed, 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 6a9b4f1690d57..186a4f53f93cb 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -204,12 +204,6 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
}];
}
-def ROCDL_WaveIdOp : ROCDL_IntrOp<"s.get.waveid.in.workgroup", [], [], [Pure], 1>,
- Arguments<(ins)> {
- let results = (outs LLVM_Type:$res);
- let assemblyFormat = "attr-dict `:` type($res)";
-}
-
//===----------------------------------------------------------------------===//
// Thread index and Block index
//===----------------------------------------------------------------------===//
>From 25ab466321f1bac97ac99c23f03650ccf8a46106 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 22 Apr 2025 21:23:33 -0400
Subject: [PATCH 7/7] remove unneeded test case.
---
mlir/test/Target/LLVMIR/rocdl.mlir | 6 ------
1 file changed, 6 deletions(-)
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index f5767dd1fc95a..3db1f7b2b6427 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -88,12 +88,6 @@ llvm.func @rocdl.lane_id() -> i32 {
llvm.return %3 : i32
}
-llvm.func @rocdl.wave_id() -> i32 {
- // CHECK: call i32 @llvm.amdgcn.s.get.waveid.in.workgroup()
- %0 = rocdl.wave_id : i32
- llvm.return %0 : i32
-}
-
llvm.func @rocdl.swizzle(%src : i32) -> i32 {
// CHECK-LABEL: rocdl.swizzle
// CHECK: call i32 @llvm.amdgcn.ds.swizzle
More information about the Mlir-commits
mailing list