Skip to content

Commit a25daa3

Browse files
durga4githubdurga4github
authored
[NVPTX] Add TMA non-tensor variant of g2s-cta intrinsic (llvm#167508)
This patch adds a TMA intrinsic for Global to shared::cta copy, which was introduced with ptx86. Also remove the NoCapture<> annotation from the pointer arguments to these intrinsics, since the copy operations are asynchronous in nature. lit tests are verified with a ptxas from cuda-12.8. Signed-off-by: Durgadoss R <[email protected]>
1 parent 7e5155a commit a25daa3

File tree

4 files changed

+106
-7
lines changed

4 files changed

+106
-7
lines changed

‎llvm/docs/NVPTXUsage.rst‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1334,6 +1334,32 @@ copied and it must be a multiple of 16.
13341334
For more information, refer PTX ISA
13351335
`<https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk>`_.
13361336

1337+
'``llvm.nvvm.cp.async.bulk.global.to.shared.cta``'
1338+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
1339+
1340+
Syntax:
1341+
"""""""
1342+
1343+
.. code-block:: llvm
1344+
1345+
declare void @llvm.nvvm.cp.async.bulk.global.to.shared.cta(ptr addrspace(3) %dst, ptr addrspace(3) %mbar, ptr addrspace(1) %src, i32 %size, i64 %ch, i1 %flag_ch)
1346+
1347+
Overview:
1348+
"""""""""
1349+
1350+
The '``@llvm.nvvm.cp.async.bulk.global.to.shared.cta``' intrinsic
1351+
corresponds to the ``cp.async.bulk.shared::cta.global.*`` family
1352+
of PTX instructions. These instructions initiate an asynchronous
1353+
copy of bulk data from global memory to shared::cta memory.
1354+
The 32-bit operand ``%size`` specifies the amount of memory to be
1355+
copied and it must be a multiple of 16. The last argument
1356+
(denoted by ``i1 %flag_ch``) is a compile-time constant. When set,
1357+
it indicates a valid cache_hint (``i64 %ch``) and generates the
1358+
``.L2::cache_hint`` variant of the PTX instruction.
1359+
1360+
For more information, refer PTX ISA
1361+
`<https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk>`_.
1362+
13371363
'``llvm.nvvm.cp.async.bulk.shared.cta.to.global``'
13381364
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
13391365

‎llvm/include/llvm/IR/IntrinsicsNVVM.td‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2716,8 +2716,19 @@ def int_nvvm_cp_async_bulk_global_to_shared_cluster
27162716
[llvm_i1_ty, // Flag for cta_mask
27172717
llvm_i1_ty], // Flag for cache_hint
27182718
[IntrConvergent, IntrArgMemOnly,
2719-
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>,
2720-
NoCapture<ArgIndex<0>>, NoCapture<ArgIndex<1>>, NoCapture<ArgIndex<2>>]>;
2719+
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>]>;
2720+
2721+
// From Global to Shared CTA
2722+
def int_nvvm_cp_async_bulk_global_to_shared_cta
2723+
: DefaultAttrsIntrinsicFlags<[],
2724+
[llvm_shared_ptr_ty, // dst_shared_cta_ptr
2725+
llvm_shared_ptr_ty, // mbarrier_ptr
2726+
llvm_global_ptr_ty, // src_gmem_ptr
2727+
llvm_i32_ty, // copy_size
2728+
llvm_i64_ty], // cache_hint
2729+
[llvm_i1_ty], // Flag for cache_hint
2730+
[IntrConvergent, IntrArgMemOnly,
2731+
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>]>;
27212732

27222733
// From Shared CTA to Shared Cluster
27232734
def int_nvvm_cp_async_bulk_shared_cta_to_cluster
@@ -2727,9 +2738,7 @@ def int_nvvm_cp_async_bulk_shared_cta_to_cluster
27272738
llvm_shared_ptr_ty, // src_smem_ptr
27282739
llvm_i32_ty], // copy_size
27292740
[IntrConvergent, IntrArgMemOnly,
2730-
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>,
2731-
NoCapture<ArgIndex<0>>, NoCapture<ArgIndex<1>>,
2732-
NoCapture<ArgIndex<2>>]>;
2741+
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>]>;
27332742

27342743
// From Shared CTA to Global memory
27352744
def int_nvvm_cp_async_bulk_shared_cta_to_global
@@ -2740,8 +2749,7 @@ def int_nvvm_cp_async_bulk_shared_cta_to_global
27402749
llvm_i64_ty], // cache_hint
27412750
[llvm_i1_ty], // Flag for cache_hint
27422751
[IntrConvergent, IntrArgMemOnly,
2743-
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<1>>,
2744-
NoCapture<ArgIndex<0>>, NoCapture<ArgIndex<1>>]>;
2752+
WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<1>>]>;
27452753

27462754
// From Shared CTA to Global memory with bytemask
27472755
def int_nvvm_cp_async_bulk_shared_cta_to_global_bytemask

‎llvm/lib/Target/NVPTX/NVPTXIntrinsics.td‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -497,6 +497,10 @@ class CpAsyncBulkStr<bit mc, bit ch, bit mask = 0> {
497497
# !if(mc, ".multicast::cluster", "")
498498
# !if(ch, ".L2::cache_hint", "");
499499

500+
// Global to Shared CTA memory
501+
string G2S_CTA = "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes"
502+
# !if(ch, ".L2::cache_hint", "");
503+
500504
// Shared CTA to Cluster memory
501505
string C2C = "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes";
502506
}
@@ -543,6 +547,21 @@ multiclass CP_ASYNC_BULK_G2S_INTR<bit has_ch> {
543547
defm CP_ASYNC_BULK_G2S : CP_ASYNC_BULK_G2S_INTR<has_ch = 0>;
544548
defm CP_ASYNC_BULK_G2S_CH : CP_ASYNC_BULK_G2S_INTR<has_ch = 1>;
545549

550+
multiclass CP_ASYNC_BULK_G2S_CTA_INTR<bit has_ch> {
551+
defvar Intr = int_nvvm_cp_async_bulk_global_to_shared_cta;
552+
553+
def "" : NVPTXInst<(outs),
554+
(ins ADDR:$dst, ADDR:$mbar, ADDR:$src,
555+
B32:$size, B64:$ch),
556+
!if(has_ch,
557+
CpAsyncBulkStr<0, 1>.G2S_CTA # " [$dst], [$src], $size, [$mbar], $ch;",
558+
CpAsyncBulkStr<0, 0>.G2S_CTA # " [$dst], [$src], $size, [$mbar];"),
559+
[(Intr addr:$dst, addr:$mbar, addr:$src, i32:$size, i64:$ch, !if(has_ch, -1, 0))]>,
560+
Requires<[hasPTX<86>, hasSM<90>]>;
561+
}
562+
defm CP_ASYNC_BULK_G2S_CTA : CP_ASYNC_BULK_G2S_CTA_INTR<has_ch = 0>;
563+
defm CP_ASYNC_BULK_G2S_CTA_CH : CP_ASYNC_BULK_G2S_CTA_INTR<has_ch = 1>;
564+
546565
def CP_ASYNC_BULK_CTA_TO_CLUSTER : NVPTXInst<(outs),
547566
(ins ADDR:$dst, ADDR:$mbar, ADDR:$src, B32:$size),
548567
CpAsyncBulkStr<0, 0>.C2C # " [$dst], [$src], $size, [$mbar];",

‎llvm/test/CodeGen/NVPTX/cp-async-bulk-ptx86.ll‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
2+
; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_90 -mattr=+ptx86| FileCheck --check-prefixes=CHECK,CHECK-PTX64 %s
3+
; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_90 -mattr=+ptx86 --nvptx-short-ptr| FileCheck --check-prefixes=CHECK,CHECK-PTX-SHARED32 %s
4+
; RUN: %if ptxas-sm_90 && ptxas-isa-8.6 %{ llc < %s -mtriple=nvptx64 -mcpu=sm_90 -mattr=+ptx86| %ptxas-verify -arch=sm_90 %}
5+
; RUN: %if ptxas-sm_90 && ptxas-isa-8.6 %{ llc < %s -mtriple=nvptx64 -mcpu=sm_90 -mattr=+ptx86 --nvptx-short-ptr| %ptxas-verify -arch=sm_90 %}
6+
7+
target triple = "nvptx64-nvidia-cuda"
8+
9+
declare void @llvm.nvvm.cp.async.bulk.global.to.shared.cta(ptr addrspace(3), ptr addrspace(3), ptr addrspace(1), i32, i64, i1)
10+
11+
define void @cp_async_bulk_g2s(ptr addrspace(1) %src, ptr addrspace(3) %bar, ptr addrspace(3) %dst, i32 %size, i64 %ch) {
12+
; CHECK-PTX64-LABEL: cp_async_bulk_g2s(
13+
; CHECK-PTX64: {
14+
; CHECK-PTX64-NEXT: .reg .b32 %r<2>;
15+
; CHECK-PTX64-NEXT: .reg .b64 %rd<5>;
16+
; CHECK-PTX64-EMPTY:
17+
; CHECK-PTX64-NEXT: // %bb.0:
18+
; CHECK-PTX64-NEXT: ld.param.b64 %rd1, [cp_async_bulk_g2s_param_0];
19+
; CHECK-PTX64-NEXT: ld.param.b64 %rd2, [cp_async_bulk_g2s_param_1];
20+
; CHECK-PTX64-NEXT: ld.param.b64 %rd3, [cp_async_bulk_g2s_param_2];
21+
; CHECK-PTX64-NEXT: ld.param.b32 %r1, [cp_async_bulk_g2s_param_3];
22+
; CHECK-PTX64-NEXT: ld.param.b64 %rd4, [cp_async_bulk_g2s_param_4];
23+
; CHECK-PTX64-NEXT: cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%rd3], [%rd1], %r1, [%rd2];
24+
; CHECK-PTX64-NEXT: cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [%rd3], [%rd1], %r1, [%rd2], %rd4;
25+
; CHECK-PTX64-NEXT: ret;
26+
;
27+
; CHECK-PTX-SHARED32-LABEL: cp_async_bulk_g2s(
28+
; CHECK-PTX-SHARED32: {
29+
; CHECK-PTX-SHARED32-NEXT: .reg .b32 %r<4>;
30+
; CHECK-PTX-SHARED32-NEXT: .reg .b64 %rd<3>;
31+
; CHECK-PTX-SHARED32-EMPTY:
32+
; CHECK-PTX-SHARED32-NEXT: // %bb.0:
33+
; CHECK-PTX-SHARED32-NEXT: ld.param.b64 %rd1, [cp_async_bulk_g2s_param_0];
34+
; CHECK-PTX-SHARED32-NEXT: ld.param.b32 %r1, [cp_async_bulk_g2s_param_1];
35+
; CHECK-PTX-SHARED32-NEXT: ld.param.b32 %r2, [cp_async_bulk_g2s_param_2];
36+
; CHECK-PTX-SHARED32-NEXT: ld.param.b32 %r3, [cp_async_bulk_g2s_param_3];
37+
; CHECK-PTX-SHARED32-NEXT: ld.param.b64 %rd2, [cp_async_bulk_g2s_param_4];
38+
; CHECK-PTX-SHARED32-NEXT: cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%r2], [%rd1], %r3, [%r1];
39+
; CHECK-PTX-SHARED32-NEXT: cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [%r2], [%rd1], %r3, [%r1], %rd2;
40+
; CHECK-PTX-SHARED32-NEXT: ret;
41+
tail call void @llvm.nvvm.cp.async.bulk.global.to.shared.cta(ptr addrspace(3) %dst, ptr addrspace(3) %bar, ptr addrspace(1) %src, i32 %size, i64 %ch, i1 0)
42+
tail call void @llvm.nvvm.cp.async.bulk.global.to.shared.cta(ptr addrspace(3) %dst, ptr addrspace(3) %bar, ptr addrspace(1) %src, i32 %size, i64 %ch, i1 1)
43+
ret void
44+
}
45+
;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
46+
; CHECK: {{.*}}

0 commit comments

Comments
 (0)