Some matmul nvfuser_splitk benchmark fails

[xwang233]

on A100 and H100

root@65627d6877e4:/opt/pytorch/nvfuser# bin/nvfuser_bench --benchmark_filter='NvFuserScheduler_Matmul_Manual/nvfuser_splitk_NT/M:1024/N:2048/K:4096/warps:8/stages:3/splitk_factor:5/smem_epilogue:1/manual_time'
......
terminate called after throwing an instance of 'nvfuser::nvfError'
  what():  Fusion returns wrong results! The result tensor has shape [..., 1024,2048]. Mismatch happens at region result[...,526:783,0:328]
Exception raised from checkMatch at /opt/pytorch/nvfuser/benchmarks/cpp/matmul.cpp:104 (most recent call first):
frame #0: <unknown function> + 0x7e55f (0x5604dbe8b55f in bin/nvfuser_bench)
frame #1: <unknown function> + 0x6f653c (0x5604dc50353c in bin/nvfuser_bench)
frame #2: <unknown function> + 0x6fb670 (0x5604dc508670 in bin/nvfuser_bench)
frame #3: <unknown function> + 0x763ef8 (0x5604dc570ef8 in bin/nvfuser_bench)
frame #4: <unknown function> + 0x7434cb (0x5604dc5504cb in bin/nvfuser_bench)
frame #5: <unknown function> + 0x7450d5 (0x5604dc5520d5 in bin/nvfuser_bench)
frame #6: <unknown function> + 0x7456f9 (0x5604dc5526f9 in bin/nvfuser_bench)
frame #7: <unknown function> + 0x73482b (0x5604dc54182b in bin/nvfuser_bench)
frame #8: <unknown function> + 0x736c6e (0x5604dc543c6e in bin/nvfuser_bench)
frame #9: <unknown function> + 0x736f0e (0x5604dc543f0e in bin/nvfuser_bench)
frame #10: <unknown function> + 0x11acb4 (0x5604dbf27cb4 in bin/nvfuser_bench)
frame #11: <unknown function> + 0x29d90 (0x7f87a25d4d90 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #12: __libc_start_main + 0x80 (0x7f87a25d4e40 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #13: <unknown function> + 0x137355 (0x5604dbf44355 in bin/nvfuser_bench)

Aborted (core dumped)

several other nvfuser_bench tests fail with similar error messages

cc @naoyam @jacobhinkle

[jacobhinkle]

I noticed these and have been trying to repro. I suspect this is due to ATen matmul evaluation, which would cause half-precision reduction for split-K. The reduced precision could then cause the closeness check to fail.

[jacobhinkle]

It could also have to do with either #1534 or #1827, since those were both merged before 86c254d, which I think is where these benchmark errors were observed. #1978 should have prevented using ATen in the matmul benchmarks and that was merged before the others... Once I can locally repro then I should be able to get to the bottom of it.

[jacobhinkle]

Looks like accuracy dropped in #1534. I am not sure why yet.

[jacobhinkle]

I've looked in more detail at NvFuserScheduler_Matmul_Manual/nvfuser_splitk_NT/M:1024/N:2048/K:4096/warps:8/stages:3/splitk_factor:5/smem_epilogue:1/manual_time. The failure is random and changes run to run, even though all seeds are set, so it's probably a data race. Sometimes there are no inaccurate entries even in this problem with 2M output elements. Other times there are a hundred. I suspect this was triggered by the increased vectorization width from 2 to 4 that was achieved with #1534; however, I tried disabling vectorization though and I still hit the error. I suspect the predication is incorrect, but I need to look into it more. We could disable the benchmarks having smem_epilogue:1 temporarily, or we could revert #1534 until I get to the bottom of this.

[jacobhinkle]

Problem goes away when run in compute-sanitizer, because instrumentation slows down the kernel enough that the data race is no longer encountered. Scheduling params migth influence this too which could explain why we don't see it in the tests. It also does not appear when disabling index hoisting and expression simplification, but again this could just be hiding the problem by slowing down the kernel.

[jacobhinkle]

This is caused by #2000. We have a possible workaround in #2001.