Name CUDA kernels based on stack trace in auto_launch

ext/cuda/cuda_utils.jl

@@ -2,8 +2,10 @@ import CUDA
 import ClimaCore.Fields
 import ClimaCore.DataLayouts
 import ClimaCore.DataLayouts: empty_kernel_stats
+import ClimaCore.DebugOnly: name_kernels_from_stack_trace
 
 const reported_stats = Dict()
+const kernel_names = Dict{String, AbstractString}()
 # Call via ClimaCore.DataLayouts.empty_kernel_stats()
 empty_kernel_stats(::ClimaComms.CUDADevice) = empty!(reported_stats)
 collect_kernel_stats() = false
@@ -39,19 +41,60 @@ 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 function and types of the args
+        key = string(objectid(f!))
+        kernel_name_exists = key in keys(kernel_names)
+        if !kernel_name_exists
+            # Construct the kernel name, ignoring modules we don't care about
+            uninteresting_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] ∉ uninteresting_modules
+                )
+            end
+            if !isnothing(first_relevant_index)
+                frame = stack[first_relevant_index]
+                name_str =
+                    string(frame.func) *
+                    "_" *
+                    string(frame.linfo.def.file) *
+                    "_" *
+                    string(frame.line)
+                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...)
+            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