Skip to content

Commit ac08000

Browse files
Wolfram70Wolfram70
committed
[MLIR][NVVM] Add prefetch Ops
This change adds `prefetch.L1`, `prefetch.L2`, and `prefetch.L1.uniform` Ops to the NVVM dialect for the `prefetch` and `prefetchu` group of instructions. PTX Spec Reference: https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu
1 parent 7786ac0 commit ac08000

File tree

6 files changed

+290
-0
lines changed

6 files changed

+290
-0
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,12 +36,16 @@ constexpr int kSharedMemoryAlignmentBit = 128;
3636

3737
/// NVVM memory space identifiers.
3838
enum NVVMMemorySpace {
39+
/// Generic memory space identifier.
40+
kGenericMemorySpace = 0,
3941
/// Global memory space identifier.
4042
kGlobalMemorySpace = 1,
4143
/// Shared memory space identifier.
4244
kSharedMemorySpace = 3,
4345
/// Constant memory space identifier.
4446
kConstantMemorySpace = 4,
47+
/// Local memory space identifier.
48+
kLocalMemorySpace = 5,
4549
/// Tensor memory space identifier.
4650
/// Tensor memory is available only in arch-accelerated
4751
/// variants from sm100 onwards.

‎mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ include "mlir/Dialect/LLVMIR/LLVMTypes.td"
2525
def LLVM_PointerGeneric : LLVM_PointerInAddressSpace<0>;
2626
def LLVM_PointerGlobal : LLVM_PointerInAddressSpace<1>;
2727
def LLVM_PointerShared : LLVM_PointerInAddressSpace<3>;
28+
def LLVM_PointerLocal : LLVM_PointerInAddressSpace<5>;
2829
def LLVM_PointerTensor : LLVM_PointerInAddressSpace<6>;
2930
def LLVM_PointerSharedCluster : LLVM_PointerInAddressSpace<7>;
3031

@@ -118,6 +119,25 @@ class NVVM_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
118119
let mnemonic = attrMnemonic;
119120
}
120121

122+
// Cache Eviction Priority enum definitions
123+
def EvictNormal : I32EnumCase<"EvictNormal", 0, "evict_normal">;
124+
def EvictFirst : I32EnumCase<"EvictFirst", 1, "evict_first">;
125+
def EvictLast : I32EnumCase<"EvictLast", 2, "evict_last">;
126+
def EvictUnchanged : I32EnumCase<"EvictUnchanged", 3, "evict_unchanged">;
127+
def NoAllocate : I32EnumCase<"NoAllocate", 4, "no_allocate">;
128+
129+
def CacheEvictionPriority : I32Enum<"CacheEvictionPriority",
130+
"NVVM Cache Eviction Priority",
131+
[EvictNormal, EvictFirst, EvictLast,
132+
EvictUnchanged, NoAllocate]> {
133+
let cppNamespace = "::mlir::NVVM";
134+
}
135+
136+
def CacheEvictionPriorityAttr : EnumAttr<NVVM_Dialect, CacheEvictionPriority,
137+
"cache_eviction_priority"> {
138+
let assemblyFormat = "$value";
139+
}
140+
121141
//===----------------------------------------------------------------------===//
122142
// NVVM intrinsic operations
123143
//===----------------------------------------------------------------------===//
@@ -2333,6 +2353,75 @@ def NVVM_CpAsyncBulkTensorSharedCTAToGlobalOp :
23332353
let hasVerifier = 1;
23342354
}
23352355

