Name CUDA kernels based on stack trace in auto_launch

ext/cuda/cuda_utils.jl

@@ -2,8 +2,11 @@ import CUDA
 import ClimaCore.Fields
 import ClimaCore.DataLayouts
 import ClimaCore.DataLayouts: empty_kernel_stats
+import ClimaCore.DebugOnly: name_kernels_from_stack_trace
+import CUDA.GPUCompiler: methodinstance
 
 const reported_stats = Dict()
+const kernel_names = IdDict()
 # Call via ClimaCore.DataLayouts.empty_kernel_stats()
 empty_kernel_stats(::ClimaComms.CUDADevice) = empty!(reported_stats)
 collect_kernel_stats() = false
@@ -39,19 +42,85 @@ function auto_launch!(
     always_inline = true,
     caller = :unknown,
 ) where {F!}
+    # If desired, compute a kernel name from the stack trace and store in
+    # a global Dict, which serves as an in memory cache
+    kernel_name = nothing
+    if name_kernels_from_stack_trace()
+        # Create a key from the method instance and types of the args
+        key = objectid(methodinstance(typeof(f!), typeof(args)))
+        kernel_name_exists = key in keys(kernel_names)
+        if !kernel_name_exists
+            # Construct the kernel name, ignoring modules we don't care about
+            ignore_modules = [
+                :Base,
+                :Core,
+                :GPUCompiler,
+                :CUDA,
+                :NVTX,
+                :ClimaCoreCUDAExt,
+                :ClimaCore,
+            ]
+            stack = stacktrace()
+            first_relevant_index = findfirst(stack) do frame
+                frame.linfo isa Core.MethodInstance && (
+                    fullname(frame.linfo.def.module)[1] ∉ ignore_modules
+                )
+            end
+            if !isnothing(first_relevant_index)
+                # Don't include file if this is inside an NVTX annotation
+                frame = stack[first_relevant_index]
+                func_name = string(frame.func)
+                if contains(func_name, "#")
+                    func_name = split(func_name, "#")[1]
+                end
+                file_path = frame.linfo.def.file
+                fp_split = split(string(file_path), "/")
+                if "NVTX" in fp_split
+                    fp_string = "_NVTX"
+                    line_string = ""
+                else
+                    # Trim base directory off of file path to shorten
+                    package_index = findfirst(fp_split) do part
+                        startswith(part, "Clima")
+                    end
+                    if isnothing(package_index)
+                        package_index = findfirst(p -> p == ".julia", fp_split)
+                    end
+                    if isnothing(package_index)
+                        package_index = findfirst(p -> p == "src", fp_split)
+                    end
+                    if isnothing(package_index)
+                        package_index = 1
+                    end
+                    fp_string =
+                        "_FILE_" *
+                        string(joinpath(fp_split[package_index:end]...))
+                    line_string = "_L" * string(frame.line)
+                end
+                name_str = string(func_name) * fp_string * line_string
+                kernel_name = replace(name_str, r"[^A-Za-z0-9]" => "_")
+            end
+            @debug "Using kernel name: $kernel_name"
+            kernel_names[key] = kernel_name
+        end
+        kernel_name = kernel_names[key]
+    end
+
     if auto
         @assert !isnothing(nitems)
         if nitems ≥ 0
-            kernel = CUDA.@cuda always_inline = true launch = false f!(args...)
+            # Note: `name = nothing` here will revert to default behavior
+            kernel = CUDA.@cuda name = kernel_name always_inline = true launch =
+                false f!(args...)
             config = CUDA.launch_configuration(kernel.fun)
             threads = min(nitems, config.threads)
             blocks = cld(nitems, threads)
             kernel(args...; threads, blocks) # This knows to use always_inline from above.
         end
     else
         kernel =
-            CUDA.@cuda always_inline = always_inline threads = threads_s blocks =
-                blocks_s f!(args...)
+            CUDA.@cuda name = kernel_name always_inline = always_inline threads =
+                threads_s blocks = blocks_s f!(args...)
     end
 
     if collect_kernel_stats() # only for development use

src/DebugOnly/DebugOnly.jl

@@ -171,4 +171,6 @@ function allow_mismatched_spaces_unsafe()
     return false
 end
 
+name_kernels_from_stack_trace() = false
+
 end