/*
 * 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 src/ir/transform.cc
 * \brief Infrastructure for transformation passes.
 */
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/extra/visit_error_context.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ir/transform.h>
#include <tvm/relax/expr.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/logging.h>

#include <optional>
#include <sstream>
#include <stack>
#include <string>

namespace tvm {
namespace transform {

using tvm::ffi::Any;

TVM_REGISTER_PASS_CONFIG_OPTION("testing.immutable_module", bool);

struct PassContextThreadLocalEntry {
  /*! \brief The default pass context. */
  PassContext default_context;

  /*! \brief The current pass context. */
  std::stack<PassContext> context_stack;

  PassContextThreadLocalEntry() {
    default_context = PassContext(ffi::make_object<PassContextNode>());
  }
};

/*! \brief Thread local store to hold the pass context. */
static PassContextThreadLocalEntry* PassContextThreadLocalStoreGet() {
  static thread_local PassContextThreadLocalEntry inst;
  return &inst;
}

void PassContext::EnterWithScope() {
  InstrumentEnterPassContext();

  PassContextThreadLocalEntry* entry = PassContextThreadLocalStoreGet();
  entry->context_stack.push(*this);
}

void PassContext::ExitWithScope() {
  PassContextThreadLocalEntry* entry = PassContextThreadLocalStoreGet();
  TVM_FFI_ICHECK(!entry->context_stack.empty());
  TVM_FFI_ICHECK(entry->context_stack.top().same_as(*this));
  entry->context_stack.pop();

  InstrumentExitPassContext();
}

PassContext PassContext::Current() {
  PassContextThreadLocalEntry* entry = PassContextThreadLocalStoreGet();
  if (!entry->context_stack.empty()) {
    return entry->context_stack.top();
  } else {
    return entry->default_context;
  }
}

// linearly scan the pass array to match pass_name
bool PassArrayContains(const ffi::Array<ffi::String>& pass_array, const std::string& pass_name) {
  for (auto x : pass_array) {
    if (x == pass_name) return true;
  }
  return false;
}

bool PassContext::PassEnabled(const PassInfo& info) const {
  if (PassArrayContains(operator->()->disabled_pass, info->name)) {
    return false;
  }

  if (PassArrayContains(operator->()->required_pass, info->name)) {
    return true;
  }

  return operator->()->opt_level >= info->opt_level;
}

class PassConfigManager {
 public:
  void Register(std::string key, ffi::String value_type_str,
                std::function<ffi::Any(ffi::Any)> legalization) {
    TVM_FFI_ICHECK_EQ(key2vtype_.count(key), 0U);
    ValueTypeInfo info;
    info.type_str = value_type_str;
    info.legalization = legalization;
    key2vtype_[key] = info;
  }

  // Trying to validate and legalize a config.
  void Legalize(ffi::Map<ffi::String, ffi::Any>* config) {
    std::vector<std::pair<std::string, ffi::Any>> update;
    for (auto [key, value] : *config) {
      auto it = key2vtype_.find(key);
      if (it == key2vtype_.end()) {
        std::ostringstream os;
        os << "Invalid config option \'" << key << "\' candidates are:";
        int counter = 0;
        for (const auto& [key, value] : key2vtype_) {
          os << ' ';
          if (counter++ != 0) os << ',';
          os << key;
        }
        TVM_FFI_THROW(AttributeError) << os.str();
      }
      const auto& info = it->second;

      TVM_FFI_CHECK(value != nullptr, AttributeError) << key << " is None";

      TVM_FFI_CHECK(info.legalization, AttributeError)
          << "Config option \'" << key << "\' was defined without a legalization function.";
      auto legalized = info.legalization(value);
      if (!legalized.same_as(value)) {
        update.emplace_back(key, legalized);
      }
    }
    for (auto&& kv : update) {
      config->Set(kv.first, kv.second);
    }
  }

  ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> ListConfigs() {
    ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> configs;
    for (const auto& kv : key2vtype_) {
      ffi::Map<ffi::String, ffi::String> metadata;
      metadata.Set("type", kv.second.type_str);
      configs.Set(kv.first, metadata);
    }
    return configs;
  }

