Add trim tests

.github/workflows/CI_NonlinearSolve.yml

@@ -36,6 +36,7 @@ jobs:
           - downstream
           - wrappers
           - misc
+          - trim
         version:
           - "1"
           - "lts"

Project.toml

@@ -171,4 +171,3 @@ path = "lib/SimpleNonlinearSolve"
 
 [targets]
 test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"]
-

test/runtests.jl

@@ -1,32 +1,42 @@
-using ReTestItems, NonlinearSolve, Hwloc, InteractiveUtils, Pkg
+using ReTestItems, Hwloc, InteractiveUtils, Pkg
 
 @info sprint(InteractiveUtils.versioninfo)
 
 const GROUP = lowercase(get(ENV, "GROUP", "All"))
 
-const EXTRA_PKGS = Pkg.PackageSpec[]
-if GROUP == "all" || GROUP == "downstream"
-    push!(EXTRA_PKGS, Pkg.PackageSpec("ModelingToolkit"))
-    push!(EXTRA_PKGS, Pkg.PackageSpec("SymbolicIndexingInterface"))
-end
-length(EXTRA_PKGS) ≥ 1 && Pkg.add(EXTRA_PKGS)
+if GROUP != "trim"
+    using NonlinearSolve  # trimming uses NonlinearSolve from a custom environment
+
+    const EXTRA_PKGS = Pkg.PackageSpec[]
+    if GROUP == "all" || GROUP == "downstream"
+        push!(EXTRA_PKGS, Pkg.PackageSpec("ModelingToolkit"))
+        push!(EXTRA_PKGS, Pkg.PackageSpec("SymbolicIndexingInterface"))
+    end
+    length(EXTRA_PKGS) ≥ 1 && Pkg.add(EXTRA_PKGS)
 
-const RETESTITEMS_NWORKERS = parse(
-    Int, get(ENV, "RETESTITEMS_NWORKERS",
-        string(min(ifelse(Sys.iswindows(), 0, Hwloc.num_physical_cores()), 4))
+    const RETESTITEMS_NWORKERS = parse(
+        Int, get(
+            ENV, "RETESTITEMS_NWORKERS",
+            string(min(ifelse(Sys.iswindows(), 0, Hwloc.num_physical_cores()), 4))
+        )
     )
-)
-const RETESTITEMS_NWORKER_THREADS = parse(Int,
-    get(
-        ENV, "RETESTITEMS_NWORKER_THREADS",
-        string(max(Hwloc.num_virtual_cores() ÷ max(RETESTITEMS_NWORKERS, 1), 1))
+    const RETESTITEMS_NWORKER_THREADS = parse(
+        Int,
+        get(
+            ENV, "RETESTITEMS_NWORKER_THREADS",
+            string(max(Hwloc.num_virtual_cores() ÷ max(RETESTITEMS_NWORKERS, 1), 1))
+        )
     )
-)
 
-@info "Running tests for group: $(GROUP) with $(RETESTITEMS_NWORKERS) workers"
+    @info "Running tests for group: $(GROUP) with $(RETESTITEMS_NWORKERS) workers"
 
-ReTestItems.runtests(
-    NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]),
-    nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS,
-    testitem_timeout = 3600
-)
+    ReTestItems.runtests(
+        NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]),
+        nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS,
+        testitem_timeout = 3600
+    )
+elseif GROUP == "trim" && VERSION >= v"1.12.0-rc1"  # trimming has been introduced in julia 1.12
+    Pkg.activate(joinpath(dirname(@__FILE__), "trim"))
+    Pkg.instantiate()
+    include("trim/runtests.jl")
+end

test/trim/Project.toml

