[ir] Deprecate FrontendAtomicStmt

taichi/inc/statements.inc.h

@@ -7,7 +7,6 @@ PER_STATEMENT(FrontendBreakStmt)
 PER_STATEMENT(FrontendContinueStmt)
 PER_STATEMENT(FrontendAllocaStmt)
 PER_STATEMENT(FrontendAssignStmt)
-PER_STATEMENT(FrontendAtomicStmt)
 PER_STATEMENT(FrontendEvalStmt)
 PER_STATEMENT(FrontendSNodeOpStmt)  // activate, deactivate, append, clear
 PER_STATEMENT(FrontendAssertStmt)

taichi/ir/expr.cpp

@@ -127,17 +127,17 @@ Expr Expr::eval() const {
 
 void Expr::operator+=(const Expr &o) {
   if (this->atomic) {
-    current_ast_builder().insert(Stmt::make<FrontendAtomicStmt>(
-        AtomicOpType::add, ptr_if_global(*this), load_if_ptr(o)));
+    (*this) = Expr::make<AtomicOpExpression>(
+        AtomicOpType::add, ptr_if_global(*this), load_if_ptr(o));
   } else {
     (*this) = (*this) + o;
   }
 }
 
 void Expr::operator-=(const Expr &o) {
   if (this->atomic) {
-    current_ast_builder().insert(Stmt::make<FrontendAtomicStmt>(
-        AtomicOpType::add, *this, -load_if_ptr(o)));
+    (*this) = Expr::make<AtomicOpExpression>(
+        AtomicOpType::sub, ptr_if_global(*this), load_if_ptr(o));
   } else {
     (*this) = (*this) - o;
   }

taichi/ir/ir.cpp

@@ -553,12 +553,6 @@ FrontendAssignStmt::FrontendAssignStmt(const Expr &lhs, const Expr &rhs)
   TI_ASSERT(lhs->is_lvalue());
 }
 
-FrontendAtomicStmt::FrontendAtomicStmt(AtomicOpType op_type,
-                                       const Expr &dest,
-                                       const Expr &val)
-    : op_type(op_type), dest(dest), val(val) {
-}
-
 IRNode *FrontendContext::root() {
   return static_cast<IRNode *>(root_node.get());
 }
@@ -845,6 +839,28 @@ std::string AtomicOpExpression::serialize() {
   }
 }
 
+void AtomicOpExpression::flatten(FlattenContext *ctx) {
+  // replace atomic sub with negative atomic add
+  if (op_type == AtomicOpType::sub) {
+    val.set(Expr::make<UnaryOpExpression>(UnaryOpType::neg, val));
+    op_type = AtomicOpType::add;
+  }
+  // expand rhs
+  auto expr = val;
+  expr->flatten(ctx);
+  if (dest.is<IdExpression>()) {  // local variable
+    // emit local store stmt
+    auto alloca = ctx->current_block->lookup_var(dest.cast<IdExpression>()->id);
+    ctx->push_back<AtomicOpStmt>(op_type, alloca, expr->stmt);
+  } else {  // global variable
+    TI_ASSERT(dest.is<GlobalPtrExpression>());
+    auto global_ptr = dest.cast<GlobalPtrExpression>();
+    global_ptr->flatten(ctx);
+    ctx->push_back<AtomicOpStmt>(op_type, ctx->back_stmt(), expr->stmt);
+  }
+  stmt = ctx->back_stmt();
+}
+
 std::string SNodeOpExpression::serialize() {
   if (value.expr) {
     return fmt::format("{}({}, [{}], {})", snode_op_type_name(op_type),

taichi/ir/ir.h

@@ -1336,16 +1336,6 @@ class Block : public IRNode {
   DEFINE_ACCEPT
 };
 
-class FrontendAtomicStmt : public Stmt {
- public:
-  AtomicOpType op_type;
-  Expr dest, val;
-
-  FrontendAtomicStmt(AtomicOpType op_type, const Expr &dest, const Expr &val);
-
-  DEFINE_ACCEPT
-};
-
 class FrontendSNodeOpStmt : public Stmt {
  public:
   SNodeOpType op_type;
@@ -1944,11 +1934,8 @@ class IdExpression : public Expression {
   }
 };
 
-// This is just a wrapper class of FrontendAtomicStmt, so that we can turn
-// ti.atomic_op() into an expression (with side effect).
+// ti.atomic_*() is an expression with side effect.
 class AtomicOpExpression : public Expression {
-  // TODO(issue#332): Flatten this into AtomicOpStmt directly, then we can
-  // deprecate FrontendAtomicStmt.
  public:
   AtomicOpType op_type;
   Expr dest, val;
@@ -1959,13 +1946,7 @@ class AtomicOpExpression : public Expression {
 
   std::string serialize() override;
 
-  void flatten(FlattenContext *ctx) override {
-    // FrontendAtomicStmt is the correct place to flatten sub-exprs like |dest|
-    // and |val| (See LowerAST). This class only wraps the frontend atomic_op()
-    // stmt as an expression.
-    ctx->push_back<FrontendAtomicStmt>(op_type, dest, val);
-    stmt = ctx->back_stmt();
-  }
+  void flatten(FlattenContext *ctx) override;
 };
 
 class SNodeOpExpression : public Expression {

taichi/transforms/ir_printer.cpp

@@ -150,12 +150,6 @@ class IRPrinter : public IRVisitor {
           stmt->op2->name(), stmt->op3->name());
   }
 
-  void visit(FrontendAtomicStmt *stmt) override {
-    print("{}{} = atomic {}({}, {})", stmt->type_hint(), stmt->name(),
-          atomic_op_type_name(stmt->op_type), stmt->dest->serialize(),
-          stmt->val->serialize());
-  }
-
   void visit(AtomicOpStmt *stmt) override {
     print("{}{} = atomic {}({}, {})", stmt->type_hint(), stmt->name(),
           atomic_op_type_name(stmt->op_type), stmt->dest->name(),

taichi/transforms/lower_ast.cpp

@@ -314,31 +314,6 @@ class LowerAST : public IRVisitor {
     throw IRModified();
   }
 
-  void visit(FrontendAtomicStmt *stmt) override {
-    // replace atomic sub with negative atomic add
-    if (stmt->op_type == AtomicOpType::sub) {
-      stmt->val.set(Expr::make<UnaryOpExpression>(UnaryOpType::neg, stmt->val));
-      stmt->op_type = AtomicOpType::add;
-    }
-    // expand rhs
-    auto expr = stmt->val;
-    auto fctx = make_flatten_ctx();
-    expr->flatten(&fctx);
-    if (stmt->dest.is<IdExpression>()) {  // local variable
-      // emit local store stmt
-      auto alloca =
-          stmt->parent->lookup_var(stmt->dest.cast<IdExpression>()->id);
-      fctx.push_back<AtomicOpStmt>(stmt->op_type, alloca, expr->stmt);
-    } else {  // global variable
-      TI_ASSERT(stmt->dest.is<GlobalPtrExpression>());
-      auto global_ptr = stmt->dest.cast<GlobalPtrExpression>();
-      global_ptr->flatten(&fctx);
-      fctx.push_back<AtomicOpStmt>(stmt->op_type, fctx.back_stmt(), expr->stmt);
-    }
-    stmt->parent->replace_with(stmt, std::move(fctx.stmts));
-    throw IRModified();
-  }
-
   void visit(FrontendSNodeOpStmt *stmt) override {
     // expand rhs
     Stmt *val_stmt = nullptr;