Skip to content

[NVPTX][NFC] Move more TMA lowering to tablegen #140914

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 29, 2025
Merged

[NVPTX][NFC] Move more TMA lowering to tablegen #140914

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 0 additions & 69 deletions llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2686,69 +2686,6 @@ void NVPTXDAGToDAGISel::SelectCpAsyncBulkTensorReduceCommon(SDNode *N,
ReplaceNode(N, CurDAG->getMachineNode(Opcode, DL, N->getVTList(), Ops));
}

void NVPTXDAGToDAGISel::SelectCpAsyncBulkG2S(SDNode *N) {
// We have {Chain, Intrinsic-ID} followed by the actual intrisic args:
// {dst, mbar, src, size, multicast, cache_hint,
// multicast_flag, cache_hint_flag}
// NumOperands = {Chain, IID} + {Actual intrinsic args}
// = {2} + {8}
size_t NumOps = N->getNumOperands();
bool IsCacheHint = N->getConstantOperandVal(NumOps - 1) == 1;
bool IsMultiCast = N->getConstantOperandVal(NumOps - 2) == 1;
size_t NumBaseArgs = 4; // dst, mbar, src, size
size_t MultiCastIdx = NumBaseArgs + 2; // for Chain and IID

SDLoc DL(N);
SmallVector<SDValue, 8> Ops(N->ops().slice(2, NumBaseArgs));

// Push MultiCast operand, if available
if (IsMultiCast)
Ops.push_back(N->getOperand(MultiCastIdx));

// Push CacheHint operand, if available
if (IsCacheHint)
Ops.push_back(N->getOperand(MultiCastIdx + 1));

// Finally, the chain operand
Ops.push_back(N->getOperand(0));

bool IsShared32 =
CurDAG->getDataLayout().getPointerSizeInBits(ADDRESS_SPACE_SHARED) == 32;
unsigned Opcode = [&]() {
if (IsMultiCast && IsCacheHint)
return IsShared32 ? NVPTX::CP_ASYNC_BULK_G2S_SHARED32_MC_CH
: NVPTX::CP_ASYNC_BULK_G2S_MC_CH;
if (IsMultiCast)
return IsShared32 ? NVPTX::CP_ASYNC_BULK_G2S_SHARED32_MC
: NVPTX::CP_ASYNC_BULK_G2S_MC;
if (IsCacheHint)
return IsShared32 ? NVPTX::CP_ASYNC_BULK_G2S_SHARED32_CH
: NVPTX::CP_ASYNC_BULK_G2S_CH;
return IsShared32 ? NVPTX::CP_ASYNC_BULK_G2S_SHARED32
: NVPTX::CP_ASYNC_BULK_G2S;
}();
ReplaceNode(N, CurDAG->getMachineNode(Opcode, DL, N->getVTList(), Ops));
}

void NVPTXDAGToDAGISel::SelectCpAsyncBulkPrefetchL2(SDNode *N) {
// We have {Chain, Intrinsic-ID} followed by the actual intrisic args:
// src, size, cache_hint, cache_hint_flag
// NumOperands = {Chain, IID} + {Actual intrinsic args}
// = {2} + {4}
size_t NumOps = N->getNumOperands();
bool IsCacheHint = N->getConstantOperandVal(NumOps - 1) == 1;
size_t NumArgs = IsCacheHint ? 3 : 2; // src, size, cache_hint

SDLoc DL(N);
SmallVector<SDValue, 4> Ops(N->ops().slice(2, NumArgs));
Ops.push_back(N->getOperand(0)); // Chain operand

unsigned Opcode = IsCacheHint
? NVPTX::CP_ASYNC_BULK_PREFETCH_CH
: NVPTX::CP_ASYNC_BULK_PREFETCH;
ReplaceNode(N, CurDAG->getMachineNode(Opcode, DL, N->getVTList(), Ops));
}

