[Snippets][CPU] Introduce MatMul tokenization config, do not tokenize Transpose after MatMul on ARM64 #32592
Conversation
| auto label = ov::pass::pattern::any_input([config](const ov::Output<ov::Node>& out) { | ||
| const auto n = out.get_node_shared_ptr(); | ||
| // Config-aware gating: optionally reject specific Transpose cases around MatMul | ||
| if (ov::is_type<ov::op::v1::Transpose>(n)) { |
There was a problem hiding this comment.
Let's move this check to is_supported_transpose lambda in is_supported_op. We will need to change their signatures, but I think it is okay
| // Transpose before MatMul on input 0 | ||
| bool is_supported_transpose_a = true; | ||
| // Transpose before MatMul on input 1 | ||
| bool is_supported_transpose_b = true; | ||
| // Transpose after MatMul on its output | ||
| bool is_supported_transpose_c = true; |
There was a problem hiding this comment.
I suppose just boolean values are not enough here: not all types of transposes are actually supported even on X86. Ideally, we need to have lambdas here which will replace is_supported_transpose in collapse_subgraph.cpp and is_valid_transpose in mha_tokenization.cpp
There was a problem hiding this comment.
That's a good point to discuss it. Right now, these flags are just being used in is_supported_transpose. It seemed for me that this way we will get less code duplication here and the logic seems to look simpler. What do you think?
There was a problem hiding this comment.
But the original ideas of this PR are to avoid unsupported transposes tokenization + make supported transposes configuration device specific. And the current behavior doesn't fully solves the both problems:
- We still tokenize transposes in MHATokenization and then we have to move them out of the Subgraph
- The supported transposes checks in
MHATokenization, related to X64, are performed for ARM as well
There was a problem hiding this comment.
Disabled all transposes on ARM
| // Disable Transpose after MatMul in general tokenization on ARM64 by default. | ||
| // Keep Transpose before MatMul inputs allowed for flexibility. | ||
| ov::snippets::pass::MatMulConfig mm_cfg; | ||
| mm_cfg.is_supported_transpose_c = false; |
There was a problem hiding this comment.
Shouldn't is_supported_transpose_a/is_supported_transpose_b be false as well on ARM?
There was a problem hiding this comment.
They are being converted to GemmCopyB, actually. Output one is the one that is troublesome right now
There was a problem hiding this comment.
Not really:
- Transpose a is placed on A, not B input. So we don't even insert GemmCopyB there
- Transpose b = true is not supported by GemmCopyB. Please see ExplicitTransposeMatMulInputs callback for the details
There was a problem hiding this comment.
Fair enough, fixed
|
Decided to switch to an alternative approach. PR can be found here: #32676 |
### Details: Introduce a callback function for `ExtractUnsupportedTransposes` pass as a part of `CommonOptimizations::Config` to customize pass behavior depending on Transpose support. For example, ARM64 platform supports transpose decomposition, but MatMul with Transpose A/B is not supported so far. Rest of the (potential) platforms mark Transpose as not supported completely An alternative approach for #32592 ### Tickets: - 176061
2 participants
Details:
Tickets: