Add mixed GPU RP path for EScalar

This commit is contained in:
2026-05-02 18:27:26 +08:00
parent 59a216ad93
commit f638cbc4e8
5 changed files with 280 additions and 23 deletions

View File

@@ -546,6 +546,98 @@ bool cuda_direct_unpack_segment(double *buffer,
return ok;
}
bool cuda_direct_pack_bssn_prefix_to_host(double *buffer,
const Parallel::gridseg *src,
const Parallel::gridseg *dst,
int type,
MyList<var> *VarLists,
int Symmetry)
{
#if USE_CUDA_BSSN
if (!buffer || !src || !dst || !src->Bg || !dst->Bg || !VarLists)
return false;
if (!cuda_can_direct_pack(src, dst, type, VarLists))
return false;
double *views[BSSN_CUDA_STATE_COUNT];
double soa_flat[3 * BSSN_CUDA_STATE_COUNT];
if (!cuda_build_bssn_host_views(src->Bg, VarLists, BSSN_CUDA_STATE_COUNT, views) ||
!cuda_build_state_soa(VarLists, BSSN_CUDA_STATE_COUNT, soa_flat))
return false;
const double t0 = sync_profile_enabled() ? MPI_Wtime() : 0.0;
bool ok = false;
if (type == 1)
{
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);
ok = bssn_cuda_pack_state_batch_to_host_buffer_for_host_views(
src->Bg, views, BSSN_CUDA_STATE_COUNT, buffer, src->Bg->shape,
i0, j0, k0,
dst->shape[0], dst->shape[1], dst->shape[2]) == 0;
}
else if (type == 2)
{
int first_fine[3];
if (!cuda_cell_gw3_restrict_params(src, dst, first_fine))
return false;
ok = bssn_cuda_restrict_state_batch_to_host_buffer_for_host_views(
src->Bg, views, BSSN_CUDA_STATE_COUNT, buffer, src->Bg->shape,
dst->shape[0], dst->shape[1], dst->shape[2],
first_fine[0], first_fine[1], first_fine[2],
soa_flat) == 0;
}
else if (type == 3)
{
int first_fine_ii[3], coarse_lb[3];
if (!cuda_cell_gw3_prolong_params(src, dst, first_fine_ii, coarse_lb))
return false;
ok = bssn_cuda_prolong_state_batch_to_host_buffer_for_host_views(
src->Bg, views, BSSN_CUDA_STATE_COUNT, buffer, src->Bg->shape,
dst->shape[0], dst->shape[1], dst->shape[2],
first_fine_ii[0], first_fine_ii[1], first_fine_ii[2],
coarse_lb[0], coarse_lb[1], coarse_lb[2],
soa_flat) == 0;
}
if (sync_profile_enabled())
sync_profile_stats().direct_pack_sec += MPI_Wtime() - t0;
(void)Symmetry;
return ok;
#else
(void)buffer; (void)src; (void)dst; (void)type; (void)VarLists; (void)Symmetry;
return false;
#endif
}
bool cuda_direct_unpack_bssn_prefix_from_host(double *buffer,
const Parallel::gridseg *dst,
int type,
MyList<var> *VarListd)
{
#if USE_CUDA_BSSN
if (!buffer || !dst || !dst->Bg || !VarListd)
return false;
if (!cuda_can_direct_unpack(dst, type, VarListd))
return false;
double *views[BSSN_CUDA_STATE_COUNT];
if (!cuda_build_bssn_host_views(dst->Bg, VarListd, BSSN_CUDA_STATE_COUNT, views))
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);
const bool ok = bssn_cuda_unpack_state_batch_from_host_buffer_for_host_views(
dst->Bg, views, BSSN_CUDA_STATE_COUNT, buffer, dst->Bg->shape,
i0, j0, k0,
dst->shape[0], dst->shape[1], dst->shape[2]) == 0;
if (sync_profile_enabled())
sync_profile_stats().direct_unpack_sec += MPI_Wtime() - t0;
return ok;
#else
(void)buffer; (void)dst; (void)type; (void)VarListd;
return false;
#endif
}
bool cuda_aware_mpi_enabled()
{
static int enabled = -1;
@@ -5276,6 +5368,7 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
{
#if USE_CUDA_BSSN || USE_CUDA_Z4C
bool handled_by_cuda = false;
int cuda_handled_count = state_count;
if (dir == PACK && (type == 1 || s_cuda_aware_pack_active) &&
cuda_state_count_direct_supported(state_count) &&
cuda_can_direct_pack(src->data, dst->data, type, VarLists))
@@ -5306,6 +5399,28 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
MPI_Abort(MPI_COMM_WORLD, 1);
}
}
#if USE_CUDA_BSSN
else if (!s_cuda_aware_pack_active &&
state_idx == 0 &&
state_count > BSSN_CUDA_STATE_COUNT &&
dir == PACK &&
cuda_direct_pack_bssn_prefix_to_host(data + size_out, src->data, dst->data,
type, VarLists, Symmetry))
{
handled_by_cuda = true;
cuda_handled_count = BSSN_CUDA_STATE_COUNT;
}
else if (!s_cuda_aware_pack_active &&
state_idx == 0 &&
state_count > BSSN_CUDA_STATE_COUNT &&
dir == UNPACK &&
cuda_direct_unpack_bssn_prefix_from_host(data + size_out, dst->data,
type, VarListd))
{
handled_by_cuda = true;
cuda_handled_count = BSSN_CUDA_STATE_COUNT;
}
#endif
if (!handled_by_cuda)
{
#if USE_CUDA_BSSN || USE_CUDA_Z4C
@@ -5408,8 +5523,8 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
}
else
{
size_out += (state_count - 1) * dst->data->shape[0] * dst->data->shape[1] * dst->data->shape[2];
while (varls->next && varld->next)
size_out += (cuda_handled_count - 1) * dst->data->shape[0] * dst->data->shape[1] * dst->data->shape[2];
for (int skipped = 1; skipped < cuda_handled_count && varls->next && varld->next; ++skipped)
{
varls = varls->next;
varld = varld->next;
@@ -7270,14 +7385,15 @@ void Parallel::prepare_inter_time_level(Patch *Pat,
if (myrank == cg->rank)
{
#if USE_CUDA_BSSN
bool bssn_prefix_done = false;
double *src1_views[BSSN_CUDA_STATE_COUNT];
double *src2_views[BSSN_CUDA_STATE_COUNT];
double *dst_views[BSSN_CUDA_STATE_COUNT];
const int state_count = cuda_state_var_count(VarList1, VarList2);
if (state_count == BSSN_CUDA_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) &&
if (state_count >= BSSN_CUDA_STATE_COUNT &&
cuda_build_bssn_host_views(cg, VarList1, BSSN_CUDA_STATE_COUNT, src1_views) &&
cuda_build_bssn_host_views(cg, VarList2, BSSN_CUDA_STATE_COUNT, src2_views) &&
cuda_build_bssn_host_views(cg, VarList3, BSSN_CUDA_STATE_COUNT, dst_views) &&
bssn_cuda_has_resident_state(cg) &&
bssn_cuda_resident_state_matches(cg, src1_views) &&
bssn_cuda_resident_state_matches(cg, src2_views) &&
@@ -7285,15 +7401,30 @@ void Parallel::prepare_inter_time_level(Patch *Pat,
src1_views, src2_views, 0, dst_views,
2, tindex) == 0)
{
if (BP == Pat->ble)
break;
BP = BP->next;
continue;
if (state_count == BSSN_CUDA_STATE_COUNT)
{
if (BP == Pat->ble)
break;
BP = BP->next;
continue;
}
bssn_prefix_done = true;
}
#endif
varl1 = VarList1;
varl2 = VarList2;
varl3 = VarList3;
#if USE_CUDA_BSSN
if (bssn_prefix_done)
{
for (int i = 0; i < BSSN_CUDA_STATE_COUNT && varl1 && varl2 && varl3; ++i)
{
varl1 = varl1->next;
varl2 = varl2->next;
varl3 = varl3->next;
}
}
#endif
while (varl1)
{
if (tindex == 0)
@@ -7347,16 +7478,17 @@ void Parallel::prepare_inter_time_level(Patch *Pat,
if (myrank == cg->rank)
{
#if USE_CUDA_BSSN
bool bssn_prefix_done = false;
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];
const int state_count = cuda_state_var_count(VarList1, VarList2);
if (state_count == BSSN_CUDA_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) &&
if (state_count >= BSSN_CUDA_STATE_COUNT &&
cuda_build_bssn_host_views(cg, VarList1, BSSN_CUDA_STATE_COUNT, src1_views) &&
cuda_build_bssn_host_views(cg, VarList2, BSSN_CUDA_STATE_COUNT, src2_views) &&
cuda_build_bssn_host_views(cg, VarList3, BSSN_CUDA_STATE_COUNT, src3_views) &&
cuda_build_bssn_host_views(cg, VarList4, BSSN_CUDA_STATE_COUNT, dst_views) &&
bssn_cuda_has_resident_state(cg) &&
bssn_cuda_resident_state_matches(cg, src1_views) &&
bssn_cuda_resident_state_matches(cg, src2_views) &&
@@ -7365,16 +7497,32 @@ void Parallel::prepare_inter_time_level(Patch *Pat,
src1_views, src2_views, src3_views, dst_views,
3, tindex) == 0)
{
if (BP == Pat->ble)
break;
BP = BP->next;
continue;
if (state_count == BSSN_CUDA_STATE_COUNT)
{
if (BP == Pat->ble)
break;
BP = BP->next;
continue;
}
bssn_prefix_done = true;
}
#endif
varl1 = VarList1;
varl2 = VarList2;
varl3 = VarList3;
varl4 = VarList4;
#if USE_CUDA_BSSN
if (bssn_prefix_done)
{
for (int i = 0; i < BSSN_CUDA_STATE_COUNT && varl1 && varl2 && varl3 && varl4; ++i)
{
varl1 = varl1->next;
varl2 = varl2->next;
varl3 = varl3->next;
varl4 = varl4->next;
}
}
#endif
while (varl1)
{
if (tindex == 0)