DeepGPR 0.0.16__tar.gz → 0.0.17__tar.gz

{deepgpr-0.0.16 → deepgpr-0.0.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DeepGPR
-Version: 0.0.16
+Version: 0.0.17
 Summary: PyTorch and CUDA for GPR FWI
 Author-email: Lei Liu <liulei990222@gmail.com>
 License-Expression: MIT

{deepgpr-0.0.16 → deepgpr-0.0.17}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "DeepGPR"
-version = "0.0.16"
+version = "0.0.17"
 authors = [
     { name = "Lei Liu", email = "liulei990222@gmail.com" }
 ]

{deepgpr-0.0.16 → deepgpr-0.0.17}/src/DeepGPR/lib/deepgpr.cu

@@ -175,6 +175,139 @@ __device__ __forceinline__ float load_wavefield_value_device(
     } \
 } while (0)
 
+static int clamp_cpml_index(int value, int extent)
+{
+    if (value < 0) return 0;
+    if (value > extent) return extent;
+    return value;
+}
+
+/*
+ * Partition the CPML shell into six disjoint boxes. The x boxes own edges and
+ * corners, the y boxes exclude the x boxes, and the z boxes exclude both.
+ * Return the number of compact boundary cells without passing a large layout
+ * object to CUDA kernels, which is unreliable with CUDA 11.8's old frontend.
+ */
+static long long cpml_cells_per_shot(
+    int NX, int NY, int NZ,
+    int pml0, int pml1, int pml2, int pml3, int pml4, int pml5,
+    int electric)
+{
+    int x_low_end = pml0 > 0
+        ? clamp_cpml_index(pml0 + (electric ? 1 : 0), NX) : 0;
+    int y_low_end = pml2 > 0
+        ? clamp_cpml_index(pml2 + (electric ? 1 : 0), NY) : 0;
+    int z_low_end = pml4 > 0
+        ? clamp_cpml_index(pml4 + (electric ? 1 : 0), NZ) : 0;
+    int x_high_begin = pml1 > 0
+        ? clamp_cpml_index(NX - 1 - pml1, NX) : NX;
+    int y_high_begin = pml3 > 0
+        ? clamp_cpml_index(NY - 1 - pml3, NY) : NY;
+    int z_high_begin = pml5 > 0
+        ? clamp_cpml_index(NZ - 1 - pml5, NZ) : NZ;
+
+    int x_mid_end = x_high_begin > x_low_end ? x_high_begin : x_low_end;
+    int y_mid_end = y_high_begin > y_low_end ? y_high_begin : y_low_end;
+    int z_mid_end = z_high_begin > z_low_end ? z_high_begin : z_low_end;
+    long long x_mid = x_mid_end - x_low_end;
+    long long y_mid = y_mid_end - y_low_end;
+
+    long long cells = (long long)x_low_end * NY * NZ;
+    if (pml1 > 0) cells += (long long)(NX - x_mid_end) * NY * NZ;
+    cells += x_mid * y_low_end * NZ;
+    if (pml3 > 0) cells += x_mid * (NY - y_mid_end) * NZ;
+    cells += x_mid * y_mid * z_low_end;
+    if (pml5 > 0) cells += x_mid * y_mid * (NZ - z_mid_end);
+    return cells;
+}
+
+template<int ELECTRIC>
+__device__ inline bool map_cpml_work(
+    long long compact_work, int step, long long cells_per_shot,
+    int NX, int NY, int NZ,
+    int pml0, int pml1, int pml2, int pml3, int pml4, int pml5,
+    int* shot, long long* i, long long* j, long long* k)
+{
+    if (cells_per_shot == 0 ||
+        compact_work >= (long long)step * cells_per_shot) {
+        return false;
+    }
+
+    int x_low_end = pml0 > 0 ? pml0 + ELECTRIC : 0;
+    int y_low_end = pml2 > 0 ? pml2 + ELECTRIC : 0;
+    int z_low_end = pml4 > 0 ? pml4 + ELECTRIC : 0;
+    if (x_low_end > NX) x_low_end = NX;
+    if (y_low_end > NY) y_low_end = NY;
+    if (z_low_end > NZ) z_low_end = NZ;
+
+    int x_high_begin = pml1 > 0 ? NX - 1 - pml1 : NX;
+    int y_high_begin = pml3 > 0 ? NY - 1 - pml3 : NY;
+    int z_high_begin = pml5 > 0 ? NZ - 1 - pml5 : NZ;
+    if (x_high_begin < 0) x_high_begin = 0;
+    if (y_high_begin < 0) y_high_begin = 0;
+    if (z_high_begin < 0) z_high_begin = 0;
+
+    int x_mid_end = x_high_begin > x_low_end ? x_high_begin : x_low_end;
+    int y_mid_end = y_high_begin > y_low_end ? y_high_begin : y_low_end;
+    int z_mid_end = z_high_begin > z_low_end ? z_high_begin : z_low_end;
+    int x_mid = x_mid_end - x_low_end;
+    int y_mid = y_mid_end - y_low_end;
+
+    long long local = compact_work % cells_per_shot;
+    long long region_size = (long long)x_low_end * NY * NZ;
+    int i0, j0, nj, k0, nk;
+    if (local < region_size) {
+        i0 = 0;
+        j0 = 0; nj = NY;
+        k0 = 0; nk = NZ;
+    } else {
+        local -= region_size;
+        region_size = pml1 > 0 ? (long long)(NX - x_mid_end) * NY * NZ : 0;
+        if (local < region_size) {
+            i0 = x_mid_end;
+            j0 = 0; nj = NY;
+            k0 = 0; nk = NZ;
+        } else {
+            local -= region_size;
+            region_size = (long long)x_mid * y_low_end * NZ;
+            if (local < region_size) {
+                i0 = x_low_end;
+                j0 = 0; nj = y_low_end;
+                k0 = 0; nk = NZ;
+            } else {
+                local -= region_size;
+                region_size = pml3 > 0 ? (long long)x_mid * (NY - y_mid_end) * NZ : 0;
+                if (local < region_size) {
+                    i0 = x_low_end;
+                    j0 = y_mid_end; nj = NY - y_mid_end;
+                    k0 = 0; nk = NZ;
+                } else {
+                    local -= region_size;
+                    region_size = (long long)x_mid * y_mid * z_low_end;
+                    if (local < region_size) {
+                        i0 = x_low_end;
+                        j0 = y_low_end; nj = y_mid;
+                        k0 = 0; nk = z_low_end;
+                    } else {
+                        local -= region_size;
+                        i0 = x_low_end;
+                        j0 = y_low_end; nj = y_mid;
+                        k0 = z_mid_end; nk = NZ - z_mid_end;
+                    }
+                }
+            }
+        }
+    }
+
+    *shot = (int)(compact_work / cells_per_shot);
+    long long jk = (long long)nj * nk;
+    *i = i0 + local / jk;
+    long long remainder = local % jk;
+    *j = j0 + remainder / nk;
+    *k = k0 + remainder % nk;
+    return true;
+}
+
 static int g_fdtd_order = 2;
 
 DEEPGPR_API int deepgpr_abi_version(void)
