[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

Author: Wolfram70

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

@@ -42,6 +42,8 @@ enum NVVMMemorySpace {
   kSharedMemorySpace = 3,
   /// Constant memory space identifier.
   kConstantMemorySpace = 4,
+  /// Local memory space identifier.
+  kLocalMemorySpace = 5,
   /// Tensor memory space identifier.
   /// Tensor memory is available only in arch-accelerated
   /// variants from sm100 onwards.

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

@@ -25,6 +25,7 @@ include "mlir/Dialect/LLVMIR/LLVMTypes.td"
 def LLVM_PointerGeneric : LLVM_PointerInAddressSpace<0>;
 def LLVM_PointerGlobal : LLVM_PointerInAddressSpace<1>;
 def LLVM_PointerShared : LLVM_PointerInAddressSpace<3>;
+def LLVM_PointerLocal : LLVM_PointerInAddressSpace<5>;
 def LLVM_PointerTensor : LLVM_PointerInAddressSpace<6>;
 def LLVM_PointerSharedCluster : LLVM_PointerInAddressSpace<7>;
 
@@ -118,6 +119,25 @@ class NVVM_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
   let mnemonic = attrMnemonic;
 }
 
+// Cache Eviction Priority enum definitions
+def EvictNormal : I32EnumCase<"EvictNormal", 0, "evict_normal">;
+def EvictFirst : I32EnumCase<"EvictFirst", 1, "evict_first">;
+def EvictLast : I32EnumCase<"EvictLast", 2, "evict_last">;
+def EvictUnchanged : I32EnumCase<"EvictUnchanged", 3, "evict_unchanged">;
+def NoAllocate : I32EnumCase<"NoAllocate", 4, "no_allocate">;
+
+def CacheEvictionPriority : I32Enum<"CacheEvictionPriority",
+                                    "NVVM Cache Eviction Priority",
+                                    [EvictNormal, EvictFirst, EvictLast,   
+                                     EvictUnchanged, NoAllocate]> {
+  let cppNamespace = "::mlir::NVVM";
+}
+
+def CacheEvictionPriorityAttr : EnumAttr<NVVM_Dialect, CacheEvictionPriority, 
+                                         "cache_eviction_priority"> {
+  let assemblyFormat = "$value";
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM intrinsic operations
 //===----------------------------------------------------------------------===//
@@ -2333,6 +2353,79 @@ def NVVM_CpAsyncBulkTensorSharedCTAToGlobalOp :
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM Prefetch Ops
+//===----------------------------------------------------------------------===//
+
+def NVVM_PrefetchL1Op : NVVM_Op<"prefetch.L1"> {
+  let summary = "Brings the cache line containing the specified address into L1 cache";
+  let description = [{
+    Brings the cache line containing the specified address into L1 cache.
+
+    Operand `addr` can be a global, local or generic address pointer.
+    No operation is performed if `addr` maps to a `shared` memory location.
+
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
+  }];
+  let arguments = (ins AnyTypeOf<[LLVM_PointerGlobal,
+                                  LLVM_PointerLocal,
+                                  LLVM_PointerGeneric]>:$addr);
+  let assemblyFormat = "$addr attr-dict `:` type($addr)";
+
+  let extraClassDeclaration = [{
+    static llvm::Intrinsic::ID getIntrinsicID(Operation &op);
+  }];
+  let llvmBuilder = [{
+    auto intId = NVVM::PrefetchL1Op::getIntrinsicID(*op);
+    createIntrinsicCall(builder, intId, $addr);
+  }];
+}
+
+def NVVM_PrefetchL2Op : NVVM_Op<"prefetch.L2"> {
+  let summary = "Brings the cache line containing the specified address into L2 cache";
+  let description = [{
+    Brings the cache line containing the specified address into L2 cache.
+
+    Operand `addr` can be a global, local or generic address pointer.
+    No operation is performed if `addr` maps to a `shared` memory location.
+
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
+  }];
+  let arguments = (ins AnyTypeOf<[LLVM_PointerGlobal,
+                                  LLVM_PointerLocal,
+                                  LLVM_PointerGeneric]>:$addr,
+                       OptionalAttr<CacheEvictionPriorityAttr>:$evictPriority);
+  let assemblyFormat = "$addr (`,` `evict_priority` `=` $evictPriority^)? attr-dict `:` type($addr)";
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    static llvm::Intrinsic::ID getIntrinsicID(Operation &op);
+  }];
+  let llvmBuilder = [{
+    auto intId = NVVM::PrefetchL2Op::getIntrinsicID(*op);
+    createIntrinsicCall(builder, intId, $addr);
+  }];
+}
+
+def NVVM_PrefetchL1UniformOp : NVVM_Op<"prefetch.L1.uniform"> {
+  let summary = "Brings the cache line containing the specified address into L1 uniform cache";
+  let description = [{
+    Brings the cache line containing the specified address into L1 uniform 
+    cache.
+
+    Operand `addr` must be a generic address pointer and no operation is 
+    performed if `addr` maps to a `const`, `local`, or `shared` memory location.
+
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
+  }];
+  let arguments = (ins LLVM_PointerGeneric:$addr);
+  let assemblyFormat = "$addr attr-dict `:` type($addr)";
+
+  let llvmBuilder = [{
+    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_prefetchu_L1, $addr);
+  }];
+}
+
 def NVVM_PrefetchTensorMapOp : NVVM_Op<"prefetch.tensormap",
                     [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
   Arguments<(ins LLVM_AnyPointer:$tmaDescriptor, PtxPredicate:$predicate)> {

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

@@ -1205,6 +1205,20 @@ LogicalResult NVVM::VoteSyncOp::verify() {
   return success();
 }
 
+LogicalResult NVVM::PrefetchL2Op::verify() {
+  auto evictPriority = getEvictPriority();
+  if (evictPriority &&
+      (llvm::cast<LLVM::LLVMPointerType>(getAddr().getType())
+           .getAddressSpace() != NVVM::NVVMMemorySpace::kGlobalMemorySpace))
+    return emitOpError(
+        "prefetch with cache eviction priority requires a global pointer");
+  if (evictPriority &&
+      *evictPriority != NVVM::CacheEvictionPriority::EvictNormal &&
+      *evictPriority != NVVM::CacheEvictionPriority::EvictLast)
+    return emitOpError("invalid cache eviction priority");
+  return success();
+}
+
 /// Packs the given `field` into the `result`.
 /// The `result` is 64-bits and each `field` can be 32-bits or narrower.
 static llvm::Value *
@@ -1712,6 +1726,46 @@ NVVM::IDArgPair DotAccumulate4WayOp::getIntrinsicIDAndArgs(
   return {ids[type], args};
 }
 
+llvm::Intrinsic::ID PrefetchL1Op::getIntrinsicID(Operation &op) {
+  auto curOp = llvm::cast<NVVM::PrefetchL1Op>(op);
+  switch (llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
+              .getAddressSpace()) {
+  case 0:
+    return llvm::Intrinsic::nvvm_prefetch_L1;
+  case NVVM::NVVMMemorySpace::kGlobalMemorySpace:
+    return llvm::Intrinsic::nvvm_prefetch_global_L1;
+  case NVVM::NVVMMemorySpace::kLocalMemorySpace:
+    return llvm::Intrinsic::nvvm_prefetch_local_L1;
+  default:
+    llvm_unreachable("Invalid pointer address space");
+  }
+}
+
+llvm::Intrinsic::ID PrefetchL2Op::getIntrinsicID(Operation &op) {
+  auto curOp = llvm::cast<NVVM::PrefetchL2Op>(op);
+  auto evictPriority = curOp.getEvictPriority();
+
+  switch (llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
+              .getAddressSpace()) {
+  case 0:
+    return llvm::Intrinsic::nvvm_prefetch_L2;
+  case NVVM::NVVMMemorySpace::kGlobalMemorySpace:
+    if (evictPriority) {
+      if (*evictPriority == NVVM::CacheEvictionPriority::EvictLast)
+        return llvm::Intrinsic::nvvm_prefetch_global_L2_evict_last;
+      else if (*evictPriority == NVVM::CacheEvictionPriority::EvictNormal)
+        return llvm::Intrinsic::nvvm_prefetch_global_L2_evict_normal;
+      else
+        llvm_unreachable("Invalid cache eviction priority");
+    }
+    return llvm::Intrinsic::nvvm_prefetch_global_L2;
+  case NVVM::NVVMMemorySpace::kLocalMemorySpace:
+    return llvm::Intrinsic::nvvm_prefetch_local_L2;
+  default:
+    llvm_unreachable("Invalid pointer address space");
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // NVVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/LLVMIR/nvvm.mlir

@@ -587,6 +587,29 @@ func.func @dot_accumulate_4way(%a_vec: vector<4xi8>, %b_vec: vector<4xi8>, %c: i
   return
 }
 
+// CHECK-LABEL: @prefetch
+func.func @prefetch(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr: !llvm.ptr<1>) {
+  // CHECK:   nvvm.prefetch.L1 %{{.*}}
+  nvvm.prefetch.L1 %gen_ptr : !llvm.ptr<0>
+  // CHECK:   nvvm.prefetch.L1 %{{.*}}
+  nvvm.prefetch.L1 %local_ptr : !llvm.ptr<5>
+  // CHECK:   nvvm.prefetch.L1 %{{.*}}
+  nvvm.prefetch.L1 %global_ptr : !llvm.ptr<1>
+  // CHECK:   nvvm.prefetch.L2 %{{.*}}
+  nvvm.prefetch.L2 %gen_ptr : !llvm.ptr<0>
+  // CHECK:   nvvm.prefetch.L2 %{{.*}}
+  nvvm.prefetch.L2 %local_ptr : !llvm.ptr<5>
+  // CHECK:   nvvm.prefetch.L2 %{{.*}}
+  nvvm.prefetch.L2 %global_ptr : !llvm.ptr<1>
+  // CHECK:   nvvm.prefetch.L2 %{{.*}}
+  nvvm.prefetch.L2 %global_ptr, evict_priority = evict_last : !llvm.ptr<1>
+  // CHECK:   nvvm.prefetch.L2 %{{.*}}
+  nvvm.prefetch.L2 %global_ptr, evict_priority = evict_normal : !llvm.ptr<1>
+  // CHECK:   nvvm.prefetch.L1.uniform %{{.*}}
+  nvvm.prefetch.L1.uniform %gen_ptr : !llvm.ptr
+  return
+}
+
 // -----
 
 // Just check these don't emit errors.

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

@@ -0,0 +1,47 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @prefetch_L1(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr: !llvm.ptr<1>) {
+  // CHECK-LABEL: define void @prefetch_L1(ptr %0, ptr addrspace(5) %1, ptr addrspace(1) %2) {
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.L1(ptr %0)
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.local.L1(ptr addrspace(5) %1)
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L1(ptr addrspace(1) %2)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.prefetch.L1 %gen_ptr : !llvm.ptr<0>
+  nvvm.prefetch.L1 %local_ptr : !llvm.ptr<5>
+  nvvm.prefetch.L1 %global_ptr : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @prefetch_L2(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr: !llvm.ptr<1>) {
+  // CHECK-LABEL: define void @prefetch_L2(ptr %0, ptr addrspace(5) %1, ptr addrspace(1) %2) {
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.L2(ptr %0)
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.local.L2(ptr addrspace(5) %1)
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L2(ptr addrspace(1) %2)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.prefetch.L2 %gen_ptr : !llvm.ptr<0>
+  nvvm.prefetch.L2 %local_ptr : !llvm.ptr<5>
+  nvvm.prefetch.L2 %global_ptr : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @prefetch_L2_eviction_priority(%global_ptr: !llvm.ptr<1>) {
+  // CHECK-LABEL: define void @prefetch_L2_eviction_priority(ptr addrspace(1) %0) {
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L2.evict.last(ptr addrspace(1) %0)
+  // CHECK-NEXT: call void @llvm.nvvm.prefetch.global.L2.evict.normal(ptr addrspace(1) %0)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.prefetch.L2 %global_ptr, evict_priority = evict_last : !llvm.ptr<1>
+  nvvm.prefetch.L2 %global_ptr, evict_priority = evict_normal : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @prefetch_L1_uniform(%gen_ptr: !llvm.ptr) {
+  // CHECK-LABEL: define void @prefetch_L1_uniform(ptr %0) {
+  // CHECK-NEXT: call void @llvm.nvvm.prefetchu.L1(ptr %0)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.prefetch.L1.uniform %gen_ptr : !llvm.ptr
+  llvm.return
+}

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

@@ -248,3 +248,19 @@ llvm.func @nvvm_cvt_bf16x2_to_f8x2_invalid_rounding(%src : vector<2xbf16>) {
   %res = nvvm.convert.bf16x2.to.f8x2 <ue8m0> %src {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16> -> i16
   llvm.return
 }
+
+// -----
+
+llvm.func @nvvm_prefetch_L2_with_evict_invalid_addr_space(%local_ptr: !llvm.ptr<5>) {
+  // expected-error @below {{prefetch with cache eviction priority requires a global pointer}}
+  nvvm.prefetch.L2 %local_ptr, evict_priority = evict_last : !llvm.ptr<5>
+  llvm.return
+}
+
+// -----
+
+llvm.func @nvvm_prefetch_L2_with_evict_invalid_evict_priority(%global_ptr: !llvm.ptr<1>) {
+  // expected-error @below {{invalid cache eviction priority}}
+  nvvm.prefetch.L2 %global_ptr, evict_priority = evict_first : !llvm.ptr<1>
+  llvm.return
+}