Merge pull request #732 from control-toolbox/ocots-patch-1

Modify timing to use CUDA and track GPU memory

Author: ocots

test/suite/builders/test_options_forwarding.jl

@@ -54,7 +54,7 @@ function test_options_forwarding()
             if is_cuda_on()
                 Test.@testset "backend (CUDA)" begin
                     Test.@test begin
-                        modeler = OptimalControl.Exa(backend=CUDA.CUDABackend())
+                        modeler = OptimalControl.Exa{OptimalControl.GPU}(backend=CUDA.CUDABackend())
                         nlp = OptimalControl.nlp_model(docp, normalized_init, modeler)
                         # With CUDA backend, x0 should be CUDA array
                         nlp.meta.x0 isa CUDA.CuArray

test/suite/solve/test_canonical.jl

@@ -206,7 +206,7 @@ function test_canonical()
                             gpu_solver[2];
                             display=false,
                         ) # warmup
-                        timed_result = @timed begin
+                        timed_result = CUDA.@timed begin
                             OptimalControl.solve(
                                 pb.ocp,
                                 normalized_init,
@@ -220,7 +220,7 @@ function test_canonical()
                         # Extract results
                         solve_result = timed_result.value
                         solve_time = timed_result.time
-                        memory_bytes = timed_result.bytes
+                        memory_bytes = timed_result.cpu_bytes + timed_result.gpu_bytes
 
                         success = OptimalControl.successful(solve_result)
                         obj = success ? OptimalControl.objective(solve_result) : 0.0