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

[Durgadoss R via llvm-commits]

================
@@ -1312,6 +1312,25 @@ class TMA_IM2COL_OFFSETS_UTIL<int dim, string mode> {
     true                     : 0);
 }
 
+// Tile::gather4/scatter4 operate on a 2D tensor (5 co-ordinates), so they
+// always use a "2d" suffix regardless of dim.
+class TMA_DIMS_UTIL<int dim, string mode> {
+  int num_dims = !cond(
+    !eq(mode, "tile_gather4")  : 2,
+    !eq(mode, "tile_scatter4") : 2,
+    true                       : dim);
+}
+
+// Modes applicable for a given tensor dimensionality. The im2col modes are
+// only defined for 3D and higher tensors; for dim < 3 only "tile" is
+// supported. `includeW128` selects whether "im2col_w_128" mode is included
+// in the mode list.
+class TMA_MODES_UTIL<int dim, bit includeW128 = 0> {
+  list<string> all_modes = !listconcat(["tile", "im2col", "im2col_w"],
----------------
durga4github wrote:

let us use a defvar here so that this is not visible outside..

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