Access profiler from cpp

test/cpp/jit/no-gtest.cpp

@@ -8,6 +8,7 @@ namespace torch {
 namespace jit {
 std::string runJITCPPTests() {
   std::stringstream out;
+  testAutogradProfiler();
   testADFormulas();
   testArgumentSpec();
   testAttributes();

test/cpp/jit/test_misc.h

@@ -1751,6 +1751,36 @@ void testDynamicDAG() {
   testContractEdgeBasic();
   testContractEdgeCycleDetection();
 }
+
+void testAutogradProfiler() {
+  constexpr int batch_size = 4;
+  constexpr int input_size = 256;
+  constexpr int seq_len = 32;
+
+  int hidden_size = 2 * input_size;
+  auto input = torch::randn({seq_len, batch_size, input_size}, at::kCPU);
+  auto hx = torch::randn({batch_size, hidden_size}, at::kCPU);
+  auto cx = torch::randn({batch_size, hidden_size}, at::kCPU);
+  auto w_ih = t_def(torch::randn({4 * hidden_size, input_size}, at::kCPU));
+  auto w_hh = t_def(torch::randn({4 * hidden_size, hidden_size}, at::kCPU));
+
+  std::stringstream ss;
+  {
+    autograd::profiler::RecordProfile guard(ss);
+    for (size_t i = 0; i < 100; ++i) {
+      std::tie(hx, cx) = lstm(input[0], hx, cx, w_ih, w_hh);
+    }
+  }
+
+  std::string result = ss.str();
+  size_t count = 0;
+  for (size_t pos = 0; (pos = result.find("tanh", pos)) != std::string::npos;
+       count++, pos++) {
+  }
+  AT_CHECK(count == 200);
+}
+
+
 } // namespace
 } // namespace jit
 } // namespace torch

torch/csrc/autograd/profiler.cpp

@@ -2,6 +2,7 @@
 #include <torch/csrc/autograd/function.h>
 
 #include <sstream>
+#include <fstream>
 
 namespace torch { namespace autograd { namespace profiler {
 
@@ -192,4 +193,79 @@ double Event::cuda_elapsed_us(const Event & e) {
 
 CUDAStubs::~CUDAStubs() = default;
 
+
+static jit::CodeTemplate event_template(R"(
+{
+  "name": "${name}",
+  "ph": "X",
+  "ts": ${ts},
+  "dur": ${dur},
+  "tid": ${tid},
+  "pid": "CPU Functions",
+  "args": {}
+})");
+
+
+RecordProfile::RecordProfile(std::ostream& out)
+: out_(out) {
+  init();
+}
+
+RecordProfile::RecordProfile(const std::string& filename)
+: file_(new std::ofstream(filename)), out_(*file_) {
+  init();
+}
+
+void RecordProfile::init() {
+  enableProfiler(ProfilerState::CPU);
+}
+
+RecordProfile::~RecordProfile() {
+  thread_event_lists event_lists = disableProfiler();
+  std::vector<Event*> events;
+  for(auto& l : event_lists) {
+    for(auto& e : l) {
+        events.push_back(&e);
+    }
+  }
+  processEvents(events);
+  if (file_){
+    file_->close();
+  }
+}
+
+void RecordProfile::processEvents(const std::vector<Event*>& events) {
+  AT_CHECK(out_, "could not open file");
+  Event* start = nullptr;
+  for (Event* e : events) {
+    if(0 == strcmp(e->name(), "__start_profile")) {
+      start = e;
+      break;
+    }
+  }
+  AT_CHECK(start, "could not find start?");
+  std::vector<Event*> stack;
+  out_ << "[\n";
+  bool first = true;
+  for(Event* e : events) {
+    if(e->kind() == "push") {
+      stack.push_back(e);
+    } else if(e->kind() == "pop") {
+      if(!first) {
+        out_ << ",\n";
+      }
+      first = false;
+      Event* e_start = stack.back();
+      stack.pop_back();
+      jit::TemplateEnv env;
+      env.s("name", e_start->name());
+      env.d("ts", start->cpu_elapsed_us(*e_start));
+      env.d("dur", e_start->cpu_elapsed_us(*e));
+      env.d("tid", e_start->thread_id());
+      out_ << event_template.format(env);
+    }
+  }
+  out_ << "]\n";
+}
+
 }}}

torch/csrc/autograd/profiler.h

@@ -17,6 +17,8 @@
 #include <ctime>
 #endif
 
+#include <torch/csrc/jit/code_template.h>
+
 typedef struct CUevent_st* CUDAEventStub;
 
 namespace torch { namespace autograd {
@@ -221,5 +223,25 @@ using thread_event_lists = std::vector<std::vector<Event>>;
 TORCH_API void enableProfiler(ProfilerState new_state);
 TORCH_API thread_event_lists disableProfiler();
 
+
+// Usage:
+//   {
+//     RecordProfile guard("filename.trace");
+//     // code you want to profile
+//   }
+// Then open filename.trace in chrome://tracing
+struct TORCH_API RecordProfile {
+  RecordProfile(std::ostream& out);
+  RecordProfile(const std::string& filename);
+
+  ~RecordProfile();
+private:
+  void init();
+  std::unique_ptr<std::ofstream> file_;
+  std::ostream& out_;
+  void processEvents(const std::vector<Event*>& events);
+};
+
+
 } // namespace profiler
 }} // namespace torch::autograd