saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix transpose optimizer on GPU EP (#15988)

Browse source 
### Description
because of #15618 , the default allocator changed to device allocator,
which will be GPU instead of CPU. in transpose optimizer we expect to
read data from initializers so a CPU allocator is required here.

this change fixes transpose optimizer on GPU EP

Fixes the issue referred to in #15869, #15796
This commit is contained in:
Yulong Wang 2023-05-19 14:33:45 -07:00 committed by GitHub
parent 4324d2173b
commit dc06c255b4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 31 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
onnxruntime/core/optimizer/transpose_optimizer/optimizer_api_impl.cc
View file

@ -911,9 +911,8 @@ PostLayoutTransformCostCheck(const api::GraphRef& graph, const api::NodeRef& nod
}
Status TransformLayoutForEP(Graph& graph, bool& modified, const IExecutionProvider& execution_provider,
const DebugGraphFn& debug_graph_fn) {
// sub graph recurse will be added later
auto api_graph = MakeApiGraph(graph, execution_provider.GetAllocator(OrtMemTypeDefault), nullptr);
AllocatorPtr cpu_allocator, const DebugGraphFn& debug_graph_fn) {
auto api_graph = MakeApiGraph(graph, cpu_allocator, nullptr);
const auto& layout_sensitive_ops = GetORTLayoutSensitiveOps();
for (auto& node : api_graph->Nodes()) {

3
onnxruntime/core/optimizer/transpose_optimizer/optimizer_utils.h
View file

@ -84,11 +84,12 @@ const std::unordered_set<std::string_view>& GetORTLayoutSensitiveOps();
/// <param name="graph">graph to transform</param>
/// <param name="modified">indicates whether the graph is modified during transformation</param>
/// <param name="execution_provider">execution provider for which the transformation needs to be performed</param>
/// <param name="cpu_allocator">a CPU allocator used in layout transformation.
/// <param name="debug_graph_fn">Optional functor to debug the graph produced during layout transformation.
/// This is called after layout transformation if new nodes are inserted, and again after those are optimized.
/// </param>
Status TransformLayoutForEP(Graph& graph, bool& modified, const IExecutionProvider& execution_provider,
const DebugGraphFn& debug_graph_fn = {});
AllocatorPtr cpu_allocator, const DebugGraphFn& debug_graph_fn = {});
/// <summary>
/// Checks if the opset of the Graph is supported by the layout transformer.

18
onnxruntime/core/session/inference_session.cc
View file

@ -942,7 +942,7 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
ORT_RETURN_IF_ERROR_SESSIONID_(
layout_transformer::TransformLayoutForEP(graph_to_transform, modified, execution_provider,
debug_graph_fn));
execution_providers_.GetDefaultCpuAllocator(), debug_graph_fn));
if (modified) {
ORT_RETURN_IF_ERROR_SESSIONID_(
@ -1266,13 +1266,19 @@ Status PartitionOrtFormatModel(onnxruntime::Graph& graph,
const ExecutionProviders& providers,
KernelRegistryManager& kernel_registry_manager,
SessionState& session_state) {
layout_transformer::TransformLayoutFunction transform_layout_fn = nullptr;
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
// only provide NCWH to NHWC layout transformer if supported
layout_transformer::TransformLayoutFunction transform_layout_fn = layout_transformer::IsSupportedOpset(graph)
? layout_transformer::TransformLayoutForEP
: nullptr;
#else // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
layout_transformer::TransformLayoutFunction transform_layout_fn{};
if (layout_transformer::IsSupportedOpset(graph)) {
transform_layout_fn =
[&providers](Graph& graph_to_transform, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
return layout_transformer::TransformLayoutForEP(graph_to_transform, modified, execution_provider,
providers.GetDefaultCpuAllocator(), debug_graph_fn);
};
}
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
GraphPartitioner partitioner(kernel_registry_manager, providers);

18
onnxruntime/test/framework/session_state_test.cc
View file

@ -157,7 +157,11 @@ TEST_P(SessionStateTestP, TestInitializerProcessing) {
GraphPartitioner partitioner(krm, execution_providers);
status = partitioner.Partition(graph, session_state.GetMutableFuncMgr(),
layout_transformer::TransformLayoutForEP);
[&execution_providers](Graph& graph, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
return layout_transformer::TransformLayoutForEP(graph, modified, execution_provider, execution_providers.GetDefaultCpuAllocator(), debug_graph_fn);
});
ASSERT_TRUE(status.IsOK()) << status;
ASSERT_STATUS_OK(session_state.FinalizeSessionState(oss.str(), krm));
@ -231,7 +235,11 @@ TEST(SessionStateTest, TestInitializerMemoryAllocatedUsingNonArenaMemory) {
// Partition the graph
GraphPartitioner partitioner(krm, execution_providers);
status = partitioner.Partition(graph, session_state.GetMutableFuncMgr(),
layout_transformer::TransformLayoutForEP);
[&execution_providers](Graph& graph, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
return layout_transformer::TransformLayoutForEP(graph, modified, execution_provider, execution_providers.GetDefaultCpuAllocator(), debug_graph_fn);
});
ASSERT_TRUE(status.IsOK()) << status;
ASSERT_STATUS_OK(session_state.FinalizeSessionState(oss.str(), krm));
@ -282,7 +290,11 @@ TEST(SessionStateTest, TestInitializerMemoryAllocatedUsingNonArenaMemory) {
// Partition the graph
GraphPartitioner partitioner(krm, execution_providers);
status = partitioner.Partition(graph, session_state.GetMutableFuncMgr(),
layout_transformer::TransformLayoutForEP);
[&execution_providers](Graph& graph, bool& modified,
const IExecutionProvider& execution_provider,
const layout_transformer::DebugGraphFn& debug_graph_fn) -> Status {
return layout_transformer::TransformLayoutForEP(graph, modified, execution_provider, execution_providers.GetDefaultCpuAllocator(), debug_graph_fn);
});
ASSERT_TRUE(status.IsOK()) << status;
// Finalize the session state