Skip to content

Commit 3fe6daa

Browse files
committed
move functions in NVVMToLLVMIRTranslation.cpp to NVVMDialect.cpp
1 parent 81d3f9a commit 3fe6daa

File tree

3 files changed

+690
-436
lines changed

3 files changed

+690
-436
lines changed

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

Lines changed: 99 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -410,9 +410,16 @@ def NVVM_ReduxOp :
410410

411411
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-redux-sync)
412412
}];
413+
let extraClassDeclaration = [{
414+
static NVVM::IIDArgsMaybeWithTypes
415+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
416+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
417+
}];
413418
string llvmBuilder = [{
414-
auto intId = getReduxIntrinsicId($_resultType, $kind, $abs, $nan);
415-
$res = createIntrinsicCall(builder, intId, {$val, $mask_and_clamp});
419+
auto [id, args, types] =
420+
NVVM::ReduxOp::getIntrinsicIDAndArgsMaybeWithTypes(
421+
*op, moduleTranslation, builder);
422+
$res = createIntrinsicCall(builder, id, args);
416423
}];
417424
let assemblyFormat = [{
418425
$kind $val `,` $mask_and_clamp attr-dict `:` type($val) `->` type($res)
@@ -876,11 +883,17 @@ def NVVM_FenceProxyAcquireOp : NVVM_Op<"fence.proxy.acquire">,
876883
}];
877884

878885
let assemblyFormat = "$scope $addr `,` $size (`from_proxy` `=` $fromProxy^)? (`to_proxy` `=` $toProxy^)? attr-dict";
886+
887+
let extraClassDeclaration = [{
888+
static NVVM::IIDArgsMaybeWithTypes
889+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
890+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
891+
}];
879892
let llvmBuilder = [{
880-
createIntrinsicCall(
881-
builder,
882-
getUnidirectionalFenceProxyID($fromProxy, $toProxy, $scope, false),
883-
{$addr, $size});
893+
auto [intId, args, types] =
894+
NVVM::FenceProxyAcquireOp::getIntrinsicIDAndArgsMaybeWithTypes(
895+
*op, moduleTranslation, builder);
896+
createIntrinsicCall(builder, intId, args);
884897
}];
885898

886899
let hasVerifier = 1;
@@ -904,9 +917,16 @@ def NVVM_FenceProxyReleaseOp : NVVM_Op<"fence.proxy.release">,
904917
}];
905918

906919
let assemblyFormat = "$scope (`from_proxy` `=` $fromProxy^)? (`to_proxy` `=` $toProxy^)? attr-dict";
920+
921+
let extraClassDeclaration = [{
922+
static NVVM::IIDArgsMaybeWithTypes
923+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
924+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
925+
}];
907926
let llvmBuilder = [{
908-
createIntrinsicCall(builder, getUnidirectionalFenceProxyID(
909-
$fromProxy, $toProxy, $scope, true));
927+
auto [intId, args, types] = NVVM::FenceProxyReleaseOp::getIntrinsicIDAndArgsMaybeWithTypes(
928+
*op, moduleTranslation, builder);
929+
createIntrinsicCall(builder, intId, args);
910930
}];
911931

912932
let hasVerifier = 1;
@@ -985,11 +1005,15 @@ def NVVM_ShflOp :
9851005

