commit 891f438c27f201ae0dde3e8408eefc6fca80bc5f
author Max191 <44243577+Max191@users.noreply.github.com> Fri Sep 20 13:51:09 2024 -0700
committer GitHub <noreply@github.com> Fri Sep 20 16:51:09 2024 -0400
tree a8b2faf8432edaf3fb203839946358216e7f92b3
parent d834aa7357179e0d806f3634d2efe3af2fa45171

[Codegen] Add control options in pack unpack decomposition (#18469)

Some early patterns in pack/unpack decomposition avoid generating
reshape ops for unit dim packs and unpacks. However, the TileAndFuse
pipeline uses these reshape ops to propagate expanded shapes to other
fusable ops. This PR adds an option to the DecomposePackUnPackOps pass
to create reshape ops anyway for unit dim cases.

The reason these unit dims show up right now is that the
iree_linalg_ext.im2col op of a unit-batched conv will have a unit
dimension in the batch dim. Ultimately, it would be good to allow for
batchless im2col ops, but in general it is good to support ops that have
required unit dimensions. When prototyping new ops, it can be easiest to
not support rank-reducing cases at first (winograd ops are another
example), so these unit dims may appear again in the future.

This PR also adds an optional control function to the pass options,
which controls which packs and unpacks get decomposed. The control
function is currently expected to be used when the `useOnlyReshapes`
option is true, since there is no control function in some upstream
patterns yet, but adding the control function upstream and fixing this
is left as a TODO.

---------

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>