Launching a kernel with a new signature compiles a lot of signature-independent host code #550

Launching a kernel with a new argument signature compiles a lot of host code that doesn't depend on that signature. With pocl on an M1, launching an already-compiled kernel for a new element type costs ~75 ms (@opencl) to ~120 ms (KernelAbstractions) of host-side compilation, excluding GPUCompiler and the program build. JuliaGPU/Metal.jl#1014 fixed the same problems in Metal.jl (4-10x less host compilation per signature); this lists which of them also apply here.

Median host compile time per new signature, Julia 1.13.1 / 1.12.7, best of 3 runs. "Prototype" is the changes below, applied with @eval on top of 499e7f4:

The KA numbers include converting arguments of the new element type, which is compiled once per type rather than per kernel.

Suggested changes, largest effect first (line numbers at 0e4812c):

1. clfunction (src/compiler/execution.jl:244-274) is compiled per kernel, but only methodinstance(F, tt) and the HostKernel constructor depend on F/TT. The lock, config, compile_or_lookup and link_kernel can move into a @noinline function taking the MethodInstance and the keyword arguments. In addition, Julia doesn't specialize on source=f because it's a function-typed argument that is only passed through, so the keyword body is compiled twice per kernel and HostKernel{F,typeof(source),tt}(...) (:272) is a dynamic call; source::S=f ... where S avoids that. This alone saves ~20 ms per signature.
2. The KernelAbstractions launch (obj::KA.Kernel{OpenCLBackend})(args...; ...) (src/OpenCLKernels.jl:132) is compiled for Vararg{Any}, by the same heuristic, so every call in it (@opencl launch=false, the launch itself, ...) is dispatched at run time and compiled separately. args::Vararg{Any,N} ... where N fixes that. The parts that don't depend on the arguments (launch_config, mkcontext, workgroup size tuning) can go into @noinline functions as well.
3. The launch chain: the generated AbstractKernel call (src/compiler/execution.jl:182-226) → launch_with_exception_mailbox (src/compiler/exceptions.jl:176-199) → cl.call (lib/cl/kernel.jl:263-328) → set_args! (lib/cl/kernel.jl:418), each a keyword call fully specialized on the arguments, while most of their bodies don't depend on them:
   • the mailbox lookup and KernelState/launch ID bookkeeping can be @noinline functions that don't take the kernel arguments;
   • in cl.call, only setting the arguments needs to be specialized; the indirect memory handling, clSetKernelExecInfo and the enqueue can be a @noinline function of the kernel, the number of arguments and the keyword arguments;
   • set_args! iterates enumerate over a heterogeneous tuple, which is type unstable; unrolling it avoids that.

After launching 4 signatures, there are 19 specializations of clfunction's keyword body, 13 of launch_with_exception_mailbox's, 14 of cl.set_args! and 8 of the KA launch.

Two problems found in Metal.jl don't apply here: no closure on the launch path captures the kernel or its arguments, and arguments are rooted with GC.@preserve rather than collected in a Vector{Any}.

Note that a reproducer needs to compile the launch inside the timed region, i.e., not let the timing function specialize on (and thus infer) the closure it times:

using pocl_jll, OpenCL, KernelAbstractions, GPUCompiler, Printf

# time spent compiling kernels, to subtract from Julia's compile time
const excluded = Ref(0)
@eval function OpenCL.compile_to_obj(job::OpenCL.OpenCLCompilerJob)
    t0 = Base.cumulative_compile_time_ns()[1]
    try invoke(OpenCL.compile_to_obj, Tuple{GPUCompiler.CompilerJob}, job)
    finally excluded[] += Base.cumulative_compile_time_ns()[1] - t0 end
end
@eval function OpenCL.link_kernel(job::OpenCL.OpenCLCompilerJob, obj::Vector{UInt8}, entry::String)
    t0 = Base.cumulative_compile_time_ns()[1]
    try invoke(OpenCL.link_kernel, Tuple{GPUCompiler.CompilerJob, Vector{UInt8}, String}, job, obj, entry)
    finally excluded[] += Base.cumulative_compile_time_ns()[1] - t0 end
end

function k_native(ys, xs, a)
    i = get_global_id()
    @inbounds ys[i] = xs[i] * a
    return
end
@kernel function k_ka(ys, @Const(xs), a)
    i = @index(Global)
    @inbounds ys[i] = xs[i] * a
end
run_native(ys, xs, a) = @opencl global_size=64 k_native(ys, xs, a)
run_ka(ys, xs, a) = k_ka(OpenCLBackend())(ys, xs, a; ndrange=64)

function measure(label, @nospecialize(f))
    Base.cumulative_compile_timing(true)
    excluded[] = 0
    t0 = Base.cumulative_compile_time_ns()[1]
    f()
    host = (Base.cumulative_compile_time_ns()[1] - t0 - excluded[]) / 1e9
    @printf("%-20s host compile %5.3f s\n", label, host)
end

mk(T) = (CLArray(zeros(T, 64)), CLArray(ones(T, 64)), T(2))
let (ys, xs, a) = mk(Float32)
    run_ka(ys, xs, a); run_native(ys, xs, a); cl.finish(cl.queue())
end
for T in (Int32, Int16, Float16, UInt8)
    ys, xs, a = mk(T)
    # KA first, so that @opencl doesn't get the per-element-type work
    measure("KA $T", () -> (run_ka(ys, xs, a); cl.finish(cl.queue())))
    measure("@opencl $T", () -> (run_native(ys, xs, a); cl.finish(cl.queue())))
end