fix graph partitioning for nested functions (#5755)

* fix graph partitioning for nested functions

* enable broken test for SCE
This commit is contained in:
Ashwini Khade 2020-11-11 11:38:27 -08:00 committed by GitHub
parent bc1768c7f1
commit 496fa18c96
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 16 deletions

View file

@ -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<KernelRegistry> fused_kernel_registry = std::make_shared<KernelRegistry>();
// 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);

View file

@ -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",