@@ -663,20 +796,19 @@ __global__ void cpml_e_gpu(
     float* __restrict__ y0EPhi1, float* __restrict__ y0EPhi2,
     float* __restrict__ ymEPhi1, float* __restrict__ ymEPhi2,
     float* __restrict__ z0EPhi1, float* __restrict__ z0EPhi2,
-    float* __restrict__ zmEPhi1, float* __restrict__ zmEPhi2)
+    float* __restrict__ zmEPhi1, float* __restrict__ zmEPhi2,
+    long long compact_cells_per_shot)
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ_FIELDS;
-    long long k = rem % NZ_FIELDS;
+    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
+    int s;
+    long long i, j, k;
+    if (!map_cpml_work<1>(compact_work, step, compact_cells_per_shot,
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, &s, &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ_FIELDS + k;
+    long long work = (long long)s * field_stride + idx;
 
     bool in_x0 = (pml0 > 0 && i > 0 && i <= pml0 && j < NY_FIELDS && k < NZ_FIELDS);
     bool in_xm = (pml1 > 0 && i >= NX_FIELDS - 1 - pml1 && i < NX_FIELDS - 1 && j < NY_FIELDS && k < NZ_FIELDS);
@@ -957,20 +1089,19 @@ __global__ void cpml_h_gpu(
     float* __restrict__ y0HPhi1, float* __restrict__ y0HPhi2,
     float* __restrict__ ymHPhi1, float* __restrict__ ymHPhi2,
     float* __restrict__ z0HPhi1, float* __restrict__ z0HPhi2,
-    float* __restrict__ zmHPhi1, float* __restrict__ zmHPhi2)
+    float* __restrict__ zmHPhi1, float* __restrict__ zmHPhi2,
+    long long compact_cells_per_shot)
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ_FIELDS;
-    long long k = rem % NZ_FIELDS;
+    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
+    int s;
+    long long i, j, k;
+    if (!map_cpml_work<0>(compact_work, step, compact_cells_per_shot,
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, &s, &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ_FIELDS + k;
+    long long work = (long long)s * field_stride + idx;
 
     bool in_x0 = (pml0 > 0 && i < pml0 && j < NY_FIELDS && k < NZ_FIELDS);
     bool in_xm = (pml1 > 0 && i >= NX_FIELDS - 1 - pml1 && i < NX_FIELDS - 1 && j < NY_FIELDS && k < NZ_FIELDS);
@@ -1111,18 +1242,19 @@ __global__ void adjoint_cpml_e_gpu(
     const float* z0R, const float* zmR, const float* update,
     float* x0P1, float* x0P2, float* xmP1, float* xmP2,
     float* y0P1, float* y0P2, float* ymP1, float* ymP2,
-    float* z0P1, float* z0P2, float* zmP1, float* zmP2)
+    float* z0P1, float* z0P2, float* zmP1, float* zmP2,
+    long long compact_cells_per_shot)
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ;
-    long long k = rem % NZ;
+    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
+    int s;
+    long long i, j, k;
+    if (!map_cpml_work<1>(compact_work, step, compact_cells_per_shot,
+        NX, NY, NZ, pml0, pml1, pml2, pml3, pml4, pml5,
+        &s, &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ + k;
+    long long work = (long long)s * field_stride + idx;
     float upd = update[idx];
 
 #define APPLY_E_PML_GPU(R, P, p, q, stride, coord, n, spacing, field, source, sign) \
@@ -1211,18 +1343,19 @@ __global__ void adjoint_cpml_h_gpu(
     const float* z0R, const float* zmR, const float* update,
     float* x0P1, float* x0P2, float* xmP1, float* xmP2,
     float* y0P1, float* y0P2, float* ymP1, float* ymP2,
-    float* z0P1, float* z0P2, float* zmP1, float* zmP2)
+    float* z0P1, float* z0P2, float* zmP1, float* zmP2,
+    long long compact_cells_per_shot)
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ;
-    long long k = rem % NZ;
+    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
+    int s;
+    long long i, j, k;
+    if (!map_cpml_work<0>(compact_work, step, compact_cells_per_shot,
+        NX, NY, NZ, pml0, pml1, pml2, pml3, pml4, pml5,
+        &s, &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ + k;
+    long long work = (long long)s * field_stride + idx;
     float upd = update[idx];
 
 #define APPLY_H_PML_GPU(R, P, p, q, stride, coord, n, spacing, field, source, sign) \
@@ -1666,6 +1799,16 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
     long long total_fields = (long long)NX_FIELDS * NY_FIELDS * NZ_FIELDS;
     dim3 grid_material(CEIL_DIV(total_fields, blockSize));
     dim3 grid_fields(CEIL_DIV((long long)step * total_fields, blockSize));
+    long long cpml_e_cells = cpml_cells_per_shot(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 1);
+    long long cpml_h_cells = cpml_cells_per_shot(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 0);
+    int has_cpml_e = has_cpml && cpml_e_cells > 0;
+    int has_cpml_h = has_cpml && cpml_h_cells > 0;
+    dim3 grid_cpml_e(CEIL_DIV((long long)step * cpml_e_cells, blockSize));
+    dim3 grid_cpml_h(CEIL_DIV((long long)step * cpml_h_cells, blockSize));
 
     build_update_coeffs_gpu<<<grid_material, blockSize, 0, stream_comp>>>(eps_r_pad, sigma_pad, mu_r_pad, ce_hist, ce_curl, ce_rhs, ch_hist, ch_curl, ch_rhs, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dt, dx);
     CUDA_CHECK_LAST();
@@ -1717,58 +1860,64 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
         if (fdtd_order == 8) {
             update_h_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(ch_hist, ch_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_h_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_h) {
+                cpml_h_gpu<8><<<grid_cpml_h, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
-                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
+                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2,
+                    cpml_h_cells);
                 CUDA_CHECK_LAST();
             }
             update_e_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(ce_hist, ce_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_e_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_e) {
+                cpml_e_gpu<8><<<grid_cpml_e, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
-                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
+                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2,
+                    cpml_e_cells);
                 CUDA_CHECK_LAST();
             }
         } else if (fdtd_order == 4) {
             update_h_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(ch_hist, ch_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_h_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_h) {
+                cpml_h_gpu<4><<<grid_cpml_h, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
-                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
+                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2,
+                    cpml_h_cells);
                 CUDA_CHECK_LAST();
             }
             update_e_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(ce_hist, ce_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_e_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_e) {
+                cpml_e_gpu<4><<<grid_cpml_e, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
-                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
+                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2,
+                    cpml_e_cells);
                 CUDA_CHECK_LAST();
             }
         } else {
             update_h_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(ch_hist, ch_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_h_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_h) {
+                cpml_h_gpu<2><<<grid_cpml_h, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
-                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
+                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2,
+                    cpml_h_cells);
                 CUDA_CHECK_LAST();
             }
             update_e_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(ce_hist, ce_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_e_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_e) {
+                cpml_e_gpu<2><<<grid_cpml_e, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
-                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
+                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2,
+                    cpml_e_cells);
                 CUDA_CHECK_LAST();
             }
         }