9861006
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync)
9871007
}];
1008+
let extraClassDeclaration = [{
1009+
static NVVM::IIDArgsMaybeWithTypes
1010+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
1011+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
1012+
}];
9881013
string llvmBuilder = [{
989-
auto intId = getShflIntrinsicId(
990-
$_resultType, $kind, static_cast<bool>($return_value_and_is_valid));
991-
$res = createIntrinsicCall(builder,
992-
intId, {$thread_mask, $val, $offset, $mask_and_clamp});
1014+
auto [intId, args, types] = NVVM::ShflOp::getIntrinsicIDAndArgsMaybeWithTypes(
1015+
*op, moduleTranslation, builder);
1016+
$res = createIntrinsicCall(builder, intId, args);
9931017
}];
9941018
let assemblyFormat = [{
9951019
$kind $thread_mask `,` $val `,` $offset `,` $mask_and_clamp attr-dict
@@ -1035,9 +1059,16 @@ def NVVM_VoteSyncOp
10351059

10361060
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync)
10371061
}];
1062+
let extraClassDeclaration = [{
1063+
static NVVM::IIDArgsMaybeWithTypes
1064+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
1065+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
1066+
}];
10381067
string llvmBuilder = [{
1039-
auto intId = getVoteSyncIntrinsicId($kind);
1040-
$res = createIntrinsicCall(builder, intId, {$mask, $pred});
1068+
auto [intId, args, types] =
1069+
NVVM::VoteSyncOp::getIntrinsicIDAndArgsMaybeWithTypes(
1070+
*op, moduleTranslation, builder);
1071+
$res = createIntrinsicCall(builder, intId, args);
10411072
}];
10421073
let assemblyFormat = "$kind $mask `,` $pred attr-dict `->` type($res)";
10431074
let hasVerifier = 1;
@@ -2106,10 +2137,16 @@ def NVVM_StMatrixOp: NVVM_Op<"stmatrix">,
21062137

21072138
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-store-instruction-stmatrix)
21082139
}];
2140+
let extraClassDeclaration = [{
2141+
static NVVM::IIDArgsMaybeWithTypes
2142+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
2143+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
2144+
}];
21092145
string llvmBuilder = [{
2110-
auto operands = moduleTranslation.lookupValues(opInst.getOperands());
2111-
auto intId = getStMatrixIntrinsicId($layout, $sources.size(), $shape, $eltType);
2112-
createIntrinsicCall(builder, intId, operands, operands[0]->getType());
2146+
auto [intId, args, types] =
2147+
NVVM::StMatrixOp::getIntrinsicIDAndArgsMaybeWithTypes(
2148+
*op, moduleTranslation, builder);
2149+
createIntrinsicCall(builder, intId, args, types);
21132150
}];
21142151
let assemblyFormat = "$ptr `,` $sources attr-dict `:` type(operands)";
21152152
let hasVerifier = 1;
@@ -2124,10 +2161,16 @@ def NVVM_LdMatrixOp: NVVM_Op<"ldmatrix">,
21242161

21252162
let summary = "cooperative matrix load";
21262163

2164+
let extraClassDeclaration = [{
2165+
static NVVM::IIDArgsMaybeWithTypes
2166+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
2167+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
2168+
}];
21272169
string llvmBuilder = [{
2128-
auto operands = moduleTranslation.lookupValues(opInst.getOperands());
2129-
auto intId = getLdMatrixIntrinsicId($layout, $num, $shape, $eltType);
2130-
$res = createIntrinsicCall(builder, intId, operands, {operands[0]->getType()});
2170+
auto [intId, args, types] =
2171+
NVVM::LdMatrixOp::getIntrinsicIDAndArgsMaybeWithTypes(
2172+
*op, moduleTranslation, builder);
2173+
$res = createIntrinsicCall(builder, intId, args, types);
21312174
}];
21322175

21332176
string baseDescription = [{
@@ -3238,11 +3281,16 @@ def NVVM_MatchSyncOp : NVVM_Op<"match.sync">,
32383281

32393282
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-match-sync)
32403283
}];
3284+
let extraClassDeclaration = [{
3285+
static NVVM::IIDArgsMaybeWithTypes
3286+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
3287+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
3288+
}];
32413289
string llvmBuilder = [{
3242-
auto intId = getMatchSyncIntrinsicId(
3243-
op.getVal().getType(), $kind);
3244-
$res = createIntrinsicCall(builder,
3245-
intId, {$thread_mask, $val});
3290+
auto [intId, args, types] =
3291+
NVVM::MatchSyncOp::getIntrinsicIDAndArgsMaybeWithTypes(
3292+
*op, moduleTranslation, builder);
3293+
$res = createIntrinsicCall(builder, intId, args);
32463294
}];
32473295
let assemblyFormat = "$kind $thread_mask `,` $val attr-dict `:` type($val) `->` type($res)";
32483296
let hasVerifier = 1;
@@ -3266,11 +3314,16 @@ def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
32663314
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-st-bulk)
32673315
}];
32683316

