Merge remote-tracking branch 'upstream/main' into retire-format-sh

Author: XuehaiPan

.github/workflows/ci.yml

@@ -226,7 +226,6 @@ jobs:
             echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" | tee -a "${GITHUB_ENV}"
             echo "UV_INDEX=${UV_INDEX}" | tee -a "${GITHUB_ENV}"
           fi
-
           export CLANG_TIDY_CMAKE_OPTIONS="${CLANG_TIDY_CMAKE_OPTIONS} -DUSE_METAL=ON"
 
           echo "USE_METAL=ON" | tee -a "${GITHUB_ENV}"

examples/flash_attention/example_gqa_bwd_tma_reduce_varlen.py

@@ -8,8 +8,6 @@
 from tilelang.profiler import do_bench
 from varlen_utils import generate_random_padding_mask, generate_qkv
 
-tilelang.disable_cache()
-
 
 def attention_ref(
         q,

examples/flash_attention/example_gqa_fwd_varlen.py

@@ -58,8 +58,12 @@ AtomicAdd::AtomicAdd(Array<PrimExpr> args, BufferMap vmap) {
   if (args.size() >= 3) {
     node->use_tma = Downcast<IntImm>(args[2]);
   }
+  node->memory_order = IntImm(0);
   if (args.size() >= 4) {
-    node->coalesced_width = Downcast<IntImm>(args[3]);
+    node->memory_order = Downcast<IntImm>(args[3]);
+  }
+  if (args.size() >= 5) {
+    node->coalesced_width = Downcast<IntImm>(args[4]);
   }
   data_ = std::move(node);
 }
@@ -285,6 +289,7 @@ For AtomicAddNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
 
   new_args.push_back(dst_value);
   new_args.push_back(src_value);
+  new_args.push_back(memory_order);
 
   Call atomicadd_call =
       tvm::tir::Call(dst->dtype, atomicadd_elem_op(), new_args);

src/op/atomic_add.cc

@@ -22,6 +22,7 @@ class AtomicAddNode : public TileOperatorNode {
       dst_range;          ///< Access ranges for source and destination
   IntImm use_tma;         ///< Whether to use TMA for memory operations
   IntImm coalesced_width; ///< Width for memory coalescing optimization
+  IntImm memory_order;    ///< Memory order for atomic operations
 
   mutable ParallelOp par_op_; ///< Associated parallel operation
   static constexpr const char *_type_key = "tl.AtomicAdd";
@@ -41,15 +42,17 @@ class AtomicAddNode : public TileOperatorNode {
         .def_ro("src_range", &AtomicAddNode::src_range)
         .def_ro("dst_range", &AtomicAddNode::dst_range)
         .def_ro("use_tma", &AtomicAddNode::use_tma)
-        .def_ro("coalesced_width", &AtomicAddNode::coalesced_width);
+        .def_ro("coalesced_width", &AtomicAddNode::coalesced_width)
+        .def_ro("memory_order", &AtomicAddNode::memory_order);
   }
 
   bool SEqualReduce(const AtomicAddNode *other, SEqualReducer equal) const {
     return equal(src, other->src) && equal(dst, other->dst) &&
            equal(src_range, other->src_range) &&
            equal(dst_range, other->dst_range) &&
            equal(use_tma, other->use_tma) &&
-           equal(coalesced_width, other->coalesced_width);
+           equal(coalesced_width, other->coalesced_width) &&
+           equal(memory_order, other->memory_order);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
@@ -59,6 +62,7 @@ class AtomicAddNode : public TileOperatorNode {
     hash_reduce(dst_range);
     hash_reduce(use_tma);
     hash_reduce(coalesced_width);
+    hash_reduce(memory_order);
   }
 
   static constexpr bool _type_has_method_sequal_reduce = true;

src/op/atomic_add.h

@@ -296,7 +296,7 @@ TIR_DEFINE_TL_BUILTIN(increase_descriptor_offset)
                                Integer(CallEffectKind::kOpaque));
 
 TIR_DEFINE_TL_BUILTIN(atomicadd_elem_op)
-    .set_num_inputs(2)
+    .set_num_inputs(3)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kOpaque));