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

[Rajat Bajpai via llvm-commits]

================
@@ -1739,6 +1771,222 @@ prefetched in terms of bytes and it must be a multiple of 16.
 
 For more information, refer [PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-prefetch).
 
+#### '`llvm.nvvm.cp.async.bulk.prefetch.evict.priority`'
+
+##### Syntax:
+
+```llvm
+declare void @llvm.nvvm.cp.async.bulk.prefetch.evict.priority(ptr addrspace(1) %src, i32 %size, i32 %evict_policy)
+```
+
+##### Overview:
+
+The '`@llvm.nvvm.cp.async.bulk.prefetch.evict.priority`' intrinsic
+asynchronously prefetches `%size` bytes from global memory into the L2 cache
+using the selected eviction priority. It lowers to the
+`cp.async.bulk.prefetch.L2.global` PTX instruction with an L2 eviction-priority
+qualifier. Unlike `@llvm.nvvm.cp.async.bulk.prefetch.L2`, it takes an eviction
+policy instead of a cache hint. `%size` must be a multiple of 16.
+
+The trailing `i32 %evict_policy` is an immediate (`immarg`) argument that
+selects the L2 eviction policy.
+
+| Value | Name           | PTX qualifier       |
+|-------|----------------|---------------------|
+| 0     | `evict_normal` | (none, the default) |
+| 1     | `evict_last`   | `.L2::evict_last`   |
+
+This intrinsic requires `sm_107f` and PTX 9.4 or later.
+
+#### '`llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.[1-5]d`'
+
+##### Syntax:
+
+```llvm
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.1d(ptr %tensor_map, i32 %d0, i32 %evict_policy)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.2d(..., i32 %d0, i32 %d1, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.3d(..., i32 %d0, i32 %d1, i32 %d2, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.4d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.5d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, i32 %d4, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.tile.gather4.2d(ptr %tensor_map, i32 %x0, i32 %y0, i32 %y1, i32 %y2, i32 %y3, i32 %evict_policy)
+```
+
+##### Overview:
+
+These intrinsics asynchronously prefetch tile-mode tensor data from global
+memory into the L2 cache using the selected eviction priority. They are the
+eviction-priority forms of the corresponding
+`@llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.*` intrinsics, and the trailing
+`i32 %evict_policy` uses the encoding described for
+`@llvm.nvvm.cp.async.bulk.prefetch.evict.priority`. They do not take cache-hint
+operands and require `sm_107f` and PTX 9.4 or later.
+
+For more information, refer to the
+[PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-prefetch-tensor).
+
+#### '`llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.[3-5]d`'
+
+##### Syntax:
+
+```llvm
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.3d(ptr %tensor_map, i32 %d0, i32 %d1, i32 %d2, i16 %im2col0, i32 %evict_policy)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.4d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, i16 %im2col0, i16 %im2col1, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.5d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, i32 %d4, i16 %im2col0, i16 %im2col1, i16 %im2col2, ...)
+
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.w.3d(ptr %tensor_map, i32 %d0, i32 %d1, i32 %d2, i16 %wHalo, i16 %wOffset, i32 %evict_policy)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.w.4d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.w.5d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, i32 %d4, ...)
+
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.w.128.3d(ptr %tensor_map, i32 %d0, i32 %d1, i32 %d2, i16 %wHalo, i16 %wOffset, i32 %evict_policy)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.w.128.4d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.evict.priority.im2col.w.128.5d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, i32 %d4, ...)
+```
+
+##### Overview:
+
+These intrinsics asynchronously prefetch im2col-mode tensor data from global
+memory into the L2 cache using the selected eviction priority. An N-dimensional
+form takes N `i32` tensor coordinates. The `im2col` forms take N - 2 `i16`
+im2col offsets, while the `im2col.w` and `im2col.w.128` forms take
+`i16 %wHalo` and `i16 %wOffset`.
+
+The trailing `i32 %evict_policy` uses the encoding described for
+`@llvm.nvvm.cp.async.bulk.prefetch.evict.priority`. These forms do not take
+cache-hint operands and require `sm_107f` and PTX 9.4 or later.
+
+For more information, refer to the
+[PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-prefetch-tensor).
+
+#### '`llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override*.[1-5]d`'
+
+##### Syntax:
+
+```llvm
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.1d(ptr %tensor_map, ptr addrspace(1) %override_addr, i32 %d0, i64 %cache_hint, i1 %flag_cache_hint)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.2d(..., i32 %d0, i32 %d1, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.3d(..., i32 %d0, i32 %d1, i32 %d2, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.4d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.5d(..., i32 %d0, i32 %d1, i32 %d2, i32 %d3, i32 %d4, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.override.addr.2d(ptr %tensor_map, ptr addrspace(1) %override_addr, i32 %x0, i32 %y0, i32 %y1, i32 %y2, i32 %y3, i64 %cache_hint, i1 %flag_cache_hint)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.dim.1d(ptr %tensor_map, ptr addrspace(1) %override_addr, i16 %ts0, i32 %d0, i64 %cache_hint, i1 %flag_cache_hint)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.dim.stride.2d(ptr %tensor_map, ptr addrspace(1) %override_addr, i16 %ts0, i16 %ts1, i32 %stride0, i16 %upper_stride, i32 %d0, i32 %d1, i64 %cache_hint, i1 %flag_cache_hint)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.dim.stride.3d(..., i16 %ts0, i16 %ts1, i16 %ts2, i32 %stride0, i32 %stride1, i16 %upper_stride, i32 %d0, i32 %d1, i32 %d2, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.dim.stride.4d(..., i16 %ts0, i16 %ts1, i16 %ts2, i16 %ts3, i32 %stride0, i32 %stride1, i32 %stride2, i16 %upper_stride, i32 %d0, i32 %d1, i32 %d2, i32 %d3, ...)
+declare void @llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.override.addr.dim.stride.5d(..., i16 %ts0, i16 %ts1, i16 %ts2, i16 %ts3, i16 %ts4, i32 %stride0, i32 %stride1, i32 %stride2, i32 %stride3, i16 %upper_stride, i32 %d0, i32 %d1, i32 %d2, i32 %d3, i32 %d4, ...)
+```
+
+##### Overview:
+
+These intrinsics asynchronously prefetch tile-mode tensor data from global
+memory into the L2 cache while overriding tensor-map properties with explicit
+operands. They lower to the
+`cp.async.bulk.prefetch.tensor.<N>d.L2.global.tile` PTX instructions qualified
+with `.override::*`.
+
+The supported override variants and operands are described in
+[Tensor-Map Property Overrides](#tensor-map-property-overrides).
+
+- The last argument is a compile-time boolean indicating whether `%cache_hint`
+  is valid. When set, it
----------------
rajatbajpai wrote:

fixed

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