Refine GPU runtime controls and input checker

This commit is contained in:
2026-05-18 01:02:55 +08:00
parent f2264989d8
commit a99534d2f3
8 changed files with 514 additions and 67 deletions

View File

@@ -75,6 +75,13 @@ def _input_or_env(input_name, env_name, default=None):
return getattr(input_data, input_name, default)
def _input_env_passthrough(runtime_env, env_name):
if env_name in runtime_env:
return
if hasattr(input_data, env_name):
runtime_env[env_name] = str(getattr(input_data, env_name))
def _start_cuda_mps_if_requested(runtime_env):
if input_data.GPU_Calculation != "yes":
return False
@@ -193,6 +200,40 @@ def _gpu_runtime_env():
for key, value in defaults.items():
runtime_env.setdefault(key, value)
passthrough_envs = [
"AMSS_CUDA_RESIDENT_SYNC",
"AMSS_CUDA_BSSN_RESIDENT_SYNC",
"AMSS_CUDA_EM_RESIDENT_SYNC",
"AMSS_CUDA_ESCALAR_RESIDENT_SYNC",
"AMSS_CUDA_BH_INTERP_RESIDENT",
"AMSS_CUDA_KEEP_RESIDENT_AFTER_STEP",
"AMSS_CUDA_KEEP_ALL_LEVELS",
"AMSS_CUDA_EM_KEEP_RESIDENT_AFTER_STEP",
"AMSS_CUDA_EM_KEEP_ALL_LEVELS",
"AMSS_CUDA_ESCALAR_KEEP_RESIDENT_AFTER_STEP",
"AMSS_CUDA_ESCALAR_KEEP_ALL_LEVELS",
"AMSS_CUDA_AMR_HOST_STAGED",
"AMSS_CUDA_AMR_RESTRICT_DEVICE",
"AMSS_CUDA_AMR_RESTRICT_BATCH",
"AMSS_CUDA_DEVICE_SEGMENT_BATCH",
"AMSS_CUDA_UNCACHED_DEVICE_BUFFERS",
"AMSS_CUDA_EM_CACHE_SOURCES",
"AMSS_CUDA_EM_ZERO_FASTPATH",
"AMSS_CUDA_AWARE_MPI",
"AMSS_CUDA_REGRID_FLUSH_ALWAYS",
"AMSS_Z4C_CUDA_RESIDENT",
"AMSS_SHELL_FAST_INTERP",
"AMSS_SHELL_PARALLEL_INTERP",
"AMSS_SHELL_CUDA_INTERP",
"AMSS_SHELL_INTERP_THREADS",
"AMSS_EM_ZERO_ANALYSIS_FASTPATH",
"AMSS_EM_ZERO_RESIDENT_DOWNLOAD_FASTPATH",
"AMSS_INTERP_FAST",
"AMSS_INTERP_GPU",
]
for env_name in passthrough_envs:
_input_env_passthrough(runtime_env, env_name)
optional_overrides = {
"AMSS_INTERP_FAST_COMPARE": "AMSS_Interp_Fast_Compare",
"AMSS_INTERP_FAST_COMPARE_LIMIT": "AMSS_Interp_Fast_Compare_Limit",
@@ -221,11 +262,13 @@ def makefile_ABE():
print( " Compiling the AMSS-NCKU executable file ABE/ABEGPU " )
print( )
z4c_mrbd = int(getattr(input_data, "AMSS_Z4C_MRBD", 0))
## Build command with CPU binding to nohz_full cores
if (input_data.GPU_Calculation == "no"):
makefile_command = f"{NUMACTL_CPU_BIND} make -j{BUILD_JOBS} INTERP_LB_MODE=off USE_CUDA_BSSN=0 USE_CUDA_Z4C=0 ABE"
makefile_command = f"{NUMACTL_CPU_BIND} env AMSS_Z4C_MRBD={z4c_mrbd} make -j{BUILD_JOBS} INTERP_LB_MODE=off USE_CUDA_BSSN=0 USE_CUDA_Z4C=0 ABE"
elif (input_data.GPU_Calculation == "yes"):
makefile_command = f"{NUMACTL_CPU_BIND} make -j{BUILD_JOBS} INTERP_LB_MODE=off USE_CUDA_BSSN=1 USE_CUDA_Z4C=1 ABE_CUDA"
makefile_command = f"{NUMACTL_CPU_BIND} env AMSS_Z4C_MRBD={z4c_mrbd} make -j{BUILD_JOBS} INTERP_LB_MODE=off USE_CUDA_BSSN=1 USE_CUDA_Z4C=1 ABE_CUDA"
else:
print( " CPU/GPU numerical calculation setting is wrong " )
print( )