Utility and example for custom op expansion

Author: gramalingam

examples/custom_op_expansion.py

@@ -0,0 +1,52 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""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
+from onnxscript import script, FLOAT, opset22 as op
+
+
+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))
+
+import onnxscript.utils.replace as replace
+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,30 @@
+# 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)