Use InboundsArrays so we can optimize without needing `--check-bounds…

…=no` Introduces type-alias `MKArray` in `type_definitions.jl` to make it (hopefully) simple to change the array type we want to use in future. In particular it is easy to revert to using `Array` just by redefining `MKArray`.
mabarnes · Jan 29, 2025 · 512b44c · 512b44c
1 parent e1dcf65
commit 512b44c
Show file tree

Hide file tree

Showing 77 changed files with 1,176 additions and 1,032 deletions.
diff --git a/.github/workflows/parallel_test.yml b/.github/workflows/parallel_test.yml
@@ -21,15 +21,15 @@ jobs:
       - uses: julia-actions/cache@v2
       - run: |
           touch Project.toml
-          julia --project -O3 --check-bounds=no -e 'import Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler"]); using MPIPreferences; MPIPreferences.use_jll_binary("OpenMPI_jll")'
-          julia --project -O3 --check-bounds=no -e 'using MPI; MPI.install_mpiexecjl(; destdir=".")'
-          julia --project -O3 --check-bounds=no -e 'import Pkg; Pkg.add(["NCDatasets", "Random", "SpecialFunctions", "StatsBase", "Test"]); Pkg.develop(path="moment_kinetics/")'
-          julia --project -O3 --check-bounds=no -e 'import Pkg; Pkg.precompile()'
-          julia --project -O3 --check-bounds=no precompile.jl
+          julia --project -O3 -e 'import Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler"]); using MPIPreferences; MPIPreferences.use_jll_binary("OpenMPI_jll")'
+          julia --project -O3 -e 'using MPI; MPI.install_mpiexecjl(; destdir=".")'
+          julia --project -O3 -e 'import Pkg; Pkg.add(["NCDatasets", "Random", "SpecialFunctions", "StatsBase", "Test"]); Pkg.develop(path="moment_kinetics/")'
+          julia --project -O3 -e 'import Pkg; Pkg.precompile()'
+          julia --project -O3 precompile.jl
           # Need to use openmpi so that we can use `--oversubscribe` to allow using more MPI ranks than physical cores
-          ./mpiexecjl -np 3 --oversubscribe julia -J moment_kinetics.so --project -O3 --check-bounds=no moment_kinetics/test/runtests.jl --ci --debug 1
-          ./mpiexecjl -np 4 --oversubscribe julia -J moment_kinetics.so --project -O3 --check-bounds=no moment_kinetics/test/runtests.jl --ci --debug 1
-          ./mpiexecjl -np 2 --oversubscribe julia -J moment_kinetics.so --project -O3 --check-bounds=no moment_kinetics/test/runtests.jl --ci --debug 1 --long
+          ./mpiexecjl -np 3 --oversubscribe julia -J moment_kinetics.so --project -O3 moment_kinetics/test/runtests.jl --ci --debug 1
+          ./mpiexecjl -np 4 --oversubscribe julia -J moment_kinetics.so --project -O3 moment_kinetics/test/runtests.jl --ci --debug 1
+          ./mpiexecjl -np 2 --oversubscribe julia -J moment_kinetics.so --project -O3 moment_kinetics/test/runtests.jl --ci --debug 1 --long
           # Note: MPI.jl's default implementation is mpich, which has a similar option
           # `--with-device=ch3:sock`, but that needs to be set when compiling mpich.
         shell: bash
@@ -47,13 +47,13 @@ jobs:
       - uses: julia-actions/cache@v2
       - run: |
           touch Project.toml
-          julia --project -O3 --check-bounds=no -e 'import Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler"]); using MPIPreferences; MPIPreferences.use_jll_binary("OpenMPI_jll")'
-          julia --project -O3 --check-bounds=no -e 'using MPI; MPI.install_mpiexecjl(; destdir=".")'
-          julia --project -O3 --check-bounds=no -e 'import Pkg; Pkg.add(["NCDatasets", "Random", "SpecialFunctions", "StatsBase", "Test"]); Pkg.develop(path="moment_kinetics/")'
-          julia --project -O3 --check-bounds=no -e 'import Pkg; Pkg.precompile()'
-          julia --project -O3 --check-bounds=no precompile.jl
+          julia --project -O3 -e 'import Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler"]); using MPIPreferences; MPIPreferences.use_jll_binary("OpenMPI_jll")'
+          julia --project -O3 -e 'using MPI; MPI.install_mpiexecjl(; destdir=".")'
+          julia --project -O3 -e 'import Pkg; Pkg.add(["NCDatasets", "Random", "SpecialFunctions", "StatsBase", "Test"]); Pkg.develop(path="moment_kinetics/")'
+          julia --project -O3 -e 'import Pkg; Pkg.precompile()'
+          julia --project -O3 precompile.jl
           # Need to use openmpi so that we can use `--oversubscribe` to allow using more MPI ranks than physical cores
