Graph/Function/IR mutation, onnxscript, and rewriter

[gramalingam]

It would be worth considering reusing (trace-mode) onnxscript (style/approach) for graph/function mutation in the IR, and using this in the rewriter API.

For example: would be nice to be able to do something like below:

def add_square_node(model: ir.Model, x: ir.Value) -> None:
    # Following returns a "builder" object for the given opset domain "".
    # It fails if model does not yet have an opset-import for "".
    op = model.op_builder("")
    # The following returns a "builder" object for the given opset domain "com.microsoft".
    # It will also add the domain with specified version, to the model imports.
    # It fails if a different version opset-import for "com.microsoft" is already present.
    msft_op = model.op_builder("com.microsoft", 1)

    # The builder objects can be used to add new nodes to the model graph,
    # as below, in onnxscript style.
    x_square = op.Mul(x, x)

By default, the new node created above can be added to the end of the graph/function. However, we should be able to specify the insertion point (as before/after a specific node, for example) via some api.

For a rewrite-rule, we can make the replacement-function take a Replacement Context object as an extra parameter (which may be useful for future extensions as well). Via the replacement-context object (which can contain information about the top-level graph/function with the match), it can get an op_builder object, and use it, perhaps as below:

import onnxscript.rewriter as rewriter

def replace_gelu_pattern(context: rewriter.Context, x: ir.Value) -> ir.Value:
    op = context.op_builder("com.microsoft", 1)
   # or: op = context.graph_or_function().op_builder("com.microsoft", 1)
    return op.Gelu(x)

This may get a bit more complicated for multi-output patterns ... at least, if we need to insert different subgraphs at different points. Straightforward if we allow the user to make the choices (by setting the insertion point), but more complicated if we want to automate it.

[gramalingam]

See https://mlir.llvm.org/doxygen/classmlir_1_1OpBuilder.html for related ideas in MLIR's op builder infrastructure

[gramalingam]

This is currently supported in some form in the rewriter.