[DRAFT: WIP] First version of attention fusion

[gramalingam]

First version of attention fusion.

Limitations:

• Targetting only static shapes for now. Dynamic shapes will alter the pattern.
• Targetting only MHA and the new onnx Attention op.

[codecov]

❌ 57 Tests Failed:

Tests completed	Failed	Passed	Skipped
8388	57	8331	3751

View the top 3 failed tests by shortest run time

::onnxscript.rewriter.onnxruntime.xformers._test_models

Stack Traces | 0s run time

No failure message available

onnxscript.rewriter.generic_pattern_test.GenericPatternTest_0::test_shared_root_value_test

Stack Traces | 0.001s run time

onnxscript/rewriter/generic_pattern_test.py:282: in test_shared_root_value_test
    rule.apply_to_model(ir_model)
onnxscript/rewriter/pattern.py:1364: in apply_to_model
    return RewriteRuleSet([self], commute=commute).apply_to_model(model, verbose=verbose)
onnxscript/rewriter/pattern.py:1535: in apply_to_model
    count = self._apply_to_graph_or_function(model, model.graph, verbose=verbose)
onnxscript/rewriter/pattern.py:1509: in _apply_to_graph_or_function
    delta = rule.try_rewrite(model, graph_or_function, node, verbose=verbose)
onnxscript/rewriter/pattern.py:1334: in try_rewrite
    match = self._matcher.match(
E   TypeError: match() got an unexpected keyword argument 'remove_nodes'

onnxscript.rewriter.generic_pattern_test.GenericPatternTest_0::test_transpose_transpose_onnxscript

Stack Traces | 0.001s run time

onnxscript/rewriter/generic_pattern_test.py:594: in test_transpose_transpose_onnxscript
    rule.apply_to_model(ir_model)
onnxscript/rewriter/pattern.py:1364: in apply_to_model
    return RewriteRuleSet([self], commute=commute).apply_to_model(model, verbose=verbose)
onnxscript/rewriter/pattern.py:1535: in apply_to_model
    count = self._apply_to_graph_or_function(model, model.graph, verbose=verbose)
onnxscript/rewriter/pattern.py:1509: in _apply_to_graph_or_function
    delta = rule.try_rewrite(model, graph_or_function, node, verbose=verbose)
onnxscript/rewriter/pattern.py:1334: in try_rewrite
    match = self._matcher.match(
E   TypeError: match() got an unexpected keyword argument 'remove_nodes'

To view more test analytics, go to the Test Analytics Dashboard
📢 Thoughts on this report? Let us know!

[github-advanced-security]

lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.