diff --git a/tools/ci_build/exclude_unused_ops_and_types.py b/tools/ci_build/exclude_unused_ops_and_types.py index c905bf6a4e..dfcd4e0d61 100644 --- a/tools/ci_build/exclude_unused_ops_and_types.py +++ b/tools/ci_build/exclude_unused_ops_and_types.py @@ -41,7 +41,7 @@ class ExcludeOpsAndTypesRegistrationProcessor(op_registration_utils.Registration return True def process_registration(self, lines: typing.List[str], constant_for_domain: str, operator: str, - start_version: int, end_version: int = None, input_type: str = None): + 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 @@ -51,12 +51,12 @@ class ExcludeOpsAndTypesRegistrationProcessor(op_registration_utils.Registration exclude = self._should_exclude_op(domain, operator, start_version, end_version) # see if a specific typed registration can be excluded - if not exclude and input_type and self._op_types_usage_manager: - exclude = not self._op_types_usage_manager.is_typed_registration_needed(domain, operator, input_type) + 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(input_type) if input_type else '')) + '<{}>'.format(type) if type else '')) for line in lines: self._output_file.write('// ' + line) diff --git a/tools/ci_build/op_registration_utils.py b/tools/ci_build/op_registration_utils.py index c4693045de..c0835e43a4 100644 --- a/tools/ci_build/op_registration_utils.py +++ b/tools/ci_build/op_registration_utils.py @@ -72,7 +72,7 @@ class RegistrationProcessor: ''' def process_registration(self, lines: typing.List[str], domain: str, operator: str, - start_version: int, end_version: int = None, input_type: str = None): + start_version: int, end_version: int = None, type: str = None): ''' Process lines that contain a kernel registration. :param lines: Array containing the original lines containing the kernel registration. @@ -80,13 +80,13 @@ class RegistrationProcessor: :param operator: Operator type :param start_version: Start version :param end_version: End version or None if unversioned registration - :param input_type: Type if this is a typed registration + :param type: Type used in registration, if this is a typed registration ''' pass def process_other_line(self, line): ''' - Process a line that does not contain a kernel registation + Process a line that does not contain a kernel registration :param line: Original line ''' pass @@ -157,10 +157,10 @@ def _process_lines(lines: typing.List[str], offset: int, registration_processor: # e.g. BuildKernelCreateInfo, trim_at = code_line.index(onnx_typed_op) + onnx_typed_op_len + 1 - *_, domain, start_version, input_type, op_type = \ + *_, domain, start_version, type, op_type = \ [arg.strip() for arg in code_line[trim_at: -len(end_mark)].split(',')] registration_processor.process_registration(lines_to_process, domain, op_type, - int(start_version), None, input_type) + int(start_version), None, type) elif onnx_versioned_op in code_line: # e.g. BuildKernelCreateInfo, trim_at = code_line.index(onnx_versioned_typed_op) + onnx_versioned_typed_op_len + 1 - *_, domain, start_version, end_version, input_type, op_type = \ + *_, domain, start_version, end_version, type, op_type = \ [arg.strip() for arg in code_line[trim_at: -len(end_mark)].split(',')] registration_processor.process_registration(lines_to_process, domain, op_type, - int(start_version), int(end_version), input_type) + int(start_version), int(end_version), type) return offset + 1 diff --git a/tools/ci_build/op_registration_validator.py b/tools/ci_build/op_registration_validator.py index c424cb9713..84a201274a 100644 --- a/tools/ci_build/op_registration_validator.py +++ b/tools/ci_build/op_registration_validator.py @@ -33,7 +33,7 @@ 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, input_type: str = None): + start_version: int, end_version: int = None, type: 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/python/util/ort_format_model/operator_type_usage_processors.py b/tools/python/util/ort_format_model/operator_type_usage_processors.py index 035f553744..3c179e3355 100644 --- a/tools/python/util/ort_format_model/operator_type_usage_processors.py +++ b/tools/python/util/ort_format_model/operator_type_usage_processors.py @@ -182,6 +182,18 @@ class DefaultTypeUsageProcessor(TypeUsageProcessor): self._output_types[int(o_str)] = set(values) +class Output0TypedRegistrationProcessor(DefaultTypeUsageProcessor): + ''' + Processor for operators where the first output type is used in a typed kernel registration. + ''' + def __init__(self, domain: str, optype: str): + # init with tracking of output 0 only. + super().__init__(domain, optype, inputs=[], outputs=[0]) + + def is_typed_registration_needed(self, type_in_registration: str): + return type_in_registration in self._output_types[0] + + class OneHotProcessor(TypeUsageProcessor): ''' Processor for the OneHot operator, which requires custom logic as the type registration key is a concatenation of @@ -276,23 +288,23 @@ def _create_operator_type_usage_processors(): # add(DefaultTypeUsageProcessor('ai.onnx', 'Cast', inputs=[0], outputs=[0])) # track input0 and output0 - # Gather and GatherElements have switching on both the data type (input0) and indices type (input1) + # Operators that switch on the type of input 0 and 1 add(DefaultTypeUsageProcessor('ai.onnx', 'Gather', inputs=[0, 1])) add(DefaultTypeUsageProcessor('ai.onnx', 'GatherElements', inputs=[0, 1])) - - # Pow dispatches on base and exponential types add(DefaultTypeUsageProcessor('ai.onnx', 'Pow', inputs=[0, 1])) - # ConstantOfShape switches on size of output type + # Operators that switch on output type add(DefaultTypeUsageProcessor('ai.onnx', 'ConstantOfShape', inputs=[], outputs=[0])) + add(DefaultTypeUsageProcessor('com.microsoft', 'DynamicQuantizeMatMul', inputs=[], outputs=[0])) # Random generator ops produce new data so we track the output type onnx_random_ops = ['RandomNormal', 'RandomNormalLike', 'RandomUniform', 'RandomUniformLike', 'Multinomial'] [add(DefaultTypeUsageProcessor('ai.onnx', op, inputs=[], outputs=[0])) for op in onnx_random_ops] # we only support 'float' as input for [Dynamic]QuantizeLinear so just track the output type - add(DefaultTypeUsageProcessor('ai.onnx', 'QuantizeLinear', inputs=[], outputs=[0])) - add(DefaultTypeUsageProcessor('ai.onnx', 'DynamicQuantizeLinear', inputs=[], outputs=[0])) + # as that's what is used in the typed registration + add(Output0TypedRegistrationProcessor('ai.onnx', 'QuantizeLinear')) + add(Output0TypedRegistrationProcessor('ai.onnx', 'DynamicQuantizeLinear')) # OneHot concatenates type strings into a triple in the typed registration # e.g. float_int64_t_int64_t