  static PassConfigManager* Global() {
    static auto* inst = new PassConfigManager();
    return inst;
  }

 private:
  struct ValueTypeInfo {
    std::string type_str;
    std::function<ffi::Any(ffi::Any)> legalization;
  };

  std::unordered_map<std::string, ValueTypeInfo> key2vtype_;
};

void PassContext::RegisterConfigOption(const char* key, ffi::String value_type_str,
                                       std::function<ffi::Any(ffi::Any)> legalization) {
  PassConfigManager::Global()->Register(key, value_type_str, legalization);
}

ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> PassContext::ListConfigs() {
  return PassConfigManager::Global()->ListConfigs();
}

PassContext PassContext::Create() { return PassContext(ffi::make_object<PassContextNode>()); }

namespace {
struct ClearOnError {
  ffi::Array<instrument::PassInstrument>* instruments{nullptr};

  ~ClearOnError() {
    if (instruments) {
      LOG(INFO) << "Pass instrumentation enter/exti failed.";
      LOG(INFO) << "Disabling pass instrumentation.";
      instruments->clear();
    }
  }
};
struct ExitContextOnError {
  std::vector<instrument::PassInstrument> successes;

  ~ExitContextOnError() {
    for (auto it = successes.rbegin(); it != successes.rend(); it++) {
      LOG(INFO) << (*it)->name << " exiting PassContext ...";
      (*it)->ExitPassContext();
      LOG(INFO) << (*it)->name << " exited PassContext.";
    }
  }
};
}  // namespace

void PassContext::InstrumentEnterPassContext() {
  auto pass_ctx_node = this->operator->();
  if (pass_ctx_node->instruments.defined()) {
    ClearOnError clear_context{&pass_ctx_node->instruments};
    ExitContextOnError exit_context;
    for (instrument::PassInstrument pi : pass_ctx_node->instruments) {
      pi->EnterPassContext();
      exit_context.successes.push_back(pi);
    }
    exit_context.successes.clear();
    clear_context.instruments = nullptr;
  }
}

namespace {

struct ExitPassSuccesses {
  ~ExitPassSuccesses() {
    if (all_initialized) {
      return;
    }

    LOG(INFO) << "Pass instrumentation entering pass context failed.";
    LOG(INFO) << "Disable pass instrumentation.";
    instruments->clear();

    for (auto it = successes.rbegin(); it != successes.rend(); it++) {
      LOG(INFO) << (*it)->name << " exiting PassContext ...";
      (*it)->ExitPassContext();
      LOG(INFO) << (*it)->name << " exited PassContext.";
    }
  }

