Add opt-in BSSN CUDA resident AMR path

2026-04-29 19:15:37 +08:00
parent 090d8657ae
commit 68eab03bac
4 changed files with 449 additions and 42 deletions
--- a/AMSS_NCKU_source/Parallel.C
+++ b/AMSS_NCKU_source/Parallel.C
@@ -173,6 +173,96 @@ int cuda_state_var_count(MyList<var> *src_vars, MyList<var> *dst_vars)
 }

 #if USE_CUDA_BSSN || USE_CUDA_Z4C
+int fortran_idint(double x)
+{
+  return (int)x;
+}
+
+bool cuda_cell_gw3_restrict_params(const Parallel::gridseg *src,
+                                   const Parallel::gridseg *dst,
+                                   int first_fine[3])
+{
+#if USE_CUDA_BSSN && defined(Cell) && (ghost_width == 3)
+  if (!src || !dst || !src->Bg || !dst->Bg)
+    return false;
+  for (int d = 0; d < dim; ++d)
+  {
+    const double llbc = dst->llb[d];
+    const double uubc = dst->uub[d];
+    const double llbf = src->Bg->bbox[d];
+    const double uubf = src->Bg->bbox[d + dim];
+    const double CD = (uubc - llbc) / (double)dst->shape[d];
+    const double FD = (uubf - llbf) / (double)src->Bg->shape[d];
+    if (fabs(CD - 2.0 * FD) > 1.0e-10)
+      return false;
+    double base;
+    if (llbc <= llbf)
+      base = llbc;
+    else
+    {
+      const int j = fortran_idint((llbc - llbf) / FD + 0.4);
+      base = ((j / 2) * 2 == j) ? llbf : (llbf - CD / 2.0);
+    }
+    const int lbc = fortran_idint((llbc - base) / CD + 0.4) + 1;
+    const int lbf = fortran_idint((llbf - base) / FD + 0.4) + 1;
+    first_fine[d] = 2 * lbc - lbf - 1;
+    if (first_fine[d] - 2 < 0)
+      return false;
+    if (first_fine[d] + 2 * (dst->shape[d] - 1) + 3 >= src->Bg->shape[d])
+      return false;
+  }
+  return true;
+#else
+  (void)src; (void)dst; (void)first_fine;
+  return false;
+#endif
+}
+
+bool cuda_cell_gw3_prolong_params(const Parallel::gridseg *src,
+                                  const Parallel::gridseg *dst,
+                                  int first_fine_ii[3],
+                                  int coarse_lb[3])
+{
+#if USE_CUDA_BSSN && defined(Cell) && (ghost_width == 3)
+  if (!src || !dst || !src->Bg || !dst->Bg)
+    return false;
+  for (int d = 0; d < dim; ++d)
+  {
+    const double llbc = src->Bg->bbox[d];
+    const double uubc = src->Bg->bbox[d + dim];
+    const double llbf = dst->llb[d];
+    const double uubf = dst->uub[d];
+    const double CD = (uubc - llbc) / (double)src->Bg->shape[d];
+    const double FD = (uubf - llbf) / (double)dst->shape[d];
+    if (fabs(CD - 2.0 * FD) > 1.0e-10)
+      return false;
+    double base;
+    if (llbc <= llbf)
+      base = llbc;
+    else
+    {
+      const int j = fortran_idint((llbc - llbf) / FD + 0.4);
+      base = ((j / 2) * 2 == j) ? llbf : (llbf - CD / 2.0);
+    }
+    const int lbf = fortran_idint((llbf - base) / FD + 0.4) + 1;
+    const int lbc = fortran_idint((llbc - base) / CD + 0.4) + 1;
+    first_fine_ii[d] = lbf;
+    coarse_lb[d] = lbc;
+    const int first_coarse = first_fine_ii[d] / 2 - coarse_lb[d];
+    const int last_fine_ii = first_fine_ii[d] + dst->shape[d] - 1;
+    const int last_coarse = last_fine_ii / 2 - coarse_lb[d];
+    if (first_coarse - 2 < 0)
+      return false;
+    if (last_coarse + 3 >= src->Bg->shape[d])
+      return false;
+  }
+  return true;
+#else
+  (void)src; (void)dst; (void)first_fine_ii; (void)coarse_lb;
+  return false;
+#endif
+}
+
 bool cuda_state_count_direct_supported(int state_count)
 {
 #if USE_CUDA_Z4C && (ABEtype == 2)
@@ -187,22 +277,36 @@ bool cuda_state_count_direct_supported(int state_count)

 bool cuda_can_direct_pack(const Parallel::gridseg *src, const Parallel::gridseg *dst, int type)
 {
-  if (type != 1 || !src || !dst || !src->Bg)
+  if (!src || !dst || !src->Bg)
    return false;
 #if USE_CUDA_Z4C && (ABEtype == 2)
+  if (type != 1)
+    return false;
  return z4c_cuda_has_resident_state(src->Bg) != 0;
 #elif USE_CUDA_BSSN
-  return bssn_cuda_has_resident_state(src->Bg) != 0;
+  if (bssn_cuda_has_resident_state(src->Bg) == 0)
+    return false;
+  if (type == 1)
+    return true;
+  int a[3], b[3];
+  if (type == 2)
+    return cuda_cell_gw3_restrict_params(src, dst, a);
+  if (type == 3)
+    return cuda_cell_gw3_prolong_params(src, dst, a, b);
+  return false;
 #else
+  (void)type;
  return false;
 #endif
 }

 bool cuda_can_direct_unpack(const Parallel::gridseg *dst, int type)
 {
-  if (type != 1 || !dst || !dst->Bg)
+  if (type < 1 || type > 3 || !dst || !dst->Bg)
    return false;
 #if USE_CUDA_Z4C && (ABEtype == 2)
+  if (type != 1)
+    return false;
  return z4c_cuda_has_resident_state(dst->Bg) != 0;
 #elif USE_CUDA_BSSN
  return bssn_cuda_has_resident_state(dst->Bg) != 0;
@@ -336,24 +440,50 @@ void ensure_device_comm_buffer(double **buffers, int *caps, int idx, int length)
 bool cuda_direct_pack_segment_to_device(double *buffer,
                                        const Parallel::gridseg *src,
                                        const Parallel::gridseg *dst,
-                                        int state_count)
+                                        int state_count,
+                                        int type)
 {
 #if USE_CUDA_BSSN
  if (state_count <= 0 || state_count > BSSN_CUDA_STATE_COUNT)
    return false;
  const double t0 = sync_profile_enabled() ? MPI_Wtime() : 0.0;
-  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);
-  const bool ok = bssn_cuda_pack_state_batch_to_device_buffer(
-                      src->Bg, state_count, buffer, src->Bg->shape,
-                      i0, j0, k0,
-                      dst->shape[0], dst->shape[1], dst->shape[2]) == 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_device_buffer(
+             src->Bg, 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_device_buffer(
+             src->Bg, state_count, buffer, src->Bg->shape,
+             dst->shape[0], dst->shape[1], dst->shape[2],
+             first_fine[0], first_fine[1], first_fine[2]) == 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_device_buffer(
+             src->Bg, 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]) == 0;
+  }
  if (sync_profile_enabled())
    sync_profile_stats().direct_pack_sec += MPI_Wtime() - t0;
  return ok;
 #else
