Page MenuHomePhabricator

[mlir][linalg] Restrict distribution to parallel dims
ClosedPublic

Authored by antiagainst on May 7 2021, 9:30 AM.

Details

Summary

According to the API contract, LinalgLoopDistributionOptions
expects to work on parallel iterators. When getting processor
information, only loop ranges for parallel dimensions should
be fed in. But right now after generating scf.for loop nests,
we feed in *all* loops, including the ones materialized for
reduction iterators. This can cause unexpected distribution
of reduction dimensions. This commit fixes it.

Diff Detail

Event Timeline

antiagainst created this revision.May 7 2021, 9:30 AM
antiagainst requested review of this revision.May 7 2021, 9:30 AM
mravishankar requested changes to this revision.May 7 2021, 1:49 PM

I am not sure how this would interfact with the case where some loops are not tiled. Would be good to try out examples from here where some loops are not generated as scf.parallel when the tile size is set to 0. In the same way, if the tile size is set to 0, then then loop wont be tiled and therefore not distributed.

Essentially trying to enforce the same contract of distribution that happens with scf.parallel. The proc_id are used for every scf.parallel encountered in the generated tiled loop nest. So if the tile size was set to 0, since it wont be lowered to an scf.parallel it wont be distributed. That does not seem to be the same contract above.

This revision now requires changes to proceed.May 7 2021, 1:49 PM
mravishankar accepted this revision.May 7 2021, 1:52 PM

Actually, I take it back. I think this makes sense. Since the tile size is set to 0, the loop wont be generated. Then the non-parallel loops are filtered out.

This revision is now accepted and ready to land.May 7 2021, 1:52 PM
This revision was automatically updated to reflect the committed changes.