/*
 * 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/instrument.cc
 * \brief Infrastructure for instrumentation.
 */
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/instrument.h>
#include <tvm/ir/transform.h>
#include <tvm/runtime/logging.h>

#include <chrono>
#include <iomanip>
#include <stack>

namespace tvm {
namespace instrument {

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

/*!
 * \brief Base PassInstrument implementation
 * \sa BasePassInstrument
 */
class BasePassInstrumentNode : public PassInstrumentNode {
 public:
  /*! \brief Callback to run when entering PassContext. */
  ffi::TypedFunction<void()> enter_pass_ctx_callback;
  /*! \brief Callback to run when exiting PassContext. */
  ffi::TypedFunction<void()> exit_pass_ctx_callback;

  /*! \brief Callback determines whether to run a pass or not. */
  ffi::TypedFunction<bool(const IRModule&, const transform::PassInfo&)> should_run_callback;

  /*! \brief Callback to run before a pass. */
  ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)> run_before_pass_callback;
  /*! \brief Callback to run after a pass. */
  ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)> run_after_pass_callback;

  /*! \brief Instrument when entering PassContext. */
  void EnterPassContext() const final;

  /*! \brief Instrument when exiting PassContext. */
  void ExitPassContext() const final;

  /*!
   * \brief Determine whether to run the pass or not.
   * \param mod The module that an optimization pass runs on.
   * \param info The pass information.
   *
   * \return true to run the pass; false to skip the pass.
   */
  bool ShouldRun(const IRModule&, const transform::PassInfo& info) const final;

  /*!
   * \brief Instrument before pass run.
   * \param mod The module that an optimization pass runs on.
   * \param info The pass information.
   */
  void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const final;

  /*!
   * \brief Instrument after pass run.
   *
   * \param mod The module that an optimization pass runs on.
   * \param info The pass information.
   */
  void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const final;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("instrument.PassInstrument", BasePassInstrumentNode,
                                    PassInstrumentNode);
};

/*!
 * \brief Managed reference class for BasePassInstrumentNode
 * \sa BasePassInstrumentNode
 */
class BasePassInstrument : public PassInstrument {
 public:
  /*!
   * \brief Constructor
   *
   * \param name Name for this instrumentation.
   *
   *
   * \param enter_pass_ctx_callback Callback to call when entering pass context.
   * \param exit_pass_ctx_callback Callback to call when exiting pass context.
   *
   * \param should_run_callback Callback to determine whether pass should run. (return true: enable;
   *                            return false: disable)
   *
   * \param run_before_pass_callback Callback to call before a pass run.
   * \param run_after_pass_callback Callback to call after a pass run.
   */
  TVM_DLL BasePassInstrument(
      ffi::String name, ffi::TypedFunction<void()> enter_pass_ctx_callback,
      ffi::TypedFunction<void()> exit_pass_ctx_callback,
      ffi::TypedFunction<bool(const IRModule&, const transform::PassInfo&)> should_run_callback,
      ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)>
          run_before_pass_callback,
      ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)>
          run_after_pass_callback);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BasePassInstrument, PassInstrument,
                                             BasePassInstrumentNode);
};

BasePassInstrument::BasePassInstrument(
    ffi::String name, ffi::TypedFunction<void()> enter_pass_ctx_callback,
    ffi::TypedFunction<void()> exit_pass_ctx_callback,
    ffi::TypedFunction<bool(const IRModule&, const transform::PassInfo&)> should_run_callback,
    ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)> run_before_pass_callback,
    ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)> run_after_pass_callback) {
  auto pi = ffi::make_object<BasePassInstrumentNode>();
  pi->name = std::move(name);

  pi->enter_pass_ctx_callback = std::move(enter_pass_ctx_callback);
  pi->exit_pass_ctx_callback = std::move(exit_pass_ctx_callback);

  pi->should_run_callback = std::move(should_run_callback);

  pi->run_before_pass_callback = std::move(run_before_pass_callback);
  pi->run_after_pass_callback = std::move(run_after_pass_callback);

  data_ = std::move(pi);
}

void BasePassInstrumentNode::EnterPassContext() const {
  if (enter_pass_ctx_callback != nullptr) {
    enter_pass_ctx_callback();
  }
}

void BasePassInstrumentNode::ExitPassContext() const {
  if (exit_pass_ctx_callback != nullptr) {
    exit_pass_ctx_callback();
  }
}

bool BasePassInstrumentNode::ShouldRun(const IRModule& ir_module,
                                       const transform::PassInfo& pass_info) const {
  if (should_run_callback == nullptr) {
    return true;
  }

  return should_run_callback(ir_module, pass_info);
}

void BasePassInstrumentNode::RunBeforePass(const IRModule& ir_module,
                                           const transform::PassInfo& pass_info) const {
  if (run_before_pass_callback != nullptr) {
    run_before_pass_callback(ir_module, pass_info);
  }
}

void BasePassInstrumentNode::RunAfterPass(const IRModule& ir_module,
                                          const transform::PassInfo& pass_info) const {
  if (run_after_pass_callback != nullptr) {
    run_after_pass_callback(ir_module, pass_info);
  }
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def(
      "instrument.PassInstrument",
      [](ffi::String name, ffi::TypedFunction<void()> enter_pass_ctx,
         ffi::TypedFunction<void()> exit_pass_ctx,
         ffi::TypedFunction<bool(const IRModule&, const transform::PassInfo&)> should_run,
         ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)> run_before_pass,
         ffi::TypedFunction<void(const IRModule&, const transform::PassInfo&)> run_after_pass) {
        return BasePassInstrument(name, enter_pass_ctx, exit_pass_ctx, should_run, run_before_pass,
                                  run_after_pass);
      });
}

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

