[llvm] [NVPTX] Add Bulk Copy Prefetch Intrinsics (PR #123226)
Abhilash Majumder via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 17 01:31:25 PST 2025
https://github.com/abhilash1910 updated https://github.com/llvm/llvm-project/pull/123226
>From a6cc8f3f13f5b3476e2e2822c0cfc0a8e7d631dc Mon Sep 17 00:00:00 2001
From: abmajumder <abmajumder at nvidia.com>
Date: Thu, 16 Jan 2025 23:01:27 +0530
Subject: [PATCH 1/2] add nvptx for cp.async.bulk.prefetch
---
llvm/docs/NVPTXUsage.rst | 32 +++++++++++++++++++++
llvm/include/llvm/IR/IntrinsicsNVVM.td | 11 +++++++
llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp | 27 +++++++++++++++++
llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h | 1 +
llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 19 ++++++++++++
llvm/test/CodeGen/NVPTX/cp-async-bulk.ll | 19 ++++++++++++
6 files changed, 109 insertions(+)
diff --git a/llvm/docs/NVPTXUsage.rst b/llvm/docs/NVPTXUsage.rst
index 25a230f65fd3dd..bb1f0ee9df8a0a 100644
--- a/llvm/docs/NVPTXUsage.rst
+++ b/llvm/docs/NVPTXUsage.rst
@@ -553,6 +553,38 @@ it must be a multiple of 16.
For more information, refer PTX ISA
`<https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk>`_.
+'``llvm.nvvm.cp.async.bulk.prefetch.L2``'
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+.. code-block:: llvm
+
+ declare void @llvm.nvvm.cp.async.bulk.prefetch.L2(ptr addrspace(1) %src, i32 %size, i64 %ch, i1 %flag_ch)
+
+Overview:
+"""""""""
+
+The '``@llvm.nvvm.cp.async.bulk.prefetch.L2``' intrinsic
+corresponds to the ``cp.async.bulk.prefetch.L2.*`` family
+of PTX instructions. These instructions initiate an asynchronous
+prefetch of bulk data from global memory to the L2 cache.
+The 32-bit operand ``%size`` specifies the amount of memory to be
+prefetched in terms of bytes and it must be a multiple of 16.
+
+* The last argument to these intrinsics is boolean flag indicating
+ support for cache_hint. These flag argument must be compile-time
+ constant. The backend looks through this flag and lowers the
+ intrinsic appropriately.
+
+* The Nth argument (denoted by ``i1 %flag_ch``) when set, indicates
+ a valid cache_hint (``i64 %ch``) and generates the ``.L2::cache_hint``
+ variant of the PTX instruction.
+
+For more information, refer PTX ISA
+`<https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-prefetch>`_.
+
'``llvm.nvvm.cp.async.bulk.tensor.g2s.tile.[1-5]d``'
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index 00a76018d8415d..c02d77057cf1f8 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -5033,4 +5033,15 @@ def int_nvvm_cp_async_bulk_shared_cta_to_global
NoCapture<ArgIndex<0>>, NoCapture<ArgIndex<1>>,
ImmArg<ArgIndex<4>>]>;
+// Intrinsics for Bulk Copy Prefetch L2
+def int_nvvm_cp_async_bulk_prefetch_L2
+ : DefaultAttrsIntrinsic<[],
+ [llvm_global_ptr_ty, // src_smem_ptr
+ llvm_i32_ty, // copy_size
+ llvm_i64_ty, // cache_hint
+ llvm_i1_ty], // Flag for cache_hint
+ [IntrConvergent, IntrArgMemOnly,
+ NoCapture<ArgIndex<0>>, ReadOnly<ArgIndex<0>>,
+ ImmArg<ArgIndex<3>>]>;
+
} // let TargetPrefix = "nvvm"
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
index 2e66b67dfdcc76..f841f21768cea6 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
@@ -3105,6 +3105,30 @@ void NVPTXDAGToDAGISel::SelectCpAsyncBulkG2S(SDNode *N) {
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, 8> Ops(N->ops().slice(2, NumArgs));
+ Ops.push_back(N->getOperand(0)); // Chain operand
+ //if (IsCacheHint) {
+ // Ops.push_back(N->getOperand(2));
+ //}
+
+ unsigned Opcode;
+ if (IsCacheHint)
+ Opcode = NVPTX::CP_ASYNC_BULK_PREFETCH_CH;
+ else
+ Opcode = NVPTX::CP_ASYNC_BULK_PREFETCH;
+ ReplaceNode(N, CurDAG->getMachineNode(Opcode, DL, N->getVTList(), Ops));
+}
+
bool NVPTXDAGToDAGISel::tryIntrinsicVoid(SDNode *N) {
unsigned IID = N->getConstantOperandVal(1);
using TMARedTy = llvm::nvvm::TMAReductionOp;
@@ -3118,6 +3142,9 @@ bool NVPTXDAGToDAGISel::tryIntrinsicVoid(SDNode *N) {
case Intrinsic::nvvm_cp_async_bulk_shared_cta_to_global:
SelectCpAsyncBulkS2G(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:
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h
index 8cadde8a822647..c673c83beba0f2 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h
+++ b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h
@@ -92,6 +92,7 @@ class LLVM_LIBRARY_VISIBILITY NVPTXDAGToDAGISel : public SelectionDAGISel {
void SelectI128toV2I64(SDNode *N);
void SelectCpAsyncBulkG2S(SDNode *N);
void SelectCpAsyncBulkS2G(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);
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 2d6ee2e28b4df7..1af3c88573272e 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -547,6 +547,25 @@ multiclass CP_ASYNC_BULK_CTA_TO_CLUSTER<NVPTXRegClass rc> {
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>;
+//------------------------------
+// Bulk Copy Prefetch Functions
+//------------------------------
+multiclass CP_ASYNC_BULK_PREFETCH_INTR {
+ defvar prefetch = "cp.async.bulk.prefetch.L2.global";
+ def "": NVPTXInst<(outs),
+ (ins Int64Regs:$src, Int32Regs:$size),
+ !strconcat(prefetch," [$src], $size;"),
+ []>,
+ Requires<[hasPTX<80>, hasSM<90>]>;
+ def _CH: NVPTXInst<(outs),
+ (ins Int64Regs:$src, Int32Regs:$size, Int64Regs:$ch),
+ !strconcat(prefetch,".L2::cache_hint [$src], $size, $ch;"),
+ []>,
+ Requires<[hasPTX<80>, hasSM<90>]>;
+}
+
+defm CP_ASYNC_BULK_PREFETCH : CP_ASYNC_BULK_PREFETCH_INTR;
+
//-------------------------------------
// TMA Async Bulk Tensor Copy Functions
//-------------------------------------
diff --git a/llvm/test/CodeGen/NVPTX/cp-async-bulk.ll b/llvm/test/CodeGen/NVPTX/cp-async-bulk.ll
index aefd18a0632a08..cbb53df4a49b09 100644
--- a/llvm/test/CodeGen/NVPTX/cp-async-bulk.ll
+++ b/llvm/test/CodeGen/NVPTX/cp-async-bulk.ll
@@ -9,6 +9,7 @@ target triple = "nvptx64-nvidia-cuda"
declare void @llvm.nvvm.cp.async.bulk.global.to.shared.cluster(ptr addrspace(3), ptr addrspace(3), ptr addrspace(1), i32, i16, i64, i1, i1)
declare void @llvm.nvvm.cp.async.bulk.shared.cta.to.global(ptr addrspace(1), ptr addrspace(3), i32, i64, i1)
declare void @llvm.nvvm.cp.async.bulk.shared.cta.to.cluster(ptr addrspace(3), ptr addrspace(3), ptr addrspace(3), i32)
+declare void @llvm.nvvm.cp.async.bulk.prefetch.L2(ptr addrspace(1), i32, i64, i1)
define void @cp_async_bulk_g2s(ptr addrspace(1) %src, ptr addrspace(3) %bar, ptr addrspace(3) %dst, i32 %size, i16 %mc, i64 %ch) {
; CHECK-PTX64-LABEL: cp_async_bulk_g2s(
@@ -116,3 +117,21 @@ define void @cp_async_bulk_cta_to_cluster(ptr addrspace(3) %src, ptr addrspace(3
tail call void @llvm.nvvm.cp.async.bulk.shared.cta.to.cluster(ptr addrspace(3) %dst, ptr addrspace(3) %bar, ptr addrspace(3) %src, i32 %size)
ret void
}
+
+define void @cp_async_bulk_prefetch(ptr addrspace(1) %src, i32 %size, i64 %ch) {
+; CHECK-PTX64-LABEL: cp_async_bulk_prefetch(
+; CHECK-PTX64: {
+; CHECK-PTX64-NEXT: .reg .b32 %r<2>;
+; CHECK-PTX64-NEXT: .reg .b64 %rd<3>;
+; CHECK-PTX64-EMPTY:
+; CHECK-PTX64-NEXT: // %bb.0:
+; CHECK-PTX64-NEXT: ld.param.u64 %rd1, [cp_async_bulk_prefetch_param_0];
+; CHECK-PTX64-NEXT: ld.param.u32 %r1, [cp_async_bulk_prefetch_param_1];
+; CHECK-PTX64-NEXT: ld.param.u64 %rd2, [cp_async_bulk_prefetch_param_2];
+; CHECK-PTX64-NEXT: cp.async.bulk.prefetch.L2.global.L2::cache_hint [%rd1], %r1, %rd2;
+; CHECK-PTX64-NEXT: cp.async.bulk.prefetch.L2.global [%rd1], %r1;
+; CHECK-PTX64-NEXT: ret;
+ tail call void @llvm.nvvm.cp.async.bulk.prefetch.L2(ptr addrspace(1) %src, i32 %size, i64 %ch, i1 1)
+ tail call void @llvm.nvvm.cp.async.bulk.prefetch.L2(ptr addrspace(1) %src, i32 %size, i64 0, i1 0)
+ ret void
+}
>From b530b81d13c60f6b6575ae77cb379082d9e7ae26 Mon Sep 17 00:00:00 2001
From: abmajumder <abmajumder at nvidia.com>
Date: Fri, 17 Jan 2025 15:00:48 +0530
Subject: [PATCH 2/2] refine comments
---
llvm/docs/NVPTXUsage.rst | 8 ++------
llvm/include/llvm/IR/IntrinsicsNVVM.td | 2 +-
llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp | 5 +----
llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 4 ++--
4 files changed, 6 insertions(+), 13 deletions(-)
diff --git a/llvm/docs/NVPTXUsage.rst b/llvm/docs/NVPTXUsage.rst
index bb1f0ee9df8a0a..3c032c65943eba 100644
--- a/llvm/docs/NVPTXUsage.rst
+++ b/llvm/docs/NVPTXUsage.rst
@@ -575,12 +575,8 @@ prefetched in terms of bytes and it must be a multiple of 16.
* The last argument to these intrinsics is boolean flag indicating
support for cache_hint. These flag argument must be compile-time
- constant. The backend looks through this flag and lowers the
- intrinsic appropriately.
-
-* The Nth argument (denoted by ``i1 %flag_ch``) when set, indicates
- a valid cache_hint (``i64 %ch``) and generates the ``.L2::cache_hint``
- variant of the PTX instruction.
+ constant. When set, indicates a valid cache_hint (``i64 %ch``)
+ and generates the ``.L2::cache_hint`` variant of the PTX instruction.
For more information, refer PTX ISA
`<https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-prefetch>`_.
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index c02d77057cf1f8..00c441920bfa1c 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -5036,7 +5036,7 @@ def int_nvvm_cp_async_bulk_shared_cta_to_global
// Intrinsics for Bulk Copy Prefetch L2
def int_nvvm_cp_async_bulk_prefetch_L2
: DefaultAttrsIntrinsic<[],
- [llvm_global_ptr_ty, // src_smem_ptr
+ [llvm_global_ptr_ty, // src_gmem_ptr
llvm_i32_ty, // copy_size
llvm_i64_ty, // cache_hint
llvm_i1_ty], // Flag for cache_hint
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
index f841f21768cea6..8ad904985695b9 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
@@ -3115,11 +3115,8 @@ void NVPTXDAGToDAGISel::SelectCpAsyncBulkPrefetchL2(SDNode *N) {
size_t NumArgs = IsCacheHint ? 3 : 2; // src, size, cache_hint
SDLoc DL(N);
- SmallVector<SDValue, 8> Ops(N->ops().slice(2, NumArgs));
+ SmallVector<SDValue, 4> Ops(N->ops().slice(2, NumArgs));
Ops.push_back(N->getOperand(0)); // Chain operand
- //if (IsCacheHint) {
- // Ops.push_back(N->getOperand(2));
- //}
unsigned Opcode;
if (IsCacheHint)
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 1af3c88573272e..72097f7dd765ba 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -552,12 +552,12 @@ defm CP_ASYNC_BULK_CTA_TO_CLUSTER_SHARED32 : CP_ASYNC_BULK_CTA_TO_CLUSTER<Int32R
//------------------------------
multiclass CP_ASYNC_BULK_PREFETCH_INTR {
defvar prefetch = "cp.async.bulk.prefetch.L2.global";
- def "": NVPTXInst<(outs),
+ def NAME: NVPTXInst<(outs),
(ins Int64Regs:$src, Int32Regs:$size),
!strconcat(prefetch," [$src], $size;"),
[]>,
Requires<[hasPTX<80>, hasSM<90>]>;
- def _CH: NVPTXInst<(outs),
+ def NAME # _CH: NVPTXInst<(outs),
(ins Int64Regs:$src, Int32Regs:$size, Int64Regs:$ch),
!strconcat(prefetch,".L2::cache_hint [$src], $size, $ch;"),
[]>,
More information about the llvm-commits
mailing list