[llvm] [NVVM][NVPTX] Add tensor map override support in S2G and Reduce intrinsics (PR #215503)

Rajat Bajpai via llvm-commits llvm-commits at lists.llvm.org
Thu Aug 13 02:57:46 PDT 2026


================
@@ -642,21 +642,27 @@ class TMA_DIMS_UTIL<int dim, string mode = ""> {
 }
 
 class TMA_IM2COL_UTIL<int dim, string mode> {
-  // For im2col_w/w_128 modes, number of offsets is always 2.
-  // For im2col mode, offsets is (dim - 2).
-  // For non-im2col modes (i.e. tile) there are no offsets.
-  int offsets = !cond(
-                  !eq(mode, "im2col") : !sub(dim, 2),
-                  !eq(mode, "im2col_w") : 2,
-                  !eq(mode, "im2col_w_128") : 2,
-                  true : 0); // for all other modes
+  int offsets = TMA_IM2COL_OFFSETS_UTIL<dim, mode>.num;
 
   dag ins_dag = !if(!gt(offsets, 0),
     !dag(ins, !listsplat(B16, offsets), !foreach(i, !range(offsets), "im2col" # i)),
     (ins));
   string base_str = !interleave(!foreach(i, !range(offsets), "$im2col" # i), ", ");
 }
 
+// Maps a TMA tensor copy "mode" to equivalent assembly mnemonic.
+//   tile          -> tile
+//   im2col        -> im2col_no_offs
+//   im2col_w      -> im2col_no_offs::w
+//   tile_scatter4 -> tile::scatter4
+class TMA_MODE_ASM_UTIL<string mode> {
----------------
rajatbajpai wrote:

Sure, makes sense.

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


More information about the llvm-commits mailing list