@@ -0,0 +1,39 @@
+name = "TrimTest"
+uuid = "7e54ada7-ece5-4046-aa01-512d530850d8"
+
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
+PolyesterWeave = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+
+[sources]
+# Remove assert that triggers false positive for JET. Tracked at https://github.com/aviatesk/JET.jl/issues/736.
+ForwardDiff = {url = "https://github.com/RomeoV/ForwardDiff.jl", rev="rv/remove-quote-assert-string-interpolation"}
+# Remove special path code path for ForwarDiff of LinearSolve, which uses untrimmable `deepcopy`.
+# Use of `deepcopy` is tracked at https://github.com/SciML/LinearSolve.jl/issues/648.
+LinearSolve = {url = "https://github.com/RomeoV/LinearSolve.jl", rev="rv/remove-linsolve-forwarddiff-special-path"}
+NonlinearSolveFirstOrder = {path = "../../lib/NonlinearSolveFirstOrder"}
+# Fix a type instability. Tracked at https://github.com/SciML/SciMLBase.jl/pull/1074.
+SciMLBase = {url = "https://github.com/AayushSabharwal/SciMLBase.jl", rev="as/fix-jet-opt"}
+
+
+[compat]
+ADTypes = "1.15.0"
+CPUSummary = "0.2.7"
+DiffEqBase = "6.179.0"
+ForwardDiff = "1.0.1"
+LinearAlgebra = "1.12.0"
+NonlinearSolveFirstOrder = "1.6.0"
+Polyester = "0.7.18"
+PolyesterWeave = "0.2.2"
+StaticArrays = "1.9.0"

test/trim/main_clean.jl

@@ -0,0 +1,8 @@
+using TrimTest
+
+function (@main)(argv::Vector{String})::Cint
+    λ = parse(Float64, argv[2])
+    sol = TrimTest.TestModuleClean.minimize(λ)
+    println(Core.stdout, sum(sol.u))
+    return 0
+end

test/trim/main_segfault.jl

@@ -0,0 +1,10 @@
+module MyModule
+include("./optimization_trimmable.jl")
+end
+
+function (@main)(argv::Vector{String})::Cint
+    λ = parse(Float64, argv[2])
+    sol = MyModule.TestModuleTrimmable.minimize(λ)
+    println(Core.stdout, sum(sol.u))
+    return 0
+end

test/trim/main_trimmable.jl

@@ -0,0 +1,8 @@
+using TrimTest
+
+function (@main)(argv::Vector{String})::Cint
+    λ = parse(Float64, argv[2])
+    sol = TrimTest.TestModuleTrimmable.minimize(λ)
+    println(Core.stdout, sum(sol.u))
+    return 0
+end

test/trim/optimization_clean.jl

@@ -0,0 +1,34 @@
+module TestModuleClean
+using NonlinearSolveFirstOrder
+using ADTypes: AutoForwardDiff
+using ForwardDiff
+using LinearAlgebra
+using StaticArrays
+using LinearSolve
+const LS = LinearSolve
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    rhs = (u .* u .- p.λ)
+    # there are some issues currently with LinearSolve and triangular matrices,
+    # so we just make `L` dense here.
+    linprob = LinearProblem(Matrix(L), rhs)
+    alg = LS.GenericLUFactorization()
+    sol = LinearSolve.solve(linprob, alg)
+    return sol.u
+end
+
+struct MyParams{T, M}
+    λ::T
+    Σ::M
+end
+
+function minimize(x)
+    autodiff = AutoForwardDiff(; chunksize = 1)
+    alg = TrustRegion(; autodiff, linsolve = LS.CholeskyFactorization())
+    ps = MyParams(rand(), hermitianpart(rand(2, 2) + 2I))
+    prob = NonlinearLeastSquaresProblem{false}(f, rand(2), ps)
+    sol = solve(prob, alg)
+    return sol
+end
+end

test/trim/optimization_trimmable.jl

@@ -0,0 +1,43 @@
+module TestModuleTrimmable
+using NonlinearSolveFirstOrder
+using DiffEqBase
+using ADTypes: AutoForwardDiff
+using ForwardDiff
+using LinearAlgebra
+using StaticArrays
+using LinearSolve
+import SciMLBase
+const LS = LinearSolve
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    rhs = (u .* u .- p.λ)
+    # there are some issues currently with LinearSolve and triangular matrices,
+    # so we just make `L` dense here.
+    linprob = LinearProblem(Matrix(L), rhs)
+    alg = LS.GenericLUFactorization()
+    sol = LinearSolve.solve(linprob, alg)
+    return sol.u
+end
+
+struct MyParams{T, M}
+    λ::T
+    Σ::M
+end
+
+const autodiff = AutoForwardDiff(; chunksize = 1)
+const alg = TrustRegion(; autodiff, linsolve = LS.CholeskyFactorization())
+const prob = NonlinearLeastSquaresProblem{false}(
+    f,
+    rand(2),
+    MyParams(rand(), hermitianpart(rand(2, 2) + 2I))
+)
+const cache = init(prob, alg)
+
+function minimize(x)
+    ps = MyParams(x, hermitianpart(rand(2, 2) + 2I))
+    reinit!(cache, rand(2); p = ps)
+    solve!(cache)
+    return cache
+end
+end

