From 3fe38b04b37d6f2377f061ade2b424f648ab6cf8 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 28 Oct 2026 13:58:52 +0100
Subject: [PATCH] Precompile common operations in the CUDA extension
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CUDACore's workload doesn't use the GPU, so the first reduction, scan or
predicate of a session compiles AcceleratedKernels' host code and
kernels. CUDACore couldn't cache them for the tunings set here anyway,
as this extension isn't loaded while CUDACore precompiles.

Run them in a workload of this extension when CUDA is functional while
precompiling, on arrays just large enough to select the kernels that a
vector of 10^6 elements and a 1000×1000 matrix do. Like CUDACore's own
workload, it needs Julia 1.12.

On an RTX 5080 with Julia 1.13, and CUDA.jl's KernelAbstractions methods
moved into CUDACore (JuliaGPU/CUDA.jl#3493), the first `sum` of 10^6
elements takes 6 ms instead of 1.1 s, `cumsum` 0.18 s instead of 1.6 s,
`any` and `all` 55 ms instead of 1.2 s, and `sum` along either dimension
of a 1000×1000 matrix 0.52 s instead of 1.6 s. Without that change,
loading CUDACore's KernelAbstractions extension invalidates part of this
code, and these take 0.6 to 1 s. The extension's image grows from 0.03
to 11 MB, and its precompilation from 1.6 to 5 s.
---
 Project.toml       |  2 ++
 ext/CUDACoreExt.jl | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/Project.toml b/Project.toml
index 1777ec22..52ff225d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,6 +9,7 @@ ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -33,5 +34,6 @@ GPUArraysCore = "0.2.1"
 KernelAbstractions = "0.9.43, 0.10"
 Metal = "1.10.1"
 OpenCL = "0.10, 1"
+PrecompileTools = "1"
 julia = "1.10"
 oneAPI = "1, 2"
diff --git a/ext/CUDACoreExt.jl b/ext/CUDACoreExt.jl
index b582963f..5bd5fdb5 100644
--- a/ext/CUDACoreExt.jl
+++ b/ext/CUDACoreExt.jl
@@ -2,6 +2,7 @@ module CUDACoreExt
 
 using CUDACore
 import AcceleratedKernels as AK
+using PrecompileTools: @setup_workload, @compile_workload
 
 # Groups of threads read runs of up to 4096 Float32s: on an RTX 5080, summing a 4096×4096 matrix
 # along `dims=1` takes 17% less time than with a block per output.
@@ -21,4 +22,31 @@ AK.sort_tuning(::CUDABackend, ::Type{T}) where {T} =
                   radix_min_slice_len=1 << 19) :
     AK.SortTuning(bitonic_max_len=512, radix_min_len=1 << 23, bitonic_max_slice_len=256)
 
+
+# CUDACore's own workload can't cache these operations: it doesn't use the GPU, and this
+# extension isn't loaded while CUDACore precompiles, so its kernels would use the default
+# tunings rather than the ones above. The arrays are just large enough to select the kernels
+# that a vector of 10^6 elements and a 1000×1000 matrix do. Without a functional GPU while
+# precompiling, these operations are only compiled when first used. Like CUDACore's own
+# workload, this needs 1.12: before, GPU code compiled while precompiling leaks into the host
+# code of the image.
+@static if VERSION >= v"1.12-"
+@setup_workload begin
+    if CUDACore.functional()
+        @compile_workload begin
+            m = CuArray(ones(Float32, 256, 32))
+            v = vec(m)
+            sum(v)
+            sum(m; dims=1)
+            sum(m; dims=2)
+            cumsum(v)
+            any(iszero, v)
+            all(iszero, v)
+            sum(CuArray(ones(Int, 8192)))
+            synchronize()
+        end
+    end
+end
+end
+
 end