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

Rajat Bajpai via llvm-commits llvm-commits at lists.llvm.org
Sun Aug 23 03:14:09 PDT 2026


================
@@ -1146,6 +1179,180 @@ foreach dim = 3...5 in {
 defm TMA_TENSOR_PF_TILE_GATHER4_2D : TMA_TENSOR_PREFETCH_INTR<5, "tile_gather4",
                                      [hasTMABlackwellSupport]>;
 
+// TMA Tensor Prefetch from Global memory with Override Global Address
+// (with cache hint and eviction priority forms).
+multiclass NVPTX_TMA_TENSOR_PREFETCH_OVERRIDE_INTR<int dim, string mode,
+                                                   string override> {
+  defvar base = TMA_TENSOR_OVERRIDE_BASE<dim, mode, override>;
+  defvar num_dims = TMA_DIMS_UTIL<dim, mode>.num_dims;
+
+  defvar prefix = "cp.async.bulk.prefetch.tensor." # num_dims # "d.L2.global";
+  defvar suffix = "." # !subst("_", "::", mode) # base.override_suffix;
+  defvar asm_str = " [$tmap, $override_addr, " # base.tensor_size_component #
+                   "{{" # base.dims_str # "}}]" # base.im2col_component;
+
+  defvar intr =
+      !cast<Intrinsic>("int_nvvm_cp_async_bulk_tensor_prefetch_" #
+      mode # "_" # override # "_" # num_dims # "d");
+  defvar intr_evict =
+      !cast<Intrinsic>("int_nvvm_cp_async_bulk_tensor_prefetch_evict_priority_" #
+      mode # "_" # override # "_" # num_dims # "d");
----------------
rajatbajpai wrote:

done

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


More information about the llvm-commits mailing list