microsoft / onnxscript

ONNX Script enables developers to naturally author ONNX functions and models using a subset of Python.
https://onnxscript.ai/
MIT License
241 stars 45 forks source link

Pattern matching strategy: Interface matching #1618

Open justinchuby opened 1 week ago

justinchuby commented 1 week ago

In matching big subgraphs like attention, it is usually the case that we only care about re-routing the in-coming and out-going edges of the subgraph. Therefore, if we have a method to identify these edges, we can proceed to replacement without having to match the exact content of the subgraph.

image

This strategy can be seen as matching only the "interface" of the subgraph. It would be possible to just specify (a) sub-patterns the characterize the input and output edges, (b) maybe leveraging metadata as well, for this pattern matcher so we limit the search space (aka do match bounded by a known subgraph, e.g. nn.Module).

E.g.

def pattern(a, b):
    c = op.Mul(a, b)
    d = op.DontCare(c)
    e = op.Relu(d)
    return e

DontCare can be greedy or something different that can be specified.

justinchuby commented 1 week ago

cc @gramalingam