Name CUDA kernels based on stack trace in auto_launch

[petebachant]

This PR adds stack trace-based CUDA kernel naming to make performance benchmarks more informative. Kernel names are computed once per session and cached in a module-level Dict.

• Code follows the style guidelines OR N/A.
• Unit tests are included OR N/A.
• Code is exercised in an integration test OR N/A.
• Documentation has been added/updated OR N/A.

TODO:

• Confirm we get useful results, i.e., that we don't spend too much time computing the kernel name and wash out the timeline.

Example results from ClimaAtmos Buildkite

From here

[petebachant]

Added some example results, which look suspiciously different. Reminds me of Simon's comment about the CUDA API not doing much in the previous implementation that computed the name every time a kernel was created. Do we need to somehow hydrate the kernel name cache initially to get a useful result here?

[imreddyTeja]

ook suspiciously different. Reminds me of Simon's comment about the CUDA API not doing much in the previous implementation that computed the name every time a kernel was created. Do we need to somehow hydrate the kernel name cache initially to get a useful result here?

Are these profile results from the first step, or a later step? The kernel renaming example takes significantly longer, which makes me think stacktrace is getting called.

[petebachant]

I think I figured this out. The compile step was running before kernel renaming was enabled. Nope, still taking way too long when this feature is enabled.

[imreddyTeja]

I just tried using this in ClimaLand, and it worked. This is so much easier than dev`ing GPUCompiler.jl

[petebachant]

The files don't make much sense in ClimaAtmos's buildkite runs since I guess they handle dev installs differently:

Maybe I should just split on src as a last resort.

[petebachant]

Breaking on src if we can't find any other tokens:

[imreddyTeja]

Suggested change

	key = objectid(methodinstance(typeof(f!), typeof(args)))
	key = objectid(methodinstance(F!, typeof(args)))

The function signature could also be changed to

function auto_launch!(
    f!::F!,
    args::ARGS,
    nitems::Union{Integer, Nothing} = nothing;
    auto = false,
    threads_s = nothing,
    blocks_s = nothing,
    always_inline = true,
    caller = :unknown,
) where {F!, ARGS}
.
.
.

 key = objectid(methodinstance(F!, ARGS))

but I think that forced specialization of the method (I'm not sure if that is desirable or not here)

[dennisYatunin]

One hangup I have is the redefine-a-function-to-activate way of using this versus an environmental variable. The latter seems nicer since you can do it without changing your code (driver script?), but the former seems to be the typical pattern. What is the upside of the redefined function approach?

Both of these approaches—redefining a method and changing an environment variables—involve recompiling ClimaCore's CUDA extension to activate this feature, since the feature is disabled in the default version of ClimaCore we precompile for the CI depot. The difference is in how this recompilation is achieved:

• When a method is redefined, all other methods that depend on it (either directly or indirectly) are immediately recompiled. The compiler achieves this by tracing through the "backedges" that connect different MethodInstance objects, starting from the MethodInstance that was redefined.
• When an environment variable is changed, Julia does not have any mechanism for automatically detecting that change. So, if ClimaCore has a constant flag determined from an environment variable, the variable needs to be changed before ClimaCore is compiled in the current Julia session. Because the compiler is not able to redefine constant variables without starting a new Julia session, the flag will always remain unchanged after ClimaCore is loaded for the first time.

Since starting a new Julia session requires recompiling the Base library, the environment variable approach can end up needing more recompilation than the other approach. So, in the interest of minimizing total compilation time, we typically use method redefinitions to activate optional package features.

On the other hand, if your setup requires you to start a new Julia session anyway, it doesn't really matter which option you choose. Also, the environment variable approach I described assumes that your flag needs to be a compile-time constant; if you're okay with performing an environment lookup every time you access the flag, then the environment variable approach will probably be simpler to implement.

[petebachant]

Thanks @dennisYatunin. What do you think of the current implementation with the const? Is that acceptable for now?

[petebachant]

I was finally able to get the JET tests to pass while retaining the ability to set this via environmental variable. It's not pretty, but it saves a bit of code downstream...

[imreddyTeja]

I was finally able to get the JET tests to pass while retaining the ability to set this via environmental variable. It's not pretty, but it saves a bit of code downstream...

Does this allow a user to turn renaming on/off while a repl is active, and without restarting it?

[petebachant]

Does this allow a user to turn renaming on/off while a repl is active, and without restarting it?

Yep, but you'll need to load the extension and redefine the method if you're in an active session:

ext = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt)
ext.name_kernels_from_stack_trace() = false

[imreddyTeja]

We should make documentation/a guide on how to use this at some point. As this is a developer tool, it's probably fine to merge without docs. Could you please squash the commits before merging?