mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-18 21:21:17 +00:00
395 lines
16 KiB
Python
395 lines
16 KiB
Python
'''exclude unused ops from build'''
|
|
# !/usr/bin/env python3
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import sys
|
|
import os
|
|
import argparse
|
|
import shutil
|
|
import onnx
|
|
from onnx import AttributeProto as AP
|
|
from logger import log
|
|
|
|
|
|
domain_map = {'': 'kOnnxDomain',
|
|
'ai.onnx': 'kOnnxDomain',
|
|
'ai.onnx.ml': 'kMLDomain',
|
|
'ai.onnx.training': 'ai.onnx.training', # we don't have a constant for the training domains currently
|
|
'ai.onnx.preview.training': 'ai.onnx.preview.training',
|
|
'com.microsoft': 'kMSDomain',
|
|
'com.microsoft.nchwc': 'kMSNchwcDomain',
|
|
'com.microsoft.mlfeaturizers': 'kMSFeaturizersDomain',
|
|
'com.microsoft.dml': 'kMSDmlDomain',
|
|
'com.intel.ai': 'kNGraphDomain',
|
|
'com.xilinx': 'kVitisAIDomain'}
|
|
|
|
|
|
def _map_domain(domain):
|
|
|
|
if domain in domain_map:
|
|
return domain_map[domain]
|
|
|
|
log.warning("Attempt to map unknown domain of {}".format(domain))
|
|
return 'UnknownDomain'
|
|
|
|
|
|
def _extract_ops_from_config(file_path, required_ops):
|
|
'''extract ops from config file of format: domain;opset;op1,op2...'''
|
|
|
|
if not file_path:
|
|
return required_ops
|
|
|
|
if not os.path.isfile(file_path):
|
|
# exit. to continue may result in unexpectedly disabling all kernels.
|
|
log.error('Configuration file {} does not exist'.format(file_path))
|
|
sys.exit(-1)
|
|
|
|
with open(file_path, 'r') as file_to_read:
|
|
|
|
for stripped_line in [line.strip() for line in
|
|
file_to_read.readlines()]:
|
|
|
|
if not stripped_line: # skip empty lines
|
|
continue
|
|
|
|
if stripped_line.startswith("#"): # skip comments
|
|
continue
|
|
|
|
raw_domain, raw_opset, raw_ops =\
|
|
[segment.strip() for segment in stripped_line.split(';')]
|
|
|
|
domain = _map_domain(raw_domain)
|
|
opset = int(raw_opset)
|
|
operators = set([raw_op.strip() for raw_op in raw_ops.split(',')])
|
|
|
|
if domain not in required_ops:
|
|
required_ops[domain] = {opset: operators}
|
|
|
|
elif opset not in required_ops[domain]:
|
|
required_ops[domain][opset] = operators
|
|
|
|
else:
|
|
required_ops[domain][opset].update(operators)
|
|
|
|
return required_ops # end of extract_ops_from_file(...)
|
|
|
|
|
|
def _extract_ops_from_model(model_path, required_ops):
|
|
'''extract ops from models under model_path and return a diction'''
|
|
|
|
if not model_path:
|
|
return required_ops
|
|
|
|
if not os.path.isdir(model_path):
|
|
# exit. to continue may result in unexpectedly disabling all kernels.
|
|
log.error('Directory containing models {} does not exist'.format(model_path))
|
|
sys.exit(-1)
|
|
|
|
def extract_ops_from_graph(graph, operators, domain_opset_map):
|
|
'''extract ops from graph and all subgraphs'''
|
|
|
|
for operator in graph.node:
|
|
|
|
mapped_domain = _map_domain(operator.domain)
|
|
|
|
if mapped_domain not in operators or\
|
|
mapped_domain not in domain_opset_map:
|
|
continue
|
|
|
|
operators[mapped_domain][domain_opset_map[mapped_domain]].add(operator.op_type)
|
|
|
|
for attr in operator.attribute:
|
|
if attr.type == AP.GRAPH: # process subgraph
|
|
extract_ops_from_graph(attr.g, operators, domain_opset_map)
|
|
|
|
# end of extract_ops_from_graph(...)
|
|
|
|
for root, _, files in os.walk(model_path):
|
|
for file in files:
|
|
|
|
if file.endswith('.onnx'):
|
|
model_path = os.path.join(root, file)
|
|
model = onnx.load(model_path)
|
|
domain_opset_map = {}
|
|
|
|
if len(model.opset_import) == 0:
|
|
continue
|
|
|
|
for opset in model.opset_import:
|
|
|
|
mapped_domain = _map_domain(opset.domain)
|
|
domain_opset_map[mapped_domain] = opset.version
|
|
|
|
if mapped_domain not in required_ops:
|
|
required_ops[mapped_domain] = {opset.version: set()}
|
|
|
|
elif opset.version not in required_ops[mapped_domain]:
|
|
required_ops[mapped_domain][opset.version] = set()
|
|
|
|
extract_ops_from_graph(model.graph, required_ops, domain_opset_map)
|
|
|
|
return required_ops # end of extract_ops_from_model(...)
|
|
|
|
|
|
def _exclude_unused_ops_in_provider(operators, provider_path):
|
|
'''rewrite provider file to exclude unused ops'''
|
|
|
|
if not os.path.isfile(provider_path):
|
|
log.warning('File {} does not exist'.format(provider_path))
|
|
return
|
|
|
|
log.info("Processing {}".format(provider_path))
|
|
onnx_op = 'ONNX_OPERATOR_KERNEL_CLASS_NAME'
|
|
onnx_op_len = len(onnx_op)
|
|
onnx_typed_op = 'ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME'
|
|
onnx_typed_op_len = len(onnx_typed_op)
|
|
onnx_versioned_op = 'ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME'
|
|
onnx_versioned_op_len = len(onnx_versioned_op)
|
|
onnx_versioned_typed_op = 'ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME'
|
|
onnx_versioned_typed_op_len = len(onnx_versioned_typed_op)
|
|
end_marks = tuple([');', ')>', ')>,', ')>,};', ')>};'])
|
|
|
|
def should_exclude_op(domain, op_type, opset_from, opset_to=None):
|
|
'''check if should exclude the op from build'''
|
|
|
|
if domain not in operators:
|
|
return True
|
|
|
|
for opset in operators[domain]:
|
|
if opset >= int(opset_from) and (opset_to is None or opset <= int(opset_to)):
|
|
if op_type in operators[domain][opset]:
|
|
return False # found a match, do not exclude
|
|
|
|
return True # end of should_exclude_op(...)
|
|
|
|
def process_lines(lines, offset):
|
|
'''extract op info from a logic code line start from offset to the line end
|
|
with any of end_marks, then trigger callback(op_type, opset_from, opset_to, domain)
|
|
return next line offset and whether current lines are disabled
|
|
'''
|
|
|
|
end_mark = ''
|
|
lines_to_process = []
|
|
|
|
while True: # collect the logical code line
|
|
|
|
lines_to_process.append(lines[offset])
|
|
stripped = lines[offset].strip()
|
|
line_end = False
|
|
|
|
for mark in end_marks:
|
|
if stripped.endswith(mark):
|
|
end_mark = mark
|
|
line_end = True
|
|
break
|
|
|
|
if line_end:
|
|
break
|
|
|
|
offset += 1 # end of while
|
|
|
|
code_line = ''.join([line.strip() for line in lines_to_process])
|
|
|
|
trim_at = 0
|
|
should_exclude = False
|
|
if onnx_op in code_line:
|
|
# e.g. class ONNX_OPERATOR_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 1, Transpose);
|
|
# e.g. BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 7, Cos)>,
|
|
trim_at = code_line.index(onnx_op) + onnx_op_len + 1
|
|
*_, domain, opset, op_type =\
|
|
[arg.strip() for arg in code_line[trim_at: -len(end_mark)].split(',')]
|
|
should_exclude = should_exclude_op(domain, op_type, opset, None)
|
|
|
|
elif onnx_typed_op in code_line:
|
|
# e.g. class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 8, float, Expand);
|
|
# e.g. BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 7, double, Sin)>,
|
|
trim_at = code_line.index(onnx_typed_op) + onnx_typed_op_len + 1
|
|
*_, domain, opset, _, op_type =\
|
|
[arg.strip() for arg in code_line[trim_at: -len(end_mark)].split(',')]
|
|
should_exclude = should_exclude_op(domain, op_type, opset, None)
|
|
|
|
elif onnx_versioned_op in code_line:
|
|
# e.g. class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 1, 10, Unsqueeze);
|
|
# e.g. BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 1, 10, Hardmax)>,
|
|
trim_at = code_line.index(onnx_versioned_op) + onnx_versioned_op_len + 1
|
|
*_, domain, opset_from, opset_to, op_type =\
|
|
[arg.strip() for arg in code_line[trim_at: -len(end_mark)].split(',')]
|
|
should_exclude = should_exclude_op(domain, op_type, opset_from, opset_to)
|
|
|
|
elif onnx_versioned_typed_op in code_line:
|
|
# e.g. class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 7, 9, float, Upsample);
|
|
# e.g. BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(
|
|
# kCpuExecutionProvider, kOnnxDomain, 1, 10, float, LogSoftmax)>,
|
|
trim_at = code_line.index(onnx_versioned_typed_op) + onnx_versioned_typed_op_len + 1
|
|
*_, domain, opset_from, opset_to, _, op_type =\
|
|
[arg.strip() for arg in code_line[trim_at: -len(end_mark)].split(',')]
|
|
should_exclude = should_exclude_op(domain, op_type, opset_from, opset_to)
|
|
|
|
if should_exclude:
|
|
log.info('Disabling: {}'.format(code_line[trim_at: -len(end_mark)]))
|
|
|
|
return offset + 1, should_exclude # end of process_lines(...)
|
|
|
|
lines = []
|
|
with open(provider_path, 'r') as file_to_read:
|
|
lines = file_to_read.readlines()
|
|
|
|
backup_path = provider_path + '~'
|
|
if not os.path.isfile(backup_path):
|
|
shutil.move(provider_path, backup_path)
|
|
|
|
with open(provider_path, 'w') as file_to_write:
|
|
line_offset = 0
|
|
|
|
while line_offset < len(lines):
|
|
|
|
line = lines[line_offset]
|
|
stripped = line.strip()
|
|
|
|
if stripped.startswith('class ONNX_OPERATOR') or\
|
|
stripped.startswith('BuildKernelCreateInfo<ONNX'):
|
|
|
|
next_line_offset, disabled = process_lines(lines,
|
|
line_offset)
|
|
|
|
for index in range(line_offset, next_line_offset):
|
|
if disabled: # comment out unused
|
|
if lines[index].rstrip().endswith('};'):
|
|
file_to_write.write('/*' + lines[index].rstrip() + '*/};\n')
|
|
else:
|
|
file_to_write.write('//' + lines[index])
|
|
|
|
else: # leave as it is
|
|
file_to_write.write(lines[index])
|
|
|
|
line_offset = next_line_offset
|
|
|
|
else: # leave as it is
|
|
file_to_write.write(line)
|
|
line_offset += 1
|
|
|
|
# end of rewrite_cpu_provider(...)
|
|
|
|
|
|
def _exclude_unused_ops_in_providers(required_operators, provider_paths):
|
|
'''rewrite multiple provider files'''
|
|
|
|
for provider_path in provider_paths:
|
|
_exclude_unused_ops_in_provider(required_operators, provider_path)
|
|
|
|
|
|
def _get_provider_paths(ort_root=None, use_cuda=False):
|
|
'''return paths to cpu and cuda providers'''
|
|
|
|
if not ort_root:
|
|
ort_root = os.path.dirname(os.path.abspath(__file__)) + '/../..'
|
|
|
|
provider_path = ort_root + '/onnxruntime/core/providers/{ep}/{ep}_execution_provider.cc'
|
|
contrib_provider_path = ort_root + '/onnxruntime/contrib_ops/{ep}/{ep}_contrib_kernels.cc'
|
|
provider_paths = [provider_path.format(ep='cpu'),
|
|
contrib_provider_path.format(ep='cpu')]
|
|
|
|
if use_cuda:
|
|
provider_paths.append(provider_path.format(ep='cuda'))
|
|
provider_paths.append(contrib_provider_path.format(ep='cuda'))
|
|
|
|
return provider_paths # end of get_provider_paths
|
|
|
|
|
|
def _create_config_file_with_required_ops(required_operators, model_path, config_path, output_file):
|
|
|
|
directory, filename = os.path.split(output_file)
|
|
if not filename:
|
|
log.error("Invalid path to write final config to. {}".format(output_file))
|
|
sys.exit(-1)
|
|
|
|
if not os.path.exists(directory):
|
|
os.makedirs(directory)
|
|
|
|
with open(output_file, 'w') as out:
|
|
model_path_info = '--model_path {} '.format(model_path) if model_path else ''
|
|
config_path_info = '--config_path {}'.format(config_path) if config_path else ''
|
|
out.write("# Generated from {}{}\n".format(model_path_info, config_path_info))
|
|
|
|
for domain in sorted(required_operators.keys()):
|
|
if domain == 'UnknownDomain':
|
|
continue
|
|
|
|
# reverse the mapping of the domain. entry must exist given we created the required_operators dictionary.
|
|
# also need to handle ai.onnx being special-cased as an empty string
|
|
orig_domain = [key for (key, value) in domain_map.items() if value == domain][0]
|
|
if not orig_domain:
|
|
orig_domain = 'ai.onnx'
|
|
|
|
for opset in sorted(required_operators[domain].keys()):
|
|
ops = required_operators[domain][opset]
|
|
if ops:
|
|
out.write("{};{};{}\n".format(orig_domain, opset, ','.join(sorted(ops))))
|
|
|
|
log.info("Wrote set of required operators to {}".format(output_file))
|
|
|
|
|
|
def exclude_unused_ops(models_path, config_path, ort_root=None, use_cuda=True, output_config_path=None):
|
|
'''Determine operators that are used, and either exclude them or create a configuration file that will.
|
|
Note that this called directly from build.py'''
|
|
|
|
if not models_path and not config_path:
|
|
log.error('Please specify model_path and/or config_path.')
|
|
sys.exit(-1)
|
|
|
|
if not ort_root and not output_config_path:
|
|
log.info('ort_root was not specified. Inferring ONNX Runtime repository root from location of this script.')
|
|
|
|
required_ops = _extract_ops_from_config(config_path, _extract_ops_from_model(models_path, {}))
|
|
|
|
if not output_config_path:
|
|
_exclude_unused_ops_in_providers(required_ops, _get_provider_paths(ort_root, use_cuda))
|
|
else:
|
|
_create_config_file_with_required_ops(required_ops, models_path, config_path, output_config_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Script to exclude unused operator kernels by disabling their registration in ONNXRuntime."
|
|
"See /docs/Reduced_Operator_Kernel_build.md for more information",
|
|
usage="Provide model_path, config_path, or both, to disable kernel registration of unused kernels.")
|
|
|
|
parser.add_argument(
|
|
"--model_path", type=str, help="Path to folder containing one or more ONNX models")
|
|
|
|
parser.add_argument(
|
|
"--config_path", type=str, help="Path to configuration file with format of 'domain;opset;op1,op2...'")
|
|
|
|
parser.add_argument(
|
|
"--ort_root", type=str, help="Path to ONNXRuntime repository root. "
|
|
"Inferred from the location of this script if not provided.")
|
|
|
|
parser.add_argument(
|
|
"--write_combined_config_to", type=str,
|
|
help="Optional path to create a configuration file with the combined set of required kernels "
|
|
"from processing --model_path and/or --config_path. If provided, a configuration file will be created "
|
|
"and NO updates will be made to the kernel registrations."
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
models_path = os.path.abspath(args.model_path) if args.model_path else ''
|
|
config_path = os.path.abspath(args.config_path) if args.config_path else ''
|
|
ort_root = os.path.abspath(args.ort_root) if args.ort_root else ''
|
|
|
|
if not models_path and not config_path:
|
|
log.error('Please specify at least either model path or file path.')
|
|
parser.print_help()
|
|
sys.exit(-1)
|
|
|
|
exclude_unused_ops(models_path, config_path, ort_root, use_cuda=True,
|
|
output_config_path=args.write_combined_config_to)
|