-          ./mpiexecjl -np 4 --oversubscribe julia -J moment_kinetics.so --project -O3 --check-bounds=no moment_kinetics/test/runtests.jl --ci --debug 1
+          ./mpiexecjl -np 4 --oversubscribe julia -J moment_kinetics.so --project -O3 moment_kinetics/test/runtests.jl --ci --debug 1
           # Note: MPI.jl's default implementation is mpich, which has a similar option
           # `--with-device=ch3:sock`, but that needs to be set when compiling mpich.
         shell: bash
diff --git a/machines/archer/jobscript-precompile-no-run.template b/machines/archer/jobscript-precompile-no-run.template
@@ -22,6 +22,6 @@ export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
 
 echo "precompiling $(date)"
 
-bin/julia --project -O3 --check-bounds=no precompile-no-run.jl
+bin/julia --project -O3 precompile-no-run.jl
 
 echo "finished! $(date)"
diff --git a/machines/archer/jobscript-precompile.template b/machines/archer/jobscript-precompile.template
@@ -22,6 +22,6 @@ export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
 
 echo "precompiling $(date)"
 
-bin/julia --project -O3 --check-bounds=no precompile.jl
+bin/julia --project -O3 precompile.jl
 
 echo "finished! $(date)"
diff --git a/machines/archer/jobscript-restart.template b/machines/archer/jobscript-restart.template
@@ -22,6 +22,6 @@ export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
 
 echo "running INPUTFILE $(date)"
 
-srun --distribution=block:block --hint=nomultithread --ntasks=$SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 --check-bounds=no run_moment_kinetics.jl --restart INPUTFILE RESTARTFROM
+srun --distribution=block:block --hint=nomultithread --ntasks=$SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 run_moment_kinetics.jl --restart INPUTFILE RESTARTFROM
 
 echo "finished INPUTFILE $(date)"
diff --git a/machines/archer/jobscript-run.template b/machines/archer/jobscript-run.template
@@ -22,6 +22,6 @@ export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
 
 echo "running INPUTFILE $(date)"
 
-srun --distribution=block:block --hint=nomultithread --ntasks=$SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 --check-bounds=no run_moment_kinetics.jl INPUTFILE
+srun --distribution=block:block --hint=nomultithread --ntasks=$SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 run_moment_kinetics.jl INPUTFILE
 
 echo "finished INPUTFILE $(date)"
diff --git a/machines/generic-batch-template/jobscript-precompile-no-run.template b/machines/generic-batch-template/jobscript-precompile-no-run.template
@@ -15,6 +15,6 @@ source julia.env
 
 echo "precompiling $(date)"
 
-bin/julia --project -O3 --check-bounds=no precompile-no-run.jl
+bin/julia --project -O3 precompile-no-run.jl
 
 echo "finished! $(date)"
diff --git a/machines/generic-batch-template/jobscript-precompile.template b/machines/generic-batch-template/jobscript-precompile.template
@@ -15,6 +15,6 @@ source julia.env
 
 echo "precompiling $(date)"
 
-bin/julia --project -O3 --check-bounds=no precompile.jl
+bin/julia --project -O3 precompile.jl
 
 echo "finished! $(date)"
diff --git a/machines/generic-batch-template/jobscript-restart.template b/machines/generic-batch-template/jobscript-restart.template
@@ -19,6 +19,6 @@ source julia.env
 echo "running INPUTFILE $(date)"
 
 # May need to change this if mpirun` is not what should be used on your system
-mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 --check-bounds=no run_moment_kinetics.jl --restart INPUTFILE RESTARTFROM
+mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 run_moment_kinetics.jl --restart INPUTFILE RESTARTFROM
 
 echo "finished INPUTFILE $(date)"
diff --git a/machines/generic-batch-template/jobscript-run.template b/machines/generic-batch-template/jobscript-run.template
@@ -19,6 +19,6 @@ source julia.env
 echo "running INPUTFILE $(date)"
 
 # May need to change this if mpirun` is not what should be used on your system
-mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 --check-bounds=no run_moment_kinetics.jl INPUTFILE
+mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 run_moment_kinetics.jl INPUTFILE
 
 echo "finished INPUTFILE $(date)"
diff --git a/machines/machine_setup.sh b/machines/machine_setup.sh
@@ -271,7 +271,7 @@ SEPARATE_POSTPROC_PROJECTS=$(bin/julia --project machines/shared/get_mk_preferen
 if [[ $BATCH_SYSTEM -eq 0 || $SEPARATE_POSTPROC_PROJECTS == "y" ]]; then
   # Batch systems can (conveniently) use different optimization flags for
   # running simulations and for post-processing.
-  OPTIMIZATION_FLAGS="-O3 --check-bounds=no"
+  OPTIMIZATION_FLAGS="-O3"
   POSTPROC_OPTIMIZATION_FLAGS="-O3"
 else
   # On interactive systems which use the same project for running simulations

diff --git a/machines/marconi/jobscript-precompile-no-run.template b/machines/marconi/jobscript-precompile-no-run.template
@@ -15,6 +15,6 @@ source julia.env
 
 echo "precompiling $(date)"
 
-bin/julia --project -O3 --check-bounds=no precompile-no-run.jl
+bin/julia --project -O3 precompile-no-run.jl
 
 echo "finished! $(date)"
diff --git a/machines/marconi/jobscript-precompile.template b/machines/marconi/jobscript-precompile.template
@@ -15,6 +15,6 @@ source julia.env
 
 echo "precompiling $(date)"
 
-bin/julia --project -O3 --check-bounds=no precompile.jl
+bin/julia --project -O3 precompile.jl
 
 echo "finished! $(date)"
diff --git a/machines/marconi/jobscript-restart.template b/machines/marconi/jobscript-restart.template
@@ -18,6 +18,6 @@ source julia.env
 
 echo "running INPUTFILE $(date)"
 
-mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 --check-bounds=no run_moment_kinetics.jl --restart INPUTFILE RESTARTFROM
+mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 run_moment_kinetics.jl --restart INPUTFILE RESTARTFROM
 
 echo "finished INPUTFILE $(date)"
diff --git a/machines/marconi/jobscript-run.template b/machines/marconi/jobscript-run.template
@@ -18,6 +18,6 @@ source julia.env
 
 echo "running INPUTFILE $(date)"
 
-mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 --check-bounds=no run_moment_kinetics.jl INPUTFILE
+mpirun -np $SLURM_NTASKS bin/julia -Jmoment_kinetics.so --project -O3 run_moment_kinetics.jl INPUTFILE
 
 echo "finished INPUTFILE $(date)"
diff --git a/machines/shared/submit-restart.sh b/machines/shared/submit-restart.sh
@@ -123,7 +123,7 @@ sed -e "s|NODES|$NODES|" -e "s|RUNTIME|$RUNTIME|" -e "s|ACCOUNT|$ACCOUNT|" -e "s
 
 if [[ "$WARN_OLD_SYSIMAGE" -eq 0 ]]; then
   # Check that source code has not been changed since moment_kinetics.so was created
-  bin/julia --project -O3 --check-bounds=no moment_kinetics/src/check_so_newer_than_code.jl moment_kinetics.so
+  bin/julia --project -O3 moment_kinetics/src/check_so_newer_than_code.jl moment_kinetics.so
 fi
 
 if [[ $SUBMIT -eq 0 ]]; then

diff --git a/machines/shared/submit-run.sh b/machines/shared/submit-run.sh
@@ -118,7 +118,7 @@ sed -e "s|NODES|$NODES|" -e "s|RUNTIME|$RUNTIME|" -e "s|ACCOUNT|$ACCOUNT|" -e "s
 
 if [[ "$WARN_OLD_SYSIMAGE" -eq 0 ]]; then
   # Check that source code has not been changed since moment_kinetics.so was created
-  bin/julia --project -O3 --check-bounds=no moment_kinetics/src/check_so_newer_than_code.jl moment_kinetics.so
+  bin/julia --project -O3 moment_kinetics/src/check_so_newer_than_code.jl moment_kinetics.so
 fi
 
 if [[ $SUBMIT -eq 0 ]]; then

diff --git a/makie_post_processing/makie_post_processing/src/generic_plot_functions.jl b/makie_post_processing/makie_post_processing/src/generic_plot_functions.jl
@@ -2376,6 +2376,11 @@ function put_legend_right(fig, ax; kwargs...)
     return Legend(fig[end,end+1], ax; kwargs...)
 end
 
+if !(AbstractMKVector <: AbstractArray)
+    # Method to pass wrapped array through to matrix_grid(), since this is an internal
+    # function of Makie that does not have methods added by InboundsArrays.jl
+    @inline Makie.matrix_grid(t, x::AbstractMKArray, y::AbstractMKArray, z::AbstractMKArray) = Makie.matrix_grid(t, x.a, y.a, z.a)
+end
 """
     curvilinear_grid_mesh(xs, ys, zs, colors)
 
@@ -2589,13 +2594,13 @@ function irregular_heatmap!(ax, xs, ys, zs; kwargs...)
         end
     end
 
-    vertices, faces, colors = curvilinear_grid_mesh(xs, ys, zeros(nx, ny), zs)
+    vertices, faces, colors = curvilinear_grid_mesh(xs, ys, mk_zeros(nx, ny), zs)
 
     return mesh!(ax, vertices, faces; color = colors, shading = NoShading, kwargs...)
 end
 
 """
-    grid_points_to_faces(coord::AbstractVector)
+    grid_points_to_faces(coord::AbstractMKVector)
     grid_points_to_faces(coord::Observable{T} where T <: AbstractVector)
     grid_points_to_faces(coord::AbstractMatrix)
     grid_points_to_faces(coord::Observable{T} where T <: AbstractMatrix)
@@ -2608,7 +2613,7 @@ points between grid points.
 """
 function grid_points_to_faces end
 
-function grid_points_to_faces(coord::AbstractVector)
+function grid_points_to_faces(coord::AbstractMKVector)
     n = length(coord)
     faces = allocate_float(n+1)
     faces[1] = coord[1]
@@ -2620,7 +2625,7 @@ function grid_points_to_faces(coord::AbstractVector)
     return faces
 end
 
-function grid_points_to_faces(coord::Observable{T} where T <: AbstractVector)
+function grid_points_to_faces(coord::Observable{T} where T <: AbstractMKVector)
     n = length(coord.val)
     faces = allocate_float(n+1)
     faces[1] = coord.val[1]
@@ -2632,7 +2637,7 @@ function grid_points_to_faces(coord::Observable{T} where T <: AbstractVector)
     return faces
 end
 
-function grid_points_to_faces(coord::AbstractMatrix)
+function grid_points_to_faces(coord::AbstractMKMatrix)
     ni, nj = size(coord)
     faces = allocate_float(ni+1, nj+1)
     faces[1,1] = coord[1,1]
@@ -2656,7 +2661,7 @@ function grid_points_to_faces(coord::AbstractMatrix)
     return faces
 end
 
-function grid_points_to_faces(coord::Observable{T} where T <: AbstractMatrix)
+function grid_points_to_faces(coord::Observable{T} where T <: AbstractMKMatrix)
     ni, nj = size(coord.val)
     faces = allocate_float(ni+1, nj+1)
     faces[1,1] = coord.val[1,1]

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -33,7 +33,7 @@ using moment_kinetics.load_data: get_variable, timestep_diagnostic_variables,
 using moment_kinetics.type_definitions
 
 using LaTeXStrings
-using NaNMath
+import NaNMath
 using OrderedCollections
 using StatsBase
 

diff --git a/makie_post_processing/makie_post_processing/src/mms.jl b/makie_post_processing/makie_post_processing/src/mms.jl
@@ -171,7 +171,7 @@ function compare_moment_symbolic_test(run_info, plot_prefix, field_label, field_
     if !input.calculate_error_norms
         field_norm = nothing
     else
-        field_norm = zeros(mk_float,nt)
+        field_norm = mk_zeros(mk_float,nt)
         for it in 1:nt
             dummy = 0.0
             #dummy_N = 0.0
@@ -576,7 +576,7 @@ function compare_ion_pdf_symbolic_test(run_info, plot_prefix; io=nothing,
                 end
             end
         end
-        field_norm = zeros(mk_float,nt)
+        field_norm = mk_zeros(mk_float,nt)
         for it in 1:nt
             dummy = 0.0
             #dummy_N = 0.0
@@ -765,7 +765,7 @@ function compare_neutral_pdf_symbolic_test(run_info, plot_prefix; io=nothing,
                 end
             end
         end
-        field_norm = zeros(mk_float,nt)
+        field_norm = mk_zeros(mk_float,nt)
         for it in 1:nt
             dummy = 0.0
             #dummy_N = 0.0

diff --git a/makie_post_processing/makie_post_processing/src/moment_constraints.jl b/makie_post_processing/makie_post_processing/src/moment_constraints.jl
@@ -32,7 +32,7 @@ function check_moment_constraints(run_info, is_neutral; input, plot_prefix)
     if is_neutral
         fn = get_variable(run_info, "f_neutral")
         if run_info.evolve_density
-            moment = zeros(run_info.z.n, run_info.r.n, run_info.nt)
+            moment = mk_zeros(run_info.z.n, run_info.r.n, run_info.nt)
             for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n, iz ∈ 1:run_info.z.n
                 moment[iz,ir,it] = integrate_over_neutral_vspace(
                     @view(fn[:,:,:,iz,ir,is,it]), run_info.vz.grid, 0, run_info.vz.wgts,
@@ -45,7 +45,7 @@ function check_moment_constraints(run_info, is_neutral; input, plot_prefix)
         end
 
         if run_info.evolve_upar
-            moment = zeros(run_info.z.n, run_info.r.n, run_info.nt)
+            moment = mk_zeros(run_info.z.n, run_info.r.n, run_info.nt)
             for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n, iz ∈ 1:run_info.z.n
                 moment[iz,ir,it] = integrate_over_neutral_vspace(
                     @view(fn[:,:,:,iz,ir,is,it]), run_info.vz.grid, 1, run_info.vz.wgts,
@@ -58,7 +58,7 @@ function check_moment_constraints(run_info, is_neutral; input, plot_prefix)
         end
 
         if run_info.evolve_ppar
-            moment = zeros(run_info.z.n, run_info.r.n, run_info.nt)
+            moment = mk_zeros(run_info.z.n, run_info.r.n, run_info.nt)
             for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n, iz ∈ 1:run_info.z.n
                 moment[iz,ir,it] = integrate_over_neutral_vspace(
                     @view(fn[:,:,:,iz,ir,is,it]), run_info.vz.grid, 2, run_info.vz.wgts,
@@ -72,7 +72,7 @@ function check_moment_constraints(run_info, is_neutral; input, plot_prefix)
     else
         f = get_variable(run_info, "f")
         if run_info.evolve_density
-            moment = zeros(run_info.z.n, run_info.r.n, run_info.nt)
+            moment = mk_zeros(run_info.z.n, run_info.r.n, run_info.nt)
             for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n, iz ∈ 1:run_info.z.n
                 moment[iz,ir,it] = integrate_over_vspace(
                     @view(f[:,:,iz,ir,is,it]), run_info.vpa.grid, 0, run_info.vpa.wgts,
@@ -84,7 +84,7 @@ function check_moment_constraints(run_info, is_neutral; input, plot_prefix)
         end
 
         if run_info.evolve_upar
-            moment = zeros(run_info.z.n, run_info.r.n, run_info.nt)
+            moment = mk_zeros(run_info.z.n, run_info.r.n, run_info.nt)
             for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n, iz ∈ 1:run_info.z.n
                 moment[iz,ir,it] = integrate_over_vspace(
                     @view(f[:,:,iz,ir,is,it]), run_info.vpa.grid, 1, run_info.vpa.wgts,
@@ -96,7 +96,7 @@ function check_moment_constraints(run_info, is_neutral; input, plot_prefix)
         end
 
         if run_info.evolve_ppar
-            moment = zeros(run_info.z.n, run_info.r.n, run_info.nt)
+            moment = mk_zeros(run_info.z.n, run_info.r.n, run_info.nt)
             for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n, iz ∈ 1:run_info.z.n
                 moment[iz,ir,it] = integrate_over_vspace(
                     @view(f[:,:,iz,ir,is,it]), run_info.vpa.grid, 2, run_info.vpa.wgts,

diff --git a/makie_post_processing/makie_post_processing/src/shared_utils.jl b/makie_post_processing/makie_post_processing/src/shared_utils.jl
@@ -55,7 +55,7 @@ function calculate_and_write_frequencies(run_name, ntime, time, z, itime_min, it
         phase = 0.0
         shifted_time = allocate_float(ntime)
         @. shifted_time = time - time[itime_min]
-        fitted_delta_phi = zeros(ntime)
+        fitted_delta_phi = mk_zeros(ntime)
 
     end
     return frequency, growth_rate, shifted_time, fitted_delta_phi