diff --git a/src/qonnx/core/execute_custom_node.py b/src/qonnx/core/execute_custom_node.py index 7acf3792..cd6bb605 100644 --- a/src/qonnx/core/execute_custom_node.py +++ b/src/qonnx/core/execute_custom_node.py @@ -27,10 +27,9 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import qonnx.custom_op.registry as registry -from qonnx.util.basic import get_preferred_onnx_opset -def execute_custom_node(node, context, graph, onnx_opset_version=get_preferred_onnx_opset()): +def execute_custom_node(node, context, graph, onnx_opset_version): """Call custom implementation to execute a single custom node. Input/output provided via context.""" op_type = node.op_type diff --git a/src/qonnx/core/modelwrapper.py b/src/qonnx/core/modelwrapper.py index c82a0fee..2ba2984a 100644 --- a/src/qonnx/core/modelwrapper.py +++ b/src/qonnx/core/modelwrapper.py @@ -39,6 +39,7 @@ import qonnx.util.basic as util import qonnx.util.onnx as onnxutil from qonnx.core.datatype import DataType +from qonnx.custom_op.registry import getCustomOp from qonnx.transformation.double_to_single_float import DoubleToSingleFloat from qonnx.transformation.general import ( RemoveStaticGraphInputs, @@ -737,3 +738,24 @@ def set_tensor_sparsity(self, tensor_name, sparsity_dict): qa.tensor_name = tensor_name qa.quant_parameter_tensor_names.append(dt) qnt_annotations.append(qa) + + def get_opset_imports(self): + """Returns a list of imported opsets as a {domain, version} dictionary.""" + return {opset.domain: opset.version for opset in self._model_proto.opset_import} + + def get_customop_wrapper(self, node, fallback_customop_version=util.get_preferred_qonnx_opset()): + """Return CustomOp instance for given node, respecting the + imported opset version in the model protobuf. If the node's domain + is not found in the model's opset imports, fallback_customop_version + will be used.""" + opset_imports = self.get_opset_imports() + try: + opset_import = opset_imports[node.domain] + return getCustomOp(node, onnx_opset_version=opset_import) + except KeyError: + # domain not found in imports, use fallback version + warnings.warn( + f"Domain {node.domain} not found in model opset imports, " + f"using fallback_customop_version={fallback_customop_version}" + ) + return getCustomOp(node, onnx_opset_version=fallback_customop_version) diff --git a/src/qonnx/core/onnx_exec.py b/src/qonnx/core/onnx_exec.py index a8f4774c..ecb808be 100644 --- a/src/qonnx/core/onnx_exec.py +++ b/src/qonnx/core/onnx_exec.py @@ -36,7 +36,7 @@ import qonnx.analysis.topology as ta import qonnx.core.execute_custom_node as ex_cu_node from qonnx.util.basic import ( - get_preferred_onnx_opset, + get_preferred_qonnx_opset, get_sanitize_quant_tensors, is_finn_op, qonnx_make_model, @@ -44,7 +44,7 @@ ) -def execute_node(node, context, graph, return_full_exec_context=False, opset_version=get_preferred_onnx_opset()): +def execute_node(node, context, graph, opset_version, return_full_exec_context=False): """Executes a single node by using onnxruntime or with a custom function. Input/output provided via context.""" @@ -158,7 +158,7 @@ def execute_onnx(model, input_dict, return_full_exec_context=False, start_node=N model_exec_mode = model.get_metadata_prop("exec_mode") if (model_exec_mode is None) or (model_exec_mode == ""): # extract opset version for node-by-node execution - opset_version = model.model.opset_import[0].version + opset_imports = model.get_opset_imports() # execute the model node by node # we can simply walk down the list since the ONNX spec guarantees that it is # topologically sorted @@ -176,7 +176,11 @@ def execute_onnx(model, input_dict, return_full_exec_context=False, start_node=N if get_sanitize_quant_tensors() != 0: # round input values to match quantization annotation execution_context = sanitize_quant_values(model, node.input, execution_context) - execute_node(node, execution_context, graph, return_full_exec_context, opset_version) + if node.domain in opset_imports: + opset_version = opset_imports[node.domain] + else: + opset_version = get_preferred_qonnx_opset() + execute_node(node, execution_context, graph, opset_version, return_full_exec_context) if get_sanitize_quant_tensors() != 0: # round output values to quantization annotation execution_context = sanitize_quant_values(model, node.output, execution_context) diff --git a/src/qonnx/custom_op/channels_last/__init__.py b/src/qonnx/custom_op/channels_last/__init__.py index f1d7c39b..9ffd4e54 100644 --- a/src/qonnx/custom_op/channels_last/__init__.py +++ b/src/qonnx/custom_op/channels_last/__init__.py @@ -2,8 +2,28 @@ from qonnx.custom_op.channels_last.conv import Conv from qonnx.custom_op.channels_last.max_pool import MaxPool -custom_op = dict() +# channels-last ops are defined by the underlying ONNX standard op +# thus, we can define them for any version of the original op +# so we emulate a custom op dictionary that mimics the support for any +# {ChannelsLastOp}_vX instead of hardcoding what versions are supported -custom_op["Conv"] = Conv -custom_op["MaxPool"] = MaxPool -custom_op["BatchNormalization"] = BatchNormalization + +class ChannelsLastCustomOpDict(dict): + def __init__(self): + self._custom_ops = {"Conv": Conv, "MaxPool": MaxPool, "BatchNormalization": BatchNormalization} + + def __getitem__(self, key): + base_key = key.split("_v")[0] # Extract base key (e.g., Conv from Conv_v13) + if base_key in self._custom_ops: + return self._custom_ops[base_key] + raise KeyError(f"Channels-last CustomOp '{key}' not found.") + + def __contains__(self, key): + base_key = key.split("_v")[0] + return base_key in self._custom_ops + + def keys(self): + return self._custom_ops.keys() + + +custom_op = ChannelsLastCustomOpDict() diff --git a/src/qonnx/custom_op/general/__init__.py b/src/qonnx/custom_op/general/__init__.py index 9b14ea8a..e125cbf8 100644 --- a/src/qonnx/custom_op/general/__init__.py +++ b/src/qonnx/custom_op/general/__init__.py @@ -52,3 +52,16 @@ custom_op["Trunc"] = Trunc custom_op["BipolarQuant"] = BipolarQuant custom_op["FloatQuant"] = FloatQuant + +custom_op["DebugMarker_v1"] = DebugMarker +custom_op["QuantAvgPool2d_v1"] = QuantAvgPool2d +custom_op["MaxPoolNHWC_v1"] = MaxPoolNHWC +custom_op["GenericPartition_v1"] = GenericPartition +custom_op["MultiThreshold_v1"] = MultiThreshold +custom_op["XnorPopcountMatMul_v1"] = XnorPopcountMatMul +custom_op["Im2Col_v1"] = Im2Col +custom_op["IntQuant_v1"] = IntQuant +custom_op["Quant_v1"] = IntQuant +custom_op["Trunc_v1"] = Trunc +custom_op["BipolarQuant_v1"] = BipolarQuant +custom_op["FloatQuant_v1"] = FloatQuant diff --git a/src/qonnx/custom_op/general/quantavgpool2d.py b/src/qonnx/custom_op/general/quantavgpool2d.py index c0e24071..00617dcf 100644 --- a/src/qonnx/custom_op/general/quantavgpool2d.py +++ b/src/qonnx/custom_op/general/quantavgpool2d.py @@ -33,7 +33,7 @@ from qonnx.core.datatype import DataType from qonnx.custom_op.base import CustomOp from qonnx.custom_op.general.maxpoolnhwc import compute_pool_output_dim -from qonnx.util.basic import qonnx_make_model +from qonnx.util.basic import get_preferred_onnx_opset, qonnx_make_model class QuantAvgPool2d(CustomOp): @@ -132,7 +132,7 @@ def execute_node(self, context, graph): outputs=[outp], ) - opset_version = self.onnx_opset_version + opset_version = get_preferred_onnx_opset() opset_imports = [helper.make_opsetid("", opset_version)] onnx_kwargs = {"opset_imports": opset_imports} model_avgpool = qonnx_make_model(graph_avgpool, **onnx_kwargs) diff --git a/src/qonnx/custom_op/registry.py b/src/qonnx/custom_op/registry.py index 3540bb5a..442089c3 100644 --- a/src/qonnx/custom_op/registry.py +++ b/src/qonnx/custom_op/registry.py @@ -28,11 +28,12 @@ import importlib -from qonnx.util.basic import get_preferred_onnx_opset - -def getCustomOp(node, onnx_opset_version=get_preferred_onnx_opset(), brevitas_exception=True): - "Return a QONNX CustomOp instance for the given ONNX node, if it exists." +def getCustomOp(node, onnx_opset_version=None, brevitas_exception=True): + "Return a QONNX CustomOp wrapper for the given ONNX node and given opset version," + "if it exists. If opset version is None, the default handler for the op type will be used. " + "If version is specified but the exact version match isn't available, the highest available version " + "smaller than the requested version will be used." op_type = node.op_type domain = node.domain if brevitas_exception: @@ -40,11 +41,31 @@ def getCustomOp(node, onnx_opset_version=get_preferred_onnx_opset(), brevitas_ex domain = domain.replace("onnx.brevitas", "qonnx.custom_op.general") try: opset_module = importlib.import_module(domain) - assert type(opset_module.custom_op) is dict, "custom_op dict not found in Python module %s" % domain - inst_wrapper = opset_module.custom_op[op_type] + assert isinstance(opset_module.custom_op, dict), "custom_op dict not found in Python module %s" % domain + if onnx_opset_version is None: + inst_wrapper = opset_module.custom_op[op_type] + else: + op_type_with_version = op_type + "_v" + str(onnx_opset_version) + if op_type_with_version in opset_module.custom_op: + # priority: if it exists, load the versioned CustomOp wrapper + inst_wrapper = opset_module.custom_op[op_type_with_version] + else: + # when the exact version match is not found + # version handling: use highest available version smaller than requested version + available_versions = [ + int(k.split("_v")[-1]) for k in opset_module.custom_op.keys() if k.startswith(op_type + "_v") + ] + suitable_versions = [v for v in available_versions if v <= onnx_opset_version] + if suitable_versions: + highest_version = max(suitable_versions) + inst_wrapper = opset_module.custom_op[f"{op_type}_v{highest_version}"] + else: + raise Exception( + "Op %s version %s not found in custom opset %s" % (op_type, str(onnx_opset_version), domain) + ) inst = inst_wrapper(node, onnx_opset_version=onnx_opset_version) return inst except ModuleNotFoundError: raise Exception("Could not load custom opset %s, check your PYTHONPATH" % domain) except KeyError: - raise Exception("Op %s not found in custom opset %s" % (op_type, domain)) + raise Exception("Op %s version %s not found in custom opset %s" % (op_type, str(onnx_opset_version), domain)) diff --git a/src/qonnx/transformation/channels_last.py b/src/qonnx/transformation/channels_last.py index 175af058..c352238c 100644 --- a/src/qonnx/transformation/channels_last.py +++ b/src/qonnx/transformation/channels_last.py @@ -270,8 +270,13 @@ def apply(self, model): # Attach to original node n.output[i] = outp_trans_in - # Modify domain + # Modify node domain n.domain = "qonnx.custom_op.channels_last" + opset_imports = model.get_opset_imports() + # Ensure channels_last domain is imported in model + if "qonnx.custom_op.channels_last" not in opset_imports: + onnx_opset = opset_imports[""] + model.model.opset_import.append(helper.make_opsetid("qonnx.custom_op.channels_last", onnx_opset)) # Set modified flag graph_modified = True diff --git a/src/qonnx/transformation/fixedpt_quantize.py b/src/qonnx/transformation/fixedpt_quantize.py index 894d7ea6..ff0c11db 100644 --- a/src/qonnx/transformation/fixedpt_quantize.py +++ b/src/qonnx/transformation/fixedpt_quantize.py @@ -48,7 +48,8 @@ class FixedPointQuantizeParamsFromDict(Transformation): data type or its canonical name rounding_mode: Rounding mode used for conversion into fixed point. Default is "ROUND", - possible values: ["ROUND", "HALF_EVEN", "CEIL", "FLOOR", "UP", "DOWN", "HALF_UP", "HALF_DOWN"] + possible values: ["ROUND", "HALF_EVEN", "CEIL", "FLOOR", "UP", "DOWN", + "HALF_UP", "HALF_DOWN"] """ def __init__(self, fixedpt_dict, rounding_mode="ROUND"): diff --git a/src/qonnx/util/basic.py b/src/qonnx/util/basic.py index 3a3ce2af..e756366d 100644 --- a/src/qonnx/util/basic.py +++ b/src/qonnx/util/basic.py @@ -51,11 +51,19 @@ def get_preferred_onnx_opset(): return 11 +def get_preferred_qonnx_opset(): + "Return preferred ONNX opset version for QONNX" + return 1 + + def qonnx_make_model(graph_proto, **kwargs): "Wrapper around ONNX make_model with preferred qonnx opset version" opset_imports = kwargs.pop("opset_imports", None) if opset_imports is None: - opset_imports = [make_opsetid("", get_preferred_onnx_opset())] + opset_imports = [ + make_opsetid("", get_preferred_onnx_opset()), + make_opsetid("qonnx.custom_op.general", get_preferred_qonnx_opset()), + ] kwargs["opset_imports"] = opset_imports else: kwargs["opset_imports"] = opset_imports diff --git a/tests/core/test_custom_onnx_exec.py b/tests/core/test_custom_onnx_exec.py index 8eec7156..54b71754 100644 --- a/tests/core/test_custom_onnx_exec.py +++ b/tests/core/test_custom_onnx_exec.py @@ -32,6 +32,8 @@ import qonnx.core.execute_custom_node as ex_cu_node from qonnx.custom_op.registry import getCustomOp +mt_node_version = 1 + def test_execute_custom_node_multithreshold(): inputs = np.ndarray( @@ -155,7 +157,7 @@ def test_execute_custom_node_multithreshold(): execution_context["v"] = inputs execution_context["thresholds"] = threshold_values - ex_cu_node.execute_custom_node(node_def, execution_context, graph_def) + ex_cu_node.execute_custom_node(node_def, execution_context, graph_def, mt_node_version) outputs = np.ndarray( shape=(6, 3, 2, 2), @@ -250,7 +252,7 @@ def test_execute_custom_node_multithreshold(): ) graph_def = helper.make_graph([node_def], "test_model", [v, thresholds], [out]) - ex_cu_node.execute_custom_node(node_def, execution_context, graph_def) + ex_cu_node.execute_custom_node(node_def, execution_context, graph_def, mt_node_version) outputs_scaled = 2.0 * outputs - 1.0 assert (execution_context["out"] == outputs_scaled).all() @@ -270,7 +272,7 @@ def test_execute_custom_node_multithreshold(): execution_context["v"] = inputs_nhwc graph_def = helper.make_graph([node_def], "test_model", [v_nhwc, thresholds], [out_nhwc]) - ex_cu_node.execute_custom_node(node_def, execution_context, graph_def) + ex_cu_node.execute_custom_node(node_def, execution_context, graph_def, mt_node_version) assert (execution_context["out"] == outputs_nhwc).all() # check the set of allowed values op_inst = getCustomOp(node_def) diff --git a/tests/core/test_modelwrapper.py b/tests/core/test_modelwrapper.py index 722f0fb1..995bcb17 100644 --- a/tests/core/test_modelwrapper.py +++ b/tests/core/test_modelwrapper.py @@ -68,6 +68,7 @@ def test_modelwrapper(): inp_sparsity = {"dw": {"kernel_shape": [3, 3]}} model.set_tensor_sparsity(first_conv_iname, inp_sparsity) assert model.get_tensor_sparsity(first_conv_iname) == inp_sparsity + assert model.get_opset_imports() == {"": 8} def test_modelwrapper_set_get_rm_initializer(): diff --git a/tests/custom_op/test_customop_version.py b/tests/custom_op/test_customop_version.py new file mode 100644 index 00000000..5364df61 --- /dev/null +++ b/tests/custom_op/test_customop_version.py @@ -0,0 +1,134 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of qonnx nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import onnx.parser as oprs + +import qonnx.custom_op.general as general +from qonnx.core.modelwrapper import ModelWrapper +from qonnx.custom_op.base import CustomOp +from qonnx.custom_op.registry import getCustomOp + + +class VerTestOp_v1(CustomOp): + def get_nodeattr_types(self): + my_attrs = {"v1_attr": ("i", True, 0)} + return my_attrs + + def make_shape_compatible_op(self, model): + ishape = model.get_tensor_shape(self.onnx_node.input[0]) + return super().make_const_shape_op(ishape) + + def infer_node_datatype(self, model): + node = self.onnx_node + # data type stays the same + dtype = model.get_tensor_datatype(node.input[0]) + model.set_tensor_datatype(node.output[0], dtype) + + def execute_node(self, context, graph): + node = self.onnx_node + context[node.output[0]] = context[node.input[0]] + + def verify_node(self): + pass + + +class VerTestOp_v2(VerTestOp_v1): + def get_nodeattr_types(self): + my_attrs = {"v2_attr": ("i", True, 0)} + return my_attrs + + +class VerTestOp_v3(VerTestOp_v2): + def get_nodeattr_types(self): + my_attrs = {"v3_attr": ("i", True, 0)} + return my_attrs + + +def make_vertest_model(vertest_ver, no_opset_import): + ishp = (1, 10) + oshp = ishp + ishp_str = str(list(ishp)) + oshp_str = str(list(oshp)) + if no_opset_import: + opset_import = "" + else: + opset_import = f', "qonnx.custom_op.general" : {vertest_ver}' + input = f""" + < + ir_version: 7, + opset_import: ["" : 9{opset_import}] + > + agraph (float{ishp_str} in0) => (float{oshp_str} out0) + {{ + out0 = qonnx.custom_op.general.VerTestOp< + v{vertest_ver}_attr={vertest_ver} + >(in0) + }} + """ + model = oprs.parse_model(input) + model = ModelWrapper(model) + return model + + +def test_customop_version(): + # unspecified version defaults to v1 implementation + general.custom_op["VerTestOp"] = VerTestOp_v1 + # v1 version is also explicitly registered + general.custom_op["VerTestOp_v1"] = VerTestOp_v1 + general.custom_op["VerTestOp_v2"] = VerTestOp_v2 + general.custom_op["VerTestOp_v3"] = VerTestOp_v3 + + # if onnx is lacking the opset import, should default to v1 handler + # (since we set custom_op["VerTestOp"] = VerTestOp_v1) + model = make_vertest_model(1, True) + inst = getCustomOp(model.graph.node[0]) + assert isinstance(inst, VerTestOp_v1) + # alternatively, when using ModelWrapper.get_customop_wrapper and onnx is + # lacking the opset import, should fall back to the specified version + inst = model.get_customop_wrapper(model.graph.node[0], fallback_customop_version=2) + assert isinstance(inst, VerTestOp_v2) + + for ver in [1, 2, 3]: + model = make_vertest_model(ver, False) + # use ModelWrapper.get_customop_wrapper for implicit + # fetching of op version + inst = model.get_customop_wrapper(model.graph.node[0]) + assert inst.get_nodeattr(f"v{ver}_attr") == ver + # explicitly specify onnx_opset_version in getCustomOp + # note: new code should avoid calling getCustomOp directly like this + # and instead use ModelWrapper.get_customop_wrapper + inst = getCustomOp(model.graph.node[0], onnx_opset_version=ver) + assert inst.get_nodeattr(f"v{ver}_attr") == ver + # unspecified version getCustomOp should default to v1 handler + model = make_vertest_model(1, False) + inst = getCustomOp(model.graph.node[0]) + assert isinstance(inst, VerTestOp_v1) + # requesting v4 should return largest available version (v3 in this case) + model = make_vertest_model(3, False) + inst = getCustomOp(model.graph.node[0], onnx_opset_version=4) + assert isinstance(inst, VerTestOp_v3)