/*!
 * \file eliminate_storage_sync_for_mbarrier.cc
 */
#include "./storage_access.h"
#include "arith/ir_mutator_with_analyzer.h"
#include "arith/ir_visitor_with_analyzer.h"
#include <tvm/runtime/registry.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/transform.h>

namespace tvm {
namespace tl {

using namespace tir;
using arith::IRMutatorWithAnalyzer;
using arith::IRVisitorWithAnalyzer;

class Eliminator : public IRMutatorWithAnalyzer {
public:
  static Stmt Substitute(Stmt stmt, bool skip_thread_partition = false) {
    arith::Analyzer analyzer;
    Eliminator transformer(&analyzer);
    return transformer.VisitStmt(stmt);
  }

  Eliminator(arith::Analyzer *analyzer) : IRMutatorWithAnalyzer(analyzer) {
    in_mbarrier_region_ = false;
  }

  Stmt VisitStmt_(const AttrStmtNode *op) final {
    if (op->attr_key == "thread_extent") {
      const VarNode *var = nullptr;
      if (op->node->IsInstance<VarNode>()) {
        var = static_cast<const VarNode *>(op->node.get());
        if (var->name_hint == "threadIdx.x") {
          thread_extent_ = op;
        }
      }
    }
    return IRMutatorWithAnalyzer::VisitStmt_(op);
  }

  Stmt VisitStmt_(const EvaluateNode *op) final {
    const CallNode *call = nullptr;
    if (op->value->IsInstance<CallNode>()) {
      call = static_cast<const CallNode *>(op->value.get());
      if (call->op.same_as(builtin::tvm_storage_sync())) {
        // Skip storage sync if we're in a region with mbarrier operations
        if (in_mbarrier_region_) {
          return Stmt();
        }
      } else if (call->op.same_as(builtin::ptx_arrive_barrier()) ||
                 call->op.same_as(builtin::ptx_wait_barrier())) {
        in_mbarrier_region_ = true;
      }
    }
    return IRMutatorWithAnalyzer::VisitStmt_(op);
  }

  Stmt VisitStmt_(const IfThenElseNode *op) final {
    bool old_in_mbarrier = in_mbarrier_region_;
    Stmt then_case = VisitStmt(op->then_case);

    Stmt ret;
    if (op->else_case.defined()) {
      in_mbarrier_region_ = old_in_mbarrier;
      Stmt else_case = VisitStmt(op->else_case.value());
      in_mbarrier_region_ = old_in_mbarrier || in_mbarrier_region_;
      ret = IfThenElse(VisitExpr(op->condition), then_case, else_case);
    } else {
      in_mbarrier_region_ = old_in_mbarrier || in_mbarrier_region_;
      ret = IfThenElse(VisitExpr(op->condition), then_case, Stmt());
    }
    return ret;
  }

private:
  bool in_mbarrier_region_;
  const AttrStmtNode *thread_extent_{nullptr};
};
using namespace tir::transform;

namespace transform {

tvm::transform::Pass EliminateStorageSyncForMBarrier() {
  auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
    auto *n = f.CopyOnWrite();
    n->body = Eliminator::Substitute(std::move(n->body));
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.EliminateStorageSyncForMBarrier",
                            {});
}

TVM_REGISTER_GLOBAL("tl.transform.EliminateStorageSyncForMBarrier")
    .set_body_typed(EliminateStorageSyncForMBarrier);

} // namespace transform
} // namespace tl
} // namespace tvm
