diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py index b8e1d03808..915f69cb3a 100644 --- a/onnxscript/optimizer/__init__.py +++ b/onnxscript/optimizer/__init__.py @@ -12,6 +12,7 @@ "inline", "optimize_ir", "optimize", + "register_constant_folder", "remove_unused_nodes", ] @@ -20,8 +21,14 @@ import onnxscript.optimizer._constant_folding as constant_folding from onnxscript import ir -from onnxscript.optimizer._constant_folding import FOLDED_FROM_KEY, basic_constant_propagation +from onnxscript.optimizer._constant_folding import ( + FOLDED_FROM_KEY, + basic_constant_propagation, +) from onnxscript.optimizer._constant_folding import fold_constants as fold_constants_ir +from onnxscript.optimizer._constant_folding import ( + register as register_constant_folder, +) from onnxscript.optimizer._optimizer import optimize_ir _ModelProtoOrIr = TypeVar("_ModelProtoOrIr", onnx.ModelProto, ir.Model) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 003e9174b6..d48e286ea9 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -258,6 +258,30 @@ def lookup_evaluators(self, domain: str, opname: str, version: int): def register( self, opname: str, domain: str = "", version=None ) -> Callable[[PartialEvaluatorFunction], PartialEvaluatorFunction]: + """Register a custom constant folding rule for an operator. + + The decorated function must have the signature: + `(node: ir.Node, op: OptimizerContext, state: OptimizerState) -> ReturnValue` + + Args: + opname: The name of the operator to fold (e.g., "Add"). + domain: The domain of the operator. Defaults to "" (standard ONNX). + version: The opset version or version range for which this rule is valid. + Can be an integer (e.g., 13) or a tuple of (min_version, max_version) + where None indicates no limit (e.g., (13, None)). + + Returns: + A decorator that registers the function as a constant folder for the op. + + Example:: + + .. code-block:: python + + @register_constant_folder("CustomOp", domain="my.domain") + def fold_custom_op(node: ir.Node, op: OptimizerContext, state: OptimizerState) -> ReturnValue: + # custom folding logic + return op.Constant(value_int=42) + """ if (domain, opname) in self.op_evaluators: evaluator_list = self.op_evaluators[(domain, opname)] else: diff --git a/onnxscript/optimizer/_constant_folding_test.py b/onnxscript/optimizer/_constant_folding_test.py index e4e92619e2..44c70d6a62 100644 --- a/onnxscript/optimizer/_constant_folding_test.py +++ b/onnxscript/optimizer/_constant_folding_test.py @@ -879,6 +879,32 @@ def test_initializer_as_graph_output_is_not_removed(self): self.assertIn("y", output_names) self.assertIn("z", output_names) + def test_register_constant_folder(self): + @optimizer.register_constant_folder("CustomAdd", domain="test.custom") + def fold_custom_add(node: ir.Node, op, state: _constant_folding.OptimizerState): + return op.Constant(value_int=42) + + try: + model = ir.from_onnx_text( + """ + + agraph (float[N] x) => (int64 z) { + z = test.custom.CustomAdd(x) + } + """ + ) + + result = _constant_folding.fold_constants(model) + self.assertTrue(result.modified) + self.assertEqual(len(model.graph), 1) + self.assertEqual(model.graph[0].op_type, "Constant") + z_value = model.graph.outputs[0] + self.assertIsNotNone(z_value.const_value) + np.testing.assert_equal(z_value.const_value.numpy(), np.array(42, dtype=np.int64)) + finally: + if ("test.custom", "CustomAdd") in _constant_folding.registry.op_evaluators: + del _constant_folding.registry.op_evaluators[("test.custom", "CustomAdd")] + def _all_value_names_unique(model: ir.Model) -> bool: """Return True if all named values in the top-level graph have unique names."""