[Feature]: Add test for atomicadd auto vectorize and remove useless code (#1019)

* update

* format

* rabbit

src/op/atomic_add.cc

@@ -272,7 +272,6 @@ For AtomicAddNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
   PrimExpr dst_predicate = MakePredicate(analyzer, loop_vars, dst->shape, 1);
 
   Array<PrimExpr> new_args;
-  new_args.push_back(StringImm("AtomicAdd"));
 
   PrimExpr src_value = BufferLoad(src, src_indices);
   if (src->dtype != dst->dtype)
@@ -288,7 +287,7 @@ For AtomicAddNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
   new_args.push_back(src_value);
 
   Call atomicadd_call =
-      tvm::tir::Call(dst->dtype, builtin::call_extern(), new_args);
+      tvm::tir::Call(dst->dtype, atomicadd_elem_op(), new_args);
 
   Stmt body = tvm::tir::Evaluate(atomicadd_call);
 
@@ -325,10 +324,6 @@ For AtomicAddNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
  */
 LayoutMap AtomicAddNode::InferLayout(const LayoutInferArgs &T,
                                      InferLevel level) const {
-  if (!par_op_.defined()) {
-    arith::Analyzer analyzer;
-    par_op_ = ParallelOp(MakeSIMTLoop(&analyzer));
-  }
   if (T.layout_map.count(src) && T.layout_map.count(dst)) {
     if (src.scope() == "local.fragment" && dst.scope() == "local.fragment") {
       const FragmentNode *src_layout = T.layout_map[src].as<FragmentNode>();
@@ -342,7 +337,7 @@ LayoutMap AtomicAddNode::InferLayout(const LayoutInferArgs &T,
       }
     }
   }
-  return par_op_->InferLayout(T, level);
+  return {};
 }
 
 /**

src/op/builtin.cc

@@ -295,5 +295,10 @@ TIR_DEFINE_TL_BUILTIN(increase_descriptor_offset)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kOpaque));
 
+TIR_DEFINE_TL_BUILTIN(atomicadd_elem_op)
+    .set_num_inputs(2)
+    .set_attr<TCallEffectKind>("TCallEffectKind",
+                               Integer(CallEffectKind::kOpaque));
+
 } // namespace tl
 } // namespace tvm

src/op/builtin.h

@@ -501,6 +501,13 @@ TVM_DLL const Op &initialize_descriptor();
  * tilelang.
  */
 TVM_DLL const Op &increase_descriptor_offset();
+/*!
+ * \brief tilelang intrinsic for element-wise atomic addition.
+ *
+ *  This op is used to represent an element-wise atomic add operation in
+ * tilelang.
+ */
+TVM_DLL const Op &atomicadd_elem_op();
 
 } // namespace tl
 } // namespace tvm

src/transform/atomicadd_vectorize.cc

