Paged attention fusion issues

[IanWood1]

We want to be able to fuse https://gist.github.com/Groverkss/36503c2f3163a23b2136c81170c9143c but there are several reshape problems blocking this.

• (1) Fold reshape ops with hal.tensor.barrier and don't bubble reshapes used by hal.tensor.export ops. If the inputs/outputs contain unit dims, they will get propagated throughout the program. For example:

%expanded = tensor.expand_shape %19 [[0, 1], [2], [3]] output_shape [1, 64, 10, 128] : tensor<64x10x128xf16> into tensor<1x64x10x128xf16>
%20 = hal.tensor.barrier join(%expanded : tensor<1x64x10x128xf16>) => %arg5 : !hal.fence
%21 = hal.tensor.export %20 : tensor<1x64x10x128xf16> -> !hal.buffer_view

• (2) Simplify expand of collapse + bubble the expand through the collapse (see Add pattern to bubble expand through collapse #18341 which is blocked by codegen issues).

%collapsed_0 = tensor.collapse_shape %13 [[0, 1], [2], [3, 4]] : tensor<16x4x128x128x32xf16> into tensor<64x128x4096xf16>
%expanded_1 = tensor.expand_shape %collapsed_0 [[0, 1], [2], [3]] output_shape [4, 16, 128, 4096] : tensor<64x128x4096xf16> into tensor<4x16x128x4096xf16>

• (3) Fold away reshapes that expand out unit dim then subsequently collapses them. This should be solved by [Dispatch] Fold collapse(expand) unit dims #19357 (it also contains some examples).
• (4) Re-land pattern to fold broadcast with attention ops. Reverted by Revert "Support fusing broadcast transposes with attention" #19835 (failing dispatch https://gist.github.com/IanWood1/791d392da72cf3b8293596fb17b4ef6d)

@MaheshRavishankar @Groverkss

[MaheshRavishankar]

@Groverkss or @manupak could one of you look at issue 4 here?

[IanWood1]

8b3ba8d "fixes" the problem with (4) but im unsure if its correct in general

[IanWood1]

(4) should be complete after #19838 and #19962