  bool all_initialized{false};
  std::vector<instrument::PassInstrument> successes;
  ffi::Array<instrument::PassInstrument>* instruments{nullptr};
};
}  // namespace

void PassContext::InstrumentExitPassContext() {
  auto pass_ctx_node = this->operator->();
  if (pass_ctx_node->instruments.defined()) {
    ClearOnError clear_context{&pass_ctx_node->instruments};
    for (instrument::PassInstrument pi : pass_ctx_node->instruments) {
      pi->ExitPassContext();
    }
    clear_context.instruments = nullptr;
  }
}

bool PassContext::InstrumentBeforePass(const IRModule& ir_module, const PassInfo& pass_info) const {
  auto pass_ctx_node = this->operator->();
  if (!pass_ctx_node->instruments.defined()) {
    return true;
  }

  const bool pass_required = PassArrayContains(pass_ctx_node->required_pass, pass_info->name);
  bool should_run = true;
  if (!pass_required) {
    for (instrument::PassInstrument pi : pass_ctx_node->instruments) {
      should_run &= pi->ShouldRun(ir_module, pass_info);
    }
  }

  if (should_run) {
    for (instrument::PassInstrument pi : pass_ctx_node->instruments) {
      pi->RunBeforePass(ir_module, pass_info);
    }
  }
  return should_run;
}

void PassContext::InstrumentAfterPass(const IRModule& ir_module, const PassInfo& pass_info) const {
  auto pass_ctx_node = this->operator->();
  if (pass_ctx_node->instruments.defined()) {
    for (instrument::PassInstrument pi : pass_ctx_node->instruments) {
      pi->RunAfterPass(ir_module, pass_info);
    }
  }
}

IRModule Pass::operator()(IRModule mod) const {
  return this->operator()(std::move(mod), PassContext::Current());
}

namespace {
/*! \brief Marker line that precedes the TVMScript-rendered location block. */
constexpr const char* kLocationMarker = "Location (TVMScript):";

/*!
 * \brief Render `node` as TVMScript with `path` underlined, via the script
 *        printer global registry, so src/ir does not link against the printer.
 *
 * Returns the rendered snippet, or std::nullopt if the printer entry is
 * unavailable or throws.
 */
std::optional<std::string> RenderScriptWithUnderline(const ffi::ObjectRef& node,
                                                     const ffi::reflection::AccessPath& path) {
  namespace ffi = tvm::ffi;
  auto config_fn = ffi::Function::GetGlobal("node.PrinterConfig");
  auto script_fn = ffi::Function::GetGlobal("node.TVMScriptPrinterScript");
  if (!config_fn.has_value() || !script_fn.has_value()) return std::nullopt;
  try {
    ffi::Map<ffi::String, ffi::Any> config_dict;
    config_dict.Set("path_to_underline", ffi::Array<ffi::reflection::AccessPath>{path});
    // Show enough context that a small function renders in full (no
    // "(... N lines skipped ...)" marker), while still bounding a large
    // module. A small TIR/Relax function is ~8-15 lines; 10 lines of context
    // on each side of the underline covers it end-to-end.
    config_dict.Set("num_context_lines", static_cast<int>(10));
    config_dict.Set("render_invisible_path_info", true);
    ffi::Any cfg = (*config_fn)(config_dict);
    ffi::Any rendered = (*script_fn)(node, cfg);
    return rendered.cast<ffi::String>().operator std::string();
  } catch (...) {
    return std::nullopt;
  }
}
}  // namespace

ffi::Error EnrichPassErrorWithContext(const ffi::Error& err, const IRModule& mod,
                                      ffi::String pass_name, ffi::Optional<GlobalVar> func) {
  namespace ffi = tvm::ffi;
  // No visit-context payload => nothing to resolve; return the error untouched.
  // (The enriched error below drops the payload, so an outer catch that
  // re-enriches lands here and returns unchanged — no double location block.)
  auto visit_ctx = ffi::VisitErrorContext::TryGetFromError(err);
  if (!visit_ctx) return err;

  // Resolve + render against the local function when requested, else the module.
  // Fall back to the module root if the function is missing or yields no path.
  ffi::ObjectRef root = mod;
  if (func.has_value()) {
    if (auto bf = mod->functions.Get(func.value())) {
      root = bf.value();
    } else {
      func = ffi::Optional<GlobalVar>(std::nullopt);
    }
  }

  auto try_render = [&](const ffi::ObjectRef& search_root) -> std::optional<std::string> {
    auto paths = ffi::VisitErrorContext::FindAccessPaths(search_root, visit_ctx.value(),
                                                         /*allow_prefix_match=*/true);
    if (paths.empty()) return std::nullopt;
    return RenderScriptWithUnderline(search_root, paths[0]);
  };

  std::optional<std::string> snippet = try_render(root);
  if (!snippet && func.has_value()) {
    // Function-local resolution failed; retry against the whole module.
    snippet = try_render(mod);
  }
  if (!snippet) return err;

  // Append the pass name + underlined TVMScript location. Preserve kind,
  // original message, and backtrace; drop the VisitErrorContext payload so an
  // outer catch does not re-enrich.
  std::ostringstream suffix;
  suffix << "\n\nError in pass: " << pass_name << "\n"
         << kLocationMarker << "\n"
         << snippet.value();
  return ffi::Error(err.kind(), err.message() + suffix.str(), err.backtrace());
}

IRModule Pass::operator()(IRModule mod, const PassContext& pass_ctx) const {
  const PassNode* node = operator->();
  TVM_FFI_ICHECK(node != nullptr);
  const PassInfo& pass_info = node->Info();
  if (!pass_ctx.InstrumentBeforePass(mod, pass_info)) {
    DLOG(INFO) << "Skipping pass : " << pass_info->name
               << " with opt level: " << pass_info->opt_level;
    return mod;
  }
  IRModule ret;
  // Error enrichment happens in the leaf pass executors (ModulePassNode /
  // FunctionPassNode), not here — guarding the generic funnel (which also wraps
  // Sequential) would re-resolve FindAccessPaths against a mutated post-pass
  // module and double-wrap the message.
  if (pass_ctx->GetConfig<bool>("testing.immutable_module", false).value()) {
    ret = Pass::AssertImmutableModule(mod, node, pass_ctx);
  } else {
    ret = node->operator()(std::move(mod), pass_ctx);
  }
  pass_ctx.InstrumentAfterPass(ret, pass_info);
  return ret;
}

IRModule Pass::AssertImmutableModule(const IRModule& mod, const PassNode* node,
                                     const PassContext& pass_ctx) {
  size_t before_pass_hash = ffi::StructuralHash()(mod);
  IRModule copy_mod = mod;
  IRModule ret = node->operator()(mod, pass_ctx);
  size_t after_pass_hash = ffi::StructuralHash()(copy_mod);
  if (before_pass_hash != after_pass_hash) {
    // The chance of getting a hash conflict between a module and the same module but mutated
    // must be very low.
    LOG_FATAL << "Immutable module has been modified in pass: " << node->Info()->name;
  }
  return ret;
}

/*!
 * \brief Module-level passes are designed to implement global
 * analysis/optimizations, i.e. interprocedural optimizations (IPO), etc. Passes
 * at this level have the full control of a given Relax program including
 * addition and deletion of functions.
 */
class ModulePassNode : public PassNode {
 public:
  /* \brief The pass meta data.*/
  PassInfo pass_info;

