Add support for a non-backtracking version of pattern disjunction

[gramalingam]

Several fusions need to support multiple variants of a pattern (such as the optional presence of an Add or some such op). This PR adds support for a non-backtracking version of pattern disjunction. We can now use an "Or" between variants such as "Add(...)" and "MatMul(...)", for example.

Supporting unrestricted Or patterns is more complicated, since failure of one alternative will require backtracking, which will require unbinding any bindings added during the unsuccessful partial search. (We can consider that later, if it seems useful.)

[codecov]

❌ 4 Tests Failed:

Tests completed	Failed	Passed	Skipped
15190	4	15186	2522

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0460_test_hardmax_default_axis

Stack Traces | 0.005s run time

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

The above exception was the direct cause of the following exception:
.nox\test_ort_nightly\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_hardmax_default_axis' (e=No module named 'tests.onnx_backend_test_code.test_hardmax_default_axis') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_hardmax_default_axis.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_hardmax_default_axis.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 FLOAT
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_hardmax_default_axis(x: FLOAT[3,4,5]) -> (FLOAT[3,4,5]):
E       y = opset13.Hardmax(x)
E       return y

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1230_test_top_k

Stack Traces | 0.005s run time

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

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_top_k' (e=No module named 'tests.onnx_backend_test_code.test_top_k') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_top_k.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_top_k.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 FLOAT, INT64
E   from onnxscript.onnx_opset import opset11
E   
E   @script()
E   def bck_test_top_k(x: FLOAT[3,4], k: INT64[1]) -> (FLOAT[3,3], INT64[3,3]):
E       values, indices = opset11.TopK(x, k, axis=1)
E       return values, indices

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0336_test_clip_default_int8_max

Stack Traces | 0.006s run time

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

The above exception was the direct cause of the following exception:
.nox\test_torch_nightly\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_clip_default_int8_max' (e=No module named 'tests.onnx_backend_test_code.test_clip_default_int8_max') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_clip_default_int8_max.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_clip_default_int8_max.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 INT8
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_clip_default_int8_max(x: INT8[3,4,5], max: INT8) -> (INT8[3,4,5]):
E       y = opset13.Clip(x, None, max)
E       return y

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[shubhambhokare1]

Thanks! This will be super useful in avoiding bloating the patterns with parameters. Just mentioning a couple clarifying questions.

[justinchuby]

Could you add an entry to https://github.com/microsoft/onnxscript/blob/main/docs/api/rewriter_pattern.md ? Thanks

[shubhambhokare1]

Thanks, will create a follow up PR converting some of the patterns in ort-fusion to utilize OrPatterns