-  (void)buffer; (void)src; (void)dst; (void)state_count;
+  (void)buffer; (void)src; (void)dst; (void)state_count; (void)type;
  return false;
 #endif
 }
@@ -382,6 +512,56 @@ bool cuda_direct_unpack_segment_from_device(double *buffer,
 #endif
 }

+bool cuda_download_resident_subset_to_host(Block *block,
+                                           MyList<var> *vars,
+                                           int state_count)
+{
+#if USE_CUDA_BSSN
+  if (!block || state_count <= 0 || state_count > BSSN_CUDA_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];
+  MyList<var> *v = vars;
+  for (int i = 0; i < state_count; ++i)
+  {
+    if (!v)
+      return false;
+    indices[i] = i;
+    views[i] = block->fgfs[v->data->sgfn];
+    v = v->next;
+  }
+  return bssn_cuda_download_state_subset(block, block->shape, state_count, indices, views) == 0;
+#else
+  (void)block; (void)vars; (void)state_count;
+  return false;
+#endif
+}
+
+bool cuda_unpack_host_region_to_resident(Block *block,
+                                         int state_index,
+                                         double *buffer,
+                                         const Parallel::gridseg *dst)
+{
+#if USE_CUDA_BSSN
+  if (!block || !dst || state_index < 0 || state_index >= BSSN_CUDA_STATE_COUNT)
+    return false;
+  if (bssn_cuda_has_resident_state(block) == 0)
+    return true;
+  const int i0 = cuda_seg_begin(dst, block, 0);
+  const int j0 = cuda_seg_begin(dst, block, 1);
+  const int k0 = cuda_seg_begin(dst, block, 2);
+  return bssn_cuda_unpack_state_region_from_host_buffer(
+             block, state_index, buffer, block->shape,
+             i0, j0, k0,
+             dst->shape[0], dst->shape[1], dst->shape[2]) == 0;
+#else
+  (void)block; (void)state_index; (void)buffer; (void)dst;
+  return false;
+#endif
+}
+
 bool cuda_device_state_count_supported(int state_count)
 {
 #if USE_CUDA_BSSN
@@ -499,11 +679,12 @@ int cuda_data_packer_device_batched(double *data,
 }
 #endif

-bool cuda_segments_same_level(MyList<Parallel::gridseg> *src,
-                              MyList<Parallel::gridseg> *dst,
-                              int rank_in,
-                              int dir,
-                              int myrank)
+bool cuda_segments_device_eligible(MyList<Parallel::gridseg> *src,
+                                   MyList<Parallel::gridseg> *dst,
+                                   int rank_in,
+                                   int dir,
+                                   int myrank,
+                                   int state_count)
 {
  bool has_work = false;
  while (src && dst)
@@ -512,13 +693,30 @@ bool cuda_segments_same_level(MyList<Parallel::gridseg> *src,
        (dir == UNPACK && src->data->Bg->rank == rank_in && dst->data->Bg->rank == myrank))
    {
      has_work = true;
-      if (!src->data || !dst->data || !src->data->Bg || !dst->data->Bg ||
-          src->data->Bg->lev != dst->data->Bg->lev)
+      if (!src->data || !dst->data || !src->data->Bg || !dst->data->Bg)
        return false;
+      int type;
+      if (src->data->Bg->lev == dst->data->Bg->lev)
+        type = 1;
+      else if (src->data->Bg->lev > dst->data->Bg->lev)
+        type = 2;
+      else
+        type = 3;
+      if (dir == PACK)
+      {
+        if (!cuda_can_direct_pack(src->data, dst->data, type))
+          return false;
+      }
+      else
+      {
+        if (!cuda_can_direct_unpack(dst->data, type))
+          return false;
+      }
    }
    src = src->next;
    dst = dst->next;
  }
+  (void)state_count;
  return has_work;
 }

