[LinalgExt] Update scatter to allow dropping unit dims #19704
Conversation
IanWood1
force-pushed
the
remove_scatter_unit_index_depth
branch
from
92776de to
dc05f7a
Compare
compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/convert_to_loops.mlir
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
Outdated
Show resolved
Hide resolved
| const auto originalSliceRank = originalType.getRank() - indexDepth; | ||
| if (originalSliceRank < 0) { | ||
| return op->emitOpError( | ||
| "expected original rank to be greater or equal to index depth"); |
There was a problem hiding this comment.
Maybe this error should be more explicit in terms of slize rank being >= 0.
compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Looks ok overall. Is there a sequencing thing that you dropped the implementation of the DropScatter*Unit* ? Your message says so above, so assuming its just a sequencing thing. If you do want to drop it, just drop it from the pattern list completely and add it back later instead of making the implementation dummy.
Signed-off-by: Ian Wood <[email protected]>
also change expected result for e2e correctness (maybe delete?) Signed-off-by: Ian Wood <[email protected]>
Signed-off-by: Ian Wood <[email protected]>
Signed-off-by: Ian Wood <[email protected]>
Signed-off-by: Ian Wood <[email protected]>
Yeah sorry that was confusing. The pattern I added and the one I deleted are 2 separate ones (the pre-existing one is a noop because there can no longer be unit dims between |
Signed-off-by: Ian Wood <[email protected]>
|
Running onnx tests https://github.com/nod-ai/SHARK-TestSuite/actions/runs/12816694583 |
Signed-off-by: Ian Wood <[email protected]>
Changes
scatterto enforce thatupdatesis only batch dims + contiguous slice. This allows us to drop fold the unit dimension fromindiceswhen index_depth is 1.