Checkpoint Z4C CUDA optimization progress

This commit is contained in:
2026-05-02 08:55:25 +08:00
parent a5c8188305
commit fcd98649f6
4 changed files with 180 additions and 24 deletions

View File

@@ -193,6 +193,28 @@ bool cuda_build_bssn_host_views(Block *block,
}
#endif
#if USE_CUDA_Z4C && (ABEtype == 2)
bool cuda_build_z4c_host_views(Block *block,
MyList<var> *vars,
int state_count,
double **views)
{
if (!block || !vars || !views || state_count != Z4C_CUDA_STATE_COUNT)
return false;
MyList<var> *v = vars;
for (int i = 0; i < Z4C_CUDA_STATE_COUNT; ++i)
{
if (!v)
return false;
views[i] = block->fgfs[v->data->sgfn];
if (!views[i])
return false;
v = v->next;
}
return v == 0;
}
#endif
bool cuda_build_state_soa(MyList<var> *vars,
int state_count,
double *soa_flat)
@@ -220,6 +242,8 @@ int fortran_idint(double x)
bool cuda_amr_restrict_device_enabled();
bool cuda_amr_prolong_device_enabled();
bool cuda_z4c_amr_prolong_device_enabled();
bool cuda_z4c_amr_unpack_device_enabled();
bool cuda_amr_restrict_compare_enabled();
bool cuda_amr_restrict_batch_enabled();
bool cuda_device_segment_batch_enabled();
@@ -324,7 +348,8 @@ bool cuda_state_count_direct_supported(int state_count)
#endif
}
bool cuda_can_direct_pack(const Parallel::gridseg *src, const Parallel::gridseg *dst, int type)
bool cuda_can_direct_pack(const Parallel::gridseg *src, const Parallel::gridseg *dst,
int type, MyList<var> *VarLists = 0)
{
if (!src || !dst || !src->Bg)
return false;
@@ -342,7 +367,7 @@ bool cuda_can_direct_pack(const Parallel::gridseg *src, const Parallel::gridseg
}
if (type == 2 && !cuda_amr_restrict_device_enabled())
return false;
if (type == 3 && !cuda_amr_prolong_device_enabled())
if (type == 3 && (!cuda_amr_prolong_device_enabled() || !cuda_z4c_amr_prolong_device_enabled()))
return false;
if (type == 2) {
int a[3];
@@ -354,7 +379,17 @@ bool cuda_can_direct_pack(const Parallel::gridseg *src, const Parallel::gridseg
if (!cuda_cell_gw3_prolong_params(src, dst, a, b))
return false;
}
return z4c_cuda_has_resident_state(src->Bg) != 0;
if (z4c_cuda_has_resident_state(src->Bg) == 0)
return false;
if (type != 1 && VarLists)
{
double *view_ptrs[Z4C_CUDA_STATE_COUNT];
if (!cuda_build_z4c_host_views(src->Bg, VarLists, Z4C_CUDA_STATE_COUNT, view_ptrs))
return false;
if (z4c_cuda_resident_state_matches(src->Bg, view_ptrs) == 0)
return false;
}
return true;
#elif USE_CUDA_BSSN
if (bssn_cuda_has_resident_state(src->Bg) == 0)
return false;
@@ -380,12 +415,24 @@ bool cuda_can_direct_pack(const Parallel::gridseg *src, const Parallel::gridseg
#endif
}
bool cuda_can_direct_unpack(const Parallel::gridseg *dst, int type)
bool cuda_can_direct_unpack(const Parallel::gridseg *dst, int type, MyList<var> *VarListd = 0)
{
if (type < 1 || type > 3 || !dst || !dst->Bg)
return false;
#if USE_CUDA_Z4C && (ABEtype == 2)
return z4c_cuda_has_resident_state(dst->Bg) != 0;
if (type != 1 && !cuda_z4c_amr_unpack_device_enabled())
return false;
if (z4c_cuda_has_resident_state(dst->Bg) == 0)
return false;
if (type != 1 && VarListd)
{
double *view_ptrs[Z4C_CUDA_STATE_COUNT];
if (!cuda_build_z4c_host_views(dst->Bg, VarListd, Z4C_CUDA_STATE_COUNT, view_ptrs))
return false;
if (z4c_cuda_resident_state_matches(dst->Bg, view_ptrs) == 0)
return false;
}
return true;
#elif USE_CUDA_BSSN
return bssn_cuda_has_resident_state(dst->Bg) != 0;
#else
@@ -507,6 +554,28 @@ bool cuda_amr_prolong_device_enabled()
return enabled != 0;
}
bool cuda_z4c_amr_prolong_device_enabled()
{
static int enabled = -1;
if (enabled < 0)
{
const char *env = getenv("AMSS_CUDA_Z4C_AMR_PROLONG_DEVICE");
enabled = (env && atoi(env) != 0) ? 1 : 0;
}
return enabled != 0;
}
bool cuda_z4c_amr_unpack_device_enabled()
{
static int enabled = -1;
if (enabled < 0)
{
const char *env = getenv("AMSS_CUDA_Z4C_AMR_UNPACK_DEVICE");
enabled = (env && atoi(env) != 0) ? 1 : 0;
}
return enabled != 0;
}
bool cuda_amr_restrict_compare_enabled()
{
static int enabled = -1;
@@ -1109,8 +1178,8 @@ int cuda_data_packer_device_batched(double *data,
return -1;
Block *block = (dir == PACK) ? src->data->Bg : dst->data->Bg;
if ((dir == PACK && !cuda_can_direct_pack(src->data, dst->data, type)) ||
(dir == UNPACK && !cuda_can_direct_unpack(dst->data, type)))
if ((dir == PACK && !cuda_can_direct_pack(src->data, dst->data, type, VarLists)) ||
(dir == UNPACK && !cuda_can_direct_unpack(dst->data, type, VarListd)))
return -1;
if (batch_block && (batch_block != block || batch_type != type))
@@ -1195,7 +1264,9 @@ bool cuda_segments_device_eligible(MyList<Parallel::gridseg> *src,
int rank_in,
int dir,
int myrank,
int state_count)
int state_count,
MyList<var> *VarLists,
MyList<var> *VarListd)
{
bool has_work = false;
while (src && dst)
@@ -1215,12 +1286,12 @@ bool cuda_segments_device_eligible(MyList<Parallel::gridseg> *src,
type = 3;
if (dir == PACK)
{
if (!cuda_can_direct_pack(src->data, dst->data, type))
if (!cuda_can_direct_pack(src->data, dst->data, type, VarLists))
return false;
}
else
{
if (!cuda_can_direct_unpack(dst->data, type))
if (!cuda_can_direct_unpack(dst->data, type, VarListd))
return false;
}
}
@@ -1318,11 +1389,13 @@ bool cuda_pack_to_device_eligible(MyList<Parallel::gridseg> *src,
MyList<Parallel::gridseg> *dst,
int rank_in,
int state_count,
int myrank)
int myrank,
MyList<var> *VarLists,
MyList<var> *VarListd)
{
if (!cuda_aware_mpi_enabled() || !cuda_device_state_count_supported(state_count))
return false;
if (!cuda_segments_device_eligible(src, dst, rank_in, PACK, myrank, state_count))
if (!cuda_segments_device_eligible(src, dst, rank_in, PACK, myrank, state_count, VarLists, VarListd))
return false;
return true;
}
@@ -1331,11 +1404,13 @@ bool cuda_recv_to_device_eligible(MyList<Parallel::gridseg> *src,
MyList<Parallel::gridseg> *dst,
int rank_in,
int state_count,
int myrank)
int myrank,
MyList<var> *VarLists,
MyList<var> *VarListd)
{
if (!cuda_aware_mpi_enabled() || !cuda_device_state_count_supported(state_count))
return false;
if (!cuda_segments_device_eligible(src, dst, rank_in, UNPACK, myrank, state_count))
if (!cuda_segments_device_eligible(src, dst, rank_in, UNPACK, myrank, state_count, VarLists, VarListd))
return false;
return true;
}
@@ -5133,7 +5208,7 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
bool handled_by_cuda = false;
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))
cuda_can_direct_pack(src->data, dst->data, type, VarLists))
{
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);
@@ -5148,7 +5223,7 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
}
else if (dir == UNPACK && (type == 1 || s_cuda_aware_pack_active) &&
cuda_state_count_direct_supported(state_count) &&
cuda_can_direct_unpack(dst->data, type))
cuda_can_direct_unpack(dst->data, type, VarListd))
{
if (s_cuda_aware_pack_active) {
handled_by_cuda = cuda_direct_unpack_segment_from_device(data + size_out, dst->data, state_count, VarListd);
@@ -5229,6 +5304,24 @@ 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)
{
#if USE_CUDA_Z4C && (ABEtype == 2)
const char *z4c_amr_env = getenv("AMSS_CUDA_Z4C_AMR_DEVICE");
if (state_count == Z4C_CUDA_STATE_COUNT && type != 1 &&
z4c_amr_env && atoi(z4c_amr_env) != 0)
{
double *views[Z4C_CUDA_STATE_COUNT];
if (cuda_build_z4c_host_views(dst->data->Bg, VarListd, state_count, views) &&
z4c_cuda_resident_state_matches(dst->data->Bg, views) != 0)
{
if (!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);
}
}
}
else
#endif
if (!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;
@@ -5906,8 +5999,8 @@ void Parallel::transfer_cached(MyList<Parallel::gridseg> **src, MyList<Parallel:
{
for (int n = 0; n < cpusize; n++)
{
cache.send_buf_is_dev[n] = cuda_pack_to_device_eligible(src[myrank], dst[myrank], n, state_count, myrank) ? 1 : 0;
cache.recv_buf_is_dev[n] = cuda_recv_to_device_eligible(src[n], dst[n], n, state_count, myrank) ? 1 : 0;
cache.send_buf_is_dev[n] = cuda_pack_to_device_eligible(src[myrank], dst[myrank], n, state_count, myrank, VarList1, VarList2) ? 1 : 0;
cache.recv_buf_is_dev[n] = cuda_recv_to_device_eligible(src[n], dst[n], n, state_count, myrank, VarList1, VarList2) ? 1 : 0;
}
cache.recv_buf_is_dev[myrank] = (cache.send_buf_is_dev[myrank] && cache.recv_buf_is_dev[myrank]) ? 1 : 0;
for (int n = 0; n < cpusize; n++)
@@ -6225,8 +6318,8 @@ void Parallel::Sync_start(MyList<Patch> *PatL, MyList<var> *VarList, int Symmetr
{
for (int n = 0; n < cpusize; n++)
{
cache.send_buf_is_dev[n] = cuda_pack_to_device_eligible(src[myrank], dst[myrank], n, state_count, myrank) ? 1 : 0;
cache.recv_buf_is_dev[n] = cuda_recv_to_device_eligible(src[n], dst[n], n, state_count, myrank) ? 1 : 0;
cache.send_buf_is_dev[n] = cuda_pack_to_device_eligible(src[myrank], dst[myrank], n, state_count, myrank, VarList, VarList) ? 1 : 0;
cache.recv_buf_is_dev[n] = cuda_recv_to_device_eligible(src[n], dst[n], n, state_count, myrank, VarList, VarList) ? 1 : 0;
}
cache.recv_buf_is_dev[myrank] = (cache.send_buf_is_dev[myrank] && cache.recv_buf_is_dev[myrank]) ? 1 : 0;
for (int n = 0; n < cpusize; n++)

View File

@@ -679,7 +679,11 @@ void Z4c_class::Step(int lev, int YN)
MPI_Abort(MPI_COMM_WORLD, 1);
}
Parallel::Sync_cached(GH->PatL[lev], SynchList_pre, Symmetry, sync_cache_pre[lev]);
{
Parallel::AsyncSyncState async_pre;
Parallel::Sync_start(GH->PatL[lev], SynchList_pre, Symmetry, sync_cache_pre[lev], async_pre);
Parallel::Sync_finish(sync_cache_pre[lev], async_pre, SynchList_pre, Symmetry);
}
if ((lev == a_lev) && (LastAnas + dT_lev >= AnasTime))
z4c_cuda_download_level_state(GH->PatL[lev], SynchList_pre, myrank, false);
@@ -778,7 +782,11 @@ void Z4c_class::Step(int lev, int YN)
MPI_Abort(MPI_COMM_WORLD, 1);
}
Parallel::Sync_cached(GH->PatL[lev], SynchList_cor, Symmetry, sync_cache_cor[lev]);
{
Parallel::AsyncSyncState async_cor;
Parallel::Sync_start(GH->PatL[lev], SynchList_cor, Symmetry, sync_cache_cor[lev], async_cor);
Parallel::Sync_finish(sync_cache_cor[lev], async_cor, SynchList_cor, Symmetry);
}
if (BH_num > 0 && lev == GH->levels - 1)
{

View File

@@ -462,25 +462,28 @@ struct StepContext {
std::array<double *, BSSN_STATE_COUNT> d_state_curr;
std::array<double *, BSSN_STATE_COUNT> d_state_next;
std::array<double *, BSSN_MATTER_COUNT> d_matter;
std::array<double *, BSSN_STATE_COUNT> resident_host;
size_t cap_all;
size_t cap_comm;
bool h_comm_pinned;
size_t cap_h_comm;
bool matter_ready;
bool state_ready;
bool resident_host_valid;
StepContext()
: d_state0_mem(nullptr), d_accum_mem(nullptr),
d_state_curr_mem(nullptr), d_state_next_mem(nullptr),
d_matter_mem(nullptr), d_comm_mem(nullptr), h_comm_mem(nullptr),
cap_all(0), cap_comm(0), h_comm_pinned(false), cap_h_comm(0),
matter_ready(false), state_ready(false)
matter_ready(false), state_ready(false), resident_host_valid(false)
{
d_state0.fill(nullptr);
d_accum.fill(nullptr);
d_state_curr.fill(nullptr);
d_state_next.fill(nullptr);
d_matter.fill(nullptr);
resident_host.fill(nullptr);
}
};
@@ -544,6 +547,8 @@ static StepAllocation detach_step_allocation(StepContext &ctx)
ctx.d_state_curr.fill(nullptr);
ctx.d_state_next.fill(nullptr);
ctx.d_matter.fill(nullptr);
ctx.resident_host.fill(nullptr);
ctx.resident_host_valid = false;
return alloc;
}
@@ -562,6 +567,8 @@ static void attach_step_allocation(StepContext &ctx, const StepAllocation &alloc
ctx.cap_h_comm = alloc.cap_h_comm;
ctx.matter_ready = false;
ctx.state_ready = false;
ctx.resident_host.fill(nullptr);
ctx.resident_host_valid = false;
}
static void recycle_step_allocation(StepAllocation &alloc)
@@ -5794,6 +5801,37 @@ static bool has_resident_state(void *block_tag)
return it != g_step_ctx.end() && it->second.state_ready;
}
static bool resident_key_usable(double **host_key)
{
if (!host_key) return false;
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
if (!host_key[i]) return false;
}
return true;
}
static bool resident_key_matches(const StepContext &ctx, double **host_key)
{
if (!ctx.state_ready || !ctx.resident_host_valid || !resident_key_usable(host_key))
return false;
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
if (ctx.resident_host[i] != host_key[i]) return false;
}
return true;
}
static void set_resident_key(StepContext &ctx, double **host_key)
{
if (!resident_key_usable(host_key)) {
ctx.resident_host.fill(nullptr);
ctx.resident_host_valid = false;
return;
}
for (int i = 0; i < BSSN_STATE_COUNT; ++i)
ctx.resident_host[i] = host_key[i];
ctx.resident_host_valid = true;
}
#define pow2(x) ((x) * (x))
@@ -7786,10 +7824,13 @@ extern "C" int z4c_cuda_rk4_substep(void *block_tag,
bind_state_input_slots(ctx.d_state_curr);
bind_state_output_slots(ctx.d_state_next);
}
double t0 = profile ? cuda_profile_now_ms() : 0.0;
if (!use_resident_state || !ctx.state_ready) {
upload_state_inputs(state_host_in, all);
if (use_resident_state) {
ctx.state_ready = true;
set_resident_key(ctx, state_host_in);
}
}
if (apply_enforce_ga) {
kern_enforce_ga_cuda<<<grid(all), BLK>>>(g_buf.slot[S_dxx], g_buf.slot[S_gxy], g_buf.slot[S_gxz],
@@ -7849,6 +7890,7 @@ extern "C" int z4c_cuda_rk4_substep(void *block_tag,
std::swap(ctx.d_state_curr_mem, ctx.d_state_next_mem);
ctx.d_state_curr.swap(ctx.d_state_next);
ctx.state_ready = true;
set_resident_key(ctx, state_host_out);
} else {
download_state_outputs(state_host_out, all);
}
@@ -8154,6 +8196,17 @@ extern "C" int z4c_cuda_has_resident_state(void *block_tag)
return has_resident_state(block_tag) ? 1 : 0;
}
extern "C" int z4c_cuda_resident_state_matches(void *block_tag,
double **state_host_key)
{
using namespace z4c_cuda;
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
auto it = g_step_ctx.find(block_tag);
if (it == g_step_ctx.end()) return 0;
return resident_key_matches(it->second, state_host_key) ? 1 : 0;
}
extern "C" void z4c_cuda_release_step_ctx(void *block_tag)
{
using namespace z4c_cuda;

View File

@@ -142,6 +142,8 @@ int z4c_cuda_download_constraint_outputs(int *ex,
double **constraint_host_out);
int z4c_cuda_has_resident_state(void *block_tag);
int z4c_cuda_resident_state_matches(void *block_tag,
double **state_host_key);
void z4c_cuda_release_step_ctx(void *block_tag);