From 4ab7d410aeacc4bd390d47a86384896d7d2aa599 Mon Sep 17 00:00:00 2001 From: cao lei Date: Fri, 26 May 2023 04:54:42 -0700 Subject: [PATCH] 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 --- .../core/optimizer/graph_transformer_utils.cc | 8 +++++--- .../core/optimizer/optimizer_execution_frame.cc | 4 ++-- .../core/optimizer/optimizer_execution_frame.h | 5 ----- .../qnn/builder/opbuilder/base_op_builder.cc | 2 +- .../qnn/builder/opbuilder/base_op_builder.h | 10 +++------- .../qnn/builder/opbuilder/conv_op_builder.cc | 4 ++-- .../qnn/builder/opbuilder/gemm_op_builder.cc | 1 - .../core/providers/qnn/builder/qnn_model.cc | 3 +-- .../core/providers/qnn/builder/qnn_model.h | 13 +------------ .../providers/qnn/builder/qnn_model_wrapper.h | 15 ++------------- .../core/providers/qnn/qnn_execution_provider.cc | 13 +------------ .../core/providers/qnn/qnn_execution_provider.h | 1 - onnxruntime/core/session/inference_session.cc | 12 +++++++----- 13 files changed, 25 insertions(+), 66 deletions(-) diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 7196465e5d..f8a48fabb3 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -234,7 +234,9 @@ InlinedVector> 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(); transformers.emplace_back(std::make_unique(std::move(cpu_allocator))); } break; @@ -349,7 +351,7 @@ InlinedVector> GenerateTransformers( if (MlasNchwcGetBlockSize() > 1) { transformers.emplace_back(std::make_unique()); } - auto cpu_allocator = cpu_execution_provider.GetAllocator(OrtMemTypeDefault); + AllocatorPtr cpu_allocator = std::make_shared(); auto cpu_registry = cpu_execution_provider.GetKernelRegistry(); auto nhwc_transformer = std::make_unique(std::move(cpu_allocator), std::move(cpu_registry)); if (nhwc_transformer->IsActive()) { @@ -425,7 +427,7 @@ InlinedVector> 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(); auto cpu_registry = cpu_execution_provider.GetKernelRegistry(); auto nhwc_transformer = std::make_unique(std::move(cpu_allocator), std::move(cpu_registry)); if (nhwc_transformer->IsActive()) { diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.cc b/onnxruntime/core/optimizer/optimizer_execution_frame.cc index 5a90efb51d..fc7e694b6a 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.cc +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.cc @@ -35,7 +35,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, const std::function& 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(); ORT_ENFORCE(allocator_ptr_, "Failed to get allocator for optimizer"); ORT_THROW_IF_ERROR(data_transfer_mgr_.RegisterDataTransfer(std::make_unique())); @@ -89,7 +89,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, const std::function& 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(); ORT_ENFORCE(allocator_ptr_, "Failed to get allocator for optimizer"); ORT_THROW_IF_ERROR(data_transfer_mgr_.RegisterDataTransfer(std::make_unique())); diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.h b/onnxruntime/core/optimizer/optimizer_execution_frame.h index bac146c96a..e1b8a91545 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.h +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.h @@ -34,10 +34,6 @@ class OptimizerExecutionFrame final : public IExecutionFrame { const std::function& 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 diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc index 5b50b85486..f31c571f13 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc @@ -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& perm, - const AllocatorPtr& cpu_allocator, std::vector& 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(); Tensor in_tensor = Tensor(tensor_dtype, tensor_shape, cpu_allocator); auto rank = perm.size(); diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h index 98fb49eb86..febf5c7e4d 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h @@ -218,33 +218,29 @@ class BaseOpBuilder : public IOpBuilder { Status TransposeInitializer(const QnnModelWrapper& qnn_model_wrapper, const onnx::TensorProto& initializer, const std::vector& perm, - const AllocatorPtr& cpu_allocator, std::vector& transposed_data) const; Status TransposeFromNchwToHwcn(const QnnModelWrapper& qnn_model_wrapper, const onnx::TensorProto& initializer, - const AllocatorPtr& cpu_allocator, std::vector& 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& 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& data_shape, const onnx::TensorProto& initializer, - const AllocatorPtr& cpu_allocator, std::vector& transposed_data) const { auto tmp = data_shape[0]; data_shape[0] = data_shape[1]; data_shape[1] = tmp; std::vector 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 { diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/conv_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/conv_op_builder.cc index dc7168f7d9..5da223fa23 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/conv_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/conv_op_builder.cc @@ -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()); } diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/gemm_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/gemm_op_builder.cc index 067fc2aa4f..a4ea16555b 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/gemm_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/gemm_op_builder.cc @@ -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)); diff --git a/onnxruntime/core/providers/qnn/builder/qnn_model.cc b/onnxruntime/core/providers/qnn/builder/qnn_model.cc index 21fe6a361b..9af7d7d0e0 100644 --- a/onnxruntime/core/providers/qnn/builder/qnn_model.cc +++ b/onnxruntime/core/providers/qnn/builder/qnn_model.cc @@ -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) { diff --git a/onnxruntime/core/providers/qnn/builder/qnn_model.h b/onnxruntime/core/providers/qnn/builder/qnn_model.h index 52b1b401f0..8797dc9277 100644 --- a/onnxruntime/core/providers/qnn/builder/qnn_model.h +++ b/onnxruntime/core/providers/qnn/builder/qnn_model.h @@ -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& 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& 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 graph_info_; QnnBackendManager* qnn_backend_manager_ = nullptr; diff --git a/onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.h b/onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.h index 6682d33093..93a2860951 100644 --- a/onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.h +++ b/onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.h @@ -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& input_index_map, const std::unordered_map& output_index_map, - const std::unordered_set& initializer_lookup, - const onnxruntime::AllocatorPtr& cpu_allocator) + const std::unordered_set& 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& input_index_map_; const std::unordered_map& output_index_map_; const std::unordered_set& initializer_lookup_; - onnxruntime::AllocatorPtr cpu_allocator_; const std::vector nchw2hwcn_perm_{2, 3, 1, 0}; const std::vector cnhw2hwcn_perm_{2, 3, 0, 1}; }; // QnnModelWrapper diff --git a/onnxruntime/core/providers/qnn/qnn_execution_provider.cc b/onnxruntime/core/providers/qnn/qnn_execution_provider.cc index 236e660a94..6677c1e97f 100644 --- a/onnxruntime/core/providers/qnn/qnn_execution_provider.cc +++ b/onnxruntime/core/providers/qnn/qnn_execution_provider.cc @@ -5,7 +5,6 @@ #include #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(OrtMemoryInfo(QNN, OrtAllocatorType::OrtDeviceAllocator)); - }); - - cpu_allocator_ = CreateAllocator(device_info); - InsertAllocator(cpu_allocator_); - qnn_backend_manager_ = std::make_unique(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 qnn_model = std::make_unique(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& fused if (load_from_cached_context) { std::unique_ptr qnn_model = std::make_unique(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)); diff --git a/onnxruntime/core/providers/qnn/qnn_execution_provider.h b/onnxruntime/core/providers/qnn/qnn_execution_provider.h index b804f41ee3..2f6166d46e 100644 --- a/onnxruntime/core/providers/qnn/qnn_execution_provider.h +++ b/onnxruntime/core/providers/qnn/qnn_execution_provider.h @@ -66,7 +66,6 @@ class QNNExecutionProvider : public IExecutionProvider { qnn::HtpPerformanceMode htp_performance_mode_ = qnn::HtpPerformanceMode::kHtpDefault; std::unique_ptr qnn_backend_manager_; std::unordered_map> qnn_models_; - AllocatorPtr cpu_allocator_; uint32_t rpc_control_latency_ = 0; bool context_cache_enabled_ = false; std::string context_cache_path_ = ""; diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 0b71fbafec..df51dcdadf 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -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(); 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(); 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)