GEMM performance regressed after sinking unused l-invariant loads

[yzhang93]

I was looking into some performance regression cases in our downstream application and found #157559 changed the generated llvm codes and thus affect the performance.

I'm attaching the before and after .ll file for reference.
https://gist.github.com/yzhang93/f7174c373ec63845d63762997b3ee293

[VigneshwarJ]

Hey @yzhang93 can I get the input.ll prior to running opt?

[yzhang93]

Hey @yzhang93 can I get the input.ll prior to running opt?

Thanks for taking a look! Here is the input https://gist.github.com/yzhang93/f7174c373ec63845d63762997b3ee293#file-input-ll.

[VigneshwarJ]

Hey @yzhang93 can I get the input.ll prior to running opt?

Thanks for taking a look! Here is the input https://gist.github.com/yzhang93/f7174c373ec63845d63762997b3ee293#file-input-ll.

Hey @yzhang93, is there a reason why the sched.barrier is called before the store? If you move the sched.barrier below the stores in bb 56, that would give a better context for the scheduler to schedule those instructions. Could you try this and let me know? Thanks.

The LICM change sinks some load instructions from the preheader to the exit block, which is what gets unrolled here. The sunk loads brings the vmcnt closer causing the issue.

[jerryyin]

@VigneshwarJ For context of the scheduling barrier placement please refer to iree-org/iree#21151.

Back then, amdgpu backend tends to place s_waitcnt too early, sometimes as early as mfma instructions. This has caused a lot of unnecessary stalls. By placing a sched.barrier, we forced backend to place the barrier after compute is done. This has turns out to benefit a fair amount of gemm/convolutions

I assume if we move the sched.barrier down, then the backend can put the s_waitcnt too early again. So to answer your question, we may or may not be able to push down or even remove the sched.barrier, this all depends on how much of the configs we track will improve performance. Meanwhile I don't have enough of the llvm context to understand what your PR does.

The LICM change sinks some load instructions from the preheader to the exit block

I'm not exactly sure what this means? Does it mean the global load will be scheduled later - to where sched.barrier placed? I wonder what's the motivation of this? (To reduce stall or to reduce dispatch latency?) To us, we have software pipeline in place so want the global load of next iteration to be scheduled as early as possible so this can interleave with compute of the current iteration. Would your PR disrupt this? Let me if my understanding is correct.

CC: @krzysz00

[krzysz00]

So if I understand what the PR does, it takes

float x;
for (...) {
  x = *somewhere;
  ... (// doesn't use x)
}
op(x, probably some loop results);

and transforms it into

for (...) {
  ...
}
float x = *somewhere;
op(x, probably some loop results);

where previously you'd go to

float x = *somewhere;
for (...) {
  ...
}
op(x, probably some loop results);

I can see a world in which this'll be bad for latency on, say, a accumulating GEMM (because loading the old value of C can be hidden by the loop) but also where you'd want this optimization (because that value from C has to hog perfectly good registers throughout the entire loop).

I think we'll want to analyze the situation more carefully to see how impactful this is etc.

[VigneshwarJ]

Yes, exactly as Krzysztof says.
Here, the float x = *somewhere; is a buffer load
@jerryyin that makes sense why sched.barriers got added. Then my suggestion might not help.
Prior to my change it was something like this, where bufferload had more instructions so vmcnt before use waited less.

x = buffer load
workgroup barrier
for (...) {
...
}
..
mfmas
sched.barrier
waitcnt \\ till x is ready
lds write x

after my change bufferload is closer to waitcnt.

workgroup barrier
for (...) {
...
}
..
x = buffer load
mfmas
sched.barrier
waitcnt \\ till x is ready
lds write x

[jerryyin]

@krzysz00 @VigneshwarJ Excellent explanation.

@yzhang93 Could you help clarify how much of regression do we see, and on how many configs we tracked? I remember you mention this happen in backward weight convolution only. Did you see buffer load before loop situation happen in those backward weight kernels?

[krzysz00]

Higher-level question: why didn't this get LICM'd at the MLIR level?

[krzysz00]

(... which also leaves the question of whether MLIR's LICM should be fiddled with to do the same thing)

I'm also noticing that, with the example, it looks like that for loop isn't the main loop of the GEMM but some other smaller loop. Is that loop perhaps a good unrolling candidate?

[yzhang93]

@yzhang93 Could you help clarify how much of regression do we see, and on how many configs we tracked? I remember you mention this happen in backward weight convolution only. Did you see buffer load before loop situation happen in those backward weight kernels?

There are 28 out of 252 shapes regressed. The performance degraded 10%~40%. To give you a brief idea, the shape in the example https://gist.github.com/yzhang93/f7174c373ec63845d63762997b3ee293 regressed from 391us to 530us.