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

Add opt-in BSSN CUDA resident AMR path

Browse Source
This commit is contained in:
CGH0S7
2026-04-29 19:15:37 +08:00
parent 090d8657ae
commit 68eab03bac
4 changed files with 449 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

159
AMSS_NCKU_source/bssn_rhs_cuda.cu
View File

@@ -5269,6 +5269,113 @@ __global__ void kern_unpack_state_segments_batch(double * __restrict__ dst_mem,
}
}
__global__ void kern_restrict_state_region_batch(const double * __restrict__ src_mem,
double * __restrict__ dst,
int nx, int ny,
int sx, int sy, int sz,
int fi0, int fj0, int fk0,
int region_all,
int state_count,
int all)
{
const int state_index = blockIdx.y;
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];
const int src = x + y * nx + z * nx * ny;
sum += wyz * w[ox] * src_mem[(size_t)state_index * all + src];
}
}
}
dst[(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,
int sx, int sy, int sz,
int ii0, int jj0, int kk0,
int lbc_i, int lbc_j, int lbc_k,
int region_all,
int state_count,
int all)
{
const int state_index = blockIdx.y;
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];
const int src = x + y * nx + z * nx * ny;
sum += wyz * wx[ox] * src_mem[(size_t)state_index * all + src];
}
}
}
dst[(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,
@@ -5929,7 +6036,7 @@ int bssn_cuda_rk4_substep(void *block_tag,
bind_state_output_slots(ctx.d_state_next);
}
double t0 = profile ? cuda_profile_now_ms() : 0.0;
if (!use_resident_state || RK4 == 0 || !ctx.state_ready) {
if (!use_resident_state || !ctx.state_ready) {
upload_state_inputs(state_host_in, all);
}
if (apply_enforce_ga) {
@@ -6274,6 +6381,56 @@ int bssn_cuda_unpack_state_segments_from_device_buffer(void *block_tag,
return 0;
}
extern "C"
int bssn_cuda_restrict_state_batch_to_device_buffer(void *block_tag,
int state_count,
double *device_buffer,
int *ex,
int sx, int sy, int sz,
int fi0, int fj0, int fk0)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
if (state_count <= 0 || state_count > BSSN_STATE_COUNT) return 1;
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;
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
kern_restrict_state_region_batch<<<launch_grid, BLK>>>(
ctx.d_state_curr_mem, device_buffer,
ex[0], ex[1], sx, sy, sz,
fi0, fj0, fk0, region_all, state_count,
ex[0] * ex[1] * ex[2]);
return 0;
}
extern "C"
int bssn_cuda_prolong_state_batch_to_device_buffer(void *block_tag,
int state_count,
double *device_buffer,
int *ex,
int sx, int sy, int sz,
int ii0, int jj0, int kk0,
int lbc_i, int lbc_j, int lbc_k)
{
init_gpu_dispatch();
CUDA_CHECK(cudaSetDevice(g_dispatch.my_device));
if (state_count <= 0 || state_count > BSSN_STATE_COUNT) return 1;
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;
dim3 launch_grid((unsigned int)grid((size_t)region_all),
(unsigned int)state_count);
kern_prolong_state_region_batch<<<launch_grid, BLK>>>(
ctx.d_state_curr_mem, device_buffer,
ex[0], ex[1], sx, sy, sz,
ii0, jj0, kk0, lbc_i, lbc_j, lbc_k,
region_all, state_count,
ex[0] * ex[1] * ex[2]);
return 0;
}
extern "C"
int bssn_cuda_download_state_subset(void *block_tag,
int *ex,