diff --git a/onnxruntime/core/providers/op_kernel_type_control_overrides.inc b/onnxruntime/core/providers/op_kernel_type_control_overrides.inc index 9203586db1..6103d78698 100644 --- a/onnxruntime/core/providers/op_kernel_type_control_overrides.inc +++ b/onnxruntime/core/providers/op_kernel_type_control_overrides.inc @@ -47,7 +47,7 @@ namespace op_kernel_type_control { // ORT_SPECIFY_OP_KERNEL_ARG_ALLOWED_TYPES(kOnnxDomain, Cast, Input, 0, float, int64_t); // ORT_SPECIFY_OP_KERNEL_ARG_ALLOWED_TYPES(kOnnxDomain, Cast, Output, 0, float, int64_t); // Specify allowed types globally: -// ORT_SPECIFY_OP_KERNEL_GLOBAL_ALLOWED_TYPES(float, double, int32_t) +// ORT_SPECIFY_OP_KERNEL_GLOBAL_ALLOWED_TYPES(float, double, int32_t); // specify allowed types here diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py old mode 100644 new mode 100755 index c70de5329f..521e83c00b --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -448,13 +448,14 @@ def parse_arguments(): "To enable support for custom operators pass 'custom_ops' as a parameter. " "e.g. '--minimal_build custom_ops'. This can be combined with an 'extended' build by passing " "'--minimal_build extended custom_ops'") + parser.add_argument("--include_ops_by_config", type=str, - help="include ops from config file. " - "See /docs/Reduced_Operator_Kernel_build.md for more information.") + help="Include ops from config file. " + "See /docs/Reduced_Operator_Kernel_build.md for more information.") parser.add_argument("--enable_reduced_operator_type_support", action='store_true', - help='If --include_ops_by_config is specified, and the configuration file was created from ORT ' - 'format models with type reduction enabled, limit the types individual operators support ' - 'where possible to further reduce the build size. ' + help='If --include_ops_by_config is specified, and the configuration file has type reduction ' + 'information, limit the types individual operators support where possible to further ' + 'reduce the build size. ' 'See /docs/Reduced_Operator_Kernel_build.md for more information.') parser.add_argument("--disable_contrib_ops", action='store_true', @@ -475,9 +476,14 @@ def parse_arguments(): help="Generate code coverage when targetting Android (only).") parser.add_argument( "--ms_experimental", action='store_true', help="Build microsoft experimental operators.") + return parser.parse_args() +def is_reduced_ops_build(args): + return args.include_ops_by_config is not None + + def resolve_executable_path(command_or_path): """Returns the absolute path of an executable.""" executable_path = shutil.which(command_or_path) @@ -679,7 +685,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home else "OFF"), "-Donnxruntime_MINIMAL_BUILD_CUSTOM_OPS=" + ("ON" if args.minimal_build and 'custom_ops' in args.minimal_build else "OFF"), - "-Donnxruntime_REDUCED_OPS_BUILD=" + ("ON" if args.include_ops_by_config else "OFF"), + "-Donnxruntime_REDUCED_OPS_BUILD=" + ("ON" if is_reduced_ops_build(args) else "OFF"), "-Donnxruntime_MSVC_STATIC_RUNTIME=" + ("ON" if args.enable_msvc_static_runtime else "OFF"), # enable pyop if it is nightly build "-Donnxruntime_ENABLE_LANGUAGE_INTEROP_OPS=" + ("ON" if args.enable_language_interop_ops else "OFF"), @@ -1354,7 +1360,7 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): # Disable python tests in a reduced build as we don't know which ops have been included and which # models can run. - if args.include_ops_by_config or args.minimal_build is not None: + if is_reduced_ops_build(args) or args.minimal_build is not None: return if is_windows(): @@ -1740,11 +1746,12 @@ def main(): if args.skip_tests: args.test = False - if args.include_ops_by_config and args.update: - from exclude_unused_ops_and_types import exclude_unused_ops_and_types - exclude_unused_ops_and_types(args.include_ops_by_config, - args.enable_reduced_operator_type_support, - args.use_cuda) + if is_reduced_ops_build(args) and args.update: + from reduce_op_kernels import reduce_ops + reduce_ops( + config_path=args.include_ops_by_config, + enable_type_reduction=args.enable_reduced_operator_type_support, + use_cuda=args.use_cuda) if args.use_tensorrt: args.use_cuda = True diff --git a/tools/ci_build/exclude_unused_ops_and_types.py b/tools/ci_build/exclude_unused_ops_and_types.py deleted file mode 100644 index dfcd4e0d61..0000000000 --- a/tools/ci_build/exclude_unused_ops_and_types.py +++ /dev/null @@ -1,180 +0,0 @@ -# !/usr/bin/env python3 -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -import argparse -import op_registration_utils -import os -import shutil -import sys -import typing - -from logger import get_logger - -# add the path to /tools/python so we can import the config parsing and type reduction processing -script_path = os.path.dirname(os.path.realpath(__file__)) -ort_root = os.path.abspath(os.path.join(script_path, '..', '..', )) -ort_tools_py_path = os.path.abspath(os.path.join(ort_root, 'tools', 'python')) -sys.path.append(ort_tools_py_path) - -from util import parse_config # noqa -from util.ort_format_model.operator_type_usage_processors import OperatorTypeUsageManager # noqa - -log = get_logger("exclude_unused_ops_and_types") - - -class ExcludeOpsAndTypesRegistrationProcessor(op_registration_utils.RegistrationProcessor): - def __init__(self, required_ops, op_type_usage_manager, output_file): - self._required_ops = required_ops - self._op_types_usage_manager = op_type_usage_manager - self._output_file = output_file - - def _should_exclude_op(self, domain, operator, start_version, end_version): - if domain not in self._required_ops: - return True - - for opset in self._required_ops[domain]: - if opset >= start_version and (end_version is None or opset <= end_version): - if operator in self._required_ops[domain][opset]: - return False # found a match, do not exclude - - return True - - def process_registration(self, lines: typing.List[str], constant_for_domain: str, operator: str, - start_version: int, end_version: int = None, type: str = None): - # convert from the ORT constant name to the domain string used in the config - domain = op_registration_utils.map_ort_constant_to_domain(constant_for_domain) - exclude = False - - if domain: - # see if entire op is excluded - exclude = self._should_exclude_op(domain, operator, start_version, end_version) - - # see if a specific typed registration can be excluded - if not exclude and type and self._op_types_usage_manager: - exclude = not self._op_types_usage_manager.is_typed_registration_needed(domain, operator, type) - - if exclude: - log.info('Disabling {}:{}({}){}'.format(constant_for_domain, operator, start_version, - '<{}>'.format(type) if type else '')) - for line in lines: - self._output_file.write('// ' + line) - - # edge case of last entry in table where we still need the terminating }; to not be commented out - if lines[-1].rstrip().endswith('};'): - self._output_file.write('};\n') - else: - for line in lines: - self._output_file.write(line) - - def process_other_line(self, line): - self._output_file.write(line) - - def ok(self): - return True - - -def _exclude_unused_ops_and_types_in_registrations(required_operators, - op_type_usage_manager, - provider_registration_paths): - '''rewrite provider registration file to exclude unused ops''' - - for kernel_registration_file in provider_registration_paths: - if not os.path.isfile(kernel_registration_file): - raise ValueError('Kernel registration file {} does not exist'.format(kernel_registration_file)) - - log.info("Processing {}".format(kernel_registration_file)) - - backup_path = kernel_registration_file + '~' - shutil.move(kernel_registration_file, backup_path) - - # read from backup and overwrite original with commented out lines for any kernels that are not required - with open(kernel_registration_file, 'w') as file_to_write: - processor = ExcludeOpsAndTypesRegistrationProcessor(required_operators, - op_type_usage_manager, - file_to_write) - - op_registration_utils.process_kernel_registration_file(backup_path, processor) - - if not processor.ok(): - # error should have already been logged so just exit - sys.exit(-1) - - -def _generate_required_types_cpp_code(ort_root: str, op_type_usage_manager: OperatorTypeUsageManager): - ''' - Generate and insert the C++ code to specify per operator type requirements. - :param ort_root: Root of the ONNX Runtime repository - :param op_type_usage_manager: OperatorTypeUsageManager that contains the required type info - ''' - # get the C++ code to insert - cpp_lines = op_type_usage_manager.get_cpp_entries() if op_type_usage_manager else None - if not cpp_lines: - return - - target = os.path.join(ort_root, 'onnxruntime', 'core', 'providers', 'op_kernel_type_control_overrides.inc') - if not os.path.exists(target) or not os.path.isfile(target): - log.warning('Could not find {}. Skipping generation of C++ code to reduce the types supported by operators.' - .format(target)) - return - - # copy existing content to use as input - src = target + '.tmp' - shutil.copyfile(target, src) - - # find the insertion block and replace any existing content in it - inserted = False - with open(src, 'r') as input, open(target, 'w') as output: - inside_insertion_block = False - for line in input.readlines(): - if '@@insertion_point_begin(allowed_types)@@' in line: - inside_insertion_block = True - output.write(line) - [output.write('{}\n'.format(code_line)) for code_line in cpp_lines] - inserted = True - continue - elif inside_insertion_block: - if '@@insertion_point_end(allowed_types)@@' in line: - inside_insertion_block = False - else: - # we ignore any old lines within the insertion block - continue - - output.write(line) - - os.remove(src) - - if not inserted: - raise RuntimeError('Insertion point was not found in {}'.format(target)) - - # future: how will any global type limitations be provided by the user - # and added to op_kernel_type_control_overrides.inc? - # should they come from a reduced build configuration file, be specified on a build command-line, - # or manually added to op_kernel_type_control_overrides.inc? - - -def exclude_unused_ops_and_types(config_path, enable_type_reduction=False, use_cuda=True): - required_ops, op_type_usage_manager = parse_config(config_path, enable_type_reduction) - - registration_files = op_registration_utils.get_kernel_registration_files(ort_root, use_cuda) - - _exclude_unused_ops_and_types_in_registrations(required_ops, op_type_usage_manager, registration_files) - - _generate_required_types_cpp_code(ort_root, op_type_usage_manager) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Script to exclude unused operator kernels by disabling their registration in ONNX Runtime. " - "The types supported by operator kernels may also be reduced if specified in the config file.", - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - - parser.add_argument("config_path", type=str, - help="Path to configuration file. " - "Create with /tools/python/create_reduced_build_config.py and edit if needed. " - "See /docs/ONNX_Runtime_Format_Model_Usage.md for more information.") - - args = parser.parse_args() - config_path = os.path.abspath(args.config_path) - - exclude_unused_ops_and_types(config_path, enable_type_reduction=True, use_cuda=True) diff --git a/tools/ci_build/op_registration_utils.py b/tools/ci_build/op_registration_utils.py index 8e94957bca..7de949340c 100644 --- a/tools/ci_build/op_registration_utils.py +++ b/tools/ci_build/op_registration_utils.py @@ -75,7 +75,8 @@ class RegistrationProcessor: ''' def process_registration(self, lines: typing.List[str], domain: str, operator: str, - start_version: int, end_version: int = None, type: str = None): + start_version: int, end_version: typing.Optional[int] = None, + type: typing.Optional[str] = None): ''' Process lines that contain a kernel registration. :param lines: Array containing the original lines containing the kernel registration. diff --git a/tools/ci_build/op_registration_validator.py b/tools/ci_build/op_registration_validator.py index 84a201274a..f911a9d185 100644 --- a/tools/ci_build/op_registration_validator.py +++ b/tools/ci_build/op_registration_validator.py @@ -33,7 +33,8 @@ class RegistrationValidator(op_registration_utils.RegistrationProcessor): self.failed = False def process_registration(self, lines: typing.List[str], domain: str, operator: str, - start_version: int, end_version: int = None, type: str = None): + start_version: int, end_version: typing.Optional[int] = None, + type: typing.Optional[str] = None): key = domain + ':' + operator prev_start, prev_end = self.last_op_registrations[key] if key in self.last_op_registrations else (None, None) diff --git a/tools/ci_build/reduce_op_kernels.py b/tools/ci_build/reduce_op_kernels.py new file mode 100755 index 0000000000..8387a541b8 --- /dev/null +++ b/tools/ci_build/reduce_op_kernels.py @@ -0,0 +1,264 @@ +# !/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import op_registration_utils +import os +import re +import shutil +import sys +import typing + +from logger import get_logger + +# add the path to /tools/python so we can import the config parsing and type reduction processing +script_path = os.path.dirname(os.path.realpath(__file__)) +ort_root = os.path.abspath(os.path.join(script_path, '..', '..', )) +ort_tools_py_path = os.path.abspath(os.path.join(ort_root, 'tools', 'python')) +sys.path.append(ort_tools_py_path) + +from util import parse_config # noqa +from util.ort_format_model.operator_type_usage_processors import OperatorTypeUsageManager # noqa + +log = get_logger("reduce_op_kernels") + + +# valid C++ scalar types that can be specified as globally allowed types +_valid_allowed_types = { + "bool", + "uint8_t", "uint16_t", "uint32_t", "uint64_t", + "int8_t", "int16_t", "int32_t", "int64_t", + "MLFloat16", "BFloat16", # in onnxruntime namespace + "float", "double", + "string", # in std namespace +} + + +def _validated_globally_allowed_types(globally_allowed_types: typing.Collection[str]) -> typing.Set[str]: + '''Return a valid set of globally allowed types.''' + # ensure globally_allowed_types is a set + if not isinstance(globally_allowed_types, set): + globally_allowed_types = set(globally_allowed_types) + + if not globally_allowed_types <= _valid_allowed_types: + raise ValueError( + "Globally allowed types must be a subset of valid allowed types. Actual: {}, valid: {}".format( + globally_allowed_types, sorted(_valid_allowed_types))) + + return globally_allowed_types + + +def _type_re_from_globally_allowed_types(globally_allowed_types: typing.Set[str]) -> typing.re.Pattern: + '''Return a regular expression to match type registration strings to a set of globally allowed types.''' + # to keep a registration, the type should match patterns like: + # 1. T0 + # 2. T0_T1_T2 + # where Ti is a member of globally_allowed_types and multiple Ti's are delimited by "_" + # this covers both the common case (1) and special cases like OneHot registration (2) + allowed_type_subpattern = \ + "(?:" + "|".join(re.escape(allowed_type) for allowed_type in sorted(globally_allowed_types)) + ")" + return re.compile("^{0}(?:_{0})*$".format(allowed_type_subpattern)) + + +class _ExcludingRegistrationProcessor(op_registration_utils.RegistrationProcessor): + '''Registration processor that excludes registrations and writes the result to an output file.''' + def __init__(self, required_ops: dict, op_type_usage_manager: typing.Optional[OperatorTypeUsageManager], + globally_allowed_types: typing.Optional[typing.Set[str]], output_file: str): + self._required_ops = required_ops + + if op_type_usage_manager is not None and globally_allowed_types is not None: + raise ValueError("At most one of op_type_usage_manager and globally_allowed_types may be provided.") + + self._op_type_usage_manager = op_type_usage_manager + + self._enable_all_ops = globally_allowed_types is not None and not required_ops + if self._enable_all_ops: + log.info("No required ops were specified but globally allowed types were specified. " + "Globally allowed types will be used to exclude op implementations.") + + self._globally_allowed_types_re = \ + _type_re_from_globally_allowed_types(globally_allowed_types) \ + if globally_allowed_types is not None else None + + self._output_file = output_file + + def _is_op_required(self, domain: str, operator: str, + start_version: int, end_version: typing.Optional[int]) -> typing.Tuple[bool, str]: + '''See if an op should be excluded because it is not required.''' + if self._enable_all_ops: + return True + + if domain not in self._required_ops: + return False + + for opset in self._required_ops[domain]: + if opset >= start_version and (end_version is None or opset <= end_version): + if operator in self._required_ops[domain][opset]: + return True + + return False + + def process_registration(self, lines: typing.List[str], constant_for_domain: str, operator: str, + start_version: int, end_version: typing.Optional[int] = None, + type: typing.Optional[str] = None): + registration_identifier = '{}:{}({}){}'.format(constant_for_domain, operator, start_version, + '<{}>'.format(type) if type else '') + + # convert from the ORT constant name to the domain string used in the config + domain = op_registration_utils.map_ort_constant_to_domain(constant_for_domain) + + exclude = False + reason = "" + + if domain is not None: + if not self._is_op_required(domain, operator, start_version, end_version): + exclude = True + reason = "Entire op is not required." + + if not exclude and type is not None: + if self._op_type_usage_manager is not None: + if not self._op_type_usage_manager.is_typed_registration_needed(domain, operator, type): + exclude = True + reason = "Specific typed registration is not required." + + elif self._globally_allowed_types_re is not None: + if not self._globally_allowed_types_re.match(type): + exclude = True + reason = "Specific typed registration does not contain globally allowed types." + + else: + log.warning('Keeping {} registration from unknown domain: {}' + .format(registration_identifier, constant_for_domain)) + + if exclude: + log.info('Disabling {} registration: {}'.format(registration_identifier, reason)) + for line in lines: + self._output_file.write('// ' + line) + + # edge case of last entry in table where we still need the terminating }; to not be commented out + if lines[-1].rstrip().endswith('};'): + self._output_file.write('};\n') + else: + for line in lines: + self._output_file.write(line) + + def process_other_line(self, line): + self._output_file.write(line) + + def ok(self): + return True + + +def _process_provider_registrations( + ort_root: str, use_cuda: bool, + required_ops: dict, + op_type_usage_manager: typing.Optional[OperatorTypeUsageManager], + globally_allowed_types: typing.Optional[typing.Set[str]]): + '''Rewrite provider registration files.''' + kernel_registration_files = op_registration_utils.get_kernel_registration_files(ort_root, use_cuda) + + for kernel_registration_file in kernel_registration_files: + if not os.path.isfile(kernel_registration_file): + raise ValueError('Kernel registration file {} does not exist'.format(kernel_registration_file)) + + log.info("Processing {}".format(kernel_registration_file)) + + backup_path = kernel_registration_file + '~' + shutil.move(kernel_registration_file, backup_path) + + # read from backup and overwrite original with commented out lines for any kernels that are not required + with open(kernel_registration_file, 'w') as file_to_write: + processor = _ExcludingRegistrationProcessor( + required_ops, op_type_usage_manager, globally_allowed_types, file_to_write) + + op_registration_utils.process_kernel_registration_file(backup_path, processor) + + if not processor.ok(): + # error should have already been logged so just exit + sys.exit(-1) + + +def _insert_type_control_cpp_code(ort_root: str, cpp_lines: typing.Sequence[str]): + ''' + Insert the C++ code to specify operator type requirements. + :param ort_root: Root of the ONNX Runtime repository + :param cpp_lines: The C++ code to insert + ''' + if not cpp_lines: + return + + target = os.path.join(ort_root, 'onnxruntime', 'core', 'providers', 'op_kernel_type_control_overrides.inc') + if not os.path.exists(target) or not os.path.isfile(target): + log.warning('Could not find {}. Skipping generation of C++ code to reduce the types supported by operators.' + .format(target)) + return + + # copy existing content to use as input + src = target + '.tmp' + shutil.copyfile(target, src) + + # find the insertion block and replace any existing content in it + inserted = False + with open(src, 'r') as input, open(target, 'w') as output: + inside_insertion_block = False + for line in input.readlines(): + if '@@insertion_point_begin(allowed_types)@@' in line: + inside_insertion_block = True + output.write(line) + [output.write('{}\n'.format(code_line)) for code_line in cpp_lines] + inserted = True + continue + elif inside_insertion_block: + if '@@insertion_point_end(allowed_types)@@' in line: + inside_insertion_block = False + else: + # we ignore any old lines within the insertion block + continue + + output.write(line) + + os.remove(src) + + if not inserted: + raise RuntimeError('Insertion point was not found in {}'.format(target)) + + +def reduce_ops(config_path: str, enable_type_reduction: bool = False, use_cuda: bool = True): + ''' + Reduce op kernel implementations. + :param config_path: Path to configuration file that specifies the ops to include + :param enable_type_reduction: Whether per operator type reduction is enabled + :param use_cuda: Whether to reduce op kernels for the CUDA provider + ''' + required_ops, op_type_usage_manager, globally_allowed_types = parse_config(config_path, enable_type_reduction) + + if globally_allowed_types is not None: + globally_allowed_types = _validated_globally_allowed_types(globally_allowed_types) + + _process_provider_registrations(ort_root, use_cuda, required_ops, op_type_usage_manager, globally_allowed_types) + + if op_type_usage_manager is not None: + type_control_cpp_code = op_type_usage_manager.get_cpp_entries() + elif globally_allowed_types is not None: + type_control_cpp_code = ["ORT_SPECIFY_OP_KERNEL_GLOBAL_ALLOWED_TYPES({});".format( + ", ".join(sorted(globally_allowed_types)))] + else: + type_control_cpp_code = [] + + _insert_type_control_cpp_code(ort_root, type_control_cpp_code) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Reduces operator kernel implementations in ONNX Runtime. " + "Entire op implementations or op implementations for specific types may be pruned.") + + parser.add_argument("config_path", type=str, + help="Path to configuration file. " + "Create with /tools/python/create_reduced_build_config.py and edit if needed. " + "See /docs/ONNX_Runtime_Format_Model_Usage.md for more information.") + + args = parser.parse_args() + config_path = os.path.abspath(args.config_path) + reduce_ops(config_path, enable_type_reduction=True, use_cuda=True) diff --git a/tools/python/create_reduced_build_config.py b/tools/python/create_reduced_build_config.py index bb2708ddbc..26caf694ef 100644 --- a/tools/python/create_reduced_build_config.py +++ b/tools/python/create_reduced_build_config.py @@ -136,7 +136,7 @@ def main(): # Debug code to validate that the config parsing matches # from util import parse_config - # required_ops, op_type_usage_processor = parse_config(args.config_path) + # required_ops, op_type_usage_processor, _ = parse_config(args.config_path, True) # op_type_usage_processor.debug_dump() diff --git a/tools/python/util/reduced_build_config_parser.py b/tools/python/util/reduced_build_config_parser.py index 5a71b58333..f49ae637b2 100644 --- a/tools/python/util/reduced_build_config_parser.py +++ b/tools/python/util/reduced_build_config_parser.py @@ -14,8 +14,14 @@ except ImportError: def parse_config(config_file: str, enable_type_reduction: bool = False): ''' - Parse the configuration file and return the required operators dictionary, and an OperatorTypeUsageManager. - The basic configuration file format is `domain;opset;op1,op2...` + Parse the configuration file and return the required operators dictionary, and possibly either an + OperatorTypeUsageManager or a globally allowed types list. + + Configuration file lines can do the following: + 1. specify required operators + 2. specify globally allowed types for all operators + + The basic format for specifying required operators (1) is `domain;opset;op1,op2...` e.g. `ai.onnx;11;Add,Cast,Clip,... If the configuration file is generated from ORT format models it may optionally contain JSON for per-operator @@ -38,16 +44,27 @@ def parse_config(config_file: str, enable_type_reduction: bool = False): ai.onnx.OneHot is an example of this, where 3 type names from the inputs are combined into a string. `{"custom": ["float_int64_t_int64_t", "int64_t_string_int64_t"]}` + The format for specifying globally allowed types for all operators (2) is `!globally_allowed_types;T0,T1,...` + Ti should be a C++ scalar type supported by ONNX and ORT. + At most one globally allowed types specification is allowed. + + Specifying per-operator type information and specifying globally allowed types are mutually exclusive - it is an + error to specify both. + :param config_file: Configuration file to parse :param enable_type_reduction: Set to True to use the type information in the config. If False the type information will be ignored. - If the flatbuffers module is unavailable any type information will be ignored - as the usage of the type information is via OperatorTypeUsageManager which has a - dependency on the ORT flatbuffers python schema. - :return: required_ops, op_type_usage_manager: + If the flatbuffers module is unavailable op-specific type information will be ignored + as the per-op usage of the type information is via OperatorTypeUsageManager which has + a dependency on the ORT flatbuffers python schema. + :return: required_ops, op_type_usage_manager, globally_allowed_types: Dictionary of domain:opset:[ops] for required operators OperatorTypeUsageManager manager with operator specific type usage information if available. None if - type reduction was disabled, or the flatbuffers module is not available. + type reduction was disabled, per-op type reduction was not specified, or the flatbuffers module is not + available. + List of globally allowed types. None if type reduction was disabled, or globally allowed types were not + specified. + At most one of the op_type_usage_manager and globally_allowed_types tuple elements will not be None. ''' if not os.path.isfile(config_file): @@ -55,18 +72,29 @@ def parse_config(config_file: str, enable_type_reduction: bool = False): required_ops = {} op_type_usage_manager = OperatorTypeUsageManager() if enable_type_reduction and have_flatbuffers else None + has_op_type_reduction_info = False + globally_allowed_types = None with open(config_file, 'r') as config: for line in [orig_line.strip() for orig_line in config.readlines()]: if not line or line.startswith("#"): # skip empty lines and comments continue + if line.startswith("!globally_allowed_types;"): # handle globally allowed types + if enable_type_reduction: + if globally_allowed_types is not None: + raise RuntimeError("Globally allowed types were already specified.") + globally_allowed_types = [segment.strip() for segment in line.split(';')[1].split(',')] + continue + domain, opset_str, operators_str = [segment.strip() for segment in line.split(';')] opset = int(opset_str) # any type reduction information is serialized json that starts/ends with { and }. # type info is optional for each operator. if '{' in operators_str: + has_op_type_reduction_info = True + # parse the entries in the json dictionary with type info operators = set() cur = 0 @@ -123,4 +151,11 @@ def parse_config(config_file: str, enable_type_reduction: bool = False): else: required_ops[domain][opset].update(operators) - return required_ops, op_type_usage_manager + if enable_type_reduction: + if not has_op_type_reduction_info: + op_type_usage_manager = None + if globally_allowed_types is not None and op_type_usage_manager is not None: + raise RuntimeError( + "Specifying globally allowed types and per-op type reduction info together is unsupported.") + + return required_ops, op_type_usage_manager, globally_allowed_types