A couple of extensions to rewriter

[gramalingam]

A couple of extensions to the rewriter, motivated by fusion optimization experimentation with SmoLLM.

• Support list of constants in match-pattern.
• One multi-output scenario is easy to handle with the single-output pattern-matcher (eg. defining a fusion rule for SkipNormalization): namely when the extra outputs are intermediate values used in the computation of the first value. Extend algorithm to handle this scenario using the efficient single-output matching-algorithm.

An example for the second point is the following pattern:

def skip_norm_pattern(op, input, skip, gamma, epsilon, stash_type):
    skip_sum = op.Add(input, skip)
    normalized = op.SimplifiedLayerNormalization(
        skip_sum,
        gamma,
        axis=-1,
        epsilon=epsilon,
        stash_type=stash_type,
        _domain="com.microsoft")
    return normalized, skip_sum

If we successfully find a match for normalized (which transitively finds a match for all of the pattern subgraph that leads up to normalized), we have also found a successful match for skip_sum, so no need for a multi-output match.

(Will add test-cases later, as I work through the fusion optimizations I am experimenting with.)

[codecov]

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
7230	1	7229	1612

View the top 1 failed tests by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd test_export2python_produces_correct_onnx_script_model_1249_test_triu

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\importlib\__init__.py:127: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_triu'

The above exception was the direct cause of the following exception:
.nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_triu' (e=No module named 'tests.onnx_backend_test_code.test_triu') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_triu.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_triu.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import INT64
E   from onnxscript.onnx_opset import opset14
E   
E   @script()
E   def bck_test_triu(x: INT64[4,5]) -> (INT64[4,5]):
E       y = opset14.Trilu(x)
E       return y

To view individual test run time comparison to the main branch, go to the Test Analytics Dashboard

[justinchuby]

I am still a little confused by the second case. Would it be helpful to make the example more concrete in the PR description? Thanks!

[gramalingam]

I am still a little confused by the second case. Would it be helpful to make the example more concrete in the PR description? Thanks!

Added an example