/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file remove_no_op.cc
 * \brief Remove no op from the stmt
 */
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>

#include <unordered_map>

#include "../../arith/const_fold.h"
#include "../../arith/ir_mutator_with_analyzer.h"
#include "../analysis/var_use_def_analysis.h"
#include "ir_utils.h"

namespace tvm {
namespace tirx {

struct RemoveNoOpConfigNode : public ffi::Object {
  int64_t max_simplification_steps;
  bool ignore_profiler_call;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RemoveNoOpConfigNode>()
        .def_ro("max_simplification_steps", &RemoveNoOpConfigNode::max_simplification_steps,
                "If non-zero, RewriteSimplifier will throw an error "
                "after the number of steps specified.  "
                "For use in debug and testing purposes.",
                refl::DefaultValue(0))
        .def_ro("ignore_profiler_call", &RemoveNoOpConfigNode::ignore_profiler_call,
                "If true, profiler calls are rendered as no-ops.", refl::DefaultValue(false));
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.RemoveNoOpConfig", RemoveNoOpConfigNode,
                                    ffi::Object);
};

class RemoveNoOpConfig : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RemoveNoOpConfig, ffi::ObjectRef,
                                                RemoveNoOpConfigNode);
};

TVM_FFI_STATIC_INIT_BLOCK() { RemoveNoOpConfigNode::RegisterReflection(); }

TVM_REGISTER_PASS_CONFIG_OPTION("tirx.RemoveNoOp", RemoveNoOpConfig);

// Mark the statement of each stage.
class NoOpRemover : public arith::IRMutatorWithAnalyzer {
 public:
  static Stmt Apply(Stmt stmt, const arith::Analyzer& analyzer, bool ignore_profiler_call = false) {
    NoOpRemover visitor(analyzer, ignore_profiler_call);
    return visitor(std::move(stmt));
  }

 private:
  using Parent = IRMutatorWithAnalyzer;
  using Parent::VisitStmt;
  using Parent::VisitStmt_;

  NoOpRemover(const arith::Analyzer& analyzer, bool ignore_profiler_call = false)
      : Parent(analyzer), ignore_profiler_call_(ignore_profiler_call) {}

  Stmt VisitStmt_(const BindNode* op) final {
    // Simply mutate the value and return.
    // Unused Bind elimination can be done later via a separate two-pass approach.
    return Parent::VisitStmt_(op);
  }
  Stmt VisitStmt_(const AttrStmtNode* op) final {
    if (op->attr_key == "pragma_debug_skip_region") {
      return MakeEvaluate(0);
    } else if (op->attr_key == s_tir::attr::async_wait_queue_scope) {
      auto wait_attrs = GetAsyncWaitAttributes(op);
      auto wait_cnt = wait_attrs.second;
      arith::Analyzer ana;
      if (ana->CanProve(wait_cnt < 0)) {
        // A negative wait count can arise if it depends on a loop variable.
        // For example, a wait count 1 - i can be negative after loop unrolling.
        // We assume that such wait is a nop.
        auto inner = op->body.as<AttrStmtNode>();
        TVM_FFI_ICHECK(inner);
        return Parent::VisitStmt(inner->body);
      }
    }

    Stmt stmt = Parent::VisitStmt_(op);
    op = stmt.as<AttrStmtNode>();
    return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
  }
  Stmt VisitStmt_(const IfThenElseNode* op) final {
    Stmt stmt = Parent::VisitStmt_(op);
    op = stmt.as<IfThenElseNode>();
    // Sometimes the condition can be statically determined,
    // in which the type of the `stmt` will not be IfThenElseNode.
    if (!op) {
      return stmt;
    }
    if (op->else_case) {
      bool no_op_else = is_no_op(op->else_case.value());
      bool no_op_then = is_no_op(op->then_case);
      if (no_op_else && no_op_then) {
        return MakeEvaluate(op->condition);
      } else if (no_op_else) {
        return IfThenElse(op->condition, op->then_case);
      } else if (no_op_then) {
        return IfThenElse(!op->condition, op->else_case.value());
      } else {
        return stmt;
      }
    } else {
      if (is_no_op(op->then_case)) {
        return MakeEvaluate(op->condition);
      } else {
        return stmt;
      }
    }
  }
  Stmt VisitStmt_(const ForNode* op) final {
    auto extent_range = arith::EvalSet(op->extent, var_range_map_);
    if (!arith::is_neg_inf(extent_range.max()) && !arith::is_pos_inf(extent_range.max()) &&
        analyzer_->CanProve(extent_range.max() <= 0)) {
      return Evaluate(0);
    }
    var_range_map_[op->loop_var.get()] = arith::IntSet::FromMinExtent(op->min, op->extent);
    Stmt stmt = Parent::VisitStmt_(op);
    var_range_map_.erase(op->loop_var.get());
    op = stmt.as<ForNode>();
    if (is_zero(op->extent)) {
      return Evaluate(0);
    }
    return is_no_op(op->body) ? MakeEvaluate({op->min, op->extent}) : stmt;
  }

  Stmt VisitStmt_(const AllocBufferNode* op) final { return StmtMutator::VisitStmt_(op); }

  Stmt VisitStmt_(const EvaluateNode* op) final {
    if (HasSideEffect(op->value)) {
      return ffi::GetRef<Stmt>(op);
    } else {
      return Evaluate(0);
    }
  }

