ExecutionProvider API refactor - Deattach allocator from EP by creating local cpu allocator instead (#16084)

### Description
ExecutionProvider API refactor - Detach allocator from EP by creating
local cpu allocator instead



### Motivation and Context
This is PR is a refactor to create local CPU allocator instead of
getting allocator from ExecutionProvider, which the final goal is to
totally detach allocators from ExecutionProvider, and put them in
session level indexed by OrtDevice
This commit is contained in:
cao lei 2023-05-26 04:54:42 -07:00 committed by GitHub
parent 4bfb8d3303
commit 4ab7d410ae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 25 additions and 66 deletions

View file

@ -234,7 +234,9 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
// run TransposeOptimizer last as it works in a slightly different way by moving Transpose nodes around.
// shouldn't affect the end result - just easier to debug any issue if it's last.
auto cpu_allocator = cpu_execution_provider.GetAllocator(OrtMemTypeDefault);
// local CPU allocator is enough as this allocator is finally passed to a local tensor.
// We will also benefit by using a local allocator as we don't need to pass allocator as parameter for EP API refactor
AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
transformers.emplace_back(std::make_unique<TransposeOptimizer>(std::move(cpu_allocator)));
} break;
@ -349,7 +351,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
if (MlasNchwcGetBlockSize() > 1) {
transformers.emplace_back(std::make_unique<NchwcTransformer>());
}
auto cpu_allocator = cpu_execution_provider.GetAllocator(OrtMemTypeDefault);
AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
auto cpu_registry = cpu_execution_provider.GetKernelRegistry();
auto nhwc_transformer = std::make_unique<NhwcTransformer>(std::move(cpu_allocator), std::move(cpu_registry));
if (nhwc_transformer->IsActive()) {
@ -425,7 +427,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformersForMinimalB
// currently the only level 3 optimizer is the NhwcTransformer which is fully supported at runtime
if (!saving) {
#ifndef DISABLE_CONTRIB_OPS
auto cpu_allocator = cpu_execution_provider.GetAllocator(OrtMemTypeDefault);
AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
auto cpu_registry = cpu_execution_provider.GetKernelRegistry();
auto nhwc_transformer = std::make_unique<NhwcTransformer>(std::move(cpu_allocator), std::move(cpu_registry));
if (nhwc_transformer->IsActive()) {

View file

@ -35,7 +35,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
const std::function<bool(const std::string&)>& is_sparse_initializer_func)
: execution_provider_(execution_provider),
is_sparse_initializer_func_(is_sparse_initializer_func) {
allocator_ptr_ = execution_provider_.GetAllocator(mem_type_);
allocator_ptr_ = std::make_shared<CPUAllocator>();
ORT_ENFORCE(allocator_ptr_, "Failed to get allocator for optimizer");
ORT_THROW_IF_ERROR(data_transfer_mgr_.RegisterDataTransfer(std::make_unique<CPUDataTransfer>()));
@ -89,7 +89,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
const std::function<bool(const std::string&)>& is_sparse_initializer_func)
: execution_provider_(execution_provider),
is_sparse_initializer_func_(is_sparse_initializer_func) {
allocator_ptr_ = execution_provider_.GetAllocator(mem_type_);
allocator_ptr_ = std::make_shared<CPUAllocator>();
ORT_ENFORCE(allocator_ptr_, "Failed to get allocator for optimizer");
ORT_THROW_IF_ERROR(data_transfer_mgr_.RegisterDataTransfer(std::make_unique<CPUDataTransfer>()));

View file

@ -34,10 +34,6 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
const std::function<bool(const std::string&)>& is_sparse_initializer_func);
~Info() = default;
AllocatorPtr GetAllocator(const OrtMemoryInfo& info) const {
return execution_provider_.GetAllocator(info.mem_type);
}
const AllocatorPtr& GetAllocator() const {
return allocator_ptr_;
}
@ -68,7 +64,6 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
}
private:
const OrtMemType mem_type_{OrtMemTypeDefault};
AllocatorPtr allocator_ptr_;
DataTransferManager data_transfer_mgr_;
// MLValues for optimizer

View file

@ -258,11 +258,11 @@ Status BaseOpBuilder::ProcessOutputs(QnnModelWrapper& qnn_model_wrapper,
Status BaseOpBuilder::TransposeInitializer(const QnnModelWrapper& qnn_model_wrapper,
const onnx::TensorProto& initializer,
const std::vector<size_t>& perm,
const AllocatorPtr& cpu_allocator,
std::vector<uint8_t>& transposed_data) const {
const DataTypeImpl* tensor_dtype = DataTypeImpl::TensorTypeFromONNXEnum(initializer.data_type())->GetElementType();
const auto tensor_shape_dims = onnxruntime::utils::GetTensorShapeFromTensorProto(initializer);
TensorShape tensor_shape{tensor_shape_dims};
AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
Tensor in_tensor = Tensor(tensor_dtype, tensor_shape, cpu_allocator);
auto rank = perm.size();

View file

@ -218,33 +218,29 @@ class BaseOpBuilder : public IOpBuilder {
Status TransposeInitializer(const QnnModelWrapper& qnn_model_wrapper,
const onnx::TensorProto& initializer,
const std::vector<size_t>& perm,
const AllocatorPtr& cpu_allocator,
std::vector<uint8_t>& transposed_data) const;
Status TransposeFromNchwToHwcn(const QnnModelWrapper& qnn_model_wrapper,
const onnx::TensorProto& initializer,
const AllocatorPtr& cpu_allocator,
std::vector<uint8_t>& transposed_data) const {
return TransposeInitializer(qnn_model_wrapper, initializer, nchw2hwcn_perm, cpu_allocator, transposed_data);
return TransposeInitializer(qnn_model_wrapper, initializer, nchw2hwcn_perm, transposed_data);
}
Status TransposeFromCnhwToHwcn(const QnnModelWrapper& qnn_model_wrapper,
const onnx::TensorProto& initializer,
const AllocatorPtr& cpu_allocator,
std::vector<uint8_t>& transposed_data) const {
return TransposeInitializer(qnn_model_wrapper, initializer, cnhw2hwcn_perm, cpu_allocator, transposed_data);
return TransposeInitializer(qnn_model_wrapper, initializer, cnhw2hwcn_perm, transposed_data);
}
Status TwoDimensionTranspose(const QnnModelWrapper& qnn_model_wrapper,
std::vector<uint32_t>& data_shape,
const onnx::TensorProto& initializer,
const AllocatorPtr& cpu_allocator,
std::vector<uint8_t>& transposed_data) const {
auto tmp = data_shape[0];
data_shape[0] = data_shape[1];
data_shape[1] = tmp;
std::vector<size_t> two_dim_trans_perm{1, 0};
return TransposeInitializer(qnn_model_wrapper, initializer, two_dim_trans_perm, cpu_allocator, transposed_data);
return TransposeInitializer(qnn_model_wrapper, initializer, two_dim_trans_perm, transposed_data);
}
void InitializeQuantizeParam(Qnn_QuantizeParams_t& quantize_param, bool is_quantized_model, float scale = 0.0f, int32_t offset = 0) const {

View file

@ -155,9 +155,9 @@ Status ConvOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
const auto& input_tensor = qnn_model_wrapper.GetInitializerTensors().at(input_name);
if (1 == input_i) { // qnn Conv weight requires HWCN
if (node_unit.OpType() == "Conv") {
ORT_RETURN_IF_ERROR(TransposeFromNchwToHwcn(qnn_model_wrapper, *input_tensor, qnn_model_wrapper.GetAllocator(), unpacked_tensor));
ORT_RETURN_IF_ERROR(TransposeFromNchwToHwcn(qnn_model_wrapper, *input_tensor, unpacked_tensor));
} else if (node_unit.OpType() == "ConvTranspose") {
ORT_RETURN_IF_ERROR(TransposeFromCnhwToHwcn(qnn_model_wrapper, *input_tensor, qnn_model_wrapper.GetAllocator(), unpacked_tensor));
ORT_RETURN_IF_ERROR(TransposeFromCnhwToHwcn(qnn_model_wrapper, *input_tensor, unpacked_tensor));
} else {
ORT_THROW("Unexpected operator %s", node_unit.OpType());
}

View file

@ -117,7 +117,6 @@ Status GemmOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
ORT_RETURN_IF_ERROR(TwoDimensionTranspose(qnn_model_wrapper,
input_shape,
*input_tensor,
qnn_model_wrapper.GetAllocator(),
unpacked_tensor));
} else {
ORT_RETURN_IF_ERROR(qnn_model_wrapper.UnpackInitializerData(*input_tensor, unpacked_tensor));

View file

@ -87,7 +87,6 @@ const NodeUnit& QnnModel::GetNodeUnit(const Node* node,
Status QnnModel::ComposeGraph(const GraphViewer& graph_viewer,
const onnxruntime::Node& fused_node) {
LOGS(logger_, VERBOSE) << "ComposeGraph Graph name: " << graph_viewer.Name();
const onnxruntime::AllocatorPtr& cpu_allocator = GetAllocator();
// Holder for the NodeUnits in the graph, this will guarantee the NodeUnits is
// valid throughout the lifetime of the ModelBuilder
@ -103,7 +102,7 @@ Status QnnModel::ComposeGraph(const GraphViewer& graph_viewer,
qnn_backend_manager_->GetQnnBackendHandle(),
model_input_index_map_,
model_output_index_map_,
initializer_inputs_, cpu_allocator);
initializer_inputs_);
bool rt = true;
rt = qnn_model_wrapper.CreateQnnGraph(qnn_backend_manager_->GetQnnContext(), graph_name);
if (!rt) {

View file

@ -4,7 +4,6 @@
#pragma once
#include "core/common/status.h"
#include "core/framework/allocator.h"
#include "core/graph/graph_viewer.h"
#include "core/providers/qnn/builder/qnn_def.h"
#include "core/providers/qnn/builder/qnn_model_wrapper.h"
@ -18,10 +17,8 @@ class QnnModel {
public:
QnnModel(const logging::Logger& logger,
QnnBackendManager* qnn_backend_manager,
const onnxruntime::AllocatorPtr& cpu_allocator,
bool is_quantized_model = true)
: cpu_allocator_(cpu_allocator),
logger_(logger),
: logger_(logger),
qnn_backend_manager_(qnn_backend_manager),
is_quantized_model_(is_quantized_model) {
}
@ -82,13 +79,6 @@ class QnnModel {
const std::unordered_map<const Node*, const NodeUnit*>& node_unit_map) const;
bool GetGraphInfoFromModel(QnnModelWrapper& model_wrapper);
onnxruntime::AllocatorPtr GetAllocator() {
if (cpu_allocator_ == nullptr) {
LOGS_DEFAULT(ERROR) << "cpu_allocator is null!";
}
return cpu_allocator_;
}
Status GetQnnTensorDataLength(const std::vector<uint32_t>& dims,
Qnn_DataType_t data_type,
size_t& data_length) const;
@ -102,7 +92,6 @@ class QnnModel {
return it->second.index_;
}
onnxruntime::AllocatorPtr cpu_allocator_;
const logging::Logger& logger_;
std::unique_ptr<GraphInfo> graph_info_;
QnnBackendManager* qnn_backend_manager_ = nullptr;

View file

@ -12,7 +12,6 @@
#include "qnn_def.h"
#include "core/common/logging/logging.h"
#include "core/graph/graph_viewer.h"
#include "core/framework/allocator.h"
#include "core/providers/shared/node_unit/node_unit.h"
#include "core/providers/shared/utils/utils.h"
@ -27,16 +26,14 @@ class QnnModelWrapper {
const Qnn_BackendHandle_t& backend_handle,
const std::unordered_map<std::string, size_t>& input_index_map,
const std::unordered_map<std::string, size_t>& output_index_map,
const std::unordered_set<std::string>& initializer_lookup,
const onnxruntime::AllocatorPtr& cpu_allocator)
const std::unordered_set<std::string>& initializer_lookup)
: graph_viewer_(graph_viewer),
logger_(logger),
qnn_interface_(qnn_interface),
backend_handle_(backend_handle),
input_index_map_(input_index_map),
output_index_map_(output_index_map),
initializer_lookup_(initializer_lookup),
cpu_allocator_(cpu_allocator) {
initializer_lookup_(initializer_lookup) {
}
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(QnnModelWrapper);
@ -108,13 +105,6 @@ class QnnModelWrapper {
return input_index_map_.find(tensor_name) != input_index_map_.end();
}
onnxruntime::AllocatorPtr GetAllocator() const {
if (cpu_allocator_ == nullptr) {
LOGS_DEFAULT(ERROR) << "cpu_allocator is null!";
}
return cpu_allocator_;
}
Status AddTransposeNode(NodeIndex node_index,
const std::string& input_name,
const std::string& output_name,
@ -214,7 +204,6 @@ class QnnModelWrapper {
const std::unordered_map<std::string, size_t>& input_index_map_;
const std::unordered_map<std::string, size_t>& output_index_map_;
const std::unordered_set<std::string>& initializer_lookup_;
onnxruntime::AllocatorPtr cpu_allocator_;
const std::vector<uint32_t> nchw2hwcn_perm_{2, 3, 1, 0};
const std::vector<uint32_t> cnhw2hwcn_perm_{2, 3, 0, 1};
}; // QnnModelWrapper

View file

@ -5,7 +5,6 @@
#include <filesystem>
#include "core/providers/common.h"
#include "core/framework/allocatormgr.h"
#include "core/framework/compute_capability.h"
#include "core/graph/graph_viewer.h"
#include "core/session/onnxruntime_session_options_config_keys.h"
@ -128,14 +127,6 @@ QNNExecutionProvider::QNNExecutionProvider(const ProviderOptions& provider_optio
ParseHtpPerformanceMode(htp_performance_mode_pos->second);
}
AllocatorCreationInfo device_info(
[](int) {
return std::make_unique<CPUAllocator>(OrtMemoryInfo(QNN, OrtAllocatorType::OrtDeviceAllocator));
});
cpu_allocator_ = CreateAllocator(device_info);
InsertAllocator(cpu_allocator_);
qnn_backend_manager_ = std::make_unique<qnn::QnnBackendManager>(backend_path_,
profiling_level_,
rpc_control_latency_,
@ -246,7 +237,7 @@ QNNExecutionProvider::GetSupportedNodes(const GraphViewer& graph_viewer,
qnn_backend_manager_->GetQnnBackendHandle(),
model_input_index_map,
model_output_index_map,
initializer_input_lookup, cpu_allocator_);
initializer_input_lookup);
for (const auto& node : graph_viewer.Nodes()) {
const NodeUnit* node_unit = node_unit_map.at(&node);
@ -436,7 +427,6 @@ Status QNNExecutionProvider::CompileFromOrtGraph(const std::vector<FusedNodeAndG
std::unique_ptr<qnn::QnnModel> qnn_model = std::make_unique<qnn::QnnModel>(logger,
qnn_backend_manager_.get(),
cpu_allocator_,
qnn_backend_manager_->IsNpuBackend());
ORT_RETURN_IF_ERROR(qnn_model->ComposeGraph(graph_viewer, fused_node));
@ -480,7 +470,6 @@ Status QNNExecutionProvider::Compile(const std::vector<FusedNodeAndGraph>& fused
if (load_from_cached_context) {
std::unique_ptr<qnn::QnnModel> qnn_model = std::make_unique<qnn::QnnModel>(logger,
qnn_backend_manager_.get(),
cpu_allocator_,
is_npu_backend);
ORT_RETURN_IF_ERROR(qnn_backend_manager_->LoadCachedQnnContext(context_cache_pathstring, *(qnn_model.get())));
ORT_RETURN_IF_ERROR(qnn_model->SetGraphInputOutputInfo(graph_viewer, fused_node));

View file

@ -66,7 +66,6 @@ class QNNExecutionProvider : public IExecutionProvider {
qnn::HtpPerformanceMode htp_performance_mode_ = qnn::HtpPerformanceMode::kHtpDefault;
std::unique_ptr<qnn::QnnBackendManager> qnn_backend_manager_;
std::unordered_map<std::string, std::unique_ptr<qnn::QnnModel>> qnn_models_;
AllocatorPtr cpu_allocator_;
uint32_t rpc_control_latency_ = 0;
bool context_cache_enabled_ = false;
std::string context_cache_path_ = "";

View file

@ -940,9 +940,10 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
transform_layout_fn = [this](Graph& graph_to_transform, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
ORT_RETURN_IF_ERROR_SESSIONID_(
layout_transformer::TransformLayoutForEP(graph_to_transform, modified, execution_provider,
execution_providers_.GetDefaultCpuAllocator(), debug_graph_fn));
std::move(cpu_allocator), debug_graph_fn));
if (modified) {
ORT_RETURN_IF_ERROR_SESSIONID_(
@ -1272,11 +1273,12 @@ Status PartitionOrtFormatModel(onnxruntime::Graph& graph,
// only provide NCWH to NHWC layout transformer if supported
if (layout_transformer::IsSupportedOpset(graph)) {
transform_layout_fn =
[&providers](Graph& graph_to_transform, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
[](Graph& graph_to_transform, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
return layout_transformer::TransformLayoutForEP(graph_to_transform, modified, execution_provider,
providers.GetDefaultCpuAllocator(), debug_graph_fn);
std::move(cpu_allocator), debug_graph_fn);
};
}
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)