Skip to content

Commit

Permalink
[LLVM][NVPTX] Add cp.async.bulk.commit/wait intrinsics (#78698)
Browse files Browse the repository at this point in the history
This patch adds NVVM intrinsics and NVPTX codegen for the bulk variants
of the async-copy commit/wait instructions.
lit tests are added to verify the generated PTX.

PTX Doc link:

https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group

Signed-off-by: Durgadoss R <durgadossr@nvidia.com>
  • Loading branch information
durga4github committed Jan 19, 2024
1 parent 42b1603 commit 43531e7
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 0 deletions.
10 changes: 10 additions & 0 deletions llvm/include/llvm/IR/IntrinsicsNVVM.td
Original file line number Diff line number Diff line change
Expand Up @@ -1454,6 +1454,16 @@ def int_nvvm_cp_async_wait_all :
ClangBuiltin<"__nvvm_cp_async_wait_all">,
Intrinsic<[],[],[]>;

// cp.async.bulk variants of the commit/wait group
def int_nvvm_cp_async_bulk_commit_group :
Intrinsic<[],[],[]>;

def int_nvvm_cp_async_bulk_wait_group :
Intrinsic<[],[llvm_i32_ty],[ImmArg<ArgIndex<0>>]>;

def int_nvvm_cp_async_bulk_wait_group_read :
Intrinsic<[],[llvm_i32_ty],[ImmArg<ArgIndex<0>>]>;

// mbarrier
def int_nvvm_mbarrier_init : ClangBuiltin<"__nvvm_mbarrier_init">,
Intrinsic<[],[llvm_ptr_ty, llvm_i32_ty],[IntrConvergent, IntrNoCallback]>;
Expand Down
16 changes: 16 additions & 0 deletions llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,22 @@ def CP_ASYNC_WAIT_ALL :
[(int_nvvm_cp_async_wait_all)]>,
Requires<[hasPTX<70>, hasSM<80>]>;

// cp.async.bulk variants of the commit/wait group
def CP_ASYNC_BULK_COMMIT_GROUP :
NVPTXInst<(outs), (ins), "cp.async.bulk.commit_group;",
[(int_nvvm_cp_async_bulk_commit_group)]>,
Requires<[hasPTX<80>, hasSM<90>]>;

def CP_ASYNC_BULK_WAIT_GROUP :
NVPTXInst<(outs), (ins i32imm:$n), "cp.async.bulk.wait_group $n;",
[(int_nvvm_cp_async_bulk_wait_group (i32 timm:$n))]>,
Requires<[hasPTX<80>, hasSM<90>]>;

def CP_ASYNC_BULK_WAIT_GROUP_READ :
NVPTXInst<(outs), (ins i32imm:$n), "cp.async.bulk.wait_group.read $n;",
[(int_nvvm_cp_async_bulk_wait_group_read (i32 timm:$n))]>,
Requires<[hasPTX<80>, hasSM<90>]>;

//-----------------------------------
// MBarrier Functions
//-----------------------------------
Expand Down
28 changes: 28 additions & 0 deletions llvm/test/CodeGen/NVPTX/intrinsics-sm90.ll
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,31 @@ define void @test_barrier_cluster_aligned() {
ret void
}

; CHECK-LABEL: test_cp_async_bulk_commit_group(
define void @test_cp_async_bulk_commit_group() {
; CHECK: cp.async.bulk.commit_group;
call void @llvm.nvvm.cp.async.bulk.commit.group()
ret void
}

; CHECK-LABEL: test_cp_async_bulk_wait_group(
define void @test_cp_async_bulk_wait_group() {
; CHECK: cp.async.bulk.wait_group 8;
call void @llvm.nvvm.cp.async.bulk.wait.group(i32 8)
; CHECK: cp.async.bulk.wait_group 0;
call void @llvm.nvvm.cp.async.bulk.wait.group(i32 0)
ret void
}

; CHECK-LABEL: test_cp_async_bulk_wait_group_read(
define void @test_cp_async_bulk_wait_group_read() {
; CHECK: cp.async.bulk.wait_group.read 8;
call void @llvm.nvvm.cp.async.bulk.wait.group.read(i32 8)
; CHECK: cp.async.bulk.wait_group.read 0;
call void @llvm.nvvm.cp.async.bulk.wait.group.read(i32 0)
ret void
}

declare i1 @llvm.nvvm.isspacep.shared.cluster(ptr %p);
declare ptr @llvm.nvvm.mapa(ptr %p, i32 %r);
declare ptr addrspace(3) @llvm.nvvm.mapa.shared.cluster(ptr addrspace(3) %p, i32 %r);
Expand All @@ -167,3 +192,6 @@ declare void @llvm.nvvm.barrier.cluster.arrive.aligned()
declare void @llvm.nvvm.barrier.cluster.arrive.relaxed.aligned()
declare void @llvm.nvvm.barrier.cluster.wait.aligned()
declare void @llvm.nvvm.fence.sc.cluster()
declare void @llvm.nvvm.cp.async.bulk.commit.group()
declare void @llvm.nvvm.cp.async.bulk.wait.group(i32)
declare void @llvm.nvvm.cp.async.bulk.wait.group.read(i32)

0 comments on commit 43531e7

Please sign in to comment.