(tvm) branch main updated: [Metal] Support metal device profiling (#17025)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 7463b37b88 [Metal] Support metal device profiling (#17025)
7463b37b88 is described below

commit 7463b37b88b488bf1cf8696632765c51760fe3be
Author: Siyuan Feng <hzfen...@sjtu.edu.cn>
AuthorDate: Fri May 24 18:51:27 2024 +0800

    [Metal] Support metal device profiling (#17025)
    
    Enable native metal device profiling through API `sampleTimestamps`
---
 src/runtime/metal/metal_device_api.mm | 37 +++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/src/runtime/metal/metal_device_api.mm 
b/src/runtime/metal/metal_device_api.mm
index 37fb9dc347..42dd249630 100644
--- a/src/runtime/metal/metal_device_api.mm
+++ b/src/runtime/metal/metal_device_api.mm
@@ -21,6 +21,7 @@
  * \file metal_device_api.mm
  */
 #include <dmlc/thread_local.h>
+#include <tvm/runtime/profiling.h>
 #include <tvm/runtime/registry.h>
 #include "metal_common.h"
 
@@ -366,6 +367,42 @@ 
TVM_REGISTER_GLOBAL("metal.ResetGlobalState").set_body_typed([]() {
   MetalWorkspace::Global()->ReinitializeDefaultStreams();
 });
 
+class MetalTimerNode : public TimerNode {
+ public:
+  MetalTimerNode() {}
+  explicit MetalTimerNode(Device dev) : dev_(dev) {
+    mtl_dev_ = MetalWorkspace::Global()->GetDevice(dev_);
+  }
+
+  virtual void Start() {
+    [mtl_dev_ sampleTimestamps:&start_cpu_time_ gpuTimestamp:&start_gpu_time_];
+  }
+  virtual void Stop() {
+    auto ws = MetalWorkspace::Global();
+    ws->StreamSync(dev_, ws->GetCurrentStream(dev_));
+    [mtl_dev_ sampleTimestamps:&stop_cpu_time_ gpuTimestamp:&stop_gpu_time_];
+  }
+  virtual int64_t SyncAndGetElapsedNanos() { return stop_gpu_time_ - 
start_gpu_time_; }
+
+  static constexpr const char* _type_key = "MetalTimerNode";
+  TVM_DECLARE_FINAL_OBJECT_INFO(MetalTimerNode, TimerNode);
+
+ private:
+  Device dev_;
+  id<MTLDevice> mtl_dev_;
+
+  MTLTimestamp start_cpu_time_;
+  MTLTimestamp start_gpu_time_;
+  MTLTimestamp stop_cpu_time_;
+  MTLTimestamp stop_gpu_time_;
+};
+
+TVM_REGISTER_OBJECT_TYPE(MetalTimerNode);
+
+TVM_REGISTER_GLOBAL("profiling.timer.metal").set_body_typed([](Device dev) {
+  return Timer(make_object<MetalTimerNode>(dev));
+});
+
 }  // namespace metal
 }  // namespace runtime
 }  // namespace tvm