Make CUDA a weak dependency using package extensions

.buildkite/pipeline.yml

@@ -2,25 +2,25 @@ 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"
+  - label: "Julia v1.10 - GPU Tests"
     plugins:
       - JuliaCI/julia#v1:
           version: "1.10"
-      - JuliaCI/julia-test#v1: ~
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
+    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
     agents:
       queue: "juliagpu"
       cuda: "*"
     if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 60
-  - label: "Julia v1.11"
+  - label: "Julia v1.11 - GPU Tests"
     plugins:
       - JuliaCI/julia#v1:
           version: "1.11"
-      - JuliaCI/julia-test#v1: ~
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
+    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
     agents:
       queue: "juliagpu"
       cuda: "*"

.gitignore

@@ -3,8 +3,11 @@
 *.jl.mem
 /Manifest.toml
 test/Manifest.toml
+ext/**/test/Manifest.toml
 
 docs/build/
 docs/pdf/build/
 docs/site/
 docs/Manifest.toml
+
+.worktrees/

Project.toml

@@ -8,7 +8,6 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Acquisition = "d4bbf60b-102b-5ffb-8f97-a7ea5817e69f"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -20,6 +19,12 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TrackingLoopFilters = "0814aff9-93cb-554c-9fff-9bf1cfdb5efa"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+TrackingCUDAExt = "CUDA"
+
 [compat]
 Accessors = "0.1"
 Acquisition = "0.1.1"

README.md

@@ -24,12 +24,25 @@ Moreover, Tracking.jl allows tracking of signals from phased antenna arrays mean
 * Secondary code detection
 * Bit detection
 * Phased array tracking
-* GPU acceleration (CUDA)
+* GPU acceleration (requires `using CUDA`)
 
 ## Getting started
 
 Install:
 ```julia
 julia> ]
 pkg> add Tracking
+```
+
+### GPU Support
+
+To use GPU acceleration features, you need to explicitly load CUDA:
+
+```julia
+using Tracking
+using CUDA  # Activates GPU functionality
+
+# Access GPU types via the extension
+ext = Base.get_extension(Tracking, :TrackingCUDAExt)
+gpu_correlator = ext.GPUDownconvertAndCorrelator(...)
 ```

ext/TrackingCUDAExt/TrackingCUDAExt.jl

@@ -0,0 +1,28 @@
+module TrackingCUDAExt
+
+using Tracking
+using CUDA
+
+# Import dependencies needed by GPU code
+using GNSSSignals
+using StaticArrays
+using Dictionaries
+using Accessors
+using DocStringExtensions
+using Unitful: Hz
+import Adapt
+
+# Import functions from Tracking that we need to extend or use
+import Tracking: downconvert_and_correlate, downconvert_and_correlate!, get_num_samples, update
+
+# Include the GPU implementation
+include("downconvert_and_correlate_gpu.jl")
+
+# Export GPU types from the extension
+# Users access these via: Base.get_extension(Tracking, :TrackingCUDAExt).GPUDownconvertAndCorrelator
+export GPUSatDownconvertAndCorrelator,
+       GPUSystemDownconvertAndCorrelator,
+       GPUDownconvertAndCorrelator,
+       convert_code_to_texture_memory
+
+end

ext/TrackingCUDAExt/test/Project.toml

@@ -0,0 +1,8 @@
+[deps]
+Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
+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"

ext/TrackingCUDAExt/test/runtests.jl

