mlir: [tensor dialect] moved unrolling to after bufferization

[liamsemeria]

Motivation

#138 showed that the MLIR backend had excessive compile times due to unrolling causing the bufferization passes to take a long time.

Description

Unrolling was moved after bufferization, which was pretty straight forward due to there already being a post-bufferization transform stage.
Moving the unrolling fixes the compile time issue and bumps performance up to tvm.

with tvm:

time loop-explore  --threads 1 --peak-flops 65.6e9 --operator conv2d --backends tvm --use-tensor --strategy tile7d --test 1 1 1 2 14 1 2 1 28 2 1 16 7 1 1 --save-temps
Schedule: tvm: [1; 1; 1; 2; 14; 1; 2; 1; 28; 2; 1; 16; 7; 1; 1]: time: 1.99 msecs, peak perf: 90.30%

real	0m4,571s
user	0m8,219s
sys	0m0,462s

with unrolling before bufferization:

time loop-explore  --threads 1 --peak-flops 65.6e9 --operator conv2d --backends mlir --use-tensor --strategy tile7d --test 1 1 1 2 14 1 2 1 28 2 1 16 7 1 1 --save-temps
Schedule: mlir: [1; 1; 1; 2; 14; 1; 2; 1; 28; 2; 1; 16; 7; 1; 1]: time: 2.42 msecs, peak perf: 74.47%

real	0m49,846s
user	0m53,662s
sys	0m0,435s

with unrolling after bufferization:

time loop-explore  --threads 1 --peak-flops 65.6e9 --operator conv2d --backends mlir --use-tensor --strategy tile7d --test 1 1 1 2 14 1 2 1 28 2 1 16 7 1 1 --save-temps
Schedule: mlir: [1; 1; 1; 2; 14; 1; 2; 1; 28; 2; 1; 16; 7; 1; 1]: time: 2.03 msecs, peak perf: 88.61%

real	0m3,465s
user	0m7,323s
sys	0m0,294s

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[guillon]

Great!