Optimize BSSN CUDA resident AMR prolong path

This commit is contained in:
2026-04-30 10:58:15 +08:00
parent 1ee229a91f
commit 18e9c9cc50
3 changed files with 778 additions and 78 deletions

View File

@@ -102,6 +102,14 @@ struct RhsStageProfileStats {
double ms[RHS_STAGE_COUNT];
};
struct CudaAuxProfileStats {
long long prepare_calls;
long long writeback_calls;
double prepare_ms;
double writeback_ms;
double writeback_gb;
};
static CudaProfileStats &cuda_profile_stats() {
static CudaProfileStats stats = {
0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
@@ -110,6 +118,45 @@ static CudaProfileStats &cuda_profile_stats() {
return stats;
}
static CudaAuxProfileStats &cuda_aux_profile_stats() {
static CudaAuxProfileStats stats = {};
return stats;
}
static bool cuda_aux_profile_enabled() {
static int enabled = -1;
if (enabled < 0) {
const char *env = getenv("AMSS_PROFILE_CUDA_AUX");
enabled = (env && atoi(env) != 0) ? 1 : 0;
}
return enabled != 0;
}
static int cuda_aux_profile_every() {
static int every = -1;
if (every < 0) {
const char *env = getenv("AMSS_PROFILE_CUDA_AUX_EVERY");
every = (env && atoi(env) > 0) ? atoi(env) : 100;
}
return every;
}
static void cuda_aux_profile_maybe_log() {
if (!cuda_aux_profile_enabled()) return;
CudaAuxProfileStats &stats = cuda_aux_profile_stats();
const long long calls = stats.prepare_calls + stats.writeback_calls;
if (calls <= 0 || calls % cuda_aux_profile_every() != 0) return;
fprintf(stderr,
"[AMSS-CUDA-AUX][rank %d][dev %d] prepare=%lld avg_prepare=%.3f ms writebacks=%lld avg_writeback=%.3f ms writeback_GB=%.3f\n",
g_dispatch.my_rank, g_dispatch.my_device,
stats.prepare_calls,
stats.prepare_calls ? stats.prepare_ms / (double)stats.prepare_calls : 0.0,
stats.writeback_calls,
stats.writeback_calls ? stats.writeback_ms / (double)stats.writeback_calls : 0.0,
stats.writeback_gb);
fflush(stderr);
}
static RhsStageProfileStats &rhs_stage_profile_stats() {
static RhsStageProfileStats stats = {};
return stats;
@@ -480,6 +527,7 @@ static const int k_lk_rhs_slots[BSSN_LK_FIELD_COUNT] = {
};
__constant__ int d_subset_state_indices[BSSN_STATE_COUNT];
__constant__ double d_comm_state_soa[3 * BSSN_STATE_COUNT];
static const int k_lk_soa_signs[3 * BSSN_LK_FIELD_COUNT] = {
1, 1, 1,
@@ -523,6 +571,7 @@ struct StepContext {
std::array<double *, BSSN_STATE_COUNT> d_state_next;
std::array<std::array<double *, BSSN_STATE_COUNT>, BSSN_RESIDENT_BANK_COUNT> d_resident;
std::array<std::array<double *, BSSN_STATE_COUNT>, BSSN_RESIDENT_BANK_COUNT> resident_host;
std::array<std::array<unsigned char, BSSN_STATE_COUNT>, BSSN_RESIDENT_BANK_COUNT> resident_host_clean;
std::array<unsigned long long, BSSN_RESIDENT_BANK_COUNT> resident_age;
std::array<bool, BSSN_RESIDENT_BANK_COUNT> resident_valid;
std::array<double *, BSSN_MATTER_COUNT> d_matter;
@@ -552,6 +601,7 @@ struct StepContext {
for (int b = 0; b < BSSN_RESIDENT_BANK_COUNT; ++b) {
d_resident[b].fill(nullptr);
resident_host[b].fill(nullptr);
resident_host_clean[b].fill(0);
}
resident_age.fill(0);
resident_valid.fill(false);
@@ -634,6 +684,7 @@ static StepAllocation detach_step_allocation(StepContext &ctx)
for (int b = 0; b < BSSN_RESIDENT_BANK_COUNT; ++b) {
ctx.d_resident[b].fill(nullptr);
ctx.resident_host[b].fill(nullptr);
ctx.resident_host_clean[b].fill(0);
}
ctx.resident_age.fill(0);
ctx.resident_valid.fill(false);
@@ -661,6 +712,7 @@ static void attach_step_allocation(StepContext &ctx, const StepAllocation &alloc
ctx.resident_clock = 0;
for (int b = 0; b < BSSN_RESIDENT_BANK_COUNT; ++b) {
ctx.resident_host[b].fill(nullptr);
ctx.resident_host_clean[b].fill(0);
}
ctx.resident_age.fill(0);
ctx.resident_valid.fill(false);
@@ -843,6 +895,25 @@ static int *ensure_comm_segment_meta_buffer(size_t needed_ints)
return g_comm_segment_meta;
}
static void upload_comm_state_soa(const double *state_soa, int state_count)
{
double soa[3 * BSSN_STATE_COUNT];
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
soa[3 * i + 0] = 1.0;
soa[3 * i + 1] = 1.0;
soa[3 * i + 2] = 1.0;
}
if (state_soa) {
const int n = (state_count < BSSN_STATE_COUNT) ? state_count : BSSN_STATE_COUNT;
for (int i = 0; i < n; ++i) {
soa[3 * i + 0] = state_soa[3 * i + 0];
soa[3 * i + 1] = state_soa[3 * i + 1];
soa[3 * i + 2] = state_soa[3 * i + 2];
}
}
CUDA_CHECK(cudaMemcpyToSymbol(d_comm_state_soa, soa, sizeof(soa)));
}
static void upload_grid_params_if_needed(const GridParams &gp)
{
if (!g_gp_host_cache_valid ||
@@ -4906,6 +4977,42 @@ static bool resident_key_usable(double **host_key)
return true;
}
static void set_resident_host_clean(StepContext &ctx, int bank, bool clean)
{
if (bank < 0 || bank >= BSSN_RESIDENT_BANK_COUNT) return;
ctx.resident_host_clean[bank].fill(clean ? 1 : 0);
}
static bool resident_host_subset_clean(const StepContext &ctx,
int bank,
int subset_count,
const int *state_indices)
{
if (bank < 0 || bank >= BSSN_RESIDENT_BANK_COUNT) return false;
for (int i = 0; i < subset_count; ++i) {
const int state_index = state_indices ? state_indices[i] : i;
if (state_index < 0 || state_index >= BSSN_STATE_COUNT)
return false;
if (!ctx.resident_host_clean[bank][state_index])
return false;
}
return true;
}
static void mark_resident_host_subset_clean(StepContext &ctx,
int bank,
int subset_count,
const int *state_indices,
bool clean)
{
if (bank < 0 || bank >= BSSN_RESIDENT_BANK_COUNT) return;
for (int i = 0; i < subset_count; ++i) {
const int state_index = state_indices ? state_indices[i] : i;
if (state_index >= 0 && state_index < BSSN_STATE_COUNT)
ctx.resident_host_clean[bank][state_index] = clean ? 1 : 0;
}
}
static void mark_resident_current_bank(StepContext &ctx, int bank)
{
if (bank < 0 || bank >= BSSN_RESIDENT_BANK_COUNT) return;
@@ -4942,6 +5049,8 @@ static void writeback_resident_bank(StepContext &ctx, int bank, size_t all)
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
if (!ctx.resident_host[bank][i]) return;
}
const bool profile = cuda_aux_profile_enabled();
const double t0 = profile ? cuda_profile_now_ms() : 0.0;
const size_t bytes = all * sizeof(double);
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
CUDA_CHECK(cudaMemcpyAsync(ctx.resident_host[bank][i],
@@ -4949,6 +5058,14 @@ static void writeback_resident_bank(StepContext &ctx, int bank, size_t all)
bytes, cudaMemcpyDeviceToHost));
}
CUDA_CHECK(cudaDeviceSynchronize());
set_resident_host_clean(ctx, bank, true);
if (profile) {
CudaAuxProfileStats &stats = cuda_aux_profile_stats();
stats.writeback_calls++;
stats.writeback_ms += cuda_profile_now_ms() - t0;
stats.writeback_gb += (double)((size_t)BSSN_STATE_COUNT * bytes) / 1.0e9;
cuda_aux_profile_maybe_log();
}
}
static int choose_resident_bank_for_reuse(StepContext &ctx, int avoid_bank, size_t all)
@@ -4971,6 +5088,7 @@ static int choose_resident_bank_for_reuse(StepContext &ctx, int avoid_bank, size
writeback_resident_bank(ctx, best, all);
ctx.resident_valid[best] = false;
ctx.resident_host[best].fill(nullptr);
ctx.resident_host_clean[best].fill(0);
ctx.resident_age[best] = 0;
if (ctx.current_bank == best) {
ctx.current_bank = -1;
@@ -4986,6 +5104,7 @@ static void assign_resident_key(StepContext &ctx, int bank, double **host_key)
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
ctx.resident_host[bank][i] = host_key[i];
}
set_resident_host_clean(ctx, bank, false);
ctx.resident_age[bank] = ++ctx.resident_clock;
}
@@ -5008,6 +5127,7 @@ static int ensure_resident_bank(StepContext &ctx,
bind_state_input_slots(ctx.d_resident[bank]);
upload_state_inputs(host_key, all);
ctx.resident_valid[bank] = true;
set_resident_host_clean(ctx, bank, true);
}
return bank;
}
@@ -5018,8 +5138,10 @@ static int ensure_resident_bank(StepContext &ctx,
bind_state_input_slots(ctx.d_resident[bank]);
upload_state_inputs(host_key, all);
ctx.resident_valid[bank] = true;
set_resident_host_clean(ctx, bank, true);
} else {
ctx.resident_valid[bank] = false;
set_resident_host_clean(ctx, bank, false);
}
update_state_ready(ctx);
return bank;
@@ -5076,6 +5198,7 @@ static int choose_resident_bank_for_reuse_avoiding(StepContext &ctx,
writeback_resident_bank(ctx, best, all);
ctx.resident_valid[best] = false;
ctx.resident_host[best].fill(nullptr);
ctx.resident_host_clean[best].fill(0);
ctx.resident_age[best] = 0;
if (ctx.current_bank == best) {
ctx.current_bank = -1;
@@ -5599,6 +5722,12 @@ __global__ void kern_unpack_state_segments_batch(double * __restrict__ dst_mem,
}
}
__device__ __forceinline__ double load_comm_state_cell_sym(const double * __restrict__ src_mem,
int state_index,
int x, int y, int z,
int nx, int ny,
int all);
__global__ void kern_restrict_state_region_batch(const double * __restrict__ src_mem,
double * __restrict__ dst,
int nx, int ny,
@@ -5635,18 +5764,93 @@ __global__ void kern_restrict_state_region_batch(const double * __restrict__ src
{
const int y = fc_j + offs[oy];
const double wyz = wz * w[oy];
for (int ox = 0; ox < 6; ++ox)
{
const int x = fc_i + offs[ox];
const int src = x + y * nx + z * nx * ny;
sum += wyz * w[ox] * src_mem[(size_t)state_index * all + src];
}
for (int ox = 0; ox < 6; ++ox)
{
const int x = fc_i + offs[ox];
sum += wyz * w[ox] *
load_comm_state_cell_sym(src_mem, state_index, x, y, z, nx, ny, all);
}
}
}
dst[(size_t)state_index * region_all + local] = sum;
}
}
__device__ __forceinline__ double load_comm_state_cell_sym(const double * __restrict__ src_mem,
int state_index,
int x, int y, int z,
int nx, int ny,
int all)
{
double s = 1.0;
if (x < 0) {
x = -x - 1;
s *= d_comm_state_soa[3 * state_index + 0];
}
if (y < 0) {
y = -y - 1;
s *= d_comm_state_soa[3 * state_index + 1];
}
if (z < 0) {
z = -z - 1;
s *= d_comm_state_soa[3 * state_index + 2];
}
const int src = x + y * nx + z * nx * ny;
return s * src_mem[(size_t)state_index * all + src];
}
__global__ void kern_restrict_state_segments_batch(const double * __restrict__ src_mem,
double * __restrict__ dst,
int nx, int ny,
const int * __restrict__ meta,
int state_count,
int all)
{
const int segment = blockIdx.z;
const int state_index = blockIdx.y;
const int *m = meta + segment * 8;
const int sx = m[0], sy = m[1];
const int region_all = m[3];
const int offset = m[4];
const int fi0 = m[5], fj0 = m[6], fk0 = m[7];
if (state_index >= state_count) return;
const double c1 = 3.0 / 256.0;
const double c2 = -25.0 / 256.0;
const double c3 = 75.0 / 128.0;
const int offs[6] = {-2, -1, 0, 1, 2, 3};
const double w[6] = {c1, c2, c3, c3, c2, c1};
for (int local = blockIdx.x * blockDim.x + threadIdx.x;
local < region_all;
local += blockDim.x * gridDim.x)
{
const int ii = local % sx;
const int jj = (local / sx) % sy;
const int kk = local / (sx * sy);
const int fc_i = fi0 + 2 * ii;
const int fc_j = fj0 + 2 * jj;
const int fc_k = fk0 + 2 * kk;
double sum = 0.0;
for (int oz = 0; oz < 6; ++oz)
{
const int z = fc_k + offs[oz];
const double wz = w[oz];
for (int oy = 0; oy < 6; ++oy)
{
const int y = fc_j + offs[oy];
const double wyz = wz * w[oy];
for (int ox = 0; ox < 6; ++ox)
{
const int x = fc_i + offs[ox];
sum += wyz * w[ox] *
load_comm_state_cell_sym(src_mem, state_index, x, y, z, nx, ny, all);
}
}
}
dst[(size_t)offset + (size_t)state_index * region_all + local] = sum;
}
}
__global__ void kern_prolong_state_region_batch(const double * __restrict__ src_mem,
double * __restrict__ dst,
int nx, int ny,
@@ -5697,8 +5901,8 @@ __global__ void kern_prolong_state_region_batch(const double * __restrict__ src_
for (int ox = 0; ox < 6; ++ox)
{
const int x = ci + offs[ox];
const int src = x + y * nx + z * nx * ny;
sum += wyz * wx[ox] * src_mem[(size_t)state_index * all + src];
sum += wyz * wx[ox] *
load_comm_state_cell_sym(src_mem, state_index, x, y, z, nx, ny, all);
}
}
}
@@ -5706,6 +5910,69 @@ __global__ void kern_prolong_state_region_batch(const double * __restrict__ src_
}
}
__global__ void kern_prolong_state_segments_batch(const double * __restrict__ src_mem,
double * __restrict__ dst,
int nx, int ny,
const int * __restrict__ meta,
int state_count,
int all)
{
const int segment = blockIdx.z;
const int state_index = blockIdx.y;
const int *m = meta + segment * 11;
const int sx = m[0], sy = m[1];
const int region_all = m[3];
const int offset = m[4];
const int ii0 = m[5], jj0 = m[6], kk0 = m[7];
const int lbc_i = m[8], lbc_j = m[9], lbc_k = m[10];
if (state_index >= state_count) return;
const double c1 = 77.0 / 8192.0;
const double c2 = -693.0 / 8192.0;
const double c3 = 3465.0 / 4096.0;
const double c4 = 1155.0 / 4096.0;
const double c5 = -495.0 / 8192.0;
const double c6 = 63.0 / 8192.0;
const int offs[6] = {-2, -1, 0, 1, 2, 3};
const double wl[6] = {c1, c2, c3, c4, c5, c6};
const double wr[6] = {c6, c5, c4, c3, c2, c1};
for (int local = blockIdx.x * blockDim.x + threadIdx.x;
local < region_all;
local += blockDim.x * gridDim.x)
{
const int ii = local % sx;
const int jj = (local / sx) % sy;
const int kk = local / (sx * sy);
const int fine_i = ii0 + ii;
const int fine_j = jj0 + jj;
const int fine_k = kk0 + kk;
const int ci = fine_i / 2 - lbc_i;
const int cj = fine_j / 2 - lbc_j;
const int ck = fine_k / 2 - lbc_k;
const double *wx = ((fine_i / 2) * 2 == fine_i) ? wl : wr;
const double *wy = ((fine_j / 2) * 2 == fine_j) ? wl : wr;
const double *wz = ((fine_k / 2) * 2 == fine_k) ? wl : wr;
double sum = 0.0;
for (int oz = 0; oz < 6; ++oz)
{
const int z = ck + offs[oz];
const double wzv = wz[oz];
for (int oy = 0; oy < 6; ++oy)
{
const int y = cj + offs[oy];
const double wyz = wzv * wy[oy];
for (int ox = 0; ox < 6; ++ox)
{
const int x = ci + offs[ox];
sum += wyz * wx[ox] *
load_comm_state_cell_sym(src_mem, state_index, x, y, z, nx, ny, all);
}
}
}
dst[(size_t)offset + (size_t)state_index * region_all + local] = sum;
}
}
__global__ void kern_pack_state_subset(const double * __restrict__ src_mem,
double * __restrict__ dst,
int subset_count,
@@ -5777,6 +6044,8 @@ static void copy_state_region_cuda(void *block_tag,
ctx.resident_age[bank] = ++ctx.resident_clock;
mark_resident_current_bank(ctx, bank);
update_state_ready(ctx);
} else {
ctx.resident_host_clean[bank][state_index] = 1;
}
}
@@ -5821,6 +6090,8 @@ static void copy_state_region_packed_cuda(void *block_tag,
ctx.resident_age[bank] = ++ctx.resident_clock;
mark_resident_current_bank(ctx, bank);
update_state_ready(ctx);
} else {
ctx.resident_host_clean[bank][state_index] = 1;
}
}
@@ -5855,6 +6126,10 @@ static void copy_state_region_packed_batch_cuda(void *block_tag,
CUDA_CHECK(cudaMemcpy(host_buffer, d_comm,
total_doubles * sizeof(double),
cudaMemcpyDeviceToHost));
if (sx == ex[0] && sy == ex[1] && sz == ex[2] &&
i0 == 0 && j0 == 0 && k0 == 0) {
mark_resident_host_subset_clean(ctx, bank, state_count, nullptr, true);
}
} else {
CUDA_CHECK(cudaMemcpy(d_comm, host_buffer,
total_doubles * sizeof(double),
@@ -5890,11 +6165,14 @@ static void download_resident_state(void *block_tag, int *ex, double **state_hos
direct_download = env ? ((atoi(env) != 0) ? 1 : 0) : 1;
}
if (direct_download) {
if (resident_host_subset_clean(ctx, bank, BSSN_STATE_COUNT, nullptr))
return;
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
CUDA_CHECK(cudaMemcpyAsync(state_host_out[i], ctx.d_resident[bank][i],
bytes, cudaMemcpyDeviceToHost));
}
CUDA_CHECK(cudaDeviceSynchronize());
set_resident_host_clean(ctx, bank, true);
if (profile) {
CudaProfileStats &stats = cuda_profile_stats();
stats.resident_download_calls++;
@@ -5903,12 +6181,15 @@ static void download_resident_state(void *block_tag, int *ex, double **state_hos
}
return;
}
if (resident_host_subset_clean(ctx, bank, BSSN_STATE_COUNT, nullptr))
return;
CUDA_CHECK(cudaMemcpy(g_buf.h_stage, ctx.d_resident_mem[bank],
(size_t)BSSN_STATE_COUNT * bytes,
cudaMemcpyDeviceToHost));
for (int i = 0; i < BSSN_STATE_COUNT; ++i) {
std::memcpy(state_host_out[i], g_buf.h_stage + (size_t)i * all, bytes);
}
set_resident_host_clean(ctx, bank, true);
if (profile) {
CudaProfileStats &stats = cuda_profile_stats();
stats.resident_download_calls++;
@@ -5939,6 +6220,9 @@ static void copy_state_subset(void *block_tag,
for (int i = 0; i < subset_count; ++i) {
const int state_index = state_indices[i];
if (state_index < 0 || state_index >= BSSN_STATE_COUNT) continue;
if (kind == cudaMemcpyDeviceToHost &&
ctx.resident_host_clean[bank][state_index])
continue;
if (!state_host[i]) continue;
active_state_indices[active_count] = state_index;
active_state_host[active_count] = state_host[i];
@@ -5965,6 +6249,8 @@ static void copy_state_subset(void *block_tag,
h_comm + (size_t)i * all,
bytes);
}
mark_resident_host_subset_clean(ctx, bank, active_count,
active_state_indices, true);
} else {
for (int i = 0; i < active_count; ++i) {
std::memcpy(h_comm + (size_t)i * all,
@@ -5980,6 +6266,8 @@ static void copy_state_subset(void *block_tag,
ctx.resident_valid[bank] = true;
ctx.resident_age[bank] = ++ctx.resident_clock;
mark_resident_current_bank(ctx, bank);
mark_resident_host_subset_clean(ctx, bank, active_count,
active_state_indices, true);
update_state_ready(ctx);
}
}
@@ -6423,6 +6711,8 @@ int bssn_cuda_rk4_substep(void *block_tag,
g_buf.slot[S_dyy], g_buf.slot[S_gyz], g_buf.slot[S_dzz],
g_buf.slot[S_Axx], g_buf.slot[S_Axy], g_buf.slot[S_Axz],
g_buf.slot[S_Ayy], g_buf.slot[S_Ayz], g_buf.slot[S_Azz]);
if (use_resident_state && input_bank >= 0)
set_resident_host_clean(ctx, input_bank, false);
}
if (profile) {
cuda_profile_sync();
@@ -6493,6 +6783,7 @@ int bssn_cuda_rk4_substep(void *block_tag,
if (use_resident_state) {
ctx.resident_valid[output_bank] = true;
ctx.resident_age[output_bank] = ++ctx.resident_clock;
set_resident_host_clean(ctx, output_bank, false);
mark_resident_current_bank(ctx, output_bank);
update_state_ready(ctx);
} else {
@@ -6683,7 +6974,7 @@ static void copy_state_device_batch(void *block_tag,
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const int bank = active_or_keyed_bank(ctx, state_host_key, all,
pack_not_unpack == 0);
pack_not_unpack == 0 || state_host_key != nullptr);
double *base_mem = ctx.d_resident_mem[bank];
const int region_all = sx * sy * sz;
dim3 launch_grid((unsigned int)grid((size_t)region_all),
@@ -6704,6 +6995,7 @@ static void copy_state_device_batch(void *block_tag,
ctx.resident_valid[bank] = true;
ctx.resident_age[bank] = ++ctx.resident_clock;
mark_resident_current_bank(ctx, bank);
set_resident_host_clean(ctx, bank, false);
update_state_ready(ctx);
}
}
@@ -6731,7 +7023,7 @@ static void copy_state_device_segments(void *block_tag,
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const int bank = active_or_keyed_bank(ctx, state_host_key, all,
pack_not_unpack == 0);
pack_not_unpack == 0 || state_host_key != nullptr);
double *base_mem = ctx.d_resident_mem[bank];
int *d_meta = ensure_comm_segment_meta_buffer((size_t)segment_count * 8);
CUDA_CHECK(cudaMemcpy(d_meta, segment_meta,
@@ -6758,6 +7050,84 @@ static void copy_state_device_segments(void *block_tag,
}
}
static void restrict_state_device_segments(void *block_tag,
int state_count,
double *device_buffer,
const int *ex,
int segment_count,
const int *segment_meta,
double **state_host_key = nullptr,
const double *state_soa = nullptr)
{
if (state_count <= 0 || state_count > BSSN_STATE_COUNT) return;
if (segment_count <= 0 || !segment_meta || !device_buffer) return;
int max_region_all = 0;
for (int s = 0; s < segment_count; ++s) {
const int *m = segment_meta + s * 8;
if (m[0] <= 0 || m[1] <= 0 || m[2] <= 0 || m[3] <= 0) return;
if (m[3] > max_region_all) max_region_all = m[3];
}
if (max_region_all <= 0) return;
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const int bank = active_or_keyed_bank(ctx, state_host_key, all,
state_host_key != nullptr);
int *d_meta = ensure_comm_segment_meta_buffer((size_t)segment_count * 8);
CUDA_CHECK(cudaMemcpy(d_meta, segment_meta,
(size_t)segment_count * 8 * sizeof(int),
cudaMemcpyHostToDevice));
upload_comm_state_soa(state_soa, state_count);
dim3 launch_grid((unsigned int)grid((size_t)max_region_all),
(unsigned int)state_count,
(unsigned int)segment_count);
kern_restrict_state_segments_batch<<<launch_grid, BLK>>>(
ctx.d_resident_mem[bank], device_buffer,
ex[0], ex[1], d_meta, state_count,
ex[0] * ex[1] * ex[2]);
}
static void prolong_state_device_segments(void *block_tag,
int state_count,
double *device_buffer,
const int *ex,
int segment_count,
const int *segment_meta,
double **state_host_key = nullptr,
const double *state_soa = nullptr)
{
if (state_count <= 0 || state_count > BSSN_STATE_COUNT) return;
if (segment_count <= 0 || !segment_meta || !device_buffer) return;
int max_region_all = 0;
for (int s = 0; s < segment_count; ++s) {
const int *m = segment_meta + s * 11;
if (m[0] <= 0 || m[1] <= 0 || m[2] <= 0 || m[3] <= 0) return;
if (m[3] > max_region_all) max_region_all = m[3];
}
if (max_region_all <= 0) return;
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const int bank = active_or_keyed_bank(ctx, state_host_key, all,
state_host_key != nullptr);
int *d_meta = ensure_comm_segment_meta_buffer((size_t)segment_count * 11);
CUDA_CHECK(cudaMemcpy(d_meta, segment_meta,
(size_t)segment_count * 11 * sizeof(int),
cudaMemcpyHostToDevice));
upload_comm_state_soa(state_soa, state_count);
dim3 launch_grid((unsigned int)grid((size_t)max_region_all),
(unsigned int)state_count,
(unsigned int)segment_count);
kern_prolong_state_segments_batch<<<launch_grid, BLK>>>(
ctx.d_resident_mem[bank], device_buffer,
ex[0], ex[1], d_meta, 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,
@@ -6882,6 +7252,70 @@ int bssn_cuda_unpack_state_segments_from_device_buffer_for_host_views(void *bloc
return 0;
}
extern "C"
int bssn_cuda_restrict_state_segments_to_device_buffer(void *block_tag,
int state_count,
double *device_buffer,
int *ex,
int segment_count,
const int *segment_meta)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
restrict_state_device_segments(block_tag, state_count, device_buffer, ex,
segment_count, segment_meta);
return 0;
}
extern "C"
int bssn_cuda_restrict_state_segments_to_device_buffer_for_host_views(void *block_tag,
double **state_host_key,
int state_count,
double *device_buffer,
int *ex,
int segment_count,
const int *segment_meta,
const double *state_soa)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
restrict_state_device_segments(block_tag, state_count, device_buffer, ex,
segment_count, segment_meta, state_host_key, state_soa);
return 0;
}
extern "C"
int bssn_cuda_prolong_state_segments_to_device_buffer(void *block_tag,
int state_count,
double *device_buffer,
int *ex,
int segment_count,
const int *segment_meta)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
prolong_state_device_segments(block_tag, state_count, device_buffer, ex,
segment_count, segment_meta);
return 0;
}
extern "C"
int bssn_cuda_prolong_state_segments_to_device_buffer_for_host_views(void *block_tag,
double **state_host_key,
int state_count,
double *device_buffer,
int *ex,
int segment_count,
const int *segment_meta,
const double *state_soa)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
prolong_state_device_segments(block_tag, state_count, device_buffer, ex,
segment_count, segment_meta, state_host_key, state_soa);
return 0;
}
extern "C"
int bssn_cuda_restrict_state_batch_to_device_buffer(void *block_tag,
int state_count,
@@ -6896,6 +7330,7 @@ int bssn_cuda_restrict_state_batch_to_device_buffer(void *block_tag,
if (!device_buffer || sx <= 0 || sy <= 0 || sz <= 0) return 1;
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const int region_all = sx * sy * sz;
upload_comm_state_soa(nullptr, state_count);
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
kern_restrict_state_region_batch<<<launch_grid, BLK>>>(
@@ -6913,7 +7348,8 @@ int bssn_cuda_restrict_state_batch_to_device_buffer_for_host_views(void *block_t
double *device_buffer,
int *ex,
int sx, int sy, int sz,
int fi0, int fj0, int fk0)
int fi0, int fj0, int fk0,
const double *state_soa)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
@@ -6923,6 +7359,7 @@ int bssn_cuda_restrict_state_batch_to_device_buffer_for_host_views(void *block_t
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const int bank = active_or_keyed_bank(ctx, state_host_key, all, true);
const int region_all = sx * sy * sz;
upload_comm_state_soa(state_soa, state_count);
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
kern_restrict_state_region_batch<<<launch_grid, BLK>>>(
@@ -6948,6 +7385,7 @@ int bssn_cuda_prolong_state_batch_to_device_buffer(void *block_tag,
if (!device_buffer || sx <= 0 || sy <= 0 || sz <= 0) return 1;
StepContext &ctx = ensure_step_ctx(block_tag, (size_t)ex[0] * ex[1] * ex[2]);
const int region_all = sx * sy * sz;
upload_comm_state_soa(nullptr, state_count);
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
kern_prolong_state_region_batch<<<launch_grid, BLK>>>(
@@ -6967,7 +7405,8 @@ int bssn_cuda_prolong_state_batch_to_device_buffer_for_host_views(void *block_ta
int *ex,
int sx, int sy, int sz,
int ii0, int jj0, int kk0,
int lbc_i, int lbc_j, int lbc_k)
int lbc_i, int lbc_j, int lbc_k,
const double *state_soa)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
@@ -6977,6 +7416,7 @@ int bssn_cuda_prolong_state_batch_to_device_buffer_for_host_views(void *block_ta
const size_t all = (size_t)ex[0] * ex[1] * ex[2];
const int bank = active_or_keyed_bank(ctx, state_host_key, all, true);
const int region_all = sx * sy * sz;
upload_comm_state_soa(state_soa, state_count);
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
kern_prolong_state_region_batch<<<launch_grid, BLK>>>(
@@ -7028,6 +7468,8 @@ int bssn_cuda_prepare_inter_time_level(void *block_tag,
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
const bool profile = cuda_aux_profile_enabled();
const double t0 = profile ? cuda_profile_now_ms() : 0.0;
if (source_count != 2 && source_count != 3) return 1;
if (!resident_key_usable(src1_host_key) ||
!resident_key_usable(src2_host_key) ||
@@ -7074,10 +7516,19 @@ int bssn_cuda_prepare_inter_time_level(void *block_tag,
(source_count == 3) ? ctx.d_resident_mem[src3_bank] : nullptr,
ctx.d_resident_mem[dst_bank],
c1, c2, c3, BSSN_STATE_COUNT, (int)all);
if (profile)
cuda_profile_sync();
ctx.resident_valid[dst_bank] = true;
ctx.resident_age[dst_bank] = ++ctx.resident_clock;
set_resident_host_clean(ctx, dst_bank, false);
mark_resident_current_bank(ctx, dst_bank);
update_state_ready(ctx);
if (profile) {
CudaAuxProfileStats &stats = cuda_aux_profile_stats();
stats.prepare_calls++;
stats.prepare_ms += cuda_profile_now_ms() - t0;
cuda_aux_profile_maybe_log();
}
return 0;
}