@@ -1964,6 +2113,16 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
     long long total_fields = (long long)NX_FIELDS * NY_FIELDS * NZ_FIELDS;
     dim3 grid_material(CEIL_DIV(total_fields, blockSize));
     dim3 grid_fields(CEIL_DIV((long long)step * total_fields, blockSize));
+    long long cpml_e_cells = cpml_cells_per_shot(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 1);
+    long long cpml_h_cells = cpml_cells_per_shot(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 0);
+    int has_cpml_e = has_cpml && cpml_e_cells > 0;
+    int has_cpml_h = has_cpml && cpml_h_cells > 0;
+    dim3 grid_cpml_e(CEIL_DIV((long long)step * cpml_e_cells, blockSize));
+    dim3 grid_cpml_h(CEIL_DIV((long long)step * cpml_h_cells, blockSize));
 
     build_update_coeffs_gpu<<<grid_material, blockSize, 0, stream_comp>>>(eps_r_pad, sigma_pad, mu_r_pad, ce_hist, ce_curl, ce_rhs, ch_hist, ch_curl, ch_rhs, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dt, dx);
     CUDA_CHECK_LAST();
@@ -2032,24 +2191,24 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
         }
 
         /* E CPML^T -> E update^T -> H CPML^T -> H update^T. */
-        if (has_cpml) {
-            LAUNCH_ORDER_KERNEL(adjoint_cpml_e_gpu, grid_fields, blockSize, stream_comp, fdtd_order,
+        if (has_cpml_e) {
+            LAUNCH_ORDER_KERNEL(adjoint_cpml_e_gpu, grid_cpml_e, blockSize, stream_comp, fdtd_order,
                 lambda_ex, lambda_ey, lambda_ez, lambda_hx, lambda_hy, lambda_hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                 pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
                 x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2,
-                z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
+                z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2, cpml_e_cells);
             CUDA_CHECK_LAST();
         }
         LAUNCH_ORDER_KERNEL(adjoint_e_gpu, grid_fields, blockSize, stream_comp, fdtd_order,
             ce_hist, ce_curl, lambda_ex, lambda_ey, lambda_ez, lambda_hx, lambda_hy, lambda_hz,
             step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
         CUDA_CHECK_LAST();
-        if (has_cpml) {
-            LAUNCH_ORDER_KERNEL(adjoint_cpml_h_gpu, grid_fields, blockSize, stream_comp, fdtd_order,
+        if (has_cpml_h) {
+            LAUNCH_ORDER_KERNEL(adjoint_cpml_h_gpu, grid_cpml_h, blockSize, stream_comp, fdtd_order,
                 lambda_ex, lambda_ey, lambda_ez, lambda_hx, lambda_hy, lambda_hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                 pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
                 x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2,
-                z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
+                z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2, cpml_h_cells);
             CUDA_CHECK_LAST();
         }
         LAUNCH_ORDER_KERNEL(adjoint_h_gpu, grid_fields, blockSize, stream_comp, fdtd_order,

{deepgpr-0.0.16 → deepgpr-0.0.17}/src/DeepGPR.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DeepGPR
-Version: 0.0.16
+Version: 0.0.17
 Summary: PyTorch and CUDA for GPR FWI
 Author-email: Lei Liu <liulei990222@gmail.com>
 License-Expression: MIT

{deepgpr-0.0.16 → deepgpr-0.0.17}/tests/test_discrete_adjoint.py

@@ -67,23 +67,24 @@ def _locations(shape):
     return torch.tensor([[center]], dtype=torch.int32)
 
 
-def _state_problem(order, pml):
+def _state_problem(order, pml, device="cpu"):
     torch.manual_seed(1100 + order + pml)
+    device = torch.device(device)
     shape = (9, 10, 8)
     nx, ny, nz = shape
-    x = torch.linspace(0.0, 1.0, nx)[:, None, None]
-    y = torch.linspace(0.0, 1.0, ny)[None, :, None]
-    z = torch.linspace(0.0, 1.0, nz)[None, None, :]
+    x = torch.linspace(0.0, 1.0, nx, device=device)[:, None, None]
+    y = torch.linspace(0.0, 1.0, ny, device=device)[None, :, None]
+    z = torch.linspace(0.0, 1.0, nz, device=device)[None, None, :]
     eps_r = 3.5 + 0.7 * x + 0.2 * y + 0.1 * z
     sigma = 1.0e-4 + 1.0e-4 * (x + y + z) / 3.0
     mu_r = 1.0 + 0.15 * x + 0.0 * y + 0.05 * z
-    location = _locations(shape)
-    source = torch.zeros((1, 1, 1), dtype=torch.float32)
+    location = _locations(shape).to(device)
+    source = torch.zeros((1, 1, 1), dtype=torch.float32, device=device)
     spacing = (0.020, 0.017, 0.013)
     dt = 1.0e-11
 
     initial_e, initial_h, initial_cpml = DeepGPR.checkpoint_initial_field(
-        device="cpu",
+        device=device,
         dx=spacing,
         dt=dt,
         source_amplitudes=source,
@@ -102,7 +103,7 @@ def _state_problem(order, pml):
     ]
     state = [tensor.clone().requires_grad_(True) for tensor in base_state]
     result = DeepGPR.compute(
-        device="cpu",
+        device=device,
         dx=spacing,
         dt=dt,
         source_amplitudes=source,
@@ -269,6 +270,15 @@ class DiscreteAdjointTests(unittest.TestCase):
     def test_native_cpml_single_step_dot_product_all_faces(self):
         self.assertLess(_state_problem(order=4, pml=2), 3.0e-5)
 
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
+    def test_cuda_cpml_single_step_dot_products_orders_2_4_8(self):
+        for order in (2, 4, 8):
+            with self.subTest(order=order):
+                self.assertLess(
+                    _state_problem(order=order, pml=2, device="cuda"),
+                    2.0e-4,
+                )
+
     def test_source_waveform_gradient_matches_finite_difference(self):
         nx, ny, nt = 12, 14, 60
         eps_r = torch.full((nx, ny), 4.0)