Add TestCompilation, a device-free CPU/GPU compilation checker

compilation_reports and @test_compilation analyze a call in four stages,
none of which need a GPU: JET over the CPU argument types (equivalent to
JET.@test_opt), JET over the CuArray-backed types seen in GPU CI jobs,
GPUCompiler IR validation plus JET over CUDA's device method table for
the kernel-side representation (derived with the same KernelAdaptor rules
as real launches), and a scan for host arrays that survive adaptation,
which would become illegal host pointers inside kernels. Includes unit
tests and a Buildkite step for them.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: dennisYatunin

.buildkite/pipeline.yml

@@ -101,6 +101,11 @@ steps:
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Utilities/unit_auto_broadcaster.jl"
 
+      - label: "Unit: test compilation checker"
+        key: unit_test_compilation
+        retry: *retry_policy
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Utilities/unit_test_compilation.jl"
+
   - group: "Unit: DataLayouts"
     steps:
 

test/TestUtilities/test_compilation.jl

@@ -0,0 +1,309 @@
+"""
+    TestCompilation
+
+Device-free compilation checking for CPU and GPU code paths. No GPU is needed
+for any check in this module (`CUDA.functional()` may be `false`), so tests
+built on it can guarantee GPU compilation without requesting CUDA devices.
+
+Given a call `f(args...)` with CPU (`Array`-backed) arguments, this module can
+run four analyses:
+
+ 1. `:cpu` — JET's optimization analysis of the call itself, equivalent to
+    `JET.@test_opt f(args...)`: reports every runtime dispatch or optimization
+    failure over the CPU argument types.
+ 2. `:host` — the same JET analysis over the argument types as they would
+    appear on a machine with a GPU: every `Array` becomes a `CuArray`, CPU
+    devices become `ClimaComms.CUDADevice`, and `DataLayouts` scopes are
+    recomputed. This is what `JET.@test_opt` sees in GPU CI jobs, and it
+    catches host-side instabilities in kernel launch code.
+ 3. `:kernel` — GPU device code analysis. The arguments are converted with the
+    same `Adapt`/`CUDA.KernelAdaptor` rules that real kernel launches use
+    (with `Array` leaves standing in for `CuArray`s), and the call is treated
+    as a kernel body. Two analyses run on it:
+      - GPUCompiler's LLVM IR validation, which catches `InvalidIRError`s
+        (dynamic dispatch, GPU-illegal operations) at the stage right before
+        the IR would be compiled to PTX; and
+      - JET's optimization analysis over CUDA's device method table, so that
+        device intrinsics like `threadIdx` are not treated as dead code.
+ 4. `:pointers` — a scan of the adapted arguments for host arrays that
+    survived adaptation. A field that remains an `Array` after the
+    `KernelAdaptor` runs corresponds to a host pointer inside a kernel
+    argument on a real GPU, which causes an illegal memory access at runtime
+    (this cannot be caught by compilation alone).
+
+The `:kernel` check relies on scope-based dispatch: converting a `DataLayout`
+to its kernel-side representation gives it the `ThisKernel` scope, so calling
+the same user-facing function on adapted arguments follows the device
+implementation. Host functions that are not scope-dispatched (e.g. ones that
+call CUDA APIs directly) should be checked with `stages = (:cpu, :host)`.
+
+# Usage
+
+    using .TestCompilation
+
+    # In a @testset (with args constructed on the CPU):
+    @test_compilation fill!(data, value)
+    @test_compilation stages = (:cpu, :host) column_integral_definite!(∫u, u)
+
+    # Programmatically:
+    ok, issues = compilation_reports(fill!, (data, value))
+    ok, issues = compilation_reports(f, args; stages = (:cpu,), ignored_modules = (...,))
+
+Each entry of `issues` is prefixed with the stage that produced it: `[cpu]`,
+`[host]`, `[kernel IR]`, `[kernel JET]`, or `[pointers]`. Keyword arguments
+other than `stages` are forwarded to `JET.report_opt` (e.g. `function_filter`
+and `ignored_modules`).
+"""
+module TestCompilation
+
+import Adapt
+import CUDA
+import ClimaComms
+import ClimaCore
+import ClimaCore: DataLayouts
+import JET
+import Test
+
+const CC = Core.Compiler
+const GC = CUDA.GPUCompiler
+
+export compilation_reports, @test_compilation
+
+# Frames that every ClimaCore JET test ignores: kernel launches unavoidably
+# pass through dynamic code in CUDA.jl and the CUDA extension (kernel caching,
+# argument conversion), and thread launches pass through dynamic error paths in
+# Base.Threads' task-spawning code. Functions parallelized over threads are
+# still fully analyzed, since inference also follows the branch that runs them
+# without spawning tasks (used for nested threaded loops).
+default_ignored_modules() = (
+    JET.AnyFrameModule(CUDA),
+    JET.AnyFrameModule(Base.get_extension(ClimaCore, :ClimaCoreCUDAExt)),
+    JET.AnyFrameModule(Base.Threads),
+)
+
+# =============================================================================
+# Argument conversion
+# =============================================================================
+
+# Stand-in for CuArray -> CuDeviceArray on the real launch path: a null-pointer
+# CuDeviceArray is a plain isbits struct, so it can be constructed without a
+# GPU, and going through the genuine KernelAdaptor applies every
+# package-specific adapt rule (e.g. grids, spaces, and limiter internals).
+struct KernelArrayStandIn end
+Adapt.adapt_storage(::KernelArrayStandIn, a::Array{T, N}) where {T, N} =
+    CUDA.CuDeviceArray{T, N, CUDA.AS.Global}(
+        reinterpret(Core.LLVMPtr{T, CUDA.AS.Global}, C_NULL),
+        size(a),
+    )
+Adapt.adapt_storage(to::KernelArrayStandIn, x) =
+    Adapt.adapt_storage(CUDA.KernelAdaptor(), x)
+
+kernel_arguments(args) =
+    map(args) do arg
+        standin = Adapt.adapt(KernelArrayStandIn(), arg)
+        Adapt.adapt(CUDA.KernelAdaptor(), standin)
+    end
+
+# Rewrite a host type to the type it would have on a machine with a GPU:
+# Arrays become CuArrays, CPU devices become CUDADevice, and DataLayouts scopes
+# are recomputed from the new parent array types.
+function host_gpu_type(@nospecialize(T))
+    T isa Type || return T
+    T <: ClimaComms.AbstractCPUDevice && return ClimaComms.CUDADevice
+    if T isa DataType && T <: Array
+        return CUDA.CuArray{eltype(T), ndims(T), CUDA.DeviceMemory}
+    elseif T isa DataType &&
+           T <: DataLayouts.DataLayout &&
+           length(T.parameters) >= 2
+        params = collect(T.parameters)
+        A = host_gpu_type(params[end])
+        S = typeof(DataLayouts.DataScope(A))
+        return T.name.wrapper{map(host_gpu_type, params[1:(end - 2)])..., S, A}
+    elseif T isa DataType && !isempty(T.parameters)
+        try
+            return T.name.wrapper{map(host_gpu_type, T.parameters)...}
+        catch
+            return T
+        end
+    else
+        return T
+    end
+end
+
+# =============================================================================
+# JET over the CUDA device method table (for kernel-side analysis)
+# =============================================================================
+
+# JET's default OptAnalyzer infers with the native method table, where CUDA
+# intrinsics like threadIdx resolve to host definitions that just throw, so
+# every kernel body looks like dead code. This report pass behaves identically
+# to JET's OptAnalysisPass but routes inference through CUDA's device method
+# table (and gets its own analysis cache, since JET caches per report pass).
+struct DeviceOptPass <: JET.ReportPass end
+(::DeviceOptPass)(T::Type{<:JET.InferenceErrorReport}, args...) =
+    JET.OptAnalysisPass()(T, args...)
+method_table_for(::JET.OptAnalysisPass, world::UInt) = CC.InternalMethodTable(world)
+method_table_for(::DeviceOptPass, world::UInt) =
+    GC.get_method_table_view(world, CUDA.method_table)
+CC.method_table(analyzer::JET.OptAnalyzer) =
+    method_table_for(JET.ReportPass(analyzer), JET.get_inference_world(analyzer))
+
+function jet_issues(prefix, f, tt; device = false, jetconfigs...)
+    result =
+        device ?
+        JET.report_opt(f, tt; report_pass = DeviceOptPass(), jetconfigs...) :
+        JET.report_opt(f, tt; jetconfigs...)
+    return map(JET.get_reports(result)) do report
+        message = sprint(JET.print_report, report)
+        location =
+            isempty(report.vst) ? "" :
+            string(" @ ", last(report.vst).file, ':', last(report.vst).line)
+        string(prefix, ' ', message, location)
+    end
+end
+
+# =============================================================================
+# GPUCompiler IR validation (for kernel-side analysis)
+# =============================================================================
+
+const PTX_CAP = v"7.0"
+const PTX_ISA = v"7.8"
+
+# Missing symbols that only resolve during a real kernel launch: libdevice
+# math functions, GPU runtime helpers, and the kernel state intrinsic.
+is_benign_ir_error(e) =
+    e[1] == GC.UNKNOWN_FUNCTION &&
+    e[3] isa AbstractString &&
+    (
+        startswith(e[3], "__nv") ||
+        startswith(e[3], "gpu_") ||
+        occursin("state_getter", e[3])
+    )
+
+function ir_issues(f, tt)
+    config = GC.CompilerConfig(
+        GC.PTXCompilerTarget(; cap = PTX_CAP, ptx = PTX_ISA),
+        CUDA.CUDACompilerParams(; cap = PTX_CAP, ptx = PTX_ISA);
+        kernel = true,
+        libraries = false,
+        always_inline = true,
+    )
+    try
+        job = GC.CompilerJob(GC.methodinstance(typeof(f), tt), config)
+        GC.JuliaContext() do _
+            GC.compile(:llvm, job)
+        end
+        return String[]
+    catch e
+        e isa GC.InvalidIRError || rethrow()
+        errors = filter(!is_benign_ir_error, e.errors)
+        return unique(map(errors) do (kind, _, meta)
+            string("[kernel IR] ", kind, meta isa Nothing ? "" : " [$meta]")
+        end)
+    end
+end
+
+# =============================================================================
+# Host pointer scan (for kernel arguments after adaptation)
+# =============================================================================
+
+function host_pointer_issues!(issues, x, path)
+    if x isa Array || x isa Ptr
+        push!(
+            issues,
+            "[pointers] $path is a host $(typeof(x).name.wrapper), \
+             which would cause an illegal memory access in a kernel",
+        )
+    elseif x isa CUDA.CuDeviceArray
+        return issues
+    elseif !isbits(x) && (isstructtype(typeof(x)) || x isa Tuple)
+        foreach(1:fieldcount(typeof(x))) do i
+            name = x isa Tuple ? "[$i]" : string('.', fieldname(typeof(x), i))
+            isdefined(x, i) &&
+                host_pointer_issues!(issues, getfield(x, i), path * name)
+        end
+    end
+    return issues
+end
+
+# =============================================================================
+# Combined checks
+# =============================================================================
+
+"""
+    compilation_reports(f, args::Tuple; stages, jetconfigs...)
+        -> (ok::Bool, issues::Vector{String})
+
+Run the [`TestCompilation`](@ref) analyses of `f(args...)` selected by
+`stages` (any subset of `(:cpu, :host, :kernel, :pointers)`; all by default).
+`ok` is `true` when every selected analysis finds no issues. All other keyword
+arguments are forwarded to `JET.report_opt`.
+"""
+function compilation_reports(
+    f,
+    args::Tuple;
+    stages = (:cpu, :host, :kernel, :pointers),
+    ignored_modules = default_ignored_modules(),
+    jetconfigs...,
+)
+    jetconfigs = (; ignored_modules, jetconfigs...)
+    issues = String[]
+    cpu_tt = Tuple{map(typeof, args)...}
+    :cpu in stages && append!(issues, jet_issues("[cpu]", f, cpu_tt; jetconfigs...))
+    :host in stages && append!(
+        issues,
+        jet_issues("[host]", f, host_gpu_type(cpu_tt); jetconfigs...),
+    )
+    if :kernel in stages || :pointers in stages
+        adapted = kernel_arguments(args)
+        if :pointers in stages
+            foreach(enumerate(adapted)) do (i, arg)
+                host_pointer_issues!(issues, arg, "args[$i]")
+            end
+        end
+        if :kernel in stages
+            kernel_f = (args...) -> (f(args...); nothing)
+            kernel_tt = Tuple{map(typeof, adapted)...}
+            append!(issues, ir_issues(kernel_f, kernel_tt))
+            append!(
+                issues,
+                jet_issues(
+                    "[kernel JET]",
+                    kernel_f,
+                    kernel_tt;
+                    device = true,
+                    jetconfigs...,
+                ),
+            )
+        end
+    end
+    return isempty(issues), issues
+end
+
+"""
+    @test_compilation [stages = ...] [kwarg = ...] f(args...)
+
+Assert (via `Test.@test`) that `f(args...)` passes the selected
+[`compilation_reports`](@ref) analyses. On failure, the issues are logged.
+
+    @test_compilation fill!(data, value)
+    @test_compilation stages = (:cpu, :host) column_integral_definite!(∫u, u)
+"""
+macro test_compilation(args...)
+    call = args[end]
+    kwargs = map(args[1:(end - 1)]) do kwarg
+        @assert Meta.isexpr(kwarg, :(=), 2) "expected keyword arguments before the call"
+        Expr(:kw, kwarg.args[1], esc(kwarg.args[2]))
+    end
+    @assert Meta.isexpr(call, :call) "expected a function call as the last argument"
+    f = esc(call.args[1])
+    call_args = map(esc, call.args[2:end])
+    quote
+        local ok, issues =
+            compilation_reports($f, ($(call_args...),); $(kwargs...))
+        ok || @info "Compilation issues for $($(string(call)))" issues
+        Test.@test ok
+    end
+end
+
+end # module TestCompilation