callback function definition changes, and handle the error raised by poplar SDK

Author: jackzipu

ODLA/platforms/odla_popart/odla_compute.cc

@@ -86,7 +86,7 @@ odla_status odla_SetContextItem(odla_context context, odla_item_type type,
   switch (type) {
     case ODLA_ASYNC_CALLBACK_FUNC:
       context->async_callback_func =
-          reinterpret_cast<void (*)(void*, odla_status)>(value);
+          reinterpret_cast<int (*)(void*, odla_status)>(value);
       break;
     case ODLA_ASYNC_CALLBACK_ARG:
       context->async_callback_arg = reinterpret_cast<void*>(value);

ODLA/platforms/odla_popart/odla_pipeline.cc

@@ -195,6 +195,16 @@ odla_context ContextQueues::get_ctx_by_tensor(const popart::TensorId& id) {
   tensor_to_idx_[id] = (idx + 1) % capacity_;
   return ctx;
 }
+
+void ContextQueues::handle_error() {
+  auto idx = wait_;
+  while (idx != tail_) {
+    odla_context ctx = buffer_[idx];
+    if (ctx != nullptr) ctx->notify();
+    idx = (idx + 1) % capacity_;
+  }
+}
+
 /*------------------------------------------------------------------------*/
 LockFreeQueue::LockFreeQueue() : head_(0), tail_(0), wait_(0) {}
 
@@ -344,6 +354,15 @@ odla_context LockFreeQueue::get_ctx_by_tensor(const popart::TensorId& id) {
   return ctx;
 }
 
+void LockFreeQueue::handle_error() {
+  auto idx = wait_;
+  while (idx != tail_.load()) {
+    odla_context ctx = buffer_[idx].load();
+    if (ctx != nullptr) ctx->notify();
+    idx = (idx + 1) % capacity_;
+  }
+}
+
 /*======================================== step io callbacks
  * =========================================*/
 popart::StepIOCallback::InputCallback input_callback =

ODLA/platforms/odla_popart/odla_pipeline.h

@@ -45,6 +45,7 @@ class Queue {
   virtual void pop_input(odla_context ctx) = 0;
   virtual void pop_output(odla_context ctx) = 0;
   virtual std::size_t size() = 0;
+  virtual void handle_error() = 0;
 };
 
 class ContextQueues : public Queue {
@@ -72,6 +73,7 @@ class ContextQueues : public Queue {
   void pop_input(odla_context ctx) final;
   void pop_output(odla_context ctx) final;
   std::size_t size() final { return (tail_ - wait_ + capacity_) % capacity_; }
+  void handle_error() final;
 };
 
 class LockFreeQueue : public Queue {
@@ -100,6 +102,7 @@ class LockFreeQueue : public Queue {
   std::size_t size() final {
     return (tail_.load() - wait_ + capacity_) % capacity_;
   }
+  void handle_error() final;
 };
 
 class QManager {
@@ -115,7 +118,10 @@ class QManager {
   void createQ(std::string queueType);
   void deleteQ();
   inline Queue* getQ() { return queue_; }
-  inline void set_status(odla_status status) { status_ = status; }
+  inline void set_status(odla_status status) {
+    status_ = status;
+    if (ODLA_SUCCESS != status_ && queue_) queue_->handle_error();
+  }
   inline odla_status get_status() { return status_; }
   static inline QManager* instance() { return instance_; }
 };
@@ -216,8 +222,7 @@ struct _odla_pipeline_async_context : public _odla_pipeline_context {
           this);
       throw std::invalid_argument("async_callback_arg is null");
     }
-    async_callback_func(async_callback_arg,
-                        ODLA_SUCCESS); // FIXME: notify the status
+    async_callback_func(async_callback_arg, QManager::instance()->get_status());
   }
   bool hold(const std::string& function_name) override { return true; }
 };

ODLA/platforms/odla_popart/odla_popart.h

@@ -176,7 +176,7 @@ struct _odla_context {
   odla_computation comp;
   std::map<popart::TensorId, std::unique_ptr<popart::IArray>> inputs;
   std::map<popart::TensorId, std::unique_ptr<popart::IArray>> outputs;
-  void (*async_callback_func)(void*, odla_status) = nullptr;
+  int (*async_callback_func)(void*, odla_status) = nullptr;
   void* async_callback_arg = nullptr;
   _odla_context(odla_computation c) : comp(c) {}
   std::thread::id thread_id_of_holder;