[llvm] [mlir] [LLVM][NVPTX] Add async bulk copy global to shared extensions (PR #222323)

[Rajat Bajpai via llvm-commits]

================
@@ -3519,31 +3529,75 @@ let IntrProperties = [NoCapture<ArgIndex<0>>, ImmArg<ArgIndex<1>>, IntrHasSideEf
 }
 
 // Intrinsics for Bulk Copy using TMA (non-tensor)
-// From Global to Shared Cluster
+// From Global to Shared Cluster (weak memory ordering).
 def int_nvvm_cp_async_bulk_global_to_shared_cluster
   : DefaultAttrsIntrinsicFlags<[],
       [llvm_shared_cluster_ptr_ty, // dst_shared_cluster_ptr
-       llvm_shared_ptr_ty,         // mbarrier_ptr
-       llvm_global_ptr_ty,         // src_gmem_ptr
-       llvm_i32_ty,                // copy_size
-       llvm_i16_ty,                // cta_mask
-       llvm_i64_ty],               // cache_hint
+        llvm_shared_ptr_ty,        // mbarrier_ptr
+        llvm_global_ptr_ty,        // src_gmem_ptr
+        llvm_i32_ty,               // copy_size
+        llvm_anyint_ty,            // cta_mask (i16 or i32 overload)
+        llvm_i64_ty],              // cache_hint
       [llvm_i1_ty,                 // Flag for cta_mask
-       llvm_i1_ty],                // Flag for cache_hint
-      [IntrConvergent, IntrArgMemOnly,
-       WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>]>;
+        llvm_i1_ty,                // Flag for cache_hint
+        llvm_i32_ty],              // validate_pattern
+      !listconcat([IntrConvergent, IntrArgMemOnly,
+                   WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>],
+                   NVVM_TMA_VALIDATE_PATTERN_ARGPROP<8>.Prop)>;
+
+// From Global to Shared Cluster (relaxed memory ordering semantics).
+def int_nvvm_cp_async_bulk_global_to_shared_cluster_relaxed
+  : DefaultAttrsIntrinsicFlags<[],
+      [llvm_shared_cluster_ptr_ty, // dst_shared_cluster_ptr
+        llvm_shared_ptr_ty,        // mbarrier_ptr
+        llvm_global_ptr_ty,        // src_gmem_ptr
+        llvm_i32_ty,               // copy_size
+        llvm_anyint_ty,            // cta_mask (i16 or i32 overload)
+        llvm_i64_ty],              // cache_hint
+      [llvm_i1_ty,                 // Flag for cta_mask
+        llvm_i1_ty,                // Flag for cache_hint
+        llvm_i32_ty,               // scope (cta/cluster/gpu/sys)
+        llvm_i32_ty],              // validate_pattern
+      !listconcat([IntrConvergent, IntrArgMemOnly,
+                   WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>],
+                   NVVM_MEM_SCOPE_ARGPROP<8>.Prop,
+                   NVVM_TMA_VALIDATE_PATTERN_ARGPROP<9>.Prop)>;
 
-// From Global to Shared CTA
+// From Global to Shared CTA (weak memory ordering).
 def int_nvvm_cp_async_bulk_global_to_shared_cta
   : DefaultAttrsIntrinsicFlags<[],
       [llvm_shared_ptr_ty, // dst_shared_cta_ptr
        llvm_shared_ptr_ty, // mbarrier_ptr
        llvm_global_ptr_ty, // src_gmem_ptr
        llvm_i32_ty,        // copy_size
+       llvm_i32_ty,        // ignore_bytes_left (controlled by flag_oob)
+       llvm_i32_ty,        // ignore_bytes_right (controlled by flag_oob)
        llvm_i64_ty],       // cache_hint
-      [llvm_i1_ty],        // Flag for cache_hint
-      [IntrConvergent, IntrArgMemOnly,
-       WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>]>;
+      [llvm_i1_ty,         // Flag for ignore_oob
+       llvm_i1_ty,         // Flag for cache_hint
+       llvm_i32_ty],       // validate_pattern
+      !listconcat([IntrConvergent, IntrArgMemOnly,
+                   WriteOnly<ArgIndex<0>>, ReadOnly<ArgIndex<2>>],
+                   NVVM_TMA_VALIDATE_PATTERN_ARGPROP<9>.Prop)>;
+
+// From Global to Shared CTA (relaxed memory ordering semantics).
+def int_nvvm_cp_async_bulk_global_to_shared_cta_relaxed
+  : DefaultAttrsIntrinsicFlags<[],
+      [llvm_shared_ptr_ty, // dst_shared_cta_ptr
+       llvm_shared_ptr_ty, // mbarrier_ptr
+       llvm_global_ptr_ty, // src_gmem_ptr
+       llvm_i32_ty,        // copy_size
+       llvm_i32_ty,        // ignore_bytes_left (controlled by flag_oob)
+       llvm_i32_ty,        // ignore_bytes_right (controlled by flag_oob)
----------------
rajatbajpai wrote:

Yes, that’s one alternative, but I didn’t go with that approach because I see report validity and OOB as the same kind of optional secondary behavior on the base operation. If one is modeled as a flag, I think the other should be as well, unless there is a strong reason to treat them differently. I didn’t find one.

Suppose we model both OOB and validate patterns as separate records. That might be fine for the bulk copy variants (less number of intrinsics), but then we would need to model the tensor variants similarly to keep the design consistent. There, this approach would blow up because of the various modes and dimensions.

Given this implication, I think using flags for both secondary functionalities gives us a more compact and consistent design.


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