saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support disabling a typed kernel registration that uses the output type (#6530)

Browse source 
* Update infrastructure to support disabling a typed kernel registration that uses output 0 for the type (vs. the normal use case of input 0).
This commit is contained in:
Scott McKay 2021-02-03 14:22:32 +10:00 committed by GitHub
parent 8d53ef69e5
commit 6cb8f8c812
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 30 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tools/ci_build/exclude_unused_ops_and_types.py
View file

@ -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)

14
tools/ci_build/op_registration_utils.py
View file

@ -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<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, 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<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(
@ -175,10 +175,10 @@ def _process_lines(lines: typing.List[str], offset: int, registration_processor:
# 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, 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

2
tools/ci_build/op_registration_validator.py
View file

@ -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)

24
tools/python/util/ort_format_model/operator_type_usage_processors.py
View file

@ -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