  /*! \brief The pass function sketches the real optimization. For example,
   * we may need to perform dead code elimination on the module level. We could
   * implement the algorithm in the `pass_func` and let it run on a module. It
   * will then remove the dead code including the unused functions in the module.
   */
  std::function<IRModule(IRModule, PassContext)> pass_func;

  ModulePassNode() = default;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ModulePassNode>().def_ro("pass_info", &ModulePassNode::pass_info);
  }

  /*!
   * \brief Run a module pass on given pass context.
   *
   * \param mod The module that an optimization pass is applied on.
   * \param mod The context that an optimization pass executes on.
   *
   * \return Return the updated module.
   */
  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;

  /*!
   * \brief Get the pass information/meta data.
   */
  PassInfo Info() const override { return pass_info; }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.ModulePass", ModulePassNode, PassNode);
};

class ModulePass : public Pass {
 public:
  ModulePass(std::function<IRModule(IRModule, PassContext)> pass_func, PassInfo pass_info);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ModulePass, Pass, ModulePassNode);
};

PassInfo::PassInfo(int opt_level, ffi::String name, tvm::ffi::Array<ffi::String> required,
                   bool traceable) {
  auto pass_info = ffi::make_object<PassInfoNode>();
  pass_info->opt_level = opt_level;
  pass_info->name = std::move(name);
  pass_info->required = std::move(required);
  pass_info->traceable = std::move(traceable);
  data_ = std::move(pass_info);
}

