search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.56k stars 569 forks source link

Missing cases where dimensions can be collapsed #17948

Closed IanWood1 closed 1 week ago

IanWood1 commented 1 month ago

There are cases where linalg.generic ops inside of a dispatch region could have their dimensions collapsed. But, collapsing the dimensions of the root elementwise op for a dispatch containing elementwise -> reduction -> elementwise -> reduction -> elementwise causes GPU compilation failures. Because of this, dims are not collapsed when the producer is a LinalgOp

Repro: Use this mlir and run iree-compile --output-format=vm-bytecode --compile-from=executable-sources ~/repro.mlir -o /tmp/a.vmfb

A bit more context here: https://github.com/iree-org/iree/issues/17594

IanWood1 commented 1 month ago

For more context, VectorReductionToGPU works by putting all ops in the function body into vector.warp_execute_on_lane_0 op and then sinking/hoisting out distributable ops. However, the element-wise op at the end cant be distributed (see https://gist.github.com/IanWood1/8f145cc0d36ea118e84775c663aa6775) which causes a crash in WarpOpScfForOp because it doesnt check to make sure the scf.for op has distributable types.

Executable source

https://gist.github.com/IanWood1/827aee042ff6968070d2d0c6beccae4a

Compile command

iree-compile --output-format=vm-bytecode --compile-from=executable-sources ~/repro.mlir -o /tmp/a.vmfb

IR dump befor/after all

https://gist.github.com/IanWood1/f4bcb96c1f044a026fb33322e7e3fa99

@bangtianliu @Groverkss would you be able to take a look?

Another issue:

Download https://gist.github.com/IanWood1/ab88d6e2b8a4f4a2a2ab68c9ecdf22b1 Checkout https://github.com/IanWood1/iree/tree/collapse_softmax (or remove reduction check in isEligibleForCollapse) Run iree-compile --output-format=vm-bytecode --mlir-print-op-on-diagnostic=false --iree-hal-target-backends=vulkan-spirv /home/ianwood/iree/tests/e2e/regression/softmax.mlir