Generate broadcast for matrix multiplication where one of the matrix is filled with ones

[umangyadav]

DOR (Definition of Ready)

Ready, no blockers

Description

Let's say i have gemm C = A * B

where A is of shape [B, M, 1]. A is filled with ones.

Let's say B is of shape [B, 1, N].

in this case C would be of shape [B, M, N] .

Because A is filled with ones, C would actually just be broadcasted B.

C = broadcast(B) to shape [B, M, N]

This could be useful transformation for seqLenKV=1 in attention. It is unusual case but in general this transformation can be useful for further graph simplifications.

DOD (Definition of Done)

Unittests are added
Show perf improvement

[pfultz2]

There is a more general approach here though. When K=1, we can rewrite the gemm into just a broadcast+mul. If A or B are all ones we will remove the mul leaving the broadcasted tensor.