diff --git a/onnxscript/_internal/irbuilder.py b/onnxscript/_internal/irbuilder.py index 1ae3c7bdb1..4b6fca188e 100644 --- a/onnxscript/_internal/irbuilder.py +++ b/onnxscript/_internal/irbuilder.py @@ -77,8 +77,8 @@ def append_parameter(self, parameter: ir.Value | ir.Attr) -> None: def add_nested_function(self, fun: IRFunction) -> None: self.nested_functions[fun.name] = fun - def get_called_functions(self) -> dict[str, values.OnnxFunction]: - called_functions: dict[str, values.OnnxFunction] = {} + def get_called_functions(self) -> dict[ir.OperatorIdentifier, values.OnnxFunction]: + called_functions: dict[ir.OperatorIdentifier, values.OnnxFunction] = {} def visit(function_ir: IRFunction): for node in ir.traversal.RecursiveGraphIterator(function_ir.graph): @@ -87,9 +87,16 @@ def visit(function_ir: IRFunction): add(callee) def add(f: values.OnnxFunction): - if f.name in called_functions: + identifier = f.function_ir.identifier() + existing = called_functions.get(identifier) + if existing is f: return - called_functions[f.name] = f + if existing is not None: + raise ValueError( + "Multiple distinct ONNX functions use the same identifier " + f"{identifier!r}. Give each function a unique domain, name, or overload." + ) + called_functions[identifier] = f visit(f.function_ir) visit(self) diff --git a/onnxscript/_internal/main.py b/onnxscript/_internal/main.py index 804dbfd135..201cd48842 100644 --- a/onnxscript/_internal/main.py +++ b/onnxscript/_internal/main.py @@ -41,6 +41,8 @@ def script_check( def script( opset: Optional[values.Opset] = None, default_opset: Optional[values.Opset] = None, + *, + op_type: str | None = None, **kwargs: Any, ) -> Callable[[Callable[_P, _R]], onnxscript.OnnxFunction[_P, _R]]: """Main decorator. Declares a function as an onnx function. @@ -48,6 +50,8 @@ def script( Args: opset: Opset the function belongs to (see :ref:`l-api-opsets`). default_opset: Opset to use for operators not in the function's opset. + op_type: Optional custom ONNX operator name. It must be a non-empty + string. When omitted or ``None``, the Python function name is used. kwargs: Additional keyword arguments. Returns: @@ -77,6 +81,10 @@ def log2(x): raise TypeError( "Script parameter must be an opset. Did you use @script instead of @script()?" ) + if op_type is not None and not isinstance(op_type, str): + raise TypeError(f"op_type must be a string or None, got {type(op_type).__name__}") + if isinstance(op_type, str) and not op_type.strip(): + raise ValueError("op_type must be a non-empty string") def transform(f: Callable[_P, _R]) -> onnxscript.OnnxFunction[_P, _R]: if not inspect.isfunction(f): @@ -92,6 +100,9 @@ def transform(f: Callable[_P, _R]) -> onnxscript.OnnxFunction[_P, _R]: env = module.__dict__.copy() env.update(closure.nonlocals) result = script_check(f_ast, opset, env, src, default_opset=default_opset) + if isinstance(op_type, str): + result.name = op_type + result.graph.name = op_type # TODO: add transformations. return onnxscript.OnnxFunction(opset, f, result, src, kwargs) diff --git a/onnxscript/_internal/main_test.py b/onnxscript/_internal/main_test.py new file mode 100644 index 0000000000..b7f3085030 --- /dev/null +++ b/onnxscript/_internal/main_test.py @@ -0,0 +1,153 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import inspect +import unittest + +import onnx + +from onnxscript import FLOAT, script +from onnxscript.onnx_opset import opset18 as op +from onnxscript.values import Opset + + +class ScriptCustomOpTypeTest(unittest.TestCase): + def test_custom_op_type_is_used_consistently(self): + custom_opset = Opset("com.example.custom", 1) + + @script(custom_opset, op_type="MY_NEW_NAME_OP", producer_name="onnxscript-test") + def custom_op(x: FLOAT) -> FLOAT: + return op.Abs(x) + + @script() + def caller(x: FLOAT) -> FLOAT: + return custom_op(x) + + self.assertEqual(custom_op.name, "MY_NEW_NAME_OP") + self.assertEqual(custom_op.__name__, "custom_op") + self.assertEqual(custom_op.op_signature.name, "MY_NEW_NAME_OP") + self.assertEqual(custom_op.function_ir.graph.name, "MY_NEW_NAME_OP") + self.assertEqual(custom_op.to_function_proto().name, "MY_NEW_NAME_OP") + + model = caller.to_model_proto() + call_node = model.graph.node[0] + self.assertEqual(call_node.op_type, "MY_NEW_NAME_OP") + self.assertEqual(call_node.domain, "com.example.custom") + self.assertEqual( + {(function.domain, function.name) for function in model.functions}, + {("com.example.custom", "MY_NEW_NAME_OP")}, + ) + self.assertNotIn("op_type", custom_op.kwargs) + self.assertEqual( + custom_op.to_model_proto(io_types=FLOAT).producer_name, "onnxscript-test" + ) + onnx.checker.check_model(model) + + def test_issue_example_uses_default_local_opset(self): + @script(op_type="MY_NEW_NAME_OP") + def custom_op(x: FLOAT) -> FLOAT: + return op.Abs(x) + + self.assertEqual(custom_op.name, "MY_NEW_NAME_OP") + self.assertEqual(custom_op.to_function_proto().name, "MY_NEW_NAME_OP") + + def test_default_op_type_is_python_function_name(self): + @script() + def default_name(x: FLOAT) -> FLOAT: + return op.Abs(x) + + self.assertEqual(default_name.name, "default_name") + self.assertEqual(default_name.to_function_proto().name, "default_name") + + def test_none_uses_default_name_and_public_signature_is_stable(self): + @script(op_type=None) + def default_name(x: FLOAT) -> FLOAT: + return op.Abs(x) + + self.assertIsNone(inspect.signature(script).parameters["op_type"].default) + self.assertEqual(default_name.name, "default_name") + + def test_custom_name_can_be_reused_in_different_domains(self): + first_opset = Opset("com.example.first", 1) + second_opset = Opset("com.example.second", 1) + + @script(first_opset, op_type="SharedName") + def first(x: FLOAT) -> FLOAT: + return op.Abs(x) + + @script(second_opset, op_type="SharedName") + def second(x: FLOAT) -> FLOAT: + return op.Neg(x) + + @script() + def intermediate(x: FLOAT) -> FLOAT: + return first(x) + + @script() + def caller(x: FLOAT) -> tuple[FLOAT, FLOAT]: + return intermediate(x), second(x) + + model = caller.to_model_proto() + self.assertTrue( + { + ("com.example.first", "SharedName"), + ("com.example.second", "SharedName"), + }.issubset({(function.domain, function.name) for function in model.functions}) + ) + onnx.checker.check_model(model) + + def test_distinct_functions_cannot_share_an_identifier(self): + custom_opset = Opset("com.example.custom", 1) + + @script(custom_opset, op_type="Duplicate") + def first(x: FLOAT) -> FLOAT: + return op.Abs(x) + + @script(custom_opset, op_type="Duplicate") + def second(x: FLOAT) -> FLOAT: + return op.Neg(x) + + @script() + def caller(x: FLOAT) -> tuple[FLOAT, FLOAT]: + return first(x), second(x) + + with self.assertRaisesRegex(ValueError, "same identifier"): + caller.to_model_proto() + + def test_repeated_calls_to_same_function_are_deduplicated(self): + custom_opset = Opset("com.example.custom", 1) + + @script(custom_opset, op_type="CalledTwice") + def custom_op(x: FLOAT) -> FLOAT: + return op.Abs(x) + + @script() + def caller(x: FLOAT) -> tuple[FLOAT, FLOAT]: + return custom_op(x), custom_op(x) + + model = caller.to_model_proto() + self.assertEqual(len(model.functions), 1) + + def test_invalid_custom_op_types_are_rejected(self): + invalid_values = ("", " ", 42) + + for value in invalid_values: + with self.subTest(value=value), self.assertRaises((TypeError, ValueError)): + + @script(op_type=value) + def invalid(x: FLOAT) -> FLOAT: + return op.Abs(x) + + def test_non_c_identifier_names_are_supported(self): + for value in ("1LeadingDigit", "Name-With-Punctuation", "Name With Spaces"): + with self.subTest(value=value): + + @script(op_type=value) + def custom_op(x: FLOAT) -> FLOAT: + return op.Abs(x) + + self.assertEqual(custom_op.name, value) + + +if __name__ == "__main__": + unittest.main()