diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 96e513c8a7bc9..36ba8db9bdc75 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1252,7 +1252,7 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) onnx_test_runner_common onnxruntime_test_utils onnxruntime_common onnxruntime onnxruntime_flatbuffers onnx_test_data_proto ${onnxruntime_EXTERNAL_LIBRARIES} - ${GETOPT_LIB_WIDE} ${SYS_PATH_LIB} ${CMAKE_DL_LIBS}) + absl::flags absl::flags_parse ${SYS_PATH_LIB} ${CMAKE_DL_LIBS}) if(NOT WIN32) if(onnxruntime_USE_SNPE) list(APPEND onnxruntime_perf_test_libs onnxruntime_providers_snpe) @@ -1272,7 +1272,7 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) target_link_libraries(onnxruntime_perf_test PRIVATE debug dbghelp advapi32) endif() else() - target_link_libraries(onnxruntime_perf_test PRIVATE onnx_test_runner_common ${GETOPT_LIB_WIDE} ${onnx_test_libs}) + target_link_libraries(onnxruntime_perf_test PRIVATE onnx_test_runner_common absl::flags absl::flags_parse ${onnx_test_libs}) endif() set_target_properties(onnxruntime_perf_test PROPERTIES FOLDER "ONNXRuntimeTest") diff --git a/include/onnxruntime/core/framework/execution_provider.h b/include/onnxruntime/core/framework/execution_provider.h index 7df3368ad4e0b..1bb7f219c9a45 100644 --- a/include/onnxruntime/core/framework/execution_provider.h +++ b/include/onnxruntime/core/framework/execution_provider.h @@ -179,7 +179,12 @@ class IExecutionProvider { /** Get the device id of current execution provider */ - virtual int GetDeviceId() const { return default_device_.Id(); }; + virtual int GetDeviceId() const { return default_device_.Id(); } + + /** + * Get the OrtDevice the execution provider was registered with. + */ + const OrtDevice& GetDevice() const { return default_device_; } /** Get execution provider's configuration options. diff --git a/include/onnxruntime/core/framework/ortdevice.h b/include/onnxruntime/core/framework/ortdevice.h index 536d641b4eef9..fea970b84fd84 100644 --- a/include/onnxruntime/core/framework/ortdevice.h +++ b/include/onnxruntime/core/framework/ortdevice.h @@ -150,6 +150,13 @@ struct OrtDevice { return alignment < other.alignment; } + bool EqualIgnoringAlignment(const OrtDevice& other) const { + return device_type == other.device_type && + memory_type == other.memory_type && + vendor_id == other.vendor_id && + device_id == other.device_id; + } + private: // Device type. int32_t device_type : 8; diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index bd694f7b3b23c..866892979b749 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -1220,7 +1220,10 @@ class Graph { // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi #endif #if !defined(ORT_MINIMAL_BUILD) - /** Gets the GraphProto representation of this Graph only. */ + /** Gets the GraphProto representation of this Graph only. + * This does not remove in-memory tags for graph initializers. + * Use ToGraphProto() const to get a GraphProto that can be serialized externally. + */ const ONNX_NAMESPACE::GraphProto& ToGraphProto(); /// @@ -1439,6 +1442,27 @@ class Graph { // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi return Resolve(default_options); } + /// + /// This function converts all the graph TensorProto initializers into OrtValues + /// and creates a in-memory external data reference for each OrtValue. + /// + /// + Status ConvertInitializersIntoOrtValues(); + + /** + * @brief Converts a subset of graph TensorProto initializers into OrtValues and updates the graph proto. + * + * This function converts specified TensorProto initializers in the graph into OrtValues and + * creates in-memory external data references for each OrtValue. It then updates the provided + * GraphProto with the modified initializers. + * + * @param iterators Span of iterators pointing to the initializers and the order that should be processed + * @param output_graph_proto The GraphProto to be updated with the modified initializers + * @return Status Returns a Status object indicating success or any errors that occurred during conversion + */ + Status RegenerateInitializersAndReplaceInMemory(gsl::span iterators, + ONNX_NAMESPACE::GraphProto& output_graph_proto) const; + const std::unordered_set& GetOuterScopeNodeArgNames() const noexcept { return outer_scope_node_arg_names_; } @@ -1595,20 +1619,25 @@ class Graph { // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi /// This function is used by ToGraphProto() to ensure in-memory external data references /// don't leak externally since they are non-standard. /// - /// It handles two scenarios: - /// - When GraphSynchronizationNeeded() is false: GraphProto is simply copied + /// It is used when GraphSynchronizationNeeded() is false: GraphProto is simply copied /// from graph_proto_ by ToGraphProto(). This copy includes both main graph /// and subgraph initializers. This function examines all initializers /// and inlines any in-memory data references. - /// - When GraphSynchronizationNeeded() is true: ToGraphProto() generates a new GraphProto - /// using ToGraphProtoInternal(). This doesn't transfer main graph initializers, which are - /// copied and inlined by ToGraphProto() itself. This function processes only the subgraph initializers - /// as needed. /// /// The GraphProto to process - /// Whether to process the main graph initializers - /// Status indicating success or failure /// - Status ProcessSubgraphsInMemoryData(ONNX_NAMESPACE::GraphProto& output_graph_proto, bool process_main) const; + /// Status indicating success or failure + Status ProcessSubgraphsInMemoryData(ONNX_NAMESPACE::GraphProto& output_graph_proto) const; + + /// + /// This function replaces all of the initializers within output_graph_proto + /// from this Graph instance. All in memory initializers are regenerated and inlined. + /// This is necessary even if the graph_proto_ is already up to date because initializers() may + /// contain obsolete initializers that are no longer in use due to optimizations and contain obsolete + /// references to OrtValues that may no longer be around (since we like appending rather than replacing). + /// + /// Destination GraphProto to receive the updated initializers. + /// Status indicating success or failure. + Status RegenerateInitializersAndReplaceInMemory(ONNX_NAMESPACE::GraphProto& output_graph_proto) const; /// /// This function traverses the graph bottom up and externalizes diff --git a/include/onnxruntime/core/session/environment.h b/include/onnxruntime/core/session/environment.h index 306f81df38e48..89467f5238fa9 100644 --- a/include/onnxruntime/core/session/environment.h +++ b/include/onnxruntime/core/session/environment.h @@ -106,6 +106,15 @@ class Environment { return shared_allocators_; } + /** + * Returns an AllocatorPtr for a shared IAllocator based allocator if it matches the memory info. + * The OrtMemoryInfo name and whether it's an arena or device allocator is ignored in the lookup, as is the + * alignment. + * The user calling this function is not expected to know the alignment, and we expect the allocator instance to be + * created with a valid alignment for the device. + */ + AllocatorPtr GetRegisteredSharedAllocator(const OrtMemoryInfo& mem_info) const; + /** * Removes registered allocator that was previously registered for sharing between multiple sessions. */ @@ -171,7 +180,7 @@ class Environment { std::unique_ptr inter_op_thread_pool_; bool create_global_thread_pools_{false}; - std::mutex mutex_; + mutable std::mutex mutex_; // shared allocators from various sources. // CreateAndRegisterAllocator[V2]: IAllocator allocators created by ORT diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 8ae7535da4413..e4f8cd6df678e 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -666,12 +666,16 @@ void Node::ToProto(NodeProto& proto, bool update_subgraphs) const { // Set attributes. proto.clear_attribute(); - for (const auto& attribute : attributes_) { + for (const auto& [name, attribute] : attributes_) { const gsl::not_null attr{proto.add_attribute()}; - *attr = attribute.second; // copy - if (update_subgraphs && attr->has_g()) { + *attr = attribute; // copy + if (update_subgraphs && utils::HasGraph(*attr)) { + auto find_hit = attr_to_subgraph_map_.find(name); + // Force ToGraphProto() const to be called so + // that any in-memory TensorProto initializers go back to being inlined + const Graph& subgraph = *find_hit->second; attr->clear_g(); - *attr->mutable_g() = attr_to_subgraph_map_.find(attribute.first)->second->ToGraphProto(); + *attr->mutable_g() = subgraph.ToGraphProto(); } } @@ -3381,7 +3385,12 @@ Status Graph::Resolve(const ResolveOptions& options) { return Status::OK(); }; - ORT_RETURN_IF_ERROR(ForThisAndAllSubgraphs(all_subgraphs, finalize_func)); + return ForThisAndAllSubgraphs(all_subgraphs, finalize_func); +} + +Status Graph::ConvertInitializersIntoOrtValues() { + std::vector all_subgraphs; + FindAllSubgraphs(all_subgraphs); auto put_weights_maybe_in_memory_func = [&](Graph& graph) -> Status { // if we have any initializers that are not in memory, put them there. @@ -4308,11 +4317,47 @@ Status InlineOrCopyInitializer(const Graph& src_graph, const ONNX_NAMESPACE::Ten } return Status::OK(); } - } // namespace -Status Graph::ProcessSubgraphsInMemoryData(ONNX_NAMESPACE::GraphProto& output_graph_proto, - bool process_main) const { +Status Graph::RegenerateInitializersAndReplaceInMemory(gsl::span iterators, + ONNX_NAMESPACE::GraphProto& output_graph_proto) const { + auto& mutable_initializers = *output_graph_proto.mutable_initializer(); + +#if !defined(DISABLE_SPARSE_TENSORS) + output_graph_proto.clear_sparse_initializer(); + + const auto& model_path = ModelPath(); + const bool has_sparse_initializers = !sparse_tensor_names_.empty(); + const auto sparse_end = sparse_tensor_names_.end(); + + for (const auto& iter : iterators) { + const auto& [name, tensor_proto] = *iter; + const auto& initializer = *tensor_proto; + if (!has_sparse_initializers || sparse_end == sparse_tensor_names_.find(name)) { + ORT_RETURN_IF_ERROR(InlineOrCopyInitializer(*this, initializer, + *mutable_initializers.Add())); + } else { + auto& sparse_initializer = *output_graph_proto.add_sparse_initializer(); + if (utils::HasExternalDataInMemory(initializer)) { + ONNX_NAMESPACE::TensorProto tensor_proto_inlined; + ORT_RETURN_IF_ERROR(InlineOrCopyInitializer(*this, initializer, + tensor_proto_inlined)); + ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(tensor_proto_inlined, model_path, sparse_initializer)); + } else { + ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(initializer, model_path, sparse_initializer)); + } + } + } +#else + for (const auto& iter : iterators) { + const auto& [name, tensor_proto] = *iter; + ORT_RETURN_IF_ERROR(InlineOrCopyInitializer(*this, *tensor_proto, *mutable_initializers.Add())); + } +#endif + return Status::OK(); +} + +Status Graph::ProcessSubgraphsInMemoryData(ONNX_NAMESPACE::GraphProto& output_graph_proto) const { for (const auto& node : Nodes()) { if (node.ContainsSubgraph()) { // Let's find this node in the output_graph_proto @@ -4343,103 +4388,48 @@ Status Graph::ProcessSubgraphsInMemoryData(ONNX_NAMESPACE::GraphProto& output_gr "Subgraph ", name, " is referred to in GetAttributeNameToSubgraphMap, but not found in node ", node.Name(), " while attempting to recurse into it."); auto& result_subgraph = *sub_hit->mutable_g(); - ORT_RETURN_IF_ERROR(subgraph->ProcessSubgraphsInMemoryData(result_subgraph, process_main)); + ORT_RETURN_IF_ERROR(subgraph->ProcessSubgraphsInMemoryData(result_subgraph)); } } } - // When graph_proto is copied from graph_proto, initializers already present in the main graph - if (parent_graph_ != nullptr || process_main) { -#if !defined(DISABLE_SPARSE_TENSORS) - auto* mutable_initializers = output_graph_proto.mutable_initializer(); - const auto& model_path = ModelPath(); - const bool has_sparse_initializers = !sparse_tensor_names_.empty(); - const auto sparse_end = sparse_tensor_names_.end(); - - // We want to make sure that sparse initializers do not appear - // as dense duplicates within the initializers list. - std::optional> initializer_to_remove; - if (has_sparse_initializers) { - // We need to remove the dense initializers that are sparse tensors - initializer_to_remove.emplace(); - } - - for (auto first = mutable_initializers->begin(), end = mutable_initializers->end(); first != end; ++first) { - auto& initializer = *first; - if (utils::HasExternalDataInMemory(initializer)) { - // If the initializer has external data in memory, we need to inline it. - ORT_RETURN_IF_ERROR(InlineOrCopyInitializer(*this, initializer, initializer)); - } - if (has_sparse_initializers && sparse_end != sparse_tensor_names_.find(initializer.name())) { - auto& sparse_initializer = *output_graph_proto.add_sparse_initializer(); - ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(initializer, model_path, sparse_initializer)); - initializer_to_remove->insert(initializer.name()); - } - } - - // erase/remove dense initializers that are sparse tensors so no duplicates are present - if (initializer_to_remove && !initializer_to_remove->empty()) { - mutable_initializers->erase(std::remove_if( - mutable_initializers->begin(), mutable_initializers->end(), - [&initializer_to_remove](const ONNX_NAMESPACE::TensorProto& initializer) { - return initializer_to_remove->count(initializer.name()) > 0; - }), - mutable_initializers->end()); - } -#else - for (auto& initializer : *output_graph_proto.mutable_initializer()) { - if (utils::HasExternalDataInMemory(initializer)) { - // If the initializer has external data in memory, we need to inline it. - ORT_RETURN_IF_ERROR(InlineOrCopyInitializer(*this, initializer, initializer)); - } + // Filter in iterators for weights that are present in the name_to_initial_tensor_ map + // and preserve the order. This is needed for tests. + InlinedVector initializers_to_process; + initializers_to_process.reserve(name_to_initial_tensor_.size()); + for (const auto& tensor_proto : output_graph_proto.initializer()) { + auto hit = name_to_initial_tensor_.find(tensor_proto.name()); + if (hit != name_to_initial_tensor_.end()) { + initializers_to_process.push_back(hit); } -#endif } - return Status::OK(); + + output_graph_proto.clear_initializer(); + return RegenerateInitializersAndReplaceInMemory(initializers_to_process, output_graph_proto); } ONNX_NAMESPACE::GraphProto Graph::ToGraphProto() const { GraphProto result; if (!GraphProtoSyncNeeded()) { result = *graph_proto_; - ORT_THROW_IF_ERROR(ProcessSubgraphsInMemoryData(result, /*process_main*/ true)); + ORT_THROW_IF_ERROR(ProcessSubgraphsInMemoryData(result)); } else { + // Recursion is handled via Node::ToProto() const -> Graph::ToGraphProto() const (this method) + // so below we handle this graph only. ToGraphProtoInternal(result); - ORT_THROW_IF_ERROR(ProcessSubgraphsInMemoryData(result, /*process_main*/ false)); - - // Add initializers to parent graph by copy converting them from graph_proto_ - // ToGraphProtoInternal() does not copy initializers for the main graph - auto* mutable_initializers = result.mutable_initializer(); - -#if !defined(DISABLE_SPARSE_TENSORS) - const auto& model_path = ModelPath(); - const bool has_sparse_initializers = !sparse_tensor_names_.empty(); - const auto sparse_end = sparse_tensor_names_.end(); - - for (const auto& initializer : graph_proto_->initializer()) { - if (!has_sparse_initializers || sparse_end == sparse_tensor_names_.find(initializer.name())) { - ORT_THROW_IF_ERROR(InlineOrCopyInitializer(*this, initializer, - *mutable_initializers->Add())); - } else { - auto& sparse_initializer = *result.add_sparse_initializer(); - if (utils::HasExternalDataInMemory(initializer)) { - ONNX_NAMESPACE::TensorProto tensor_proto; - ORT_THROW_IF_ERROR(InlineOrCopyInitializer(*this, initializer, - tensor_proto)); - ORT_THROW_IF_ERROR(utils::DenseTensorToSparseTensorProto(tensor_proto, model_path, sparse_initializer)); - } else { - ORT_THROW_IF_ERROR(utils::DenseTensorToSparseTensorProto(initializer, model_path, sparse_initializer)); - } + InlinedVector initializers_to_process; + initializers_to_process.reserve(name_to_initial_tensor_.size()); + for (const auto& tensor_proto : graph_proto_->initializer()) { + auto hit = name_to_initial_tensor_.find(tensor_proto.name()); + if (hit != name_to_initial_tensor_.end()) { + initializers_to_process.push_back(hit); } } -#else - for (const auto& initializer : graph_proto_->initializer()) { - ORT_THROW_IF_ERROR(InlineOrCopyInitializer(*this, initializer, *mutable_initializers->Add())); - } -#endif - } + ORT_THROW_IF_ERROR(RegenerateInitializersAndReplaceInMemory(initializers_to_process, + result)); + } return result; } @@ -5235,23 +5225,7 @@ Status Graph::AddConstantProtoAsInitializer(const ONNX_NAMESPACE::NodeProto& nod tensor_proto.set_name(std::string(new_name.value())); } - // In the constant node, we won't have symbolic dims. - const auto tensor_shape = utils::GetTensorShapeFromTensorProto(tensor_proto); - auto ml_data = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); - const size_t size_in_bytes = Tensor::CalculateTensorStorageSize(ml_data, tensor_shape); - - if (size_in_bytes > utils::kSmallTensorExternalDataThreshold) { - OrtValue ort_value; - ORT_RETURN_IF_ERROR(utils::TensorProtoToOrtValue(Env::Default(), ModelPath(), tensor_proto, - CPUAllocator::DefaultInstance(), ort_value)); - - constexpr const bool use_tensor_buffer_true = true; - auto tensor_proto_to_add = utils::TensorToTensorProto(ort_value.Get(), tensor_proto.name(), - use_tensor_buffer_true); - ORT_RETURN_IF_ERROR(AddInitializedOrtValue(tensor_proto_to_add, ort_value)); - } else { - AddInitializedTensor(tensor_proto); - } + AddInitializedTensor(tensor_proto); if (GetNodeArg(tensor_proto.name()) == nullptr) { TypeProto t{utils::TypeProtoFromTensorProto(tensor_proto)}; diff --git a/onnxruntime/core/optimizer/attention_fusion.cc b/onnxruntime/core/optimizer/attention_fusion.cc index 616bc1257676f..3f9b58f71bd23 100644 --- a/onnxruntime/core/optimizer/attention_fusion.cc +++ b/onnxruntime/core/optimizer/attention_fusion.cc @@ -111,7 +111,7 @@ static NodeArg& MergeQkvWeights(Graph& graph, int64_t hidden_size, utils::SetRawDataInTensorProto(initializer, result.data(), gsl::narrow(element_count) * sizeof(MLFloat16)); } - return graph_utils::AddInitializerWithExternalData(graph, initializer); + return graph_utils::AddInitializer(graph, initializer); } static NodeArg* ConvertMaskToInt32(Graph& graph, NodeArg* mask_input, ProviderType provider_type, diff --git a/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc b/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc index a98d0ea6f978b..86a7a4d6afbf8 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc +++ b/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc @@ -189,7 +189,7 @@ NodeArg* CreateInitializerFromVector(Graph& graph, "total_count: ", total_count, " values.size(): ", values.size()); utils::SetRawDataInTensorProto(const_tensor, values.data(), values.size() * sizeof(int64_t)); - return &graph_utils::AddInitializerWithExternalData(graph, const_tensor); + return &graph_utils::AddInitializer(graph, const_tensor); } NodeArg* InsertNodesForValidIndices(Graph& graph, diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 3d838d8aacfbb..16e8955cb4486 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -95,7 +95,7 @@ static bool ConstantFoldShapeNode(Graph& graph, Node& node) { ONNX_NAMESPACE::TensorShapeProto result_shape; result_shape.add_dim()->set_dim_value(clamped_slice_length); constant_arg_out->SetShape(result_shape); - graph_utils::AddInitializerWithExternalData(graph, shape_constant); + graph_utils::AddInitializer(graph, shape_constant); } return is_concrete_shape; // convert to constant if this is true @@ -317,11 +317,11 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, // Build the TensorProto that corresponds to the computed OrtValue and add it as initializer to the graph. auto* constant_arg_out = node->MutableOutputDefs()[fetch_idx]; const Tensor& out_tensor = ort_value.Get(); - constexpr const bool use_tensor_buffer_true = true; + constexpr const bool use_tensor_buffer_false = false; ONNX_NAMESPACE::TensorProto out_tensorproto = utils::TensorToTensorProto( out_tensor, constant_arg_out->Name(), - use_tensor_buffer_true); + use_tensor_buffer_false); ONNX_NAMESPACE::TensorShapeProto result_shape; for (auto& dim : out_tensor.Shape().GetDims()) { @@ -329,12 +329,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, } constant_arg_out->SetShape(result_shape); - // The data is too small and has been inlined. - if (!utils::HasExternalData(out_tensorproto)) { - ORT_THROW_IF_ERROR(graph.AddInitializedOrtValue(out_tensorproto, OrtValue())); - } else { - ORT_THROW_IF_ERROR(graph.AddInitializedOrtValue(out_tensorproto, ort_value)); - } + graph.AddInitializedTensor(out_tensorproto); } } } diff --git a/onnxruntime/core/optimizer/conv_add_fusion.cc b/onnxruntime/core/optimizer/conv_add_fusion.cc index c349adfccce53..6478fa7d29d5b 100644 --- a/onnxruntime/core/optimizer/conv_add_fusion.cc +++ b/onnxruntime/core/optimizer/conv_add_fusion.cc @@ -79,7 +79,7 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie auto new_name = graph.GenerateNodeArgName("ConvAddFusion_B_" + B_input_name); new_conv_B_tensor_proto.set_name(new_name); - NodeArg& new_conv_B_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_conv_B_tensor_proto); + NodeArg& new_conv_B_node_arg = graph_utils::AddInitializer(graph, new_conv_B_tensor_proto); graph_utils::ReplaceNodeInput(node, 2, new_conv_B_node_arg); } else { @@ -94,7 +94,7 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie auto new_name = graph.GenerateNodeArgName("ConvAddFusion_Add_B_" + add_B_tensor_proto->name()); new_conv_B_tensor_proto.set_name(new_name); - NodeArg& new_add_B_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_conv_B_tensor_proto); + NodeArg& new_add_B_node_arg = graph_utils::AddInitializer(graph, new_conv_B_tensor_proto); graph_utils::AddNodeInput(node, 2, new_add_B_node_arg); } diff --git a/onnxruntime/core/optimizer/conv_bn_fusion.cc b/onnxruntime/core/optimizer/conv_bn_fusion.cc index 8bf5420baddde..a14639631d7a1 100644 --- a/onnxruntime/core/optimizer/conv_bn_fusion.cc +++ b/onnxruntime/core/optimizer/conv_bn_fusion.cc @@ -120,10 +120,10 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff new_conv_W_tensor_proto.set_name(new_W_name); new_conv_B_tensor_proto.set_name(new_B_name); - NodeArg& new_conv_W_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_conv_W_tensor_proto); + NodeArg& new_conv_W_node_arg = graph_utils::AddInitializer(graph, new_conv_W_tensor_proto); graph_utils::ReplaceNodeInput(node, 1, new_conv_W_node_arg); - auto& new_conv_B_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_conv_B_tensor_proto); + auto& new_conv_B_node_arg = graph_utils::AddInitializer(graph, new_conv_B_tensor_proto); if (conv_inputs.size() == 3) { graph_utils::ReplaceNodeInput(node, 2, new_conv_B_node_arg); diff --git a/onnxruntime/core/optimizer/conv_mul_fusion.cc b/onnxruntime/core/optimizer/conv_mul_fusion.cc index dc50a150537f7..e91a00729e9db 100644 --- a/onnxruntime/core/optimizer/conv_mul_fusion.cc +++ b/onnxruntime/core/optimizer/conv_mul_fusion.cc @@ -90,7 +90,7 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef new_conv_W_tensor_proto.set_name(new_W_name); // Replace initializers of conv node - NodeArg& new_conv_W_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_conv_W_tensor_proto); + NodeArg& new_conv_W_node_arg = graph_utils::AddInitializer(graph, new_conv_W_tensor_proto); graph_utils::ReplaceNodeInput(conv_node, 1, new_conv_W_node_arg); if (is_3d) { @@ -100,7 +100,7 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef auto new_B_name = graph.GenerateNodeArgName("ConvMulFusion_Mul_B_" + mul_B_tensor_proto->name()); new_conv_B_tensor_proto.set_name(new_B_name); - NodeArg& new_conv_B_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_conv_B_tensor_proto); + NodeArg& new_conv_B_node_arg = graph_utils::AddInitializer(graph, new_conv_B_tensor_proto); graph_utils::ReplaceNodeInput(conv_node, 2, new_conv_B_node_arg); } diff --git a/onnxruntime/core/optimizer/double_qdq_pairs_remover.cc b/onnxruntime/core/optimizer/double_qdq_pairs_remover.cc index 7f214e656e0ab..96f75f07e32e1 100644 --- a/onnxruntime/core/optimizer/double_qdq_pairs_remover.cc +++ b/onnxruntime/core/optimizer/double_qdq_pairs_remover.cc @@ -53,7 +53,7 @@ static void ApplyNewInputValue(Graph& graph, Node& node, QDQ::InputIndex index, auto new_name = graph.GenerateNodeArgName("DoubleQDQRemoved_" + node.InputDefs()[index]->Name()); new_input_tensor.set_name(new_name); new_input_tensor.add_dims(1); - NodeArg& new_input = graph_utils::AddInitializerWithExternalData(graph, new_input_tensor); + NodeArg& new_input = graph_utils::AddInitializer(graph, new_input_tensor); graph_utils::ReplaceNodeInput(node, index, new_input); } diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index ad25f95ac1186..f8fd807084d38 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -474,7 +474,7 @@ static NodeArg* ExtractEmbedding(Graph& graph, utils::SetRawDataInTensorProto(initializer, data, gsl::narrow(element_count) * sizeof(MLFloat16)); } - NodeArg& node_arg = graph_utils::AddInitializerWithExternalData(graph, initializer); + NodeArg& node_arg = graph_utils::AddInitializer(graph, initializer); modified = true; return &node_arg; } diff --git a/onnxruntime/core/optimizer/fuse_initializers_transformer.cc b/onnxruntime/core/optimizer/fuse_initializers_transformer.cc index 388ab14dd51fe..e604c688ee033 100644 --- a/onnxruntime/core/optimizer/fuse_initializers_transformer.cc +++ b/onnxruntime/core/optimizer/fuse_initializers_transformer.cc @@ -137,12 +137,8 @@ static void FuseInitializerWithNode(Graph& graph, graph.RemoveEdge(node.Index(), next_node.Index(), 0, static_cast(next_node_arg_index)); // Add the new converted Tensor in next node as initializer potentially with external data - ONNX_NAMESPACE::TensorProto dst_tensor = utils::TensorToTensorProto(new_data.Get(), new_arg_name, true); - if (!utils::HasExternalData(dst_tensor)) { - new_data = OrtValue(); // Data is inline - } - - auto& new_arg = graph_utils::AddInitializerWithExternalData(graph, dst_tensor, std::move(new_data)); + ONNX_NAMESPACE::TensorProto dst_tensor = utils::TensorToTensorProto(new_data.Get(), new_arg_name, false); + auto& new_arg = graph_utils::AddInitializer(graph, dst_tensor); graph_utils::ReplaceNodeInput(next_node, static_cast(next_node_arg_index), new_arg); } diff --git a/onnxruntime/core/optimizer/gather_fusion.cc b/onnxruntime/core/optimizer/gather_fusion.cc index 3cd06350df95d..bd730683a4c91 100644 --- a/onnxruntime/core/optimizer/gather_fusion.cc +++ b/onnxruntime/core/optimizer/gather_fusion.cc @@ -256,7 +256,7 @@ Status GatherSliceToSplitFusion::ApplyImpl(Graph& graph, bool& modified, int gra axes_initializer_proto.add_dims(static_cast(1)); axes_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); axes_initializer_proto.add_int64_data(axis); - NodeArg* axes_arg = &graph_utils::AddInitializerWithExternalData(graph, axes_initializer_proto); + NodeArg* axes_arg = &graph_utils::AddInitializer(graph, axes_initializer_proto); Node& squeeze_node = graph.AddNode(graph.GenerateNodeName("Squeeze"), "Squeeze", "Squeeze for Fused Gather nodes", {split_output_arg, axes_arg}, {original_output_arg}); @@ -272,7 +272,7 @@ Status GatherSliceToSplitFusion::ApplyImpl(Graph& graph, bool& modified, int gra split_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); split_initializer_proto.add_dims(static_cast(split_values.size())); split_initializer_proto.mutable_int64_data()->Add(split_values.begin(), split_values.end()); - NodeArg* split_initializer_arg = &graph_utils::AddInitializerWithExternalData(graph, split_initializer_proto); + NodeArg* split_initializer_arg = &graph_utils::AddInitializer(graph, split_initializer_proto); const auto split_node_name = graph.GenerateNodeName(nodes_to_fuse[0].get().Name() + "/GatherSliceToSplitFusion"); Node& split_node = graph.AddNode(split_node_name, "Split", "Split for Fused Gather nodes", {graph.GetNodeArg(node_arg->Name()), split_initializer_arg}, split_outputs); @@ -359,7 +359,7 @@ Status GatherToSliceFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le unsqueeze_axes_initializer_proto.add_dims(static_cast(1)); unsqueeze_axes_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); unsqueeze_axes_initializer_proto.add_int64_data(static_cast(0)); - NodeArg* unsqueeze_axes_arg = &graph_utils::AddInitializerWithExternalData(graph, unsqueeze_axes_initializer_proto); + NodeArg* unsqueeze_axes_arg = &graph_utils::AddInitializer(graph, unsqueeze_axes_initializer_proto); for (size_t i = 0; i < range_input_defs.size(); ++i) { Node& unsqueeze_node = graph.AddNode(graph.GenerateNodeName("Unsqueeze_" + std::to_string(i)), "Unsqueeze", @@ -386,7 +386,7 @@ Status GatherToSliceFusion::ApplyImpl(Graph& graph, bool& modified, int graph_le } else { slice_axes_initializer_proto.add_int32_data(static_cast(axis)); } - NodeArg* slice_axes_arg = &graph_utils::AddInitializerWithExternalData(graph, slice_axes_initializer_proto); + NodeArg* slice_axes_arg = &graph_utils::AddInitializer(graph, slice_axes_initializer_proto); Node& slice_node = graph.AddNode(graph.GenerateNodeName("Slice"), "Slice", "Slice for Fused Gather nodes", {gather_node.MutableInputDefs()[0], unsqueeze_outputs[0], unsqueeze_outputs[1], slice_axes_arg, unsqueeze_outputs[2]}, diff --git a/onnxruntime/core/optimizer/matmul_add_fusion.cc b/onnxruntime/core/optimizer/matmul_add_fusion.cc index 761fe1854274e..fed72db71332a 100644 --- a/onnxruntime/core/optimizer/matmul_add_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_add_fusion.cc @@ -194,7 +194,7 @@ Status MatMulAddFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, shape_initializer_proto.add_dims(static_cast(shape.size())); shape_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); utils::SetRawDataInTensorProto(shape_initializer_proto, shape.data(), shape.size() * sizeof(int64_t)); - NodeArg* shape_arg = &graph_utils::AddInitializerWithExternalData(graph, shape_initializer_proto); + NodeArg* shape_arg = &graph_utils::AddInitializer(graph, shape_initializer_proto); ONNX_NAMESPACE::TypeProto new_arg_type; const ONNX_NAMESPACE::TensorProto_DataType element_type = static_cast( gemm_input_defs[0]->TypeAsProto()->tensor_type().elem_type()); diff --git a/onnxruntime/core/optimizer/matmul_bn_fusion.cc b/onnxruntime/core/optimizer/matmul_bn_fusion.cc index 725cb3fc33f04..367fb42d7928d 100644 --- a/onnxruntime/core/optimizer/matmul_bn_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_bn_fusion.cc @@ -212,14 +212,14 @@ Status MatmulBNFusion::Apply(Graph& graph, Node& matmul_node, RewriteRuleEffect& matmul_b.ToProto(new_gemm_b_tensor); const std::string new_gemm_b_name = graph.GenerateNodeArgName("MatMulBnFusion_GemmB_" + matmul_b_tensor->name()); new_gemm_b_tensor.set_name(new_gemm_b_name); - NodeArg& new_gemm_b_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_gemm_b_tensor); + NodeArg& new_gemm_b_node_arg = graph_utils::AddInitializer(graph, new_gemm_b_tensor); // create bias tensorProto for new Gemm node from initializer. ONNX_NAMESPACE::TensorProto new_gemm_bias_tensor; bias.ToProto(new_gemm_bias_tensor); const std::string new_gemm_bias_name = graph.GenerateNodeArgName("MatMulBnFusion_GemmBias"); new_gemm_bias_tensor.set_name(new_gemm_bias_name); - NodeArg& new_gemm_bias_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_gemm_bias_tensor); + NodeArg& new_gemm_bias_node_arg = graph_utils::AddInitializer(graph, new_gemm_bias_tensor); Node& gemm_node = graph.AddNode( graph.GenerateNodeArgName("MatMulBnFusion_Gemm"), diff --git a/onnxruntime/core/optimizer/nchwc_transformer.cc b/onnxruntime/core/optimizer/nchwc_transformer.cc index 335209dbfadaf..f094a48e10c33 100644 --- a/onnxruntime/core/optimizer/nchwc_transformer.cc +++ b/onnxruntime/core/optimizer/nchwc_transformer.cc @@ -437,7 +437,7 @@ void NchwcTransformerImpl::TransformConv(Node& node) { nchwc_conv_W_tensor_proto.add_dims(conv_W_dims[i]); } - nchwc_conv_W_arg = &graph_utils::AddInitializerWithExternalData(graph_, nchwc_conv_W_tensor_proto); + nchwc_conv_W_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_W_tensor_proto); filters_map->emplace(input_defs[1], nchwc_conv_W_arg); } @@ -464,7 +464,7 @@ void NchwcTransformerImpl::TransformConv(Node& node) { nchwc_conv_B_tensor_proto.add_dims(nchwc_output_channels); - nchwc_conv_B_arg = &graph_utils::AddInitializerWithExternalData(graph_, nchwc_conv_B_tensor_proto); + nchwc_conv_B_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_B_tensor_proto); aligned_biases_.emplace(input_defs[2], nchwc_conv_B_arg); } } @@ -580,7 +580,7 @@ Node& NchwcTransformerImpl::InsertReshape(NodeArg* input_arg, } shape_tensor_proto.add_dims(split_channels ? kNchwcDims + 1 : kNchwcDims); - shape_arg = &graph_utils::AddInitializerWithExternalData(graph_, shape_tensor_proto); + shape_arg = &graph_utils::AddInitializer(graph_, shape_tensor_proto); } Node& reshape_node = graph_.AddNode(graph_.GenerateNodeName("Reshape"), @@ -892,7 +892,7 @@ void NchwcTransformerImpl::TransformBatchNormalization(Node& node) { nchwc_conv_W_tensor_proto.add_dims(1); nchwc_conv_W_tensor_proto.add_dims(1); - auto* nchwc_conv_W_arg = &graph_utils::AddInitializerWithExternalData(graph_, nchwc_conv_W_tensor_proto); + auto* nchwc_conv_W_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_W_tensor_proto); std::copy_n(bn_B.data(), channels, padded_buffer.data()); @@ -903,7 +903,7 @@ void NchwcTransformerImpl::TransformBatchNormalization(Node& node) { gsl::narrow(nchwc_channels) * sizeof(float)); nchwc_conv_B_tensor_proto.add_dims(nchwc_channels); - auto* nchwc_conv_B_arg = &graph_utils::AddInitializerWithExternalData(graph_, nchwc_conv_B_tensor_proto); + auto* nchwc_conv_B_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_B_tensor_proto); // Create the replacement node. std::string nchwc_node_name = graph_.GenerateNodeName(output_defs[0]->Name() + "_bn_nchwc"); diff --git a/onnxruntime/core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.cc b/onnxruntime/core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.cc index 42cd31b5bd7b4..42d27de632b91 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.cc @@ -130,22 +130,22 @@ static bool TryConvertDynamicQuantizeLSTM(Node& op_node, Graph& graph, const log weights_proto_u8.set_name(weight_tensor_proto->name() + "_s8_2_u8"); weights_proto_u8.mutable_dims()->CopyFrom(weight_tensor_proto->dims()); utils::SetRawDataInTensorProto(weights_proto_u8, w_temp.data(), static_cast(w_temp.size())); - input_defs[w_idx] = &graph_utils::AddInitializerWithExternalData(graph, weights_proto_u8); + input_defs[w_idx] = &graph_utils::AddInitializer(graph, weights_proto_u8); ONNX_NAMESPACE::TensorProto weight_zp_proto_u8; QDQ::Int8TensorProto2Uint8(weight_zp_tensor_proto, weight_zp_proto_u8, graph, true); - input_defs[w_zp_idx] = &graph_utils::AddInitializerWithExternalData(graph, weight_zp_proto_u8); + input_defs[w_zp_idx] = &graph_utils::AddInitializer(graph, weight_zp_proto_u8); ONNX_NAMESPACE::TensorProto r_proto_u8; r_proto_u8.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8); r_proto_u8.set_name(r_tensor_proto->name() + "_s8_2_u8"); r_proto_u8.mutable_dims()->CopyFrom(r_tensor_proto->dims()); utils::SetRawDataInTensorProto(r_proto_u8, r_temp.data(), static_cast(r_temp.size())); - input_defs[r_idx] = &graph_utils::AddInitializerWithExternalData(graph, r_proto_u8); + input_defs[r_idx] = &graph_utils::AddInitializer(graph, r_proto_u8); ONNX_NAMESPACE::TensorProto r_zp_proto_u8; QDQ::Int8TensorProto2Uint8(r_zp_tensor_proto, r_zp_proto_u8, graph, true); - input_defs[r_zp_idx] = &graph_utils::AddInitializerWithExternalData(graph, r_zp_proto_u8); + input_defs[r_zp_idx] = &graph_utils::AddInitializer(graph, r_zp_proto_u8); return true; } diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_s8_to_u8.cc b/onnxruntime/core/optimizer/qdq_transformer/qdq_s8_to_u8.cc index 98c818b0c761b..828165e99d840 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_s8_to_u8.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_s8_to_u8.cc @@ -61,7 +61,7 @@ static bool QDQ_S8_to_U8(Graph& graph, Node& q_node, Node& dq_node) { zp_tensor_proto_u8.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8); zp_tensor_proto_u8.set_name(graph.GenerateNodeArgName("qdq_s8_to_u8_zp_conversion")); utils::SetRawDataInTensorProto(zp_tensor_proto_u8, &q_zp_value, sizeof(uint8_t)); - NodeArg* zp_u8_arg = &graph_utils::AddInitializerWithExternalData(graph, zp_tensor_proto_u8); + NodeArg* zp_u8_arg = &graph_utils::AddInitializer(graph, zp_tensor_proto_u8); auto q_output_node_arg_name = graph.GenerateNodeArgName("qdq_s8_to_u8_quant"); NodeArg* q_output_arg = &graph.GetOrCreateNodeArg(q_output_node_arg_name, nullptr); diff --git a/onnxruntime/core/optimizer/qdq_transformer/s8_to_u8.cc b/onnxruntime/core/optimizer/qdq_transformer/s8_to_u8.cc index 616144c0ccde0..f094f3c199f2a 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/s8_to_u8.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/s8_to_u8.cc @@ -43,12 +43,12 @@ bool ConvertS8WeightToU8(Graph& graph, Node& op_node, // The weights fits into S7, overflow is not a problem, no need to convert to U8 return false; } - input_defs[weights_idx] = &graph_utils::AddInitializerWithExternalData(graph, weights_proto_u8); + input_defs[weights_idx] = &graph_utils::AddInitializer(graph, weights_proto_u8); // Convert weight zero point to uint8 ONNX_NAMESPACE::TensorProto weight_zp_proto_u8; Int8TensorProto2Uint8(weight_zp_tensor_proto, weight_zp_proto_u8, graph, true); - input_defs[weight_zp_idx] = &graph_utils::AddInitializerWithExternalData(graph, weight_zp_proto_u8); + input_defs[weight_zp_idx] = &graph_utils::AddInitializer(graph, weight_zp_proto_u8); return true; } diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc index dce69e2913582..34d7ba3c79775 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc @@ -439,23 +439,23 @@ Status DQMatMulToMatMulNBitsAction::ProcessNewNode(Graph& graph, } } - auto weight_T_tp = utils::TensorToTensorProto(weight_dst, weight_dst_name, true); - auto scale_T_tp = utils::TensorToTensorProto(scale_dst, scale_dst_name, true); + auto weight_T_tp = utils::TensorToTensorProto(weight_dst, weight_dst_name, false); + auto scale_T_tp = utils::TensorToTensorProto(scale_dst, scale_dst_name, false); std::optional zp_T_tp; if (zp_dst) { - zp_T_tp.emplace(utils::TensorToTensorProto(*zp_dst, zp_dst_name, true)); + zp_T_tp.emplace(utils::TensorToTensorProto(*zp_dst, zp_dst_name, false)); } auto& input_defs = replacement_node.MutableInputDefs(); - input_defs.push_back(&graph_utils::AddInitializerWithExternalData(graph, weight_T_tp, std::move(weight_dst))); + input_defs.push_back(&graph_utils::AddInitializer(graph, weight_T_tp)); replacement_node.MutableInputArgsCount().push_back(1); - input_defs.push_back(&graph_utils::AddInitializerWithExternalData(graph, scale_T_tp, std::move(scale_dst))); + input_defs.push_back(&graph_utils::AddInitializer(graph, scale_T_tp)); replacement_node.MutableInputArgsCount().push_back(1); if (zp_T_tp) { - input_defs.push_back(&graph_utils::AddInitializerWithExternalData(graph, zp_T_tp.value(), std::move(*zp_dst))); + input_defs.push_back(&graph_utils::AddInitializer(graph, zp_T_tp.value())); replacement_node.MutableInputArgsCount().push_back(1); } diff --git a/onnxruntime/core/optimizer/qdq_transformer/weight_bias_quantization.cc b/onnxruntime/core/optimizer/qdq_transformer/weight_bias_quantization.cc index aa6f9c5409de7..8caa67f266266 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/weight_bias_quantization.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/weight_bias_quantization.cc @@ -131,14 +131,14 @@ Status WeightBiasQuantization::ApplyImpl(Graph& graph, bool& modified, int graph weight_scale_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_weight_scale")); weight_scale_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); weight_scale_proto.mutable_float_data()->Add(scale); - weight_scale_arg = &graph_utils::AddInitializerWithExternalData(graph, weight_scale_proto); + weight_scale_arg = &graph_utils::AddInitializer(graph, weight_scale_proto); // Weight zero point initializer. ONNX_NAMESPACE::TensorProto weight_zp_proto; weight_zp_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_weight_zp")); weight_zp_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT8); weight_zp_proto.mutable_int32_data()->Add(static_cast(zp)); - NodeArg& weight_zp_arg = graph_utils::AddInitializerWithExternalData(graph, weight_zp_proto); + NodeArg& weight_zp_arg = graph_utils::AddInitializer(graph, weight_zp_proto); // Q from float32 to int8. ONNX_NAMESPACE::TypeProto weight_q_type_proto; diff --git a/onnxruntime/core/optimizer/relu_clip_fusion.cc b/onnxruntime/core/optimizer/relu_clip_fusion.cc index efd7022ab764b..07902fde04930 100644 --- a/onnxruntime/core/optimizer/relu_clip_fusion.cc +++ b/onnxruntime/core/optimizer/relu_clip_fusion.cc @@ -97,7 +97,7 @@ Status FuseReluClip::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff mutable_next_node->AddAttribute("min", 0.f); } else { // Add the initialized tensor to the graph - auto* replacement_min_nodearg = &graph_utils::AddInitializerWithExternalData(graph, replacement_min); + auto* replacement_min_nodearg = &graph_utils::AddInitializer(graph, replacement_min); // Replace the input def at the appropriate index of the Clip node auto& mutable_input_defs = mutable_next_node->MutableInputDefs(); diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 36213609f6b61..324905f953eec 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -438,7 +438,7 @@ bool ReshapeFusion::Fuse_Subgraph(Node& reshape, Graph& graph, const logging::Lo shape_initializer_proto.add_dims(static_cast(shape_value.size())); shape_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); utils::SetRawDataInTensorProto(shape_initializer_proto, shape_value.data(), shape_value.size() * sizeof(int64_t)); - auto& new_node_arg = graph_utils::AddInitializerWithExternalData(graph, shape_initializer_proto); + auto& new_node_arg = graph_utils::AddInitializer(graph, shape_initializer_proto); // Safely remove concat parent nodes which have only one output for (int i = 0; i < concat_input_count; ++i) { @@ -492,7 +492,7 @@ bool ReshapeFusion::FuseContiguousReshapes(Node& reshape, Graph& graph) { shape_initializer_proto.add_dims(static_cast(shape_value.size())); shape_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); utils::SetRawDataInTensorProto(shape_initializer_proto, shape_value.data(), shape_value.size() * sizeof(int64_t)); - NodeArg* shape_arg = &graph_utils::AddInitializerWithExternalData(graph, shape_initializer_proto); + NodeArg* shape_arg = &graph_utils::AddInitializer(graph, shape_initializer_proto); Node& reshape_node = graph.AddNode(graph.GenerateNodeName(name + "_new_reshape"), "Reshape", "Reshape for " + name, {contiguous_reshapes[0].get().MutableInputDefs()[0], shape_arg}, {contiguous_reshapes.back().get().MutableOutputDefs()[0]}); diff --git a/onnxruntime/core/optimizer/stft_decomposition.cc b/onnxruntime/core/optimizer/stft_decomposition.cc index 74121508132dc..5c09e5225ab9c 100644 --- a/onnxruntime/core/optimizer/stft_decomposition.cc +++ b/onnxruntime/core/optimizer/stft_decomposition.cc @@ -46,7 +46,7 @@ NodeArg* AddInitializer(Graph& graph, const char* name, const int64_t (&shape)[T proto.add_dims(shape[i]); } utils::SetRawDataInTensorProto(proto, begin, element_count * sizeof(TDataType)); - return &graph_utils::AddInitializerWithExternalData(graph, proto); + return &graph_utils::AddInitializer(graph, proto); } template diff --git a/onnxruntime/core/optimizer/transformer_memcpy.cc b/onnxruntime/core/optimizer/transformer_memcpy.cc index a320de2ee7a13..cc7682b2b418d 100644 --- a/onnxruntime/core/optimizer/transformer_memcpy.cc +++ b/onnxruntime/core/optimizer/transformer_memcpy.cc @@ -383,21 +383,7 @@ bool TransformerMemcpyImpl::ProcessInitializers(const KernelRegistryManager& ker TensorProto new_tensor_proto = *tensor_proto; *(new_tensor_proto.mutable_name()) = new_def_name; - // Query any OrtValue existing for the original initializer - // We are checking outer scope because GetInitializer is called with true, therefore, we potentially - // have references to parent graphs. - // We are doing this so the same OrtValue is re-used in subgraphs and no copies made for big items. - constexpr const bool check_outer_scope_true = true; - OrtValue ort_value; - // The initializer can be in memory with OrtValue or it can be a flatbuffer mapped. - if (utils::HasExternalDataInMemory(new_tensor_proto) && - graph_.GetOrtValueInitializer(name, ort_value, check_outer_scope_true)) { - // Re-use the same ort_value and proto that points to the same buffer - ORT_IGNORE_RETURN_VALUE(graph_utils::AddInitializerWithExternalData(graph_, new_tensor_proto, - std::move(ort_value))); - } else { - ORT_IGNORE_RETURN_VALUE(graph_utils::AddInitializer(graph_, new_tensor_proto)); - } + ORT_IGNORE_RETURN_VALUE(graph_utils::AddInitializer(graph_, new_tensor_proto)); replacements.insert(std::make_pair(provider_def, &new_def)); } diff --git a/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc b/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc index 48ea54434b805..3a95d2a53e8f5 100644 --- a/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc +++ b/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc @@ -586,10 +586,10 @@ void ApiGraph::TransposeInitializer(std::string_view name, const std::vector& shape) { @@ -622,7 +622,7 @@ void ApiGraph::ReshapeInitializer(std::string_view name, const std::vector()->Reshape(new_shape); - } - - auto& new_node_arg = graph_utils::AddInitializerWithExternalData(graph, new_tensor_proto, ort_value); + auto& new_node_arg = graph_utils::AddInitializer(graph, new_tensor_proto); graph_utils::ReplaceNodeWithInitializer(graph, node, new_node_arg); // Remove the Unsqueeze node and replace it with the initializer. diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index e8d133779f33c..51a8b13cd8261 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -734,6 +734,10 @@ struct CudaEpFactory : OrtEpFactory { } */ + // guard against bad device discovery. max devices we expect to add is num_cuda_devices. if we're attempting + // to add more than that we have duplicates in the `devices` array. + max_ep_devices = std::min(max_ep_devices, static_cast(num_cuda_devices)); + int16_t device_id = 0; for (size_t i = 0; i < num_devices && num_ep_devices < max_ep_devices; ++i) { const OrtHardwareDevice& device = *devices[i]; diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 286db9070766d..cc9d9f3da1d81 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -123,10 +123,11 @@ void* OutputAllocator::reallocateOutputAsync(char const* /*tensorName*/, void* / // even for empty tensors, so allocate a dummy byte. size = std::max(size, static_cast(1)); if (size > allocated_size) { - cudaFree(outputPtr); + alloc_->Free(alloc_, outputPtr); outputPtr = nullptr; allocated_size = 0; - if (cudaMalloc(&outputPtr, size) == cudaSuccess) { + outputPtr = alloc_->Alloc(alloc_, size); + if (outputPtr) { allocated_size = size; } } @@ -352,193 +353,6 @@ bool ApplyProfileShapesFromProviderOptions(std::vector shape values" for the INT32 shape tensor input across this inference run - * @param shape_tensor_values_int64 holds "shape tensor -> shape values" for the INT64 shape tensor input across this inference run - */ -Status ApplyProfileShapesFromInputTensorValue(std::vector& trt_profiles, - Ort::KernelContext ctx, - nvinfer1::ITensor* input, - ShapeRangesMap& shape_ranges, - const std::unordered_map& input_indexes, - std::unordered_map>& shape_tensor_values, - std::unordered_map>& shape_tensor_values_int64, - cudaStream_t stream, - bool* engine_update) { - for (size_t i = 0; i < trt_profiles.size(); i++) { - const std::string& input_name = input->getName(); - nvinfer1::Dims dims = input->getDimensions(); - int nb_dims = dims.nbDims; - - size_t input_index = 0; - const auto& iter = input_indexes.find(input_name); - if (iter != input_indexes.end()) { - input_index = iter->second; - } - - auto input_tensor = ctx.GetInput(input_index); - auto tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); - const auto tensor_shapes = tensor_info.GetShape(); - auto& shape_ranges_per_input = shape_ranges[input_name]; - - auto trt_profile = trt_profiles[i]; - - // If there are multiple profiles, for second and rest of profiles, simply copy the min/max/opt profile values from the first profile. - // Following "if statement" won't be executed since TRT EP currently only allows single profile for non-explicit profiles case. - if (i > 0) { - if (input->isShapeTensor()) { - // shape tensor - int shape_size = nb_dims == 0 ? 1 : static_cast(tensor_shapes[0]); - std::vector shapes_min(shape_size), shapes_opt(shape_size), shapes_max(shape_size); - for (int j = 0; j < shape_size; j++) { - shapes_min[j] = *(trt_profiles[0]->getShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN)); - shapes_max[j] = *(trt_profiles[0]->getShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX)); - shapes_opt[j] = *(trt_profiles[0]->getShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT)); - } - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, &shapes_min[0], shape_size); - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, &shapes_max[0], shape_size); - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, &shapes_opt[0], shape_size); - } else { - // execution tensor - nvinfer1::Dims dims_min, dims_opt, dims_max; - dims_min = trt_profiles[0]->getDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN); - dims_max = trt_profiles[0]->getDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX); - dims_opt = trt_profiles[0]->getDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, dims_min); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, dims_max); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, dims_opt); - } - continue; - } - - // Create shape profile - if (input->isShapeTensor()) { - // Get shape values for shape tensor input - const auto tensor_type = tensor_info.GetElementType(); - // The shape of the "shape tensor" is either zero dimension (scalar) or 1-dimension - int shape_size = dims.nbDims == 0 ? 1 : static_cast(tensor_shapes[0]); - // For setting TRT optimization profile. (Note: the min/opt/max profile values are still int32 even though int64 is supported after TRT 10) - std::vector values(shape_size); - - switch (tensor_type) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: { - auto buffer = std::make_unique(shape_size); - auto status = GetShapeOfShapeTensor(input_tensor, buffer.get(), shape_size, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - shape_tensor_values[input_name].resize(shape_size); - for (int j = 0; j < shape_size; ++j) { - shape_tensor_values[input_name][j] = buffer[j]; - values[j] = buffer[j]; - } - break; - } - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: { - auto buffer = std::make_unique(shape_size); - auto status = GetShapeOfShapeTensor(input_tensor, buffer.get(), shape_size, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - shape_tensor_values_int64[input_name].resize(shape_size); - for (int j = 0; j < shape_size; ++j) { - shape_tensor_values_int64[input_name][j] = buffer[j]; - values[j] = static_cast(buffer[j]); - } - break; - } - default: { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT shape tensor data type: " + std::to_string(tensor_type) + " not supported."); - } - } - - // Update shape ranges - std::vector shapes_min(shape_size), shapes_opt(shape_size), shapes_max(shape_size); - int shape_range_size = static_cast(shape_ranges_per_input.size()); - if (shape_size == shape_range_size) { - // If shape size matches, check/update shape range - for (int j = 0; j < shape_size; ++j) { - auto& shape_range = shape_ranges_per_input[j][0]; // only has one profile - shapes_min[j] = static_cast(shape_range[0]); - shapes_max[j] = static_cast(shape_range[1]); - shapes_opt[j] = static_cast(shape_range[2]); - - const auto& tensor_shape_value = values[j]; - // Update shape range lower bound - if (tensor_shape_value < shape_range[0]) { - shape_range[0] = tensor_shape_value; - shapes_min[j] = tensor_shape_value; - *engine_update = true; - } - // Update shape range upper bound - if (tensor_shape_value > shape_range[1]) { - shape_range[1] = tensor_shape_value; - shape_range[2] = tensor_shape_value; - shapes_max[j] = tensor_shape_value; - shapes_opt[j] = tensor_shape_value; - *engine_update = true; - } - } - } else { - // If shape size doesn't match, initialize shape_range with the new shape value - shape_ranges_per_input.clear(); - for (int j = 0; j < shape_size; ++j) { - const auto& tensor_shape_value = values[j]; - std::vector> profile_vector; - std::vector shape_vector{tensor_shape_value, tensor_shape_value, tensor_shape_value}; - profile_vector.push_back(shape_vector); // only one profile needed - shape_ranges_per_input[j] = profile_vector; - shapes_min[j] = tensor_shape_value; - shapes_opt[j] = tensor_shape_value; - shapes_max[j] = tensor_shape_value; - } - *engine_update = true; - } - - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, &shapes_min[0], shape_size); - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, &shapes_max[0], shape_size); - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, &shapes_opt[0], shape_size); - } else { // Execution tensor - nvinfer1::Dims dims_min(dims), dims_opt(dims), dims_max(dims); - for (int j = 0, end = nb_dims; j < end; ++j) { - const auto& tensor_shape = tensor_shapes[j]; - if (shape_ranges_per_input.find(j) != shape_ranges_per_input.end()) { - auto& shape_range = shape_ranges_per_input[j][0]; // only has one profile - dims_min.d[j] = static_cast(shape_range[0]); - dims_max.d[j] = static_cast(shape_range[1]); - dims_opt.d[j] = static_cast(shape_range[2]); - - // Update minimum dimension - if (tensor_shape < shape_range[0]) { - shape_range[0] = tensor_shape; - dims_min.d[j] = static_cast(tensor_shape); - *engine_update = true; - } - // Update maximum dimension - if (tensor_shape > shape_range[1]) { - shape_range[1] = tensor_shape; - shape_range[2] = tensor_shape; - dims_max.d[j] = static_cast(tensor_shape); - dims_opt.d[j] = static_cast(tensor_shape); - *engine_update = true; - } - } - } - - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, dims_min); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, dims_max); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, dims_opt); - } - } - return Status::OK(); -} - #define CASE_GET_INPUT_TENSOR(DATA_TYPE, SrcT) \ case DATA_TYPE: { \ auto input_tensor_ptr = input_tensor.GetTensorData(); \ @@ -554,6 +368,7 @@ Status ApplyProfileShapesFromInputTensorValue(std::vector(); \ + skip_input_binding_allowed = false; \ if (input_tensor_ptr != nullptr && elem_cnt > 0) { \ scratch_buffers.push_back(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, elem_cnt * sizeof(DstT))); \ data = scratch_buffers.back().get(); \ @@ -568,6 +383,7 @@ Status ApplyProfileShapesFromInputTensorValue(std::vector(); \ + data_ptr = output_tensor_ptr; \ if (output_tensor_ptr != nullptr && elem_cnt > 0) { \ buffers[output_name] = output_tensor_ptr; \ } else { \ @@ -580,6 +396,8 @@ Status ApplyProfileShapesFromInputTensorValue(std::vector(); \ + data_ptr = output_tensor_ptr; \ + skip_output_binding_allowed = false; \ if (output_tensor_ptr != nullptr && elem_cnt > 0) { \ scratch_buffers.push_back(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, elem_cnt * sizeof(DstT))); \ buffers[output_name] = scratch_buffers.back().get(); \ @@ -628,7 +446,8 @@ Status BindContextInput(Ort::KernelContext& ctx, std::unordered_map>& shape_tensor_values_int64, std::vector>& scratch_buffers, OrtAllocator* alloc, - cudaStream_t stream) { + cudaStream_t stream, + bool& skip_input_binding_allowed) { auto input_tensor = ctx.GetInput(input_index); auto tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); const auto tensor_shapes = tensor_info.GetShape(); @@ -647,7 +466,7 @@ Status BindContextInput(Ort::KernelContext& ctx, if (trt_engine->isShapeInferenceIO(input_name)) { // Bind "shape tensor" input buffer - + skip_input_binding_allowed = false; // Shape tensor input binding cannot be skipped // The shape of the "shape tensor" is either zero dimension (scalar) or 1-dimension int shape_size = trt_engine->getTensorShape(input_name).nbDims == 0 ? 1 : static_cast(tensor_shapes[0]); switch (tensor_type) { @@ -775,19 +594,20 @@ Status BindContextOutput(Ort::KernelContext& ctx, DDSOutputAllocatorMap& dds_output_allocator_map, std::vector>& scratch_buffers, OrtAllocator* alloc, - std::unordered_map& buffers) { + std::unordered_map& buffers, + nvinfer1::Dims& dims, + void*& data_ptr, + bool& skip_output_binding_allowed) { // Get output shape - nvinfer1::Dims dims = trt_context->getTensorShape(output_name); + dims = trt_context->getTensorShape(output_name); int nb_dims = dims.nbDims; bool is_DDS = false; - std::vector output_shapes(nb_dims); for (int j = 0, end = nb_dims; j < end; ++j) { // data-dependent shape if (dims.d[j] == -1) { is_DDS = true; break; } - output_shapes[j] = dims.d[j]; } auto known_DDS = dds_output_allocator_map.find(output_name) != dds_output_allocator_map.end(); @@ -800,16 +620,19 @@ Status BindContextOutput(Ort::KernelContext& ctx, // Otherwise, if the shape of the output tensor is known prior to the runtime, ORT will pre-allocate memory buffer for the output tensor for enqueueV3. if (is_DDS || known_DDS) { if (!known_DDS) { - auto allocatorPtr = std::make_unique(); + auto allocatorPtr = std::make_unique(alloc); trt_context->setOutputAllocator(output_name, allocatorPtr.get()); dds_output_allocator_map[output_name] = std::move(allocatorPtr); + dims.nbDims = -1; // Set to -1 to indicate that the shape is not known at this point. + data_ptr = nullptr; // Set data_ptr to nullptr for DDS output binding. } } else { - output_tensors[i] = ctx.GetOutput(output_index, output_shapes); + output_tensors[i] = ctx.GetOutput(output_index, dims.d, nb_dims); auto& output_tensor = output_tensors[i]; const auto elem_cnt = output_tensor.GetTensorTypeAndShapeInfo().GetElementCount(); switch (output_type) { + // below macros set data_ptr and skip_output_binding_allowed variables CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, float) CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, uint16_t) CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16, uint16_t) @@ -840,7 +663,6 @@ Status BindContextOutput(Ort::KernelContext& ctx, * we are waiting for ORT core to support "assign" memory address to ORT context output. Some works need to be done in ORT memory planner to be aware of this memory support. */ Status BindKernelOutput(Ort::KernelContext& ctx, - OrtMemoryInfo* /*mem_info*/, DDSOutputAllocatorMap& allocator_map, char const* output_name, size_t output_index, @@ -903,31 +725,6 @@ NvExecutionProvider::PerThreadContext::~PerThreadContext() { trt_context_map_.clear(); } -/* - * Returns true if the shape ranges maintained by the PerThreadContext is different from the shape ragnes maintained by TRT EP, meaning the - * engine is being updated and the execution context maintained by the PerThreadContext should be updated as well. Otherwise, returns false. - * - */ -bool NvExecutionProvider::PerThreadContext::CompareProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges) { - if (shape_ranges.size() > 0) { - if (input_shape_ranges_[fused_node] != shape_ranges) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] The shape ranges maintained by the PerThreadContext is different from the shape ranges maintained by TRT EP. \ - This means the engine is updated and will need to update the execution context as well."; - return true; - } - } - return false; -} - -/* - * Updates the shape ranges maintained by the PerThreadContext. - * As long as the execution context maintained by the PerThreadContext is updated, the associated shape ranges should be updated as well. - * - */ -void NvExecutionProvider::PerThreadContext::UpdateProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges) { - input_shape_ranges_[fused_node] = shape_ranges; -} - void NvExecutionProvider::PerThreadContext::ResetTensorRTContext(std::string fused_node) { auto it = trt_context_map_.find(fused_node); if (it != trt_context_map_.end()) { @@ -1081,7 +878,6 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) engine_decryption_lib_path_ = info.engine_decryption_lib_path; } force_sequential_engine_build_ = info.force_sequential_engine_build; - context_memory_sharing_enable_ = info.context_memory_sharing_enable; sparsity_enable_ = info.sparsity_enable; auxiliary_streams_ = info.auxiliary_streams; profile_min_shapes = info.profile_min_shapes; @@ -1225,7 +1021,6 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) << ", nv_engine_decryption_enable: " << engine_decryption_enable_ << ", nv_engine_decryption_lib_path: " << engine_decryption_lib_path_ << ", nv_force_sequential_engine_build: " << force_sequential_engine_build_ - << ", nv_context_memory_sharing_enable: " << context_memory_sharing_enable_ << ", nv_sparsity_enable: " << sparsity_enable_ << ", nv_auxiliary_streams: " << auxiliary_streams_ << ", nv_cuda_graph_enable: " << cuda_graph_enable_ @@ -1298,9 +1093,15 @@ void NvExecutionProvider::IncrementRegularRunCountBeforeGraphCapture() { } std::vector NvExecutionProvider::CreatePreferredAllocators() { + OrtArenaCfg arena_cfg(0, static_cast(ArenaExtendStrategy::kSameAsRequested), + -1, -1, -1, -1); AllocatorCreationInfo default_memory_info( [](OrtDevice::DeviceId device_id) { return std::make_unique(device_id, CUDA); }, - narrow(device_id_)); + narrow(device_id_), + true, + arena_cfg, + // make it stream aware + true); AllocatorCreationInfo pinned_allocator_info( [](OrtDevice::DeviceId device_id) { @@ -2244,6 +2045,96 @@ common::Status NvExecutionProvider::Compile(const std::vector return Status::OK(); } +/** + * @brief Determines whether I/O binding is required for TensorRT execution. + * + * This function optimizes TensorRT inference performance by determining when tensor + * input/output binding operations can be skipped. Binding is an expensive operation + * that involves setting up tensor pointers in the TensorRT execution context, so + * avoiding unnecessary rebinding can significantly improve inference throughput. + * + * The function implements a three-tier decision logic: + * 1. First run: Always requires binding to establish initial tensor mappings + * 2. Subsequent runs with optimization allowed: Only rebind if tensors have changed + * 3. Subsequent runs without optimization: Always rebind for safety + * + * @tparam TRTState The TensorRT state type (TensorrtFuncState or TensorrtShortFuncState) + * @param trt_state Pointer to the TensorRT execution state containing tensor cache + * and configuration flags + * @param ctx ONNX Runtime kernel context providing access to current input tensors + * + * @return true if I/O binding is required (tensors changed or safety conditions apply), + * false if binding can be safely skipped (optimization enabled and tensors unchanged) + * + * @note This function modifies trt_state by: + * - Setting is_first_run to false after first execution + * - Caching current tensor parameters in input_tensors vector + * - Updating cached tensors when changes are detected + * + * @warning The skip_io_binding_allowed flag must be carefully managed as incorrect + * usage can lead to inference with stale tensor bindings and incorrect results. + */ +template +static bool IsIOBindingRequired(TRTState* const trt_state, const Ort::KernelContext& ctx) { + // Check if input tensors have changed since the last run + // If so, we need to bind input tensors again + bool require_io_binding = false; + + if (trt_state->is_first_run) { + // If this is the first run, we always bind input tensors + require_io_binding = true; + auto input_tensor_count = ctx.GetInputCount(); + auto output_tensor_count = ctx.GetOutputCount(); + trt_state->input_tensors.resize(input_tensor_count); + trt_state->output_tensors.resize(output_tensor_count); + for (size_t input_index = 0; input_index < input_tensor_count; ++input_index) { + const auto& input_tensor = ctx.GetInput(input_index); + const auto& tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); + + trt_state->input_tensors[input_index] = TensorParams{input_tensor.GetTensorRawData(), tensor_info.GetShape()}; + } + trt_state->is_first_run = false; + } else if (trt_state->skip_io_binding_allowed) { + // If skip_io_binding_allowed is true, we can skip binding if input tensors are the same as before + auto input_tensor_count = ctx.GetInputCount(); + for (size_t input_index = 0; input_index < input_tensor_count; ++input_index) { + const auto& input_tensor = ctx.GetInput(input_index); + const auto& tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); + + TensorParams ip_tensor{input_tensor.GetTensorRawData(), tensor_info.GetShape()}; + + if (ip_tensor != trt_state->input_tensors[input_index]) { + require_io_binding = true; + trt_state->input_tensors[input_index] = ip_tensor; + } + } + } else { + // If this is not the first run and skip_io_binding_allowed is false, we need to bind input tensors + require_io_binding = true; + } + + if (!require_io_binding) { + // no need to bind inputs, but check outputs as well + auto output_tensor_count = ctx.GetOutputCount(); + + for (size_t output_index = 0; output_index < output_tensor_count; ++output_index) { + const auto& prev_output_tensor = trt_state->output_tensors[output_index]; + + if (prev_output_tensor.dims.nbDims != -1) { + const auto& new_output_tensor = ctx.GetOutput(output_index, prev_output_tensor.dims.d, prev_output_tensor.dims.nbDims); + + // different output tensor data means we need to bind outputs again + if (prev_output_tensor.data != new_output_tensor.GetTensorRawData()) { + require_io_binding = true; + break; + } + } + } + } + + return require_io_binding; +} + Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& graph_body_viewer, const Node& fused_node, std::unordered_map& input_map, @@ -2349,21 +2240,6 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr ShapeRangesMap input_explicit_shape_ranges; ShapeRangesMap input_implicit_shape_ranges; - auto tensor_is_dynamic = [&](nvinfer1::ITensor* tensor) -> bool { - if (tensor->isShapeTensor()) { - return true; - } else { - nvinfer1::Dims dims = tensor->getDimensions(); - // Execution tensor - for (int j = 0, end = dims.nbDims; j < end; ++j) { - if (dims.d[j] == -1) { - return true; - } - } - } - return false; - }; - bool has_dynamic_shape = false; // True if input tensor has dynamic shape and no explicit profile is specified, otherwise false if ((!profile_min_shapes_.empty()) && (!profile_max_shapes_.empty()) && (!profile_opt_shapes_.empty())) { has_explicit_profile = true; @@ -2375,7 +2251,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr } else { for (unsigned int i = 0, end = num_inputs; i < end; ++i) { auto input = trt_network->getInput(i); - has_dynamic_shape |= tensor_is_dynamic(input); + has_dynamic_shape |= checkTrtTensorIsDynamic(input); } if (has_dynamic_shape) { LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] No explicit optimization profile was specified. " @@ -2574,31 +2450,18 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // Build context // Note: Creating an execution context from an engine is thread safe per TRT doc // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - if (context_memory_sharing_enable_) { -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) -#endif - size_t mem_size = trt_engine->getDeviceMemorySizeV2(); -#if defined(_MSC_VER) -#pragma warning(pop) -#endif - if (mem_size > max_ctx_mem_size_) { - max_ctx_mem_size_ = mem_size; - } - trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); - } else { - trt_context = std::unique_ptr(trt_engine->createExecutionContext()); - } + trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); if (!trt_context) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Nv EP could not build execution context for fused node: " + fused_node.Name()); } + bool is_dynamic_shape_context = false; // Create input to index map for (int i = 0; i < num_inputs; ++i) { auto input = trt_network->getInput(i); const std::string& input_name = input->getName(); + is_dynamic_shape_context |= checkTrtDimIsDynamic(trt_engine->getTensorShape(input_name.c_str())); const auto& iter = input_map.find(input_name); if (iter != input_map.end()) { input_indexes[input_name] = iter->second; @@ -2639,10 +2502,9 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr input_shape_ranges_[context->node_name], &tensorrt_mu_, trt_node_name_with_precision, engine_cache_enable_, cache_path_, runtime_.get(), profiles_[context->node_name], - context_memory_sharing_enable_, &max_ctx_mem_size_, engine_decryption_enable_, engine_decryption_, engine_encryption_, detailed_build_log_, sparsity_enable_, - auxiliary_streams_, cuda_graph_enable_, cache_prefix_, cache_suffix}; + auxiliary_streams_, cuda_graph_enable_, is_dynamic_shape_context, cache_prefix_, cache_suffix}; *state = p.release(); return 0; }; @@ -2666,25 +2528,20 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr const std::unordered_map& output_indexes = (trt_state->output_info)[0]; const std::unordered_map& output_types = (trt_state->output_info)[1]; auto fused_node_name = trt_state->fused_node_name; - // This map "shape_ranges" contains the shape range info for setting TRT optimization profiles. - // The info is used for both shape tensor and execution tensor: - // tensor name->(dimension->[min, max, opt]) - auto& shape_ranges = trt_state->input_shape_ranges; + std::unordered_map> shape_tensor_values; // This map holds "shape tensor -> shape values" for the shape tensor input across this inference run std::unordered_map> shape_tensor_values_int64; // same as above but for int64 shape tensor input auto& dds_output_allocator_map = this->dds_output_allocator_maps_[fused_node_name]; auto trt_engine = trt_state->engine->get(); auto trt_context = trt_state->context->get(); auto trt_profiles = trt_state->profiles; - auto max_context_mem_size_ptr = trt_state->max_context_mem_size_ptr; - int num_inputs = static_cast(input_indexes.size()); int num_outputs = static_cast(output_indexes.size()); std::unordered_set input_names; - OrtDevice device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, - narrow(device_id_)); - OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, device); if (alloc_ == nullptr) { + OrtDevice device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, + narrow(device_id_)); + OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, device); Ort::ThrowOnError(api->KernelContext_GetAllocator(context, &mem_info, &alloc_)); } OrtAllocator* alloc = alloc_; @@ -2698,68 +2555,13 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Nv EP select an optimization profile for the current context failed"); } - // Name the engine cache based on GPU compute capacity and reduce the chance of loading an incompatible cache - // Note: Engine cache generated on a GPU with large memory might not be loadable on a GPU with smaller memory, even if they share the same compute capacity - // Prepare cache name - std::string cache_path = ""; - // Customize cache prefix if assigned - if (!cache_prefix_.empty()) { - cache_path = GetCachePath(trt_state->engine_cache_path, trt_state->cache_prefix) + trt_state->cache_suffix; - } else { - cache_path = GetCachePath(trt_state->engine_cache_path, trt_state->trt_node_name_with_precision); - } - - // Enable hardware compatility mode if assigned - std::string cache_hw_compat = "_sm" + compute_capability_; - - // Name the engine cache based on GPU compute capacity and reduce the chance of loading an incompatible cache - // Note: Engine cache generated on a GPU with large memory might not be loadable on a GPU with smaller memory, even if they share the same compute capacity - const std::string cache_path_prefix = cache_path + cache_hw_compat; - std::string engine_cache_path = cache_path_prefix + ".engine"; - const std::string encrypted_engine_cache_path = engine_cache_path + ".encrypted"; - const std::string profile_cache_path = cache_path_prefix + ".profile"; - - // If weight-stripped engine is enabled and refitted engine cache is not present, - // TRT EP will use the engine cache with ".stripped.engine" appended to the end. - const std::filesystem::path engine_cache_fs_path = engine_cache_path; - if (weight_stripped_engine_enable_ && !std::filesystem::exists(engine_cache_fs_path)) { - engine_cache_path = cache_path_prefix + ".stripped.engine"; - weight_stripped_engine_refit_ = true; - } - - // Check and update shape ranges for dynamic shape inputs. - for (int i = 0, end = num_inputs; i < end; ++i) { - auto input = trt_state->network->get()->getInput(i); - const std::string& input_name = input->getName(); - input_names.insert(input_name); - - // If there is any input tensor in shape_ranges, it means this input tensor has dynamic shape and its profile shape values have not yet resolved. - // TRT EP will help determine the min/max/opt profile values based on current input tensor value. - if (shape_ranges.find(input_name) != shape_ranges.end()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Nv EP failed to parse input tensor and generate optimization profiles."); - } - } - - if (weight_stripped_engine_refit_) { - auto status = RefitEngine(model_path_, - onnx_model_folder_path_, - engine_cache_path, - false /* path check for security */, - onnx_model_bytestream_, - onnx_model_bytestream_size_, - trt_engine, - false /* serialize refitted engine to disk */, - detailed_build_log_); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - // Check before using trt_engine if (trt_engine == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "No engine is found."); } + bool require_io_binding = IsIOBindingRequired(trt_state, ctx); + // Get input and output binding names int total_bindings = trt_engine->getNbIOTensors(); std::vector input_binding_names, output_binding_names; @@ -2776,23 +2578,25 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr /* * Set input shapes and bind input buffers */ - std::vector> scratch_buffers; - for (size_t i = 0, end = input_binding_names.size(); i < end; ++i) { - char const* input_name = input_binding_names[i]; - - size_t input_index = 0; - const auto iter = input_indexes.find(input_name); - if (iter != input_indexes.end()) { - input_index = iter->second; - } - auto input_tensor = ctx.GetInput(input_index); - auto tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); - const auto tensor_shapes = tensor_info.GetShape(); + auto& scratch_buffers = trt_state->scratch_buffers; + if (require_io_binding) { + scratch_buffers.clear(); + bool skip_input_binding_allowed = true; + for (size_t i = 0, end = input_binding_names.size(); i < end; ++i) { + char const* input_name = input_binding_names[i]; + + size_t input_index = 0; + const auto iter = input_indexes.find(input_name); + if (iter != input_indexes.end()) { + input_index = iter->second; + } - auto status = BindContextInput(ctx, trt_engine, trt_context, input_name, input_index, shape_tensor_values, shape_tensor_values_int64, scratch_buffers, alloc, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + auto status = BindContextInput(ctx, trt_engine, trt_context, input_name, input_index, shape_tensor_values, shape_tensor_values_int64, scratch_buffers, alloc, stream, skip_input_binding_allowed); + if (status != Status::OK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + } } + trt_state->skip_io_binding_allowed = skip_input_binding_allowed; } /* @@ -2806,44 +2610,51 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr std::unordered_map output_dim_sizes; output_dim_sizes.reserve(num_outputs); - for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { - char const* output_name = output_binding_names[i]; + if (require_io_binding) { + bool skip_output_binding_allowed = true; + for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { + char const* output_name = output_binding_names[i]; - size_t output_index = 0; - const auto& index_iter = output_indexes.find(output_name); - if (index_iter != output_indexes.end()) { - output_index = index_iter->second; - } + size_t output_index = 0; + const auto& index_iter = output_indexes.find(output_name); + if (index_iter != output_indexes.end()) { + output_index = index_iter->second; + } - size_t output_type = 0; - const auto type_iter = output_types.find(output_name); - if (type_iter != output_types.end()) { - output_type = type_iter->second; - } + size_t output_type = 0; + const auto type_iter = output_types.find(output_name); + if (type_iter != output_types.end()) { + output_type = type_iter->second; + } + + nvinfer1::Dims dims; + void* data_ptr = nullptr; + + Status status = BindContextOutput(ctx, trt_context, output_name, output_index, output_type, i, output_tensors, output_dim_sizes, + dds_output_allocator_map, scratch_buffers, alloc, buffers, dims, data_ptr, skip_output_binding_allowed); + if (status != Status::OK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + } - Status status = BindContextOutput(ctx, trt_context, output_name, output_index, output_type, i, output_tensors, output_dim_sizes, - dds_output_allocator_map, scratch_buffers, alloc, buffers); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + trt_state->output_tensors[output_index] = TensorParams{data_ptr, dims}; } + + trt_state->skip_io_binding_allowed = trt_state->skip_io_binding_allowed | skip_output_binding_allowed; } // Set execution context memory - if (trt_state->context_memory_sharing_enable) { -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) -#endif + if (require_io_binding) { size_t mem_size = trt_engine->getDeviceMemorySizeV2(); -#if defined(_MSC_VER) -#pragma warning(pop) -#endif - if (mem_size > *max_context_mem_size_ptr) { - *max_context_mem_size_ptr = mem_size; + if (trt_state->is_dynamic_shape) { + mem_size = trt_context->updateDeviceMemorySizeForShapes(); + } + if (trt_state->context_memory_size != mem_size) { + LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] A new context memory was allocated with size " << mem_size; + trt_state->context_memory = IAllocator::MakeUniquePtrFromOrtAllocator(alloc, mem_size, false /*use_reserve*/); + trt_state->context_memory_size = mem_size; + trt_context->setDeviceMemoryV2(trt_state->context_memory.get(), mem_size); } - trt_context->setDeviceMemory(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, *max_context_mem_size_ptr).get()); } - // Start CUDA graph capture. // Note: The reason we don't put graph capture in OnRunStart() like CUDA EP does is because // current ORT TRT doesn't get cuda stream until compute time and graph capture requires cuda stream. @@ -2894,7 +2705,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr if (index_iter != output_indexes.end()) { output_index = index_iter->second; } - auto status = BindKernelOutput(ctx, &mem_info, dds_output_allocator_map, output_name, output_index, output_type, stream); + auto status = BindKernelOutput(ctx, dds_output_allocator_map, output_name, output_index, output_type, stream); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, status.ErrorMessage()); } @@ -2961,33 +2772,19 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra // // Note: Creating an execution context from an engine is thread safe per TRT doc // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - if (context_memory_sharing_enable_) { -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) -#endif - size_t mem_size = trt_engine->getDeviceMemorySizeV2(); -#if defined(_MSC_VER) -#pragma warning(pop) -#endif - if (mem_size > max_ctx_mem_size_) { - max_ctx_mem_size_ = mem_size; - } - trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); - - } else { - trt_context = std::unique_ptr(trt_engine->createExecutionContext()); - } + trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); if (!trt_context) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Nv EP could not build execution context for fused node: " + fused_node.Name()); } + bool is_dynamic_shape_context = false; // Create input/output to index maps for (int32_t i = 0; i < trt_engine->getNbIOTensors(); ++i) { auto const& name = trt_engine->getIOTensorName(i); auto const& mode = trt_engine->getTensorIOMode(name); if (mode == nvinfer1::TensorIOMode::kINPUT) { + is_dynamic_shape_context |= checkTrtDimIsDynamic(trt_engine->getTensorShape(name)); const auto& iter = input_map.find(name); if (iter != input_map.end()) { input_indexes[name] = iter->second; @@ -3027,9 +2824,8 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra &contexts_[context->node_name], input_info_[context->node_name], output_info_[context->node_name], - context_memory_sharing_enable_, - &max_ctx_mem_size_, - &tensorrt_mu_}; + &tensorrt_mu_, + is_dynamic_shape_context}; *state = p.release(); return 0; }; @@ -3056,15 +2852,14 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra auto& dds_output_allocator_map = this->dds_output_allocator_maps_[fused_node_name]; auto trt_engine = trt_state->engine->get(); auto trt_context = trt_state->context->get(); - auto max_context_mem_size_ptr = trt_state->max_context_mem_size_ptr; int num_outputs = static_cast(output_indexes.size()); std::unordered_map> shape_tensor_values; // This map holds "shape tensor -> shape values" for the shape tensor input across this inference run std::unordered_map> shape_tensor_values_int64; // same as above but for int64 shape tensor input - OrtDevice device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, - narrow(device_id_)); - OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, device); if (alloc_ == nullptr) { + OrtDevice device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, + narrow(device_id_)); + OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, device); Ort::ThrowOnError(api->KernelContext_GetAllocator(context, &mem_info, &alloc_)); } OrtAllocator* alloc = alloc_; @@ -3078,6 +2873,8 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "No engine is found."); } + bool require_io_binding = IsIOBindingRequired(trt_state, ctx); + // Get input and output binding names int total_bindings = trt_engine->getNbIOTensors(); std::vector input_binding_names, output_binding_names; @@ -3094,20 +2891,25 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra /* * Set input shapes and bind input buffers */ - std::vector> scratch_buffers; - for (size_t i = 0, end = input_binding_names.size(); i < end; ++i) { - char const* input_name = input_binding_names[i]; - - size_t input_index = 0; - const auto iter = input_indexes.find(input_name); - if (iter != input_indexes.end()) { - input_index = iter->second; - } + auto& scratch_buffers = trt_state->scratch_buffers; + if (require_io_binding) { + scratch_buffers.clear(); + bool skip_input_binding_allowed = true; + for (size_t i = 0, end = input_binding_names.size(); i < end; ++i) { + char const* input_name = input_binding_names[i]; + + size_t input_index = 0; + const auto iter = input_indexes.find(input_name); + if (iter != input_indexes.end()) { + input_index = iter->second; + } - Status status = BindContextInput(ctx, trt_engine, trt_context, input_name, input_index, shape_tensor_values, shape_tensor_values_int64, scratch_buffers, alloc, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + Status status = BindContextInput(ctx, trt_engine, trt_context, input_name, input_index, shape_tensor_values, shape_tensor_values_int64, scratch_buffers, alloc, stream, skip_input_binding_allowed); + if (status != Status::OK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + } } + trt_state->skip_io_binding_allowed = skip_input_binding_allowed; } /* @@ -3121,44 +2923,52 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra std::unordered_map output_dim_sizes; output_dim_sizes.reserve(num_outputs); - for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { - char const* output_name = output_binding_names[i]; + if (require_io_binding) { + bool skip_output_binding_allowed = true; + for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { + char const* output_name = output_binding_names[i]; - size_t output_index = 0; - const auto& index_iter = output_indexes.find(output_name); - if (index_iter != output_indexes.end()) { - output_index = index_iter->second; - } + size_t output_index = 0; + const auto& index_iter = output_indexes.find(output_name); + if (index_iter != output_indexes.end()) { + output_index = index_iter->second; + } - size_t output_type = 0; - const auto type_iter = output_types.find(output_name); - if (type_iter != output_types.end()) { - output_type = type_iter->second; - } + size_t output_type = 0; + const auto type_iter = output_types.find(output_name); + if (type_iter != output_types.end()) { + output_type = type_iter->second; + } + + nvinfer1::Dims dims; + void* data_ptr = nullptr; + + Status status = BindContextOutput(ctx, trt_context, output_name, output_index, output_type, i, output_tensors, output_dim_sizes, + dds_output_allocator_map, scratch_buffers, alloc, buffers, dims, data_ptr, skip_output_binding_allowed); + if (status != Status::OK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + } - Status status = BindContextOutput(ctx, trt_context, output_name, output_index, output_type, i, output_tensors, output_dim_sizes, - dds_output_allocator_map, scratch_buffers, alloc, buffers); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); + trt_state->output_tensors[output_index] = TensorParams{data_ptr, dims}; } + + trt_state->skip_io_binding_allowed = trt_state->skip_io_binding_allowed | skip_output_binding_allowed; } // Set execution context memory - if (trt_state->context_memory_sharing_enable) { -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) -#endif + if (require_io_binding) { size_t mem_size = trt_engine->getDeviceMemorySizeV2(); -#if defined(_MSC_VER) -#pragma warning(pop) -#endif - if (mem_size > *max_context_mem_size_ptr) { - *max_context_mem_size_ptr = mem_size; + if (trt_state->is_dynamic_shape) { + mem_size = trt_context->updateDeviceMemorySizeForShapes(); + } + if (trt_state->context_memory_size != mem_size) { + LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] A new context memory was allocated with size " << mem_size; + trt_state->context_memory = IAllocator::MakeUniquePtrFromOrtAllocator(alloc, mem_size, false /*use_reserve*/); + // trt_state->context_memory = IAllocator::MakeUniquePtr(alloc, mem_size, false /*use_reserve*/, stream); + trt_state->context_memory_size = mem_size; + trt_context->setDeviceMemoryV2(trt_state->context_memory.get(), mem_size); } - trt_context->setDeviceMemory(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, *max_context_mem_size_ptr).get()); } - // Start CUDA graph capture. // Note: The reason we don't put graph capture in OnRunStart() like CUDA EP does is because // current ORT TRT doesn't get cuda stream until compute time and graph capture requires cuda stream. @@ -3209,7 +3019,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra if (index_iter != output_indexes.end()) { output_index = index_iter->second; } - auto status = BindKernelOutput(ctx, &mem_info, dds_output_allocator_map, output_name, output_index, output_type, stream); + auto status = BindKernelOutput(ctx, dds_output_allocator_map, output_name, output_index, output_type, stream); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, status.ErrorMessage()); } diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 7a0c47d28c81d..83b89a2e9d1fb 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -78,6 +78,9 @@ using unique_pointer = std::unique_ptr; // class OutputAllocator : public nvinfer1::IOutputAllocator { public: + OutputAllocator() = delete; + OutputAllocator(OrtAllocator* allocator) : alloc_(allocator) {}; + void* reallocateOutputAsync(char const* tensorName, void* currentMemory, uint64_t size, uint64_t alignment, cudaStream_t stream) noexcept override; void notifyShape(char const* tensorName, nvinfer1::Dims const& dims) noexcept override; @@ -95,10 +98,11 @@ class OutputAllocator : public nvinfer1::IOutputAllocator { } ~OutputAllocator() override { - cudaFree(outputPtr); + alloc_->Free(alloc_, outputPtr); } private: + OrtAllocator* alloc_; void* outputPtr{nullptr}; uint64_t allocated_size = 0; std::vector output_shapes; @@ -110,6 +114,45 @@ class OutputAllocator : public nvinfer1::IOutputAllocator { */ using ShapeRangesMap = std::unordered_map>>>; +/** + * @brief Container for tensor data and their shape. + * + */ +struct TensorParams { + const void* data{nullptr}; + nvinfer1::Dims dims; + + TensorParams() = default; + + TensorParams(const void* data_ptr, const std::vector& shape) { + // Initialize data and dims from the Ort::ConstValue + data = data_ptr; + + dims.nbDims = static_cast(shape.size()); + for (int i = 0; i < dims.nbDims; ++i) { + dims.d[i] = static_cast(shape[i]); + } + } + + TensorParams(const void* data_ptr, nvinfer1::Dims& shape) { + // Initialize data and dims from the Ort::ConstValue + data = data_ptr; + + dims = shape; + } + + bool operator!=(const TensorParams& other) const { + if (data != other.data || dims.nbDims != other.dims.nbDims) + return true; + + for (int i = 0; i < dims.nbDims; ++i) { + if (dims.d[i] != other.dims.d[i]) + return true; + } + return false; + } +}; + // Information to construct kernel function state. struct TensorrtFuncState { AllocateFunc test_allocate_func = nullptr; @@ -130,8 +173,6 @@ struct TensorrtFuncState { std::string engine_cache_path; nvinfer1::IRuntime* runtime = nullptr; std::vector profiles; - bool context_memory_sharing_enable = false; - size_t* max_context_mem_size_ptr = nullptr; bool engine_decryption_enable = false; int (*engine_decryption)(const char*, char*, size_t*) = nullptr; int (*engine_encryption)(const char*, char*, size_t) = nullptr; @@ -139,8 +180,16 @@ struct TensorrtFuncState { bool sparsity_enable = false; int auxiliary_streams = -1; bool cuda_graph_enable = 0; + bool is_dynamic_shape = false; std::string cache_prefix; std::string cache_suffix; + std::vector> scratch_buffers; + std::vector input_tensors; + std::vector output_tensors; + bool is_first_run = true; // Indicates if this is the first run of the engine + bool skip_io_binding_allowed = false; // Indicates if input/output binding can be skipped + IAllocatorUniquePtr context_memory = nullptr; + size_t context_memory_size = 0; }; // Minimum information to construct kernel function state for direct engine load code path @@ -153,9 +202,15 @@ struct TensorrtShortFuncState { std::unique_ptr* context = nullptr; std::vector> input_info; std::vector> output_info; - bool context_memory_sharing_enable = false; - size_t* max_context_mem_size_ptr = nullptr; std::mutex* tensorrt_mu_ptr = nullptr; + bool is_dynamic_shape = false; + std::vector> scratch_buffers; + std::vector input_tensors; + std::vector output_tensors; + bool is_first_run = true; // Indicates if this is the first run of the engine + bool skip_io_binding_allowed = false; // Indicates if input/output binding can be skipped + IAllocatorUniquePtr context_memory = nullptr; + size_t context_memory_size = 0; }; // Holds important information for building valid ORT graph. @@ -251,9 +306,7 @@ class NvExecutionProvider : public IExecutionProvider { std::mutex tensorrt_mu_; int device_id_; std::string compute_capability_; - bool context_memory_sharing_enable_ = false; size_t max_ctx_mem_size_ = 0; - IAllocatorUniquePtr context_memory_ = nullptr; mutable char model_path_[4096] = {}; // Reserved for max path length bool engine_decryption_enable_ = false; int (*engine_decryption_)(const char*, char*, size_t*) = nullptr; @@ -341,8 +394,6 @@ class NvExecutionProvider : public IExecutionProvider { nvinfer1::IExecutionContext& GetTensorRTContext(std::string fused_node); bool UpdateTensorRTContext(std::string fused_node, std::unique_ptr context); void ResetTensorRTContext(std::string fused_node); - bool CompareProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges); - void UpdateProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges); void InitCUDAGraph(); void SetGraphStream(cudaStream_t stream); diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h index 2a67f3c3bec4d..4d6c6fe116076 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h @@ -34,7 +34,6 @@ struct NvExecutionProviderInfo { bool engine_decryption_enable{false}; std::string engine_decryption_lib_path{""}; bool force_sequential_engine_build{false}; - bool context_memory_sharing_enable{false}; std::string timing_cache_path{""}; bool detailed_build_log{false}; bool sparsity_enable{false}; diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h index 22e5eea6924de..ea586ba445ba2 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h @@ -683,4 +683,29 @@ std::string GetCacheSuffix(const std::string& fused_node_name, const std::string } return ""; } + +/* + * Checks if there is a an element with value `-1` in nvinfer1::Dims + */ +static bool checkTrtDimIsDynamic(nvinfer1::Dims dims) { + for (int j = 0, end = dims.nbDims; j < end; ++j) { + if (dims.d[j] == -1) { + return true; + } + } + return false; +} + +/* + * Checks if an nvinfer1::ITensor signales a dynamic shape, + * either due to dynamic shapes or due to it being a shape tensor + */ +static bool checkTrtTensorIsDynamic(nvinfer1::ITensor* tensor) { + if (tensor->isShapeTensor()) { + return true; + } else { + // Execution tensor + return checkTrtDimIsDynamic(tensor->getDimensions()); + } +} } // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc index e236cccaaaa77..d23d50549b2c5 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc @@ -557,6 +557,67 @@ struct NvTensorRtRtxEpFactory : OrtEpFactory { return ORT_VERSION; } + /** + * @brief Checks if a given OrtHardwareDevice is a supported NVIDIA GPU. + * + * This function verifies if the provided hardware device corresponds to a physical + * NVIDIA GPU that meets the minimum compute capability requirements for this execution provider. + * + * The check is performed by: + * 1. Extracting the LUID (Locally Unique Identifier) from the device's metadata. + * 2. Converting the string LUID to a 64-bit integer. + * 3. Iterating through all available CUDA devices on the system. + * 4. For each CUDA device, constructing its 64-bit LUID from its properties. + * 5. Comparing the LUIDs. If a match is found, it checks if the device's + * compute capability is at least 8.0 (Ampere) or newer. + * + * @param device The OrtHardwareDevice to check. + * @return True if the device is a supported NVIDIA GPU, false otherwise. + */ + bool IsOrtHardwareDeviceSupported(const OrtHardwareDevice& device) { + const auto& metadata_entries = device.metadata.Entries(); + const auto it = metadata_entries.find("LUID"); + if (it == metadata_entries.end()) { + return false; + } + + uint64_t target_luid; + try { + target_luid = std::stoull(it->second); + } catch (const std::exception&) { + return false; + } + + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess) { + return false; + } + + for (int i = 0; i < device_count; ++i) { + cudaDeviceProp prop; + if (cudaGetDeviceProperties(&prop, i) != cudaSuccess) { + continue; + } + + // The LUID is an 8-byte value, valid on Windows when luidDeviceNodeMask is non-zero. + // We reconstruct the 64-bit integer representation from the raw bytes. + if (prop.luidDeviceNodeMask == 0) { + continue; + } + + // Ensure the LUID is 8 bytes and reinterpret it directly as a uint64_t for comparison. + static_assert(sizeof(prop.luid) == sizeof(uint64_t), "cudaDeviceProp::luid should be 8 bytes"); + uint64_t current_luid = *reinterpret_cast(prop.luid); + + if (current_luid == target_luid) { + // Ampere architecture or newer is required. + return prop.major >= 8; + } + } + + return false; + } + // Creates and returns OrtEpDevice instances for all OrtHardwareDevices that this factory supports. // An EP created with this factory is expected to be able to execute a model with *all* supported // hardware devices at once. A single instance of NvTensorRtRtx EP is not currently setup to partition a model among @@ -579,11 +640,12 @@ struct NvTensorRtRtxEpFactory : OrtEpFactory { int16_t device_id = 0; for (size_t i = 0; i < num_devices && num_ep_devices < max_ep_devices; ++i) { const OrtHardwareDevice& device = *devices[i]; + if (factory->ort_api.HardwareDevice_Type(&device) == OrtHardwareDeviceType::OrtHardwareDeviceType_GPU && - factory->ort_api.HardwareDevice_VendorId(&device) == factory->vendor_id) { + factory->ort_api.HardwareDevice_VendorId(&device) == factory->vendor_id && + factory->IsOrtHardwareDeviceSupported(device)) { OrtKeyValuePairs* ep_options = nullptr; OrtKeyValuePairs* ep_metadata = nullptr; - factory->ort_api.CreateKeyValuePairs(&ep_options); factory->ort_api.CreateKeyValuePairs(&ep_metadata); factory->ort_api.AddKeyValuePair(ep_options, "device_id", std::to_string(device_id).c_str()); diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index a22d21d8d798b..bdeea726a2cf5 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -491,16 +491,29 @@ Status ShaderHelper::GenerateSourceCode(std::string& code, std::vector& sha ss << ","; } - auto alignment = (data_type == ProgramUniformVariableDataType::Float16 && length > 4) ? "@align(16) " : ""; - ss << "\n " << alignment << name << ": "; + // The actual variable type for the uniform variable depends on the data type (T) and length (N). + // + // For T in [i32, u32, f32]: + // - If N == 1, the type is simply i32, u32, or f32. + // - If 2 < N <= 4, the type is vecN, vecN, or vecN where N is the length. + // - If N > 4, the type is array, ceil(N / 4)>. + // + // For T is f16: + // - If N == 1 or N == 2, the type is u32. + // - If 2 < N <= 8, the type is vecX where X is ceil(N / 2). + // - If N > 8, the type is array, X> where X is ceil(N / 8). + // + // Note: Using f16 type in uniforms is not generally supported on all devices. We use a u32 variable to represent + // 2 f16 values. + + if (data_type == ProgramUniformVariableDataType::Float16) { + data_type = ProgramUniformVariableDataType::Uint32; // f16 is represented as u32 + length = (length + 1) / 2; // each u32 can hold 2 f16 values + } + ss << "\n " << name << ": "; if (length > 4) { - if (data_type == ProgramUniformVariableDataType::Float16) { - size_t array_size = (length + 7) / 8; - ss << "array, " << array_size << ">"; - } else { - size_t array_size = (length + 3) / 4; - ss << "array, " << array_size << ">"; - } + size_t array_size = (length + 3) / 4; + ss << "array, " << array_size << ">"; } else if (length > 1) { ss << "vec" << length << "<" << data_type << ">"; } else { diff --git a/onnxruntime/core/providers/webgpu/shader_variable.h b/onnxruntime/core/providers/webgpu/shader_variable.h index 2aba2a59d157f..78c98ab26f5b8 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.h +++ b/onnxruntime/core/providers/webgpu/shader_variable.h @@ -17,18 +17,34 @@ template || std::is_same_v>> std::string GetElementAt(std::string_view var, const TIdx& idx, TRank rank, bool is_f16 = false) { - // "std::string::rfind(str, 0) == 0" is equivalent to "std::string::starts_with(str)" before C++20. - if (var.rfind("uniforms.", 0) == 0) { - if (rank > 4) { - if constexpr (std::is_integral_v) { - if (is_f16) { - return MakeStringWithClassicLocale(var, "[", idx / 8, "][", (idx % 8) / 4, "][", (idx % 8) % 4, "]"); + if (var.starts_with("uniforms.")) { + if (is_f16) { + if (rank > 8) { + // array, N> + if constexpr (std::is_integral_v) { + return MakeStringWithClassicLocale("bitcast>(", var, "[", idx / 8, "][", (idx % 8) / 2, "])[", (idx % 8) % 2, "]"); } else { - return MakeStringWithClassicLocale(var, "[", idx / 4, "][", idx % 4, "]"); + return MakeStringWithClassicLocale("bitcast>(", var, "[(", idx, ") / 8][((", idx, ") % 8) / 2])[((", idx, ") % 8) % 2]"); + } + } else if (rank > 2) { + // vecN + if constexpr (std::is_integral_v) { + return MakeStringWithClassicLocale("bitcast>(", var, "[", idx / 2, "])[", idx % 2, "]"); + } else { + return MakeStringWithClassicLocale("bitcast>(", var, "[(", idx, ") / 2])[(", idx, ") % 2]"); } } else { - if (is_f16) { - return MakeStringWithClassicLocale(var, "[(", idx, ") / 8][(", idx, ") % 8 / 4][(", idx, ") % 8 % 4]"); + // u32 + if constexpr (std::is_integral_v) { + return MakeStringWithClassicLocale("bitcast>(", var, ")[", idx % 2, "]"); + } else { + return MakeStringWithClassicLocale("bitcast>(", var, ")[(", idx, ") % 2]"); + } + } + } else { + if (rank > 4) { + if constexpr (std::is_integral_v) { + return MakeStringWithClassicLocale(var, "[", idx / 4, "][", idx % 4, "]"); } else { return MakeStringWithClassicLocale(var, "[(", idx, ") / 4][(", idx, ") % 4]"); } diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 4bd79a627df22..a9557f7b9aa87 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -373,26 +373,57 @@ Status WebGpuContext::Run(ComputeContext& context, const ProgramBase& program) { continue; } - bool is_f16 = uniform.data_type == ProgramUniformVariableDataType::Float16; - - size_t element_size = ProgramUniformVariableDataTypeSize[static_cast(uniform.data_type)]; + // Calculate the size and alignment of the uniform variable. + // // https://www.w3.org/TR/WGSL/#alignof - size_t base_alignment = is_f16 - ? (length > 4 ? 16 : length > 2 ? 8 - : length * element_size) - : (length > 2 ? 16 : length * element_size); - size_t struct_size = is_f16 && length <= 4 ? length * element_size : 16; - - current_offset = (current_offset + base_alignment - 1) / base_alignment * base_alignment; + // + // For f16: + // - length > 8 : array, N> (align 16) (size 16 * N, N = ceil(length / 8)) + // - length == 7 or 8: vec4 (align 16) (size 16) + // - length == 5 or 6: vec3 (align 16) (size 12) + // - length == 3 or 4: vec2 (align 8) (size 8) + // - length == 1 or 2: u32 (align 4) (size 4) + // + // For other types (i32, u32, f32): + // - length > 4 : array, N> (align 16) (size 16 * N, N = ceil(length / 4)) + // - length == 4 : vec4 (align 16) (size 16) + // - length == 3 : vec3 (align 16) (size 12) + // - length == 2 : vec2 (align 8) (size 8) + // - length == 1 : T (align 4) (size 4) + // + + const bool is_f16 = uniform.data_type == ProgramUniformVariableDataType::Float16; + + size_t variable_alignment = 4; // default alignment for scalar types + size_t variable_size = 4; // default size for scalar types + + if (is_f16) { + if (length > 6) { + variable_alignment = 16; + variable_size = 16 * ((length + 7) / 8); + } else if (length > 4) { + variable_alignment = 16; + variable_size = 12; + } else if (length > 2) { + variable_alignment = 8; + variable_size = 8; + } + } else { + if (length > 3) { + variable_alignment = 16; + variable_size = 16 * ((length + 3) / 4); + } else if (length > 2) { + variable_alignment = 16; + variable_size = 12; + } else if (length > 1) { + variable_alignment = 8; + variable_size = 8; + } + } + current_offset = (current_offset + variable_alignment - 1) / variable_alignment * variable_alignment; uniform_and_offsets.emplace_back(uniform, current_offset); - // For non-float16 type, when length > 4, the uniform variable is of type array,N>, where - // N = ceil(data.length / 4) and SizeOf(vec4) = 16. The total byte length is N * SizeOf(vec4). - // For float16 type, when length > 4, the uniform variable is of type array,N>, where - // N = ceil(data.length / 8) and SizeOf(mat2x4) = 16. The total byte length is N * SizeOf(mat2x4). - size_t element_per_struct = is_f16 ? 8 : 4; - current_offset += - length > 4 ? (length + element_per_struct - 1) / element_per_struct * struct_size : length * element_size; + current_offset += variable_size; } // Meet alignment of struct here: https://www.w3.org/TR/WGSL/#alignment-and-size. For simplicity, set diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index 2b553aecbca6c..dfb2e33f8cb32 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -72,21 +72,23 @@ ProviderInfo_CUDA& GetProviderInfo_CUDA(); #endif // defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) namespace { -// Ignore whether there is an arena wrapping the allocator by excluding OrtMemoryInfo.alloc_type from the comparison +// Ignore whether there is an arena wrapping the allocator by excluding OrtMemoryInfo.alloc_type from the comparison. static bool AreOrtMemoryInfosEquivalent( const OrtMemoryInfo& left, const OrtMemoryInfo& right, - bool match_name = true) { + bool match_name = true, + bool ignore_alignment = false) { return left.mem_type == right.mem_type && - left.device == right.device && + (ignore_alignment ? left.device.EqualIgnoringAlignment(right.device) : left.device == right.device) && (!match_name || strcmp(left.name, right.name) == 0); } std::vector::const_iterator FindExistingAllocator(const std::vector& allocators, const OrtMemoryInfo& mem_info, - bool match_name = true) { + bool match_name = true, + bool ignore_alignment = false) { auto ite = std::find_if(std::begin(allocators), std::end(allocators), - [&mem_info, match_name](const AllocatorPtr& alloc_ptr) { + [&mem_info, match_name, ignore_alignment](const AllocatorPtr& alloc_ptr) { // We want to do the equality checking of 2 OrtMemoryInfos sans the OrtAllocatorType field. // This is because we want to avoid registering two allocators for the same device that just // differ on OrtAllocatorType. @@ -96,7 +98,8 @@ std::vector::const_iterator FindExistingAllocator(const std::vecto // OrtDeviceAllocator (which is the only accepted value while registering a custom allocator). // If we allowed this, it could potentially cause a lot of confusion as to which shared allocator // to use for that device and we want to avoid having any ugly logic around this. - return AreOrtMemoryInfosEquivalent(alloc_ptr->Info(), mem_info, match_name); + return AreOrtMemoryInfosEquivalent(alloc_ptr->Info(), mem_info, + match_name, ignore_alignment); }); return ite; @@ -428,8 +431,25 @@ Status Environment::CreateAndRegisterAllocatorV2(const std::string& provider_typ } Environment::~Environment() { - // need to make sure all the OrtAllocator instances are released prior to any plugin EPs being freed + // need to make sure all the OrtAllocator instances are released prior to any plugin EPs being freed. + // this is because any entry in shared_allocators_ wrapping an OrtAllocator from a plugin EP owns the OrtAllocator + // instance and will call Release on it. If the plugin EP has been freed the Release will fail. shared_allocators_.clear(); + +#if !defined(ORT_MINIMAL_BUILD) + // unregister any remaining EP libraries so they're cleaned up in a determistic way. + while (!ep_libraries_.empty()) { + auto it = ep_libraries_.begin(); + ORT_IGNORE_RETURN_VALUE(UnregisterExecutionProviderLibrary(it->first)); + } +#endif +} + +AllocatorPtr Environment::GetRegisteredSharedAllocator(const OrtMemoryInfo& mem_info) const { + std::lock_guard lock{mutex_}; + + auto it = FindExistingAllocator(shared_allocators_, mem_info, /*match_name*/ false, /*ignore_alignment*/ true); + return it != shared_allocators_.end() ? *it : nullptr; } Status Environment::GetSharedAllocator(const OrtMemoryInfo& mem_info, OrtAllocator*& allocator) { diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index f4f76a389030e..c0900c5ad28a0 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1421,6 +1421,29 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool } } + // We choose to convert initializers into OrtValues before partitioning here so plug-in EPs could + // take advantage of the initializers being in OrtValue format and not to deal with protobuf. + // + // The initializers data is transferred to an OrtValue. The original TensorProto is replaced + // with a TensorProto that has the same data type, shape and name. However, its external data + // is used in a non-standard way. The location is set to a string constant utils::kTensorProtoMemoryAddressTag, + // The file offset is set to the address of the OrtValue's data buffer, and the length is set to the size of the + // OrtValue's data buffer. Because this external location is non-standard, onnx code can not handle it, so we choose + // to do it as late as possible but before the partitioning so type and shape inference accesses the initializers + // before they are converted to OrtValues. + // + // If any transformations are applied later, they would not introduce any in-memory initializers, + // type and shape inference would run only on any newly added nodes and any new initializers + // will be converted at session finalization time. + // + // The conversion is performed using the following steps (within ConvertInitializersIntoOrtValues()) + // constexpr const bool use_tensor_buffer_true = true; + // auto tensor_proto_to_add = utils::TensorToTensorProto(ort_value.Get(), tensor_proto.name(), + // use_tensor_buffer_true); + // ORT_RETURN_IF_ERROR(graph.ReplaceInitializedTensor(tensor_proto_to_add, ort_value)); + + ORT_RETURN_IF_ERROR_SESSIONID_(graph.ConvertInitializersIntoOrtValues()); + // Do partitioning based on execution providers' capabilities. ORT_RETURN_IF_ERROR_SESSIONID_(partitioner.Partition(graph, session_state_->GetMutableFuncMgr(), transform_layout_fn, session_options_.config_options, *session_logger_, @@ -1984,13 +2007,15 @@ static void ResolveMemoryPatternFlags(SessionState& session_state) { // For now, this function only checks for invalid combination of DML EP with other EPs. // TODO: extend this function to check for other invalid combinations of EPs. common::Status InferenceSession::HasInvalidCombinationOfExecutionProviders() const { - // DML EP is only allowed with CPU EP + // DML EP is not allowed with other GPU or NPU EPs. + // historical reason for this is unknown. relaxing the limit that it must only be used with the CPU EP to support + // scenarios where alternative EPs are CPU based (e.g. openvino). bool has_dml_ep = execution_providers_.Get(kDmlExecutionProvider) != nullptr; if (has_dml_ep) { - const auto& ep_list = execution_providers_.GetIds(); - for (const auto& ep : ep_list) { - if (ep == kDmlExecutionProvider || ep == kCpuExecutionProvider) continue; - return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "DML EP can be used with only CPU EP."); + for (const auto& ep : execution_providers_) { + if (ep->Type() != kDmlExecutionProvider && ep->GetDevice().Type() != OrtDevice::CPU) { + return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "DML EP can only be used with CPU EPs."); + } } } return Status::OK(); diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index e8e51db13bcd3..64c4ada07f28f 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -21,7 +21,7 @@ import onnxruntime -def get_ort_device_type(device_type: str, device_index) -> C.OrtDevice: +def get_ort_device_type(device_type: str) -> int: if device_type == "cuda": return C.OrtDevice.cuda() elif device_type == "cann": @@ -32,8 +32,10 @@ def get_ort_device_type(device_type: str, device_index) -> C.OrtDevice: return C.OrtDevice.dml() elif device_type == "webgpu": return C.OrtDevice.webgpu() - elif device_type == "ort": - return C.get_ort_device(device_index).device_type() + elif device_type == "gpu": + return C.OrtDevice.gpu() + elif device_type == "npu": + return C.OrtDevice.npu() else: raise Exception("Unsupported device type: " + device_type) @@ -765,7 +767,7 @@ def bind_input(self, name, device_type, device_id, element_type, shape, buffer_p self._iobinding.bind_input( name, C.OrtDevice( - get_ort_device_type(device_type, device_id), + get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id, ), @@ -812,7 +814,7 @@ def bind_output( self._iobinding.bind_output( name, C.OrtDevice( - get_ort_device_type(device_type, device_id), + get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id, ), @@ -823,7 +825,7 @@ def bind_output( self._iobinding.bind_output( name, C.OrtDevice( - get_ort_device_type(device_type, device_id), + get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id, ), @@ -889,7 +891,7 @@ def _get_c_value(self) -> C.OrtValue: return self._ortvalue @classmethod - def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0) -> OrtValue: + def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id=-1) -> OrtValue: """ Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu @@ -897,6 +899,7 @@ def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device :param numpy_obj: The Numpy object to construct the OrtValue from :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 + :param vendor_id: The device's PCI vendor id. If provided, the device_type should be "gpu" or "npu". """ # Hold a reference to the numpy object (if device_type is 'cpu') as the OrtValue # is backed directly by the data buffer of the numpy object and so the numpy object @@ -904,11 +907,7 @@ def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device return cls( C.OrtValue.ortvalue_from_numpy( numpy_obj, - C.OrtDevice( - get_ort_device_type(device_type, device_id), - C.OrtDevice.default_memory(), - device_id, - ), + OrtDevice.make(device_type, device_id, vendor_id)._get_c_device(), ), numpy_obj if device_type.lower() == "cpu" else None, ) @@ -929,7 +928,7 @@ def ortvalue_from_numpy_with_onnx_type(cls, data: np.ndarray, /, onnx_element_ty @classmethod def ortvalue_from_shape_and_type( - cls, shape: Sequence[int], element_type, device_type: str = "cpu", device_id: int = 0 + cls, shape: Sequence[int], element_type, device_type: str = "cpu", device_id: int = 0, vendor_id: int = -1 ) -> OrtValue: """ Factory method to construct an OrtValue (which holds a Tensor) from given shape and element_type @@ -938,7 +937,11 @@ def ortvalue_from_shape_and_type( :param element_type: The data type of the elements. It can be either numpy type (like numpy.float32) or an integer for onnx type (like onnx.TensorProto.BFLOAT16). :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 + :param vendor_id: If provided the device type should be "gpu" or "npu". """ + + device = OrtDevice.make(device_type, device_id, vendor_id)._get_c_device() + # Integer for onnx element type (see https://onnx.ai/onnx/api/mapping.html). # This is helpful for some data type (like TensorProto.BFLOAT16) that is not available in numpy. if isinstance(element_type, int): @@ -946,11 +949,7 @@ def ortvalue_from_shape_and_type( C.OrtValue.ortvalue_from_shape_and_onnx_type( shape, element_type, - C.OrtDevice( - get_ort_device_type(device_type, device_id), - C.OrtDevice.default_memory(), - device_id, - ), + device, ) ) @@ -958,11 +957,7 @@ def ortvalue_from_shape_and_type( C.OrtValue.ortvalue_from_shape_and_type( shape, element_type, - C.OrtDevice( - get_ort_device_type(device_type, device_id), - C.OrtDevice.default_memory(), - device_id, - ), + device, ) ) @@ -1085,14 +1080,27 @@ def _get_c_device(self): return self._ort_device @staticmethod - def make(ort_device_name, device_id): - return OrtDevice( - C.OrtDevice( - get_ort_device_type(ort_device_name, device_id), - C.OrtDevice.default_memory(), - device_id, + def make(ort_device_name, device_id, vendor_id=-1): + if vendor_id < 0: + # backwards compatibility with predefined OrtDevice names + return OrtDevice( + C.OrtDevice( + get_ort_device_type(ort_device_name), + C.OrtDevice.default_memory(), + device_id, + ) + ) + else: + # generic. use GPU or NPU for ort_device_name and provide a vendor id. + # vendor id of 0 is valid in some cases (e.g. webgpu is generic and does not have a vendor id) + return OrtDevice( + C.OrtDevice( + get_ort_device_type(ort_device_name), + C.OrtDevice.default_memory(), + vendor_id, + device_id, + ) ) - ) def device_id(self): return self._ort_device.device_id() @@ -1100,6 +1108,9 @@ def device_id(self): def device_type(self): return self._ort_device.device_type() + def device_vendor_id(self): + return self._ort_device.vendor_id() + class SparseTensor: """ diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index 958c9fc46bcd8..590e1ef3cdbdb 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -99,6 +99,44 @@ TensorShape GetShape(const py::array& arr) { return shape; } +AllocatorPtr GetSharedAllocator(const OrtDevice& device) { + auto& env = GetOrtEnv()->GetEnvironment(); + + OrtMemoryInfo mem_info("ignored", OrtDeviceAllocator, device); + return env.GetRegisteredSharedAllocator(mem_info); +} + +MemCpyFunc CreateDataTransferMemCpy([[maybe_unused]] const OrtDevice& src_device, + [[maybe_unused]] const OrtDevice& dst_device) { +#if defined(ORT_MINIMAL_BUILD) + // plugin EPs are not supported in a minimal build so there won't be any data transfers registered + return nullptr; +#else + + auto& env = GetOrtEnv()->GetEnvironment(); + const DataTransferManager& data_transfer_manager = env.GetDataTransferManager(); + const IDataTransfer* data_transfer = data_transfer_manager.GetDataTransfer(src_device, dst_device); + if (!data_transfer) { + return nullptr; + } + + const auto copy_func = [src_device, dst_device, data_transfer](void* dst, const void* src, size_t bytes) { + OrtMemoryInfo src_memory_info("ignored", OrtDeviceAllocator, src_device); + OrtMemoryInfo dst_memory_info("ignored", OrtDeviceAllocator, dst_device); + + // real shape doesn't matter as the Tensor instances here are temporary in order to be able to call CopyTensor. + // we set the shape to `bytes` and the data type to uint8_t to copy the correct number of bytes. + TensorShape shape = {narrow(bytes)}; + Tensor src_tensor{DataTypeImpl::GetType(), shape, const_cast(src), src_memory_info}; + Tensor dst_tensor{DataTypeImpl::GetType(), shape, dst, dst_memory_info}; + + ORT_THROW_IF_ERROR(data_transfer->CopyTensor(src_tensor, dst_tensor)); + }; + + return copy_func; +#endif +} + void CpuToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { memcpy(dst, src, num_bytes); } @@ -158,9 +196,10 @@ void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_CUDA().cudaMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetCudaToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, CudaToCpuMemCpy}}; +const std::unordered_map* GetCudaToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, 0}, CudaToCpuMemCpy}, + }; return ↦ } @@ -215,9 +254,10 @@ void MIGraphXToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_MIGraphX().MIGraphXMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetMIGraphXToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, MIGraphXToCpuMemCpy}}; +const std::unordered_map* GetMIGraphXToHostMemCpyFunction(const OrtDevice& device) { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::AMD, 0}, MIGraphXToCpuMemCpy}, + }; return ↦ } @@ -334,9 +374,10 @@ void DmlToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { D3D12_RESOURCE_STATE_UNORDERED_ACCESS); } -const std::unordered_map* GetDmlToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, DmlToCpuMemCpy}}; +const std::unordered_map* GetDmlToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::MICROSOFT, 0}, DmlToCpuMemCpy}, + }; return ↦ } @@ -352,9 +393,10 @@ void CannToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_CANN().cannMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetCannToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::NPU, CannToCpuMemCpy}}; +const std::unordered_map* GetCannToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::NPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::HUAWEI, 0}, CannToCpuMemCpy}, + }; return ↦ } @@ -402,9 +444,10 @@ void RocmToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_ROCM().rocmMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetRocmToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, RocmToCpuMemCpy}}; +const std::unordered_map* GetRocmToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::AMD, 0}, RocmToCpuMemCpy}, + }; return ↦ } @@ -581,7 +624,7 @@ using OrtPybindSingleUseAllocatorPtr = std::shared_ptr& p_tensor, - MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) { + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy) { CopyDataToTensor(darray, npy_type, *p_tensor, mem_cpy_to_device); } -void CopyDataToTensor(const py::array& py_array, int npy_type, Tensor& tensor, MemCpyFunc mem_cpy_to_device) { +void CopyDataToTensor(const py::array& py_array, int npy_type, Tensor& tensor, const MemCpyFunc& mem_cpy_to_device) { CopyDataToTensor(reinterpret_cast(py_array.ptr()), npy_type, tensor, mem_cpy_to_device); } @@ -656,7 +699,7 @@ void CopyDataToTensor(const py::array& py_array, int npy_type, Tensor& tensor, M // The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope static std::unique_ptr CreateTensor(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject, bool use_numpy_data_memory = true, - MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) { + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy) { PyArrayObject* darray = PyArray_GETCONTIGUOUS(pyObject); ORT_ENFORCE(darray != nullptr, "The object must be a contiguous array for input '", name_input, "'."); @@ -746,7 +789,8 @@ static void CreateSequenceOfTensors(AllocatorPtr alloc, const std::string& name_ // as the backing data buffer for the ORT Tensor where applicable (for numeric tensors) // The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope static void CreateTensorMLValue(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject, - OrtValue* p_mlvalue, bool use_numpy_data_memory = true, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) { + OrtValue* p_mlvalue, bool use_numpy_data_memory = true, + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy) { auto p_tensor = CreateTensor(alloc, name_input, pyObject, use_numpy_data_memory, mem_cpy_to_device); auto ml_tensor = DataTypeImpl::GetType(); @@ -994,9 +1038,10 @@ static void CreateGenericIterableMLValue(PyObject* iterator, AllocatorPtr alloc, // Setting `use_numpy_data_memory` to `true` will ensure that the underlying numpy array buffer is directly used // as the backing data buffer for the ORT Tensor where applicable (for numeric tensors) // The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope -void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, const std::string& name_input, - const py::object& value, OrtValue* p_mlvalue, bool accept_only_numpy_array, - bool use_numpy_data_memory, MemCpyFunc mem_cpy_to_device) { +void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, + const std::string& name_input, const py::object& value, OrtValue* p_mlvalue, + bool accept_only_numpy_array, bool use_numpy_data_memory, + const MemCpyFunc& mem_cpy_to_device) { onnx::TypeProto type_proto; if (PyObjectCheck_NumpyArray(value.ptr())) { // The most frequent case: input comes as an array. diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.h b/onnxruntime/python/onnxruntime_pybind_mlvalue.h index e9bafea2ed1b5..7b65c0aae45c1 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.h +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.h @@ -42,22 +42,27 @@ MLDataType NumpyTypeToOnnxRuntimeTensorType(int numpy_type); MLDataType OnnxTypeToOnnxRuntimeTensorType(int onnx_element_type); -using MemCpyFunc = void (*)(void*, const void*, size_t); - +using MemCpyFunc = std::function; using DataTransferAlternative = std::variant; +// helpers to get allocator and IDataTransfer from Environment for plugin EP +AllocatorPtr GetSharedAllocator(const OrtDevice& device); +MemCpyFunc CreateDataTransferMemCpy(const OrtDevice& src_device, const OrtDevice& dst_device); + void CpuToCpuMemCpy(void*, const void*, size_t); -void CopyDataToTensor(const pybind11::array& py_array, int npy_type, Tensor& tensor, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy); +void CopyDataToTensor(const pybind11::array& py_array, int npy_type, Tensor& tensor, + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy); pybind11::object AddTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions); + const std::unordered_map* mem_cpy_to_host_functions); -pybind11::object GetPyObjectFromSparseTensor(size_t pos, const OrtValue& ort_value, const DataTransferManager* data_transfer_manager); +pybind11::object GetPyObjectFromSparseTensor(size_t pos, const OrtValue& ort_value, + const DataTransferManager* data_transfer_manager); pybind11::object AddNonTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions); + const std::unordered_map* mem_cpy_to_host_functions); OrtMemoryInfo GetMemoryInfoPerDeviceType(const OrtDevice& ort_device); @@ -69,7 +74,7 @@ void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes); void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetCudaToHostMemCpyFunction(); +const std::unordered_map* GetCudaToHostMemCpyFunction(); bool IsCudaDeviceIdValid(const onnxruntime::logging::Logger& logger, int id); @@ -87,7 +92,7 @@ void CpuToDmlMemCpy(void* dst, const void* src, size_t num_bytes); void DmlToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetDmlToHostMemCpyFunction(); +const std::unordered_map* GetDmlToHostMemCpyFunction(); #endif @@ -97,7 +102,7 @@ void CpuToMIGraphXMemCpy(void* dst, const void* src, size_t num_bytes); void MIGraphXToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetMIGraphXToHostMemCpyFunction(); +const std::unordered_map* GetMIGraphXToHostMemCpyFunction(); AllocatorPtr GetMIGraphXAllocator(OrtDevice::DeviceId id); @@ -109,7 +114,7 @@ void CpuToCannMemCpy(void* dst, const void* src, size_t num_bytes); void CannToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetCannToHostMemCpyFunction(); +const std::unordered_map* GetCannToHostMemCpyFunction(); bool IsCannDeviceIdValid(const onnxruntime::logging::Logger& logger, int id); @@ -127,17 +132,18 @@ void CpuToRocmMemCpy(void* dst, const void* src, size_t num_bytes); void RocmToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetRocmToHostMemCpyFunction(); +const std::unordered_map* GetRocmToHostMemCpyFunction(); #endif void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, const std::string& name_input, const pybind11::object& value, OrtValue* p_mlvalue, - bool accept_only_numpy_array = false, bool use_numpy_data_memory = true, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy); + bool accept_only_numpy_array = false, bool use_numpy_data_memory = true, + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy); pybind11::object GetPyObjFromTensor(const OrtValue& rtensor, const DataTransferManager* data_transfer_manager = nullptr, - const std::unordered_map* mem_cpy_to_host_functions = nullptr); + const std::unordered_map* mem_cpy_to_host_functions = nullptr); // The below two functions are used to convert OrtValue to numpy arrays diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index d1d4d6f3cdad5..7234543eb14de 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -23,42 +23,57 @@ std::unique_ptr OrtValueFromShapeAndType(const std::vector& s MLDataType element_type, const OrtDevice& device) { AllocatorPtr allocator; + if (strcmp(GetDeviceName(device), CPU) == 0) { allocator = GetAllocator(); - } else if (strcmp(GetDeviceName(device), CUDA) == 0) { + } else { +#if !defined(ORT_MINIMAL_BUILD) + // prefer a shared allocator from the environment. + // these are provided by plugin EPs or custom allocators explicitly registered by the user. + allocator = GetSharedAllocator(device); +#endif + + if (!allocator) { + if (strcmp(GetDeviceName(device), CUDA) == 0) { #ifdef USE_CUDA - if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } - allocator = GetCudaAllocator(device.Id()); + if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { + throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); + } + + allocator = GetCudaAllocator(device.Id()); #else - throw std::runtime_error( - "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " - "Please use the CUDA package of OnnxRuntime to use this feature."); + throw std::runtime_error( + "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " + "Please use the CUDA package of OnnxRuntime to use this feature."); #endif - } else if (strcmp(GetDeviceName(device), HIP) == 0) { + } else if (strcmp(GetDeviceName(device), HIP) == 0) { #if USE_ROCM - if (!IsRocmDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } - allocator = GetRocmAllocator(device.Id()); + if (!IsRocmDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { + throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); + } + + allocator = GetRocmAllocator(device.Id()); #elif USE_MIGRAPHX - allocator = GetMIGraphXAllocator(device.Id()); + allocator = GetMIGraphXAllocator(device.Id()); #else - throw std::runtime_error( - "Can't allocate memory on the AMD device using this package of OnnxRuntime. " - "Please use the ROCm package of OnnxRuntime to use this feature."); + throw std::runtime_error( + "Can't allocate memory on the AMD device using this package of OnnxRuntime. " + "Please use the ROCm package of OnnxRuntime to use this feature."); #endif - } else if (strcmp(GetDeviceName(device), DML) == 0) { + } else if (strcmp(GetDeviceName(device), DML) == 0) { #if USE_DML - allocator = GetDmlAllocator(device.Id()); + allocator = GetDmlAllocator(device.Id()); #else - throw std::runtime_error( - "Can't allocate memory on the DirectML device using this package of OnnxRuntime. " - "Please use the DirectML package of OnnxRuntime to use this feature."); + throw std::runtime_error( + "Can't allocate memory on the DirectML device using this package of OnnxRuntime. " + "Please use the DirectML package of OnnxRuntime to use this feature."); #endif - } else { - throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device"); + } + } + + if (!allocator) { + throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device"); + } } auto ml_value = std::make_unique(); @@ -90,7 +105,8 @@ void addOrtValueMethods(pybind11::module& m) { if (device.Vendor() == OrtDevice::VendorIds::MICROSOFT) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in DML + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in DML CreateGenericMLValue( nullptr, GetDmlAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToDmlMemCpy); } else @@ -103,8 +119,10 @@ void addOrtValueMethods(pybind11::module& m) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in CUDA - CreateGenericMLValue(nullptr, GetCudaAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToCudaMemCpy); + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in CUDA + CreateGenericMLValue(nullptr, GetCudaAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToCudaMemCpy); } else #endif #ifdef USE_ROCM @@ -115,22 +133,34 @@ void addOrtValueMethods(pybind11::module& m) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in CUDA - CreateGenericMLValue(nullptr, GetRocmAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToRocmMemCpy); + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in ROCM + CreateGenericMLValue(nullptr, GetRocmAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToRocmMemCpy); } else #endif #if USE_MIGRAPHX if (device.Vendor() == OrtDevice::VendorIds::AMD) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in MIGraphX - CreateGenericMLValue(nullptr, GetMIGraphXAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToMIGraphXMemCpy); + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in MIGraphX + CreateGenericMLValue(nullptr, GetMIGraphXAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToMIGraphXMemCpy); } else #endif { - throw std::runtime_error( - "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " - "Please use the CUDA package of OnnxRuntime to use this feature."); + // see if we can do the copy with an allocator and IDataTransfer registered by a plugin EP + auto allocator = GetSharedAllocator(device); + auto cpu_to_device_copy_fn = allocator ? CreateDataTransferMemCpy(OrtDevice{}, device) : nullptr; + if (cpu_to_device_copy_fn) { + CreateGenericMLValue(nullptr, allocator, "", array_on_cpu, ml_value.get(), true, false, + cpu_to_device_copy_fn); + } else { + throw std::runtime_error( + "Can't allocate memory on the device using this package of OnnxRuntime. " + "Please use the appropriate package of OnnxRuntime for your hardware to use this feature."); + } } } else if (device.Type() == OrtDevice::NPU && device.Vendor() == OrtDevice::VendorIds::HUAWEI) { #ifdef USE_CANN @@ -214,8 +244,16 @@ void addOrtValueMethods(pybind11::module& m) { } else #endif { - throw std::runtime_error( - "Unsupported GPU device: Cannot find the supported GPU device."); + // see if we can do the copy with an allocator and IDataTransfer registered by a plugin EP + auto allocator = GetSharedAllocator(device); + auto cpu_to_device_copy_fn = allocator ? CreateDataTransferMemCpy(OrtDevice{}, device) : nullptr; + if (cpu_to_device_copy_fn) { + onnxruntime::python::CopyDataToTensor(py_values, values_type, *(ml_value->GetMutable()), + cpu_to_device_copy_fn); + } else { + throw std::runtime_error( + "Unsupported GPU device: Cannot find the supported GPU device."); + } } } else if (device.Type() == OrtDevice::DML) { #if USE_DML diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index acf0681cf8752..03ad0185d1394 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -205,7 +205,7 @@ void AppendLoraParametersAsInputs(const RunOptions& run_options, template static py::object AddNonTensor(const OrtValue& val, const DataTransferManager* /*data_transfer_manager*/, - const std::unordered_map* /*mem_cpy_to_host_functions*/) { + const std::unordered_map* /*mem_cpy_to_host_functions*/) { return py::cast(val.Get()); } @@ -265,39 +265,65 @@ pybind11::array PrimitiveTensorToNumpyFromDevice(const OrtValue& ort_value, cons // pretty much does what a DataTransferManager does - copy data from device(s) to the host py::object GetPyObjFromTensor(const OrtValue& ort_value, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { ORT_ENFORCE(ort_value.IsTensor(), "This function only supports tensors"); const auto& tensor = ort_value.Get(); + const auto& device = tensor.Location().device; + if (tensor.IsDataTypeString()) { - ORT_ENFORCE(tensor.Location().device.Type() == OrtDevice::CPU, "Strings can only be on CPU"); + ORT_ENFORCE(device.Type() == OrtDevice::CPU, "Strings can only be on CPU"); // Create a numpy array of strings (python objects) by copy/converting them py::array result = StringTensorToNumpyArray(tensor); return py::cast(result); } - const auto device_type = tensor.Location().device.Type(); + const auto device_type = device.Type(); // Create an numpy array on top of the OrtValue memory, no copy if (device_type == OrtDevice::CPU) { py::array result = PrimitiveTensorToNumpyOverOrtValue(ort_value); return py::cast(result); } - if (!data_transfer_manager && !mem_cpy_to_host_functions) { - throw std::runtime_error( - "GetPyObjFromTensor: Either data transfer manager or a " - "function to copy data to the host is needed to convert non-CPU tensor to numpy array"); - } - py::array result; if (data_transfer_manager != nullptr) { result = PrimitiveTensorToNumpyFromDevice(ort_value, data_transfer_manager); } else { - auto mem_cpy_to_host = mem_cpy_to_host_functions->find(device_type); - ORT_ENFORCE(mem_cpy_to_host != mem_cpy_to_host_functions->end(), - "Unable to locate a function that can copy data to the host from the device"); - result = PrimitiveTensorToNumpyFromDevice(ort_value, mem_cpy_to_host->second); + bool copied = false; + if (mem_cpy_to_host_functions) { + auto it = std::find_if(mem_cpy_to_host_functions->begin(), mem_cpy_to_host_functions->end(), + [&device](const auto& entry) { + const auto& copy_device = entry.first; + // We're ignoring OrtDevice.Id() currently for historical reasons. + // The key to mem_cpy_to_host_functions was previously the device type (CPU/GPU/NPU). + // This changed to be OrtDevice to get the vendor id. + // Assumably it would be better to also match on device id, but that was not possible + // previously and to preserve existing behavior we keep the old logic and expect the + // copy function to handle the device id correctly. + return device.Type() == copy_device.Type() && + device.MemType() == copy_device.MemType() && + device.Vendor() == copy_device.Vendor(); + }); + + if (it != mem_cpy_to_host_functions->end()) { + result = PrimitiveTensorToNumpyFromDevice(ort_value, it->second); + copied = true; + } + } + + if (!copied) { + // see if we have a shared data transfer function from a plugin EP + auto device_to_cpu_copy_func = CreateDataTransferMemCpy(device, OrtDevice{}); + if (device_to_cpu_copy_func) { + result = PrimitiveTensorToNumpyFromDevice(ort_value, device_to_cpu_copy_func); + } else { + throw std::runtime_error( + "GetPyObjFromTensor: Either data transfer manager or a " + "function to copy data to the host is needed to convert non-CPU tensor to numpy array"); + } + } } + return py::cast(result); } @@ -373,7 +399,7 @@ py::object GetPyObjectFromSparseTensor(size_t pos, const OrtValue& ort_value, co template <> py::object AddNonTensor(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { const auto& seq_tensors = val.Get(); py::list py_list; for (const auto& ort_value : seq_tensors) { @@ -389,7 +415,7 @@ py::object AddNonTensor(const OrtValue& val, py::object AddNonTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { // Should be in sync with core/framework/datatypes.h auto val_type = val.Type(); if (val_type->IsTensorSequenceType()) { @@ -429,7 +455,7 @@ py::object AddNonTensorAsPyObj(const OrtValue& val, } py::object AddTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { return GetPyObjFromTensor(val, data_transfer_manager, mem_cpy_to_host_functions); } @@ -1885,6 +1911,10 @@ void addObjectMethods(py::module& m, ExecutionProviderRegistrationFn ep_registra vendor = OrtDevice::VendorIds::NVIDIA; #elif USE_ROCM || USE_MIGRAPHX vendor = OrtDevice::VendorIds::AMD; +#endif + } else if (type == OrtDevice::NPU) { +#if USE_CANN + vendor = OrtDevice::VendorIds::HUAWEI; #endif } @@ -1894,12 +1924,15 @@ void addObjectMethods(py::module& m, ExecutionProviderRegistrationFn ep_registra .def("device_id", &OrtDevice::Id, R"pbdoc(Device Id.)pbdoc") .def("device_type", &OrtDevice::Type, R"pbdoc(Device Type.)pbdoc") .def("vendor_id", &OrtDevice::Vendor, R"pbdoc(Vendor Id.)pbdoc") + // generic device types that are typically used with a vendor id. .def_static("cpu", []() { return OrtDevice::CPU; }) + .def_static("gpu", []() { return OrtDevice::GPU; }) + .def_static("npu", []() { return OrtDevice::NPU; }) + // EP specific device types for backward compatibility. .def_static("cuda", []() { return OrtDevice::GPU; }) .def_static("cann", []() { return OrtDevice::NPU; }) - .def_static("fpga", []() { return OrtDevice::FPGA; }) - .def_static("npu", []() { return OrtDevice::NPU; }) .def_static("dml", []() { return OrtDevice::DML; }) + .def_static("fpga", []() { return OrtDevice::FPGA; }) .def_static("webgpu", []() { return OrtDevice::GPU; }) .def_static("default_memory", []() { return OrtDevice::MemType::DEFAULT; }); diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py index a236c4da1738e..85b3632c516ca 100644 --- a/onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py @@ -6,7 +6,6 @@ import json import logging import os -import textwrap from pathlib import Path import numpy as np @@ -93,16 +92,10 @@ def save_processing( if separate_encoder_and_decoder_init: return - audio_processor_json = textwrap.dedent("""\ - { + audio_processor_cfg = { "feature_extraction": { "sequence": [ - { - "operation": { - "name": "audio_decoder", - "type": "AudioDecoder" - } - }, + {"operation": {"name": "audio_decoder", "type": "AudioDecoder"}}, { "operation": { "name": "STFT", @@ -511,27 +504,23 @@ def save_processing( 0.000986635684967041, 0.0005550682544708252, 0.0002467334270477295, - 0.0000616908073425293 - ] - } + 0.0000616908073425293, + ], + }, } }, { "operation": { "name": "log_mel_spectrogram", "type": "LogMelSpectrum", - "attrs": { - "chunk_size": 30, - "hop_length": 160, - "n_fft": 400, - "n_mel": 80 - } + "attrs": {"chunk_size": 30, "hop_length": 160, "n_fft": 400, "n_mel": config.num_mel_bins}, } - } + }, ] } } - """) + audio_processor_json = json.dumps(audio_processor_cfg, indent=4) + with open(os.path.join(output_dir, "audio_processor_config.json"), "w") as f: f.write(audio_processor_json) diff --git a/onnxruntime/test/framework/cuda/fence_cuda_test.cc b/onnxruntime/test/framework/cuda/fence_cuda_test.cc index b86f3efeefafd..fced72ce3246d 100644 --- a/onnxruntime/test/framework/cuda/fence_cuda_test.cc +++ b/onnxruntime/test/framework/cuda/fence_cuda_test.cc @@ -67,7 +67,7 @@ static common::Status LoadInferenceSessionFromModel(FenceCudaTestInferenceSessio tensor_proto.set_data_type(PROTO_DATATYPE); \ for (auto v : value) tensor_proto.PROTO_ADD_DATA(v); \ tensor_proto.set_name(name); \ - return graph_utils::AddInitializerWithExternalData(graph, tensor_proto); \ + return graph_utils::AddInitializer(graph, tensor_proto); \ } CREATE_INITIALIZER_FUNC(float, TensorProto_DataType_FLOAT, add_float_data) diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index e2b54950e7b24..ca1166e19037c 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -1894,14 +1894,21 @@ TEST_F(GraphTest, AddRemoveInitializerHandling) { ASSERT_EQ(graph_proto_from_graph.initializer_size(), 2); auto validate_proto = [&](const GraphProto& proto) { + // Due to changes in a way we generate ToGraphProto() const, we can not guarantee the order of initializers + // in the generated GraphProto. auto initializers = proto.initializer(); - // we expect '2' to be before '1' due to the remove moving the last initializer into the slot of the one being - // removed in order to free memory and only move one entry - EXPECT_EQ(initializers[0].name(), init2.name()); - EXPECT_EQ(initializers[0].int32_data()[0], 2); - - EXPECT_EQ(initializers[1].name(), init.name()); - EXPECT_EQ(initializers[1].int32_data()[0], 1); + auto hit = std::find_if(initializers.begin(), initializers.end(), + [&init](const ONNX_NAMESPACE::TensorProto& t) { return t.name() == init.name(); }); + EXPECT_NE(hit, initializers.end()) + << "Initializer with name '" << init.name() << "' not found in the proto."; + EXPECT_EQ(hit->int32_data()[0], 1); + + hit = std::find_if(initializers.begin(), initializers.end(), + [&init2](const ONNX_NAMESPACE::TensorProto& t) { return t.name() == init2.name(); }); + EXPECT_NE(hit, initializers.end()) + << "Initializer with name '" << init2.name() << "' not found in the proto."; + + EXPECT_EQ(hit->int32_data()[0], 2); }; validate_proto(graph_proto_from_const_graph); diff --git a/onnxruntime/test/perftest/command_args_parser.cc b/onnxruntime/test/perftest/command_args_parser.cc index 843875a881f0a..5c81696d5c57e 100644 --- a/onnxruntime/test/perftest/command_args_parser.cc +++ b/onnxruntime/test/perftest/command_args_parser.cc @@ -4,6 +4,7 @@ // Licensed under the MIT License. #include "command_args_parser.h" +#include "utils.h" #include #include @@ -11,14 +12,6 @@ #include #include -// Windows Specific -#ifdef _WIN32 -#include "getopt.h" -#include "windows.h" -#else -#include -#endif - #include #include #include @@ -26,161 +19,163 @@ #include "test_configuration.h" #include "strings_helper.h" +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/flags/usage.h" +#include "absl/flags/usage_config.h" +#include "absl/flags/reflection.h" + +static const onnxruntime::perftest::PerformanceTestConfig& DefaultPerformanceTestConfig() { + static onnxruntime::perftest::PerformanceTestConfig default_config{}; + return default_config; +} + +ABSL_FLAG(std::string, f, "", "Specifies a free dimension by name to override to a specific value for performance optimization."); +ABSL_FLAG(std::string, F, "", "Specifies a free dimension by denotation to override to a specific value for performance optimization."); +ABSL_FLAG(std::string, m, "duration", "Specifies the test mode. Value could be 'duration' or 'times'."); +ABSL_FLAG(std::string, e, "cpu", "Specifies the provider 'cpu','cuda','dnnl','tensorrt', 'nvtensorrtrtx', 'openvino', 'dml', 'acl', 'nnapi', 'coreml', 'qnn', 'snpe', 'rocm', 'migraphx', 'xnnpack', 'vitisai' or 'webgpu'."); +ABSL_FLAG(size_t, r, DefaultPerformanceTestConfig().run_config.repeated_times, "Specifies the repeated times if running in 'times' test mode."); +ABSL_FLAG(size_t, t, DefaultPerformanceTestConfig().run_config.duration_in_seconds, "Specifies the seconds to run for 'duration' mode."); +ABSL_FLAG(std::string, p, "", "Specifies the profile name to enable profiling and dump the profile data to the file."); +ABSL_FLAG(int, x, DefaultPerformanceTestConfig().run_config.intra_op_num_threads, "Sets the number of threads used to parallelize the execution within nodes, A value of 0 means ORT will pick a default. Must >=0."); +ABSL_FLAG(int, y, DefaultPerformanceTestConfig().run_config.inter_op_num_threads, "Sets the number of threads used to parallelize the execution of the graph (across nodes), A value of 0 means ORT will pick a default. Must >=0."); +ABSL_FLAG(size_t, c, DefaultPerformanceTestConfig().run_config.concurrent_session_runs, "Specifies the (max) number of runs to invoke simultaneously."); +ABSL_FLAG(int, d, DefaultPerformanceTestConfig().run_config.cudnn_conv_algo, "Specifies CUDNN convolution algorithms: 0(benchmark), 1(heuristic), 2(default)."); +ABSL_FLAG(int, o, DefaultPerformanceTestConfig().run_config.optimization_level, "Specifies graph optimization level. Default is 99 (all). Valid values are 0 (disable), 1 (basic), 2 (extended), 3 (layout), 99 (all)."); +ABSL_FLAG(std::string, u, "", "Specifies the optimized model path for saving."); +ABSL_FLAG(std::string, i, "", + "Specifies EP specific runtime options as key-value pairs.\n Different runtime options available are: \n" + " [Usage]: -e -i '| |'\n" + "\n" + " [ACL only] [enable_fast_math]: Options: 'true', 'false', default: 'false', \n" + "\n" + " [DML only] [performance_preference]: DML device performance preference, options: 'default', 'minimum_power', 'high_performance', \n" + " [DML only] [device_filter]: DML device filter, options: 'any', 'gpu', 'npu', \n" + " [DML only] [disable_metacommands]: Options: 'true', 'false', \n" + " [DML only] [enable_graph_capture]: Options: 'true', 'false', \n" + " [DML only] [enable_graph_serialization]: Options: 'true', 'false', \n" + "\n" + " [OpenVINO only] [device_type]: Overrides the accelerator hardware type and precision with these values at runtime.\n" + " [OpenVINO only] [device_id]: Selects a particular hardware device for inference.\n" + " [OpenVINO only] [num_of_threads]: Overrides the accelerator hardware type and precision with these values at runtime.\n" + " [OpenVINO only] [cache_dir]: Explicitly specify the path to dump and load the blobs(Model caching) or cl_cache (Kernel Caching) files feature. If blob files are already present, it will be directly loaded.\n" + " [OpenVINO only] [enable_opencl_throttling]: Enables OpenCL queue throttling for GPU device(Reduces the CPU Utilization while using GPU) \n" + " [Example] [For OpenVINO EP] -e openvino -i \"device_type|CPU num_of_threads|5 enable_opencl_throttling|true cache_dir|\"\"\"\n" + "\n" + " [QNN only] [backend_type]: QNN backend type. E.g., 'cpu', 'htp'. Mutually exclusive with 'backend_path'.\n" + " [QNN only] [backend_path]: QNN backend path. E.g., '/folderpath/libQnnHtp.so', '/winfolderpath/QnnHtp.dll'. Mutually exclusive with 'backend_type'.\n" + " [QNN only] [profiling_level]: QNN profiling level, options: 'basic', 'detailed', default 'off'.\n" + " [QNN only] [profiling_file_path] : QNN profiling file path if ETW not enabled.\n" + " [QNN only] [rpc_control_latency]: QNN rpc control latency. default to 10.\n" + " [QNN only] [vtcm_mb]: QNN VTCM size in MB. default to 0(not set).\n" + " [QNN only] [htp_performance_mode]: QNN performance mode, options: 'burst', 'balanced', 'default', 'high_performance', \n" + " 'high_power_saver', 'low_balanced', 'extreme_power_saver', 'low_power_saver', 'power_saver', 'sustained_high_performance'. Default to 'default'. \n" + " [QNN only] [op_packages]: QNN UDO package, allowed format: \n" + " op_packages|::[:],::[:]. \n" + " [QNN only] [qnn_context_priority]: QNN context priority, options: 'low', 'normal', 'normal_high', 'high'. Default to 'normal'. \n" + " [QNN only] [qnn_saver_path]: QNN Saver backend path. e.g '/folderpath/libQnnSaver.so'.\n" + " [QNN only] [htp_graph_finalization_optimization_mode]: QNN graph finalization optimization mode, options: \n" + " '0', '1', '2', '3', default is '0'.\n" + " [QNN only] [soc_model]: The SoC Model number. Refer to QNN SDK documentation for specific values. Defaults to '0' (unknown). \n" + " [QNN only] [htp_arch]: The minimum HTP architecture. The driver will use ops compatible with this architecture. \n" + " Options are '0', '68', '69', '73', '75'. Defaults to '0' (none). \n" + " [QNN only] [device_id]: The ID of the device to use when setting 'htp_arch'. Defaults to '0' (for single device). \n" + " [QNN only] [enable_htp_fp16_precision]: Enable the HTP_FP16 precision so that the float32 model will be inferenced with fp16 precision. \n" + " Otherwise, it will be fp32 precision. Works for float32 model for HTP backend. Defaults to '1' (with FP16 precision.). \n" + " [QNN only] [offload_graph_io_quantization]: Offload graph input quantization and graph output dequantization to another EP (typically CPU EP). \n" + " Defaults to '0' (QNN EP handles the graph I/O quantization and dequantization). \n" + " [QNN only] [enable_htp_spill_fill_buffer]: Enable HTP spill fill buffer, used while generating QNN context binary.\n" + " [QNN only] [enable_htp_shared_memory_allocator]: Enable the QNN HTP shared memory allocator and use it for inputs and outputs. Requires libcdsprpc.so/dll to be available.\n" + " Defaults to '0' (disabled).\n" + " [Example] [For QNN EP] -e qnn -i \"backend_type|cpu\" \n" + "\n" + " [TensorRT only] [trt_max_partition_iterations]: Maximum iterations for TensorRT parser to get capability.\n" + " [TensorRT only] [trt_min_subgraph_size]: Minimum size of TensorRT subgraphs.\n" + " [TensorRT only] [trt_max_workspace_size]: Set TensorRT maximum workspace size in byte.\n" + " [TensorRT only] [trt_fp16_enable]: Enable TensorRT FP16 precision.\n" + " [TensorRT only] [trt_int8_enable]: Enable TensorRT INT8 precision.\n" + " [TensorRT only] [trt_int8_calibration_table_name]: Specify INT8 calibration table name.\n" + " [TensorRT only] [trt_int8_use_native_calibration_table]: Use Native TensorRT calibration table.\n" + " [TensorRT only] [trt_dla_enable]: Enable DLA in Jetson device.\n" + " [TensorRT only] [trt_dla_core]: DLA core number.\n" + " [TensorRT only] [trt_dump_subgraphs]: Dump TRT subgraph to onnx model.\n" + " [TensorRT only] [trt_engine_cache_enable]: Enable engine caching.\n" + " [TensorRT only] [trt_engine_cache_path]: Specify engine cache path.\n" + " [TensorRT only] [trt_engine_cache_prefix]: Customize engine cache prefix when trt_engine_cache_enable is true.\n" + " [TensorRT only] [trt_engine_hw_compatible]: Enable hardware compatibility. Engines ending with '_sm80+' can be re-used across all Ampere+ GPU (a hardware-compatible engine may have lower throughput and/or higher latency than its non-hardware-compatible counterpart).\n" + " [TensorRT only] [trt_weight_stripped_engine_enable]: Enable weight-stripped engine build.\n" + " [TensorRT only] [trt_onnx_model_folder_path]: Folder path for the ONNX model with weights.\n" + " [TensorRT only] [trt_force_sequential_engine_build]: Force TensorRT engines to be built sequentially.\n" + " [TensorRT only] [trt_context_memory_sharing_enable]: Enable TensorRT context memory sharing between subgraphs.\n" + " [TensorRT only] [trt_layer_norm_fp32_fallback]: Force Pow + Reduce ops in layer norm to run in FP32 to avoid overflow.\n" + " [Example] [For TensorRT EP] -e tensorrt -i 'trt_fp16_enable|true trt_int8_enable|true trt_int8_calibration_table_name|calibration.flatbuffers trt_int8_use_native_calibration_table|false trt_force_sequential_engine_build|false'\n" + "\n" + " [NNAPI only] [NNAPI_FLAG_USE_FP16]: Use fp16 relaxation in NNAPI EP..\n" + " [NNAPI only] [NNAPI_FLAG_USE_NCHW]: Use the NCHW layout in NNAPI EP.\n" + " [NNAPI only] [NNAPI_FLAG_CPU_DISABLED]: Prevent NNAPI from using CPU devices.\n" + " [NNAPI only] [NNAPI_FLAG_CPU_ONLY]: Using CPU only in NNAPI EP.\n" + " [Example] [For NNAPI EP] -e nnapi -i \"NNAPI_FLAG_USE_FP16 NNAPI_FLAG_USE_NCHW NNAPI_FLAG_CPU_DISABLED\"\n" + "\n" + " [CoreML only] [ModelFormat]:[MLProgram, NeuralNetwork] Create an ML Program model or Neural Network. Default is NeuralNetwork.\n" + " [CoreML only] [MLComputeUnits]:[CPUAndNeuralEngine CPUAndGPU ALL CPUOnly] Specify to limit the backend device used to run the model.\n" + " [CoreML only] [AllowStaticInputShapes]:[0 1].\n" + " [CoreML only] [EnableOnSubgraphs]:[0 1].\n" + " [CoreML only] [SpecializationStrategy]:[Default FastPrediction].\n" + " [CoreML only] [ProfileComputePlan]:[0 1].\n" + " [CoreML only] [AllowLowPrecisionAccumulationOnGPU]:[0 1].\n" + " [CoreML only] [ModelCacheDirectory]:[path../a/b/c].\n" + " [Example] [For CoreML EP] -e coreml -i \"ModelFormat|MLProgram MLComputeUnits|CPUAndGPU\"\n" + "\n" + " [SNPE only] [runtime]: SNPE runtime, options: 'CPU', 'GPU', 'GPU_FLOAT16', 'DSP', 'AIP_FIXED_TF'. \n" + " [SNPE only] [priority]: execution priority, options: 'low', 'normal'. \n" + " [SNPE only] [buffer_type]: options: 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. default: ITENSOR'. \n" + " [SNPE only] [enable_init_cache]: enable SNPE init caching feature, set to 1 to enabled it. Disabled by default. \n" + " [Example] [For SNPE EP] -e snpe -i \"runtime|CPU priority|low\" \n"); +ABSL_FLAG(int, S, DefaultPerformanceTestConfig().run_config.random_seed_for_input_data, "Given random seed, to produce the same input data. This defaults to -1(no initialize)."); +ABSL_FLAG(std::string, T, "", "Specifies intra op thread affinity string."); +ABSL_FLAG(std::string, C, "", + "Specifies session configuration entries as key-value pairs:\n -C \"| |\" \n" + "Refer to onnxruntime_session_options_config_keys.h for valid keys and values. \n" + "[Example] -C \"session.disable_cpu_ep_fallback|1 ep.context_enable|1\" \n"); +ABSL_FLAG(std::string, R, "", "Allows user to register custom op by .so or .dll file."); +ABSL_FLAG(bool, A, DefaultPerformanceTestConfig().run_config.enable_cpu_mem_arena, "Disables memory arena."); +ABSL_FLAG(bool, M, DefaultPerformanceTestConfig().run_config.enable_memory_pattern, "Disables memory pattern."); +ABSL_FLAG(bool, s, DefaultPerformanceTestConfig().run_config.f_dump_statistics, "Shows statistics result, like P75, P90. If no result_file provided this defaults to on."); +ABSL_FLAG(bool, v, DefaultPerformanceTestConfig().run_config.f_verbose, "Shows verbose information."); +ABSL_FLAG(bool, I, DefaultPerformanceTestConfig().run_config.generate_model_input_binding, "Generates tensor input binding. Free dimensions are treated as 1 unless overridden using -f."); +ABSL_FLAG(bool, P, false, "Uses parallel executor instead of sequential executor."); +ABSL_FLAG(bool, q, DefaultPerformanceTestConfig().run_config.do_cuda_copy_in_separate_stream, "[CUDA only] Uses separate stream for copy."); +ABSL_FLAG(bool, z, DefaultPerformanceTestConfig().run_config.set_denormal_as_zero, "Sets denormal as zero. When turning on this option reduces latency dramatically, a model may have denormals."); +ABSL_FLAG(bool, D, DefaultPerformanceTestConfig().run_config.disable_spinning, "Disables spinning entirely for thread owned by onnxruntime intra-op thread pool."); +ABSL_FLAG(bool, Z, DefaultPerformanceTestConfig().run_config.disable_spinning_between_run, "Disallows thread from spinning during runs to reduce cpu usage."); +ABSL_FLAG(bool, n, DefaultPerformanceTestConfig().run_config.exit_after_session_creation, "Allows user to measure session creation time to measure impact of enabling any initialization optimizations."); +ABSL_FLAG(bool, l, DefaultPerformanceTestConfig().model_info.load_via_path, "Provides file as binary in memory by using fopen before session creation."); +ABSL_FLAG(bool, g, DefaultPerformanceTestConfig().run_config.enable_cuda_io_binding, "[TensorRT RTX | TensorRT | CUDA] Enables tensor input and output bindings on CUDA before session run."); +ABSL_FLAG(bool, X, DefaultPerformanceTestConfig().run_config.use_extensions, "Registers custom ops from onnxruntime-extensions."); +ABSL_FLAG(std::string, plugin_ep_libs, "", + "Specifies a list of plugin execution provider (EP) registration names and their corresponding shared libraries to register.\n" + "[Usage]: --plugin_ep_libs \"plugin_ep_name_1|plugin_ep_1.dll plugin_ep_name_2|plugin_ep_2.dll ... \""); +ABSL_FLAG(std::string, plugin_eps, "", "Specifies a semicolon-separated list of plugin execution providers (EPs) to use."); +ABSL_FLAG(std::string, plugin_ep_options, "", + "Specifies provider options for each EP listed in --plugin_eps. Options (key-value pairs) for each EP are separated by space and EPs are separated by semicolons.\n" + "[Usage]: --plugin_ep_options \"ep_1_option_1_key|ep_1_option_1_value ...;ep_2_option_1_key|ep_2_option_1_value ...;... \" or \n" + "--plugin_ep_options \";ep_2_option_1_key|ep_2_option_1_value ...;... \" or \n" + "--plugin_ep_options \"ep_1_option_1_key|ep_1_option_1_value ...;;ep_3_option_1_key|ep_3_option_1_value ...;... \""); +ABSL_FLAG(bool, list_ep_devices, false, "Prints all available device indices and their properties (including metadata). This option makes the program exit early without performing inference.\n"); +ABSL_FLAG(std::string, select_ep_devices, "", "Specifies a semicolon-separated list of device indices to add to the session and run with."); +ABSL_FLAG(bool, h, false, "Print program usage."); + namespace onnxruntime { namespace perftest { -/*static*/ void CommandLineParser::ShowUsage() { - printf( - "perf_test [options...] model_path [result_file]\n" - "Options:\n" - "\t-m [test_mode]: Specifies the test mode. Value could be 'duration' or 'times'.\n" - "\t\tProvide 'duration' to run the test for a fix duration, and 'times' to repeated for a certain times. \n" - "\t-M: Disable memory pattern.\n" - "\t-A: Disable memory arena\n" - "\t-I: Generate tensor input binding. Free dimensions are treated as 1 unless overridden using -f.\n" - "\t-c [parallel runs]: Specifies the (max) number of runs to invoke simultaneously. Default:1.\n" - "\t-e [cpu|cuda|dnnl|tensorrt|openvino|dml|acl|nnapi|coreml|qnn|snpe|rocm|migraphx|xnnpack|vitisai|webgpu]: Specifies the provider 'cpu','cuda','dnnl','tensorrt', " - "'nvtensorrtrtx', 'openvino', 'dml', 'acl', 'nnapi', 'coreml', 'qnn', 'snpe', 'rocm', 'migraphx', 'xnnpack', 'vitisai' or 'webgpu'. " - "Default:'cpu'.\n" - "\t-b [tf|ort]: backend to use. Default:ort\n" - "\t-r [repeated_times]: Specifies the repeated times if running in 'times' test mode.Default:1000.\n" - "\t-t [seconds_to_run]: Specifies the seconds to run for 'duration' mode. Default:600.\n" - "\t-p [profile_file]: Specifies the profile name to enable profiling and dump the profile data to the file.\n" - "\t-s: Show statistics result, like P75, P90. If no result_file provided this defaults to on.\n" - "\t-S: Given random seed, to produce the same input data. This defaults to -1(no initialize).\n" - "\t-v: Show verbose information.\n" - "\t-x [intra_op_num_threads]: Sets the number of threads used to parallelize the execution within nodes, A value of 0 means ORT will pick a default. Must >=0.\n" - "\t-y [inter_op_num_threads]: Sets the number of threads used to parallelize the execution of the graph (across nodes), A value of 0 means ORT will pick a default. Must >=0.\n" - "\t-f [free_dimension_override]: Specifies a free dimension by name to override to a specific value for performance optimization. " - "Syntax is [dimension_name:override_value]. override_value must > 0\n" - "\t-F [free_dimension_override]: Specifies a free dimension by denotation to override to a specific value for performance optimization. " - "Syntax is [dimension_denotation:override_value]. override_value must > 0\n" - "\t-P: Use parallel executor instead of sequential executor.\n" - "\t-o [optimization level]: Default is 99 (all). Valid values are 0 (disable), 1 (basic), 2 (extended), 3 (layout), 99 (all).\n" - "\t\tPlease see onnxruntime_c_api.h (enum GraphOptimizationLevel) for the full list of all optimization levels.\n" - "\t-u [optimized_model_path]: Specify the optimized model path for saving.\n" - "\t-d [CUDA only][cudnn_conv_algorithm]: Specify CUDNN convolution algorithms: 0(benchmark), 1(heuristic), 2(default). \n" - "\t-q [CUDA only] use separate stream for copy. \n" - "\t-g [TensorRT RTX | TensorRT | CUDA] Enable tensor input and output bindings on CUDA before session run \n" - "\t-z: Set denormal as zero. When turning on this option reduces latency dramatically, a model may have denormals.\n" - "\t-C: Specify session configuration entries as key-value pairs: -C \"| |\" \n" - "\t Refer to onnxruntime_session_options_config_keys.h for valid keys and values. \n" - "\t [Example] -C \"session.disable_cpu_ep_fallback|1 ep.context_enable|1\" \n" - "\t-i: Specify EP specific runtime options as key value pairs. Different runtime options available are: \n" - "\t [Usage]: -e -i '| |'\n" - "\n" - "\t [ACL only] [enable_fast_math]: Options: 'true', 'false', default: 'false', \n" - "\t [DML only] [performance_preference]: DML device performance preference, options: 'default', 'minimum_power', 'high_performance', \n" - "\t [DML only] [device_filter]: DML device filter, options: 'any', 'gpu', 'npu', \n" - "\t [DML only] [disable_metacommands]: Options: 'true', 'false', \n" - "\t [DML only] [enable_graph_capture]: Options: 'true', 'false', \n" - "\t [DML only] [enable_graph_serialization]: Options: 'true', 'false', \n" - "\n" - "\t [OpenVINO only] [device_type]: Overrides the accelerator hardware type and precision with these values at runtime.\n" - "\t [OpenVINO only] [device_id]: Selects a particular hardware device for inference.\n" - "\t [OpenVINO only] [num_of_threads]: Overrides the accelerator hardware type and precision with these values at runtime.\n" - "\t [OpenVINO only] [cache_dir]: Explicitly specify the path to dump and load the blobs(Model caching) or cl_cache (Kernel Caching) files feature. If blob files are already present, it will be directly loaded.\n" - "\t [OpenVINO only] [enable_opencl_throttling]: Enables OpenCL queue throttling for GPU device(Reduces the CPU Utilization while using GPU) \n" - "\t [Example] [For OpenVINO EP] -e openvino -i \"device_type|CPU num_of_threads|5 enable_opencl_throttling|true cache_dir|\"\"\"\n" - "\n" - "\t [QNN only] [backend_type]: QNN backend type. E.g., 'cpu', 'htp'. Mutually exclusive with 'backend_path'.\n" - "\t [QNN only] [backend_path]: QNN backend path. E.g., '/folderpath/libQnnHtp.so', '/winfolderpath/QnnHtp.dll'. Mutually exclusive with 'backend_type'.\n" - "\t [QNN only] [profiling_level]: QNN profiling level, options: 'basic', 'detailed', default 'off'.\n" - "\t [QNN only] [profiling_file_path] : QNN profiling file path if ETW not enabled.\n" - "\t [QNN only] [rpc_control_latency]: QNN rpc control latency. default to 10.\n" - "\t [QNN only] [vtcm_mb]: QNN VTCM size in MB. default to 0(not set).\n" - "\t [QNN only] [htp_performance_mode]: QNN performance mode, options: 'burst', 'balanced', 'default', 'high_performance', \n" - "\t 'high_power_saver', 'low_balanced', 'extreme_power_saver', 'low_power_saver', 'power_saver', 'sustained_high_performance'. Default to 'default'. \n" - "\t [QNN only] [op_packages]: QNN UDO package, allowed format: \n" - "\t op_packages|::[:],::[:]. \n" - "\t [QNN only] [qnn_context_priority]: QNN context priority, options: 'low', 'normal', 'normal_high', 'high'. Default to 'normal'. \n" - "\t [QNN only] [qnn_saver_path]: QNN Saver backend path. e.g '/folderpath/libQnnSaver.so'.\n" - "\t [QNN only] [htp_graph_finalization_optimization_mode]: QNN graph finalization optimization mode, options: \n" - "\t '0', '1', '2', '3', default is '0'.\n" - "\t [QNN only] [soc_model]: The SoC Model number. Refer to QNN SDK documentation for specific values. Defaults to '0' (unknown). \n" - "\t [QNN only] [htp_arch]: The minimum HTP architecture. The driver will use ops compatible with this architecture. \n" - "\t Options are '0', '68', '69', '73', '75'. Defaults to '0' (none). \n" - "\t [QNN only] [device_id]: The ID of the device to use when setting 'htp_arch'. Defaults to '0' (for single device). \n" - "\t [QNN only] [enable_htp_fp16_precision]: Enable the HTP_FP16 precision so that the float32 model will be inferenced with fp16 precision. \n" - "\t Otherwise, it will be fp32 precision. Works for float32 model for HTP backend. Defaults to '1' (with FP16 precision.). \n" - "\t [QNN only] [offload_graph_io_quantization]: Offload graph input quantization and graph output dequantization to another EP (typically CPU EP). \n" - "\t Defaults to '0' (QNN EP handles the graph I/O quantization and dequantization). \n" - "\t [QNN only] [enable_htp_spill_fill_buffer]: Enable HTP spill fill buffer, used while generating QNN context binary.\n" - "\t [QNN only] [enable_htp_shared_memory_allocator]: Enable the QNN HTP shared memory allocator and use it for inputs and outputs. Requires libcdsprpc.so/dll to be available.\n" - "\t Defaults to '0' (disabled).\n" - "\t [Example] [For QNN EP] -e qnn -i \"backend_type|cpu\" \n" - "\n" - "\t [TensorRT only] [trt_max_partition_iterations]: Maximum iterations for TensorRT parser to get capability.\n" - "\t [TensorRT only] [trt_min_subgraph_size]: Minimum size of TensorRT subgraphs.\n" - "\t [TensorRT only] [trt_max_workspace_size]: Set TensorRT maximum workspace size in byte.\n" - "\t [TensorRT only] [trt_fp16_enable]: Enable TensorRT FP16 precision.\n" - "\t [TensorRT only] [trt_int8_enable]: Enable TensorRT INT8 precision.\n" - "\t [TensorRT only] [trt_int8_calibration_table_name]: Specify INT8 calibration table name.\n" - "\t [TensorRT only] [trt_int8_use_native_calibration_table]: Use Native TensorRT calibration table.\n" - "\t [TensorRT only] [trt_dla_enable]: Enable DLA in Jetson device.\n" - "\t [TensorRT only] [trt_dla_core]: DLA core number.\n" - "\t [TensorRT only] [trt_dump_subgraphs]: Dump TRT subgraph to onnx model.\n" - "\t [TensorRT only] [trt_engine_cache_enable]: Enable engine caching.\n" - "\t [TensorRT only] [trt_engine_cache_path]: Specify engine cache path.\n" - "\t [TensorRT only] [trt_engine_cache_prefix]: Customize engine cache prefix when trt_engine_cache_enable is true.\n" - "\t [TensorRT only] [trt_engine_hw_compatible]: Enable hardware compatibility. Engines ending with '_sm80+' can be re-used across all Ampere+ GPU (a hardware-compatible engine may have lower throughput and/or higher latency than its non-hardware-compatible counterpart).\n" - "\t [TensorRT only] [trt_weight_stripped_engine_enable]: Enable weight-stripped engine build.\n" - "\t [TensorRT only] [trt_onnx_model_folder_path]: Folder path for the ONNX model with weights.\n" - "\t [TensorRT only] [trt_force_sequential_engine_build]: Force TensorRT engines to be built sequentially.\n" - "\t [TensorRT only] [trt_context_memory_sharing_enable]: Enable TensorRT context memory sharing between subgraphs.\n" - "\t [TensorRT only] [trt_layer_norm_fp32_fallback]: Force Pow + Reduce ops in layer norm to run in FP32 to avoid overflow.\n" - "\t [Example] [For TensorRT EP] -e tensorrt -i 'trt_fp16_enable|true trt_int8_enable|true trt_int8_calibration_table_name|calibration.flatbuffers trt_int8_use_native_calibration_table|false trt_force_sequential_engine_build|false'\n" - "\n" - "\t [NNAPI only] [NNAPI_FLAG_USE_FP16]: Use fp16 relaxation in NNAPI EP..\n" - "\t [NNAPI only] [NNAPI_FLAG_USE_NCHW]: Use the NCHW layout in NNAPI EP.\n" - "\t [NNAPI only] [NNAPI_FLAG_CPU_DISABLED]: Prevent NNAPI from using CPU devices.\n" - "\t [NNAPI only] [NNAPI_FLAG_CPU_ONLY]: Using CPU only in NNAPI EP.\n" - "\t [Example] [For NNAPI EP] -e nnapi -i \"NNAPI_FLAG_USE_FP16 NNAPI_FLAG_USE_NCHW NNAPI_FLAG_CPU_DISABLED\"\n" - "\n" - "\t [CoreML only] [ModelFormat]:[MLProgram, NeuralNetwork] Create an ML Program model or Neural Network. Default is NeuralNetwork.\n" - "\t [CoreML only] [MLComputeUnits]:[CPUAndNeuralEngine CPUAndGPU ALL CPUOnly] Specify to limit the backend device used to run the model.\n" - "\t [CoreML only] [AllowStaticInputShapes]:[0 1].\n" - "\t [CoreML only] [EnableOnSubgraphs]:[0 1].\n" - "\t [CoreML only] [SpecializationStrategy]:[Default FastPrediction].\n" - "\t [CoreML only] [ProfileComputePlan]:[0 1].\n" - "\t [CoreML only] [AllowLowPrecisionAccumulationOnGPU]:[0 1].\n" - "\t [CoreML only] [ModelCacheDirectory]:[path../a/b/c].\n" - "\t [Example] [For CoreML EP] -e coreml -i \"ModelFormat|MLProgram MLComputeUnits|CPUAndGPU\"\n" - "\n" - "\t [SNPE only] [runtime]: SNPE runtime, options: 'CPU', 'GPU', 'GPU_FLOAT16', 'DSP', 'AIP_FIXED_TF'. \n" - "\t [SNPE only] [priority]: execution priority, options: 'low', 'normal'. \n" - "\t [SNPE only] [buffer_type]: options: 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. default: ITENSOR'. \n" - "\t [SNPE only] [enable_init_cache]: enable SNPE init caching feature, set to 1 to enabled it. Disabled by default. \n" - "\t [Example] [For SNPE EP] -e snpe -i \"runtime|CPU priority|low\" \n\n" - "\n" - "\t-T [Set intra op thread affinities]: Specify intra op thread affinity string\n" - "\t [Example]: -T 1,2;3,4;5,6 or -T 1-2;3-4;5-6 \n" - "\t\t Use semicolon to separate configuration between threads.\n" - "\t\t E.g. 1,2;3,4;5,6 specifies affinities for three threads, the first thread will be attached to the first and second logical processor.\n" - "\t\t The number of affinities must be equal to intra_op_num_threads - 1\n\n" - "\t-D [Disable thread spinning]: disable spinning entirely for thread owned by onnxruntime intra-op thread pool.\n" - "\t-Z [Force thread to stop spinning between runs]: disallow thread from spinning during runs to reduce cpu usage.\n" - "\t-n [Exit after session creation]: allow user to measure session creation time to measure impact of enabling any initialization optimizations.\n" - "\t-l Provide file as binary in memory by using fopen before session creation.\n" - "\t-R [Register custom op]: allow user to register custom op by .so or .dll file.\n" - "\t-X [Enable onnxruntime-extensions custom ops]: Registers custom ops from onnxruntime-extensions. " - "onnxruntime-extensions must have been built in to onnxruntime. This can be done with the build.py " - "'--use_extensions' option.\n" - "\t-h: help\n"); -} -#ifdef _WIN32 -static const ORTCHAR_T* overrideDelimiter = L":"; -#else -static const ORTCHAR_T* overrideDelimiter = ":"; -#endif -static bool ParseDimensionOverride(std::basic_string& dim_identifier, int64_t& override_val) { - std::basic_string free_dim_str(optarg); - size_t delimiter_location = free_dim_str.find(overrideDelimiter); +static bool ParseDimensionOverride(std::string& dim_identifier, int64_t& override_val, const char* option) { + std::basic_string free_dim_str(option); + size_t delimiter_location = free_dim_str.find(":"); if (delimiter_location >= free_dim_str.size() - 1) { return false; } dim_identifier = free_dim_str.substr(0, delimiter_location); - std::basic_string override_val_str = free_dim_str.substr(delimiter_location + 1, std::wstring::npos); + std::string override_val_str = free_dim_str.substr(delimiter_location + 1, std::string::npos); ORT_TRY { override_val = std::stoll(override_val_str.c_str()); if (override_val <= 0) { @@ -193,240 +188,326 @@ static bool ParseDimensionOverride(std::basic_string& dim_identifier, return true; } -/*static*/ bool CommandLineParser::ParseArguments(PerformanceTestConfig& test_config, int argc, ORTCHAR_T* argv[]) { - int ch; - while ((ch = getopt(argc, argv, ORT_TSTR("m:e:r:t:p:x:y:c:d:o:u:i:f:F:S:T:C:AMPIDZvhsqznlgR:X"))) != -1) { - switch (ch) { - case 'f': { - std::basic_string dim_name; - int64_t override_val; - if (!ParseDimensionOverride(dim_name, override_val)) { - return false; - } - test_config.run_config.free_dim_name_overrides[dim_name] = override_val; - break; +std::string CustomUsageMessage() { + std::ostringstream oss; + oss << "onnxruntime_perf_test [options...] model_path [result_file]\n\n"; + oss << "Note: Options may be specified with either a single dash(-option) or a double dash(--option). Both forms are accepted and treated identically.\n\n"; + oss << "Options:"; + + return oss.str(); +} + +bool CommandLineParser::ParseArguments(PerformanceTestConfig& test_config, int argc, ORTCHAR_T* argv[]) { + // Following callback is to make sure all the ABSL flags defined above will be showed up when running with "--help". + // Note: By default abseil only wants flags in binary's main. It expects the main routine to reside in .cc or -main.cc or + // _main.cc, where the is the name of the binary (without .exe on Windows). See usage_config.cc in abseil for more details. + absl::FlagsUsageConfig config; + config.contains_help_flags = [](absl::string_view filename) { + return std::filesystem::path(filename).filename() == std::filesystem::path(__FILE__).filename(); + }; + + config.normalize_filename = [](absl::string_view f) { + return std::string(f); + }; + absl::SetFlagsUsageConfig(config); + absl::SetProgramUsageMessage(CustomUsageMessage()); + + auto utf8_strings = utils::ConvertArgvToUtf8Strings(argc, argv); + auto utf8_argv = utils::CStringsFromStrings(utf8_strings); + auto positional = absl::ParseCommandLine(static_cast(utf8_argv.size()), utf8_argv.data()); + + // -f + { + const auto& dim_override_str = absl::GetFlag(FLAGS_f); + if (!dim_override_str.empty()) { + std::string dim_name; + int64_t override_val; + if (!ParseDimensionOverride(dim_name, override_val, dim_override_str.c_str())) { + return false; } - case 'F': { - std::basic_string dim_denotation; - int64_t override_val; - if (!ParseDimensionOverride(dim_denotation, override_val)) { - return false; - } - test_config.run_config.free_dim_denotation_overrides[dim_denotation] = override_val; - break; + test_config.run_config.free_dim_name_overrides[dim_name] = override_val; + } + } + + // -F + { + const auto& dim_override_str = absl::GetFlag(FLAGS_F); + if (!dim_override_str.empty()) { + std::string dim_denotation; + int64_t override_val; + if (!ParseDimensionOverride(dim_denotation, override_val, dim_override_str.c_str())) { + return false; } - case 'm': - if (!CompareCString(optarg, ORT_TSTR("duration"))) { - test_config.run_config.test_mode = TestMode::kFixDurationMode; - } else if (!CompareCString(optarg, ORT_TSTR("times"))) { - test_config.run_config.test_mode = TestMode::KFixRepeatedTimesMode; - } else { - return false; - } - break; - case 'p': - test_config.run_config.profile_file = optarg; - break; - case 'M': - test_config.run_config.enable_memory_pattern = false; - break; - case 'A': - test_config.run_config.enable_cpu_mem_arena = false; - break; - case 'e': - if (!CompareCString(optarg, ORT_TSTR("cpu"))) { - test_config.machine_config.provider_type_name = onnxruntime::kCpuExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("cuda"))) { - test_config.machine_config.provider_type_name = onnxruntime::kCudaExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("dnnl"))) { - test_config.machine_config.provider_type_name = onnxruntime::kDnnlExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("openvino"))) { - test_config.machine_config.provider_type_name = onnxruntime::kOpenVINOExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("tensorrt"))) { - test_config.machine_config.provider_type_name = onnxruntime::kTensorrtExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("qnn"))) { - test_config.machine_config.provider_type_name = onnxruntime::kQnnExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("snpe"))) { - test_config.machine_config.provider_type_name = onnxruntime::kSnpeExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("nnapi"))) { - test_config.machine_config.provider_type_name = onnxruntime::kNnapiExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("vsinpu"))) { - test_config.machine_config.provider_type_name = onnxruntime::kVSINPUExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("coreml"))) { - test_config.machine_config.provider_type_name = onnxruntime::kCoreMLExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("dml"))) { - test_config.machine_config.provider_type_name = onnxruntime::kDmlExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("acl"))) { - test_config.machine_config.provider_type_name = onnxruntime::kAclExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("armnn"))) { - test_config.machine_config.provider_type_name = onnxruntime::kArmNNExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("rocm"))) { - test_config.machine_config.provider_type_name = onnxruntime::kRocmExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("migraphx"))) { - test_config.machine_config.provider_type_name = onnxruntime::kMIGraphXExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("xnnpack"))) { - test_config.machine_config.provider_type_name = onnxruntime::kXnnpackExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("vitisai"))) { - test_config.machine_config.provider_type_name = onnxruntime::kVitisAIExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("webgpu"))) { - test_config.machine_config.provider_type_name = onnxruntime::kWebGpuExecutionProvider; - } else if (!CompareCString(optarg, ORT_TSTR("nvtensorrtrtx"))) { - test_config.machine_config.provider_type_name = onnxruntime::kNvTensorRTRTXExecutionProvider; - } else { - return false; - } - break; - case 'r': - test_config.run_config.repeated_times = static_cast(OrtStrtol(optarg, nullptr)); - if (test_config.run_config.repeated_times <= 0) { - return false; - } - test_config.run_config.test_mode = TestMode::KFixRepeatedTimesMode; - break; - case 't': - test_config.run_config.duration_in_seconds = static_cast(OrtStrtol(optarg, nullptr)); - if (test_config.run_config.repeated_times <= 0) { - return false; - } + test_config.run_config.free_dim_denotation_overrides[dim_denotation] = override_val; + } + } + + // -m + { + const auto& test_mode_str = absl::GetFlag(FLAGS_m); + if (!test_mode_str.empty()) { + if (test_mode_str == "duration") { test_config.run_config.test_mode = TestMode::kFixDurationMode; - break; - case 's': - test_config.run_config.f_dump_statistics = true; - break; - case 'S': - test_config.run_config.random_seed_for_input_data = static_cast( - OrtStrtol(optarg, nullptr)); - break; - case 'v': - test_config.run_config.f_verbose = true; - break; - case 'x': - test_config.run_config.intra_op_num_threads = static_cast(OrtStrtol(optarg, nullptr)); - if (test_config.run_config.intra_op_num_threads < 0) { - return false; - } - break; - case 'y': - test_config.run_config.inter_op_num_threads = static_cast(OrtStrtol(optarg, nullptr)); - if (test_config.run_config.inter_op_num_threads < 0) { - return false; - } - break; - case 'P': - test_config.run_config.execution_mode = ExecutionMode::ORT_PARALLEL; - break; - case 'c': - test_config.run_config.concurrent_session_runs = - static_cast(OrtStrtol(optarg, nullptr)); - if (test_config.run_config.concurrent_session_runs <= 0) { - return false; - } - break; - case 'o': { - int tmp = static_cast(OrtStrtol(optarg, nullptr)); - switch (tmp) { - case ORT_DISABLE_ALL: - test_config.run_config.optimization_level = ORT_DISABLE_ALL; - break; - case ORT_ENABLE_BASIC: - test_config.run_config.optimization_level = ORT_ENABLE_BASIC; - break; - case ORT_ENABLE_EXTENDED: - test_config.run_config.optimization_level = ORT_ENABLE_EXTENDED; - break; - case ORT_ENABLE_LAYOUT: - test_config.run_config.optimization_level = ORT_ENABLE_LAYOUT; - break; - case ORT_ENABLE_ALL: + } else if (test_mode_str == "times") { + test_config.run_config.test_mode = TestMode::KFixRepeatedTimesMode; + } else { + return false; + } + } + } + + // -p + { + const auto& profile_file = absl::GetFlag(FLAGS_p); + if (!profile_file.empty()) test_config.run_config.profile_file = ToPathString(profile_file); + } + + // -M + test_config.run_config.enable_memory_pattern = absl::GetFlag(FLAGS_M); + + // -A + test_config.run_config.enable_cpu_mem_arena = absl::GetFlag(FLAGS_A); + + // -e + { + auto const& ep = absl::GetFlag(FLAGS_e); + if (!ep.empty()) { + if (ep == "cpu") { + test_config.machine_config.provider_type_name = onnxruntime::kCpuExecutionProvider; + } else if (ep == "cuda") { + test_config.machine_config.provider_type_name = onnxruntime::kCudaExecutionProvider; + } else if (ep == "dnnl") { + test_config.machine_config.provider_type_name = onnxruntime::kDnnlExecutionProvider; + } else if (ep == "openvino") { + test_config.machine_config.provider_type_name = onnxruntime::kOpenVINOExecutionProvider; + } else if (ep == "tensorrt") { + test_config.machine_config.provider_type_name = onnxruntime::kTensorrtExecutionProvider; + } else if (ep == "qnn") { + test_config.machine_config.provider_type_name = onnxruntime::kQnnExecutionProvider; + } else if (ep == "snpe") { + test_config.machine_config.provider_type_name = onnxruntime::kSnpeExecutionProvider; + } else if (ep == "nnapi") { + test_config.machine_config.provider_type_name = onnxruntime::kNnapiExecutionProvider; + } else if (ep == "vsinpu") { + test_config.machine_config.provider_type_name = onnxruntime::kVSINPUExecutionProvider; + } else if (ep == "coreml") { + test_config.machine_config.provider_type_name = onnxruntime::kCoreMLExecutionProvider; + } else if (ep == "dml") { + test_config.machine_config.provider_type_name = onnxruntime::kDmlExecutionProvider; + } else if (ep == "acl") { + test_config.machine_config.provider_type_name = onnxruntime::kAclExecutionProvider; + } else if (ep == "armnn") { + test_config.machine_config.provider_type_name = onnxruntime::kArmNNExecutionProvider; + } else if (ep == "rocm") { + test_config.machine_config.provider_type_name = onnxruntime::kRocmExecutionProvider; + } else if (ep == "migraphx") { + test_config.machine_config.provider_type_name = onnxruntime::kMIGraphXExecutionProvider; + } else if (ep == "xnnpack") { + test_config.machine_config.provider_type_name = onnxruntime::kXnnpackExecutionProvider; + } else if (ep == "vitisai") { + test_config.machine_config.provider_type_name = onnxruntime::kVitisAIExecutionProvider; + } else if (ep == "webgpu") { + test_config.machine_config.provider_type_name = onnxruntime::kWebGpuExecutionProvider; + } else if (ep == "nvtensorrtrtx") { + test_config.machine_config.provider_type_name = onnxruntime::kNvTensorRTRTXExecutionProvider; + } else { + return false; + } + } + } + + // Helper function to check if the option is explicitly specified. + // Abseil Flags does not provide this capability by default. + // It cannot distinguish between cases where: + // - The user typed `-r 1000` (explicitly passing the default value), and + // - The user omitted `-r` entirely. + // To determine this accurately, we must inspect argv directly. + auto is_option_specified = [&](std::string option) { + for (int i = 1; i < argc; ++i) { + auto utf8_arg = ToUTF8String(argv[i]); + if (utf8_arg == ("-" + option) || utf8_arg == ("--" + option)) { + return true; + } + } + return false; + }; + + // -r + if (is_option_specified("r")) { + if (absl::GetFlag(FLAGS_r) == static_cast(0)) return false; + test_config.run_config.repeated_times = absl::GetFlag(FLAGS_r); + test_config.run_config.test_mode = TestMode::KFixRepeatedTimesMode; + } + + // -t + if (is_option_specified("t")) { + if (absl::GetFlag(FLAGS_t) <= static_cast(0)) return false; + test_config.run_config.duration_in_seconds = absl::GetFlag(FLAGS_t); + test_config.run_config.test_mode = TestMode::kFixDurationMode; + } + + // -s + test_config.run_config.f_dump_statistics = absl::GetFlag(FLAGS_s); + + // -S + test_config.run_config.random_seed_for_input_data = absl::GetFlag(FLAGS_S); + + // -v + test_config.run_config.f_verbose = absl::GetFlag(FLAGS_v); + + // -x + if (absl::GetFlag(FLAGS_x) < 0) return false; + test_config.run_config.intra_op_num_threads = absl::GetFlag(FLAGS_x); + + // -y + if (absl::GetFlag(FLAGS_y) < 0) return false; + test_config.run_config.inter_op_num_threads = absl::GetFlag(FLAGS_y); + + // -P + if (absl::GetFlag(FLAGS_P)) test_config.run_config.execution_mode = ExecutionMode::ORT_PARALLEL; + + // -c + if (absl::GetFlag(FLAGS_c) <= static_cast(0)) return false; + test_config.run_config.concurrent_session_runs = absl::GetFlag(FLAGS_c); + + // -o + { + const auto optimization_level = absl::GetFlag(FLAGS_o); + if (optimization_level != test_config.run_config.optimization_level) { + switch (optimization_level) { + case ORT_DISABLE_ALL: + test_config.run_config.optimization_level = ORT_DISABLE_ALL; + break; + case ORT_ENABLE_BASIC: + test_config.run_config.optimization_level = ORT_ENABLE_BASIC; + break; + case ORT_ENABLE_EXTENDED: + test_config.run_config.optimization_level = ORT_ENABLE_EXTENDED; + break; + case ORT_ENABLE_LAYOUT: + test_config.run_config.optimization_level = ORT_ENABLE_LAYOUT; + break; + case ORT_ENABLE_ALL: + test_config.run_config.optimization_level = ORT_ENABLE_ALL; + break; + default: { + if (optimization_level > ORT_ENABLE_ALL) { // relax constraint test_config.run_config.optimization_level = ORT_ENABLE_ALL; - break; - default: { - if (tmp > ORT_ENABLE_ALL) { // relax constraint - test_config.run_config.optimization_level = ORT_ENABLE_ALL; - } else { - return false; - } + } else { + return false; } } - break; } - case 'u': - test_config.run_config.optimized_model_path = optarg; - break; - case 'I': - test_config.run_config.generate_model_input_binding = true; - break; - case 'd': - test_config.run_config.cudnn_conv_algo = static_cast(OrtStrtol(optarg, nullptr)); - break; - case 'q': - test_config.run_config.do_cuda_copy_in_separate_stream = true; - break; - case 'z': - test_config.run_config.set_denormal_as_zero = true; - break; - case 'i': - test_config.run_config.ep_runtime_config_string = optarg; - break; - case 'T': - test_config.run_config.intra_op_thread_affinities = ToUTF8String(optarg); - break; - case 'C': { - ORT_TRY { - ParseSessionConfigs(ToUTF8String(optarg), test_config.run_config.session_config_entries); - } - ORT_CATCH(const std::exception& ex) { - ORT_HANDLE_EXCEPTION([&]() { - fprintf(stderr, "Error parsing session configuration entries: %s\n", ex.what()); - }); - return false; - } - break; + } + } + + // -u + { + const auto& optimized_model_path = absl::GetFlag(FLAGS_u); + if (!optimized_model_path.empty()) test_config.run_config.optimized_model_path = ToPathString(optimized_model_path); + } + + // -I + test_config.run_config.generate_model_input_binding = absl::GetFlag(FLAGS_I); + + // -d + if (absl::GetFlag(FLAGS_d) < 0) return false; + test_config.run_config.cudnn_conv_algo = absl::GetFlag(FLAGS_d); + + // -q + test_config.run_config.do_cuda_copy_in_separate_stream = absl::GetFlag(FLAGS_q); + + // -z + test_config.run_config.set_denormal_as_zero = absl::GetFlag(FLAGS_z); + + // -i + { + const auto& ep_options = absl::GetFlag(FLAGS_i); + if (!ep_options.empty()) test_config.run_config.ep_runtime_config_string = ToPathString(ep_options); + } + + // -T + if (!absl::GetFlag(FLAGS_T).empty()) test_config.run_config.intra_op_thread_affinities = absl::GetFlag(FLAGS_T); + + // -C + { + const auto& session_configs = absl::GetFlag(FLAGS_C); + if (!session_configs.empty()) { + ORT_TRY { + ParseSessionConfigs(session_configs, test_config.run_config.session_config_entries); } - case 'D': - test_config.run_config.disable_spinning = true; - break; - case 'Z': - test_config.run_config.disable_spinning_between_run = true; - break; - case 'n': - test_config.run_config.exit_after_session_creation = true; - break; - case 'l': - test_config.model_info.load_via_path = true; - break; - case 'R': - test_config.run_config.register_custom_op_path = optarg; - break; - case 'g': - test_config.run_config.enable_cuda_io_binding = true; - break; - case 'X': - test_config.run_config.use_extensions = true; - break; - case '?': - case 'h': - default: + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + fprintf(stderr, "Error parsing session configuration entries: %s\n", ex.what()); + }); return false; + } } } - // parse model_path and result_file_path - argc -= optind; - argv += optind; - - switch (argc) { - case 2: - test_config.model_info.result_file_path = argv[1]; - break; - case 1: - test_config.run_config.f_dump_statistics = true; - break; - default: - return false; + // -D + test_config.run_config.disable_spinning = absl::GetFlag(FLAGS_D); + + // -Z + test_config.run_config.disable_spinning_between_run = absl::GetFlag(FLAGS_Z); + + // -n + test_config.run_config.exit_after_session_creation = absl::GetFlag(FLAGS_n); + + // -l + test_config.model_info.load_via_path = absl::GetFlag(FLAGS_l); + + // -R + { + const auto& register_custom_op_path = absl::GetFlag(FLAGS_R); + if (!register_custom_op_path.empty()) test_config.run_config.register_custom_op_path = ToPathString(register_custom_op_path); } - test_config.model_info.model_file_path = argv[0]; + // -g + test_config.run_config.enable_cuda_io_binding = absl::GetFlag(FLAGS_g); + + // -X + test_config.run_config.use_extensions = absl::GetFlag(FLAGS_X); + + // --plugin_ep_libs + { + const auto& plugin_ep_names_and_libs = absl::GetFlag(FLAGS_plugin_ep_libs); + if (!plugin_ep_names_and_libs.empty()) test_config.plugin_ep_names_and_libs = ToPathString(plugin_ep_names_and_libs); + } + + // --plugin_eps + { + const auto& plugin_eps = absl::GetFlag(FLAGS_plugin_eps); + if (!plugin_eps.empty()) ParseEpList(plugin_eps, test_config.machine_config.plugin_provider_type_list); + } + + // --plugin_ep_options + { + const auto& plugin_ep_options = absl::GetFlag(FLAGS_plugin_ep_options); + if (!plugin_ep_options.empty()) test_config.run_config.ep_runtime_config_string = ToPathString(plugin_ep_options); + } + + // --list_ep_devices + if (absl::GetFlag(FLAGS_list_ep_devices)) { + test_config.list_available_ep_devices = true; + return true; + } + + // --select_ep_devices + { + const auto& select_ep_devices = absl::GetFlag(FLAGS_select_ep_devices); + if (!select_ep_devices.empty()) test_config.selected_ep_device_indices = select_ep_devices; + } + + if (positional.size() == 2) { + test_config.model_info.model_file_path = ToPathString(positional[1]); + test_config.run_config.f_dump_statistics = true; + } else if (positional.size() == 3) { + test_config.model_info.model_file_path = ToPathString(positional[1]); + test_config.model_info.result_file_path = ToPathString(positional[2]); + } else { + return false; + } return true; } diff --git a/onnxruntime/test/perftest/command_args_parser.h b/onnxruntime/test/perftest/command_args_parser.h index 86c81072233c0..5a94f99874797 100644 --- a/onnxruntime/test/perftest/command_args_parser.h +++ b/onnxruntime/test/perftest/command_args_parser.h @@ -11,7 +11,6 @@ struct PerformanceTestConfig; class CommandLineParser { public: - static void ShowUsage(); static bool ParseArguments(PerformanceTestConfig& test_config, int argc, ORTCHAR_T* argv[]); }; diff --git a/onnxruntime/test/perftest/common_utils.cc b/onnxruntime/test/perftest/common_utils.cc new file mode 100644 index 0000000000000..5cc6c240e25f0 --- /dev/null +++ b/onnxruntime/test/perftest/common_utils.cc @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test/perftest/utils.h" +#include "test/perftest/strings_helper.h" +#include + +#include + +#include + +namespace onnxruntime { +namespace perftest { +namespace utils { + +void ListEpDevices(const Ort::Env& env) { + std::vector ep_devices = env.GetEpDevices(); + + for (size_t i = 0; i < ep_devices.size(); ++i) { + auto device = ep_devices[i]; + std::string device_info_msg = "===== device id " + std::to_string(i) + " ======\n"; + device_info_msg += "name: " + std::string(device.EpName()) + "\n"; + device_info_msg += "vendor: " + std::string(device.EpVendor()) + "\n"; + + auto metadata = device.EpMetadata(); + std::unordered_map metadata_entries = metadata.GetKeyValuePairs(); + if (!metadata_entries.empty()) { + device_info_msg += "metadata:\n"; + } + + for (auto& entry : metadata_entries) { + device_info_msg += " " + entry.first + ": " + entry.second + "\n"; + } + device_info_msg += "\n"; + fprintf(stdout, "%s", device_info_msg.c_str()); + } +} + +void RegisterExecutionProviderLibrary(Ort::Env& env, PerformanceTestConfig& test_config) { + if (!test_config.plugin_ep_names_and_libs.empty()) { + std::unordered_map ep_names_to_libs; + ParseSessionConfigs(ToUTF8String(test_config.plugin_ep_names_and_libs), ep_names_to_libs); + if (ep_names_to_libs.size() > 0) { + for (auto& pair : ep_names_to_libs) { + const std::filesystem::path library_path = pair.second; + const std::string registration_name = pair.first; + Ort::Status status(Ort::GetApi().RegisterExecutionProviderLibrary(env, registration_name.c_str(), ToPathString(library_path.string()).c_str())); + if (status.IsOK()) { + test_config.registered_plugin_eps.push_back(registration_name); + } else { + fprintf(stderr, "Can't register %s plugin library: %s\n", registration_name.c_str(), status.GetErrorMessage().c_str()); + } + } + } + } +} + +void UnregisterExecutionProviderLibrary(Ort::Env& env, PerformanceTestConfig& test_config) { + for (auto& registration_name : test_config.registered_plugin_eps) { + Ort::Status status(Ort::GetApi().UnregisterExecutionProviderLibrary(env, registration_name.c_str())); + if (!status.IsOK()) { + fprintf(stderr, "%s", status.GetErrorMessage().c_str()); + } + } +} + +std::vector ConvertArgvToUtf8Strings(int argc, ORTCHAR_T* argv[]) { + std::vector utf8_args; + utf8_args.reserve(argc); + for (int i = 0; i < argc; ++i) { + std::string utf8_string = ToUTF8String(argv[i]); + + // Abseil flags doens't natively alias "-h" to "--help". + // We make "-h" alias to "--help" here. + if (utf8_string == "-h" || utf8_string == "--h") { + utf8_args.push_back("--help"); + } else { + utf8_args.push_back(utf8_string); + } + } + return utf8_args; +} + +std::vector CStringsFromStrings(std::vector& utf8_args) { + std::vector utf8_argv; + utf8_argv.reserve(utf8_args.size()); + for (auto& str : utf8_args) { + utf8_argv.push_back(&str[0]); + } + return utf8_argv; +} + +} // namespace utils +} // namespace perftest +} // namespace onnxruntime diff --git a/onnxruntime/test/perftest/main.cc b/onnxruntime/test/perftest/main.cc index 43bf54963cabb..973baf774b024 100644 --- a/onnxruntime/test/perftest/main.cc +++ b/onnxruntime/test/perftest/main.cc @@ -6,6 +6,8 @@ #include #include "command_args_parser.h" #include "performance_runner.h" +#include "utils.h" +#include "strings_helper.h" #include using namespace onnxruntime; @@ -19,7 +21,7 @@ int real_main(int argc, char* argv[]) { g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION); perftest::PerformanceTestConfig test_config; if (!perftest::CommandLineParser::ParseArguments(test_config, argc, argv)) { - perftest::CommandLineParser::ShowUsage(); + fprintf(stderr, "%s", "See 'onnxruntime_perf_test --help'."); return -1; } Ort::Env env{nullptr}; @@ -41,6 +43,30 @@ int real_main(int argc, char* argv[]) { if (failed) return -1; } + + if (!test_config.plugin_ep_names_and_libs.empty()) { + perftest::utils::RegisterExecutionProviderLibrary(env, test_config); + } + + // Unregister all registered plugin EP libraries before program exits. + // This is necessary because unregistering the plugin EP also unregisters any associated shared allocators. + // If we don't do this and program returns, the factories stored inside the environment will be destroyed when the environment goes out of scope. + // Later, when the shared allocator's deleter runs, it may cause a segmentation fault because it attempts to use the already-destroyed factory to call ReleaseAllocator. + // See "ep_device.ep_factory->ReleaseAllocator" in Environment::CreateSharedAllocatorImpl. + auto unregister_plugin_eps_at_scope_exit = gsl::finally([&]() { + if (!test_config.registered_plugin_eps.empty()) { + perftest::utils::UnregisterExecutionProviderLibrary(env, test_config); // this won't throw + } + }); + + if (test_config.list_available_ep_devices) { + perftest::utils::ListEpDevices(env); + if (test_config.registered_plugin_eps.empty()) { + fprintf(stdout, "No plugin execution provider libraries are registered. Please specify them using \"--plugin_ep_libs\"; otherwise, only CPU may be available.\n"); + } + return 0; + } + std::random_device rd; perftest::PerformanceRunner perf_runner(env, test_config, rd); diff --git a/onnxruntime/test/perftest/ort_test_session.cc b/onnxruntime/test/perftest/ort_test_session.cc index 7a210ca8482a4..7156a1eb5c347 100644 --- a/onnxruntime/test/perftest/ort_test_session.cc +++ b/onnxruntime/test/perftest/ort_test_session.cc @@ -62,6 +62,84 @@ OnnxRuntimeTestSession::OnnxRuntimeTestSession(Ort::Env& env, std::random_device : rand_engine_(rd()), input_names_(m.GetInputCount()), input_names_str_(m.GetInputCount()), input_length_(m.GetInputCount()) { Ort::SessionOptions session_options; + // Add EP devices if any (created by plugin EP) + if (!performance_test_config.registered_plugin_eps.empty()) { + std::vector ep_devices = env.GetEpDevices(); + // EP -> associated EP devices (All OrtEpDevice instances must be from the same execution provider) + std::unordered_map> added_ep_devices; + std::unordered_set added_ep_device_index_set; + + auto& ep_list = performance_test_config.machine_config.plugin_provider_type_list; + std::unordered_set ep_set(ep_list.begin(), ep_list.end()); + + // Select EP devices by provided device index + if (!performance_test_config.selected_ep_device_indices.empty()) { + std::vector device_list; + device_list.reserve(performance_test_config.selected_ep_device_indices.size()); + ParseEpDeviceIndexList(performance_test_config.selected_ep_device_indices, device_list); + for (auto index : device_list) { + if (static_cast(index) > (ep_devices.size() - 1)) { + fprintf(stderr, "%s", "The device index provided is not correct. Will skip this device id."); + continue; + } + + Ort::ConstEpDevice& device = ep_devices[index]; + if (ep_set.find(std::string(device.EpName())) != ep_set.end()) { + if (added_ep_device_index_set.find(index) == added_ep_device_index_set.end()) { + added_ep_devices[device.EpName()].push_back(device); + added_ep_device_index_set.insert(index); + fprintf(stdout, "[Plugin EP] EP Device [Index: %d, Name: %s] has been added to session.\n", index, device.EpName()); + } + } else { + std::string err_msg = "[Plugin EP] [WARNING] : The EP device index and its corresponding OrtEpDevice is not created from " + + performance_test_config.machine_config.provider_type_name + ". Will skip adding this device.\n"; + fprintf(stderr, "%s", err_msg.c_str()); + } + } + } else { + // Find and select the OrtEpDevice associated with the EP in "--plugin_eps". + for (size_t index = 0; index < ep_devices.size(); ++index) { + Ort::ConstEpDevice& device = ep_devices[index]; + if (ep_set.find(std::string(device.EpName())) != ep_set.end()) { + added_ep_devices[device.EpName()].push_back(device); + fprintf(stdout, "EP Device [Index: %d, Name: %s] has been added to session.\n", static_cast(index), device.EpName()); + } + } + } + + if (added_ep_devices.empty()) { + ORT_THROW("[ERROR] [Plugin EP]: No matching EP devices found."); + } + + std::string ep_option_string = ToUTF8String(performance_test_config.run_config.ep_runtime_config_string); + + // EP's associated provider option lists + std::vector> ep_options_list; + ParseEpOptions(ep_option_string, ep_options_list); + + // If user only provide the EPs' provider option lists for the first several EPs, + // add empty provider option lists for the rest EPs. + if (ep_options_list.size() < ep_list.size()) { + for (size_t i = ep_options_list.size(); i < ep_list.size(); ++i) { + ep_options_list.emplace_back(); // Adds a new empty map + } + } else if (ep_options_list.size() > ep_list.size()) { + ORT_THROW("[ERROR] [Plugin EP]: Too many EP provider option lists provided."); + } + + // EP -> associated provider options + std::unordered_map> ep_options_map; + for (size_t i = 0; i < ep_list.size(); ++i) { + ep_options_map.emplace(ep_list[i], ep_options_list[i]); + } + + for (auto& ep_and_devices : added_ep_devices) { + auto& ep = ep_and_devices.first; + auto& devices = ep_and_devices.second; + session_options.AppendExecutionProvider_V2(env, devices, ep_options_map[ep]); + } + } + provider_name_ = performance_test_config.machine_config.provider_type_name; std::unordered_map provider_options; if (provider_name_ == onnxruntime::kDnnlExecutionProvider) { diff --git a/onnxruntime/test/perftest/strings_helper.cc b/onnxruntime/test/perftest/strings_helper.cc index 9fd49da1d0486..f4860b35c79da 100644 --- a/onnxruntime/test/perftest/strings_helper.cc +++ b/onnxruntime/test/perftest/strings_helper.cc @@ -8,6 +8,8 @@ #include "strings_helper.h" #include "core/common/common.h" +#include "core/common/parse_string.h" +#include "core/common/string_utils.h" namespace onnxruntime { namespace perftest { @@ -53,5 +55,40 @@ void ParseSessionConfigs(const std::string& configs_string, session_configs.insert(std::make_pair(std::move(key), std::move(value))); } } + +void ParseEpOptions(const std::string& input, std::vector>& result) { + auto tokens = utils::SplitString(input, ";", true); + + for (const auto& token : tokens) { + result.emplace_back(); // Adds a new empty map + if (!token.empty()) { + ParseSessionConfigs(std::string(token), result.back()); // only parse non-empty + } + // if token is empty, we still get an empty map in `result` + } +} + +void ParseEpList(const std::string& input, std::vector& result) { + std::stringstream ss(input); + std::string token; + + while (std::getline(ss, token, ';')) { + if (!token.empty()) { + result.push_back(token); + } + } +} + +void ParseEpDeviceIndexList(const std::string& input, std::vector& result) { + std::stringstream ss(input); + std::string item; + + while (std::getline(ss, item, ';')) { + if (!item.empty()) { + int value = ParseStringWithClassicLocale(item); + result.push_back(value); + } + } +} } // namespace perftest } // namespace onnxruntime diff --git a/onnxruntime/test/perftest/strings_helper.h b/onnxruntime/test/perftest/strings_helper.h index 0d6c56709fde6..621ab746273bd 100644 --- a/onnxruntime/test/perftest/strings_helper.h +++ b/onnxruntime/test/perftest/strings_helper.h @@ -5,6 +5,7 @@ #include #include #include +#include namespace onnxruntime { namespace perftest { @@ -12,5 +13,11 @@ namespace perftest { void ParseSessionConfigs(const std::string& configs_string, std::unordered_map& session_configs, const std::unordered_set& available_keys = {}); + +void ParseEpList(const std::string& input, std::vector& result); + +void ParseEpOptions(const std::string& input, std::vector>& result); + +void ParseEpDeviceIndexList(const std::string& input, std::vector& result); } // namespace perftest } // namespace onnxruntime diff --git a/onnxruntime/test/perftest/test_configuration.h b/onnxruntime/test/perftest/test_configuration.h index 8145f5f35c3b3..29ee84dd40dac 100644 --- a/onnxruntime/test/perftest/test_configuration.h +++ b/onnxruntime/test/perftest/test_configuration.h @@ -35,6 +35,7 @@ struct ModelInfo { struct MachineConfig { Platform platform{Platform::kWindows}; std::string provider_type_name{onnxruntime::kCpuExecutionProvider}; + std::vector plugin_provider_type_list; }; struct RunConfig { @@ -59,8 +60,8 @@ struct RunConfig { bool set_denormal_as_zero{false}; std::basic_string ep_runtime_config_string; std::unordered_map session_config_entries; - std::map, int64_t> free_dim_name_overrides; - std::map, int64_t> free_dim_denotation_overrides; + std::map free_dim_name_overrides; + std::map free_dim_denotation_overrides; std::string intra_op_thread_affinities; bool disable_spinning = false; bool disable_spinning_between_run = false; @@ -74,6 +75,10 @@ struct PerformanceTestConfig { ModelInfo model_info; MachineConfig machine_config; RunConfig run_config; + std::basic_string plugin_ep_names_and_libs; + std::vector registered_plugin_eps; + std::string selected_ep_device_indices; + bool list_available_ep_devices = false; }; } // namespace perftest diff --git a/onnxruntime/test/perftest/utils.h b/onnxruntime/test/perftest/utils.h index f22abc04fa99e..9f180e2c8d942 100644 --- a/onnxruntime/test/perftest/utils.h +++ b/onnxruntime/test/perftest/utils.h @@ -2,7 +2,8 @@ // Licensed under the MIT License. #pragma once - +#include "test/perftest/test_configuration.h" +#include #include namespace onnxruntime { @@ -22,6 +23,16 @@ class ICPUUsage { std::unique_ptr CreateICPUUsage(); +std::vector ConvertArgvToUtf8Strings(int argc, ORTCHAR_T* argv[]); + +std::vector CStringsFromStrings(std::vector& utf8_args); + +void RegisterExecutionProviderLibrary(Ort::Env& env, PerformanceTestConfig& test_config); + +void UnregisterExecutionProviderLibrary(Ort::Env& env, PerformanceTestConfig& test_config); + +void ListEpDevices(const Ort::Env& env); + } // namespace utils } // namespace perftest } // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc index 0559699670c4a..19505da1bbe56 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc @@ -394,6 +394,7 @@ TYPED_TEST(NvExecutionProviderTest, IOTypeTests) { } } +#if defined(WIN32) static bool SessionHasEp(Ort::Session& session, const char* ep_name) { // Access the underlying InferenceSession. const OrtSession* ort_session = session; @@ -409,11 +410,10 @@ static bool SessionHasEp(Ort::Session& session, const char* ep_name) { return has_ep; } -#if defined(WIN32) // Tests autoEP feature to automatically select an EP that supports the GPU. // Currently only works on Windows. TEST(NvExecutionProviderTest, AutoEp_PreferGpu) { - PathString model_name = ORT_TSTR("nv_execution_provider_data_dyn_test.onnx"); + PathString model_name = ORT_TSTR("nv_execution_provider_auto_ep.onnx"); std::string graph_name = "test"; std::vector dims = {1, 3, 2}; diff --git a/onnxruntime/test/python/onnxruntime_test_python_autoep.py b/onnxruntime/test/python/onnxruntime_test_python_autoep.py index 0c52740398b7a..cb31627a87c48 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_autoep.py +++ b/onnxruntime/test/python/onnxruntime_test_python_autoep.py @@ -183,7 +183,7 @@ def test_example_plugin_ep_devices(self): Test registration of an example EP plugin and retrieval of its OrtEpDevice. """ if sys.platform != "win32": - self.skipTest("Skipping test because it device discovery is only supported on Windows") + self.skipTest("Skipping test because device discovery is only supported on Windows") ep_lib_path = "example_plugin_ep.dll" try: @@ -244,6 +244,44 @@ def test_example_plugin_ep_devices(self): del sess # Delete session before unregistering library self.unregister_execution_provider_library(ep_name) + def test_example_plugin_ep_data_transfer(self): + """ + Test usage of shared data transfer and allocator from plugin EP. + """ + if sys.platform != "win32": + self.skipTest("Skipping test because device discovery is only supported on Windows") + + if "DmlExecutionProvider" in onnxrt.get_available_providers(): + self.skipTest("Skipping because DML EP data transfer is broken if we haven't created an inference session") + + ep_lib_path = "example_plugin_ep.dll" + try: + ep_lib_path = get_name("example_plugin_ep.dll") + except FileNotFoundError: + self.skipTest(f"Skipping test because EP library '{ep_lib_path}' cannot be found") + + ep_name = "example_ep" + self.register_execution_provider_library(ep_name, os.path.realpath(ep_lib_path)) + + data = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) + data2 = data + 1 + + # the example EP pretends to use GPU memory so we can test data transfer. + # by matching its OrtDevice info we will hit its allocator and data transfer implementations. + # copy data from CPU to the fake GPU memory + gpu_value = onnxrt.OrtValue.ortvalue_from_numpy(data, "gpu", 0, 0xBE57) + # copy back to CPU + cpu_data = gpu_value.numpy() + np.testing.assert_equal(data, cpu_data) + + gpu_value.update_inplace(data2) # update the fake GPU data + cpu_data_2 = gpu_value.numpy() # copy back to CPU + np.testing.assert_equal(data2, cpu_data_2) + + gpu_value = None # Delete OrtValue before unregistering library as the allocator will be destroyed. + + self.unregister_execution_provider_library(ep_name) + if __name__ == "__main__": unittest.main(verbosity=1) diff --git a/orttraining/orttraining/core/optimizer/conv1d_replacement.cc b/orttraining/orttraining/core/optimizer/conv1d_replacement.cc index 90be9e24d3dd4..ff220fcb067b8 100644 --- a/orttraining/orttraining/core/optimizer/conv1d_replacement.cc +++ b/orttraining/orttraining/core/optimizer/conv1d_replacement.cc @@ -121,7 +121,7 @@ void Conv1dToMatmul(Graph& graph, Node& conv, const std::string transformer_name initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); InlinedVector initializer_proto_value{weight_squeeze_axis}; initializer_proto.set_raw_data(initializer_proto_value.data(), initializer_proto_value.size() * sizeof(int64_t)); - auto& axes_input = graph_utils::AddInitializerWithExternalData(graph, initializer_proto); + auto& axes_input = graph_utils::AddInitializer(graph, initializer_proto); // Squeeze node doesn't have opschema here, so we need to set input args count manually weight_squeeze.MutableInputArgsCount().resize(2); graph_utils::AddNodeInput(weight_squeeze, 1, axes_input); diff --git a/orttraining/orttraining/core/optimizer/megatron_transformer.cc b/orttraining/orttraining/core/optimizer/megatron_transformer.cc index 55286379fd273..7c429ae5cb643 100644 --- a/orttraining/orttraining/core/optimizer/megatron_transformer.cc +++ b/orttraining/orttraining/core/optimizer/megatron_transformer.cc @@ -453,15 +453,15 @@ Status MegatronTransformer::TransformGPT2MLP(Graph& graph, bool& modified, return skip_status; } - NodeArg& a_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, a_weight_initializer_partition); + NodeArg& a_weight_partition_arg = graph_utils::AddInitializer(graph, a_weight_initializer_partition); graph_utils::ReplaceNodeInput(node, 1, a_weight_partition_arg); updated_weight_names_.insert({a_weight_arg->Name(), a_weight_partition_arg.Name()}); - NodeArg& a_bias_partition_arg = graph_utils::AddInitializerWithExternalData(graph, a_bias_initializer_partition); + NodeArg& a_bias_partition_arg = graph_utils::AddInitializer(graph, a_bias_initializer_partition); graph_utils::ReplaceNodeInput(add_node, 1, a_bias_partition_arg); updated_weight_names_.insert({b_weight_arg->Name(), a_bias_partition_arg.Name()}); - NodeArg& b_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, b_weight_initializer_partition); + NodeArg& b_weight_partition_arg = graph_utils::AddInitializer(graph, b_weight_initializer_partition); graph_utils::ReplaceNodeInput(matmul2_node, 1, b_weight_partition_arg); updated_weight_names_.insert({a_bias_arg->Name(), b_weight_partition_arg.Name()}); @@ -600,15 +600,15 @@ Status MegatronTransformer::TransformBARTMLP(Graph& graph, bool& modified, return skip_status; } - NodeArg& dense_wi_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, dense_wi_weight_initializer_partition); + NodeArg& dense_wi_weight_partition_arg = graph_utils::AddInitializer(graph, dense_wi_weight_initializer_partition); graph_utils::ReplaceNodeInput(*second_op, 0, dense_wi_weight_partition_arg); updated_weight_names_.insert({dense_wi_weight_arg->Name(), dense_wi_weight_partition_arg.Name()}); - NodeArg& dense_wi_bias_partition_arg = graph_utils::AddInitializerWithExternalData(graph, dense_wi_bias_initializer_partition); + NodeArg& dense_wi_bias_partition_arg = graph_utils::AddInitializer(graph, dense_wi_bias_initializer_partition); graph_utils::ReplaceNodeInput(biasgelu_node, 1, dense_wi_bias_partition_arg); updated_weight_names_.insert({dense_wi_bias_arg->Name(), dense_wi_bias_partition_arg.Name()}); - NodeArg& dense_wo_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, dense_wo_weight_initializer_partition); + NodeArg& dense_wo_weight_partition_arg = graph_utils::AddInitializer(graph, dense_wo_weight_initializer_partition); graph_utils::ReplaceNodeInput(*transpose_op_ptr, 0, dense_wo_weight_partition_arg); updated_weight_names_.insert({dense_wo_weight_arg->Name(), dense_wo_weight_partition_arg.Name()}); @@ -814,15 +814,15 @@ Status MegatronTransformer::TransformGPT2Attention(Graph& graph, bool& modified, [](Node* node_ptr) { return node_ptr != nullptr; }); // Replace by the partition weights. - NodeArg& qkv_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, qkv_weight_initializer_partition); + NodeArg& qkv_weight_partition_arg = graph_utils::AddInitializer(graph, qkv_weight_initializer_partition); graph_utils::ReplaceNodeInput(node, 1, qkv_weight_partition_arg); updated_weight_names_.insert({qkv_weight_arg->Name(), qkv_weight_partition_arg.Name()}); - NodeArg& qkv_bias_partition_arg = graph_utils::AddInitializerWithExternalData(graph, qkv_bias_initializer_partition); + NodeArg& qkv_bias_partition_arg = graph_utils::AddInitializer(graph, qkv_bias_initializer_partition); graph_utils::ReplaceNodeInput(add_node, 1, qkv_bias_partition_arg); updated_weight_names_.insert({qkv_bias_arg->Name(), qkv_bias_partition_arg.Name()}); - NodeArg& dense_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, dense_weight_initializer_partition); + NodeArg& dense_weight_partition_arg = graph_utils::AddInitializer(graph, dense_weight_initializer_partition); graph_utils::ReplaceNodeInput(matmul_node, 1, dense_weight_partition_arg); updated_weight_names_.insert({dense_weight_arg->Name(), dense_weight_partition_arg.Name()}); @@ -849,7 +849,7 @@ Status MegatronTransformer::TransformGPT2Attention(Graph& graph, bool& modified, val_partition.insert(val_partition.end(), val, val + size); val_partition[2] /= horizontal_parallel_size_; tensor_partition.set_raw_data(val_partition.data(), size * sizeof(int64_t)); - NodeArg& node_arg_partition = graph_utils::AddInitializerWithExternalData(graph, tensor_partition); + NodeArg& node_arg_partition = graph_utils::AddInitializer(graph, tensor_partition); graph_utils::ReplaceNodeInput(*node_ptr, 1, node_arg_partition); graph.RemoveInitializedTensor(shape_arg->Name()); } @@ -1130,7 +1130,7 @@ Status MegatronTransformer::TransformBARTAttention(Graph& graph, bool& modified, size_t i = 0; for (auto trans_ptr : weight_transpose_node_ptrs) { auto weight_name = trans_ptr->MutableInputDefs()[0]->Name(); - NodeArg& qkv_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, qkv_weight_initializer_partitions[i]); + NodeArg& qkv_weight_partition_arg = graph_utils::AddInitializer(graph, qkv_weight_initializer_partitions[i]); graph_utils::ReplaceNodeInput(*trans_ptr, 0, qkv_weight_partition_arg); graph.RemoveInitializedTensor(weight_name); updated_weight_names_.insert({weight_name, qkv_weight_partition_arg.Name()}); @@ -1139,14 +1139,14 @@ Status MegatronTransformer::TransformBARTAttention(Graph& graph, bool& modified, i = 0; for (auto add_ptr : bias_add_node_ptrs) { auto bias_name = add_ptr->MutableInputDefs()[1]->Name(); - NodeArg& qkv_bias_partition_arg = graph_utils::AddInitializerWithExternalData(graph, qkv_bias_initializer_partitions[i]); + NodeArg& qkv_bias_partition_arg = graph_utils::AddInitializer(graph, qkv_bias_initializer_partitions[i]); graph_utils::ReplaceNodeInput(*add_ptr, 1, qkv_bias_partition_arg); graph.RemoveInitializedTensor(bias_name); updated_weight_names_.insert({bias_name, qkv_bias_partition_arg.Name()}); i++; } - NodeArg& dense_weight_partition_arg = graph_utils::AddInitializerWithExternalData(graph, dense_weight_initializer_partition); + NodeArg& dense_weight_partition_arg = graph_utils::AddInitializer(graph, dense_weight_initializer_partition); graph_utils::ReplaceNodeInput(*last_transpose, 0, dense_weight_partition_arg); graph.RemoveInitializedTensor(dense_weight_arg->Name()); updated_weight_names_.insert({dense_weight_arg->Name(), dense_weight_partition_arg.Name()}); @@ -1178,7 +1178,7 @@ Status MegatronTransformer::TransformBARTAttention(Graph& graph, bool& modified, val_partition.insert(val_partition.end(), val, val + size); val_partition[idx] /= horizontal_parallel_size_; tensor_partition.set_raw_data(val_partition.data(), size * sizeof(int64_t)); - NodeArg& node_arg_partition = graph_utils::AddInitializerWithExternalData(graph, tensor_partition); + NodeArg& node_arg_partition = graph_utils::AddInitializer(graph, tensor_partition); graph_utils::ReplaceNodeInput(*node_ptr, 1, node_arg_partition); graph.RemoveInitializedTensor(shape_arg->Name()); } diff --git a/orttraining/orttraining/core/optimizer/qdq_fusion.cc b/orttraining/orttraining/core/optimizer/qdq_fusion.cc index 4a5bdc1f8fcd2..42720dbbb11e5 100644 --- a/orttraining/orttraining/core/optimizer/qdq_fusion.cc +++ b/orttraining/orttraining/core/optimizer/qdq_fusion.cc @@ -45,7 +45,7 @@ int ReplaceOrCreateZeroPointInitializer(Graph& graph, Node& quantize_node) { // Since the quantize node has the zero point initializer input, replace it graph_utils::ReplaceNodeInput(quantize_node, 2, - graph_utils::AddInitializerWithExternalData(graph, zero_point_tensor_float)); + graph_utils::AddInitializer(graph, zero_point_tensor_float)); } else { // The quantize node does not have the zero point optional input. // Create the zero point initializer to be 0. @@ -55,7 +55,7 @@ int ReplaceOrCreateZeroPointInitializer(Graph& graph, Node& quantize_node) { // Since the input did not exist, add the newly created initializer as an input graph_utils::AddNodeInput(quantize_node, 2, - graph_utils::AddInitializerWithExternalData(graph, zero_point_tensor_float)); + graph_utils::AddInitializer(graph, zero_point_tensor_float)); } return zero_point_type; diff --git a/orttraining/orttraining/core/optimizer/sce_loss_grad_bias_fusion.cc b/orttraining/orttraining/core/optimizer/sce_loss_grad_bias_fusion.cc index 8c9c12ceb4497..84bf715c7c85a 100644 --- a/orttraining/orttraining/core/optimizer/sce_loss_grad_bias_fusion.cc +++ b/orttraining/orttraining/core/optimizer/sce_loss_grad_bias_fusion.cc @@ -83,7 +83,7 @@ Status SceLossGradBiasFusion::ApplyImpl(Graph& graph, bool& modified, int graph_ ignore_index_initializer_proto.set_name(graph.GenerateNodeArgName("sce_grad_ignore_index")); ignore_index_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); ignore_index_initializer_proto.add_int64_data(static_cast(-1)); - new_scegrad_node_inputs.emplace_back(&graph_utils::AddInitializerWithExternalData(graph, ignore_index_initializer_proto)); + new_scegrad_node_inputs.emplace_back(&graph_utils::AddInitializer(graph, ignore_index_initializer_proto)); } new_scegrad_node_inputs.emplace_back(bias_def); if (!p_reshape) {