ModulePass::ModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
                       PassInfo pass_info) {
  auto n = ffi::make_object<ModulePassNode>();
  n->pass_func = std::move(pass_func);
  n->pass_info = std::move(pass_info);
  data_ = std::move(n);
}

// Module -> Module optimizations.
IRModule ModulePassNode::operator()(IRModule mod, const PassContext& pass_ctx) const {
  const PassInfo& pass_info = Info();
  TVM_FFI_ICHECK(mod.defined()) << "The input module must be set.";

  VLOG_CONTEXT << pass_info->name;
  VLOG(0) << "Executing module pass with opt level: " << pass_info->opt_level;

  // The module is the access-path root for any visit-context error the pass
  // body throws. Enrich at this leaf executor so the message surfaces a precise,
  // TVMScript-rendered location.
  IRModule input_mod = mod;
  try {
    mod = pass_func(std::move(mod), pass_ctx);
  } catch (ffi::Error& err) {
    throw EnrichPassErrorWithContext(err, input_mod, pass_info->name);
  }

  TVM_FFI_ICHECK(mod.defined()) << "The return value of a module pass must be set.";

  return mod;
}

Sequential::Sequential(tvm::ffi::Array<Pass> passes, PassInfo pass_info) {
  auto n = ffi::make_object<SequentialNode>();
  n->passes = std::move(passes);
  n->pass_info = std::move(pass_info);
  data_ = std::move(n);
}

Sequential::Sequential(tvm::ffi::Array<Pass> passes, ffi::String name) {
  auto n = ffi::make_object<SequentialNode>();
  n->passes = std::move(passes);
  PassInfo pass_info = PassInfo(0, std::move(name), {}, /* traceable */ false);
  n->pass_info = std::move(pass_info);
  data_ = std::move(n);
}

const SequentialNode* Sequential::operator->() const {
  return static_cast<const SequentialNode*>(get());
}

void SequentialNode::ResolveDependency(const IRModule& mod) {
  // TODO(zhiics) Implement it.
  // 1. Consider the required passes for each pass.
  // 2. Only resolve the enabled passes.
  // 3. Build a dependency graph. Probably we need to update the pass list.
  TVM_FFI_THROW(InternalError) << "Pass dependency has not been resolved yet."
                               << "\n";
}

Pass GetPass(const ffi::String& pass_name) {
  std::optional<tvm::ffi::Function> f;
  if (pass_name.operator std::string().find("transform.") != std::string::npos) {
    f = tvm::ffi::Function::GetGlobal(pass_name);
  } else {
    f = tvm::ffi::Function::GetGlobal("transform." + pass_name);
  }
  TVM_FFI_ICHECK(f.has_value()) << "Cannot use " << pass_name << " to create the pass";
  return (*f)().cast<Pass>();
}

// TODO(zhiics): we currently only sequentially execute each pass in
// a Sequential without the consideration of their orders. The phase
// ordering problem needs to be handled in the future.
IRModule SequentialNode::operator()(IRModule mod, const PassContext& pass_ctx) const {
  for (const Pass& pass : passes) {
    VLOG(0) << "Running pass " << pass->Info()->name;
    TVM_FFI_ICHECK(pass.defined()) << "Found undefined pass for optimization.";
    const PassInfo& pass_info = pass->Info();
    if (!pass_ctx.PassEnabled(pass_info)) {
      VLOG(0) << "skipping disabled pass '" << pass_info->name << "'";
      continue;
    }

    // resolve dependencies
    for (const auto& it : pass_info->required) {
      mod = GetPass(it)(std::move(mod), pass_ctx);
    }

    mod = pass(std::move(mod), pass_ctx);
  }
  return mod;
}

Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func, int opt_level,
                      ffi::String name, tvm::ffi::Array<ffi::String> required, bool traceable) {
  PassInfo pass_info = PassInfo(opt_level, name, required, traceable);
  return ModulePass(std::move(pass_func), pass_info);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("transform.PassInfo",
           [](int opt_level, ffi::String name, tvm::ffi::Array<ffi::String> required,
              bool traceable) { return PassInfo(opt_level, name, required, traceable); })
      .def_packed("transform.Info", [](ffi::PackedArgs args, ffi::Any* ret) {
        Pass pass = args[0].cast<Pass>();
        *ret = pass->Info();
      });
}