  bool HasSideEffect(const Expr& value) {
    if (auto prim = value.as<PrimExpr>()) return HasSideEffect(prim.value());
    // Variables and string literals are pure.  Preserve non-primitive calls
    // conservatively because their effect metadata is independent of the
    // call's semantic return type.
    return value.as<CallNode>() != nullptr;
  }

  Stmt VisitStmt_(const BufferStoreNode* op) final {
    BufferStore store = ffi::GetRef<BufferStore>(op);

    // Helper function that returns a statement containing only the
    // side effects of evaluating this BufferStore, but not the store
    // itself.
    auto only_side_effects = [&]() {
      ffi::Array<Stmt> statements;
      statements.push_back(MakeEvaluate(store->value));
      for (const auto& index : store->indices) {
        statements.push_back(MakeEvaluate(index));
      }
      return this->VisitStmt(SeqStmt(statements));
    };

    // A write whose destination is known to already contain the
    // values to be written is a no-op.
    PrimExpr stores_existing_value =
        store->value - BufferLoad(store->buffer, store->indices, store->predicate) == 0;
    stores_existing_value = analyzer_->Simplify(stores_existing_value);
    if (is_one(stores_existing_value)) {
      return only_side_effects();
    }

    // If the stored value is a load from the same location, the
    // statement is a no-op, regardless of contextual information.
    if (const BufferLoadNode* load = store->value.as<BufferLoadNode>()) {
      if (load->buffer.same_as(store->buffer) &&
          analyzer_->CanProveEqual(load->buffer->elem_offset, store->buffer->elem_offset) &&
          ArrayValueEqual(load->buffer->shape, store->buffer->shape) &&
          ArrayValueEqual(load->buffer->strides, store->buffer->strides) &&
          ArrayValueEqual(load->indices, store->indices)) {
        return only_side_effects();
      }
    }

    return store;
  }

  Stmt VisitStmt_(const DeclBufferNode* op) final { return StmtMutator::VisitStmt_(op); }

 private:
  bool ArrayValueEqual(const ffi::Array<PrimExpr>& a, const ffi::Array<PrimExpr>& b) {
    if (a.size() != b.size()) {
      return false;
    }
    for (size_t i = 0; i < a.size(); i++) {
      if (!analyzer_->CanProveEqual(a[i], b[i])) {
        return false;
      }
    }
    return true;
  }

  bool HasSideEffect(const PrimExpr& value) {
    if (ignore_profiler_call_) {
      if (const CallNode* call = value.as<CallNode>()) {
        static const Op& timer_init_cuda_op = Op::Get("tirx.timer_init_cuda");
        static const Op& timer_start_cuda_op = Op::Get("tirx.timer_start_cuda");
        static const Op& timer_end_cuda_op = Op::Get("tirx.timer_end_cuda");
        static const Op& timer_finalize_cuda_op = Op::Get("tirx.timer_finalize_cuda");
        if (call->op.same_as(timer_init_cuda_op) || call->op.same_as(timer_start_cuda_op) ||
            call->op.same_as(timer_end_cuda_op) || call->op.same_as(timer_finalize_cuda_op)) {
          return false;
        }
      }
    }
    return SideEffect(value) > CallEffectKind::kReadState;
  }

  Stmt MakeEvaluate(PrimExpr value) {
    if (SideEffect(value) > CallEffectKind::kReadState) {
      return Evaluate(value);
    } else {
      return Evaluate(0);
    }
  }
  Stmt MakeEvaluate(const ffi::Array<PrimExpr>& values) {
    ffi::Array<Stmt> stmts;
    for (PrimExpr e : values) {
      if (SideEffect(e) > CallEffectKind::kReadState) {
        stmts.push_back(Evaluate(e));
      }
    }

    if (stmts.size() == 0) {
      return Evaluate(0);
    } else if (stmts.size() == 1) {
      return stmts[0];
    } else {
      return SeqStmt(stmts);
    }
  }

  std::unordered_map<const VarNode*, arith::IntSet> var_range_map_;
  bool ignore_profiler_call_{false};
};

Stmt RemoveNoOp(Stmt stmt, const arith::Analyzer& analyzer, bool ignore_profiler_call) {
  return NoOpRemover::Apply(std::move(stmt), analyzer, ignore_profiler_call);
}

namespace transform {

Pass RemoveNoOp() {
  auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
    RemoveNoOpConfig config =
        ctx->GetConfig<RemoveNoOpConfig>("tirx.RemoveNoOp")
            .value_or(tvm::transform::PassConfigWithDefaults<RemoveNoOpConfig>());

    arith::Analyzer analyzer;
    analyzer->rewrite_simplify.SetMaximumRewriteSteps(config->max_simplification_steps);

    bool ignore_profiler_call = config->ignore_profiler_call;

    {
      auto* write_ptr = f.CopyOnWrite();
      write_ptr->body =
          NoOpRemover::Apply(std::move(write_ptr->body), analyzer, ignore_profiler_call);
    }
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "tirx.RemoveNoOp", {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tirx.transform.RemoveNoOp", RemoveNoOp);
}

}  // namespace transform

}  // namespace tirx
}  // namespace tvm
