Fix OptimizationProblem for SVector/SArray: use out-of-place form

.github/workflows/FormatCheck.yml

@@ -14,6 +14,9 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v6
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: '1'
       - uses: fredrikekre/runic-action@v1
         with:
           version: '1'

.github/workflows/GPU.yml

@@ -17,7 +17,7 @@ concurrency:
 jobs:
   cuda-tests:
     name: "CUDA Tests (Julia ${{ matrix.version }})"
-    runs-on: [self-hosted, Linux, X64, gpu-t4]
+    runs-on: [self-hosted, gpu-v100]
     timeout-minutes: 120
     strategy:
       fail-fast: false
@@ -32,7 +32,7 @@ jobs:
       - run: |
           julia --project=test -e '
             import Pkg
-            Pkg.add(; name = "CUDA")'
+            Pkg.add(; name = "CUDA", version = "5.0 - 5.10")'
           rm -f test/Manifest.toml
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1

src/precompilation.jl

@@ -21,7 +21,9 @@ using PrecompileTools
         p = @SArray Float32[1.0, 100.0]
 
         # Create optimization problem with StaticArrays
-        prob = OptimizationProblem(_rosenbrock_precompile, x0, p; lb = lb, ub = ub)
+        # Use out-of-place form {false} since SVector is immutable
+        opt_f = OptimizationFunction{false}(_rosenbrock_precompile)
+        prob = OptimizationProblem(opt_f, x0, p; lb = lb, ub = ub)
 
         # Precompile SerialPSO - most commonly used CPU algorithm
         sol = solve(prob, SerialPSO(10), maxiters = 2)

test/constraints.jl

@@ -12,7 +12,8 @@ function conss(x, p)
     return SVector{3}(-x[1] + 2 * x[2] - 1, +x[1] - 2 * x[2] + 1, (x[1]^2) / 4 + x[2]^2 - 1)
 end
 
-opt_f = OptimizationFunction(objective, cons = conss)
+# Use out-of-place form {false} since SVector is immutable
+opt_f = OptimizationFunction{false}(objective, cons = conss)
 
 x0 = @SVector [1.0f0, 1.0f0]
 lb = @SVector [0.0f0, 0.0f0]

test/gpu.jl

@@ -21,7 +21,9 @@ include("./utils.jl")
     x0 = @SArray zeros(Float32, N)
     p = @SArray Float32[1.0, 100.0]
 
-    prob = OptimizationProblem(rosenbrock, x0, p; lb = lb, ub = ub)
+    # Use out-of-place form {false} since SVector is immutable
+    opt_f = OptimizationFunction{false}(rosenbrock)
+    prob = OptimizationProblem(opt_f, x0, p; lb = lb, ub = ub)
 
     n_particles = 5000
 

test/lbfgs.jl

@@ -2,11 +2,17 @@ using ParallelParticleSwarms, Optimization, StaticArrays
 
 include("./utils.jl")
 
+# Reclaim GPU memory from previous test files to avoid OOM
+if GROUP == "CUDA"
+    CUDA.reclaim()
+end
+
 function objf(x, p)
     return 1 - x[1]^2 - x[2]^2
 end
 
-optprob = OptimizationFunction(objf, Optimization.AutoEnzyme())
+# Use out-of-place form {false} since SVector is immutable
+optprob = OptimizationFunction{false}(objf, Optimization.AutoEnzyme())
 x0 = rand(2)
 x0 = SVector{2}(x0)
 prob = OptimizationProblem(optprob, x0)
@@ -27,7 +33,8 @@ function rosenbrock(x, p)
 end
 x0 = @SArray rand(Float32, N)
 p = @SArray Float32[1.0, 100.0]
-optf = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff())
+# Use out-of-place form {false} since SArray is immutable
+optf = OptimizationFunction{false}(rosenbrock, Optimization.AutoForwardDiff())
 prob = OptimizationProblem(optf, x0, p)
 l0 = rosenbrock(x0, p)
 
@@ -60,7 +67,8 @@ l0 = rosenbrock(x0, p)
 )
 @show sol.objective
 
-optf = OptimizationFunction(rosenbrock, Optimization.AutoEnzyme())
+# Use out-of-place form {false} since SArray is immutable
+optf = OptimizationFunction{false}(rosenbrock, Optimization.AutoEnzyme())
 prob = OptimizationProblem(optf, x0, p)
 l0 = rosenbrock(x0, p)
 

test/regression.jl

@@ -28,7 +28,9 @@ using QuasiMonteCarlo
         ub = ub
     )
 
-    prob = OptimizationProblem(rosenbrock, x0, p; lb = lb, ub = ub)
+    # Use out-of-place form {false} since SVector is immutable
+    opt_f = OptimizationFunction{false}(rosenbrock)
+    prob = OptimizationProblem(opt_f, x0, p; lb = lb, ub = ub)
 
     n_particles = 2000
 
@@ -224,7 +226,9 @@ end
         ub = ub
     )
 
-    prob = OptimizationProblem(rosenbrock, x0, p; lb = lb, ub = ub)
+    # Use out-of-place form {false} since SVector is immutable
+    opt_f = OptimizationFunction{false}(rosenbrock)
+    prob = OptimizationProblem(opt_f, x0, p; lb = lb, ub = ub)
 
     n_particles = 2000
 

test/reinit.jl

@@ -17,7 +17,9 @@ end
 x0 = @SArray zeros(Float32, 3)
 p = @SArray Float32[1.0, 100.0]
 
-prob = OptimizationProblem(rosenbrock, x0, p; lb = lb, ub = ub)
+# Use out-of-place form {false} since SVector is immutable
+opt_f = OptimizationFunction{false}(rosenbrock)
+prob = OptimizationProblem(opt_f, x0, p; lb = lb, ub = ub)
 
 n_particles = 2000
 

test/runtests.jl

@@ -11,6 +11,8 @@ global CI_GROUP = get(ENV, "GROUP", "CPU")
 @testset for BACKEND in unique(("CPU", CI_GROUP))
     global GROUP = BACKEND
     @testset "$(BACKEND) optimizers tests" include("./gpu.jl")
+    GC.gc(true)
     @testset "$(BACKEND) optimizers with constraints tests" include("./constraints.jl")
+    GC.gc(true)
     @testset "$(BACKEND) hybrid optimizers" include("./lbfgs.jl")
 end