#define TCGEN05_ST_OPCODE(SHAPE, NUM) \
(enableUnpack ? NVPTX::TCGEN05_ST_##SHAPE##_##NUM##_UNPACK \
: NVPTX::TCGEN05_ST_##SHAPE##_##NUM)
Expand Down Expand Up @@ -2865,12 +2802,6 @@ bool NVPTXDAGToDAGISel::tryIntrinsicVoid(SDNode *N) {
switch (IID) {
default:
return false;
case Intrinsic::nvvm_cp_async_bulk_global_to_shared_cluster:
SelectCpAsyncBulkG2S(N);
return true;
case Intrinsic::nvvm_cp_async_bulk_prefetch_L2:
SelectCpAsyncBulkPrefetchL2(N);
return true;
case Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_1d:
case Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_2d:
case Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_3d:
Expand Down
2 changes: 0 additions & 2 deletions llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,6 @@ class LLVM_LIBRARY_VISIBILITY NVPTXDAGToDAGISel : public SelectionDAGISel {
bool tryEXTRACT_VECTOR_ELEMENT(SDNode *N);
void SelectV2I64toI128(SDNode *N);
void SelectI128toV2I64(SDNode *N);
void SelectCpAsyncBulkG2S(SDNode *N);
void SelectCpAsyncBulkPrefetchL2(SDNode *N);
void SelectCpAsyncBulkTensorG2SCommon(SDNode *N, bool IsIm2Col = false);
void SelectCpAsyncBulkTensorS2GCommon(SDNode *N, bool IsIm2Col = false);
void SelectCpAsyncBulkTensorPrefetchCommon(SDNode *N, bool IsIm2Col = false);
Expand Down
80 changes: 39 additions & 41 deletions llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -544,52 +544,50 @@ multiclass CP_ASYNC_BULK_S2G_INTR<bit has_ch> {
[(int_nvvm_cp_async_bulk_shared_cta_to_global_bytemask addr:$dst, addr:$src, i32:$size, i64:$ch, !if(has_ch, -1, 0), i16:$mask)]>,
Requires<[hasPTX<86>, hasSM<100>]>;
}
defm CP_ASYNC_BULK_S2G : CP_ASYNC_BULK_S2G_INTR<0>;
defm CP_ASYNC_BULK_S2G_CH : CP_ASYNC_BULK_S2G_INTR<1>;
defm CP_ASYNC_BULK_S2G : CP_ASYNC_BULK_S2G_INTR<has_ch = 0>;
defm CP_ASYNC_BULK_S2G_CH : CP_ASYNC_BULK_S2G_INTR<has_ch = 1>;

multiclass CP_ASYNC_BULK_G2S<NVPTXRegClass rc> {
def NAME: NVPTXInst<(outs),
(ins rc:$dst, rc:$mbar, Int64Regs:$src, Int32Regs:$size),
!strconcat(CpAsyncBulkStr<0, 0>.G2S, " [$dst], [$src], $size, [$mbar];"), []>,
Requires<[hasPTX<80>, hasSM<90>]>;
def NAME # _MC: NVPTXInst<(outs),
(ins rc:$dst, rc:$mbar, Int64Regs:$src, Int32Regs:$size, Int16Regs:$mc),
!strconcat(CpAsyncBulkStr<1, 0>.G2S, " [$dst], [$src], $size, [$mbar], $mc;"), []>,
Requires<[hasPTX<80>, hasSM<90>]>;
def NAME # _CH: NVPTXInst<(outs),
(ins rc:$dst, rc:$mbar, Int64Regs:$src, Int32Regs:$size, Int64Regs:$ch),
!strconcat(CpAsyncBulkStr<0, 1>.G2S, " [$dst], [$src], $size, [$mbar], $ch;"), []>,
Requires<[hasPTX<80>, hasSM<90>]>;
def NAME # _MC_CH: NVPTXInst<(outs),
(ins rc:$dst, rc:$mbar, Int64Regs:$src, Int32Regs:$size, Int16Regs:$mc, Int64Regs:$ch),
!strconcat(CpAsyncBulkStr<1, 1>.G2S, " [$dst], [$src], $size, [$mbar], $mc, $ch;"), []>,
Requires<[hasPTX<80>, hasSM<90>]>;
multiclass CP_ASYNC_BULK_G2S_INTR<bit has_ch> {
defvar Intr = int_nvvm_cp_async_bulk_global_to_shared_cluster;

def NAME : NVPTXInst<(outs),
(ins ADDR:$dst, ADDR:$mbar, ADDR:$src,
Int32Regs:$size, Int16Regs:$mask, Int64Regs:$ch),
!if(has_ch,
CpAsyncBulkStr<0, 1>.G2S # " [$dst], [$src], $size, [$mbar], $ch;",
CpAsyncBulkStr<0, 0>.G2S # " [$dst], [$src], $size, [$mbar];"),
[(Intr addr:$dst, addr:$mbar, addr:$src, i32:$size, i16:$mask, i64:$ch, 0, !if(has_ch, -1, 0))]>,
Requires<[hasPTX<80>, hasSM<90>]>;

def NAME # _MC : NVPTXInst<(outs),
(ins ADDR:$dst, ADDR:$mbar, ADDR:$src,
Int32Regs:$size, Int16Regs:$mask, Int64Regs:$ch),
!if(has_ch,
CpAsyncBulkStr<1, 1>.G2S # " [$dst], [$src], $size, [$mbar], $mask, $ch;",
CpAsyncBulkStr<1, 0>.G2S # " [$dst], [$src], $size, [$mbar], $mask;"),
[(Intr addr:$dst, addr:$mbar, addr:$src, i32:$size, i16:$mask, i64:$ch, -1, !if(has_ch, -1, 0))]>,
Requires<[hasPTX<80>, hasSM<90>]>;
}
defm CP_ASYNC_BULK_G2S : CP_ASYNC_BULK_G2S<Int64Regs>;
defm CP_ASYNC_BULK_G2S_SHARED32 : CP_ASYNC_BULK_G2S<Int32Regs>;
defm CP_ASYNC_BULK_G2S : CP_ASYNC_BULK_G2S_INTR<has_ch = 0>;
defm CP_ASYNC_BULK_G2S_CH : CP_ASYNC_BULK_G2S_INTR<has_ch = 1>;

multiclass CP_ASYNC_BULK_CTA_TO_CLUSTER<NVPTXRegClass rc> {
def NAME: NVPTXInst<(outs),
(ins rc:$dst, rc:$mbar, rc:$src, Int32Regs:$size),
!strconcat(CpAsyncBulkStr<0, 0>.C2C, " [$dst], [$src], $size, [$mbar];"),
[(int_nvvm_cp_async_bulk_shared_cta_to_cluster rc:$dst, rc:$mbar, rc:$src, Int32Regs:$size)]>,
Requires<[hasPTX<80>, hasSM<90>]>;
def CP_ASYNC_BULK_CTA_TO_CLUSTER : NVPTXInst<(outs),
(ins ADDR:$dst, ADDR:$mbar, ADDR:$src, Int32Regs:$size),
CpAsyncBulkStr<0, 0>.C2C # " [$dst], [$src], $size, [$mbar];",
[(int_nvvm_cp_async_bulk_shared_cta_to_cluster addr:$dst, addr:$mbar, addr:$src, i32:$size)]>,
Requires<[hasPTX<80>, hasSM<90>]>;

multiclass CP_ASYNC_BULK_PREFETCH_INTR<bit has_ch> {
def NAME : NVPTXInst<(outs), (ins ADDR:$src, Int32Regs:$size, Int64Regs:$ch),
!if(has_ch,
"cp.async.bulk.prefetch.L2.global.L2::cache_hint" # " [$src], $size, $ch;",
"cp.async.bulk.prefetch.L2.global" # " [$src], $size;"),
[(int_nvvm_cp_async_bulk_prefetch_L2 addr:$src, i32:$size, i64:$ch, !if(has_ch, -1, 0))]>,
Requires<[hasPTX<80>, hasSM<90>]>;
}
defm CP_ASYNC_BULK_CTA_TO_CLUSTER : CP_ASYNC_BULK_CTA_TO_CLUSTER<Int64Regs>;
defm CP_ASYNC_BULK_CTA_TO_CLUSTER_SHARED32 : CP_ASYNC_BULK_CTA_TO_CLUSTER<Int32Regs>;
defm CP_ASYNC_BULK_PREFETCH : CP_ASYNC_BULK_PREFETCH_INTR<has_ch = 0>;
defm CP_ASYNC_BULK_PREFETCH_CH : CP_ASYNC_BULK_PREFETCH_INTR<has_ch = 1>;

//------------------------------
// Bulk Copy Prefetch Functions
//------------------------------
def CP_ASYNC_BULK_PREFETCH : NVPTXInst<(outs),
(ins Int64Regs:$src, Int32Regs:$size),
"cp.async.bulk.prefetch.L2.global [$src], $size;", []>,
Requires<[hasPTX<80>, hasSM<90>]>;

def CP_ASYNC_BULK_PREFETCH_CH : NVPTXInst<(outs),
(ins Int64Regs:$src, Int32Regs:$size, Int64Regs:$ch),
"cp.async.bulk.prefetch.L2.global.L2::cache_hint [$src], $size, $ch;", []>,
Requires<[hasPTX<80>, hasSM<90>]>;
//-------------------------------------
// TMA Async Bulk Tensor Copy Functions
//-------------------------------------
Expand Down
Loading