commit 16bdaa90e2f02769db2a8949ab88e58c9443392d
author lialan <xunli@amd.com> Tue May 28 13:31:48 2024 -0400
committer GitHub <noreply@github.com> Tue May 28 10:31:48 2024 -0700
tree f3ac5012cc1503a7b6d9ac860e77c9a36809549f
parent 6c75aa1083d6f9a1fa7f2b1ddd032decc9e87aa7

Data tiling: transpose narrow-N into narrow-M (#17446)

(This is a rebasing PR for https://github.com/iree-org/iree/pull/16890 )

This is a generic idea in the design of matrix multiplication
implementations: the M and N dimensions play symmetrical roles, so there
is this opportunity to halve the problem space by transposition. The
immediate motivation is ukernels: we have chosen to implement narrow
ukernels only for the narrow-M cases, not narrow-N, in preparation for
this. This is the reason why this PR is a 5%-9% e2e speedup on multiple
ML models with ukernels (and a > 2x speedup on matvec microbenchmarks).

The idea should be beneficial beyond ukernels though:
* With codegen (outside of ukernels), inner unit dimensions have often
caused codegen to fall off of good vectorization cases. This
transposition moves unit or generally smaller static dimensions to the
outer dimensions, which will help with that.
* When we get to serious distribution tiling (#16410), the reduction of
generality will again greatly help.

This transposition is made easier by (and was all along part of the idea
in) the RHS-transposition in `mmt4d` (the `t` in `mmt4d`), as generally
with matrix multiplication

```
B * Transpose(A) == Transpose( A * Transpose(B) )
```

so in `mmt4d` terms

```
mmt4d(B, A) == Transpose(mmt4d(A, B))
```

As `pack` and `unpack` already have enough generality to perform these
transpositions, we just directly generate the right transposing `pack`
and `unpack` ops. An earlier plan was to generate `linalg.transpose` and
rely on a later folding pattern, but it turned out to just be simpler to
directly generate the already-transposed `pack`, `unpack`.

A legitimate question was: should this transposition be implemented at
`SetEncoding` instead of at `MaterializeEncoding`? That would have been
simpler in some ways, but:
* The benefit of the transposition depends on the backend, so it doesn't
belong in Flow.
* SetEncoding must be reversible in case the back-end doesn't want to do
data-tiling. The transposition would be difficult to revert, and
generally confusing in settings where it's not wanted.
* The above mmt4d-specific trait simplifying the transposition only
helps since at MaterializeEncoding we know we are generating a mmt4d. We
couldn't so easily rely on that in SetEncoding.
* Before MaterializeEncoding we would have to handle `linalg.generic`,
not just named matmul ops.

Co-authored-by: Benoit Jacob <jacob.benoit.1@gmail.com>

benchmark-extra: x86_64-dt-only, android-cpu-dt-only

Signed-off-by: Alan Li <me@alanli.org>
Co-authored-by: Benoit Jacob <jacob.benoit.1@gmail.com>