Utility and example for custom op expansion (#2701)

A utility and an example showing how onnxscript functions can be used to
define function expansions and be used with the inliner to replace calls
to the custom function with an expanded subgraph. This is useful to
perform certain classes of graph surgery easily.

Author: gramalingam

.lintrunner.toml

@@ -39,6 +39,7 @@ include_patterns = [
 exclude_patterns = [
     'tests/**', # Skip linting test files for speed
     # FIXME: Fix typing annotations in these files
+    'examples/custom_op_expansion.py',
     'onnxscript/converter_test.py',
     'onnxscript/converter.py',
     'onnxscript/evaluator_test.py',

examples/custom_op_expansion.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ruff: noqa
+
+"""A utility and an example showing how onnxscript functions can be used to define function expansions
+and be used with the inliner to replace calls to the custom function with an expanded subgraph.
+This is useful to perform certain classes of graph surgery easily.
+"""
+
+import onnx
+
+import onnxscript
+import onnxscript.utils.replace as replace
+
+script = onnxscript.script
+FLOAT = onnxscript.FLOAT
+op = onnxscript.values.opset22
+local = onnxscript.values.Opset("local", 1)
+
+
+# Example Model: Actual models can come from ModelBuilder or Exporter or any other source.
+# Models can contain calls to custom operations (from a custom domain like 'local' here or
+# even "com.microsoft" etc.)
+@script()
+def model_script(X: FLOAT["N"], Y: FLOAT["N"]) -> FLOAT["N"]:
+    DoubleX = op.Add(X, X)
+    YSquare = op.Mul(Y, Y)
+    # Example call to a custom operation
+    Temp1 = local.CustomOp1(DoubleX, YSquare)
+    # Another call to a custom operation with an attribute
+    Temp2 = local.CustomOp2(Temp1, alp=0.9)
+    return Temp2
+
+
+# Define expansions for custom operations as onnxscript functions
+@script(opset=local)
+def CustomOp1(X: FLOAT["N"], Y: FLOAT["N"]) -> FLOAT["N"]:
+    Temp1 = op.Sub(X, Y)
+    return op.Div(Temp1, X)
+
+
+@script(opset=local)
+def CustomOp2(X: FLOAT["N"], alp: float) -> FLOAT["N"]:
+    Temp2 = op.Elu(X, alpha=alp)
+    return op.Mul(Temp2, Temp2)
+
+
+# Now, we can replace the custom operations in the model with their expansions:
+
+functions = [CustomOp1.to_function_proto(), CustomOp2.to_function_proto()]
+
+model = model_script.to_model_proto()
+
+print("Original Model with custom operations:")
+print(onnx.printer.to_text(model))
+
+
+updated_model = replace.replace_functions(model, functions)
+
+print("\nUpdated Model after replacing custom operations with their expansions:")
+print(onnx.printer.to_text(updated_model))

onnxscript/utils/replace.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""A utility function to replace custom operations in a model with their expansions"""
+
+from typing import Sequence
+
+import onnx
+import onnx_ir as ir
+import onnx_ir.passes.common as common_passes
+
+
+def replace_functions(
+    model: onnx.ModelProto, functions: Sequence[onnx.FunctionProto]
+) -> onnx.ModelProto:
+    """A utility function to replace custom operations in a model with their expansions:
+    Args:
+        model: An ONNX ModelProto possibly containing calls to custom operations.
+        functions: A sequence of FunctionProto defining the expansions for the custom operations.
+
+    Returns:
+        An updated ModelProto with custom operations replaced by their expansions.
+    """
+    irmodel = ir.from_proto(model)
+    irfunctions = [ir.from_proto(func) for func in functions]
+    model_functions = irmodel.functions
+    if len(model_functions) != 0:
+        # Since we use inlining, check that there are no model-local functions.
+        raise ValueError("Input model cannot have model-local functions.")
+    for func in irfunctions:
+        model_functions[func.identifier()] = func
+
+    # TODO (rama): Ideally, we should provide users more control over renaming strategy for inlined values.
+    common_passes.InlinePass()(irmodel)
+    common_passes.RemoveUnusedOpsetsPass()(irmodel)
+    return ir.to_proto(irmodel)