diff --git a/Project.toml b/Project.toml index da35f758..989424f5 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" @@ -66,6 +67,7 @@ Metal = "1" MuladdMacro = "0.2.4" OpenCL = "0.10" Parameters = "0.13" +PrecompileTools = "1.2.1" Random = "1" RecursiveArrayTools = "4.2" SciMLBase = "3.49" diff --git a/src/DiffEqGPU.jl b/src/DiffEqGPU.jl index 83b1fb98..cc2f47b9 100644 --- a/src/DiffEqGPU.jl +++ b/src/DiffEqGPU.jl @@ -35,6 +35,7 @@ using StaticArrays: StaticArrays using StaticArraysCore: MArray, MMatrix, SArray, SMatrix, SVector, Size, StaticMatrix, StaticVector, similar_type using Parameters: Parameters +using PrecompileTools: @compile_workload, @setup_workload using MuladdMacro: MuladdMacro, @muladd using Random: Random using Setfield: Setfield, @set, @set! @@ -247,4 +248,16 @@ export terminate! make_prob_compatible, maxthreads, maybe_prefer_blocks, lufact!, vectorized_solve, vectorized_asolve, vectorized_map_solve +@setup_workload begin + @compile_workload begin + precompile_rhs(u, p, t) = p * u + prob = ODEProblem{false}(precompile_rhs, SVector(1.0f0), (0.0f0, 0.1f0), 1.0f0) + ensemble_prob = EnsembleProblem(prob) + solve( + ensemble_prob, GPUTsit5(), EnsembleGPUKernel(CPU()); + trajectories = 2, adaptive = false, dt = 0.05f0 + ) + end +end + end # module