3317+
let extraClassDeclaration = [{
3318+
static NVVM::IIDArgsMaybeWithTypes
3319+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
3320+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
3321+
}];
32693322
string llvmBuilder = [{
3270-
auto intId = getStBulkIntrinsicId(
3271-
llvm::cast<LLVM::LLVMPointerType>(op.getAddr().getType()));
3272-
createIntrinsicCall(builder, intId,
3273-
{$addr, $size, builder.getInt64($initVal)});
3323+
auto [intId, args, types] =
3324+
NVVM::BulkStoreOp::getIntrinsicIDAndArgsMaybeWithTypes(
3325+
*op, moduleTranslation, builder);
3326+
createIntrinsicCall(builder, intId, args);
32743327
}];
32753328

32763329
let assemblyFormat = "$addr `,` `size` `=` $size (`,` `init` `=` $initVal^)? attr-dict `:` type($addr)";
@@ -3767,24 +3820,16 @@ def NVVM_Tcgen05LdOp : NVVM_Op<"tcgen05.ld", [NVVMRequiresSMa<[100, 101]>]> {
37673820

37683821
let hasVerifier = 1;
37693822

3823+
let extraClassDeclaration = [{
3824+
static NVVM::IIDArgsMaybeWithTypes
3825+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
3826+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
3827+
}];
37703828
string llvmBuilder = [{
3771-
llvm::LLVMContext &Context = moduleTranslation.getLLVMContext();
3772-
auto Pack = llvm::ConstantInt::get(Context, llvm::APInt(1, $pack));
3773-
3774-
unsigned num = $_resultType->isVectorTy()
3775-
? llvm::cast<llvm::VectorType>($_resultType)
3776-
->getElementCount()
3777-
.getFixedValue()
3778-
: 1;
3779-
3780-
auto ID = getTcgen05LdIntrinsicID($shape, num);
3781-
if (ID == llvm::Intrinsic::not_intrinsic)
3782-
llvm::report_fatal_error("unknow intrinsic signature for tcgen05.ld");
3783-
3784-
if ($offset)
3785-
$res = createIntrinsicCall(builder, ID, {$tmemAddr, $offset, Pack});
3786-
else
3787-
$res = createIntrinsicCall(builder, ID, {$tmemAddr, Pack});
3829+
auto [id, args, types] =
3830+
NVVM::Tcgen05LdOp::getIntrinsicIDAndArgsMaybeWithTypes(*op,
3831+
moduleTranslation, builder);
3832+
$res = createIntrinsicCall(builder, id, args);
37883833
}];
37893834
}
37903835

@@ -3855,24 +3900,16 @@ def NVVM_Tcgen05StOp : NVVM_Op<"tcgen05.st", [NVVMRequiresSMa<[100, 101]>]> {
38553900
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st)
38563901
}];
38573902

3903+
let extraClassDeclaration = [{
3904+
static NVVM::IIDArgsMaybeWithTypes
3905+
getIntrinsicIDAndArgsMaybeWithTypes(Operation &op,
3906+
LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
3907+
}];
38583908
string llvmBuilder = [{
3859-
llvm::LLVMContext &Context = moduleTranslation.getLLVMContext();
3860-
auto Unpack = llvm::ConstantInt::get(Context, llvm::APInt(1, $unpack));
3861-
3862-
auto valTy = $val->getType();
3863-
uint32_t num = valTy->isVectorTy() ? llvm::cast<llvm::VectorType>(valTy)
3864-
->getElementCount()
3865-
.getFixedValue()
3866-
: 1;
3867-
3868-
auto ID = getTcgen05StIntrinsicID($shape, num);
3869-
if (ID == llvm::Intrinsic::not_intrinsic)
3870-
llvm::report_fatal_error("unknow intrinsic signature for tcgen05.st");
3871-
3872-
if ($offset)
3873-
createIntrinsicCall(builder, ID, {$tmemAddr, $offset, $val, Unpack});
3874-
else
3875-
createIntrinsicCall(builder, ID, {$tmemAddr, $val, Unpack});
3909+
auto [id, args, types] =
3910+
NVVM::Tcgen05StOp::getIntrinsicIDAndArgsMaybeWithTypes(*op,
3911+
moduleTranslation, builder);
3912+
createIntrinsicCall(builder, id, args);
38763913
}];
38773914

38783915
let hasVerifier = 1;

0 commit comments

Comments
 (0)