How to use rewriter to append new nodes ?

[Johansmm]

I would like to know if it is possible to write pattern functions to update an ONNX model, in order to including new nodes but without remove the previous ones. Here is a simple example of what I would like to achieve

import onnx
from onnxscript.rewriter import pattern, rewrite


def _target_pattern(op, x):
    return op.Relu(x, _outputs=['relu'])


def _replacement_pattern(op, relu, **_):
    return op.Identity(relu)


def _validate_pattern(op, relu, **_):
    return not any(outbound.op_type in ["Identity"] for outbound in relu.consumers())


# Define model
model = onnx.parser.parse_model(
    """< ir_version: 8, opset_import: ["" : 15] >
    test_model (float[N, 3, 112, 112] X) => (float [N, ?, ?, ?] Y){
        A = Relu(X)
        Y = Relu(A)
    }"""
)

# Define transformation rules
rules = [
    pattern.RewriteRule(_target_pattern, _replacement_pattern, _validate_pattern,
                        remove_nodes=False)
]

# Apply rewrites
new_model = rewrite(model, pattern_rewrite_rules=rules)

# Expected sequence
assert [node.op_type for node in new_model.graph.node] == ["Relu", "Identity", "Relu", "Identity"]

I am trying to use the output of _target_pattern in _replacement_pattern to add the node, but I can't seem to stop the nodes from being removed.

Any idea?

[justinchuby]

I think you would have to include the original nodes in the replacement pattern. Looks like a pure matcher would be helpful for this case @gramalingam

[Johansmm]

@justinchuby yes it is exactly what I want to do, could you explain me how to do it ? I tried but I'm not sure I understand how PatternMatcher works.

[leshabirukov]

My "Pure Matcher":

  def _validate_pattern(op, relu, **_):  
      if not any(outbound.op_type in ["Identity"] for outbound in relu.consumers())  
            Remember( relu ) # (or relu.name, if we are to use `ModelProto` for surgery)  
      return False  

By the way, @Johansmm, thanks for _outputs=['relu'] insight, it is useful.
So as op.Split( X, _outputs=2), @justinchuby, please legalize _outputs, - underscore in name scares me.

[Johansmm]

@leshabirukov I do not sure to understand 'Remember', is this a missing feature ?

[leshabirukov]

Remember is pseudocode, ok, I'll clarify it:

  from onnxscript.rewriter import pattern, rewrite
  from onnxscript import ir
  from onnxscript import opset15 as op, script, FLOAT
  
  def _target_pattern(op, x):
      return op.Relu(x, _outputs=['relu'])
  
  
  def _replacement_pattern(op, relu, **_):
      return op.Identity(relu)
  
  insertion_points = []
  def _validate_pattern(op, relu, **_):
      insertion_points.append(relu.name)
      return False
  
  
  @script()
  def ir_model(X: FLOAT['N', 3, 112, 112]) -> FLOAT['N', '_', '_', '_']:
      A = op.Relu(X)
      Y = op.Relu(A)
      return Y
  
  
  # Define transformation rules
  rules = [
      pattern.RewriteRule(_target_pattern, _replacement_pattern, _validate_pattern, )
  ]
  
  ir_model= ir.serde.deserialize_model(ir_model.to_model_proto())
  # Apply rewrites
  ir_model= rewrite( ir_model, pattern_rewrite_rules=rules)
  new_model = ir.serde.serialize_model(ir_model)
  
  print(insertion_points)
  

['A', 'Y']

It is correct connector names, insert desired Identity-es there.

[justinchuby]

My "Pure Matcher":

  def _validate_pattern(op, relu, **_):  
      if not any(outbound.op_type in ["Identity"] for outbound in relu.consumers())  
            Remember( relu ) # (or relu.name, if we are to use `ModelProto` for surgery)  
      return False  

By the way, @Johansmm, thanks for _outputs=['relu'] insight, it is useful. So as op.Split( X, _outputs=2), @justinchuby, please legalize _outputs, - underscore in name scares me.

The underscore name is intended to make it appear different than node inputs. Suggestions welcomed!

[Johansmm]

@leshabirukov Ok I think I understood the global idea:

1. find the tensors where the new functions should be inserted (with onnxscript.rewriter)
2. insert the operation manually (potentially with onnxscript.ir)
    
   Thanks !

My "Pure Matcher":

  def _validate_pattern(op, relu, **_):  
      if not any(outbound.op_type in ["Identity"] for outbound in relu.consumers())  
            Remember( relu ) # (or relu.name, if we are to use `ModelProto` for surgery)  
      return False  

By the way, @Johansmm, thanks for _outputs=['relu'] insight, it is useful. So as op.Split( X, _outputs=2), @justinchuby, please legalize _outputs, - underscore in name scares me.

The underscore name is intended to make it appear different than node inputs. Suggestions welcomed!

@justinchuby what about output_values or output_ir_values, since relu is a ir.Value?

[leshabirukov]

@justinchuby

The underscore name is intended to make it appear different than node inputs. Suggestions welcomed!

Oh, so they not supposed internal or temporary, ok then.

@Johansmm, You're welcome!

[Johansmm]

I have come up with an alternative solution with the idea of “pure matcher:”

import onnx
from onnxscript.rewriter import pattern, rewrite, ir

model = onnx.parser.parse_model(
    """<ir_version: 7, opset_import: [ "" : 10 ]>
    agraph (float[N, 128] X) => (float Y){
        A = Relu(X)
        Y = Relu(A)
    }"""
)

matcher = pattern.SimplePatternMatcher(pattern=pattern._to_graph_pattern(_target_pattern))
ir_model = ir.from_proto(model)
for node in ir_model.graph._nodes:
    if match := matcher.match(None, None, node):
        x = match.bindings["relu"]
        y = ir.Value(name=x.name)
        x.name = f"{x.name}_hidden"
        identity_node = ir.Node("", "Identity", inputs=[x], outputs=[y])
        ir_model.graph.insert_after(node, [identity_node])

new_model = ir.to_proto(ir_model)
onnx.checker.check_model(new_model)

Let me know if there is an easier way to insert the node (and/or loop the nodes in ir_model).
For now, I think we can close the issue :)

[justinchuby]

To loop over the graph you can do

for node in ir_model.graph

Insertion seems fine, or you can also do node.append(identity_node)

[leshabirukov]

My contribution: code for insertion new nodes before and after the match, with avoidance of matching on newborn.

def gen_matches( matcher ):
    for node in ir_model.graph:
        if match := matcher.match(None, None, node):

            old_in = match.bindings["x"]
            new_pre = ir.Value(name=f"{old_in.name}_post")
            # new node will be at 0-th input!
            node.replace_input_with( 0, new_pre )
            yield ir.Node("", "Relu", inputs=[old_in], outputs=[new_pre], name=f'before_{node.name}')

            old_out = match.bindings["relu"]
            new_post = ir.Value(name=f"{old_out.name}_post")
            for (use_n, use_i) in old_out.uses():
                use_n.replace_input_with( use_i, new_post )
            yield ir.Node("", "Relu", inputs=[old_out], outputs=[new_post], name=f'after_{node.name}')

for change in list(gen_matches( matcher ) ):
    ir_model.graph.append( change )

ir_model.graph.sort()