From a92d21ef4e2e005df6813d8dccf09eb0081d9730 Mon Sep 17 00:00:00 2001 From: Nathan Musoke Date: Fri, 26 Jan 2024 14:09:56 -0500 Subject: [PATCH] opt: Use PrecomileTools Use PrecompileTools to precompile functions required for a very basic simulation. This should help with TTFX. Could be extended to `PencilGrids`, functions inside `simulate!` that only run intermittently, and summary statistics. --- Project.toml | 2 ++ src/UltraDark.jl | 28 ++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/Project.toml b/Project.toml index ae6a2a4..9cd8d64 100644 --- a/Project.toml +++ b/Project.toml @@ -14,6 +14,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" PencilArrays = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e" PencilFFTs = "4a48f351-57a6-4416-9ec4-c37015456aae" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67" @@ -28,6 +29,7 @@ MPI = "0.20" NPZ = "0.4" PencilArrays = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18" PencilFFTs = "0.12, 0.13, 0.14, 0.15" +PrecompileTools = "1.2" Statistics = "1.8, 1.9" Strided = "2.0.4" julia = "1.8, 1.9" diff --git a/src/UltraDark.jl b/src/UltraDark.jl index dff8a78..e6bb2bf 100644 --- a/src/UltraDark.jl +++ b/src/UltraDark.jl @@ -5,6 +5,7 @@ using AbstractFFTs: fftfreq, rfftfreq import Folds using LinearAlgebra using PencilFFTs, MPI +using PrecompileTools using FFTW @@ -273,4 +274,31 @@ function simulate!( end +@setup_workload begin + + resol = 64 + output_dir = mktempdir() + + @compile_workload begin + output_times = [0, 1e-9] + + output_config = OutputConfig( + output_dir, + output_times; + box = true, + slice = true, + npy = true, + h5 = true, + ) + + options = Config.SimulationConfig() + + grids = Grids(10.0, resol) + # Set ψx to something non-zero + grids.ψx .= (grids.x .^ 2 .+ grids.y .^ 2 .+ grids.z .^ 2) .^ 0.5 ./ 1e9 + + simulate!(grids, options, output_config) + end +end + end # module