From 20ee399be47700a9bf44688147c3df391b54ebc0 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Wed, 27 Apr 2022 10:31:02 -0700 Subject: [PATCH] Include layout transformation ops in extended minimal build and above. (#11355) --- .../core/framework/kernel_def_hash_helpers.cc | 43 +++---- .../core/providers/cpu/tensor/gather.cc | 15 ++- .../core/providers/cpu/tensor/transpose.cc | 53 +++++---- .../testdata/required_ops_and_types.config | 5 - tools/ci_build/build.py | 2 + tools/ci_build/reduce_op_kernels.py | 111 ++++++++++++++++-- 6 files changed, 170 insertions(+), 59 deletions(-) diff --git a/onnxruntime/core/framework/kernel_def_hash_helpers.cc b/onnxruntime/core/framework/kernel_def_hash_helpers.cc index a00b0dd2d6..15e02f2c23 100644 --- a/onnxruntime/core/framework/kernel_def_hash_helpers.cc +++ b/onnxruntime/core/framework/kernel_def_hash_helpers.cc @@ -13,9 +13,11 @@ std::optional GetHashValueFromStaticKernelHashMap(const std::string& // Since layout transformation can happen in an extended build, if these nodes are not picked up and compiled by // NNAPI or other compiling EPs then we need a way to get the hashes for these nodes. Since the infrastructure // as well as op_schema required to generate these hashes is not available in an extended minimal build, - // we maintain a static map of nodes to hash value. This hash value can then be used to retireive the + // we maintain a static map of nodes to hash value. This hash value can then be used to retrieve the // kernel for the given op. - static std::unordered_map static_kernel_hashes{ + static const std::unordered_map static_kernel_hashes{ + // Note: these region_begin/end markers are used by tools/ci_build/reduce_op_kernels.py + // @@region_begin(layout_transformation_required_kernels)@@ {"Transpose_1", 4324835766923221184ULL}, {"Transpose_13", 17267477159887372848ULL}, {"Squeeze_1", 12889825108950034784ULL}, @@ -31,6 +33,7 @@ std::optional GetHashValueFromStaticKernelHashMap(const std::string& {"Identity_13", 16879814636194901248ULL}, {"Identity_14", 16515685968327103576ULL}, {"Identity_16", 17661628575887109792ULL}, + // @@region_end(layout_transformation_required_kernels)@@ }; auto key = op_type + "_" + std::to_string(since_version); @@ -71,24 +74,24 @@ void UpdateHashForBackwardsCompatibility(HashValue& hash) { // onnxruntime/test/providers/kernel_def_hash_test.cc regarding how/when hashes might change and the best way to // address that. static const std::unordered_map hashes{ - // old new domain, operator, opset[, type] - {2832535737534577496ULL, 16708009824840936392ULL}, // kOnnxDomain, Dropout, 7 - {12198479371038564912ULL, 1718418059112844640ULL}, // kOnnxDomain, Scan, 9 - {2560955351529676608ULL, 3668627007850399040ULL}, // kOnnxDomain, Scan, 11 - {10232409728231027688ULL, 5212043150202938416ULL}, // kOnnxDomain, Not, 1 - {11912523891622051440ULL, 10225383741733918632ULL}, // kOnnxDomain, RoiAlign, 10, float - {18084231515768318048ULL, 17022700455473327752ULL}, // kOnnxDomain, RoiAlign, 10, double - {14033689580222898712ULL, 634727773751317256ULL}, // kOnnxDomain, GatherND, 11 - {646512416908411600ULL, 3064028185911332496ULL}, // kOnnxDomain, GatherND, 12 - {15019893097608892000ULL, 11311962292460032936ULL}, // kOnnxDomain, GatherND, 13 - {14259324427750852648ULL, 7767393334034626736ULL}, // kOnnxDomain, StringNormalizer, 10 - // contrib ops - {7642430665819070720ULL, 8620498355864235632ULL}, // kMSDomain, CropAndResize, 1 - {15019666093341768288ULL, 11924582339825775592ULL}, // kMSDomain, GridSample, 1 - {8466416990072218056ULL, 18418354579469131656ULL}, // kOnnxDomain, LayerNormalization, 1, float - {4058615579523172864ULL, 4827261308628792072ULL}, // kOnnxDomain, LayerNormalization, 1, double - {16349480652468900704ULL, 4809288790945391544ULL}, // kOnnxDomain, SimplifiedLayerNormalization, 1, float - {418129161279605176ULL, 13556035637124174064ULL}}; // kOnnxDomain, SimplifiedLayerNormalization, 1, double + // old new domain, operator, opset[, type] + {2832535737534577496ULL, 16708009824840936392ULL}, // kOnnxDomain, Dropout, 7 + {12198479371038564912ULL, 1718418059112844640ULL}, // kOnnxDomain, Scan, 9 + {2560955351529676608ULL, 3668627007850399040ULL}, // kOnnxDomain, Scan, 11 + {10232409728231027688ULL, 5212043150202938416ULL}, // kOnnxDomain, Not, 1 + {11912523891622051440ULL, 10225383741733918632ULL}, // kOnnxDomain, RoiAlign, 10, float + {18084231515768318048ULL, 17022700455473327752ULL}, // kOnnxDomain, RoiAlign, 10, double + {14033689580222898712ULL, 634727773751317256ULL}, // kOnnxDomain, GatherND, 11 + {646512416908411600ULL, 3064028185911332496ULL}, // kOnnxDomain, GatherND, 12 + {15019893097608892000ULL, 11311962292460032936ULL}, // kOnnxDomain, GatherND, 13 + {14259324427750852648ULL, 7767393334034626736ULL}, // kOnnxDomain, StringNormalizer, 10 + // contrib ops + {7642430665819070720ULL, 8620498355864235632ULL}, // kMSDomain, CropAndResize, 1 + {15019666093341768288ULL, 11924582339825775592ULL}, // kMSDomain, GridSample, 1 + {8466416990072218056ULL, 18418354579469131656ULL}, // kOnnxDomain, LayerNormalization, 1, float + {4058615579523172864ULL, 4827261308628792072ULL}, // kOnnxDomain, LayerNormalization, 1, double + {16349480652468900704ULL, 4809288790945391544ULL}, // kOnnxDomain, SimplifiedLayerNormalization, 1, float + {418129161279605176ULL, 13556035637124174064ULL}}; // kOnnxDomain, SimplifiedLayerNormalization, 1, double auto iter = hashes.find(hash); if (iter != hashes.cend()) { diff --git a/onnxruntime/core/providers/cpu/tensor/gather.cc b/onnxruntime/core/providers/cpu/tensor/gather.cc index 8861a28437..b5d6121557 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather.cc +++ b/onnxruntime/core/providers/cpu/tensor/gather.cc @@ -4,17 +4,26 @@ //https://github.com/onnx/onnx/blob/master/docs/Operators.md#Gather #include "core/providers/cpu/tensor/gather.h" #include "core/common/common.h" -#include "core/platform/threadpool.h" #include "core/framework/op_kernel_type_control_utils.h" +#include "core/platform/threadpool.h" #include "core/providers/op_kernel_type_control.h" namespace onnxruntime { +using DefaultIndexTypes = TypeList; + namespace op_kernel_type_control { -ORT_SPECIFY_OP_KERNEL_ARG_DEFAULT_TYPES_ALL_OPSETS( - kCpuExecutionProvider, kOnnxDomain, Gather, Input, 1, int32_t, int64_t); +ORT_SPECIFY_OP_KERNEL_ARG_DEFAULT_TYPE_LIST_ALL_OPSETS( + kCpuExecutionProvider, kOnnxDomain, Gather, Input, 1, DefaultIndexTypes); + +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +// enable all types for layout transformation +ORT_SPECIFY_OP_KERNEL_ARG_REQUIRED_TYPE_LIST_ALL_OPSETS( + kCpuExecutionProvider, kOnnxDomain, Gather, Input, 1, DefaultIndexTypes); +#else ORT_SPECIFY_OP_KERNEL_ARG_REQUIRED_TYPES_ALL_OPSETS( kCpuExecutionProvider, kOnnxDomain, Gather, Input, 1, int32_t, int64_t); +#endif } // namespace op_kernel_type_control using IndexTypes = ORT_OP_KERNEL_ARG_DEFAULT_TYPE_LIST_ALL_OPSETS(kCpuExecutionProvider, kOnnxDomain, diff --git a/onnxruntime/core/providers/cpu/tensor/transpose.cc b/onnxruntime/core/providers/cpu/tensor/transpose.cc index 5f0f603819..f7c956daf5 100644 --- a/onnxruntime/core/providers/cpu/tensor/transpose.cc +++ b/onnxruntime/core/providers/cpu/tensor/transpose.cc @@ -12,11 +12,22 @@ namespace onnxruntime { +namespace { +using DefaultDataTypes = element_type_lists::All; +} // namespace + namespace op_kernel_type_control { // we're using one set of types for all opsets ORT_SPECIFY_OP_KERNEL_ARG_DEFAULT_TYPE_LIST_ALL_OPSETS( kCpuExecutionProvider, kOnnxDomain, Transpose, Input, 0, - element_type_lists::All); + DefaultDataTypes); + +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +// enable all types for layout transformation +ORT_SPECIFY_OP_KERNEL_ARG_REQUIRED_TYPE_LIST_ALL_OPSETS( + kCpuExecutionProvider, kOnnxDomain, Transpose, Input, 0, + DefaultDataTypes); +#endif } // namespace op_kernel_type_control namespace { @@ -41,15 +52,15 @@ struct MultiIndex { std::vector stride; /* There is one MultiIndex instance per axis in the tensor. - * The array keeps track of the position of a pointer walking through the data. - * Any function using it creates an array of MultiIndex - * then calls function IncrementIndexAndComputeOffsetSetup - * to initialize the array. This constructor does not initialize - * anything because it would be overwritten by function - * IncrementIndexAndComputeOffsetSetup. This one calls method Init. - * Function IncrementIndexAndComputeOffset is called to increment - * the array of MultiIndex to move to the next data in the tensor. - */ + * The array keeps track of the position of a pointer walking through the data. + * Any function using it creates an array of MultiIndex + * then calls function IncrementIndexAndComputeOffsetSetup + * to initialize the array. This constructor does not initialize + * anything because it would be overwritten by function + * IncrementIndexAndComputeOffsetSetup. This one calls method Init. + * Function IncrementIndexAndComputeOffset is called to increment + * the array of MultiIndex to move to the next data in the tensor. + */ MultiIndex() : index(), upper_bound(), stride() { n_axes = 0; } void Init(size_t num_axes) { @@ -67,10 +78,10 @@ struct MultiIndex { }; /* This function initializes an array of MultiIndex of size num_axes (one instance per axis). -* target_dims is the shape of the transposed tensor, stride is linked to the tensor to -* be transposed, if source_dims is the shape, stride[i] = source_dims[i+1] * source_dims[i+2] * ... * 1. -* element_size is the size of the tensor element (sizeof(float), sizeof(double)). -*/ + * target_dims is the shape of the transposed tensor, stride is linked to the tensor to + * be transposed, if source_dims is the shape, stride[i] = source_dims[i+1] * source_dims[i+2] * ... * 1. + * element_size is the size of the tensor element (sizeof(float), sizeof(double)). + */ static void IncrementIndexAndComputeOffsetSetup(MultiIndex& mindex, size_t num_axes, gsl::span target_dims, const gsl::span& stride, size_t element_size) { mindex.Init(num_axes); @@ -86,13 +97,13 @@ static void IncrementIndexAndComputeOffsetSetup(MultiIndex& mindex, size_t num_a } /* This function increments an array of MultiIndex initialized by function IncrementIndexAndComputeOffsetSetup. -* It increments the last dimension, checks if it stays within boundary. If it stays in, it returns, -* otherwise, it reset the dimension to zero and increments the previous one. -* While doing that, every modification brought to the array of indices is applied on the -* pointer local_source. It avoids computing again local_source from the source tensor. -* At every time, the following condition is verified: -* local_source = source + (sum_i mindex[i].index * mindex[i].stride -*/ + * It increments the last dimension, checks if it stays within boundary. If it stays in, it returns, + * otherwise, it reset the dimension to zero and increments the previous one. + * While doing that, every modification brought to the array of indices is applied on the + * pointer local_source. It avoids computing again local_source from the source tensor. + * At every time, the following condition is verified: + * local_source = source + (sum_i mindex[i].index * mindex[i].stride + */ template static inline void IncrementIndexAndComputeOffset(MultiIndex& mindex, const T*& local_source) { // Increment the last dimension. diff --git a/onnxruntime/test/testdata/required_ops_and_types.config b/onnxruntime/test/testdata/required_ops_and_types.config index 3c2e3e9a9e..17687906d7 100644 --- a/onnxruntime/test/testdata/required_ops_and_types.config +++ b/onnxruntime/test/testdata/required_ops_and_types.config @@ -7,8 +7,3 @@ ai.onnx;8;MaxPool{"inputs": {"0": ["float"]}},Sum{"inputs": {"0": ["float"]}} ai.onnx;9;Cast{"inputs": {"0": ["float"]}, "outputs": {"0": ["bool"]}} ai.onnx;11;ArgMax{"inputs": {"0": ["float"]}},If,Loop ai.onnx.ml;1;ArrayFeatureExtractor,LinearClassifier,Normalizer,ZipMap - -# Note: The lines below were added manually. -# TODO find a way to avoid manual modification of this file -# also include Transpose added by layout transformation -ai.onnx;1;Transpose diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index d8e8514c52..05dae9d10e 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -2411,12 +2411,14 @@ def main(): if is_reduced_ops_build(args): from reduce_op_kernels import reduce_ops + is_extended_minimal_build_or_higher = args.minimal_build is None or "extended" in args.minimal_build for config in configs: reduce_ops( config_path=args.include_ops_by_config, build_dir=get_config_build_dir(build_dir, config), enable_type_reduction=args.enable_reduced_operator_type_support, use_cuda=args.use_cuda, + is_extended_minimal_build_or_higher=is_extended_minimal_build_or_higher, ) cmake_extra_args = [] diff --git a/tools/ci_build/reduce_op_kernels.py b/tools/ci_build/reduce_op_kernels.py index 7a142d6a36..0ed8b6f150 100755 --- a/tools/ci_build/reduce_op_kernels.py +++ b/tools/ci_build/reduce_op_kernels.py @@ -3,6 +3,8 @@ # Licensed under the MIT License. import argparse +import io +import re import shutil import sys import typing @@ -25,6 +27,71 @@ from util.ort_format_model.operator_type_usage_processors import OpTypeImplFilte log = get_logger("reduce_op_kernels") +def _adapt_filters_for_extended_minimal_build( + base_required_ops: typing.Optional[dict], base_op_type_impl_filter: typing.Optional[OpTypeImplFilterInterface] +): + """ + Adapts the values returned by parse_config() for an extended minimal build or higher. + In particular: + - Includes ONNX ops needed by layout transformation + """ + # layout transformation requires certain ONNX ops to be available + layout_transformation_required_ops = dict() # op name -> set of opset versions + layout_transformation_required_ops_file = ORT_ROOT / "onnxruntime/core/framework/kernel_def_hash_helpers.cc" + with open(layout_transformation_required_ops_file, mode="r") as f: + region_boundary_pattern = re.compile(r"@@region_(begin|end)\(layout_transformation_required_kernels\)@@") + op_to_hash_pattern = re.compile(r'\{"(\w+)_(\d+)",\s+\w+\},') + in_region = False + for line in f: + region_boundary_match = region_boundary_pattern.search(line) + if region_boundary_match: + in_region = region_boundary_match.group(1) == "begin" + continue + + if not in_region: + continue + + op_to_hash_match = op_to_hash_pattern.search(line) + if op_to_hash_match: + op_name, opset = op_to_hash_match.group(1, 2) + layout_transformation_required_ops.setdefault(op_name, set()).add(int(opset)) + + adapted_required_ops = None + if base_required_ops is not None: + adapted_required_ops = base_required_ops.copy() + required_onnx_ops = adapted_required_ops.setdefault("ai.onnx", dict()) + for op_type, opsets in layout_transformation_required_ops.items(): + for opset in opsets: + required_onnx_opset_ops = required_onnx_ops.setdefault(opset, set()) + required_onnx_opset_ops.add(op_type) + + adapted_op_type_impl_filter = None + if base_op_type_impl_filter is not None: + + class _AdaptedFilter(OpTypeImplFilterInterface): + def __init__(self, filter_to_adapt: OpTypeImplFilterInterface, required_optypes: typing.Set[str]): + self.filter_to_adapt = filter_to_adapt + self.required_optypes = required_optypes + + def is_typed_registration_needed(self, domain: str, optype: str, type_registration_str: str): + # Always require registration for ONNX ops in self.required_optypes. + if domain == "ai.onnx" and optype in self.required_optypes: + return True + return self.filter_to_adapt.is_typed_registration_needed(domain, optype, type_registration_str) + + def get_cpp_entries(self): + # The required types for ops in self.required_optypes must be specified in the C++ implementation. + # Doing that also accounts for globally allowed types. + # We don't need to do anything special with the allowed type overrides here. + return self.filter_to_adapt.get_cpp_entries() + + adapted_op_type_impl_filter = _AdaptedFilter( + base_op_type_impl_filter, set(layout_transformation_required_ops.keys()) + ) + + return (adapted_required_ops, adapted_op_type_impl_filter) + + class _ExcludingRegistrationProcessor(op_registration_utils.RegistrationProcessor): """Registration processor that excludes registrations and writes the result to an output file.""" @@ -32,7 +99,7 @@ class _ExcludingRegistrationProcessor(op_registration_utils.RegistrationProcesso self, required_ops: typing.Optional[dict], op_type_impl_filter: typing.Optional[OpTypeImplFilterInterface], - output_file: str, + output_file: io.TextIOWrapper, ): self._required_ops = required_ops self._op_type_impl_filter = op_type_impl_filter @@ -40,7 +107,7 @@ class _ExcludingRegistrationProcessor(op_registration_utils.RegistrationProcesso def _is_op_required( self, domain: str, operator: str, start_version: int, end_version: typing.Optional[int] - ) -> typing.Tuple[bool, str]: + ) -> bool: """See if an op is required.""" if self._required_ops is None: return True @@ -107,14 +174,20 @@ class _ExcludingRegistrationProcessor(op_registration_utils.RegistrationProcesso return True -def _get_op_reduction_file_path(ort_root: Path, build_dir: Path, original_path: typing.Optional[Path] = None): +def _get_op_reduction_root(build_dir: Path): """ - Return the op reduction file path corresponding to `original_path` or the op reduction file root if unspecified. - Op reduction files are in a subdirectory of `build_dir` but otherwise share the same components of `original_path` + Return the op reduction root directory which is a subdirectory of `build_dir`. + """ + return Path(build_dir, OP_REDUCTION_DIR) + + +def _get_op_reduction_file_path(ort_root: Path, build_dir: Path, original_path: Path): + """ + Return the op reduction file path corresponding to `original_path`. + Op reduction files are in the op reduction root but otherwise share the same components of `original_path` relative to `ort_root`. """ - op_reduction_root = Path(build_dir, OP_REDUCTION_DIR) - return (op_reduction_root / original_path.relative_to(ort_root)) if original_path is not None else op_reduction_root + return _get_op_reduction_root(build_dir) / original_path.relative_to(ort_root) def _generate_provider_registrations( @@ -195,21 +268,31 @@ def _generate_type_control_overrides(ort_root: Path, build_dir: Path, cpp_lines: raise RuntimeError("Insertion point was not found in {}".format(target)) -def reduce_ops(config_path: str, build_dir: str, enable_type_reduction: bool = False, use_cuda: bool = True): +def reduce_ops( + config_path: str, + build_dir: str, + enable_type_reduction: bool, + use_cuda: bool, + is_extended_minimal_build_or_higher: bool, +): """ Reduce op kernel implementations. :param config_path: Path to configuration file that specifies the ops to include :param build_dir: Path to the build directory. The op reduction files will be generated under the build directory. :param enable_type_reduction: Whether per operator type reduction is enabled :param use_cuda: Whether to reduce op kernels for the CUDA provider + :param is_extended_minimal_build_or_higher: Whether this build has at least the features of an extended minimal + build enabled. """ build_dir = Path(build_dir).resolve() build_dir.mkdir(parents=True, exist_ok=True) required_ops, op_type_impl_filter = parse_config(config_path, enable_type_reduction) + if is_extended_minimal_build_or_higher: + required_ops, op_type_impl_filter = _adapt_filters_for_extended_minimal_build(required_ops, op_type_impl_filter) # delete any existing generated files first - op_reduction_root = _get_op_reduction_file_path(ORT_ROOT, build_dir) + op_reduction_root = _get_op_reduction_root(build_dir) if op_reduction_root.is_dir(): log.info(f"Deleting existing op reduction file root directory: {op_reduction_root}") shutil.rmtree(op_reduction_root) @@ -231,7 +314,8 @@ if __name__ == "__main__": 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.", + "See https://onnxruntime.ai/docs/reference/reduced-operator-config-file.html for more " + "information.", ) parser.add_argument( @@ -241,6 +325,12 @@ if __name__ == "__main__": help="Path to the build directory. " "The op reduction files will be generated under the build directory.", ) + parser.add_argument( + "--is_extended_minimal_build_or_higher", + action="store_true", + help="Whether this build has at least the features of an extended minimal build enabled.", + ) + parser.add_argument( "--enable_type_reduction", action="store_true", help="Whether per operator type reduction is enabled." ) @@ -254,4 +344,5 @@ if __name__ == "__main__": build_dir=args.cmake_build_dir, enable_type_reduction=args.enable_type_reduction, use_cuda=args.use_cuda, + is_extended_minimal_build_or_higher=args.is_extended_minimal_build_or_higher, )