2356+
//===----------------------------------------------------------------------===//
2357+
// NVVM Prefetch Ops
2358+
//===----------------------------------------------------------------------===//
2359+
2360+
def PrefetchCacheLevelL1 : I32EnumCase<"L1", 0, "L1">;
2361+
def PrefetchCacheLevelL2 : I32EnumCase<"L2", 1, "L2">;
2362+
2363+
def PrefetchCacheLevel : I32Enum<"PrefetchCacheLevel",
2364+
"NVVM Prefetch Cache Level",
2365+
[PrefetchCacheLevelL1, PrefetchCacheLevelL2]> {
2366+
let cppNamespace = "::mlir::NVVM";
2367+
}
2368+
2369+
def PrefetchCacheLevelAttr : EnumAttr<NVVM_Dialect, PrefetchCacheLevel, "prefetch_cache_level"> {
2370+
let assemblyFormat = "$value";
2371+
}
2372+
2373+
def NVVM_PrefetchOp : NVVM_Op<"prefetch"> {
2374+
let summary = "Brings the cache line containing an address into the specified cache level";
2375+
let description = [{
2376+
Operand `addr` can be a global, local or generic address pointer. No
2377+
operation is performed if `addr` maps to a `shared` memory location.
2378+
2379+
The `cacheLevel` attribute specifies the cache level to which the cache line
2380+
containing the specified address is brought.
2381+
2382+
The `evictPriority` attribute is optional and specifies the cache eviction
2383+
priority when `cacheLevel` is L2.
2384+
2385+
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
2386+
}];
2387+
let arguments = (ins PrefetchCacheLevelAttr:$cacheLevel,
2388+
AnyTypeOf<[LLVM_PointerGlobal,
2389+
LLVM_PointerLocal,
2390+
LLVM_PointerGeneric]>:$addr,
2391+
OptionalAttr<CacheEvictionPriorityAttr>:$evictPriority);
2392+
let assemblyFormat = "`level` `=` $cacheLevel `,` $addr (`,` `evict_priority` `=` $evictPriority^)? attr-dict `:` type($addr)";
2393+
let hasVerifier = 1;
2394+
2395+
let extraClassDeclaration = [{
2396+
static llvm::Intrinsic::ID getIntrinsicID(Operation &op);
2397+
}];
2398+
let llvmBuilder = [{
2399+
auto intId = NVVM::PrefetchOp::getIntrinsicID(*op);
2400+
createIntrinsicCall(builder, intId, $addr);
2401+
}];
2402+
}
2403+
2404+
def NVVM_PrefetchUniformOp : NVVM_Op<"prefetch.uniform"> {
2405+
let summary = "Brings the cache line containing an address into the specified uniform cache level";
2406+
let description = [{
2407+
Operand `addr` must be a generic address pointer and no operation is
2408+
performed if `addr` maps to a `const`, `local`, or `shared` memory location.
2409+
2410+
The `cacheLevel` attribute specifies the cache level to which the cache line
2411+
containing the specified address is brought. The only supported level is L1.
2412+
2413+
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
2414+
}];
2415+
let arguments = (ins PrefetchCacheLevelAttr:$cacheLevel,
2416+
LLVM_PointerGeneric:$addr);
2417+
let assemblyFormat = "`level` `=` $cacheLevel `,` $addr attr-dict `:` type($addr)";
2418+
let hasVerifier = 1;
2419+
2420+
let llvmBuilder = [{
2421+
createIntrinsicCall(builder, llvm::Intrinsic::nvvm_prefetchu_L1, $addr);
2422+
}];
2423+
}
2424+
23362425
def NVVM_PrefetchTensorMapOp : NVVM_Op<"prefetch.tensormap",
23372426
[DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
23382427
Arguments<(ins LLVM_AnyPointer:$tmaDescriptor, PtxPredicate:$predicate)> {

‎mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1205,6 +1205,35 @@ LogicalResult NVVM::VoteSyncOp::verify() {
12051205
return success();
12061206
}
12071207

1208+
LogicalResult NVVM::PrefetchOp::verify() {
1209+
auto evictPriority = getEvictPriority();
1210+
1211+
if (evictPriority && getCacheLevel() != NVVM::PrefetchCacheLevel::L2)
1212+
return emitOpError(
1213+
"cache eviction priority supported only for cache level L2");
1214+
1215+
if (evictPriority &&
1216+
(llvm::cast<LLVM::LLVMPointerType>(getAddr().getType())
1217+
.getAddressSpace() != NVVM::NVVMMemorySpace::kGlobalMemorySpace))
1218+
return emitOpError("cache eviction priority requires a global pointer");
1219+
1220+
if (evictPriority &&
1221+
*evictPriority != NVVM::CacheEvictionPriority::EvictNormal &&
1222+
*evictPriority != NVVM::CacheEvictionPriority::EvictLast)
1223+
return emitOpError(
1224+
"unsupported cache eviction priority, only evict_last and "
1225+
"evict_normal are supported");
1226+
1227+
return success();
1228+
}
1229+
1230+
LogicalResult NVVM::PrefetchUniformOp::verify() {
1231+
if (getCacheLevel() != NVVM::PrefetchCacheLevel::L1)
1232+
return emitOpError(
1233+
"unsupported cache level, the only supported level is L1");
1234+
return success();
1235+
}
1236+
12081237
/// Packs the given `field` into the `result`.
12091238
/// The `result` is 64-bits and each `field` can be 32-bits or narrower.
12101239
static llvm::Value *
@@ -1712,6 +1741,48 @@ NVVM::IDArgPair DotAccumulate4WayOp::getIntrinsicIDAndArgs(
17121741
return {ids[type], args};
17131742
}
17141743

1744+
llvm::Intrinsic::ID PrefetchOp::getIntrinsicID(Operation &op) {
1745+
auto curOp = llvm::cast<NVVM::PrefetchOp>(op);
1746+
NVVM::PrefetchCacheLevel cacheLevel = curOp.getCacheLevel();
1747+
std::optional<NVVM::CacheEvictionPriority> evictPriority =
1748+
curOp.getEvictPriority();
1749+
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
1750+
.getAddressSpace();
1751+
1752+
if (cacheLevel == NVVM::PrefetchCacheLevel::L1) {
1753+
switch (as) {
1754+
case NVVM::NVVMMemorySpace::kGenericMemorySpace:
1755+
return llvm::Intrinsic::nvvm_prefetch_L1;
1756+
case NVVM::NVVMMemorySpace::kGlobalMemorySpace:
1757+
return llvm::Intrinsic::nvvm_prefetch_global_L1;
1758+
case NVVM::NVVMMemorySpace::kLocalMemorySpace:
1759+
return llvm::Intrinsic::nvvm_prefetch_local_L1;
1760+
default:
1761+
llvm_unreachable("Invalid pointer address space");
1762+
}
1763+
} else if (cacheLevel == NVVM::PrefetchCacheLevel::L2) {
1764+
switch (as) {
1765+
case NVVM::NVVMMemorySpace::kGenericMemorySpace:
1766+
return llvm::Intrinsic::nvvm_prefetch_L2;
1767+
case NVVM::NVVMMemorySpace::kGlobalMemorySpace:
1768+
if (evictPriority) {
1769+
if (*evictPriority == NVVM::CacheEvictionPriority::EvictLast)
1770+
return llvm::Intrinsic::nvvm_prefetch_global_L2_evict_last;
1771+
else if (*evictPriority == NVVM::CacheEvictionPriority::EvictNormal)
1772+
return llvm::Intrinsic::nvvm_prefetch_global_L2_evict_normal;
1773+
else
1774+
llvm_unreachable("Invalid cache eviction priority");
1775+
}
1776+
return llvm::Intrinsic::nvvm_prefetch_global_L2;
1777+
case NVVM::NVVMMemorySpace::kLocalMemorySpace:
1778+
return llvm::Intrinsic::nvvm_prefetch_local_L2;
1779+
default:
1780+
llvm_unreachable("Invalid pointer address space");
1781+
}
1782+
}
1783+
llvm_unreachable("Invalid cache level");
1784+
}
1785+
17151786
//===----------------------------------------------------------------------===//
17161787
// NVVMDialect initialization, type parsing, and registration.
17171788
//===----------------------------------------------------------------------===//

‎mlir/test/Dialect/LLVMIR/nvvm.mlir

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -587,6 +587,29 @@ func.func @dot_accumulate_4way(%a_vec: vector<4xi8>, %b_vec: vector<4xi8>, %c: i
587587
return
588588
}
589589

590+
// CHECK-LABEL: @prefetch
591+
func.func @prefetch(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr: !llvm.ptr<1>) {
592+
// CHECK: nvvm.prefetch level = L1, %{{.*}}
593+
nvvm.prefetch level = L1, %gen_ptr : !llvm.ptr<0>
594+
// CHECK: nvvm.prefetch level = L1, %{{.*}}
595+
nvvm.prefetch level = L1, %local_ptr : !llvm.ptr<5>
596+
// CHECK: nvvm.prefetch level = L1, %{{.*}}
597+
nvvm.prefetch level = L1, %global_ptr : !llvm.ptr<1>
598+
// CHECK: nvvm.prefetch level = L2, %{{.*}}
599+
nvvm.prefetch level = L2, %gen_ptr : !llvm.ptr<0>
600+
// CHECK: nvvm.prefetch level = L2, %{{.*}}
601+
nvvm.prefetch level = L2, %local_ptr : !llvm.ptr<5>
602+
// CHECK: nvvm.prefetch level = L2, %{{.*}}
603+
nvvm.prefetch level = L2, %global_ptr : !llvm.ptr<1>
604+
// CHECK: nvvm.prefetch level = L2, %{{.*}}
605+
nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_last : !llvm.ptr<1>
606+
// CHECK: nvvm.prefetch level = L2, %{{.*}}
607+
nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_normal : !llvm.ptr<1>
608+
// CHECK: nvvm.prefetch.uniform level = L1, %{{.*}}
609+
nvvm.prefetch.uniform level = L1, %gen_ptr : !llvm.ptr
610+
return
611+
}
612+
590613
// -----
591614

592615
// Just check these don't emit errors.

‎mlir/test/Target/LLVMIR/nvvm/prefetch.mlir

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
2+
3+
llvm.func @prefetch_L1(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr: !llvm.ptr<1>) {
4+
// CHECK-LABEL: define void @prefetch_L1(ptr %0, ptr addrspace(5) %1, ptr addrspace(1) %2) {
5+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.L1(ptr %0)
6+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.local.L1(ptr addrspace(5) %1)
7+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L1(ptr addrspace(1) %2)
8+
// CHECK-NEXT: ret void
9+
// CHECK-NEXT: }
10+
nvvm.prefetch level = L1, %gen_ptr : !llvm.ptr<0>
11+
nvvm.prefetch level = L1, %local_ptr : !llvm.ptr<5>
12+
nvvm.prefetch level = L1, %global_ptr : !llvm.ptr<1>
13+
llvm.return
14+
}
15+
16+
llvm.func @prefetch_L2(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr: !llvm.ptr<1>) {
17+
// CHECK-LABEL: define void @prefetch_L2(ptr %0, ptr addrspace(5) %1, ptr addrspace(1) %2) {
18+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.L2(ptr %0)
19+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.local.L2(ptr addrspace(5) %1)
20+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L2(ptr addrspace(1) %2)
21+
// CHECK-NEXT: ret void
22+
// CHECK-NEXT: }
23+
nvvm.prefetch level = L2, %gen_ptr : !llvm.ptr<0>
24+
nvvm.prefetch level = L2, %local_ptr : !llvm.ptr<5>
25+
nvvm.prefetch level = L2, %global_ptr : !llvm.ptr<1>
26+
llvm.return
27+
}
28+
29+
llvm.func @prefetch_L2_eviction_priority(%global_ptr: !llvm.ptr<1>) {
30+
// CHECK-LABEL: define void @prefetch_L2_eviction_priority(ptr addrspace(1) %0) {
31+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L2.evict.last(ptr addrspace(1) %0)
32+
// CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L2.evict.normal(ptr addrspace(1) %0)
33+
// CHECK-NEXT: ret void
34+
// CHECK-NEXT: }
35+
nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_last : !llvm.ptr<1>
36+
nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_normal : !llvm.ptr<1>
37+
llvm.return
38+
}
39+
40+
llvm.func @prefetch_L1_uniform(%gen_ptr: !llvm.ptr) {
41+
// CHECK-LABEL: define void @prefetch_L1_uniform(ptr %0) {
42+
// CHECK-NEXT: call void @llvm.nvvm.prefetchu.L1(ptr %0)
43+
// CHECK-NEXT: ret void
44+
// CHECK-NEXT: }
45+
nvvm.prefetch.uniform level = L1, %gen_ptr : !llvm.ptr
46+
llvm.return
47+
}

‎mlir/test/Target/LLVMIR/nvvmir-invalid.mlir

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,3 +248,59 @@ llvm.func @nvvm_cvt_bf16x2_to_f8x2_invalid_rounding(%src : vector<2xbf16>) {
248248
%res = nvvm.convert.bf16x2.to.f8x2 <ue8m0> %src {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16> -> i16
249249
llvm.return
250250
}
251+
252+
// -----
253+
254+
llvm.func @nvvm_prefetch_L1_with_evict_priority(%global_ptr: !llvm.ptr<1>) {
255+
// expected-error @below {{cache eviction priority supported only for cache level L2}}
256+
nvvm.prefetch level = L1, %global_ptr, evict_priority = evict_last : !llvm.ptr<1>
257+
llvm.return
258+
}
259+
260+
// -----
261+
262+
llvm.func @nvvm_prefetch_L2_with_evict_last_invalid_addr_space(%local_ptr: !llvm.ptr<5>) {
263+
// expected-error @below {{cache eviction priority requires a global pointer}}
264+
nvvm.prefetch level = L2, %local_ptr, evict_priority = evict_last : !llvm.ptr<5>
265+
llvm.return
266+
}
267+
268+
// -----
269+
270+
llvm.func @nvvm_prefetch_L2_with_evict_normal_invalid_addr_space(%local_ptr: !llvm.ptr<5>) {
271+
// expected-error @below {{cache eviction priority requires a global pointer}}
272+
nvvm.prefetch level = L2, %local_ptr, evict_priority = evict_normal : !llvm.ptr<5>
273+
llvm.return
274+
}
275+
276+
// -----
277+
278+
llvm.func @nvvm_prefetch_L2_with_invalid_evict_first(%global_ptr: !llvm.ptr<1>) {
279+
// expected-error @below {{unsupported cache eviction priority, only evict_last and evict_normal are supported}}
280+
nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_first : !llvm.ptr<1>
281+
llvm.return
282+
}
283+
284+
// -----
285+
286+
llvm.func @nvvm_prefetch_L2_with_invalid_evict_unchanged(%global_ptr: !llvm.ptr<1>) {
287+
// expected-error @below {{unsupported cache eviction priority, only evict_last and evict_normal are supported}}
288+
nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_unchanged : !llvm.ptr<1>
289+
llvm.return
290+
}
291+
292+
// -----
293+
294+
llvm.func @nvvm_prefetch_L2_with_invalid_no_allocate(%global_ptr: !llvm.ptr<1>) {
295+
// expected-error @below {{unsupported cache eviction priority, only evict_last and evict_normal are supported}}
296+
nvvm.prefetch level = L2, %global_ptr, evict_priority = no_allocate : !llvm.ptr<1>
297+
llvm.return
298+
}
299+
300+
// -----
301+
302+
llvm.func @nvvm_prefetch_uniform_with_L2(%gen_ptr: !llvm.ptr) {
303+
// expected-error @below {{unsupported cache level, the only supported level is L1}}
304+
nvvm.prefetch.uniform level = L2, %gen_ptr : !llvm.ptr
305+
llvm.return
306+
}

0 commit comments

Comments
 (0)