TVM 中的 Profiler 設計

一、基本用法

首先看 Profiler 的用法：

with ms.Profiler() as profiler:
    # .... 用戶代碼
    
print("Tuning Time:")
print(profiler.table())

二、前端接口設計

其中 Profiler 類的設計是綁定和映射到了 C++ 端的接口上。Profile 提供了 Context 語義，支持 with 語句。

@register_object("meta_schedule.Profiler")
class Profiler(Object):
    """Tuning time profiler."""

    def __init__(self) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.Profiler,  # type: ignore # pylint: disable=no-member
        )

    def get(self) -> Dict[str, float]:
        """Get the profiling results in seconds"""
        return _ffi_api.ProfilerGet(self)  # type: ignore # pylint: disable=no-member

    def table(self) -> str:
        """Get the profiling results in a table format"""
        return _ffi_api.ProfilerTable(self)  # type: ignore # pylint: disable=no-member

    def __enter__(self) -> "Profiler":
        """Entering the scope of the context manager"""
        _ffi_api.ProfilerEnterWithScope(self)  # type: ignore # pylint: disable=no-member
        return self

    def __exit__(self, ptype, value, trace) -> None:
        """Exiting the scope of the context manager"""
        _ffi_api.ProfilerExitWithScope(self)  # type: ignore # pylint: disable=no-member

    @staticmethod
    def current() -> Optional["Profiler"]:
        """Get the current profiler."""
        return _ffi_api.ProfilerCurrent()  # type: ignore # pylint: disable=no-member

    @staticmethod
    def timeit(name: str):
        """Timeit a block of code"""

        @contextmanager
        def _timeit():
            try:
                f = _ffi_api.ProfilerTimedScope(name)  # type: ignore # pylint: disable=no-member
                yield
            finally:
                if f:
                    f()

        return _timeit()

其中 enter 調用時會執行 ProfilerEnterWithScope 其負責往一個 Stack 式的結構中添加一個 Profiler 對象：

void Profiler::EnterWithScope() {
  ThreadLocalProfilers()->push_back(*this);
  (*this)->total_timer = ProfilerTimedScope("Total");
}

在退出 with 語句時，調用 exit 執行 ProfilerExitWithScope 會調用 ProfilerTimedScope() 對象返回的函數，觸發當前層計時結束。
profiler.table() 負責收集所有的耗時統計信息，並按照預定義的格式 format 展示給用戶。

// std::unordered_map<std::string, double> stats_sec; << 這傢伙是被所有的Profiler共享的
// runtime::PackedFunc total_timer; << 負責外層的整體耗時計算

String ProfilerNode::Table() const {
  CHECK(!stats_sec.empty()) << "ValueError: The stats are empty. Please run the profiler first.";
  CHECK(stats_sec.count("Total"))
      << "ValueError: The total time is not recorded. This method should be called only after "
         "exiting the profiler's with scope.";
  double total = stats_sec.at("Total");
  struct Entry {
    String name;
    double minutes;
    double percentage;
    bool operator<(const Entry& other) const { return percentage > other.percentage; }
  };
  std::vector<Entry> table_entry;
  for (const auto& kv : stats_sec) {
    table_entry.push_back(Entry{kv.first, kv.second / 60.0, kv.second / total * 100.0});
  }
  std::sort(table_entry.begin(), table_entry.end());
  support::TablePrinter p;
  p.Row() << "ID"
          << "Name"
          << "Time (min)"
          << "Percentage";
  p.Separator();
  for (int i = 0, n = table_entry.size(); i < n; ++i) {
    if (i == 0) {
      p.Row() << "" << table_entry[i].name << table_entry[i].minutes << table_entry[i].percentage;
    } else {
      p.Row() << i << table_entry[i].name << table_entry[i].minutes << table_entry[i].percentage;
    }
  }
  p.Separator();
  return p.AsStr();
}

三、後端接口設計

// TVM 框架中某個函數
void Apply(const TaskScheduler& task_scheduler, int task_id,
             const Array<MeasureCandidate>& measure_candidates,
             const Array<BuilderResult>& builder_results,
             const Array<RunnerResult>& runner_results) final {
    auto _ = Profiler::TimedScope("MeasureCallback/AddToDatabase");  // 構造時觸發計時開始
    // 框架代碼
    
  } // 出了函數的作用域，_ 對象會被析構，觸發計時結束，計算duration，放到全局的table中
 
TimedScope的實現僅僅是返回了一個ScopedTimer對象，由ScopedTimer對象的析構函數負責觸發「計時結束」。
ScopedTimer Profiler::TimedScope(String name) { return ScopedTimer(ProfilerTimedScope(name)); }


PackedFunc ProfilerTimedScope(String name) {
  if (Optional<Profiler> opt_profiler = Profiler::Current()) {           // <---- Profiler::Current()是一個「棧」設計，爲了支持「嵌套統計」功能
    return TypedPackedFunc<void()>([profiler = opt_profiler.value(),                  //
                                    tik = std::chrono::high_resolution_clock::now(),  // <--- 創建一個函數deleter回調函數，藉助lambda函數傳遞計時起始點
                                    name = std::move(name)]() {
      auto tok = std::chrono::high_resolution_clock::now();
      double duration =
          std::chrono::duration_cast<std::chrono::nanoseconds>(tok - tik).count() / 1e9;
      profiler->stats_sec[name] += duration;      // <---- 
    });
  }
  return nullptr;
}

