feat: GPU-accelerated downconvert and correlate via KernelAbstractions.jl

.buildkite/pipeline.yml

@@ -2,27 +2,40 @@ env:
   SECRET_CODECOV_TOKEN: "Q3fuMdJjaQy9h/uk43rwSqz8M6ulvlCedU2Ir0S3QLP4t9F8cf7pzrTkX+nVhkGycZ/r5FRtTOwPr445R3wK5v9mEAsJN5GMOgI5w/L8m2XDwLmW3PN8RMno+fm2JVxZyPMNNmIQqbYEmmQcBS6Q3nywW3xi0Cl5umJuwDB+NdOFbpq3wc2wrnbOAbwlBJoCJmlH+F4ncuVY6EMmsgNKAf9RqUNWQxIthG616X1cNwuYEpL4dO/PWY2GMXWXTQ8ndO/713p4b5yIlzDP0mr2MrO+1A5fhgPc7Vr+f9mUlIAx+9AsWQYPrqPTkr2L5+mfaTodVE3u2Cop877WJZQD7w==;U2FsdGVkX1/wk2jzfWlRZ66IWgionQK/5Fu0pg3u0b26hhmmMjAjOklyi7QZKhJHjjt4KjK/dJzhd3eK28S0qQ=="
 
 steps:
-  - label: "Julia v1.10 - GPU Tests"
+  - label: "Julia v1.10 - KA + CUDA Tests"
     plugins:
       - JuliaCI/julia#v1:
           version: "1.10"
     commands:
-      - julia --project=ext/TrackingCUDAExt/test -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
-      - julia --project=ext/TrackingCUDAExt/test ext/TrackingCUDAExt/test/runtests.jl
+      - julia --project -e 'using Pkg; Pkg.instantiate(); Pkg.add("CUDA")'
+      - julia --project test/ka_gpu_tests.jl
     agents:
       queue: "juliagpu"
       cuda: "*"
     if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 60
-  - label: "Julia v1.11 - GPU Tests"
+  - label: "Julia latest - KA + CUDA Tests"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.11"
+          version: "1"
     commands:
-      - julia --project=ext/TrackingCUDAExt/test -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
-      - julia --project=ext/TrackingCUDAExt/test ext/TrackingCUDAExt/test/runtests.jl
+      - julia --project -e 'using Pkg; Pkg.instantiate(); Pkg.add("CUDA")'
+      - julia --project test/ka_gpu_tests.jl
     agents:
       queue: "juliagpu"
       cuda: "*"
     if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
+    timeout_in_minutes: 60
+  - label: "Benchmark: KA-CUDA vs CPU"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+    commands:
+      - julia --project=benchmark -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+      - JULIA_NUM_THREADS=auto julia --project=benchmark benchmark/run_multi_sat_benchmarks.jl 2>&1 | tee benchmark_results.txt
+      - BODY=$$(awk '/^## Multi-Satellite/,0' benchmark_results.txt); if [ -n "$$BODY" ]; then echo "$$BODY" | buildkite-agent annotate --style "info" --context "benchmarks"; fi
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 60

.gitignore

@@ -6,8 +6,10 @@ test/Manifest.toml
 ext/**/test/Manifest.toml
 
 docs/build/
+docs/plans/
 docs/pdf/build/
 docs/site/
 docs/Manifest.toml
+docs/plans
 
 .worktrees/

Project.toml

@@ -10,35 +10,39 @@ Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+FastSinCos = "eb71ef10-bbd6-487f-97cc-2efbd12ad3c4"
 GNSSSignals = "52c80523-2a4e-5c38-8979-05588f836870"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
+SIMD = "fdea26ae-647d-5447-a871-4b548cad5224"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TrackingLoopFilters = "0814aff9-93cb-554c-9fff-9bf1cfdb5efa"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
-
-[weakdeps]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-
-[extensions]
-TrackingCUDAExt = ["Adapt", "CUDA"]
+VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
 Accessors = "0.1"
 Acquisition = "0.1.1, 0.2, 0.3, 1"
-LinearAlgebra = "1"
-Adapt = "4.3"
 Bumper = "0.7.1"
-CUDA = "5.0"
 ConstructionBase = "1.6.0"
 Dictionaries = "0.4.1"
 DocStringExtensions = "0.6, 0.7, 0.8, 0.9"
+FastSinCos = "1"
 GNSSSignals = "0.17.3, 1"
-LoopVectorization = "0.8, 0.9, 0.10, 0.11, 0.12"
+KernelAbstractions = "0.9"
+LinearAlgebra = "1"
+SIMD = "3"
 StaticArrays = "0.9, 0.10, 0.11, 0.12, 1.0"
-StructArrays = "0.4, 0.6.5, 0.7"
 TrackingLoopFilters = "0.2, 1"
 Unitful = "0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18, 1.0"
+VectorizationBase = "0.21.72"
+AMDGPU = "1"
+CUDA = "5"
 julia = "1.10"
