Optimize BSSN-EScalar CUDA path

This commit is contained in:
2026-05-05 10:47:46 +08:00
parent 06f62dee36
commit 85fe29cc2e
9 changed files with 1821 additions and 276 deletions

View File

@@ -18,6 +18,7 @@
#endif
#if USE_CUDA_BSSN
#include "bssn_rhs_cuda.h"
#define AMSS_BSSN_CUDA_MAX_STATE_COUNT BSSN_ESCALAR_CUDA_STATE_COUNT
#endif
#if USE_CUDA_Z4C
#include "z4c_rhs_cuda.h"
@@ -179,10 +180,12 @@ bool cuda_build_bssn_host_views(Block *block,
int state_count,
double **views)
{
if (!block || !vars || !views || state_count != BSSN_CUDA_STATE_COUNT)
if (!block || !vars || !views ||
(state_count != BSSN_CUDA_STATE_COUNT &&
state_count != BSSN_ESCALAR_CUDA_STATE_COUNT))
return false;
MyList<var> *v = vars;
for (int i = 0; i < BSSN_CUDA_STATE_COUNT; ++i)
for (int i = 0; i < state_count; ++i)
{
if (!v)
return false;
@@ -196,10 +199,12 @@ bool cuda_build_bssn_soa(MyList<var> *vars,
int state_count,
double *soa_flat)
{
if (!vars || !soa_flat || state_count != BSSN_CUDA_STATE_COUNT)
if (!vars || !soa_flat ||
(state_count != BSSN_CUDA_STATE_COUNT &&
state_count != BSSN_ESCALAR_CUDA_STATE_COUNT))
return false;
MyList<var> *v = vars;
for (int i = 0; i < BSSN_CUDA_STATE_COUNT; ++i)
for (int i = 0; i < state_count; ++i)
{
if (!v)
return false;
@@ -317,7 +322,7 @@ bool cuda_state_count_direct_supported(int state_count)
#if USE_CUDA_Z4C && (ABEtype == 2)
return state_count == Z4C_CUDA_STATE_COUNT;
#elif USE_CUDA_BSSN
return state_count > 0 && state_count <= BSSN_CUDA_STATE_COUNT;
return state_count > 0 && state_count <= BSSN_ESCALAR_CUDA_STATE_COUNT;
#else
(void)state_count;
return false;
@@ -372,22 +377,68 @@ bool cuda_can_direct_unpack(const Parallel::gridseg *dst, int type)
#endif
}
bool cuda_amr_host_staged_enabled();
double *alloc_device_comm_buffer(int length);
void free_device_comm_buffer(double *&ptr);
bool cuda_direct_pack_segment_to_device(double *buffer,
const Parallel::gridseg *src,
const Parallel::gridseg *dst,
int state_count,
int type,
MyList<var> *VarLists,
int Symmetry);
bool cuda_direct_pack_segment(double *buffer,
const Parallel::gridseg *src,
const Parallel::gridseg *dst,
int state_count,
MyList<var> *VarLists)
int type,
MyList<var> *VarLists,
int Symmetry)
{
#if USE_CUDA_Z4C && (ABEtype == 2)
if (state_count != Z4C_CUDA_STATE_COUNT)
return false;
#elif USE_CUDA_BSSN
if (state_count <= 0 || state_count > BSSN_CUDA_STATE_COUNT)
if (state_count <= 0 || state_count > AMSS_BSSN_CUDA_MAX_STATE_COUNT)
return false;
#else
return false;
#endif
const double t0 = sync_profile_enabled() ? MPI_Wtime() : 0.0;
if (type == 2 || type == 3)
{
#if USE_CUDA_BSSN
if (!cuda_amr_host_staged_enabled())
return false;
const int region_all = dst->shape[0] * dst->shape[1] * dst->shape[2];
const int total = state_count * region_all;
static double *stage_dev = 0;
static int stage_cap = 0;
if (total > stage_cap)
{
free_device_comm_buffer(stage_dev);
stage_dev = alloc_device_comm_buffer(total);
stage_cap = total;
}
if (!cuda_direct_pack_segment_to_device(stage_dev, src, dst, state_count, type, VarLists, Symmetry))
return false;
cudaError_t cerr = cudaMemcpy(buffer, stage_dev, (size_t)total * sizeof(double), cudaMemcpyDeviceToHost);
if (cerr != cudaSuccess)
{
fprintf(stderr, "Parallel: CUDA host-staged AMR pack cudaMemcpy failed, err=%d\n", (int)cerr);
return false;
}
if (sync_profile_enabled())
sync_profile_stats().direct_pack_sec += MPI_Wtime() - t0;
return true;
#else
return false;
#endif
}
const int i0 = cuda_seg_begin(dst, src->Bg, 0);
const int j0 = cuda_seg_begin(dst, src->Bg, 1);
const int k0 = cuda_seg_begin(dst, src->Bg, 2);
@@ -396,7 +447,7 @@ bool cuda_direct_pack_segment(double *buffer,
i0, j0, k0,
dst->shape[0], dst->shape[1], dst->shape[2]) == 0;
#else
double *views[BSSN_CUDA_STATE_COUNT];
double *views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const bool have_views = cuda_build_bssn_host_views(src->Bg, VarLists, state_count, views);
const bool ok = have_views
? bssn_cuda_pack_state_batch_to_host_buffer_for_host_views(
@@ -422,7 +473,7 @@ bool cuda_direct_unpack_segment(double *buffer,
if (state_count != Z4C_CUDA_STATE_COUNT)
return false;
#elif USE_CUDA_BSSN
if (state_count <= 0 || state_count > BSSN_CUDA_STATE_COUNT)
if (state_count <= 0 || state_count > AMSS_BSSN_CUDA_MAX_STATE_COUNT)
return false;
#else
return false;
@@ -436,7 +487,7 @@ bool cuda_direct_unpack_segment(double *buffer,
i0, j0, k0,
dst->shape[0], dst->shape[1], dst->shape[2]) == 0;
#else
double *views[BSSN_CUDA_STATE_COUNT];
double *views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const bool have_views = cuda_build_bssn_host_views(dst->Bg, VarListd, state_count, views);
const bool ok = have_views
? bssn_cuda_unpack_state_batch_from_host_buffer_for_host_views(
@@ -464,6 +515,17 @@ bool cuda_aware_mpi_enabled()
return enabled != 0;
}
bool cuda_cached_device_buffers_enabled(int state_count)
{
#if USE_CUDA_BSSN
if (state_count == BSSN_ESCALAR_CUDA_STATE_COUNT)
return false;
#else
(void)state_count;
#endif
return cuda_aware_mpi_enabled();
}
bool cuda_amr_restrict_device_enabled()
{
static int enabled = -1;
@@ -486,6 +548,17 @@ bool cuda_amr_prolong_device_enabled()
return enabled != 0;
}
bool cuda_amr_host_staged_enabled()
{
static int enabled = -1;
if (enabled < 0)
{
const char *env = getenv("AMSS_CUDA_AMR_HOST_STAGED");
enabled = (env && atoi(env) != 0) ? 1 : 0;
}
return enabled != 0;
}
bool cuda_amr_restrict_compare_enabled()
{
static int enabled = -1;
@@ -627,12 +700,12 @@ bool cuda_direct_pack_segment_to_device(double *buffer,
}
#endif
#if USE_CUDA_BSSN
if (state_count <= 0 || state_count > BSSN_CUDA_STATE_COUNT)
if (state_count <= 0 || state_count > AMSS_BSSN_CUDA_MAX_STATE_COUNT)
return false;
const double t0 = sync_profile_enabled() ? MPI_Wtime() : 0.0;
bool ok = false;
double *views[BSSN_CUDA_STATE_COUNT];
double soa_flat[3 * BSSN_CUDA_STATE_COUNT];
double *views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double soa_flat[3 * AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const bool have_views = cuda_build_bssn_host_views(src->Bg, VarLists, state_count, views);
const bool have_soa = cuda_build_bssn_soa(VarLists, state_count, soa_flat);
if (type == 1)
@@ -812,13 +885,13 @@ bool cuda_direct_unpack_segment_from_device(double *buffer,
}
#endif
#if USE_CUDA_BSSN
if (state_count <= 0 || state_count > BSSN_CUDA_STATE_COUNT)
if (state_count <= 0 || state_count > AMSS_BSSN_CUDA_MAX_STATE_COUNT)
return false;
const double t0 = sync_profile_enabled() ? MPI_Wtime() : 0.0;
const int i0 = cuda_seg_begin(dst, dst->Bg, 0);
const int j0 = cuda_seg_begin(dst, dst->Bg, 1);
const int k0 = cuda_seg_begin(dst, dst->Bg, 2);
double *views[BSSN_CUDA_STATE_COUNT];
double *views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const bool have_views = cuda_build_bssn_host_views(dst->Bg, VarListd, state_count, views);
const bool ok = have_views
? bssn_cuda_unpack_state_batch_from_device_buffer_for_host_views(
@@ -843,12 +916,12 @@ bool cuda_download_resident_subset_to_host(Block *block,
int state_count)
{
#if USE_CUDA_BSSN
if (!block || state_count <= 0 || state_count > BSSN_CUDA_STATE_COUNT)
if (!block || state_count <= 0 || state_count > AMSS_BSSN_CUDA_MAX_STATE_COUNT)
return false;
if (bssn_cuda_has_resident_state(block) == 0)
return true;
int indices[BSSN_CUDA_STATE_COUNT];
double *views[BSSN_CUDA_STATE_COUNT];
int indices[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double *views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
MyList<var> *v = vars;
for (int i = 0; i < state_count; ++i)
{
@@ -871,7 +944,7 @@ bool cuda_unpack_host_region_to_resident(Block *block,
const Parallel::gridseg *dst)
{
#if USE_CUDA_BSSN
if (!block || !dst || state_index < 0 || state_index >= BSSN_CUDA_STATE_COUNT)
if (!block || !dst || state_index < 0 || state_index >= AMSS_BSSN_CUDA_MAX_STATE_COUNT)
return false;
if (bssn_cuda_has_resident_state(block) == 0)
return true;
@@ -895,7 +968,7 @@ bool cuda_device_state_count_supported(int state_count)
return true;
#endif
#if USE_CUDA_BSSN
return state_count > 0 && state_count <= BSSN_CUDA_STATE_COUNT;
return state_count > 0 && state_count <= AMSS_BSSN_CUDA_MAX_STATE_COUNT;
#else
(void)state_count;
return false;
@@ -915,8 +988,8 @@ bool cuda_flush_device_segment_batch(Block *block,
return true;
const int stride = (dir == PACK && type == 3) ? 11 : 8;
const int segment_count = (int)(meta.size() / stride);
double *views[BSSN_CUDA_STATE_COUNT];
double soa_flat[3 * BSSN_CUDA_STATE_COUNT];
double *views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double soa_flat[3 * AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const bool have_views = cuda_build_bssn_host_views(block, vars, state_count, views);
const bool have_soa = cuda_build_bssn_soa(vars, state_count, soa_flat);
if (dir == PACK)
@@ -5022,14 +5095,17 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
{
#if USE_CUDA_BSSN || USE_CUDA_Z4C
bool handled_by_cuda = false;
if (dir == PACK && (type == 1 || s_cuda_aware_pack_active) &&
const bool host_staged_amr =
dir == PACK && !s_cuda_aware_pack_active && (type == 2 || type == 3) &&
cuda_amr_host_staged_enabled();
if (dir == PACK && (type == 1 || s_cuda_aware_pack_active || host_staged_amr) &&
cuda_state_count_direct_supported(state_count) &&
cuda_can_direct_pack(src->data, dst->data, type))
{
if (s_cuda_aware_pack_active) {
handled_by_cuda = cuda_direct_pack_segment_to_device(data + size_out, src->data, dst->data, state_count, type, VarLists, Symmetry);
} else {
handled_by_cuda = cuda_direct_pack_segment(data + size_out, src->data, dst->data, state_count, VarLists);
handled_by_cuda = cuda_direct_pack_segment(data + size_out, src->data, dst->data, state_count, type, VarLists, Symmetry);
}
if (!handled_by_cuda)
{
@@ -5037,7 +5113,7 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
MPI_Abort(MPI_COMM_WORLD, 1);
}
}
else if (dir == UNPACK && (type == 1 || s_cuda_aware_pack_active) &&
else if (dir == UNPACK && (type == 1 || s_cuda_aware_pack_active || host_staged_amr) &&
cuda_state_count_direct_supported(state_count) &&
cuda_can_direct_unpack(dst->data, type))
{
@@ -5102,7 +5178,8 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
if (cuda_state_count_direct_supported(state_count) &&
dst->data && dst->data->Bg && bssn_cuda_has_resident_state(dst->data->Bg))
{
if (!cuda_unpack_host_region_to_resident(dst->data->Bg, state_idx, data + size_out, dst->data))
if (type != 2 && type != 3 &&
!cuda_unpack_host_region_to_resident(dst->data->Bg, state_idx, data + size_out, dst->data))
{
cout << "Parallel::data_packer: CUDA resident fallback upload failed." << endl;
MPI_Abort(MPI_COMM_WORLD, 1);
@@ -5775,7 +5852,7 @@ void Parallel::transfer_cached(MyList<Parallel::gridseg> **src, MyList<Parallel:
cout << "Parallel::transfer_cached: variable lists do not match." << endl;
MPI_Abort(MPI_COMM_WORLD, 1);
}
if (cuda_aware_mpi_enabled())
if (cuda_cached_device_buffers_enabled(state_count))
{
for (int n = 0; n < cpusize; n++)
{
@@ -6094,7 +6171,7 @@ void Parallel::Sync_start(MyList<Patch> *PatL, MyList<var> *VarList, int Symmetr
cout << "Parallel::Sync_start: variable lists do not match." << endl;
MPI_Abort(MPI_COMM_WORLD, 1);
}
if (cuda_aware_mpi_enabled())
if (cuda_cached_device_buffers_enabled(state_count))
{
for (int n = 0; n < cpusize; n++)
{
@@ -6976,16 +7053,16 @@ void Parallel::prepare_inter_time_level(Patch *Pat,
if (myrank == cg->rank)
{
#if USE_CUDA_BSSN
double *src1_views[BSSN_CUDA_STATE_COUNT];
double *src2_views[BSSN_CUDA_STATE_COUNT];
double *dst_views[BSSN_CUDA_STATE_COUNT];
double *src1_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double *src2_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double *dst_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const int state_count = cuda_state_var_count(VarList1, VarList2);
if (state_count == BSSN_CUDA_STATE_COUNT &&
if (cuda_state_count_direct_supported(state_count) &&
cuda_build_bssn_host_views(cg, VarList1, state_count, src1_views) &&
cuda_build_bssn_host_views(cg, VarList2, state_count, src2_views) &&
cuda_build_bssn_host_views(cg, VarList3, state_count, dst_views) &&
bssn_cuda_has_resident_state(cg) &&
bssn_cuda_prepare_inter_time_level(cg, cg->shape,
bssn_cuda_prepare_inter_time_level(cg, cg->shape, state_count,
src1_views, src2_views, 0, dst_views,
2, tindex) == 0)
{
@@ -7051,18 +7128,18 @@ void Parallel::prepare_inter_time_level(Patch *Pat,
if (myrank == cg->rank)
{
#if USE_CUDA_BSSN
double *src1_views[BSSN_CUDA_STATE_COUNT];
double *src2_views[BSSN_CUDA_STATE_COUNT];
double *src3_views[BSSN_CUDA_STATE_COUNT];
double *dst_views[BSSN_CUDA_STATE_COUNT];
double *src1_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double *src2_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double *src3_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
double *dst_views[AMSS_BSSN_CUDA_MAX_STATE_COUNT];
const int state_count = cuda_state_var_count(VarList1, VarList2);
if (state_count == BSSN_CUDA_STATE_COUNT &&
if (cuda_state_count_direct_supported(state_count) &&
cuda_build_bssn_host_views(cg, VarList1, state_count, src1_views) &&
cuda_build_bssn_host_views(cg, VarList2, state_count, src2_views) &&
cuda_build_bssn_host_views(cg, VarList3, state_count, src3_views) &&
cuda_build_bssn_host_views(cg, VarList4, state_count, dst_views) &&
bssn_cuda_has_resident_state(cg) &&
bssn_cuda_prepare_inter_time_level(cg, cg->shape,
bssn_cuda_prepare_inter_time_level(cg, cg->shape, state_count,
src1_views, src2_views, src3_views, dst_views,
3, tindex) == 0)
{
@@ -7500,6 +7577,8 @@ void Parallel::Restrict_cached(MyList<Patch> *PatcL, MyList<Patch> *PatfL,
cache.tc_req_is_recv = new int[cache.max_reqs];
cache.tc_completed = new int[cache.max_reqs];
}
for (int i = 0; i < cpusize; i++)
cache.combined_src[i] = cache.combined_dst[i] = 0;
MyList<Parallel::gridseg> *dst = build_complete_gsl(PatcL);
for (int node = 0; node < cpusize; node++)
@@ -7561,6 +7640,8 @@ void Parallel::OutBdLow2Hi_cached(MyList<Patch> *PatcL, MyList<Patch> *PatfL,
cache.tc_req_is_recv = new int[cache.max_reqs];
cache.tc_completed = new int[cache.max_reqs];
}
for (int i = 0; i < cpusize; i++)
cache.combined_src[i] = cache.combined_dst[i] = 0;
MyList<Parallel::gridseg> *dst = build_buffer_gsl(PatfL);
for (int node = 0; node < cpusize; node++)
@@ -7613,6 +7694,8 @@ void Parallel::OutBdLow2Himix_cached(MyList<Patch> *PatcL, MyList<Patch> *PatfL,
cache.tc_req_is_recv = new int[cache.max_reqs];
cache.tc_completed = new int[cache.max_reqs];
}
for (int i = 0; i < cpusize; i++)
cache.combined_src[i] = cache.combined_dst[i] = 0;
MyList<Parallel::gridseg> *dst = build_buffer_gsl(PatfL);
for (int node = 0; node < cpusize; node++)