[Mlir-commits] [llvm] [mlir] [LLVM][NVPTX] Add Rubin extensions for G2S Tensor intrinsics (PR #220029)

[Rajat Bajpai]

================
@@ -1455,6 +1466,54 @@ class CP_ASYNC_BULK_TENSOR_PREFETCH_OVERRIDE_INTR<int dim, string mode,
          [NVVM_ARG_NAME_PROP<!size(ArgsTy), "flag_cache_hint">.Prop]));
 }
 
+class CP_ASYNC_BULK_TENSOR_G2S_OVERRIDE_INTR<int dim, string mode, string override>
+    : CP_ASYNC_BULK_TENSOR_OVERRIDE_BASE<dim, mode, override> {
+  string Name = Prefix # "g2s" # Suffix;
+  // The override_addr_dim/dim_stride variants support only 32-bit multicast masks.
+  LLVMType MulticastTy = !if(!eq(override, "override_addr"), llvm_anyint_ty,
+                             llvm_i32_ty);
+  list<LLVMType> ArgsTy = !listconcat(
+      [llvm_shared_cluster_ptr_ty, // dst_smem_ptr
+       llvm_shared_ptr_ty,         // mbarrier_smem_ptr
+       llvm_ptr_ty],               // tensormap_ptr
+      OverrideTypes,               // override_addr
+      TensorDimsTy,                // actual tensor dims
+      Im2ColOffsetsTy,             // im2col offsets
+      [MulticastTy,                // cta_mask
+       llvm_i64_ty]);              // cache_hint
+  list<LLVMType> FlagsTy = [llvm_i1_ty,   // Flag for cta_mask
+                            llvm_i1_ty,   // Flag for cache_hint
+                            llvm_i32_ty,  // cta_group
+                            llvm_i32_ty]; // validate_pattern
+  defvar CtaGroupIdx = !add(!size(ArgsTy), 2);
+  defvar ValidateIdx = !add(!size(ArgsTy), 3);
+  list<IntrinsicProperty> IntrProp = !listconcat(
+      [IntrConvergent, WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>,
+       ReadOnly<ArgIndex<3>>,
+       Range<ArgIndex<CtaGroupIdx>, 0, 3>],
----------------
rajatbajpai wrote:

We are just emitting name for cta_group in MMA. ...., `/* cta_group= */` i32 1)

Would it fine to add that as a follow-up change for all G2S variants not just new one?

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