graph: backend: dnnl: kernels: eliminate reorders in sdp_decomp

[TaoLv]

Eliminate the reorders in sdpa decomposition kernel. They were used to densify the inputs, eg. from non-trivial strides to format ab. The non-trivial strides now will be used directly to create matmul inputs. There was also a reorder at the end if user requests for non-dense output. Now the second matmul will generate non-dense output directly in this case.
For int8 cases, we still need the reorders to shift the input from u8 to s8 for matmul primitive. These reorders are kept.

To-do:

• Measure the performance impact with different cases on different CPU platforms.
• Extend the changes to sdp_decomp_training kernel as well.
• Add more tests with non-trivial strides to make sure mamtul actually supports them.

[TaoLv]

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph
disable test_device_gpu

[TaoLv]

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph
disable test_device_gpu

[TaoLv]

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph
disable test_device_gpu

[TaoLv]

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph

[dzarukin]

Do you need this one here?

[TaoLv]

I moved it to the cpp files. It's required to call dnnl_get_max_threads() and parallel_nd_ext() functions in the decomposition kernel.

[dzarukin]

Hopefully, we get CPU primitive soon and this kind of logic lives there, but until it's there, shouldn't the logic attempt to create matmul and if it can't be created or doesn't dispatch to brg_matmul indicate a reorder is required?

Would it be naive or too relax, or might actually work?

[TaoLv]

With non-unit stride, the matmul can still create and with ref:any implementation. The explicit stride check here is to make sure the subsequent matmul can dispatch to the optimized brg_matmul implementation. It's assumed that reorder + brg_matmul performs better than a single ref:any.

onednn_verbose,v1,primitive,exec,cpu,matmul,ref:any,undef,src:f16::blocked:ab:16x2:f0 wei:f16::blocked:ba:2x16:f0 dst:f32::blocked:ab::f0,attr-scratchpad:user attr-post-ops:binary_div:f16:0+binary_add:f16:3:ab,,16x8:8x16,0.173828

I can feel that trying to create the matmul with user strides and then check if the implementation string is brg_matmul sound a better idea as there may be other factors impacting the kernel selection. But it may need a large refactor as the "reorder" decision will transfer to matmul rather than sdp_decomp_reorder_t. I will try it out.

[dzarukin]

Question - is it possible to out-source the scratchpad decision on sdp_decomp_reorder_t?

[TaoLv]

This is not the scratchpad memory of the reorder. This is actually the dst memory of the order and part of the bigger scratchpad of the sdp_decomp kernel.

• sdp_decomp kernel: owns the configuration and registration of scratchpad of sdpa (via memory_planning(), prepare_sub_args()).
• sdp_decomp_reorder_t: owns the scratchpad of the reorder primitive (via scratchpad_desc()).

Checking reorder.is_alias() here makes it clearer that if the reorder is not an alias, sdp_decomp kernel needs to register a buffer to save the dense dst of it (aka the dense wei of mm1). Outsourcing the logic to sdp_decomp_reorder_t would make it less obvious and require more information about the sdpa execution order and lifetime.

This is my 2 cents. Let me what you think. Thanks!