test/trim/runtests.jl

@@ -0,0 +1,78 @@
+using SafeTestsets
+
+@safetestset "Clean implementation (non-trimmable)" begin
+    using JET
+    using SciMLBase: successful_retcode
+    include("optimization_clean.jl")
+    @test successful_retcode(TestModuleClean.minimize(1.0).retcode)
+    # can't use `@test_opt` macro here because it would try to eval before
+    # `using JET` is processed
+    test_opt(TestModuleClean.minimize, (typeof(1.0),))
+end
+
+@safetestset "Trimmable implementation" begin
+    using JET
+    using SciMLBase: successful_retcode
+    include("optimization_trimmable.jl")
+    @test successful_retcode(TestModuleTrimmable.minimize(1.0).retcode)
+    # can't use `@test_opt` macro here because it would try to eval before
+    # `using JET` is processed
+    test_opt(TestModuleTrimmable.minimize, (typeof(1.0),))
+end
+
+@safetestset "Run trim" begin
+    # https://discourse.julialang.org/t/capture-stdout-and-stderr-in-case-a-command-fails/101772/3?u=romeov
+    """
+    Run a Cmd object, returning the stdout & stderr contents plus the exit code
+    """
+    function _execute(cmd::Cmd)
+        out = Pipe()
+        err = Pipe()
+        process = run(pipeline(ignorestatus(cmd); stdout = out, stderr = err))
+        close(out.in)
+        close(err.in)
+        out = (
+            stdout = String(read(out)), stderr = String(read(err)),
+            exitcode = process.exitcode
+        )
+        return out
+    end
+
+    JULIAC = normpath(
+        joinpath(
+        Sys.BINDIR, Base.DATAROOTDIR, "julia", "juliac",
+        "juliac.jl"
+    )
+    )
+    @test isfile(JULIAC)
+
+    for (mainfile, expectedtopass) in [
+        ("main_trimmable.jl", true),
+    #= The test below should verify that we indeed can't get a trimmed binary
+    # for the "clean" implementation, but will trigger in the future if
+    # it does start working. Unfortunately, right now it hangs indefinitely
+    # so we are commenting it out. =#
+    # ("main_clean.jl", false),
+        ("main_segfault.jl", false),
+    ]
+        binpath = tempname()
+        cmd = `$(Base.julia_cmd()) --project=. --depwarn=error $(JULIAC) --experimental --trim=unsafe-warn --output-exe $(binpath) $(mainfile)`
+
+        # since we are calling Julia from Julia, we first need to clean some
+        # environment variables
+        clean_env = copy(ENV)
+        delete!(clean_env, "JULIA_PROJECT")
+        delete!(clean_env, "JULIA_LOAD_PATH")
+        # We could just check for success, but then failures are hard to debug.
+        # Instead we use `_execute` to also capture `stdout` and `stderr`.
+        # @test success(setenv(cmd, clean_env))
+        trimcall = _execute(setenv(cmd, clean_env; dir = @__DIR__))
+        if trimcall.exitcode != 0 && expectedtopass
+            @show trimcall.stdout
+            @show trimcall.stderr
+        end
+        @test trimcall.exitcode == 0 broken = !expectedtopass
+        @test isfile(binpath) broken = !expectedtopass
+        @test success(`$(binpath) 1.0`) broken = !expectedtopass
+    end
+end

test/trim/src/TrimTest.jl

@@ -0,0 +1,27 @@
+module TrimTest
+#=
+Currently, trimming only works if the target code is in a package. I.e., trying to trim
+```julia
+include("optimization_trimmable.jl")
+function (@main)(argv::Vector{String})::Cint
+    minimize(1.0)
+    return 0
+end
+```
+or even
+```julia
+mod MyMod
+    include("optimization_trimmable.jl")
+end
+function (@main)(argv::Vector{String})::Cint
+    MyMod.minimize(1.0)
+    return 0
+end
+```
+segfaults `juliac`. Looking at the segfault stacktrace it seems the culprit is
+`const cache = init(...)`. Either way, we circumvent the segfault by putting
+this below code into a package definition.
+=#
+include("../optimization_trimmable.jl")
+include("../optimization_clean.jl")
+end