Time JIT separately; update config scripts

scripts/config_defiant.sh

@@ -10,10 +10,15 @@ echo $MV_DIR
 module purge
 
 # load required modules
-module load PrgEnv-cray-amd
-module load cray-mpich
+module load craype-accel-amd-gfx908
+module load PrgEnv-cray
+module load amd
 module load julia
 
+export MPIR_CVAR_GPU_EAGER_DEVICE_MEM=0
+export MPICH_GPU_SUPPORT_ENABLED=1
+export MPICH_SMP_SINGLE_COPY_MODE=CMA
+
 # remove existing generated Manifest.toml
 rm -f $MV_DIR/Manifest.toml
 rm -f $MV_DIR/LocalPreferences.toml

scripts/config_milan0.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+
+MV_DIR=/home/4pf/MiniVATES.jl
+
+module load nvhpc
+
+#export CUDA_HOME=/sw/wombat/Nvidia_HPC_SDK/Linux_aarch64/24.5/cuda/12.4
+
+# rm -f $MV_DIR/Manifest.toml
+# rm -f $MV_DIR/LocalPreferences.toml
+
+julia --project=$MV_DIR -e 'using Pkg; Pkg.instantiate()'
+
+# MPI + CUDA
+julia --project=$MV_DIR -e ' \
+    using Pkg; \
+    if !haskey(Pkg.project().dependencies, "MPIPreferences"); \
+        Pkg.add("MPIPreferences"); \
+    end; \
+    using MPIPreferences; \
+    MPIPreferences.use_jll_binary(); \
+    using MPI; \
+    MPI.install_mpiexecjl(force = true); \
+    if !haskey(Pkg.project().dependencies, "CUDA"); \
+        Pkg.add("CUDA"); \
+    end; \
+    using CUDA; \
+    '
+
+# JACC
+julia --project=$MV_DIR -e ' \
+    using Pkg; \
+    jaccInfo = Pkg.dependencies()[Pkg.project().dependencies["JACC"]]; \
+    if jaccInfo.git_revision != "fix-cuda-thread-counts-2d"; \
+        Pkg.add(; name="JACC", url = "https://github.com/PhilipFackler/JACC.jl.git", rev = "fix-cuda-thread-counts-2d"); \
+    end; \
+    using JACC; \
+    JACC.JACCPreferences.set_backend("cuda"); \
+    '
+
+# Verify the packages are installed correctly
+julia --project=$MV_DIR -e 'using Pkg; Pkg.instantiate()'
+julia --project=$MV_DIR -e ' \
+    using Pkg; \
+    Pkg.status(); \
+    using CUDA; \
+    CUDA.versioninfo(); \
+    '

src/BinSeries.jl

@@ -55,6 +55,10 @@ tmfmt(tm::AbstractFloat) = @sprintf("%3.6fs", tm)
         @show rank, start, stop
     end
 
+    updAvgJ = 0.0
+    mdnAvgJ = 0.0
+    binAvgJ = 0.0
+
     updAvg = 0.0
     mdnAvg = 0.0
     binAvg = 0.0
@@ -100,26 +104,49 @@ tmfmt(tm::AbstractFloat) = @sprintf("%3.6fs", tm)
             end
             MPI.Barrier(comm)
         end
+
+        if fi == start
+            updAvgJ = updAvg
+            mdnAvgJ = mdnAvg
+            binAvgJ = binAvg
+            updAvg = 0.0
+            mdnAvg = 0.0
+            binAvg = 0.0
+        end
     end
+    sumJ = [updAvgJ, mdnAvgJ, binAvgJ]
+    MPI.Reduce!(sumJ, MPI.SUM, 0, comm)
     sum = [updAvg, mdnAvg, binAvg]
     MPI.Reduce!(sum, MPI.SUM, 0, comm)
     if rank == 0
+        avgJ = sumJ ./ commSize
         avg = sum ./ nFiles
         println("Averages:")
-        println("    updateEvents: ", tmfmt(avg[1]))
-        println("    mdNorm:       ", tmfmt(avg[2]))
-        println("    binEvents:    ", tmfmt(avg[3]))
+        println("    updateEvents (JIT): ", tmfmt(avgJ[1]))
+        println("    updateEvents:       ", tmfmt(avg[1]))
+        println("    mdNorm (JIT):       ", tmfmt(avgJ[2]))
+        println("    mdNorm:             ", tmfmt(avg[2]))
+        println("    binEvents (JIT):    ", tmfmt(avgJ[3]))
+        println("    binEvents:          ", tmfmt(avg[3]))
     end
 
     return (signal, eventsHist)
 end
 
 function mergeHistogramToRootProcess!(hist::Hist3)
     dur = @elapsed begin
-        weights = binweights(hist)
+        weights = Core.Array(binweights(hist))
         MPI.Reduce!(weights, MPI.SUM, 0, MPI.COMM_WORLD)
+        ret = Hist3(
+            edges(hist),
+            nbins(hist),
+            origin(hist),
+            boxLength(hist),
+            adapt_structure(JACC.Array, weights),
+        )
     end
     if MPI.Comm_rank(MPI.COMM_WORLD) == 0
         println("Reduce: ", tmfmt(dur))
     end
+    return ret
 end