iree icon indicating copy to clipboard operation
iree copied to clipboard

[Im2col] Add option to unroll decomposed im2col loops

Open Max191 opened this issue 1 year ago • 1 comments

This adds an option to the DecomposeIm2colPass to unroll the resulting loop nest of the decomposition, and sets it to true by default. This is an easier form to handle in most lowerings of the op. In particular, it is easier to generate better barrier placement on GPU lowerings, and unrolling the loops results in better overall barrier placement with the current pipeline.

Max191 avatar Aug 23 '24 21:08 Max191

I think this unrolling should happen later (after Vectorization or bufferization) because dropping the loop structure on tensors is always a little iffy IMO. Have you looked into options like that at all?

The main reason I unrolled during decomposition is because that is when we know the loops came from an im2col op. Vectorization relies on the decomposition, so it would need to unroll before vectorization (or otherwise annotate the loops to be unrolled later, which I don't like). I also noticed better barrier placement after bufferization, so it seems to be better right now to do this before bufferization (but not saying it needs to happen before bufferization, maybe there are better ways to achieve the same barrier placement).

Max191 avatar Aug 29 '24 15:08 Max191