[Mlir-commits] [mlir] [MLIR][NVVM] Update prefetch.tensormap Op (PR #153134)

[Guray Ozen]

================
@@ -2464,15 +2465,45 @@ def NVVM_PrefetchOp : NVVM_Op<"prefetch"> {
   }];
 }
 
-def NVVM_PrefetchTensorMapOp : NVVM_Op<"prefetch.tensormap",
-                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
-  Arguments<(ins LLVM_AnyPointer:$tmaDescriptor, PtxPredicate:$predicate)> {
-  let assemblyFormat = "$tmaDescriptor (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";
+def NVVM_PrefetchTensorMapOp : NVVM_Op<"prefetch.tensormap", 
+    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>, NVVMRequiresSM<90>]> {
+  let summary = "Brings the cache line containing an address from the constant (`.const`) or parameter (`.param`) state spaces for subsequent use by the `cp.async.bulk.tensor` instruction";
+  let description = [{
+    Operand `tmaDescriptor` can be a `const` or generic address 
+    pointer.
+    If it is a generic address pointer, it must map to a memory 
+    location in the [const](https://docs.nvidia.com/cuda/parallel-thread-execution/#constant-state-space) or [param](https://docs.nvidia.com/cuda/parallel-thread-execution/#parameter-state-space) state space.
+
+    If the `in_param_space` attribute is present, the `tmaDescriptor` must be a
+    generic address pointer and is treated as pointing to a memory location in 
+    the `param` state space. 
+
+    [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_PointerGeneric,
+                                  LLVM_PointerConst]>:$tmaDescriptor,
+                       PtxPredicate:$predicate,
+                       UnitAttr:$in_param_space);
+  let assemblyFormat = "(`param` $in_param_space^)? $tmaDescriptor (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";
----------------
grypp wrote:

```suggestion
  let assemblyFormat = "$in_param_space $tmaDescriptor (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";
```

https://github.com/llvm/llvm-project/pull/153134