witwin-maxwell 0.3.0__py3-none-win_amd64.whl

witwin/maxwell/fdtd/cuda/kernels/observers.cu

@@ -0,0 +1,347 @@
+#include "../launch.h"
+#include "../tensors.h"
+#include "common.cuh"
+
+namespace {
+
+template <typename IndexT>
+__global__ void accumulate_point_observer_kernel(
+    int64_t total,
+    unsigned int ny,
+    unsigned int nz,
+    const float* __restrict__ field,
+    const IndexT* __restrict__ point_i,
+    const IndexT* __restrict__ point_j,
+    const IndexT* __restrict__ point_k,
+    float weighted_cos,
+    float weighted_sin,
+    float* __restrict__ real_accum,
+    float* __restrict__ imag_accum) {
+  const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= total) {
+    return;
+  }
+  const long long field_offset = offset3d(
+      static_cast<unsigned int>(point_i[index]),
+      static_cast<unsigned int>(point_j[index]),
+      static_cast<unsigned int>(point_k[index]),
+      ny,
+      nz);
+  const float value = field[field_offset];
+  real_accum[index] += value * weighted_cos;
+  imag_accum[index] += value * weighted_sin;
+}
+
+template <int Axis>
+__global__ void accumulate_plane_observer_kernel(
+    unsigned int nx,
+    unsigned int ny,
+    unsigned int nz,
+    const float* __restrict__ field,
+    int plane_index,
+    float weighted_cos,
+    float weighted_sin,
+    float* __restrict__ real_accum,
+    float* __restrict__ imag_accum) {
+  unsigned int i = 0;
+  unsigned int j = 0;
+  unsigned int k = 0;
+  int64_t linear = 0;
+  if constexpr (Axis == 0) {
+    k = blockIdx.x * blockDim.x + threadIdx.x;
+    j = blockIdx.y * blockDim.y + threadIdx.y;
+    if (j >= ny || k >= nz) {
+      return;
+    }
+    i = static_cast<unsigned int>(plane_index);
+    linear = static_cast<int64_t>(j) * nz + k;
+  } else if constexpr (Axis == 1) {
+    k = blockIdx.x * blockDim.x + threadIdx.x;
+    i = blockIdx.y * blockDim.y + threadIdx.y;
+    if (i >= nx || k >= nz) {
+      return;
+    }
+    j = static_cast<unsigned int>(plane_index);
+    linear = static_cast<int64_t>(i) * nz + k;
+  } else {
+    j = blockIdx.x * blockDim.x + threadIdx.x;
+    i = blockIdx.y * blockDim.y + threadIdx.y;
+    if (i >= nx || j >= ny) {
+      return;
+    }
+    k = static_cast<unsigned int>(plane_index);
+    linear = static_cast<int64_t>(i) * ny + j;
+  }
+  const float value = field[offset3d(i, j, k, ny, nz)];
+  real_accum[linear] += value * weighted_cos;
+  imag_accum[linear] += value * weighted_sin;
+}
+
+// Single-block reduction of the instantaneous Poynting flux through a plane:
+//   flux = scale * sum_pq (Ea*Hb - Eb*Ha) * weight
+// The four tangential field planes are already Yee-averaged onto the common
+// (P, Q) grid and share `weights`'s layout; the scalar result is written
+// straight into out[out_index] (a preallocated time-series slot), so a sampled
+// step needs no temporaries and no host round-trip. block_size must be a power
+// of two for the tree reduction.
+__global__ void plane_flux_reduce_kernel(
+    int64_t plane_size,
+    float scale,
+    const float* __restrict__ ea,
+    const float* __restrict__ eb,
+    const float* __restrict__ ha,
+    const float* __restrict__ hb,
+    const float* __restrict__ weights,
+    float* __restrict__ out,
+    int64_t out_index) {
+  extern __shared__ float sdata[];
+  const unsigned int tid = threadIdx.x;
+  float local = 0.0f;
+  for (int64_t idx = tid; idx < plane_size; idx += blockDim.x) {
+    local += (ea[idx] * hb[idx] - eb[idx] * ha[idx]) * weights[idx];
+  }
+  sdata[tid] = local;
+  __syncthreads();
+  for (unsigned int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
+    if (tid < stride) {
+      sdata[tid] += sdata[tid + stride];
+    }
+    __syncthreads();
+  }
+  if (tid == 0) {
+    out[out_index] = scale * sdata[0];
+  }
+}
+
+}  // namespace
+
+namespace {
+
+void check_index_vector(const torch::stable::Tensor& tensor, const char* name) {
+  check_cuda_tensor(tensor, name);
+  check_contiguous_tensor(tensor, name);
+  STD_TORCH_CHECK(
+      tensor.scalar_type() == torch::headeronly::ScalarType::Int || tensor.scalar_type() == torch::headeronly::ScalarType::Long,
+      name,
+      " must be int32 or int64");
+  STD_TORCH_CHECK(tensor.dim() == 1, name, " must be rank 1");
+}
+
+template <typename IndexT>
+void launch_point_observers(
+    int64_t total,
+    unsigned int ny,
+    unsigned int nz,
+    const float* __restrict__ field,
+    const IndexT* __restrict__ point_i,
+    const IndexT* __restrict__ point_j,
+    const IndexT* __restrict__ point_k,
+    float weighted_cos,
+    float weighted_sin,
+    float* __restrict__ real_accum,
+    float* __restrict__ imag_accum,
+    int block_size) {
+  accumulate_point_observer_kernel<IndexT><<<linear_grid(total, block_size), block_size, 0, current_cuda_stream()>>>(
+      total,
+      ny,
+      nz,
+      field,
+      point_i,
+      point_j,
+      point_k,
+      weighted_cos,
+      weighted_sin,
+      real_accum,
+      imag_accum);
+}
+
+template <int Axis>
+void launch_plane_observer(
+    unsigned int nx,
+    unsigned int ny,
+    unsigned int nz,
+    const float* __restrict__ field,
+    int plane_index,
+    float weighted_cos,
+    float weighted_sin,
+    float* __restrict__ real_accum,
+    float* __restrict__ imag_accum) {
+  const dim3 block(64, 4, 1);
+  const unsigned int dim_x = Axis == 2 ? ny : nz;
+  const unsigned int dim_y = Axis == 0 ? ny : nx;
+  const dim3 grid((dim_x + block.x - 1) / block.x, (dim_y + block.y - 1) / block.y, 1);
+  accumulate_plane_observer_kernel<Axis><<<grid, block, 0, current_cuda_stream()>>>(
+      nx,
+      ny,
+      nz,
+      field,
+      plane_index,
+      weighted_cos,
+      weighted_sin,
+      real_accum,
+      imag_accum);
+}
+
+}  // namespace
+
+void accumulate_point_observers_cuda(
+    const torch::stable::Tensor& field,
+    const torch::stable::Tensor& point_i,
+    const torch::stable::Tensor& point_j,
+    const torch::stable::Tensor& point_k,
+    torch::stable::Tensor real_accum,
+    torch::stable::Tensor imag_accum,
+    double weighted_cos,
+    double weighted_sin) {
+  check_float32_tensor(field, "field");
+  check_float32_tensor(real_accum, "real_accum");
+  check_float32_tensor(imag_accum, "imag_accum");
+  check_contiguous_tensor(field, "field");
+  check_contiguous_tensor(real_accum, "real_accum");
+  check_contiguous_tensor(imag_accum, "imag_accum");
+  STD_TORCH_CHECK(field.dim() == 3, "field must be rank 3");
+  check_index_vector(point_i, "point_i");
+  check_index_vector(point_j, "point_j");
+  check_index_vector(point_k, "point_k");
+  STD_TORCH_CHECK(point_i.scalar_type() == point_j.scalar_type(), "point_j must match point_i dtype");
+  STD_TORCH_CHECK(point_i.scalar_type() == point_k.scalar_type(), "point_k must match point_i dtype");
+  check_same_cuda_device(field, point_i, "point_i");
+  check_same_cuda_device(field, point_j, "point_j");
+  check_same_cuda_device(field, point_k, "point_k");
+  check_same_cuda_device(field, real_accum, "real_accum");
+  check_same_cuda_device(field, imag_accum, "imag_accum");
+  STD_TORCH_CHECK(real_accum.sizes().equals(imag_accum.sizes()), "observer accumulators must match");
+  STD_TORCH_CHECK(point_i.numel() == real_accum.numel(), "point_i size must match accumulators");
+  STD_TORCH_CHECK(point_j.numel() == real_accum.numel(), "point_j size must match accumulators");
+  STD_TORCH_CHECK(point_k.numel() == real_accum.numel(), "point_k size must match accumulators");
+  torch::stable::accelerator::DeviceGuard guard(field.get_device_index());
+  const auto sizes = field.sizes();
+  const int64_t point_count = real_accum.numel();
+  const int block_size = point_count <= 64 ? 64 : (point_count >= 262144 ? 512 : 256);
+  if (point_i.scalar_type() == torch::headeronly::ScalarType::Int) {
+    launch_point_observers<int>(
+        point_count,
+        static_cast<unsigned int>(sizes[1]),
+        static_cast<unsigned int>(sizes[2]),
+        field.mutable_data_ptr<float>(),
+        point_i.mutable_data_ptr<int>(),
+        point_j.mutable_data_ptr<int>(),
+        point_k.mutable_data_ptr<int>(),
+        static_cast<float>(weighted_cos),
+        static_cast<float>(weighted_sin),
+        real_accum.mutable_data_ptr<float>(),
+        imag_accum.mutable_data_ptr<float>(),
+        block_size);
+  } else {
+    launch_point_observers<int64_t>(
+        point_count,
+        static_cast<unsigned int>(sizes[1]),
+        static_cast<unsigned int>(sizes[2]),
+        field.mutable_data_ptr<float>(),
+        point_i.mutable_data_ptr<int64_t>(),
+        point_j.mutable_data_ptr<int64_t>(),
+        point_k.mutable_data_ptr<int64_t>(),
+        static_cast<float>(weighted_cos),
+        static_cast<float>(weighted_sin),
+        real_accum.mutable_data_ptr<float>(),
+        imag_accum.mutable_data_ptr<float>(),
+        block_size);
+  }
+  WITWIN_CUDA_CHECK();
+}
+
+void accumulate_plane_observer_cuda(
+    const torch::stable::Tensor& field,
+    torch::stable::Tensor real_accum,
+    torch::stable::Tensor imag_accum,
+    int64_t axis,
+    int64_t plane_index,
+    double weighted_cos,
+    double weighted_sin) {
+  check_float32_tensor(field, "field");
+  check_float32_tensor(real_accum, "real_accum");
+  check_float32_tensor(imag_accum, "imag_accum");
+  check_contiguous_tensor(field, "field");
+  check_contiguous_tensor(real_accum, "real_accum");
+  check_contiguous_tensor(imag_accum, "imag_accum");
+  STD_TORCH_CHECK(field.dim() == 3, "field must be rank 3");
+  STD_TORCH_CHECK(axis >= 0 && axis <= 2, "axis must be 0, 1, or 2");
+  STD_TORCH_CHECK(plane_index >= 0 && plane_index < field.size(axis), "plane_index is out of range");
+  check_same_cuda_device(field, real_accum, "real_accum");
+  check_same_cuda_device(field, imag_accum, "imag_accum");
+  STD_TORCH_CHECK(real_accum.sizes().equals(imag_accum.sizes()), "observer accumulators must match");
+  const int64_t expected = axis == 0
+      ? field.size(1) * field.size(2)
+      : (axis == 1 ? field.size(0) * field.size(2) : field.size(0) * field.size(1));
+  STD_TORCH_CHECK(real_accum.numel() == expected, "plane observer accumulator size does not match selected plane");
+  torch::stable::accelerator::DeviceGuard guard(field.get_device_index());
+  const auto sizes = field.sizes();
+  const auto nx = static_cast<unsigned int>(sizes[0]);
+  const auto ny = static_cast<unsigned int>(sizes[1]);
+  const auto nz = static_cast<unsigned int>(sizes[2]);
+  const auto plane = static_cast<int>(plane_index);
+  const auto cos_weight = static_cast<float>(weighted_cos);
+  const auto sin_weight = static_cast<float>(weighted_sin);
+  if (axis == 0) {
+    launch_plane_observer<0>(
+        nx, ny, nz, field.mutable_data_ptr<float>(), plane, cos_weight, sin_weight,
+        real_accum.mutable_data_ptr<float>(), imag_accum.mutable_data_ptr<float>());
+  } else if (axis == 1) {
+    launch_plane_observer<1>(
+        nx, ny, nz, field.mutable_data_ptr<float>(), plane, cos_weight, sin_weight,
+        real_accum.mutable_data_ptr<float>(), imag_accum.mutable_data_ptr<float>());
+  } else {
+    launch_plane_observer<2>(
+        nx, ny, nz, field.mutable_data_ptr<float>(), plane, cos_weight, sin_weight,
+        real_accum.mutable_data_ptr<float>(), imag_accum.mutable_data_ptr<float>());
+  }
+  WITWIN_CUDA_CHECK();
+}
+
+void plane_flux_reduce_cuda(
+    const torch::stable::Tensor& ea,
+    const torch::stable::Tensor& eb,
+    const torch::stable::Tensor& ha,
+    const torch::stable::Tensor& hb,
+    const torch::stable::Tensor& weights,
+    torch::stable::Tensor out,
+    int64_t out_index,
+    double scale) {
+  check_float32_tensor(ea, "ea");
+  check_float32_tensor(eb, "eb");
+  check_float32_tensor(ha, "ha");
+  check_float32_tensor(hb, "hb");
+  check_float32_tensor(weights, "weights");
+  check_float32_tensor(out, "out");
+  check_contiguous_tensor(ea, "ea");
+  check_contiguous_tensor(eb, "eb");
+  check_contiguous_tensor(ha, "ha");
+  check_contiguous_tensor(hb, "hb");
+  check_contiguous_tensor(weights, "weights");
+  check_contiguous_tensor(out, "out");
+  const int64_t plane_size = weights.numel();
+  STD_TORCH_CHECK(ea.numel() == plane_size, "ea size must match weights");
+  STD_TORCH_CHECK(eb.numel() == plane_size, "eb size must match weights");
+  STD_TORCH_CHECK(ha.numel() == plane_size, "ha size must match weights");
+  STD_TORCH_CHECK(hb.numel() == plane_size, "hb size must match weights");
+  STD_TORCH_CHECK(out.dim() == 1, "out must be rank 1");
+  STD_TORCH_CHECK(out_index >= 0 && out_index < out.numel(), "out_index is out of range");
+  check_same_cuda_device(ea, eb, "eb");
+  check_same_cuda_device(ea, ha, "ha");
+  check_same_cuda_device(ea, hb, "hb");
+  check_same_cuda_device(ea, weights, "weights");
+  check_same_cuda_device(ea, out, "out");
+  torch::stable::accelerator::DeviceGuard guard(ea.get_device_index());
+  const int block_size = 256;
+  plane_flux_reduce_kernel<<<1, block_size, block_size * sizeof(float), current_cuda_stream()>>>(
+      plane_size,
+      static_cast<float>(scale),
+      ea.mutable_data_ptr<float>(),
+      eb.mutable_data_ptr<float>(),
+      ha.mutable_data_ptr<float>(),
+      hb.mutable_data_ptr<float>(),
+      weights.mutable_data_ptr<float>(),
+      out.mutable_data_ptr<float>(),
+      out_index);
+  WITWIN_CUDA_CHECK();
+}