TVM 中的 Profiler 是支持多重任意嵌套的，實現通過了一個 Vector<Profile> 模擬了棧了操作：

std::vector<Profiler>* ThreadLocalProfilers() {
  static thread_local std::vector<Profiler> profilers;   // <---- 支持嵌套，全局stack 結構
  return &profilers;
}

void Profiler::EnterWithScope() {
  ThreadLocalProfilers()->push_back(*this);     // 入棧
  (*this)->total_timer = ProfilerTimedScope("Total");
}

void Profiler::ExitWithScope() {
  ThreadLocalProfilers()->pop_back();   // 出棧
  if ((*this)->total_timer != nullptr) {
    (*this)->total_timer();
    (*this)->total_timer = nullptr;
  }
}

附錄：飛槳 Profiler 設計

一、基本用法

如下是一個簡單的使用樣例：

import paddle.fluid as fluid
import paddle.fluid.profiler as profiler

profiler.start_profiler('GPU')    # <---- 開始記錄
for iter in range(10):
    if iter == 2:
        profiler.reset_profiler()
    # except each iteration
profiler.stop_profiler('total', '/tmp/profile')   # <---- 結束並記錄結果到文件裏

"""
------------------------->     Profiling Report     <-------------------------

Place: CPU
Time unit: ms
Sorted by total time in descending order in the same thread
#Sorted by number of calls in descending order in the same thread
#Sorted by number of max in descending order in the same thread
#Sorted by number of min in descending order in the same thread
#Sorted by number of avg in descending order in the same thread

Event                       Calls       Total       Min.        Max.        Ave.        Ratio.
thread0::conv2d             8           129.406     0.304303    127.076     16.1758     0.983319
thread0::elementwise_add    8           2.11865     0.193486    0.525592    0.264832    0.016099
thread0::feed               8           0.076649    0.006834    0.024616    0.00958112  0.000582432

#### 2) sorted_key = None  ####
# Since the profiling results are printed in the order of first end time of Ops,
# the printed order is feed->conv2d->elementwise_add
------------------------->     Profiling Report     <-------------------------

Place: CPU
Time unit: ms
Sorted by event first end time in descending order in the same thread

Event                       Calls       Total       Min.        Max.        Ave.        Ratio.
thread0::feed               8           0.077419    0.006608    0.023349    0.00967738  0.00775934
thread0::conv2d             8           7.93456     0.291385    5.63342     0.99182     0.795243
thread0::elementwise_add    8           1.96555     0.191884    0.518004    0.245693    0.196998
"""

二、前端接口設計

上面涉及到了兩個核心的接口：start_profiler 和 stop_profiler

def start_profiler(state, tracer_option='Default'):
    if state == "GPU":
        prof_state = core.ProfilerState.kCUDA
    # ......
    if tracer_option == "Default":
        prof_tracer_option = core.TracerOption.kDefault
    # .....
    
    core.set_tracer_option(prof_tracer_option)
    core.enable_profiler(prof_state)

def stop_profiler(sorted_key=None, profile_path='/tmp/profile'):
    core.disable_profiler(key_map[sorted_key], profile_path)

這兩個接口內部實現都是通過「修改全局變量開關」來實現的，這些接口均通過 Pybind 映射到C++端：

static TracerOption g_tracer_option = TracerOption::kDefault;  // 全局靜態變量

void SetTracerOption(TracerOption option) {
  std::lock_guard<std::mutex> l(profiler_mu);
  g_tracer_option = option;
}


class ProfilerHelper {
 public:
  // The profiler state, the initial value is ProfilerState::kDisabled
  static ProfilerState g_state;
  
  // 省略
}

三、後端接口設計

C++ 端的主要用法：

  LOG(INFO) << "Usage 2: RecordEvent";
  for (int i = 1; i < 5; ++i) {
    std::string name = "evs_op_" + std::to_string(i);
    RecordEvent record_event(name);
    int counter = 1;
    while (counter != i * 1000) counter++;
  }

可以看出提供的核心數據結構類是 RecordEvent ，也是藉助對象的構造和析構來實現的。

// Default tracing level.
// It is Recommended to set the level explicitly.
static constexpr uint32_t kDefaultTraceLevel = 4;

class RecordEvent {
 public:
  static bool IsEnabled();
  
    explicit RecordEvent(
      const std::string& name,
      const TracerEventType type = TracerEventType::UserDefined,
      uint32_t level = kDefaultTraceLevel,
      const EventRole role = EventRole::kOrdinary);

    void End();

    ~RecordEvent() { End(); }
    
 private:
   std::string* name_{nullptr};
   uint64_t start_ns_;
   // 省略其他成員
 }

可以看出構造時初始化 start_ns_成員，在析構時調用 End 函數：

void RecordEvent::End() {
  uint64_t end_ns = PosixInNsec();
  HostEventRecorder<CommonEvent>::GetInstance().RecordEvent(
          shallow_copy_name_, start_ns_, end_ns, role_, type_);
}

從實現上看，框架層面所有的 Event 是保存到全局單例 HostEventRecorder 中的，包含了name、時間起點、終點、類型、角色：

template <typename EventType>
class HostEventRecorder {
 public:
  // singleton
  static HostEventRecorder &GetInstance() {
    static HostEventRecorder instance;
    return instance;
  }