|
| 1 | +# SPDX-License-Identifier: MPL-2.0 |
| 2 | +# (PMPL-1.0-or-later preferred; MPL-2.0 required for Julia ecosystem) |
| 3 | +# E2E pipeline tests for AcceleratorGate.jl |
| 4 | +# Tests full backend selection, memory tracking, and operation registry workflows. |
| 5 | + |
| 6 | +using Test |
| 7 | +using AcceleratorGate |
| 8 | + |
| 9 | +@testset "E2E Pipeline Tests" begin |
| 10 | + |
| 11 | + @testset "Full pipeline: backend selection and dispatch" begin |
| 12 | + # Start from a clean state with JuliaBackend as default |
| 13 | + set_backend!(JuliaBackend()) |
| 14 | + @test current_backend() isa JuliaBackend |
| 15 | + |
| 16 | + # Register operations on a simulated GPU backend |
| 17 | + empty!(AcceleratorGate._BACKEND_OPS) |
| 18 | + register_operation!(CUDABackend, :matmul) |
| 19 | + register_operation!(CUDABackend, :fft) |
| 20 | + register_operation!(TPUBackend, :einsum) |
| 21 | + |
| 22 | + # Verify specialties are accessible before selection |
| 23 | + @test is_specialized(CUDABackend(0), :matmul) |
| 24 | + @test is_specialized(TPUBackend(0), :einsum) |
| 25 | + |
| 26 | + # Without hardware, selection always falls back to Julia |
| 27 | + b = select_backend(:matmul, 2048) |
| 28 | + @test b isa JuliaBackend |
| 29 | + |
| 30 | + # Use @with_backend macro for a scoped override |
| 31 | + @with_backend CUDABackend(0) begin |
| 32 | + @test current_backend() isa CUDABackend |
| 33 | + @test current_backend().device == 0 |
| 34 | + end |
| 35 | + # Scope restored after block |
| 36 | + @test current_backend() isa JuliaBackend |
| 37 | + |
| 38 | + empty!(AcceleratorGate._BACKEND_OPS) |
| 39 | + end |
| 40 | + |
| 41 | + @testset "Full pipeline: memory tracking lifecycle" begin |
| 42 | + empty!(AcceleratorGate._MEMORY_USAGE) |
| 43 | + |
| 44 | + julia_b = JuliaBackend() |
| 45 | + tpu_b = TPUBackend(0) |
| 46 | + |
| 47 | + # Allocate across two backends |
| 48 | + track_allocation!(julia_b, Int64(4096)) |
| 49 | + track_allocation!(julia_b, Int64(8192)) |
| 50 | + track_allocation!(tpu_b, Int64(1_000_000)) |
| 51 | + |
| 52 | + @test memory_usage(julia_b) == Int64(12288) |
| 53 | + @test memory_usage(tpu_b) == Int64(1_000_000) |
| 54 | + |
| 55 | + # Partial deallocations |
| 56 | + track_deallocation!(julia_b, Int64(4096)) |
| 57 | + @test memory_usage(julia_b) == Int64(8192) |
| 58 | + |
| 59 | + # Full memory report |
| 60 | + report = memory_report() |
| 61 | + @test report isa Dict{String, Int64} |
| 62 | + @test report["JuliaBackend"] == Int64(8192) |
| 63 | + @test report["TPUBackend"] == Int64(1_000_000) |
| 64 | + |
| 65 | + empty!(AcceleratorGate._MEMORY_USAGE) |
| 66 | + end |
| 67 | + |
| 68 | + @testset "Full pipeline: capability report and diagnostics" begin |
| 69 | + reset_diagnostics!() |
| 70 | + |
| 71 | + # Record several diagnostic events |
| 72 | + _record_diagnostic!("cuda", "runtime_fallbacks") |
| 73 | + _record_diagnostic!("tpu", "runtime_fallbacks") |
| 74 | + _record_diagnostic!("cuda", "runtime_fallbacks") |
| 75 | + |
| 76 | + diag = runtime_diagnostics() |
| 77 | + @test diag["backends"]["cuda"]["runtime_fallbacks"] == 2 |
| 78 | + @test diag["backends"]["tpu"]["runtime_fallbacks"] == 1 |
| 79 | + |
| 80 | + # Full capability report structure |
| 81 | + report = capability_report() |
| 82 | + @test haskey(report, "generated_at") |
| 83 | + @test haskey(report, "strategy_order") |
| 84 | + @test haskey(report, "backends") |
| 85 | + @test haskey(report, "platform") |
| 86 | + @test length(report["strategy_order"]) == 12 |
| 87 | + |
| 88 | + reset_diagnostics!() |
| 89 | + @test runtime_diagnostics()["backends"]["cuda"]["runtime_fallbacks"] == 0 |
| 90 | + end |
| 91 | + |
| 92 | + @testset "Error handling: invalid backend operations" begin |
| 93 | + # Switching to a non-default backend and back should be safe |
| 94 | + set_backend!(NPUBackend(2)) |
| 95 | + @test current_backend() isa NPUBackend |
| 96 | + @test current_backend().device == 2 |
| 97 | + |
| 98 | + set_backend!(JuliaBackend()) |
| 99 | + @test current_backend() isa JuliaBackend |
| 100 | + |
| 101 | + # select_backend with all standard types excluded still returns Julia |
| 102 | + b = select_backend(:fft, 512; exclude=[CUDABackend, ROCmBackend, MetalBackend]) |
| 103 | + @test b isa JuliaBackend |
| 104 | + end |
| 105 | + |
| 106 | + @testset "Round-trip consistency: backend set/get" begin |
| 107 | + original = current_backend() |
| 108 | + |
| 109 | + for BackendType in [CUDABackend, ROCmBackend, MetalBackend, TPUBackend] |
| 110 | + b = BackendType(0) |
| 111 | + set_backend!(b) |
| 112 | + got = current_backend() |
| 113 | + @test got isa BackendType |
| 114 | + @test got.device == 0 |
| 115 | + end |
| 116 | + |
| 117 | + # Restore original |
| 118 | + set_backend!(JuliaBackend()) |
| 119 | + @test current_backend() isa JuliaBackend |
| 120 | + end |
| 121 | + |
| 122 | +end |
0 commit comments