profiling.h

Go to the documentation of this file.

/*
 * 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.
 */

#ifndef TVM_RUNTIME_PROFILING_H_
#define TVM_RUNTIME_PROFILING_H_

#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>

#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace tvm {

namespace runtime {

class TimerNode : public Object {
 public:
  virtual void Start() = 0;
  virtual void Stop() = 0;
  virtual int64_t SyncAndGetElapsedNanos() = 0;

  virtual ~TimerNode() {}

  static constexpr const char* _type_key = "TimerNode";
  TVM_DECLARE_BASE_OBJECT_INFO(TimerNode, Object);
};

class Timer : public ObjectRef {
 public:
  static TVM_DLL Timer Start(Device dev);

  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Timer, ObjectRef, TimerNode);
};

Timer DefaultTimer(Device dev);

namespace profiling {
struct DeviceWrapperNode : public Object {
  Device device;

  explicit DeviceWrapperNode(Device device) : device(device) {}

  static constexpr const char* _type_key = "runtime.profiling.DeviceWrapper";
  TVM_DECLARE_BASE_OBJECT_INFO(DeviceWrapperNode, Object);
};

class DeviceWrapper : public ObjectRef {
 public:
  explicit DeviceWrapper(Device dev) { data_ = make_object<DeviceWrapperNode>(dev); }
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(DeviceWrapper, ObjectRef, DeviceWrapperNode);
};

class ReportNode : public Object {
 public:
  Array<Map<String, ObjectRef>> calls;
  Map<String, Map<String, ObjectRef>> device_metrics;
  String AsCSV() const;
  String AsTable(bool sort = true, bool aggregate = true, bool compute_col_sums = true) const;
  String AsJSON() const;

  static constexpr const char* _type_key = "runtime.profiling.Report";
  TVM_DECLARE_FINAL_OBJECT_INFO(ReportNode, Object);
};

class Report : public ObjectRef {
 public:
  explicit Report(Array<Map<String, ObjectRef>> calls,
                  Map<String, Map<String, ObjectRef>> device_metrics);

  static Report FromJSON(String json);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Report, ObjectRef, ReportNode);
};

class MetricCollectorNode : public Object {
 public:
  virtual void Init(Array<DeviceWrapper> devs) = 0;
  virtual ObjectRef Start(Device dev) = 0;
  virtual Map<String, ObjectRef> Stop(ObjectRef obj) = 0;

  virtual ~MetricCollectorNode() {}

  static constexpr const char* _type_key = "runtime.profiling.MetricCollector";
  TVM_DECLARE_BASE_OBJECT_INFO(MetricCollectorNode, Object);
};

class MetricCollector : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(MetricCollector, ObjectRef, MetricCollectorNode);
};

struct CallFrame {
  Device dev;
  String name;
  Timer timer;
  std::unordered_map<std::string, ObjectRef> extra_metrics;
  std::vector<std::pair<MetricCollector, ObjectRef>> extra_collectors;
};

class Profiler {
 public:
  explicit Profiler(std::vector<Device> devs, std::vector<MetricCollector> metric_collectors);
  void Start();
  void Stop();
  void StartCall(String name, Device dev,
                 std::unordered_map<std::string, ObjectRef> extra_metrics = {});
  void StopCall(std::unordered_map<std::string, ObjectRef> extra_metrics = {});
  profiling::Report Report(bool aggregate = true, bool sort = true);
  bool IsRunning() const { return is_running_; }

 private:
  std::vector<Device> devs_;
  bool is_running_{false};
  std::vector<CallFrame> calls_;
  std::stack<CallFrame> in_flight_;
  std::vector<MetricCollector> collectors_;
};

/* \brief A duration in time. */
class DurationNode : public Object {
 public:
  /* The duration as a floating point number of microseconds. */
  double microseconds;

  /* \brief Construct a new duration.
   * \param a The duration in microseconds.
   */
  explicit DurationNode(double a) : microseconds(a) {}

  static constexpr const char* _type_key = "runtime.profiling.Duration";
  TVM_DECLARE_FINAL_OBJECT_INFO(DurationNode, Object);
};

/* A percentage of something */
class PercentNode : public Object {
 public:
  /* The percent as a floating point value out of 100%. i.e. if `percent` is 10 then we have 10%. */
  double percent;

  /* \brief Construct a new percentage.
   * \param a The percentage out of 100.
   */
  explicit PercentNode(double a) : percent(a) {}

  static constexpr const char* _type_key = "runtime.profiling.Percent";
  TVM_DECLARE_FINAL_OBJECT_INFO(PercentNode, Object);
};

/* A count of something */
class CountNode : public Object {
 public:
  /* The actual count */
  int64_t value;

  /* \brief Construct a new count.
   * \param a The count.
   */
  explicit CountNode(int64_t a) : value(a) {}

  static constexpr const char* _type_key = "runtime.profiling.Count";
  TVM_DECLARE_FINAL_OBJECT_INFO(CountNode, Object);
};

String ShapeString(const std::vector<NDArray>& shapes);
String ShapeString(NDArray shape, DLDataType dtype);
String ShapeString(const std::vector<int64_t>& shape, DLDataType dtype);

}  // namespace profiling
}  // namespace runtime
}  // namespace tvm

#endif  // TVM_RUNTIME_PROFILING_H_