/*! \brief PassProfile stores profiling information for a given pass and its sub-passes. */
struct PassProfile {
  // TODO(@altanh): expose PassProfile through TVM ffi::Object API
  using Clock = std::chrono::steady_clock;
  using Duration = std::chrono::duration<double, std::micro>;
  using Time = std::chrono::time_point<Clock>;

  /*! \brief The name of the pass being profiled. */
  ffi::String name;
  /*! \brief The time when the pass was entered. */
  Time start;
  /*! \brief The time when the pass completed. */
  Time end;
  /*! \brief The total duration of the pass, i.e. end - start. */
  Duration duration;
  /*! \brief PassProfiles for all sub-passes invoked during the execution of the pass. */
  std::vector<PassProfile> children;

  explicit PassProfile(ffi::String name)
      : name(name), start(Clock::now()), end(Clock::now()), children() {}

  /*! \brief Gets the PassProfile of the currently executing pass. */
  static PassProfile* Current();
  /*! \brief Pushes a new PassProfile with the given pass name. */
  static void EnterPass(ffi::String name);
  /*! \brief Pops the current PassProfile. */
  static void ExitPass();
};

struct PassProfileThreadLocalEntry {
  /*! \brief The placeholder top-level PassProfile. */
  PassProfile root;
  /*! \brief The stack of PassProfiles for nested passes currently running. */
  std::stack<PassProfile*> profile_stack;

  PassProfileThreadLocalEntry() : root("root") {}
};

/*! \brief Thread local store to hold the pass profiling data. */
static PassProfileThreadLocalEntry* PassProfileThreadLocalStoreGet() {
  static thread_local PassProfileThreadLocalEntry inst;
  return &inst;
}

void PassProfile::EnterPass(ffi::String name) {
  PassProfile* cur = PassProfile::Current();
  cur->children.emplace_back(name);
  PassProfileThreadLocalStoreGet()->profile_stack.push(&cur->children.back());
}

void PassProfile::ExitPass() {
  PassProfile* cur = PassProfile::Current();
  TVM_FFI_ICHECK_NE(cur->name, "root") << "mismatched enter/exit for pass profiling";
  cur->end = PassProfile::Clock::now();
  cur->duration = std::chrono::duration_cast<PassProfile::Duration>(cur->end - cur->start);
  PassProfileThreadLocalStoreGet()->profile_stack.pop();
}

PassProfile* PassProfile::Current() {
  PassProfileThreadLocalEntry* entry = PassProfileThreadLocalStoreGet();
  if (!entry->profile_stack.empty()) {
    return entry->profile_stack.top();
  } else {
    return &entry->root;
  }
}

ffi::String RenderPassProfiles() {
  PassProfileThreadLocalEntry* entry = PassProfileThreadLocalStoreGet();
  TVM_FFI_ICHECK(entry->profile_stack.empty())
      << "cannot print pass profile while still in a pass!";

  if (entry->root.children.empty()) {
    LOG(WARNING) << "no passes have been profiled, did you enable pass profiling?";
    return ffi::String();
  }

  // (depth, parent_duration, pass)
  std::stack<std::tuple<size_t, PassProfile::Duration, PassProfile*>> profiles;

  // push top level passes
  PassProfile::Duration top_dur(0);
  for (auto it = entry->root.children.begin(); it != entry->root.children.end(); ++it) {
    top_dur += it->duration;
  }
  for (auto it = entry->root.children.rbegin(); it != entry->root.children.rend(); ++it) {
    profiles.push(std::make_tuple(0, top_dur, &*it));
  }

  std::ostringstream os;
  os << std::fixed;

  while (profiles.size() > 0) {
    auto [depth, parent_duration, profile] = profiles.top();
    profiles.pop();

    // indent depth
    for (size_t i = 0; i < depth; ++i) {
      os << "\t";
    }

    // calculate time spent in pass itself (excluding sub-passes), and push children
    PassProfile::Duration self_duration = profile->duration;
    for (auto it = profile->children.rbegin(); it != profile->children.rend(); ++it) {
      self_duration -= it->duration;
      profiles.push(std::make_tuple(depth + 1, profile->duration, &*it));
    }

    double parent_pct = profile->duration.count() / parent_duration.count() * 100.0;
    double total_pct = profile->duration.count() / top_dur.count() * 100.0;

    os << profile->name << ": ";
    os << std::setprecision(0);
    os << profile->duration.count() << "us [" << self_duration.count() << "us] ";
    os << std::setprecision(2) << "(" << total_pct << "%; " << parent_pct << "%)\n";
  }

  return os.str();
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("instrument.RenderTimePassProfiles", RenderPassProfiles)
      .def("instrument.MakePassTimingInstrument", []() {
        auto run_before_pass = [](const IRModule&, const transform::PassInfo& pass_info) {
          PassProfile::EnterPass(pass_info->name);
          return true;
        };

        auto run_after_pass = [](const IRModule&, const transform::PassInfo& pass_info) {
          PassProfile::ExitPass();
        };

        auto exit_pass_ctx = []() { PassProfileThreadLocalStoreGet()->root.children.clear(); };

        return BasePassInstrument("PassTimingInstrument",
                                  /* enter_pass_ctx */ nullptr, exit_pass_ctx,
                                  /* should_run */ nullptr, run_before_pass, run_after_pass);
      });
}

}  // namespace instrument
}  // namespace tvm
