Deliver Glow traces through onnxifi (#2678)

pytorch · Apr 17, 2019 · 3b2b8ad · 3b2b8ad
1 parent c2f7786
commit 3b2b8ad
Show file tree

Hide file tree

Showing 8 changed files with 88 additions and 17 deletions.
diff --git a/include/glow/Backends/TraceEvents.h b/include/glow/Backends/TraceEvents.h
@@ -142,6 +142,9 @@ class TraceContext {
   /// \returns TraceEvents for the last run.
   std::vector<TraceEvent> &getTraceEvents() { return traceEvents_; }
 
+  /// \returns TraceEvents for the last run.
+  llvm::ArrayRef<TraceEvent> getTraceEvents() const { return traceEvents_; }
+
   /// \returns the integer thread id used for logged events in this context.
   int getTraceThread() const { return traceThread_; }
 

diff --git a/lib/Onnxifi/Base.cpp b/lib/Onnxifi/Base.cpp
@@ -94,9 +94,16 @@ onnxStatus Graph::setIOAndRun(uint32_t inputsCount,
                               const onnxTensorDescriptorV1 *inputDescriptors,
                               uint32_t outputsCount,
                               const onnxTensorDescriptorV1 *outputDescriptors,
-                              EventPtr outputEvent) {
+                              EventPtr outputEvent,
+                              onnxTraceEventList *traceEvents) {
   auto ctx = llvm::make_unique<ExecutionContext>();
 
+  if (traceEvents) {
+    // TODO: get thread ID
+    ctx->setTraceContext(
+        llvm::make_unique<TraceContext>(TraceLevel::STANDARD, 0));
+  }
+
   // Create tensors for input placeholders
   for (unsigned i = 0; i < inputsCount; ++i) {
     const auto &inOnnxTensor = inputDescriptors[i];
@@ -161,7 +168,48 @@ onnxStatus Graph::setIOAndRun(uint32_t inputsCount,
     ctx->getPlaceholderBindings()->insert(outPhPtr, std::move(t));
   }
 
-  return run(std::move(ctx), outputEvent, std::move(phNameToOnnxTensorOutputs));
+  return run(std::move(ctx), outputEvent, std::move(phNameToOnnxTensorOutputs),
+             traceEvents);
+}
+
+void Graph::setTraceEvents(onnxTraceEventList *traceEvents,
+                           const TraceContext &traceContext) {
+  if (!traceEvents) {
+    return;
+  }
+
+  std::vector<onnxTraceEvent *> traceEventsVec;
+  for (const auto &glowTraceEvent : traceContext.getTraceEvents()) {
+    auto *traceEvent = new onnxTraceEvent();
+    assert(
+        glowTraceEvent.type.size() == 1 &&
+        "Events with types longer than a single char not supported by onnxifi");
+    traceEvent->eventType = glowTraceEvent.type[0];
+    traceEvent->timestamp = glowTraceEvent.timestamp;
+    traceEvent->tid = glowTraceEvent.tid;
+    char *eventName = new char[glowTraceEvent.name.size() + 1];
+    assert(eventName);
+    strcpy(eventName, glowTraceEvent.name.c_str());
+    traceEvent->eventName = eventName;
+    traceEventsVec.push_back(traceEvent);
+  }
+
+  traceEvents->numEvents = traceEventsVec.size();
+  traceEvents->traceEvents = new onnxTraceEvent *[traceEventsVec.size()];
+  assert(traceEvents->traceEvents);
+  std::copy(traceEventsVec.begin(), traceEventsVec.end(),
+            traceEvents->traceEvents);
+}
+
+void Graph::releaseTraceEvents(onnxTraceEventList *traceEvents) {
+  assert(traceEvents);
+  for (uint64_t i = 0; i < traceEvents->numEvents; ++i) {
+    onnxTraceEvent *traceEvent = traceEvents->traceEvents[i];
+    delete[] traceEvent->eventName;
+    delete traceEvent;
+  }
+
+  delete[] traceEvents->traceEvents;
 }
 
 Graph::Graph(BackendPtr backendPtr) : backendPtr_(backendPtr) {}

diff --git a/lib/Onnxifi/Base.h b/lib/Onnxifi/Base.h
@@ -115,12 +115,14 @@ class Graph {
   /// Set input memory addresses for inputs based on the \p inputDescriptors.
   /// Set output memory addresses for outputs based on the \p
   /// outputDescriptors. Will async signal the \p outputEvent when run is
-  /// complete.
+  /// complete. \p traceEvents is a pointer to onnxTraceEventList, if it is not
+  /// null then it is expected that this will be populated with trace events
+  /// from the run before signalling the outputEvent.
   onnxStatus setIOAndRun(uint32_t inputsCount,
                          const onnxTensorDescriptorV1 *inputDescriptors,
                          uint32_t outputsCount,
                          const onnxTensorDescriptorV1 *outputDescriptors,
-                         EventPtr outputEvent);
+                         EventPtr outputEvent, onnxTraceEventList *traceEvents);
 
   /// Init Glow graph based on the ONNX model \p onnxModel and
   /// static trained weights \p weightDescriptors.
@@ -131,7 +133,17 @@ class Graph {
   virtual onnxStatus
   run(std::unique_ptr<ExecutionContext> ctx, EventPtr outputEvent,
       std::unordered_map<Placeholder *, onnxTensorDescriptorV1>
-          phNameToOnnxTensorOutputs) = 0;
+          phNameToOnnxTensorOutputs,
+      onnxTraceEventList *traceEvents) = 0;
+
+  /// Copy any trace events \p traceContext into \p traceEvents. If
+  /// \p traceEvents is null then do nothing.
+  static void setTraceEvents(onnxTraceEventList *traceEvents,
+                             const TraceContext &traceContext);
+
+  /// Free all memory that was allocated by setTraceEvents when creating \p
+  /// traceEvents.
+  static void releaseTraceEvents(onnxTraceEventList *traceEvents);
 
 protected:
   BackendPtr backendPtr_;

diff --git a/lib/Onnxifi/HostManagerOnnxifi.cpp b/lib/Onnxifi/HostManagerOnnxifi.cpp
@@ -78,11 +78,13 @@ onnxStatus
 HostManagerGraph::run(std::unique_ptr<ExecutionContext> ctx,
                       EventPtr outputEvent,
                       std::unordered_map<Placeholder *, onnxTensorDescriptorV1>
-                          phNameToOnnxTensorOutputs) {
+                          phNameToOnnxTensorOutputs,
+                      onnxTraceEventList *traceEvents) {
   backendPtr_->getBackendId()->runNetwork(
       this, std::move(ctx),
       [phNameToOnnxTensorOutputs = std::move(phNameToOnnxTensorOutputs),
-       outputEvent](runtime::RunIdentifierTy runId, llvm::Error err,
+       outputEvent,
+       traceEvents](runtime::RunIdentifierTy runId, llvm::Error err,
                     std::unique_ptr<ExecutionContext> ctx) {
         // If an Error occurred then log it in errToBool and signal the output
         // event.
@@ -104,6 +106,10 @@ HostManagerGraph::run(std::unique_ptr<ExecutionContext> ctx,
                  res->size() * res->getType().getElementSize());
         }
 
+        if (auto *traceContext = ctx->getTraceContext()) {
+          setTraceEvents(traceEvents, *traceContext);
+        }
+
         outputEvent->signal();
       });
 

diff --git a/lib/Onnxifi/HostManagerOnnxifi.h b/lib/Onnxifi/HostManagerOnnxifi.h
@@ -69,7 +69,8 @@ class HostManagerGraph : public Graph {
   /// copy output placeholder tensors back to the given onnxifi tensors.
   onnxStatus run(std::unique_ptr<ExecutionContext> ctx, EventPtr outputEvent,
                  std::unordered_map<Placeholder *, onnxTensorDescriptorV1>
-                     phNameToOnnxTensorOutputs) override;
+                     phNameToOnnxTensorOutputs,
+                 onnxTraceEventList *traceEvents) override;
 
   /// \returns the unique string name of the HostManagerGraph that the
   /// underlying HostManagerGraph uses to identify this network.

diff --git a/lib/Onnxifi/InlineOnnxifi.cpp b/lib/Onnxifi/InlineOnnxifi.cpp
@@ -84,8 +84,8 @@ InlineGraph::initGraph(const void *onnxModel, size_t onnxModelSize,
 onnxStatus
 InlineGraph::run(std::unique_ptr<ExecutionContext> ctx, EventPtr outputEvent,
                  std::unordered_map<Placeholder *, onnxTensorDescriptorV1>
-                     phNameToOnnxTensorOutputs) {
-
+                     phNameToOnnxTensorOutputs,
+                 onnxTraceEventList *traceEvents) {
   executionEngine_.run(*ctx);
 
   // Dump profile if requested.
@@ -97,6 +97,10 @@ InlineGraph::run(std::unique_ptr<ExecutionContext> ctx, EventPtr outputEvent,
     serializeToYaml(getProfileFile(modelHash_), QI);
   }
 
+  if (auto *traceContext = ctx->getTraceContext()) {
+    setTraceEvents(traceEvents, *traceContext);
+  }
+
   outputEvent->signal();
   return ONNXIFI_STATUS_SUCCESS;
 }

diff --git a/lib/Onnxifi/InlineOnnxifi.h b/lib/Onnxifi/InlineOnnxifi.h
@@ -44,7 +44,8 @@ class InlineGraph : public Graph {
 
   onnxStatus run(std::unique_ptr<ExecutionContext> ctx, EventPtr outputEvent,
                  std::unordered_map<Placeholder *, onnxTensorDescriptorV1>
-                     phNameToOnnxTensorOutputs) override;
+                     phNameToOnnxTensorOutputs,
+                 onnxTraceEventList *traceEvents) override;
 
 private:
   ExecutionEngine executionEngine_;

diff --git a/lib/Onnxifi/onnxifiGlow.cpp b/lib/Onnxifi/onnxifiGlow.cpp
@@ -456,10 +456,6 @@ GLOW_ONNXIFI_LIBRARY_FUNCTION_WRAPPER(onnxSetIOAndRunGraph)(
     const onnxTensorDescriptorV1 *inputDescriptors, uint32_t outputsCount,
     const onnxTensorDescriptorV1 *outputDescriptors,
     onnxMemoryFenceV1 *outputFence, onnxTraceEventList *traceEvents) {
-  if (traceEvents) {
-    llvm::errs() << "Glow doesn't support tracing yet\n";
-  }
-
   auto &manager = glow::onnxifi::GlowOnnxifiManager::get();
 
   if (!inputDescriptors || !outputDescriptors || !outputFence) {
@@ -501,7 +497,7 @@ GLOW_ONNXIFI_LIBRARY_FUNCTION_WRAPPER(onnxSetIOAndRunGraph)(
 
   // Set graph IO and run async
   return glowGraph->setIOAndRun(inputsCount, inputDescriptors, outputsCount,
-                                outputDescriptors, outputEvent);
+                                outputDescriptors, outputEvent, traceEvents);
 }
 
 /// Deinitialize an ONNXIFI graph and release associated resources.
@@ -527,7 +523,7 @@ GLOW_ONNXIFI_LIBRARY_FUNCTION_WRAPPER(onnxReleaseTraceEvents)(
   if (!traceEvents) {
     return ONNXIFI_STATUS_INVALID_POINTER;
   }
-  llvm::errs() << "onnxReleaseTraceEvents not implemented\n";
+  glow::onnxifi::Graph::releaseTraceEvents(traceEvents);
   return ONNXIFI_STATUS_INTERNAL_ERROR;
 }