// Pattern A (RM): auto-default repr from reflection for PassInfoNode.

TVM_FFI_STATIC_INIT_BLOCK() {
  PassContextNode::RegisterReflection();
  PassInfoNode::RegisterReflection();
  SequentialNode::RegisterReflection();
  ModulePassNode::RegisterReflection();
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("transform.MakeModulePass",
           [](ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>, PassContext)> pass_func,
              PassInfo pass_info) {
             auto wrapped_pass_func = [pass_func](IRModule mod, PassContext ctx) {
               return pass_func(ffi::RValueRef<IRModule>(std::move(mod)), ctx);
             };
             return ModulePass(wrapped_pass_func, pass_info);
           })
      .def("transform.RunPass",
           [](Pass pass, ffi::RValueRef<IRModule> mod) { return pass(*std::move(mod)); });
}

// Pattern A (RM): auto-default repr from reflection for ModulePassNode.

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def_packed("transform.Sequential", [](ffi::PackedArgs args, ffi::Any* ret) {
    auto passes = args[0].cast<tvm::ffi::Array<Pass>>();
    int opt_level = args[1].cast<int>();
    std::string name = args[2].cast<std::string>();
    auto required = args[3].cast<tvm::ffi::Array<ffi::String>>();
    bool traceable = args[4].cast<bool>();
    PassInfo pass_info = PassInfo(opt_level, name, required, /* traceable */ traceable);
    *ret = Sequential(passes, pass_info);
  });
}

// Pattern A (RM): auto-default repr from reflection for SequentialNode.

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def(
      "transform.PassContext",
      [](int opt_level, ffi::Array<ffi::String> required, ffi::Array<ffi::String> disabled,
         ffi::Array<instrument::PassInstrument> instruments,
         ffi::Optional<ffi::Map<ffi::String, ffi::Any>> config) {
        auto pctx = PassContext::Create();
        pctx->opt_level = opt_level;

        pctx->required_pass = std::move(required);
        pctx->disabled_pass = std::move(disabled);
        pctx->instruments = std::move(instruments);

        if (config.has_value()) {
          pctx->config = config.value();
        }
        PassConfigManager::Global()->Legalize(&(pctx->config));
        return pctx;
      });
}

// Pattern A (RM): auto-default repr from reflection for PassContextNode.

class PassContext::Internal {
 public:
  static void EnterScope(PassContext pass_ctx) { pass_ctx.EnterWithScope(); }

  static void ExitScope(PassContext pass_ctx) { pass_ctx.ExitWithScope(); }
};

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("transform.GetCurrentPassContext", PassContext::Current)
      .def("transform.EnterPassContext", PassContext::Internal::EnterScope)
      .def("transform.ExitPassContext", PassContext::Internal::ExitScope)
      .def("transform.OverrideInstruments",
           [](PassContext pass_ctx, ffi::Array<instrument::PassInstrument> instruments) {
             pass_ctx.InstrumentExitPassContext();
             pass_ctx->instruments = instruments;
             pass_ctx.InstrumentEnterPassContext();
           });
}

Pass PrintIR(ffi::String header) {
  auto pass_func = [header](IRModule mod, const PassContext& ctx) {
    LOG(INFO) << "PrintIR(" << header << "):\n" << mod;
    return mod;
  };
  return CreateModulePass(pass_func, 0, "PrintIR", {}, /* traceable */ false);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("transform.PrintIR", PrintIR)
      .def("transform.ListConfigs", PassContext::ListConfigs);
}

}  // namespace transform
}  // namespace tvm
