#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"
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.
assert(uses_map[ind] > 0);
#endif
+ // Set subtensor information
+ // Todo: move this phase outside as optimization phase
return tensor_builders;
}