witwin/maxwell/fdtd/cuda/kernels/projection.cu

@@ -0,0 +1,211 @@
+#include "../launch.h"
+#include "../tensors.h"
+#include "common.cuh"
+
+namespace {
+
+void check_field3d(const torch::stable::Tensor& field, const char* name) {
+  check_float32_tensor(field, name);
+  check_contiguous_tensor(field, name);
+  STD_TORCH_CHECK(field.dim() == 3, name, " must be a contiguous 3D float32 CUDA tensor");
+}
+
+template <int Axis>
+__device__ __forceinline__ void face_indices(
+    int coord_a,
+    int coord_b,
+    int nx,
+    int ny,
+    int nz,
+    long long& low,
+    long long& high) {
+  if constexpr (Axis == 0) {
+    const int j = coord_a;
+    const int k = coord_b;
+    low = offset3d(0, j, k, ny, nz);
+    high = low + static_cast<long long>(nx - 1) * ny * nz;
+    return;
+  }
+  if constexpr (Axis == 1) {
+    const int i = coord_a;
+    const int k = coord_b;
+    low = offset3d(i, 0, k, ny, nz);
+    high = low + static_cast<long long>(ny - 1) * nz;
+    return;
+  }
+  const int i = coord_a;
+  const int j = coord_b;
+  low = offset3d(i, j, 0, ny, nz);
+  high = low + nz - 1;
+}
+
+template <int Axis>
+__global__ void project_periodic_boundary_kernel(
+    int nx,
+    int ny,
+    int nz,
+    float* __restrict__ field) {
+  const int coord_b = blockIdx.x * blockDim.x + threadIdx.x;
+  const int coord_a = blockIdx.y * blockDim.y + threadIdx.y;
+  const int dim_b = Axis == 2 ? ny : nz;
+  const int dim_a = Axis == 0 ? ny : nx;
+  if (coord_a >= dim_a || coord_b >= dim_b) {
+    return;
+  }
+  long long low = 0;
+  long long high = 0;
+  face_indices<Axis>(coord_a, coord_b, nx, ny, nz, low, high);
+  const float average = 0.5f * (field[low] + field[high]);
+  field[low] = average;
+  field[high] = average;
+}
+
+template <int Axis, bool RealPhase>
+__global__ void project_bloch_boundary_kernel(
+    int nx,
+    int ny,
+    int nz,
+    float phase_cos,
+    float phase_sin,
+    float* __restrict__ real_field,
+    float* __restrict__ imag_field) {
+  const int coord_b = blockIdx.x * blockDim.x + threadIdx.x;
+  const int coord_a = blockIdx.y * blockDim.y + threadIdx.y;
+  const int dim_b = Axis == 2 ? ny : nz;
+  const int dim_a = Axis == 0 ? ny : nx;
+  if (coord_a >= dim_a || coord_b >= dim_b) {
+    return;
+  }
+  long long low = 0;
+  long long high = 0;
+  face_indices<Axis>(coord_a, coord_b, nx, ny, nz, low, high);
+
+  const float low_r = real_field[low];
+  const float low_i = imag_field[low];
+  const float high_r = real_field[high];
+  const float high_i = imag_field[high];
+  const float c = phase_cos;
+
+  float projected_low_r;
+  float projected_low_i;
+  float projected_high_r;
+  float projected_high_i;
+  if constexpr (RealPhase) {
+    projected_low_r = 0.5f * (low_r + c * high_r);
+    projected_low_i = 0.5f * (low_i + c * high_i);
+    projected_high_r = c * projected_low_r;
+    projected_high_i = c * projected_low_i;
+  } else {
+    const float s = phase_sin;
+    projected_low_r = 0.5f * (low_r + c * high_r + s * high_i);
+    projected_low_i = 0.5f * (low_i + c * high_i - s * high_r);
+    projected_high_r = c * projected_low_r - s * projected_low_i;
+    projected_high_i = s * projected_low_r + c * projected_low_i;
+  }
+
+  real_field[low] = projected_low_r;
+  imag_field[low] = projected_low_i;
+  real_field[high] = projected_high_r;
+  imag_field[high] = projected_high_i;
+}
+
+template <int Axis>
+void launch_periodic_projection(int nx, int ny, int nz, float* __restrict__ field) {
+  const dim3 block(16, 16, 1);
+  const unsigned int dim_b = Axis == 2 ? static_cast<unsigned int>(ny) : static_cast<unsigned int>(nz);
+  const unsigned int dim_a = Axis == 0 ? static_cast<unsigned int>(ny) : static_cast<unsigned int>(nx);
+  const dim3 grid((dim_b + block.x - 1) / block.x, (dim_a + block.y - 1) / block.y, 1);
+  project_periodic_boundary_kernel<Axis><<<grid, block, 0, current_cuda_stream()>>>(
+      nx,
+      ny,
+      nz,
+      field);
+}
+
+template <int Axis, bool RealPhase>
+void launch_bloch_projection(
+    int nx,
+    int ny,
+    int nz,
+    float phase_cos,
+    float phase_sin,
+    float* __restrict__ field_real,
+    float* __restrict__ field_imag) {
+  const dim3 block(16, 16, 1);
+  const unsigned int dim_b = Axis == 2 ? static_cast<unsigned int>(ny) : static_cast<unsigned int>(nz);
+  const unsigned int dim_a = Axis == 0 ? static_cast<unsigned int>(ny) : static_cast<unsigned int>(nx);
+  const dim3 grid((dim_b + block.x - 1) / block.x, (dim_a + block.y - 1) / block.y, 1);
+  project_bloch_boundary_kernel<Axis, RealPhase><<<grid, block, 0, current_cuda_stream()>>>(
+      nx,
+      ny,
+      nz,
+      phase_cos,
+      phase_sin,
+      field_real,
+      field_imag);
+}
+
+}  // namespace
+
+void project_periodic_boundary_cuda(torch::stable::Tensor field, int64_t axis) {
+  check_field3d(field, "field");
+  STD_TORCH_CHECK(axis >= 0 && axis < 3, "axis must be in [0, 3)");
+  const torch::stable::accelerator::DeviceGuard device_guard(field.get_device_index());
+  const int nx = static_cast<int>(field.size(0));
+  const int ny = static_cast<int>(field.size(1));
+  const int nz = static_cast<int>(field.size(2));
+  if (axis == 0) {
+    launch_periodic_projection<0>(nx, ny, nz, field.mutable_data_ptr<float>());
+  } else if (axis == 1) {
+    launch_periodic_projection<1>(nx, ny, nz, field.mutable_data_ptr<float>());
+  } else {
+    launch_periodic_projection<2>(nx, ny, nz, field.mutable_data_ptr<float>());
+  }
+  WITWIN_CUDA_CHECK();
+}
+
+void project_bloch_boundary_cuda(
+    torch::stable::Tensor field_real,
+    torch::stable::Tensor field_imag,
+    int64_t axis,
+    double phase_cos,
+    double phase_sin) {
+  check_field3d(field_real, "field_real");
+  check_field3d(field_imag, "field_imag");
+  check_same_cuda_device(field_real, field_imag, "field_imag");
+  STD_TORCH_CHECK(field_real.sizes().equals(field_imag.sizes()), "field_imag must match field_real shape");
+  STD_TORCH_CHECK(axis >= 0 && axis < 3, "axis must be in [0, 3)");
+  const torch::stable::accelerator::DeviceGuard device_guard(field_real.get_device_index());
+  const int nx = static_cast<int>(field_real.size(0));
+  const int ny = static_cast<int>(field_real.size(1));
+  const int nz = static_cast<int>(field_real.size(2));
+  const float phase_cos_f = static_cast<float>(phase_cos);
+  const float phase_sin_f = static_cast<float>(phase_sin);
+  const bool real_phase = phase_sin_f == 0.0f;
+  if (axis == 0) {
+    if (real_phase) {
+      launch_bloch_projection<0, true>(
+          nx, ny, nz, phase_cos_f, phase_sin_f, field_real.mutable_data_ptr<float>(), field_imag.mutable_data_ptr<float>());
+    } else {
+      launch_bloch_projection<0, false>(
+          nx, ny, nz, phase_cos_f, phase_sin_f, field_real.mutable_data_ptr<float>(), field_imag.mutable_data_ptr<float>());
+    }
+  } else if (axis == 1) {
+    if (real_phase) {
+      launch_bloch_projection<1, true>(
+          nx, ny, nz, phase_cos_f, phase_sin_f, field_real.mutable_data_ptr<float>(), field_imag.mutable_data_ptr<float>());
+    } else {
+      launch_bloch_projection<1, false>(
+          nx, ny, nz, phase_cos_f, phase_sin_f, field_real.mutable_data_ptr<float>(), field_imag.mutable_data_ptr<float>());
+    }
+  } else {
+    if (real_phase) {
+      launch_bloch_projection<2, true>(
+          nx, ny, nz, phase_cos_f, phase_sin_f, field_real.mutable_data_ptr<float>(), field_imag.mutable_data_ptr<float>());
+    } else {
+      launch_bloch_projection<2, false>(
+          nx, ny, nz, phase_cos_f, phase_sin_f, field_real.mutable_data_ptr<float>(), field_imag.mutable_data_ptr<float>());
+    }
+  }
+  WITWIN_CUDA_CHECK();
+}