@@ -0,0 +1,226 @@
+module TrackingCUDAExtTests
+
+using Test: @test, @testset, @inferred
+using Unitful: Hz
+using Tracking
+using CUDA: CUDA, cu
+using GNSSSignals: GPSL1, gen_code, get_code_frequency, get_code_center_frequency_ratio, get_code
+using Pkg
+using Bumper: SlabBuffer
+using Tracking:
+    CPUDownconvertAndCorrelator,
+    SystemSatsState,
+    SatState,
+    TrackState,
+    downconvert_and_correlate,
+    get_last_fully_integrated_correlator,
+    track,
+    get_code_phase,
+    get_carrier_phase,
+    get_carrier_doppler,
+    get_code_doppler,
+    get_last_fully_integrated_filtered_prompt,
+    NumAnts,
+    get_default_correlator
+
+# Access extension types via Base.get_extension
+const TrackingCUDAExt = Base.get_extension(Tracking, :TrackingCUDAExt)
+const GPUDownconvertAndCorrelator = TrackingCUDAExt.GPUDownconvertAndCorrelator
+
+# Helper function to check CUDA.jl version
+function is_cuda_below_version(version_string::String)
+    cuda_version = Pkg.installed()["CUDA"]
+    return cuda_version < VersionNumber(version_string)
+end
+
+@testset "GPU Downconvert and Correlate" begin
+    !CUDA.functional() && return
+
+    gpsl1 = GPSL1()
+    sampling_frequency = 5e6Hz
+    code_phase = 10.5
+    num_samples_signal = 5000
+    intermediate_frequency = 0.0Hz
+
+    system_sats_state = SystemSatsState(
+        gpsl1,
+        [SatState(gpsl1, 1, code_phase, 1000.0Hz), SatState(gpsl1, 2, 11.0, 500.0Hz)];
+    )
+    multiple_system_sats_state = (system_sats_state,)
+
+    downconvert_and_correlator =
+        GPUDownconvertAndCorrelator(multiple_system_sats_state, num_samples_signal)
+
+    track_state = TrackState(multiple_system_sats_state)
+
+    preferred_num_code_blocks_to_integrate = 1
+
+    signal = cu(
+        gen_code(
+            num_samples_signal,
+            gpsl1,
+            1,
+            sampling_frequency,
+            get_code_frequency(gpsl1) + 1000Hz * get_code_center_frequency_ratio(gpsl1),
+            code_phase,
+        ) .* cis.(2π * (0:(num_samples_signal-1)) * 1000.0Hz / sampling_frequency),
+    )
+
+    next_track_state = @inferred downconvert_and_correlate(
+        downconvert_and_correlator,
+        signal,
+        track_state,
+        preferred_num_code_blocks_to_integrate,
+        sampling_frequency,
+        intermediate_frequency,
+    )
+
+    # GPU uses floating point arithmetic and might differ a little with the fixed point arithmetic
+    if is_cuda_below_version("5.8.0")
+        @test real.(
+            get_last_fully_integrated_correlator(next_track_state, 1).accumulators
+        ) ≈ [2921, 4949, 2921]
+    else
+        # Unfortuntely, the texture index rounding is off with CUDA 5.9 TODO: report
+        accumulator =
+            real.(
+                get_last_fully_integrated_correlator(next_track_state, 1).accumulators
+            )
+        @test 2900 < accumulator[1] < 3000
+        @test 4900 < accumulator[2] < 5000
+        @test 2900 < accumulator[3] < 3000
+    end
+
+    signal = cu(
+        gen_code(
+            num_samples_signal,
+            gpsl1,
+            2,
+            sampling_frequency,
+            get_code_frequency(gpsl1) + 500Hz * get_code_center_frequency_ratio(gpsl1),
+            11.0,
+        ) .* cis.(2π * (0:(num_samples_signal-1)) * 500.0Hz / sampling_frequency),
+    )
+
+    next_track_state = @inferred downconvert_and_correlate(
+        downconvert_and_correlator,
+        signal,
+        track_state,
+        preferred_num_code_blocks_to_integrate,
+        sampling_frequency,
+        intermediate_frequency,
+    )
+
+    # GPU uses floating point arithmetic and might differ a little with the fixed point arithmetic
+    if is_cuda_below_version("5.8.0")
+        @test real.(
+            get_last_fully_integrated_correlator(next_track_state, 2).accumulators
+        ) ≈ [2919, 4947, 2919]
+    else
+        # Unfortuntely, the texture index rounding is off with CUDA 5.9 TODO: report
+        accumulator =
+            real.(
+                get_last_fully_integrated_correlator(next_track_state, 2).accumulators
+            )
+        @test 2900 < accumulator[1] < 3000
+        @test 4900 < accumulator[2] < 5000
+        @test 2900 < accumulator[3] < 3000
+    end
+end
+
+@testset "Track multiple signals with GPU" begin
+    !CUDA.functional() && return
+    gpsl1 = GPSL1()
+    carrier_doppler = 200Hz
+    start_code_phase = 100
+    code_frequency = carrier_doppler / 1540 + get_code_frequency(gpsl1)
+    sampling_frequency = 4e6Hz
+    prn = 1
+    range = 0:3999
+    start_carrier_phase = π / 2
+
+    num_samples = 4000
+    num_ants = NumAnts(3)
+
+    correlator = get_default_correlator(gpsl1, num_ants)
+    sat_states =
+        [SatState(gpsl1, 1, start_code_phase, carrier_doppler - 20Hz; num_ants, correlator)]
+
+    system_sats_state = SystemSatsState(gpsl1, sat_states)
+
+    # TODO: Why doesn't @inferred work here?
+    downconvert_and_correlator =
+        GPUDownconvertAndCorrelator((system_sats_state,), num_samples)
+
+    track_state = @inferred TrackState(system_sats_state)
+
+    signal =
+        cis.(2π .* carrier_doppler .* range ./ sampling_frequency .+ start_carrier_phase) .*
+        get_code.(
+            gpsl1,
+            code_frequency .* range ./ sampling_frequency .+ start_code_phase,
+            prn,
+        )
+    signal_mat = cu(repeat(signal; outer = (1, 3)))
+
+    track_state = @inferred track(
+        signal_mat,
+        track_state,
+        sampling_frequency;
+        downconvert_and_correlator,
+    )
+
+    iterations = 2000
+    code_phases = zeros(iterations)
+    carrier_phases = zeros(iterations)
+    tracked_code_phases = zeros(iterations)
+    tracked_carrier_phases = zeros(iterations)
+    tracked_code_dopplers = zeros(iterations)
+    tracked_carrier_dopplers = zeros(iterations)
+    tracked_prompts = zeros(ComplexF64, iterations)
+    for i = 1:iterations
+        carrier_phase =
+            mod2pi(
+                2π * carrier_doppler * 4000 * i / sampling_frequency +
+                start_carrier_phase +
+                π,
+            ) - π
+        code_phase =
+            mod(code_frequency * 4000 * i / sampling_frequency + start_code_phase, 1023)
+        signal =
+            cis.(2π .* carrier_doppler .* range ./ sampling_frequency .+ carrier_phase) .*
+            get_code.(
+                gpsl1,
+                code_frequency .* range ./ sampling_frequency .+ code_phase,
+                prn,
+            )
+        signal_mat = cu(repeat(signal; outer = (1, 3)))
+        track_state = @inferred track(
+            signal_mat,
+            track_state,
+            sampling_frequency;
+            downconvert_and_correlator,
+        )
+        comp_carrier_phase =
+            mod2pi(
+                2π * carrier_doppler * 4000 * (i + 1) / sampling_frequency +
+                start_carrier_phase +
+                π,
+            ) - π
+        comp_code_phase = mod(
+            code_frequency * 4000 * (i + 1) / sampling_frequency + start_code_phase,
+            1023,
+        )
+        tracked_code_phases[i] = get_code_phase(track_state)
+        tracked_carrier_phases[i] = get_carrier_phase(track_state)
+        tracked_carrier_dopplers[i] = get_carrier_doppler(track_state) / Hz
+        tracked_code_dopplers[i] = get_code_doppler(track_state) / Hz
+        tracked_prompts[i] = get_last_fully_integrated_filtered_prompt(track_state)
+        code_phases[i] = comp_code_phase
+        carrier_phases[i] = comp_carrier_phase
+    end
+    @test tracked_code_phases[end] ≈ code_phases[end] atol = 1e-2
+    @test tracked_carrier_phases[end] + π ≈ carrier_phases[end] atol = 5e-5
+end
+
+end