Skip to content

[NVPTX] Add prefetch tensormap variant #146203

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+50 −1
Open

[NVPTX] Add prefetch tensormap variant #146203

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
8 changes: 7 additions & 1 deletion llvm/docs/NVPTXUsage.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -971,6 +971,9 @@ Syntax:
declare void @llvm.nvvm.prefetch.L1(ptr %ptr)
declare void @llvm.nvvm.prefetch.L2(ptr %ptr)

declare void @llvm.nvvm.prefetch.tensormap.p0(ptr %ptr)
declare void @llvm.nvvm.prefetch.tensormap.p4(ptr addrspace(4) %const_ptr)

declare void @llvm.nvvm.prefetch.global.L2.evict.normal(ptr addrspace(1) %global_ptr)
declare void @llvm.nvvm.prefetch.global.L2.evict.last(ptr addrspace(1) %global_ptr)

Expand All @@ -983,7 +986,10 @@ The '``@llvm.nvvm.prefetch.*``' and '``@llvm.nvvm.prefetchu.*``' intrinsic
correspond to the '``prefetch.*``;' and '``prefetchu.*``' family of PTX instructions.
The '``prefetch.*``' instructions bring the cache line containing the
specified address in global or local memory address space into the
specified cache level (L1 or L2). The '`prefetchu.*``' instruction brings the cache line
specified cache level (L1 or L2). If the '``.tensormap``' qualifier is specified then the
prefetch instruction brings the cache line containing the specified address in the
'``.const``' or '``.param memory``' state space for subsequent use by the '``cp.async.bulk.tensor``'
instruction.The '`prefetchu.*``' instruction brings the cache line
containing the specified generic address into the specified uniform cache level.
If no address space is specified, it is assumed to be generic address. The intrinsic
uses and eviction priority which can be accessed by the '``.level::eviction_priority``' modifier.
Expand Down
4 changes: 4 additions & 0 deletions llvm/include/llvm/IR/IntrinsicsNVVM.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@

def llvm_global_ptr_ty : LLVMQualPointerType<1>; // (global)ptr
def llvm_shared_ptr_ty : LLVMQualPointerType<3>; // (shared)ptr
def llvm_constant_ptr_ty: LLVMQualPointerType<4>; // (const)ptr
def llvm_local_ptr_ty : LLVMQualPointerType<5>; // (local)ptr
def llvm_tmem_ptr_ty : LLVMQualPointerType<6>; // (tensor memory)ptr
def llvm_shared_cluster_ptr_ty : LLVMQualPointerType<7>; // (shared_cluster)ptr
Expand Down Expand Up @@ -2092,6 +2093,9 @@ let IntrProperties = [IntrArgMemOnly, ReadOnly<ArgIndex<0>>, NoCapture<ArgIndex<
def int_nvvm_prefetch_local_ # level : Intrinsic<[], [llvm_local_ptr_ty]>;
}

def int_nvvm_prefetch_tensormap_p0: DefaultAttrsIntrinsic<[], [llvm_ptr_ty]>;
def int_nvvm_prefetch_tensormap_p4 : DefaultAttrsIntrinsic<[], [llvm_constant_ptr_ty]>;

foreach eviction_priority = ["evict_normal", "evict_last"] in
def int_nvvm_prefetch_global_L2_ # eviction_priority : Intrinsic<[], [llvm_global_ptr_ty]>;

Expand Down
9 changes: 9 additions & 0 deletions llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -760,6 +760,15 @@ def PREFETCH_GLOBAL_L1 : PREFETCH_INTRS<"prefetch.global.L1">;
def PREFETCH_LOCAL_L1 : PREFETCH_INTRS<"prefetch.local.L1">;
def PREFETCH_GLOBAL_L2 : PREFETCH_INTRS<"prefetch.global.L2">;
def PREFETCH_LOCAL_L2 : PREFETCH_INTRS<"prefetch.local.L2">;
def PREFETCH_CONST_TENSORMAP : BasicNVPTXInst<(outs), (ins ADDR:$addr),
"prefetch.const.tensormap",
[(int_nvvm_prefetch_tensormap_p4 addr:$addr)]>,
Requires<[hasPTX<80>, hasSM<90>]>;

def PREFETCH_GENERIC_TENSORMAP : BasicNVPTXInst<(outs), (ins ADDR:$addr),
"prefetch.tensormap",
[(int_nvvm_prefetch_tensormap_p0 addr:$addr)]>,
Requires<[hasPTX<80>, hasSM<90>]>;

def PREFETCH_GLOBAL_L2_EVICT_NORMAL : BasicNVPTXInst<(outs), (ins ADDR:$addr),
"prefetch.global.L2::evict_normal",
Expand Down
30 changes: 30 additions & 0 deletions llvm/test/CodeGen/NVPTX/prefetch.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ declare void @llvm.nvvm.prefetch.local.L2(ptr addrspace(5) %local_ptr)
declare void @llvm.nvvm.prefetch.L1(ptr %ptr)
declare void @llvm.nvvm.prefetch.L2(ptr %ptr)

declare void @llvm.nvvm.prefetch.tensormap.p0(ptr %ptr)
declare void @llvm.nvvm.prefetch.tensormap.p4(ptr addrspace(4) %const_ptr)

declare void @llvm.nvvm.prefetch.global.L2.evict.normal(ptr addrspace(1) %global_ptr)
declare void @llvm.nvvm.prefetch.global.L2.evict.last(ptr addrspace(1) %global_ptr)

Expand Down Expand Up @@ -67,6 +70,33 @@ define void @prefetch_(ptr %ptr) {
ret void
}


define void @prefetch_generic_tensormap(ptr %ptr) {
; CHECK-PTX64-LABEL: prefetch_generic_tensormap(
; CHECK-PTX64: {
; CHECK-PTX64-NEXT: .reg .b64 %rd<2>;
; CHECK-PTX64-EMPTY:
; CHECK-PTX64-NEXT: // %bb.0:
; CHECK-PTX64-NEXT: ld.param.b64 %rd1, [prefetch_generic_tensormap_param_0];
; CHECK-PTX64-NEXT: prefetch.tensormap [%rd1];
; CHECK-PTX64-NEXT: ret;
tail call void @llvm.nvvm.prefetch.tensormap.p0(ptr %ptr)
ret void
}

define void @prefetch_const_tensormap(ptr addrspace(4) %const_ptr) {
; CHECK-PTX64-LABEL: prefetch_const_tensormap(
; CHECK-PTX64: {
; CHECK-PTX64-NEXT: .reg .b64 %rd<2>;
; CHECK-PTX64-EMPTY:
; CHECK-PTX64-NEXT: // %bb.0:
; CHECK-PTX64-NEXT: ld.param.b64 %rd1, [prefetch_const_tensormap_param_0];
; CHECK-PTX64-NEXT: prefetch.const.tensormap [%rd1];
; CHECK-PTX64-NEXT: ret;
tail call void @llvm.nvvm.prefetch.tensormap.p4(ptr addrspace(4) %const_ptr)
ret void
}

define void @prefetchu_l1(ptr %ptr) {
; CHECK-PTX64-LABEL: prefetchu_l1(
; CHECK-PTX64: {
Expand Down
Loading