Binary size reduction in KernelTypeStrResolver and GraphPartitioner (#13172)

Reduce binary size for minimal Android builds.
- reduce places where Status objects are created in KernelTypeStrResolver::LoadFromOrtFormat()
- remove some unused parameters (in a base minimal build) and code in graph_partitioner.cc
This commit is contained in:
Edward Chen 2022-09-30 13:50:39 -07:00 committed by GitHub
parent b76a65c784
commit aae35f2759
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 96 additions and 37 deletions

View file

@ -3,6 +3,7 @@
#include "core/framework/graph_partitioner.h"
#include <cassert>
#include <functional>
#include "core/framework/compute_capability.h"
@ -44,7 +45,6 @@ NonCudaOps non_cuda;
namespace onnxruntime {
namespace {
// contains some common parameters used by the partitioning helper functions
struct PartitionParams {
std::reference_wrapper<Graph> graph;
@ -103,7 +103,7 @@ static bool TryAssignSingleNode(Graph& graph,
const std::string& provider_type) {
// The provider can run a single node in the <graph> if not using meta-defs.
// A fused kernel is not supported in this case.
ORT_ENFORCE(1 == indexed_sub_graph.nodes.size());
assert(indexed_sub_graph.GetMetaDef() == nullptr && indexed_sub_graph.nodes.size() == 1);
auto* node = graph.GetNode(indexed_sub_graph.nodes[0]);
if (nullptr != node && node->GetExecutionProviderType().empty()) {
@ -115,17 +115,31 @@ static bool TryAssignSingleNode(Graph& graph,
return false;
}
static Status GetCapabilityForEP(Graph& graph, KernelRegistryManager& kernel_registry_mgr,
IExecutionProvider& current_ep, GraphPartitioner::Mode mode,
std::vector<std::unique_ptr<ComputeCapability>>& capabilities,
TransformLayoutFunction transform_layout) {
namespace {
struct GetCapabilityForEPParams {
std::reference_wrapper<Graph> graph;
std::reference_wrapper<const KernelRegistryManager> kernel_registry_mgr;
std::reference_wrapper<IExecutionProvider> current_ep;
std::reference_wrapper<std::vector<std::unique_ptr<ComputeCapability>>> capabilities;
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
GraphPartitioner::Mode mode;
TransformLayoutFunction transform_layout;
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
};
} // namespace
static Status GetCapabilityForEP(const GetCapabilityForEPParams& params) {
auto& current_ep = params.current_ep.get();
const auto& ep_type = current_ep.Type();
if (current_ep.GetPreferredLayout() == DataLayout::NHWC && !transform_layout) {
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
if (current_ep.GetPreferredLayout() == DataLayout::NHWC && !params.transform_layout) {
LOGS_DEFAULT(WARNING) << ep_type << " cannot be used with this model due to its ONNX opset not being supported by "
"the layout transformer.";
return Status::OK();
}
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
auto get_capabilities = [](const IExecutionProvider& ep,
const GraphViewer& graph_viewer,
@ -142,11 +156,15 @@ static Status GetCapabilityForEP(Graph& graph, KernelRegistryManager& kernel_reg
return capabilities;
};
const auto& kernel_registry_mgr = params.kernel_registry_mgr.get();
const auto kernel_registries_for_ep = kernel_registry_mgr.GetKernelRegistriesByProviderType(ep_type);
const KernelLookup kernel_lookup{ep_type,
kernel_registries_for_ep,
kernel_registry_mgr.GetKernelTypeStrResolver()};
auto& graph = params.graph.get();
auto& capabilities = params.capabilities.get();
{
const GraphViewer graph_viewer(graph);
capabilities = get_capabilities(current_ep, graph_viewer, kernel_lookup);
@ -159,7 +177,7 @@ static Status GetCapabilityForEP(Graph& graph, KernelRegistryManager& kernel_reg
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
// Run layout transformer only for EPs other than CPU EP and provided the preferred layout is NHWC
// CPU EP layout transformation happens later when level 3 transformers are run.
if (mode != GraphPartitioner::Mode::kAssignOnly &&
if (params.mode != GraphPartitioner::Mode::kAssignOnly &&
current_ep.GetPreferredLayout() == DataLayout::NHWC) {
for (auto& capability : capabilities) {
TryAssignNodes(graph, *capability->sub_graph, ep_type);
@ -169,7 +187,7 @@ static Status GetCapabilityForEP(Graph& graph, KernelRegistryManager& kernel_reg
// Perform layout transformation on the specific EP assigned graph
bool modified = false;
ORT_RETURN_IF_ERROR(transform_layout(graph, modified, current_ep));
ORT_RETURN_IF_ERROR(params.transform_layout(graph, modified, current_ep));
// It is possible some new nodes are introduced during transformation. These nodes can be either existing nodes
// which are reconstructed to update domain or completely new nodes which are necessary for layout transformation.
@ -213,8 +231,6 @@ static Status GetCapabilityForEP(Graph& graph, KernelRegistryManager& kernel_reg
}
}
}
#else // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
ORT_UNUSED_PARAMETER(mode);
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
return Status::OK();
@ -354,8 +370,15 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
// run the function by SessionOption, we should create a function kernel for it and
// delegate the compute to the functions inside the dlls.
std::vector<std::unique_ptr<ComputeCapability>> capabilities;
ORT_RETURN_IF_ERROR(GetCapabilityForEP(graph, kernel_registry_mgr, current_ep, mode, capabilities,
transform_layout_function));
const auto get_capability_params = GetCapabilityForEPParams{
std::ref(graph),
std::cref(kernel_registry_mgr),
std::ref(current_ep),
std::ref(capabilities),
mode,
transform_layout_function,
};
ORT_RETURN_IF_ERROR(GetCapabilityForEP(get_capability_params));
if (capabilities.empty()) {
return Status::OK();
}
@ -570,16 +593,20 @@ static Status PartitionOrtFormatModelImpl(const PartitionParams& partition_param
}
}
const std::string& type = current_ep.Type();
std::vector<std::unique_ptr<ComputeCapability>> capabilities;
// clang-format off
const auto get_capability_params = GetCapabilityForEPParams{
std::ref(graph),
std::cref(kernel_registry_mgr),
std::ref(current_ep),
std::ref(capabilities),
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
TransformLayoutFunction transform_layout_function = partition_params.transform_layout_function;
#else
TransformLayoutFunction transform_layout_function{};
#endif
ORT_RETURN_IF_ERROR(GetCapabilityForEP(graph, kernel_registry_mgr, current_ep,
GraphPartitioner::Mode::kOrtFormatLoad, capabilities,
transform_layout_function));
GraphPartitioner::Mode::kOrtFormatLoad,
partition_params.transform_layout_function,
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
};
// clang-format on
ORT_RETURN_IF_ERROR(GetCapabilityForEP(get_capability_params));
if (capabilities.empty()) {
return Status::OK();
}
@ -594,6 +621,7 @@ static Status PartitionOrtFormatModelImpl(const PartitionParams& partition_param
compilation_entries.reserve(capabilities.size());
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
const std::string& type = current_ep.Type();
for (const auto& capability : capabilities) {
const IndexedSubGraph& indexed_sub_graph = *capability->sub_graph;
const IndexedSubGraph::MetaDef* metadef = indexed_sub_graph.GetMetaDef();

View file

@ -171,56 +171,87 @@ Status KernelTypeStrResolver::SaveToOrtFormat(
}
#endif // !defined(ORT_MINIMAL_BUILD)
Status KernelTypeStrResolver::LoadFromOrtFormat(const fbs::KernelTypeStrResolver& fbs_kernel_type_str_resolver) {
// returns an error message string which is empty if successful
static std::string LoadFromOrtFormatImpl(const fbs::KernelTypeStrResolver& fbs_kernel_type_str_resolver,
OpKernelTypeStrMap& op_kernel_type_str_map_out) {
const auto* fbs_op_kernel_type_str_args = fbs_kernel_type_str_resolver.op_kernel_type_str_args();
ORT_FORMAT_RETURN_IF_NULL(fbs_op_kernel_type_str_args, "op_kernel_type_str_args");
if (!fbs_op_kernel_type_str_args) {
return "op_kernel_type_str_args is null.";
}
OpKernelTypeStrMap op_kernel_type_str_map{};
op_kernel_type_str_map.reserve(fbs_op_kernel_type_str_args->size());
for (const auto* fbs_op_kernel_type_str_args_entry : *fbs_op_kernel_type_str_args) {
ORT_FORMAT_RETURN_IF_NULL(fbs_op_kernel_type_str_args_entry, "op_kernel_type_str_args entry");
if (!fbs_op_kernel_type_str_args_entry) {
return "op_kernel_type_str_args entry is null.";
}
const auto* fbs_op_id = fbs_op_kernel_type_str_args_entry->op_id();
ORT_FORMAT_RETURN_IF_NULL(fbs_op_id, "op_id");
if (!fbs_op_id) {
return "op_id is null.";
}
const auto* fbs_kernel_type_str_args = fbs_op_kernel_type_str_args_entry->kernel_type_str_args();
ORT_FORMAT_RETURN_IF_NULL(fbs_kernel_type_str_args, "kernel_type_str_args");
if (!fbs_kernel_type_str_args) {
return "kernel_type_str_args is null.";
}
KernelTypeStrToArgsMap kernel_type_str_map{};
kernel_type_str_map.reserve(fbs_kernel_type_str_args->size());
for (const auto* fbs_kernel_type_str_args_entry : *fbs_kernel_type_str_args) {
ORT_FORMAT_RETURN_IF_NULL(fbs_kernel_type_str_args_entry, "kernel_type_str_args entry");
if (!fbs_kernel_type_str_args_entry) {
return "kernel_type_str_args entry is null.";
}
const auto* fbs_kernel_type_str = fbs_kernel_type_str_args_entry->kernel_type_str();
ORT_FORMAT_RETURN_IF_NULL(fbs_kernel_type_str, "kernel_type_str");
if (!fbs_kernel_type_str) {
return "kernel_type_str is null.";
}
const auto* fbs_args = fbs_kernel_type_str_args_entry->args();
ORT_FORMAT_RETURN_IF_NULL(fbs_args, "args");
if (!fbs_args) {
return "args is null.";
}
InlinedVector<ArgTypeAndIndex> args{};
args.reserve(fbs_args->size());
for (const auto* fbs_arg : *fbs_args) {
ORT_FORMAT_RETURN_IF_NULL(fbs_arg, "args entry");
if (!fbs_arg) {
return "args entry is null.";
}
args.push_back(ArgTypeAndIndex{
fbs_arg->arg_type() == fbs::ArgType::INPUT ? ArgType::kInput : ArgType::kOutput,
fbs_arg->index()});
}
const auto [it, inserted] = kernel_type_str_map.try_emplace(fbs_kernel_type_str->str(), std::move(args));
ORT_RETURN_IF_NOT(inserted, "Duplicate entry for kernel type str: ", it->first, ". ",
fbs::utils::kInvalidOrtFormatModelMessage);
if (!inserted) {
return MakeString("Duplicate entry for kernel type str: ", it->first, ".");
}
}
OpIdentifier op_id;
ORT_RETURN_IF_ERROR(fbs::utils::LoadOpIdentifierOrtFormat(*fbs_op_id, op_id));
const auto load_op_id_status = fbs::utils::LoadOpIdentifierOrtFormat(*fbs_op_id, op_id);
if (!load_op_id_status.IsOK()) {
return load_op_id_status.ErrorMessage();
}
const auto [it, inserted] = op_kernel_type_str_map.try_emplace(std::move(op_id),
std::move(kernel_type_str_map));
ORT_RETURN_IF_NOT(inserted, "Duplicate entry for op id: ", it->first, ". ",
fbs::utils::kInvalidOrtFormatModelMessage);
if (!inserted) {
return MakeString("Duplicate entry for op id: ", it->first, ".");
}
}
op_kernel_type_str_map_ = std::move(op_kernel_type_str_map);
return Status::OK();
op_kernel_type_str_map_out = std::move(op_kernel_type_str_map);
return "";
}
Status KernelTypeStrResolver::LoadFromOrtFormat(const fbs::KernelTypeStrResolver& fbs_kernel_type_str_resolver) {
const auto error_message = LoadFromOrtFormatImpl(fbs_kernel_type_str_resolver, op_kernel_type_str_map_);
return error_message.empty() ? Status::OK()
: ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, error_message, " ",
fbs::utils::kInvalidOrtFormatModelMessage);
}
void KernelTypeStrResolver::Merge(KernelTypeStrResolver src) {