[LinalgExt] Scatter fusion by expansion 3/3 #19588
Conversation
IanWood1
changed the title
IanWood1
force-pushed
the
scatter_reshape_fusion
branch
from
c6c862a to
7cd1fad
Compare
| @@ -391,6 +392,197 @@ static std::optional<SmallVector<Value>> fuseAttentionWithReshapeByExpansion( | |||
| return resultVals; | |||
| } | |||
|
|
|||
| namespace { | |||
| class ScatterExpansionInfo { | |||
There was a problem hiding this comment.
I am wondering if we need to replicate all this logic for each op. It seems with a cursory glance that the CollapseInfo (which can be reused for expansion as well) can just use methods defined by the FusionInterface? Then we should be able to just use that as a tool to get the reassociations we need. The actual math is fairly independent of the op?
There was a problem hiding this comment.
The problem I had with that is that the original operand has a nullptr AffineMap which breaks the logic of CollapseInfo. We need a way to represent the link between the innermost dims of original and updates while also expressing the fact that the outermost dims of original are not indexed in an AffineMap represent-able way.
I know CustomOp uses affine symbols which could possibly be useful here? I'm not to familiar with it so I'll need to check.
Signed-off-by: Ian Wood <[email protected]>
IanWood1
force-pushed
the
scatter_reshape_fusion
branch
from
7cd1fad to
8b870e6
Compare
Signed-off-by: Ian Wood <[email protected]>
Signed-off-by: Ian Wood <[email protected]>
Also, clean up some duplicated logic Signed-off-by: Ian Wood <[email protected]>
Implements fusion with reshapes by expansion for
LinalgExt::ScatterOp.TODO: fix >1 dynamic dim (also a problem with attention).
See main issue #19091