@@ -23,38 +23,39 @@ AtomicAddVectorizePlanner::Plan(const For &node, int compute_capability) {
 
   PostOrderVisit(node, [&](const ObjectRef &obj) {
     if (const auto *call = obj.as<CallNode>()) {
-      if (call->op == builtin::call_extern() && call->args.size() >= 2) {
-        const auto *func_name = call->args[0].as<StringImmNode>();
-        if (!func_name)
+      if (call->op == atomicadd_elem_op()) {
+        if (call->args.size() < 2) {
+          // Fallback: unexpected arity
+          vectorize_size_max = 1;
+          DLOG(WARNING) << "[AtomicAddVectorizePlanner] atomicadd_elem_op "
+                           "expects 2 args, got "
+                        << call->args.size() << "; Fallback to no vectorize";
           return;
-        if (func_name->value == "AtomicAdd") {
-          DataType dtype;
-          if (const auto *load = call->args[1].as<BufferLoadNode>()) {
-            dtype = load->dtype;
+        }
+        DataType dtype;
+        if (const auto *load = call->args[0].as<BufferLoadNode>()) {
+          dtype = load->dtype;
+          vectorize_size_max = GetVectorizeSizeMax(compute_capability, dtype);
+        } else if (const auto *ite = call->args[0].as<IfThenElseNode>()) {
+          if (const auto *then_load = ite->then_case.as<BufferLoadNode>()) {
+            dtype = then_load->dtype;
+            vectorize_size_max = GetVectorizeSizeMax(compute_capability, dtype);
+          } else if (const auto *else_load =
+                         ite->else_case.as<BufferLoadNode>()) {
+            dtype = else_load->dtype;
             vectorize_size_max = GetVectorizeSizeMax(compute_capability, dtype);
-          } else if (const auto *ite = call->args[1].as<IfThenElseNode>()) {
-            if (const auto *then_load = ite->then_case.as<BufferLoadNode>()) {
-              dtype = then_load->dtype;
-              vectorize_size_max =
-                  GetVectorizeSizeMax(compute_capability, dtype);
-            } else if (const auto *else_load =
-                           ite->else_case.as<BufferLoadNode>()) {
-              dtype = else_load->dtype;
-              vectorize_size_max =
-                  GetVectorizeSizeMax(compute_capability, dtype);
-            } else {
-              // fallback
-              vectorize_size_max = 1;
-              DLOG(WARNING) << "[AtomicAddVectorizePlanner] IfThenElse case "
-                               "has no BufferLoad; Fallback to no vectorize";
-            }
           } else {
             // fallback
             vectorize_size_max = 1;
-            DLOG(WARNING) << "[AtomicAddVectorizePlanner] Unexpected arg1 type "
-                          << call->args[1]->GetTypeKey()
-                          << "; Fallback to no vectorize";
+            DLOG(WARNING) << "[AtomicAddVectorizePlanner] IfThenElse case "
+                             "has no BufferLoad; Fallback to no vectorize";
           }
+        } else {
+          // fallback
+          vectorize_size_max = 1;
+          DLOG(WARNING) << "[AtomicAddVectorizePlanner] Unexpected arg1 type "
+                        << call->args[1]->GetTypeKey()
+                        << "; Fallback to no vectorize";
         }
       }
     }
@@ -75,22 +76,19 @@ void AtomicAddVectorizePlanner::VisitStmt_(const ForNode *node) {
 }
 
 void AtomicAddVectorizePlanner::VisitExpr_(const CallNode *node) {
-  if (node->op == builtin::call_extern() && node->args.size() >= 2) {
-    if (const auto *func_name = node->args[0].as<StringImmNode>()) {
-      if (func_name->value == "AtomicAdd") {
-        const BufferLoadNode *buffer_load_dst =
-            node->args[1].as<BufferLoadNode>();
-        const BufferLoadNode *buffer_load_src =
-            node->args[2].as<BufferLoadNode>();
-        if (buffer_load_src && buffer_load_src->buffer.defined() &&
-            buffer_load_dst && buffer_load_dst->buffer.defined()) {
-          Buffer dst_buffer = buffer_load_dst->buffer;
-          UpdateVectorSize(buffer_load_dst->indices, dst_buffer);
+  if (node->op == atomicadd_elem_op() && !node->args.empty()) {
+    if (node->args.size() < 2) {
+      return arith::IRVisitorWithAnalyzer::VisitExpr_(node);
+    }
+    const BufferLoadNode *buffer_load_dst = node->args[0].as<BufferLoadNode>();
+    const BufferLoadNode *buffer_load_src = node->args[1].as<BufferLoadNode>();
+    if (buffer_load_src && buffer_load_src->buffer.defined() &&
+        buffer_load_dst && buffer_load_dst->buffer.defined()) {
+      Buffer dst_buffer = buffer_load_dst->buffer;
+      UpdateVectorSize(buffer_load_dst->indices, dst_buffer);
 
-          Buffer src_buffer = buffer_load_src->buffer;
-          UpdateVectorSize(buffer_load_src->indices, src_buffer);
-        }
-      }
+      Buffer src_buffer = buffer_load_src->buffer;
+      UpdateVectorSize(buffer_load_src->indices, src_buffer);
     }
   }
   return arith::IRVisitorWithAnalyzer::VisitExpr_(node);
@@ -188,6 +186,8 @@ private:
   Stmt VisitStmt_(const ForNode *node) final {
     inner_for_ = node;
     auto ret = StmtExprMutator::VisitStmt_(node);
+    if (vector_size_ == 1)
+      return ret;
     if (inner_for_ == node) {
       For fnode = ret.as<For>().value();
       auto old_var = fnode->loop_var;
@@ -210,47 +210,54 @@ private:
   }
 
   PrimExpr VisitExpr_(const CallNode *node) final {
-    if (dynamic_) {
-      return StmtExprMutator::VisitExpr_(node);
+    bool legal_vectorize = true;
+    if (dynamic_)
+      legal_vectorize = false;
+    if (!(node->op == atomicadd_elem_op()))
+      legal_vectorize = false;
+    if (node->args.size() < 2)
+      legal_vectorize = false;
+    if (legal_vectorize) {
+      const BufferLoadNode *temp_dst_node = node->args[0].as<BufferLoadNode>();
+      const BufferLoadNode *temp_value_node =
+          node->args[1].as<BufferLoadNode>();
+      if (!temp_dst_node || !temp_value_node)
+        legal_vectorize = false;
     }
-    if (vector_size_ == 2 || vector_size_ == 4) {
-      if (node->op == builtin::call_extern() && node->args.size() >= 2) {
-        if (const auto *func_name = node->args[0].as<StringImmNode>()) {
-          if (func_name->value == "AtomicAdd") {
-            const BufferLoadNode *temp_dst_node =
-                node->args[1].as<BufferLoadNode>();
-            const BufferLoadNode *temp_value_node =
-                node->args[2].as<BufferLoadNode>();
-            if (!temp_dst_node || !temp_value_node) {
-              return StmtExprMutator::VisitExpr_(node);
-            }
-            const BufferLoad dst_node =
-                Downcast<BufferLoad>(node->args[1].as<BufferLoadNode>());
-            const BufferLoad value_node =
-                Downcast<BufferLoad>(node->args[2].as<BufferLoadNode>());
+    if (legal_vectorize) {
+      const BufferLoad dst_node = Downcast<BufferLoad>(node->args[0]);
+      const BufferLoad value_node = Downcast<BufferLoad>(node->args[1]);
 
-            Call address_of_dst =
-                Call(DataType::Handle(), builtin::address_of(), {dst_node});
-            Call address_of_value =
-                Call(DataType::Handle(), builtin::address_of(), {value_node});
-            Array<PrimExpr> new_args;
-            if (vector_size_ == 2) {
-              new_args.push_back(StringImm("AtomicAddx2"));
-            } else {
-              new_args.push_back(StringImm("AtomicAddx4"));
-            }
-            new_args.push_back(address_of_dst);
-            new_args.push_back(address_of_value);
+      Call address_of_dst =
+          Call(DataType::Handle(), builtin::address_of(), {dst_node});
+      Call address_of_value =
+          Call(DataType::Handle(), builtin::address_of(), {value_node});
+      Array<PrimExpr> new_args;
+      if (vector_size_ == 4) {
+        new_args.push_back(StringImm("AtomicAddx4"));
+      } else if (vector_size_ == 2) {
+        new_args.push_back(StringImm("AtomicAddx2"));
+      } else {
+        new_args.push_back(StringImm("AtomicAdd"));
+      }
+      new_args.push_back(address_of_dst);
+      new_args.push_back(address_of_value);
 
-            Call new_call =
-                tvm::tir::Call(node->dtype, builtin::call_extern(), new_args);
+      Call new_call =
+          tvm::tir::Call(node->dtype, builtin::call_extern(), new_args);
 
-            return new_call;
-          }
-        }
-      }
+      return new_call;
+    } else {
+      Array<PrimExpr> new_args;
+      new_args.push_back(StringImm("AtomicAdd"));
+      for (auto x : node->args)
+        new_args.push_back(x);
+
+      Call new_call =
+          tvm::tir::Call(node->dtype, builtin::call_extern(), new_args);
+
+      return new_call;
     }
-    return StmtExprMutator::VisitExpr_(node);
   }
 
   const ForNode *inner_for_;
@@ -263,9 +270,6 @@ For VectorizeAtomicAdd(const For &for_node, int compute_capability) {
   AtomicAddVectorizePlanResult res = {1, false, 0};
   AtomicAddVectorizePlanner planner;
   res = planner.Plan(for_node, compute_capability);
-  int vectorize_hint = res.vector_size;
-  if (vectorize_hint == 1)
-    return for_node;
   auto rewriter = AtomicAddVectorizeRewriter(res);
   return Downcast<For>(rewriter(for_node));
 }

src/transform/atomicadd_vectorize.h

@@ -8,6 +8,7 @@
 
 #include "../layout/layout.h"
 #include "../layout/utils.h"
+#include "../op/builtin.h"
 #include "arith/int_operator.h"
 #include "arith/ir_visitor_with_analyzer.h"
 #include "atomicadd_vectorize.h"