[Language] Support atomic add with ret (#870)

* Add atomic operations for CUDA templates in new atomic.h file

- Introduced atomic functions including AtomicMax, AtomicMin, AtomicAdd, and their return variants for various data types.
- Implemented support for half, bfloat16, and float types with appropriate memory ordering.
- Moved atomic-related utilities from common.h to the new atomic.h file for better organization.
- Added Python bindings for atomic operations in tilelang, including atomic_max, atomic_min, atomic_add, and their vectorized counterparts.
- Updated customize.py to utilize the new atomic functions, enhancing modularity and maintainability.

* Refactor atomic operations in CUDA templates for improved readability

- Reformatted atomic operation implementations in atomic.h for better code clarity.
- Adjusted function signatures in tilelang's atomic.py to enhance readability by aligning parameters.
- Cleaned up unnecessary whitespace and comments in customize.py to streamline the codebase.

* Add thread storage synchronization configuration option

- Introduced a new configuration option `tl.disable_thread_storage_sync` to control the automatic insertion of thread synchronization barriers in shared memory access.
- Updated the `ThreadSync` pass to check this configuration and bypass synchronization if disabled.
- Enhanced documentation in `builtin.h` and `pass_config.py` to clarify the purpose and usage of the new option.

* Refactor thread storage sync configuration retrieval

- Simplified the retrieval of the thread storage sync configuration in the `ThreadSync` pass by removing unnecessary intermediate variables.
- Ensured that the inclusion of `builtin.h` is consistent by moving it to the appropriate location in the file.

* test fix

* Update atomic operations and tests for improved functionality

- Updated atomic operations in CUDA templates to remove unnecessary address_of calls, enhancing performance and readability.
- Refactored atomic operation signatures in tilelang's atomic.py to accept references instead of pointers.
- Added new atomic operations and corresponding test cases for atomic add, max, min, and load/store functionalities in the testing suite.
- Updated the TVM subproject to the latest commit for better compatibility.

* Update attention sink examples to use 32 heads

- Modified the `heads` parameter in both `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py` from 1 to 32 to enhance performance in attention mechanisms.
- Ensured consistency across example scripts for improved usability and testing.

* Refactor atomic add handling in vectorization

- Simplified the extraction of buffer loads for atomic add operations by removing unnecessary address_of calls, improving code clarity and performance.
- Updated the data type retrieval for vectorization size calculation to directly access the buffer load node, enhancing efficiency.

* Add loop break functionality and enhance thread synchronization

- Introduced a new `loop_break` function in `customize.py` to allow breaking out of loops, returning a call to the `tl.loop_break` intrinsic.
- Updated the `sync_threads` function in `builtin.py` to accept optional parameters for `barrier_id` and `arrive_count`, improving its flexibility for thread synchronization.
- Added necessary imports in `__init__.py` to include the new `loop_break` function for broader accessibility.

* test fix

Author: LeiWang1999

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

@@ -366,9 +366,9 @@ def gen_inputs(B, H, Sq, Skv, D,
 
 def main(
     batch: int = 1,
-    heads: int = 64,
-    seq_q: int = 4096,
-    seq_kv: int = 4096,
+    heads: int = 32,
+    seq_q: int = 256,
+    seq_kv: int = 256,
     dim: int = 128,
     groups: int = 8,
     window_size: int | None = None,

examples/attention_sink/example_mha_sink_fwd_bhsd.py

@@ -229,10 +229,10 @@ def gen_inputs(B, H, Sq, Skv, D) -> tuple[torch.Tensor, torch.Tensor, torch.Tens
     return query, key, value, sinks
 
 
-def main(batch: int = 8,
-         heads: int = 32,
-         seq_q: int = 4096,
-         seq_kv: int = 4096,
+def main(batch: int = 1,
+         heads: int = 1,
+         seq_q: int = 256,
+         seq_kv: int = 256,
          dim: int = 128,
          window_size: int | None = None,
          tune: bool = False):

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py

@@ -354,10 +354,10 @@ def gen_inputs(B, H, Sq, Skv, D) -> tuple[torch.Tensor, torch.Tensor, torch.Tens
     return query, key, value, sinks
 
 
-def main(batch: int = 8,
+def main(batch: int = 1,
          heads: int = 32,
-         seq_q: int = 4096,
-         seq_kv: int = 4096,
+         seq_q: int = 256,
+         seq_kv: int = 256,
          dim: int = 128,
          window_size: int | None = None,
          tune: bool = False):

src/op/atomic_add.cc

@@ -293,10 +293,7 @@ For AtomicAddNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
   if (dst_predicate.defined())
     dst_value = if_then_else(dst_predicate, dst_value, make_zero(dst->dtype));
 
-  Call address_of_value =
-      tvm::tir::Call(DataType::Handle(), builtin::address_of(), {dst_value});
-
-  new_args.push_back(address_of_value);
+  new_args.push_back(dst_value);
   new_args.push_back(src_value);
 
   Call atomicadd_call =

src/op/builtin.cc

@@ -20,6 +20,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kDebugMergeSharedMemoryAllocations, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableTMALower, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableSafeMemoryLegalize, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableWarpSpecialized, Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION(kDisableThreadStorageSync, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kConfigIndexBitwidth, Integer);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableDynamicTailSplit, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDynamicAlignment, Integer);

src/op/builtin.h

@@ -55,6 +55,20 @@ static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
 static constexpr const char *kDisableDynamicTailSplit =
     "tl.disable_dynamic_tail_split";
 
+/*!
+ * \brief Whether to disable thread storage synchronization
+ *
+ * When enabled, disables the automatic insertion of thread synchronization
+ * barriers (e.g., __syncthreads()) for shared memory access coordination.
+ * This can be useful for performance optimization in cases where manual
+ * synchronization is preferred or when synchronization is not needed.
+ *
+ * kDisableThreadStorageSync = "tl.disable_thread_storage_sync"
+ *
+ */
+static constexpr const char *kDisableThreadStorageSync =
+    "tl.disable_thread_storage_sync";
+
 /*!
  * \brief The size of the vectorized dimension in buffer, designed by user
  *

src/tl_templates/cuda/atomic.h

@@ -0,0 +1,189 @@
+#pragma once
+
+#ifndef __CUDACC_RTC__
+#include <cuda_runtime.h>
+#endif
+
+#include <cuda/atomic>
+#include <cutlass/numeric_types.h>
+
+using cutlass::bfloat16_t;
+using cutlass::half_t;
+
+#define TL_DEVICE __forceinline__ __device__
+
+template <typename T> struct normalize_atomic_type {
+  using type = T;
+};
+
+template <> struct normalize_atomic_type<half_t> {
+  using type = half;
+};
+
+#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 750))
+template <> struct normalize_atomic_type<bfloat16_t> {
+  using type = __nv_bfloat16;
+};
+#endif
+
+template <typename T1, typename T2> TL_DEVICE T1 cuda_cast(T2 val) {
+  return T1(val);
+}
+
+template <> TL_DEVICE half cuda_cast<half, float>(float val) {
+  return __float2half(val);
+}
+
+#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 750))
+template <> TL_DEVICE __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) {
+  return __float2bfloat16(val);
+}
+#endif
+
+template <typename T1, typename T2>
+TL_DEVICE void AtomicMax(T1 &ref, T2 val,
+                         int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  T1 *address = &ref;
+  if constexpr (std::is_same_v<NT1, half> ||
+                std::is_same_v<NT1, __nv_bfloat16>) {
+    atomicMax(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val));
+  } else {
+    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
+    aref.fetch_max(cuda_cast<NT1>(val), cuda::memory_order(memory_order));
+  }
+}
+
+template <typename T1, typename T2>
+TL_DEVICE T1 AtomicMaxRet(T1 &ref, T2 val,
+                          int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  T1 *address = &ref;
+  if constexpr (std::is_same_v<NT1, half> ||
+                std::is_same_v<NT1, __nv_bfloat16>) {
+    return static_cast<T1>(
+        atomicMax(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val)));
+  } else {
+    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
+    return static_cast<T1>(
+        aref.fetch_max(cuda_cast<NT1>(val), cuda::memory_order(memory_order)));
+  }
+}
+
+template <typename T1, typename T2>
+TL_DEVICE void AtomicMin(T1 &ref, T2 val,
+                         int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  T1 *address = &ref;
+  if constexpr (std::is_same_v<NT1, half> ||
+                std::is_same_v<NT1, __nv_bfloat16>) {
+    atomicMin(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val));
+  } else {
+    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
+    aref.fetch_min(cuda_cast<NT1>(val), cuda::memory_order(memory_order));
+  }
+}
+
+template <typename T1, typename T2>
+TL_DEVICE T1 AtomicMinRet(T1 &ref, T2 val,
+                          int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  T1 *address = &ref;
+  if constexpr (std::is_same_v<NT1, half> ||
+                std::is_same_v<NT1, __nv_bfloat16>) {
+    return static_cast<T1>(
+        atomicMin(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val)));
+  } else {
+    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
+    return static_cast<T1>(
+        aref.fetch_min(cuda_cast<NT1>(val), cuda::memory_order(memory_order)));
+  }
+}
+
+template <typename T1, typename T2>
+TL_DEVICE void AtomicAdd(T1 &ref, T2 val,
+                         int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  T1 *address = &ref;
+  if constexpr (std::is_same_v<NT1, half> ||
+                std::is_same_v<NT1, __nv_bfloat16>) {
+    atomicAdd(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val));
+  } else {
+    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
+    aref.fetch_add(cuda_cast<NT1>(val), cuda::memory_order(memory_order));
+  }
+}
+
+template <typename T1, typename T2>
+TL_DEVICE T1 AtomicAddRet(T1 &ref, T2 val,
+                          int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  T1 *address = &ref;
+  if constexpr (std::is_same_v<NT1, half> ||
+                std::is_same_v<NT1, __nv_bfloat16>) {
+    return static_cast<T1>(
+        atomicAdd(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val)));
+  } else {
+    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
+    return static_cast<T1>(
+        aref.fetch_add(cuda_cast<NT1>(val), cuda::memory_order(memory_order)));
+  }
+}
+
+TL_DEVICE void AtomicAddx2(half_t *ref, half_t *val) {
+  atomicAdd(reinterpret_cast<half2 *>(ref),
+            static_cast<half2>(*reinterpret_cast<half2 *>(val)));
+}
+
+TL_DEVICE half2 AtomicAddx2Ret(half_t *ref, half_t *val) {
+  return atomicAdd(reinterpret_cast<half2 *>(ref),
+                   static_cast<half2>(*reinterpret_cast<half2 *>(val)));
+}
+
+#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 750))
+TL_DEVICE void AtomicAddx2(bfloat16_t *ref, bfloat16_t *val) {
+  atomicAdd(
+      reinterpret_cast<__nv_bfloat162 *>(ref),
+      static_cast<__nv_bfloat162>(*reinterpret_cast<__nv_bfloat162 *>(val)));
+}
+
+TL_DEVICE __nv_bfloat162 AtomicAddx2Ret(bfloat16_t *ref, bfloat16_t *val) {
+  return atomicAdd(
+      reinterpret_cast<__nv_bfloat162 *>(ref),
+      static_cast<__nv_bfloat162>(*reinterpret_cast<__nv_bfloat162 *>(val)));
+}
+#endif
+
+#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ >= 900))
+TL_DEVICE void AtomicAddx2(float *ref, float *val) {
+  atomicAdd(reinterpret_cast<float2 *>(ref),
+            static_cast<float2>(*reinterpret_cast<float2 *>(val)));
+}
+
+TL_DEVICE float2 AtomicAddx2Ret(float *ref, float *val) {
+  return atomicAdd(reinterpret_cast<float2 *>(ref),
+                   static_cast<float2>(*reinterpret_cast<float2 *>(val)));
+}
+
+TL_DEVICE void AtomicAddx4(float *ref, float *val) {
+  atomicAdd(reinterpret_cast<float4 *>(ref),
+            static_cast<float4>(*reinterpret_cast<float4 *>(val)));
+}
+
+TL_DEVICE float4 AtomicAddx4Ret(float *ref, float *val) {
+  return atomicAdd(reinterpret_cast<float4 *>(ref),
+                   static_cast<float4>(*reinterpret_cast<float4 *>(val)));
+}
+#endif
+
+template <typename T> TL_DEVICE T AtomicLoad(T &ref, int memory_order) {
+  cuda::atomic_ref<T, cuda::thread_scope_device> aref(ref);
+  return aref.load(cuda::memory_order(memory_order));
+}
+
+template <typename T1, typename T2>
+TL_DEVICE void AtomicStore(T1 &ref, T2 value, int memory_order) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(ref);
+  aref.store(cuda_cast<NT1>(value), cuda::memory_order(memory_order));
+}