+
+[extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[targets]
+test = ["AMDGPU", "CUDA"]

ext/TrackingCUDAExt/test/Project.toml → benchmark/Project.toml

@@ -1,8 +1,7 @@
 [deps]
-Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GNSSSignals = "52c80523-2a4e-5c38-8979-05588f836870"
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracking = "10b2438b-ffd4-5096-aa58-44041d5c8f3b"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"

benchmark/bench_cpu.jl

@@ -0,0 +1,139 @@
+# CPU downconvert-and-correlate + track benchmarks
+
+using BenchmarkTools
+using GNSSSignals
+using Unitful: Hz
+using Tracking
+using Tracking:
+    CPUDownconvertAndCorrelator,
+    EarlyPromptLateCorrelator,
+    NumAnts,
+    SystemSatsState,
+    SatState,
+    TrackState,
+    downconvert_and_correlate,
+    get_correlator_sample_shifts,
+    get_code_type,
+    gen_code_replica!
+using StaticArrays
+
+# ── Helper: set up common benchmark state ──────────────────────────────────
+
+function setup_benchmark(;
+    signal_type = Float32,
+    num_samples = 2000,
+    sampling_frequency = 5e6Hz,
+    system = GPSL1(),
+    num_ants = 1,
+)
+    code_phase = 10.5
+    carrier_doppler = 1000.0Hz
+    code_doppler = carrier_doppler * GNSSSignals.get_code_center_frequency_ratio(system)
+    code_frequency = code_doppler + get_code_frequency(system)
+
+    correlator = EarlyPromptLateCorrelator(; num_ants = NumAnts(num_ants))
+    static_shifts = get_correlator_sample_shifts(correlator, sampling_frequency, code_frequency)
+    dynamic_shifts = collect(static_shifts)
+
+    signal = num_ants == 1 ?
+        rand(Complex{signal_type}, num_samples) :
+        rand(Complex{signal_type}, num_samples, num_ants)
+
+    code_replica = Vector{get_code_type(system)}(
+        undef, num_samples + maximum(static_shifts) - minimum(static_shifts),
+    )
+    gen_code_replica!(
+        code_replica, system, code_frequency, sampling_frequency,
+        code_phase, 1, num_samples, static_shifts, 1, Val(sampling_frequency),
+    )
+
+    return (;
+        correlator, signal, code_replica,
+        static_shifts, dynamic_shifts,
+        sampling_frequency, carrier_doppler, code_phase,
+        system, num_samples,
+    )
+end
+
+# ── High-level downconvert_and_correlate (full pipeline) ───────────────────
+
+function bench_downconvert_and_correlate(;
+    signal_type = Float32,
+    num_samples = 2000,
+    sampling_frequency = 5e6Hz,
+    system = GPSL1(),
+    num_ants = 1,
+)
+    downconvert_and_correlator = CPUDownconvertAndCorrelator(Val(sampling_frequency))
+    system_sats_state = SystemSatsState(
+        system,
+        [SatState(system, 1, 10.5, 1000.0Hz; num_ants = NumAnts(num_ants))],
+    )
+    track_state = TrackState((system_sats_state,))
+    signal = num_ants == 1 ?
+        rand(Complex{signal_type}, num_samples) :
+        rand(Complex{signal_type}, num_samples, num_ants)
+
+    @benchmarkable Tracking.downconvert_and_correlate(
+        $downconvert_and_correlator, $signal, $track_state, 1,
+        $sampling_frequency, $(0.0Hz),
+    )
+end
+
+# ── Fused kernel microbenchmarks ───────────────────────────────────────────
+
+function bench_fused_kernel(;
+    signal_type = Float32,
+    num_samples = 2000,
+    num_ants = 1,
+    shifts = :static,
+)
+    s = setup_benchmark(; signal_type, num_samples, num_ants)
+    sample_shifts = shifts == :static ? s.static_shifts : s.dynamic_shifts
+    # Warmup to trigger compilation
+    Tracking.downconvert_and_correlate_fused!(
+        s.correlator, s.signal, s.code_replica, sample_shifts,
+        s.carrier_doppler, s.sampling_frequency, 0.0, 1, s.num_samples,
+    )
+    @benchmarkable Tracking.downconvert_and_correlate_fused!(
+        $(s.correlator), $(s.signal), $(s.code_replica), $sample_shifts,
+        $(s.carrier_doppler), $(s.sampling_frequency), 0.0, 1, $(s.num_samples),
+    )
+end
+
+function cpu_suite()
+    suite = BenchmarkGroup()
+
+    # Full pipeline: CPU, various signal types
+    foreach((Int16, Int32, Float32, Float64)) do signal_type
+        suite["downconvert and correlate"]["CPU"][string(signal_type)] =
+            bench_downconvert_and_correlate(; signal_type)
+    end
+
+    # Full pipeline: multi-antenna
+    suite["downconvert and correlate"]["CPU"]["Float32 4ant"] =
+        bench_downconvert_and_correlate(; num_ants = 4)
+    suite["downconvert and correlate"]["CPU"]["Int16 4ant"] =
+        bench_downconvert_and_correlate(; signal_type = Int16, num_ants = 4)
+
+    # Full pipeline: track
+    system = GPSL1()
+    sampling_frequency = 5e6Hz
+    downconvert_and_correlator = CPUDownconvertAndCorrelator(Val(sampling_frequency))
+    track_state = TrackState(system, [SatState(system, 1, 0.0, 1000Hz)])
+    signal = rand(ComplexF32, 2000)
+    suite["track"]["Float32"] = @benchmarkable track(
+        $signal, $track_state, $sampling_frequency;
+        downconvert_and_correlator = $downconvert_and_correlator,
+    )
+
+    # Fused kernel microbenchmarks (only available on branches with the fused kernel)
+    if isdefined(Tracking, :downconvert_and_correlate_fused!)
+        suite["fused kernel"]["1-ant static taps"]  = bench_fused_kernel(; shifts = :static)
+        suite["fused kernel"]["1-ant dynamic taps"] = bench_fused_kernel(; shifts = :dynamic)
+        suite["fused kernel"]["4-ant static taps"]  = bench_fused_kernel(; num_ants = 4, shifts = :static)
+        suite["fused kernel"]["4-ant dynamic taps"] = bench_fused_kernel(; num_ants = 4, shifts = :dynamic)
+    end
+
+    suite
+end