Profiling JAX code and the mysterious redzone_checker kernel

[emprice]

Following the instructions in the JAX documentation, I successfully captured a profile of my long-running script in an attempt to find bottlenecks. To my surprise, a little over 90% of my runtime is (according to TensorBoard) consumed by the redzone_checker kernel. I found very little information about this kernel online, and, since it is consuming a lot of my runtime, I'd like to know

1. What is redzone_checker doing?
2. Is there a way to reduce the amount of time spent on this kernel?

[hawkinsp]

redzone_checker is a kernel XLA/GPU uses during convolution autotuning to check for bugs in CuDNN. The idea is that we allocate a slightly larger region than CuDNN requested, and then launch a GPU kernel that checks that nothing was written into the out-of-bounds region.

If it is taking a significant amount of time, that's something that we'd want to look at. But perhaps that means we are spending a lot of time compiling? It's hard to say without a small runnable reproduction, and perhaps details of your CUDA/CuDNN versions and GPU model.

[emprice]

Thanks for getting back to me! I suppose it could be due to a high compile time, but it would be nice to know for sure. I created this MWE gist (sorry, it is a little long, but these two functions are part of a larger code, which I've tried to minimize as best I can). Here is the kernel_stats view from TensorBoard:

Your explanation about redzone_checker makes sense, but I'm curious: Why would convolution autotuning be needed at all in a script that doesn't convolve anything? Is it being used behind the scenes somewhere?

Finally, the versions and hardware:

• jax: 0.2.10
• jaxlib: 0.1.61+cuda112
• CUDA: 11.2
• CuDNN: 8.1.1.33-1+cuda11.2
• GPU: NVIDIA Quadro RTX 5000

[emprice]

@hawkinsp Do you have any ideas for this, or is there other information I can provide that would be helpful?

[hawkinsp]

I'm actually wondering if this simply means that you are tracing the wrong part of the computation, e.g. you have sampled a period that mostly covers compile time, not run time. I'd look at the trace viewer part of the Tensorboard view to figure this out: is there an initial period full of things like the redzone checker, followed by what looks like meaningful computation? If you collect a longer trace, does it contain relatively less redzone checking?

Another suggestion: try running the computation twice and then tracing/profiling the second run. This should rule out one-off things that happen at compile time.

You might find it helpful to add:

cores[0].block_until_ready()

to the end of your script, incidentally.

One deficiency in the current jax.profiler interface is it is a sampling profiler only. It is a little hard to programmatically control when the tracing starts and stops around a particular set of code. I believe @skye is looking at fixing this.