[llvm] [LLVM][NVPTX] Add support for TMA prefetch Rubin extensions (PR #217630)

[Rajat Bajpai via llvm-commits]

================
@@ -1357,10 +1363,85 @@ class CP_ASYNC_BULK_TENSOR_REDUCE_OVERRIDE_INTR<int dim, string mode, string ove
         Range<ArgIndex<RedOpIdx>, 0, 8>,
         ArgInfo<ArgIndex<RedOpIdx>, [ArgName<"red_op">,
                                      ImmArgPrinter<"printTMAReductionOp">]>,
-        ArgInfo<ArgIndex<!add(RedOpIdx, 1)>, [ArgName<"flag_cache_hint">]>
+        NVVM_CACHE_HINT_ARGPROP<!add(RedOpIdx, 1)>.Prop
+  ];
+}
+
+// Helper that emits the intrinsic-property set for the evict_policy argument.
+// `upperBound` is the exclusive upper bound of the allowed policy value range.
+// Prefetch accepts both EVICT_NORMAL and EVICT_LAST (default 2).
+// Applypriority only accepts EVICT_NORMAL (upperBound = 1).
+class NVVM_EVICT_POLICY_ARGPROP<int argNo, int upperBound = 2> {
+  list<IntrinsicProperty> Prop = [
+    ImmArg<ArgIndex<argNo>>,
+    Range<ArgIndex<argNo>, 0, upperBound>,
+    ArgInfo<ArgIndex<argNo>, [ArgName<"evict_policy">,
+                              ImmArgPrinter<"printEvictPolicyType">]>
   ];
 }
 
+class CP_ASYNC_BULK_TENSOR_PREFETCH_OVERRIDE_INTR<int dim, string mode,
+                                                  string override,
+                                                  bit IsEvict = 0>
+    : CP_ASYNC_BULK_TENSOR_OVERRIDE_BASE<dim, mode, override> {
+  string Name = Prefix # "prefetch" # !if(IsEvict, "_evict_priority", "") # Suffix;
+
+  list<LLVMType> ArgsTy = !listconcat(
+      [llvm_ptr_ty],       // tensormap_ptr
+      OverrideTypes,       // override_addr (src)
+      TensorDimsTy,        // actual tensor dims
+      Im2ColOffsetsTy,     // im2col offsets
+      !if(IsEvict, [llvm_i32_ty], [llvm_i64_ty]));
+
+  list<LLVMType> FlagsTy = !if(IsEvict, [], [llvm_i1_ty]);
+
+  defvar PolicyIdx = !sub(!size(ArgsTy), 1);
+  list<IntrinsicProperty> IntrProp = !listconcat(
+      [IntrConvergent, ReadOnly<ArgIndex<0>>, ReadOnly<ArgIndex<1>>],
+      !if(IsEvict,
+          NVVM_EVICT_POLICY_ARGPROP<PolicyIdx>.Prop,
+         [NVVM_CACHE_HINT_ARGPROP<!size(ArgsTy)>.Prop]));
+}
+
+// Applypriority TMA variants.
+class APPLYPRIORITY_BULK_TENSOR_INTR<int dim, string mode> {
+  string Name = "int_nvvm_applypriority_async_bulk_tensor_evict_priority_" #
+                mode # "_" # !if(!eq(mode, "tile_gather4"), 2, dim) # "d";
----------------
rajatbajpai wrote:

Not directly, utils are part of backend TD.

I moved them to frontend TD.

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