64-BitBrainstorm_2026/AMSS-NCKU
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize BSSN CUDA resident state and CUDA-aware MPI

Browse Source
This commit is contained in:
CGH0S7
2026-04-29 17:05:10 +08:00
parent a0dab90bcb
commit 22c1e7168b
7 changed files with 7461 additions and 6727 deletions
Download Patch File Download Diff File Expand all files Collapse all files

188
AMSS_NCKU_source/bssn_rhs_cuda.cu
View File

@@ -76,11 +76,36 @@ struct CudaProfileStats {
double output_ms;
};
enum RhsStageId {
RHS_STAGE_PREP = 0,
RHS_STAGE_DERIV1,
RHS_STAGE_METRIC,
RHS_STAGE_GAUGE_DERIV,
RHS_STAGE_GAMMA_CONTRACT,
RHS_STAGE_RICCI_DIFF,
RHS_STAGE_RICCI_FUSED,
RHS_STAGE_CHI,
RHS_STAGE_GAUGE_RHS,
RHS_STAGE_KODIS,
RHS_STAGE_CONSTRAINTS,
RHS_STAGE_COUNT
};
struct RhsStageProfileStats {
long long calls;
double ms[RHS_STAGE_COUNT];
};
static CudaProfileStats &cuda_profile_stats() {
static CudaProfileStats stats = {0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
return stats;
}
static RhsStageProfileStats &rhs_stage_profile_stats() {
static RhsStageProfileStats stats = {};
return stats;
}
static bool cuda_profile_enabled() {
static int enabled = -1;
if (enabled < 0) {
@@ -99,6 +124,24 @@ static int cuda_profile_every() {
return every;
}
static bool rhs_stage_timing_enabled() {
static int enabled = -1;
if (enabled < 0) {
const char *env = getenv("AMSS_GPU_STAGE_TIMING");
enabled = (env && atoi(env) != 0) ? 1 : 0;
}
return enabled != 0;
}
static int rhs_stage_timing_every() {
static int every = -1;
if (every < 0) {
const char *env = getenv("AMSS_GPU_STAGE_TIMING_EVERY");
every = (env && atoi(env) > 0) ? atoi(env) : cuda_profile_every();
}
return every;
}
static double cuda_profile_now_ms() {
using clock = std::chrono::steady_clock;
return std::chrono::duration<double, std::milli>(
@@ -131,6 +174,36 @@ static void cuda_profile_maybe_log() {
fflush(stderr);
}
static void rhs_stage_profile_accumulate(const double *stage_ms) {
if (!rhs_stage_timing_enabled()) return;
RhsStageProfileStats &stats = rhs_stage_profile_stats();
stats.calls++;
for (int i = 0; i < RHS_STAGE_COUNT; ++i) {
stats.ms[i] += stage_ms[i];
}
if (stats.calls <= 0 || stats.calls % rhs_stage_timing_every() != 0) return;
fprintf(stderr,
"[AMSS-CUDA-STAGE][rank %d][dev %d] calls=%lld"
" prep=%.3f deriv1=%.3f metric=%.3f gauge_deriv=%.3f"
" gamma_contract=%.3f ricci_diff=%.3f ricci_fused=%.3f"
" chi=%.3f gauge_rhs=%.3f kodis=%.3f constraints=%.3f ms\n",
g_dispatch.my_rank, g_dispatch.my_device, stats.calls,
stats.ms[RHS_STAGE_PREP] / (double)stats.calls,
stats.ms[RHS_STAGE_DERIV1] / (double)stats.calls,
stats.ms[RHS_STAGE_METRIC] / (double)stats.calls,
stats.ms[RHS_STAGE_GAUGE_DERIV] / (double)stats.calls,
stats.ms[RHS_STAGE_GAMMA_CONTRACT] / (double)stats.calls,
stats.ms[RHS_STAGE_RICCI_DIFF] / (double)stats.calls,
stats.ms[RHS_STAGE_RICCI_FUSED] / (double)stats.calls,
stats.ms[RHS_STAGE_CHI] / (double)stats.calls,
stats.ms[RHS_STAGE_GAUGE_RHS] / (double)stats.calls,
stats.ms[RHS_STAGE_KODIS] / (double)stats.calls,
stats.ms[RHS_STAGE_CONSTRAINTS] / (double)stats.calls);
fflush(stderr);
}
/* ------------------------------------------------------------------ */
/* Error checking */
/* ------------------------------------------------------------------ */
@@ -4643,6 +4716,20 @@ static void compute_patch_boundary_flags(int *ex,
static void upload_state_inputs(double **state_host, size_t all)
{
const size_t bytes = all * sizeof(double);
static int direct_upload = -1;
if (direct_upload < 0) {
const char *env = getenv("AMSS_CUDA_DIRECT_STATE_UPLOAD");
const char *pin_env = getenv("AMSS_CUDA_PIN_GRIDFUNCS");
direct_upload = env ? ((atoi(env) != 0) ? 1 : 0)
: ((pin_env && atoi(pin_env) != 0) ? 1 : 0);
}
if (direct_upload) {
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
CUDA_CHECK(cudaMemcpyAsync(g_buf.slot[k_state_input_slots[i]], state_host[i],
bytes, cudaMemcpyHostToDevice));
}
return;
}
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
std::memcpy(g_buf.h_stage + (size_t)i * all, state_host[i], bytes);
}
@@ -4697,11 +4784,24 @@ static void launch_rhs_pipeline(int all, double eps, int co)
{
const double SYM = 1.0;
const double ANTI = -1.0;
const bool stage_timing = rhs_stage_timing_enabled();
double stage_ms[RHS_STAGE_COUNT] = {};
double stage_t0 = stage_timing ? cuda_profile_now_ms() : 0.0;
#define D(s) g_buf.slot[s]
#define MARK_RHS_STAGE(stage_id) do { \
if (stage_timing) { \
cuda_profile_sync(); \
const double stage_t1 = cuda_profile_now_ms(); \
stage_ms[(stage_id)] += stage_t1 - stage_t0; \
stage_t0 = stage_t1; \
} \
} while (0)
kern_phase1_prep<<<grid(all),BLK>>>(
D(S_Lap), D(S_chi), D(S_dxx), D(S_dyy), D(S_dzz),
D(S_alpn1), D(S_chin1), D(S_gxx), D(S_gyy), D(S_gzz));
MARK_RHS_STAGE(RHS_STAGE_PREP);
{
double *src_fields[] = {
@@ -4742,6 +4842,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
src_fields, fx_fields, fy_fields, fz_fields,
soa_signs, all);
}
MARK_RHS_STAGE(RHS_STAGE_DERIV1);
kern_phase2_metric_rhs<<<grid(all),BLK>>>(
D(S_alpn1), D(S_chin1),
@@ -4799,6 +4900,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_Gamzxx), D(S_Gamzxy), D(S_Gamzxz),
D(S_Gamzyy), D(S_Gamzyz), D(S_Gamzzz),
D(S_Gamx_rhs), D(S_Gamy_rhs), D(S_Gamz_rhs));
MARK_RHS_STAGE(RHS_STAGE_METRIC);
{
double *src_fields[] = {D(S_betax), D(S_betay), D(S_betaz)};
@@ -4832,6 +4934,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
src_fields, fx_fields, fy_fields, fz_fields,
soa_signs, all);
}
MARK_RHS_STAGE(RHS_STAGE_GAUGE_DERIV);
kern_phase8_9_gamma_rhs_contract_fused<<<grid(all),BLK>>>(
D(S_gupxx), D(S_gupxy), D(S_gupxz),
@@ -4854,6 +4957,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_gxxx),D(S_gxyx),D(S_gxzx),D(S_gyyx),D(S_gyzx),D(S_gzzx),
D(S_gxxy),D(S_gxyy),D(S_gxzy),D(S_gyyy),D(S_gyzy),D(S_gzzy),
D(S_gxxz),D(S_gxyz),D(S_gxzz),D(S_gyyz),D(S_gyzz),D(S_gzzz));
MARK_RHS_STAGE(RHS_STAGE_GAMMA_CONTRACT);
{
double *src_fields[] = {D(S_dxx), D(S_dyy), D(S_dzz), D(S_gxy), D(S_gxz), D(S_gyz)};
@@ -4870,6 +4974,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_gupyy), D(S_gupyz), D(S_gupzz),
src_fields, dst_fields, soa_signs, all);
}
MARK_RHS_STAGE(RHS_STAGE_RICCI_DIFF);
kern_phase11_ricci_fused<<<grid(all),BLK>>>(
D(S_gxx),D(S_gxy),D(S_gxz),D(S_gyy),D(S_gyz),D(S_gzz),
@@ -4889,6 +4994,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_gxxz),D(S_gxyz),D(S_gxzz),D(S_gyyz),D(S_gyzz),D(S_gzzz),
D(S_Rxx),D(S_Rxy),D(S_Rxz),
D(S_Ryy),D(S_Ryz),D(S_Rzz));
MARK_RHS_STAGE(RHS_STAGE_RICCI_FUSED);
kern_phase12_13_chi_correction_fused<<<grid((size_t)all),BLK>>>(
D(S_chi), D(S_chin1),
@@ -4904,6 +5010,7 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_Gamzyy), D(S_Gamzyz), D(S_Gamzzz),
D(S_Rxx), D(S_Rxy), D(S_Rxz),
D(S_Ryy), D(S_Ryz), D(S_Rzz));
MARK_RHS_STAGE(RHS_STAGE_CHI);
kern_phase15_trK_Aij_gauge<<<grid(all),BLK>>>(
D(S_alpn1), D(S_chin1),
@@ -4936,8 +5043,10 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_betax_rhs), D(S_betay_rhs), D(S_betaz_rhs),
D(S_Gamx_rhs), D(S_Gamy_rhs), D(S_Gamz_rhs),
D(S_f_arr), D(S_S_arr));
MARK_RHS_STAGE(RHS_STAGE_GAUGE_RHS);
gpu_lopsided_kodis_state_batch(eps, all);
MARK_RHS_STAGE(RHS_STAGE_KODIS);
if (co == 0) {
{
@@ -4982,7 +5091,10 @@ static void launch_rhs_pipeline(int all, double eps, int co)
D(S_gzzx), D(S_gzzy), D(S_gzzz),
D(S_ham_Res), D(S_movx_Res), D(S_movy_Res), D(S_movz_Res));
}
MARK_RHS_STAGE(RHS_STAGE_CONSTRAINTS);
rhs_stage_profile_accumulate(stage_ms);
#undef MARK_RHS_STAGE
#undef D
}
@@ -5196,6 +5308,21 @@ static void download_resident_state(void *block_tag, int *ex, double **state_hos
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const size_t bytes = all * sizeof(double);
StepContext &ctx = ensure_step_ctx(block_tag, all);
static int direct_download = -1;
if (direct_download < 0) {
const char *env = getenv("AMSS_CUDA_DIRECT_STATE_DOWNLOAD");
const char *pin_env = getenv("AMSS_CUDA_PIN_GRIDFUNCS");
direct_download = env ? ((atoi(env) != 0) ? 1 : 0)
: ((pin_env && atoi(pin_env) != 0) ? 1 : 0);
}
if (direct_download) {
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
CUDA_CHECK(cudaMemcpyAsync(state_host_out[i], ctx.d_state_curr[i],
bytes, cudaMemcpyDeviceToHost));
}
CUDA_CHECK(cudaDeviceSynchronize());
return;
}
CUDA_CHECK(cudaMemcpy(g_buf.h_stage, ctx.d_state_curr_mem,
(size_t)BSSN_STATE_COUNT * bytes,
cudaMemcpyDeviceToHost));
@@ -5902,6 +6029,67 @@ int bssn_cuda_unpack_state_batch_from_host_buffer(void *block_tag,
return 0;
}
static void copy_state_device_batch(void *block_tag,
int state_count,
double *device_buffer,
const int *ex,
int i0, int j0, int k0,
int sx, int sy, int sz,
int pack_not_unpack)
{
if (state_count <= 0 || state_count > BSSN_STATE_COUNT) return;
if (sx <= 0 || sy <= 0 || sz <= 0) return;
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const int region_all = sx * sy * sz;
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
if (pack_not_unpack) {
kern_pack_state_region_batch<<<launch_grid, BLK>>>(
ctx.d_state_curr_mem, device_buffer,
ex[0], ex[1], i0, j0, k0, sx, sy, sz,
region_all, state_count,
ex[0] * ex[1] * ex[2]);
} else {
kern_unpack_state_region_batch<<<launch_grid, BLK>>>(
ctx.d_state_curr_mem, device_buffer,
ex[0], ex[1], i0, j0, k0, sx, sy, sz,
region_all, state_count,
ex[0] * ex[1] * ex[2]);
}
}
extern "C"
int bssn_cuda_pack_state_batch_to_device_buffer(void *block_tag,
int state_count,
double *device_buffer,
int *ex,
int i0, int j0, int k0,
int sx, int sy, int sz)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
copy_state_device_batch(block_tag, state_count, device_buffer, ex,
i0, j0, k0, sx, sy, sz, 1);
return 0;
}
extern "C"
int bssn_cuda_unpack_state_batch_from_device_buffer(void *block_tag,
int state_count,
double *device_buffer,
int *ex,
int i0, int j0, int k0,
int sx, int sy, int sz)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
copy_state_device_batch(block_tag, state_count, device_buffer, ex,
i0, j0, k0, sx, sy, sz, 0);
return 0;
}
extern "C"
int bssn_cuda_download_state_subset(void *block_tag,
int *ex,