merge prefetch.uniform into prefetch

Author: Wolfram70

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

@@ -2378,18 +2378,22 @@ def NVVM_PrefetchOp : NVVM_Op<"prefetch"> {
 
     The `cacheLevel` attribute specifies the cache level to which the cache line
     containing the specified address is brought.
+    
+    `uniform` can be specified after the `cacheLevel` to indicate that the 
+    prefetch is performed to the specified uniform cache level. If `uniform` is specified, `addr` must be a generic address pointer and no operation is performed if `addr` maps to a `const`, `local`, or `shared` memory location.
 
     The `evictPriority` attribute is optional and specifies the cache eviction
     priority when `cacheLevel` is L2.
 
     [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
   }];
   let arguments = (ins PrefetchCacheLevelAttr:$cacheLevel,
+                       UnitAttr:$uniform,
                        AnyTypeOf<[LLVM_PointerGlobal,
                                   LLVM_PointerLocal,
                                   LLVM_PointerGeneric]>:$addr,
                        OptionalAttr<CacheEvictionPriorityAttr>:$evictPriority);
-  let assemblyFormat = "`level` `=` $cacheLevel `,` $addr (`,` `evict_priority` `=` $evictPriority^)? attr-dict `:` type($addr)";
+  let assemblyFormat = "`level` `=` $cacheLevel (`uniform` $uniform^)? `,` $addr (`,` `evict_priority` `=` $evictPriority^)? attr-dict `:` type($addr)";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
@@ -2401,27 +2405,6 @@ def NVVM_PrefetchOp : NVVM_Op<"prefetch"> {
   }];
 }
 
-def NVVM_PrefetchUniformOp : NVVM_Op<"prefetch.uniform"> {
-  let summary = "Brings the cache line containing an address into the specified uniform cache level";
-  let description = [{
-    Operand `addr` must be a generic address pointer and no operation is 
-    performed if `addr` maps to a `const`, `local`, or `shared` memory location.
-    
-    The `cacheLevel` attribute specifies the cache level to which the cache line
-    containing the specified address is brought. The only supported level is L1.
-
-    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
-  }];
-  let arguments = (ins PrefetchCacheLevelAttr:$cacheLevel, 
-                       LLVM_PointerGeneric:$addr);
-  let assemblyFormat = "`level` `=` $cacheLevel `,` $addr attr-dict `:` type($addr)";
-  let hasVerifier = 1;
-
-  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

@@ -1206,15 +1206,27 @@ LogicalResult NVVM::VoteSyncOp::verify() {
 }
 
 LogicalResult NVVM::PrefetchOp::verify() {
+  unsigned addressSpace =
+      llvm::cast<LLVM::LLVMPointerType>(getAddr().getType()).getAddressSpace();
   auto evictPriority = getEvictPriority();
 
+  if (getUniform()) {
+    if (!(getCacheLevel() == NVVM::PrefetchCacheLevel::L1)) {
+      return emitOpError("unsupported cache level, the only supported uniform "
+                         "cache level is L1");
+    }
+    if (addressSpace != NVVM::NVVMMemorySpace::kGenericMemorySpace) {
+      return emitOpError(
+          "prefetch to uniform cache requires a generic pointer");
+    }
+  }
+
   if (evictPriority && getCacheLevel() != NVVM::PrefetchCacheLevel::L2)
     return emitOpError(
         "cache eviction priority supported only for cache level L2");
 
   if (evictPriority &&
-      (llvm::cast<LLVM::LLVMPointerType>(getAddr().getType())
-           .getAddressSpace() != NVVM::NVVMMemorySpace::kGlobalMemorySpace))
+      (addressSpace != NVVM::NVVMMemorySpace::kGlobalMemorySpace))
     return emitOpError("cache eviction priority requires a global pointer");
 
   if (evictPriority &&
@@ -1227,13 +1239,6 @@ LogicalResult NVVM::PrefetchOp::verify() {
   return success();
 }
 
-LogicalResult NVVM::PrefetchUniformOp::verify() {
-  if (getCacheLevel() != NVVM::PrefetchCacheLevel::L1)
-    return emitOpError(
-        "unsupported cache level, the only supported level is L1");
-  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 *
@@ -1771,6 +1776,13 @@ llvm::Intrinsic::ID PrefetchOp::getIntrinsicID(Operation &op) {
   unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
                     .getAddressSpace();
 
+  if (curOp.getUniform()) {
+    if (cacheLevel == NVVM::PrefetchCacheLevel::L1)
+      return llvm::Intrinsic::nvvm_prefetchu_L1;
+    else
+      llvm_unreachable("Invalid uniform cache level");
+  }
+
   if (cacheLevel == NVVM::PrefetchCacheLevel::L1) {
     switch (as) {
     case NVVM::NVVMMemorySpace::kGenericMemorySpace:

mlir/test/Dialect/LLVMIR/nvvm.mlir

@@ -614,8 +614,8 @@ func.func @prefetch(%gen_ptr: !llvm.ptr, %local_ptr: !llvm.ptr<5>, %global_ptr:
   nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_last : !llvm.ptr<1>
   // CHECK:   nvvm.prefetch level = L2, %{{.*}}
   nvvm.prefetch level = L2, %global_ptr, evict_priority = evict_normal : !llvm.ptr<1>
-  // CHECK:   nvvm.prefetch.uniform level = L1, %{{.*}}
-  nvvm.prefetch.uniform level = L1, %gen_ptr : !llvm.ptr
+  // CHECK:   nvvm.prefetch level = L1 uniform, %{{.*}}
+  nvvm.prefetch level = L1 uniform, %gen_ptr : !llvm.ptr
   return
 }
 

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

@@ -42,6 +42,6 @@ llvm.func @prefetch_L1_uniform(%gen_ptr: !llvm.ptr) {
   // CHECK-NEXT: call void @llvm.nvvm.prefetchu.L1(ptr %0)
   // CHECK-NEXT: ret void
   // CHECK-NEXT: }
-  nvvm.prefetch.uniform level = L1, %gen_ptr : !llvm.ptr
+  nvvm.prefetch level = L1 uniform, %gen_ptr : !llvm.ptr
   llvm.return
 }

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

@@ -300,7 +300,15 @@ llvm.func @nvvm_prefetch_L2_with_invalid_no_allocate(%global_ptr: !llvm.ptr<1>)
 // -----
 
 llvm.func @nvvm_prefetch_uniform_with_L2(%gen_ptr: !llvm.ptr) {
-  // expected-error @below {{unsupported cache level, the only supported level is L1}}
-  nvvm.prefetch.uniform level = L2, %gen_ptr : !llvm.ptr
+  // expected-error @below {{unsupported cache level, the only supported uniform cache level is L1}}
+  nvvm.prefetch level = L2 uniform, %gen_ptr : !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+llvm.func @nvvm_prefetch_uniform_with_invalid_addr_space(%global_ptr: !llvm.ptr<1>) {
+  // expected-error @below {{prefetch to uniform cache requires a generic pointer}}
+  nvvm.prefetch level = L1 uniform, %global_ptr : !llvm.ptr<1>
   llvm.return
 }