Optimize Z4C GPU runtime defaults

This commit is contained in:
2026-05-07 15:37:09 +08:00
parent 83afaf19ce
commit 96829d0441
4 changed files with 114 additions and 45 deletions

View File

@@ -160,6 +160,11 @@ def _gpu_runtime_env():
"AMSS_CUDA_DEVICE_SEGMENT_BATCH": "0",
"AMSS_CUDA_UNCACHED_DEVICE_BUFFERS": "1",
}
if getattr(input_data, "Equation_Class", "") == "Z4C":
defaults.update({
"AMSS_Z4C_CUDA_RESIDENT": "1",
"AMSS_CONSTRAINT_OUT_EVERY": "1000000",
})
for key, value in defaults.items():
runtime_env.setdefault(key, value)
@@ -271,16 +276,26 @@ def run_ABE():
mpi_env = None
started_mps = False
mpi_processes = int(input_data.MPI_processes)
if (input_data.GPU_Calculation == "yes" and
getattr(input_data, "Equation_Class", "") == "Z4C"):
z4c_env_np = os.environ.get("AMSS_Z4C_GPU_MPI_PROCESSES")
if z4c_env_np and int(z4c_env_np) > 0:
mpi_processes = int(z4c_env_np)
elif mpi_processes < 4:
mpi_processes = 4
if (input_data.GPU_Calculation == "no"):
mpi_command = NUMACTL_CPU_BIND + " mpirun -np " + str(input_data.MPI_processes) + " ./ABE"
mpi_command = NUMACTL_CPU_BIND + " mpirun -np " + str(mpi_processes) + " ./ABE"
#mpi_command = " mpirun -np " + str(input_data.MPI_processes) + " ./ABE"
mpi_command_outfile = "ABE_out.log"
elif (input_data.GPU_Calculation == "yes"):
mpi_command = NUMACTL_CPU_BIND + " I_MPI_OFFLOAD=1 I_MPI_OFFLOAD_IPC=0 mpirun -np " + str(input_data.MPI_processes) + " ./ABE_CUDA"
mpi_command = NUMACTL_CPU_BIND + " I_MPI_OFFLOAD=1 I_MPI_OFFLOAD_IPC=0 mpirun -np " + str(mpi_processes) + " ./ABE_CUDA"
mpi_command_outfile = "ABEGPU_out.log"
mpi_env = _gpu_runtime_env()
started_mps = _start_cuda_mps_if_requested(mpi_env)
print(" GPU optimized runtime switches:")
print(f" MPI processes={mpi_processes}")
print(f" AMSS_INTERP_FAST={mpi_env.get('AMSS_INTERP_FAST', '')}")
print(f" AMSS_INTERP_GPU={mpi_env.get('AMSS_INTERP_GPU', '')}")
print(f" AMSS_ANALYSIS_MAP_EVERY={mpi_env.get('AMSS_ANALYSIS_MAP_EVERY', '')}")
@@ -300,6 +315,8 @@ def run_ABE():
print(f" AMSS_CUDA_AMR_RESTRICT_BATCH={mpi_env.get('AMSS_CUDA_AMR_RESTRICT_BATCH', '')}")
print(f" AMSS_CUDA_DEVICE_SEGMENT_BATCH={mpi_env.get('AMSS_CUDA_DEVICE_SEGMENT_BATCH', '')}")
print(f" AMSS_CUDA_UNCACHED_DEVICE_BUFFERS={mpi_env.get('AMSS_CUDA_UNCACHED_DEVICE_BUFFERS', '')}")
print(f" AMSS_Z4C_CUDA_RESIDENT={mpi_env.get('AMSS_Z4C_CUDA_RESIDENT', '')}")
print(f" AMSS_CONSTRAINT_OUT_EVERY={mpi_env.get('AMSS_CONSTRAINT_OUT_EVERY', '')}")
if "CUDA_MPS_PIPE_DIRECTORY" in mpi_env:
print(f" CUDA_MPS_PIPE_DIRECTORY={mpi_env['CUDA_MPS_PIPE_DIRECTORY']}")