From 496fa18c9686dc194d98843f126bac8f64052fb9 Mon Sep 17 00:00:00 2001 From: Ashwini Khade Date: Wed, 11 Nov 2020 11:38:27 -0800 Subject: [PATCH] fix graph partitioning for nested functions (#5755) * fix graph partitioning for nested functions * enable broken test for SCE --- .../core/framework/graph_partitioner.cc | 24 +++++++++---------- .../onnx_backend_test_series_filters.jsonc | 4 ---- 2 files changed, 12 insertions(+), 16 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 478ab2a918..69c9096480 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -231,6 +231,8 @@ static Status InlineNodes(Graph& graph, bool& modified_graph) { modified_graph = true; } + ORT_RETURN_IF_ERROR(graph.Resolve()); + return Status::OK(); } @@ -251,23 +253,21 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f // It is only visible for current session. std::shared_ptr fused_kernel_registry = std::make_shared(); - // process full graph with each EP - for (const auto& ep : providers_) { - ORT_RETURN_IF_ERROR(PartitionImpl(graph, export_dll, func_mgr, kernel_registry_mgr_, - *fused_kernel_registry, *ep)); - } - bool modified_graph = false; - ORT_RETURN_IF_ERROR(InlineNodes(graph, modified_graph)); - - // Resolve and rerun graph partition - if (modified_graph) { - ORT_RETURN_IF_ERROR(graph.Resolve()); + do { + // process full graph with each EP for (const auto& ep : providers_) { ORT_RETURN_IF_ERROR(PartitionImpl(graph, export_dll, func_mgr, kernel_registry_mgr_, *fused_kernel_registry, *ep)); } - } + + modified_graph = false; + // expand any nodes that have an ONNX function definition + // but no matching ORT kernel + ORT_RETURN_IF_ERROR(InlineNodes(graph, modified_graph)); + // rerun graph partition to assign nodes added as part of + // function expansion. + } while (modified_graph); if (!fused_kernel_registry->IsEmpty()) { kernel_registry_mgr_.RegisterKernelRegistry(fused_kernel_registry); diff --git a/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc b/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc index e53a6e2c2d..5d61909136 100644 --- a/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc +++ b/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc @@ -76,10 +76,6 @@ "^test_resize_upsample_sizes_nearest_cpu", "^test_resize_upsample_sizes_nearest_floor_align_corners_cpu", "^test_resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric_cpu", - "^test_sce_NCd1_mean_weight_negative_ii_cpu", // NOT_IMPLEMENTED : Could not find an implementation for the node NegaticeLogLikelihoodLoss(13) - "^test_sce_NCd1_mean_weight_negative_ii_expanded_cpu", - "^test_sce_NCd1_mean_weight_negative_ii_log_prob_cpu", - "^test_sce_NCd1_mean_weight_negative_ii_log_prob_expanded_cpu", "^test_softmax_axis_0_cpu", // NOT_IMPLEMENTED : Could not find an implementation for the node Softmax(13) "^test_softmax_axis_0_expanded_cpu", "^test_softmax_axis_1_cpu",