[Handler] Fix triv's weight offset setting

This patch fixes triv's weight offset setting.

Signed-off-by: Dongju Chae <dongju.chae@samsung.com>

diff --git a/src/core/ne-handler.cc b/src/core/ne-handler.cc
index 2eb46a6..0682587 100644 (file)
--- a/src/core/ne-handler.cc
+++ b/src/core/ne-handler.cc
@@ -956,7 +956,9 @@ TrinityVision::setModel (const generic_buffer *model_buf, Model ** model_ptr)
   if (model_buf == nullptr || model_ptr == nullptr)
     return -EINVAL;
 
-  Model *model;
+  Model *model = nullptr;
+  HWmem * hwmem_prog = nullptr;
+  HWmem * hwmem_weight = nullptr;
   int status;
 
   /** In TRIV1, model data (including program/weight) should be contiguous */
@@ -993,7 +995,7 @@ TrinityVision::setModel (const generic_buffer *model_buf, Model ** model_ptr)
 
   /** allocate program (optional; NOP) */
   if (model->getMetadata()->getProgramSize() > 0) {
-    HWmem * hwmem_prog = new HWmem (new HWmemChunk);
+    hwmem_prog = new HWmem (new HWmemChunk);
     model->setProgramData (hwmem_prog);
 
     hwmem_prog->setParent (model);
@@ -1003,23 +1005,11 @@ TrinityVision::setModel (const generic_buffer *model_buf, Model ** model_ptr)
       logerr (TAG, "Failed to allocate program\n");
       goto delete_exit;
     }
-
-    /** register this model to the driver */
-    model_config_t config;
-    config.dbuf_fd = hwmem_prog->getDmabuf ();
-    config.program_size = hwmem_prog->getSize ();
-    config.program_offset_addr = hwmem_prog->getOffset ();
-
-    status = api_->registerModel (&config);
-    if (status != 0)
-      goto delete_exit;
-
-    model->setInternalID(config.id);
   }
 
   /** allocate weight (optional) */
   if (model->getMetadata()->getWeightSize() > 0) {
-    HWmem * hwmem_weight = new HWmem (new HWmemChunk);
+    hwmem_weight = new HWmem (new HWmemChunk);
     model->setWeightData (hwmem_weight);
 
     hwmem_weight->setParent (model);
@@ -1032,6 +1022,22 @@ TrinityVision::setModel (const generic_buffer *model_buf, Model ** model_ptr)
     }
   }
 
+  if (hwmem_prog != nullptr) {
+    /** register this model to the driver */
+    model_config_t config;
+    config.dbuf_fd = hwmem_prog->getDmabuf ();
+    config.program_size = hwmem_prog->getSize ();
+    config.program_offset_addr = hwmem_prog->getOffset ();
+    if (hwmem_weight != nullptr)
+      config.weight_offset_addr = hwmem_weight->getOffset ();
+
+    status = api_->registerModel (&config);
+    if (status != 0)
+      goto delete_exit;
+
+    model->setInternalID(config.id);
+  }
+
   *model_ptr = model;
   return status;