diff options
Diffstat (limited to 'runtimes/neurun')
| -rw-r--r-- | runtimes/neurun/src/linear/Linear.cc | 37 |
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; } |