@@ -530,16 +728,8 @@ bool cuda_pack_to_device_eligible(MyList<Parallel::gridseg> *src,
 {
  if (!cuda_aware_mpi_enabled() || !cuda_device_state_count_supported(state_count))
    return false;
-  if (!cuda_segments_same_level(src, dst, rank_in, PACK, myrank))
+  if (!cuda_segments_device_eligible(src, dst, rank_in, PACK, myrank, state_count))
    return false;
-  while (src && dst)
-  {
-    if (dst->data->Bg->rank == rank_in && src->data->Bg->rank == myrank &&
-        !cuda_can_direct_pack(src->data, dst->data, 1))
-      return false;
-    src = src->next;
-    dst = dst->next;
-  }
  return true;
 }

@@ -551,16 +741,8 @@ bool cuda_recv_to_device_eligible(MyList<Parallel::gridseg> *src,
 {
  if (!cuda_aware_mpi_enabled() || !cuda_device_state_count_supported(state_count))
    return false;
-  if (!cuda_segments_same_level(src, dst, rank_in, UNPACK, myrank))
+  if (!cuda_segments_device_eligible(src, dst, rank_in, UNPACK, myrank, state_count))
    return false;
-  while (src && dst)
-  {
-    if (src->data->Bg->rank == rank_in && dst->data->Bg->rank == myrank &&
-        !cuda_can_direct_unpack(dst->data, 1))
-      return false;
-    src = src->next;
-    dst = dst->next;
-  }
  return true;
 }

@@ -4355,11 +4537,12 @@ 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 && cuda_state_count_direct_supported(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))
          {
            if (s_cuda_aware_pack_active) {
-              handled_by_cuda = cuda_direct_pack_segment_to_device(data + size_out, src->data, dst->data, state_count);
+              handled_by_cuda = cuda_direct_pack_segment_to_device(data + size_out, src->data, dst->data, state_count, type);
            } else {
              handled_by_cuda = cuda_direct_pack_segment(data + size_out, src->data, dst->data, state_count);
            }
@@ -4369,7 +4552,8 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
              MPI_Abort(MPI_COMM_WORLD, 1);
            }
          }
-          else if (dir == UNPACK && cuda_state_count_direct_supported(state_count) &&
+          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))
          {
            if (s_cuda_aware_pack_active) {
@@ -4393,6 +4577,17 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
              MPI_Abort(MPI_COMM_WORLD, 1);
            }
 #endif
+#endif
+#if USE_CUDA_BSSN
+          if (dir == PACK && state_idx == 0 && cuda_state_count_direct_supported(state_count) &&
+              src->data && src->data->Bg && bssn_cuda_has_resident_state(src->data->Bg))
+          {
+            if (!cuda_download_resident_subset_to_host(src->data->Bg, VarLists, state_count))
+            {
+              cout << "Parallel::data_packer: CUDA resident fallback download failed." << endl;
+              MPI_Abort(MPI_COMM_WORLD, 1);
+            }
+          }
 #endif
          if (dir == PACK)
            switch (type)
@@ -4414,9 +4609,22 @@ int Parallel::data_packer(double *data, MyList<Parallel::gridseg> *src, MyList<P
                         dst->data->llb, dst->data->uub, varls->data->SoA, Symmetry);
            }
          if (dir == UNPACK) // from target data to corresponding grid
+          {
            f_copy(DIM, dst->data->Bg->bbox, dst->data->Bg->bbox + dim, dst->data->Bg->shape, dst->data->Bg->fgfs[varld->data->sgfn],
                   dst->data->llb, dst->data->uub, dst->data->shape, data + size_out,
                   dst->data->llb, dst->data->uub);
+#if USE_CUDA_BSSN
+            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))
+              {
+                cout << "Parallel::data_packer: CUDA resident fallback upload failed." << endl;
+                MPI_Abort(MPI_COMM_WORLD, 1);
+              }
+            }
+#endif
+          }
 #if USE_CUDA_BSSN || USE_CUDA_Z4C
          }
          else