Register subtensor info to tensor builder (#3519)

Register subtensor info to tensor builder if
- Tensor is subtensor of other tensor
- Backend support subtensor allocation

Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>

1 files changed, 31 insertions, 6 deletions

diff --git a/runtimes/neurun/src/linear/Linear.cc b/runtimes/neurun/src/linear/Linear.cc
index e2c33154e..59a00505d 100644
--- a/runtimes/neurun/src/linear/Linear.cc
+++ b/runtimes/neurun/src/linear/Linear.cc
@@ -21,6 +21,8 @@
 #include "graph/operation/LowerInfo.h"
 #include "backend/interface/IStageGenerator.h"
 #include "internal/Convert.h"
+#include "backend/interface/IConfig.h"
+#include "backend/common/operand/SubTensorInfo.h"
 
 #include "logging.h"
 
@@ -85,13 +87,34 @@ backend::TensorBuilderSet Linear::planTensors()
           uses_map[ind]++;
         }
 
-        // Prepare tensor builders to be returned
-        const auto info = ::internal::asTensorInfo(obj.shape(), obj.typeInfo());
-        iterTensorBuilders(ind, [&tensor_builders, &info](const graph::operand::Index &ind,
-                                                          ITensorBuilderPtr tensor_builder) {
-          tensor_builder->registerTensorInfo(ind, info);
+        for (auto backend : obj.lower_info()->def_backends())
+        {
+          bool isSubTensor = false;
+          auto tensor_builder = backend->tensor_builder();
+
+          if (backend->config()->SupportSubTensorAlloc())
+          {
+            const auto parentInfo = obj.parent_info();
+            if (parentInfo != nullptr)
+            {
+              isSubTensor = true;
+            }
+          }
+
+          if (isSubTensor)
+          {
+            const backend::operand::SubTensorInfo info(obj);
+            tensor_builder->registerSubTensorInfo(ind, info);
+          }
+          else
+          {
+            const auto info = ::internal::asTensorInfo(obj.shape(), obj.typeInfo());
+            tensor_builder->registerTensorInfo(ind, info);
+          }
+
+          // Prepare tensor builders to be returned
           tensor_builders.insert(tensor_builder);
-        });
+        }
       });
 
   // If a tensor is model output, increase the use of the tensor.
@@ -156,6 +179,8 @@ backend::TensorBuilderSet Linear::planTensors()
     assert(uses_map[ind] > 0);
 #endif
 
+  // Set subtensor information
+  // Todo: move this phase outside as optimization phase
   return tensor_builders;
 }