From 446a5b5ec606a6eb9d4db28f04964d33af99e6a4 Mon Sep 17 00:00:00 2001
From: Nick Gibson <nickg@fb.com>
Date: Wed, 1 May 2019 16:04:14 -0700
Subject: [PATCH] Add TraceEvents to Habana backend and simplify (#2822)

---
 lib/Backends/Habana/Habana.cpp              | 14 ++++++++++++++
 lib/Backends/Habana/Habana.h                |  2 +-
 lib/Backends/Habana/HabanaDeviceManager.cpp |  5 +++++
 lib/Backends/Habana/HabanaDeviceManager.h   |  4 ++--
 4 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/lib/Backends/Habana/Habana.cpp b/lib/Backends/Habana/Habana.cpp
index 86deea6176..faf95c0ac6 100644
--- a/lib/Backends/Habana/Habana.cpp
+++ b/lib/Backends/Habana/Habana.cpp
@@ -442,6 +442,9 @@ void HabanaFunction::setupRuns() {}
 void HabanaFunction::beforeRun(const PlaceholderBindings &ctx) {}
 
 void HabanaFunction::execute(ExecutionContext *context) {
+  auto *tc = context->getTraceContext();
+  TRACE_EVENT_BEGIN(tc, "execute");
+
   uint32_t deviceId =
       static_cast<HabanaBindings *>(context->getDeviceBindings())
           ->getDeviceId();
@@ -456,6 +459,7 @@ void HabanaFunction::execute(ExecutionContext *context) {
   std::vector<EnqueueTensorInfo> outputInfo;
 
   // Set up input buffers and record bindings for enqueuing.
+  TRACE_EVENT_BEGIN(tc, "copyInputs");
   auto *bindings = context->getPlaceholderBindings();
   for (auto *P : getInputs()) {
     Tensor *T = bindings->get(P);
@@ -474,7 +478,9 @@ void HabanaFunction::execute(ExecutionContext *context) {
     // Copy from the tensor into the designated IO buffer.
     memcpy(eti.pTensorData, T->getUnsafePtr(), eti.tensorSize);
   }
+  TRACE_EVENT_END(tc, "copyInputs");
 
+  TRACE_EVENT_BEGIN(tc, "registerOutputs");
   // Set up output buffers and record bindings for enqueuing.
   for (auto *P : getOutputs()) {
     Tensor *T = bindings->get(P);
@@ -493,21 +499,29 @@ void HabanaFunction::execute(ExecutionContext *context) {
   }
 
   EnqueueTensorInfo noInputEti = {"unused", (char *)nullptr, 0};
+  TRACE_EVENT_END(tc, "registerOutputs");
 
   // Enqueue the run and wait for it to come back.
   synWaitHandle handle;
   {
     // Activate and enqueue need to be atomic.
+    TRACE_EVENT_BEGIN(tc, "getSynapseLock");
     std::lock_guard<std::mutex> g(synapseLock);
+    TRACE_EVENT_END(tc, "getSynapseLock");
+    TRACE_EVENT_BEGIN(tc, "synActivateTopology");
     chk(synActivateTopology(deviceId, topologyId));
+    TRACE_EVENT_END(tc, "synActivateTopology");
+    TRACE_EVENT_BEGIN(tc, "synEnqueue");
     chk(synEnqueueByName(
         deviceId, inputInfo.empty() ? &noInputEti : inputInfo.data(),
         inputInfo.size(), outputInfo.data(), outputInfo.size(), &handle));
+    TRACE_EVENT_END(tc, "synEnqueue");
   }
 
   static_cast<HabanaBindings *>(context->getDeviceBindings())
       ->setHandle(HabanaWaitHandle(deviceId, handle, std::move(inputInfo),
                                    std::move(outputInfo)));
+  TRACE_EVENT_END(tc, "execute");
 }
 
 void HabanaFunction::afterRun(const PlaceholderBindings &ctx) {}
diff --git a/lib/Backends/Habana/Habana.h b/lib/Backends/Habana/Habana.h
index ee2ea0f2c2..be8926e417 100644
--- a/lib/Backends/Habana/Habana.h
+++ b/lib/Backends/Habana/Habana.h
@@ -108,7 +108,7 @@ class HabanaIOBufferPool {
   uint8_t *buffer_;
   /// The number of buffers in the pool.
   unsigned numBuffers_{kDefaultNumBuffers};
-  static constexpr unsigned kDefaultNumBuffers{3};
+  static constexpr unsigned kDefaultNumBuffers{10};
 
   /// Queue of HabanaIOBuffers and a mutex and condition variable for
   /// synchronized access.
diff --git a/lib/Backends/Habana/HabanaDeviceManager.cpp b/lib/Backends/Habana/HabanaDeviceManager.cpp
index aaffdedae9..9f521f54e8 100644
--- a/lib/Backends/Habana/HabanaDeviceManager.cpp
+++ b/lib/Backends/Habana/HabanaDeviceManager.cpp
@@ -177,6 +177,9 @@ void HabanaDeviceManager::addNetwork(const Module *module,
       readyCB(module, MAKE_ERR("Unable to add function"));
       return;
     }
+
+    // Optimistically activate the topology
+    synActivateTopology(deviceId_, topologyId);
   }
 
   lk.unlock();
@@ -292,11 +295,13 @@ void HabanaDeviceManager::runFunctionImpl(RunIdentifierTy runId,
                      functionName = std::move(functionName),
                      ctx = std::move(ctx),
                      resultCB = std::move(resultCB)]() mutable {
+    TRACE_EVENT_BEGIN(ctx->getTraceContext(), "wait");
     auto &habanaHandle =
         static_cast<HabanaBindings *>(ctx->getDeviceBindings())->getHandle();
     bool ok = habanaHandle.wait();
     std::unique_ptr<HabanaIOBuffer> ioBuffer =
         static_cast<HabanaBindings *>(ctx->getDeviceBindings())->getIOBuffer();
+    TRACE_EVENT_END(ctx->getTraceContext(), "wait");
 
     // Notify anything waiting for a topo switch.
     {
diff --git a/lib/Backends/Habana/HabanaDeviceManager.h b/lib/Backends/Habana/HabanaDeviceManager.h
index d7c8472e87..db23fd2587 100644
--- a/lib/Backends/Habana/HabanaDeviceManager.h
+++ b/lib/Backends/Habana/HabanaDeviceManager.h
@@ -49,9 +49,9 @@ class HabanaDeviceManager : public DeviceManager {
   /// Thread pool for waiting on the results of executing functions.
   std::unique_ptr<ThreadPool> waitPool_;
   /// The default number of workers in run pool (overridable).
-  constexpr static unsigned kNumRunners = 3;
+  constexpr static unsigned kNumRunners = 1;
   /// The default number of workers in wait pool (overridable).
-  constexpr static unsigned kNumWaiters = 3;
+  constexpr static unsigned kNumWaiters = 1;
   /// The number of workers in run pool.
   unsigned numRunners_{kNumRunners};
   /// The number of workers in wait pool.