Skip to content

Commit b4855f9

Browse files
krzysz00krzysz00
committed
[mlir][ROCDL] Annotate lane ID functions with noundef, ranges
Now that we have general support for setting argument and result attributes on LLVM intrinsics, extend the definitions of mbcnt.lo and mbcnt.hi to carry such attributes. With that, update the construction of the mbcnt.lo/mbcnt.hi calls used to get the lane ID to be `noundef` (since the lane ID is always defined) and to be annotated with the correct ranges (so that generic LLVM passes can correctly optimized based on the fact that there are never more than 32/64 lanes). (Also, handle a pattern that wasn't using getLaneId() and get rid of a dead argument)
1 parent d74d841 commit b4855f9

File tree

4 files changed

+51
-35
lines changed

4 files changed

+51
-35
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -93,19 +93,22 @@ class ROCDL_IntrPure1Op<string mnemonic> :
9393

9494
class ROCDL_IntrOp<string mnemonic, list<int> overloadedResults,
9595
list<int> overloadedOperands, list<Trait> traits, int numResults,
96-
int requiresAccessGroup = 0, int requiresAliasAnalysis = 0, list<int> immArgPositions = [],
96+
int requiresAccessGroup = 0, int requiresAliasAnalysis = 0,
97+
int requiresArgAndResultAttrs = 0,
98+
list<int> immArgPositions = [],
9799
list<string> immArgAttrNames = []> :
98100
LLVM_IntrOpBase<ROCDL_Dialect, mnemonic,
99101
"amdgcn_" # !subst(".", "_", mnemonic), overloadedResults,
100102
overloadedOperands, traits, numResults, requiresAccessGroup,
101-
requiresAliasAnalysis, 0, 0, 0, immArgPositions, immArgAttrNames>;
103+
requiresAliasAnalysis, 0, requiresArgAndResultAttrs, 0,
104+
immArgPositions, immArgAttrNames>;
102105

103106
// Subclass to save typing and ease readibility when there aren't overloaded
104107
// operands or memory accesses.
105108
class ROCDL_ConcreteNonMemIntrOp<string mnemonic, list<Trait> traits,
106109
int numResults, list<int> immArgPositions = [],
107110
list<string> immArgNames = []>
108-
: ROCDL_IntrOp<mnemonic, [], [], traits, numResults, 0, 0,
111+
: ROCDL_IntrOp<mnemonic, [], [], traits, numResults, 0, 0, 0,
109112
immArgPositions, immArgNames>;
110113
//===----------------------------------------------------------------------===//
111114
// ROCDL special register op definitions
@@ -148,8 +151,11 @@ class ROCDL_DimGetterFunctionOp<string mnemonic, string device_function,
148151
//===----------------------------------------------------------------------===//
149152

150153
class ROCDL_MbcntOp<string mnemonic> :
151-
ROCDL_IntrPure1Op<"mbcnt." # mnemonic>,
152-
Arguments<(ins I32:$in0, I32:$in1)> {
154+
ROCDL_IntrOp<"mbcnt." # mnemonic, [], [], [Pure], 1,
155+
0, 0, /*requiresArgAndResultAttrs=*/1> {
156+
dag args = (ins I32:$in0, I32:$in1);
157+
let arguments = !con(args, baseArgs);
158+
let results = (outs I32:$res);
153159
let assemblyFormat = [{
154160
$in0 `,` $in1 attr-dict `:` `(` type($in0) `,` type($in1) `)` `->` type($res)
155161
}];
@@ -501,7 +507,7 @@ def ROCDL_ds_read_tr16_b64 : ROCDL_LDS_Read_Tr_IntrOp<"ds.read.tr16.b64">;
501507
//===---------------------------------------------------------------------===//
502508

503509
def ROCDL_LoadToLDSOp :
504-
ROCDL_IntrOp<"load.to.lds", [], [0], [], 0, 0, 1, [2, 3, 4], ["size", "offset", "aux"]> {
510+
ROCDL_IntrOp<"load.to.lds", [], [0], [], 0, 0, 1, 0, [2, 3, 4], ["size", "offset", "aux"]> {
505511
dag args = (ins Arg<LLVM_AnyPointer, "", [MemRead]>:$globalPtr,
506512
Arg<ROCDLBufferLDS, "", [MemWrite]>:$ldsPtr,
507513
I32Attr:$size,
@@ -520,7 +526,7 @@ def ROCDL_LoadToLDSOp :
520526
}
521527

522528
def ROCDL_GlobalLoadLDSOp :
523-
ROCDL_IntrOp<"global.load.lds", [], [], [], 0, 0, 1, [2, 3, 4], ["size", "offset", "aux"]> {
529+
ROCDL_IntrOp<"global.load.lds", [], [], [], 0, 0, 1, 0, [2, 3, 4], ["size", "offset", "aux"]> {
524530
dag args = (ins Arg<ROCDLGlobalBuffer, "", [MemRead]>:$globalPtr,
525531
Arg<ROCDLBufferLDS, "", [MemWrite]>:$ldsPtr,
526532
I32Attr:$size,
@@ -734,7 +740,7 @@ def ROCDL_RawBufferAtomicUMinOp :
734740

735741
// DPP Update intrinsic
736742
def ROCDL_DPPUpdateOp : ROCDL_IntrOp<"update.dpp", [], [0],
737-
[AllTypesMatch<["res", "src", "old"]>], 1, 0, 0,
743+
[AllTypesMatch<["res", "src", "old"]>], 1, 0, 0, 0,
738744
[2, 3, 4, 5], ["dppCtrl", "rowMask", "bankMask", "boundCtrl"]>,
739745
Arguments<(ins LLVM_Type:$old, LLVM_Type:$src, I32Attr:$dppCtrl, I32Attr:$rowMask,
740746
I32Attr:$bankMask, I1Attr:$boundCtrl)> {
@@ -746,7 +752,7 @@ def ROCDL_DPPUpdateOp : ROCDL_IntrOp<"update.dpp", [], [0],
746752

747753
// PermLaneX16 intrinsic operation
748754
def ROCDL_PermlaneX16Op : ROCDL_IntrOp<"permlanex16", [], [0],
749-
[AllTypesMatch<["res", "old", "src0"]>, AllTypesMatch<["src1", "src2"]>], 1, 0, 0,
755+
[AllTypesMatch<["res", "old", "src0"]>, AllTypesMatch<["src1", "src2"]>], 1, 0, 0, 0,
750756
[4, 5], ["fi", "boundControl"]>,
751757
Arguments<(ins LLVM_Type:$old, LLVM_Type:$src0, LLVM_Type:$src1, LLVM_Type:$src2,
752758
I1Attr:$fi, I1Attr:$boundControl)> {

‎mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp‎

Lines changed: 29 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -79,17 +79,30 @@ static bool canBeCalledWithBarePointers(gpu::GPUFuncOp func) {
7979
return canBeBare;
8080
}
8181

82-
static Value getLaneId(ConversionPatternRewriter &rewriter, Location loc,
83-
const unsigned indexBitwidth) {
82+
static Value getLaneId(RewriterBase &rewriter, Location loc) {
8483
auto int32Type = IntegerType::get(rewriter.getContext(), 32);
8584
Value zero = arith::ConstantIntOp::create(rewriter, loc, 0, 32);
8685
Value minus1 = arith::ConstantIntOp::create(rewriter, loc, -1, 32);
87-
Value mbcntLo = ROCDL::MbcntLoOp::create(rewriter, loc, int32Type,
88-
ValueRange{minus1, zero});
89-
Value laneId = ROCDL::MbcntHiOp::create(rewriter, loc, int32Type,
90-
ValueRange{minus1, mbcntLo});
86+
NamedAttribute noundef = rewriter.getNamedAttr(
87+
LLVM::LLVMDialect::getNoUndefAttrName(), rewriter.getUnitAttr());
88+
NamedAttribute lowRange = rewriter.getNamedAttr(
89+
LLVM::LLVMDialect::getRangeAttrName(),
90+
LLVM::ConstantRangeAttr::get(rewriter.getContext(), APInt::getZero(32),
91+
APInt(32, 32)));
92+
NamedAttribute highRange = rewriter.getNamedAttr(
93+
LLVM::LLVMDialect::getRangeAttrName(),
94+
LLVM::ConstantRangeAttr::get(rewriter.getContext(), APInt::getZero(32),
95+
APInt(32, 64)));
96+
Value mbcntLo = ROCDL::MbcntLoOp::create(
97+
rewriter, loc, int32Type, minus1, zero, /*arg_attrs=*/{},
98+
/*res_attrs=*/
99+
rewriter.getArrayAttr(rewriter.getDictionaryAttr({noundef, lowRange})));
100+
Value laneId = ROCDL::MbcntHiOp::create(
101+
rewriter, loc, int32Type, minus1, mbcntLo, /*arg_attrs=*/{},
102+
rewriter.getArrayAttr(rewriter.getDictionaryAttr({noundef, highRange})));
91103
return laneId;
92104
}
105+
93106
static constexpr StringLiteral amdgcnDataLayout =
94107
"e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32"
95108
"-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:"
@@ -104,18 +117,16 @@ struct GPULaneIdOpToROCDL : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
104117
LogicalResult
105118
matchAndRewrite(gpu::LaneIdOp op, gpu::LaneIdOp::Adaptor adaptor,
106119
ConversionPatternRewriter &rewriter) const override {
107-
auto loc = op->getLoc();
120+
Location loc = op.getLoc();
108121
MLIRContext *context = rewriter.getContext();
109-
// convert to: %mlo = call @llvm.amdgcn.mbcnt.lo(-1, 0)
110-
// followed by: %lid = call @llvm.amdgcn.mbcnt.hi(-1, %mlo)
111-
112-
Type intTy = IntegerType::get(context, 32);
113-
Value zero = arith::ConstantIntOp::create(rewriter, loc, 0, 32);
114-
Value minus1 = arith::ConstantIntOp::create(rewriter, loc, -1, 32);
115-
Value mbcntLo = ROCDL::MbcntLoOp::create(rewriter, loc, intTy,
116-
ValueRange{minus1, zero});
117-
Value laneId = ROCDL::MbcntHiOp::create(rewriter, loc, intTy,
118-
ValueRange{minus1, mbcntLo});
122+
// convert to:
123+
// %mlo = call noundef range(i32 0, 32)
124+
// @llvm.amdgcn.mbcnt.lo(-1, 0)
125+
// followed by:
126+
// %lid = call noundef range(i32 0, 64)
127+
// @llvm.amdgcn.mbcnt.hi(-1, %mlo)
128+
129+
Value laneId = getLaneId(rewriter, loc);
119130
// Truncate or extend the result depending on the index bitwidth specified
120131
// by the LLVMTypeConverter options.
121132
const unsigned indexBitwidth = getTypeConverter()->getIndexTypeBitwidth();
@@ -185,8 +196,7 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
185196
Location loc = op->getLoc();
186197
Value initShflValue = adaptor.getValue();
187198

188-
const unsigned indexBitwidth = getTypeConverter()->getIndexTypeBitwidth();
189-
Value srcLaneId = getLaneId(rewriter, loc, indexBitwidth);
199+
Value srcLaneId = getLaneId(rewriter, loc);
190200

191201
auto int32Type = IntegerType::get(rewriter.getContext(), 32);
192202
Value width = adaptor.getWidth();

‎mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ gpu.module @test_module {
5454
// CHECK: = llvm.sext %{{.*}} : i32 to i64
5555
%gDimZ = gpu.grid_dim z
5656

57-
// CHECK: = rocdl.mbcnt.lo %{{.*}}, %{{.*}} : (i32, i32) -> i32
58-
// CHECK: = rocdl.mbcnt.hi %{{.*}}, %{{.*}} : (i32, i32) -> i32
57+
// CHECK: = rocdl.mbcnt.lo %{{.*}}, %{{.*}} {res_attrs = [{llvm.noundef, llvm.range = #llvm.constant_range<i32, 0, 32>}]} : (i32, i32) -> i32
58+
// CHECK: = rocdl.mbcnt.hi %{{.*}}, %{{.*}} {res_attrs = [{llvm.noundef, llvm.range = #llvm.constant_range<i32, 0, 64>}]} : (i32, i32) -> i32
5959
// CHECK: = llvm.sext %{{.*}} : i32 to i64
6060
%laneId = gpu.lane_id
6161

@@ -701,7 +701,7 @@ gpu.module @test_module {
701701
// CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32
702702
// CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32
703703
// CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32
704-
%shfli, %predi = gpu.shuffle idx %arg0, %arg1, %arg2 : f32
704+
%shfli, %predi = gpu.shuffle idx %arg0, %arg1, %arg2 : f32
705705
// *** UP mode shuffle ***
706706
// CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi
707707
// CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32

‎mlir/test/Target/LLVMIR/rocdl.mlir‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -86,12 +86,12 @@ llvm.func @kernel_func_unsafe_fp_atomics()
8686
}
8787

8888
llvm.func @rocdl.lane_id() -> i32 {
89-
// CHECK: [[mbcntlo:%.+]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
90-
// CHECK-NEXT: call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[mbcntlo]])
89+
// CHECK: [[mbcntlo:%.+]] = call noundef range(i32 0, 32) i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
90+
// CHECK-NEXT: call noundef range(i32 0, 64) i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[mbcntlo]])
9191
%0 = llvm.mlir.constant(-1 : i32) : i32
9292
%1 = llvm.mlir.constant(0 : i32) : i32
93-
%2 = rocdl.mbcnt.lo %0, %1 : (i32, i32) -> i32
94-
%3 = rocdl.mbcnt.hi %0, %2 : (i32, i32) -> i32
93+
%2 = rocdl.mbcnt.lo %0, %1 {res_attrs = [{llvm.noundef, llvm.range = #llvm.constant_range<i32, 0, 32>}]} : (i32, i32) -> i32
94+
%3 = rocdl.mbcnt.hi %0, %2 {res_attrs = [{llvm.noundef, llvm.range = #llvm.constant_range<i32, 0, 64>}]} : (i32, i32) -> i32
9595
llvm.return %3 : i32
9696
}
9797

0 commit comments

Comments
 (0)