witwin-maxwell 0.3.0__py3-none-win_amd64.whl

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

@@ -0,0 +1,297 @@
+#include <algorithm>
+
+#include "../launch.h"
+#include "../tensors.h"
+#include "common.cuh"
+
+namespace {
+
+__global__ void accumulate_dft_batched_kernel(
+    int64_t ex_numel,
+    int64_t ey_numel,
+    int64_t ez_numel,
+    const float* __restrict__ ex,
+    const float* __restrict__ ey,
+    const float* __restrict__ ez,
+    const float* __restrict__ weighted_cos,
+    const float* __restrict__ weighted_sin,
+    float* __restrict__ ex_real,
+    float* __restrict__ ex_imag,
+    float* __restrict__ ey_real,
+    float* __restrict__ ey_imag,
+    float* __restrict__ ez_real,
+    float* __restrict__ ez_imag) {
+  const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t freq = static_cast<int64_t>(blockIdx.y);
+  const float cos_value = weighted_cos[freq];
+  const float sin_value = weighted_sin[freq];
+  if (index < ex_numel) {
+    const int64_t linear = freq * ex_numel + index;
+    const float value = ex[index];
+    ex_real[linear] += value * cos_value;
+    ex_imag[linear] += value * sin_value;
+  }
+  if (index < ey_numel) {
+    const int64_t linear = freq * ey_numel + index;
+    const float value = ey[index];
+    ey_real[linear] += value * cos_value;
+    ey_imag[linear] += value * sin_value;
+  }
+  if (index < ez_numel) {
+    const int64_t linear = freq * ez_numel + index;
+    const float value = ez[index];
+    ez_real[linear] += value * cos_value;
+    ez_imag[linear] += value * sin_value;
+  }
+}
+
+__global__ void accumulate_dft_single_frequency_kernel(
+    int64_t ex_numel,
+    int64_t ey_numel,
+    int64_t ez_numel,
+    const float* __restrict__ ex,
+    const float* __restrict__ ey,
+    const float* __restrict__ ez,
+    const float* __restrict__ weighted_cos,
+    const float* __restrict__ weighted_sin,
+    float* __restrict__ ex_real,
+    float* __restrict__ ex_imag,
+    float* __restrict__ ey_real,
+    float* __restrict__ ey_imag,
+    float* __restrict__ ez_real,
+    float* __restrict__ ez_imag) {
+  const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const float cos_value = weighted_cos[0];
+  const float sin_value = weighted_sin[0];
+  if (index < ex_numel) {
+    const float value = ex[index];
+    ex_real[index] += value * cos_value;
+    ex_imag[index] += value * sin_value;
+  }
+  if (index < ey_numel) {
+    const float value = ey[index];
+    ey_real[index] += value * cos_value;
+    ey_imag[index] += value * sin_value;
+  }
+  if (index < ez_numel) {
+    const float value = ez[index];
+    ez_real[index] += value * cos_value;
+    ez_imag[index] += value * sin_value;
+  }
+}
+
+__global__ void accumulate_dft_frequency_coarsened_kernel(
+    int64_t max_field_numel,
+    int64_t frequency_count,
+    int64_t ex_numel,
+    int64_t ey_numel,
+    int64_t ez_numel,
+    const float* __restrict__ ex,
+    const float* __restrict__ ey,
+    const float* __restrict__ ez,
+    const float* __restrict__ weighted_cos,
+    const float* __restrict__ weighted_sin,
+    float* __restrict__ ex_real,
+    float* __restrict__ ex_imag,
+    float* __restrict__ ey_real,
+    float* __restrict__ ey_imag,
+    float* __restrict__ ez_real,
+    float* __restrict__ ez_imag) {
+  const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= max_field_numel) {
+    return;
+  }
+  if (index < ex_numel) {
+    const float value = ex[index];
+    for (int64_t freq = 0; freq < frequency_count; ++freq) {
+      const int64_t linear = freq * ex_numel + index;
+      ex_real[linear] += value * weighted_cos[freq];
+      ex_imag[linear] += value * weighted_sin[freq];
+    }
+  }
+  if (index < ey_numel) {
+    const float value = ey[index];
+    for (int64_t freq = 0; freq < frequency_count; ++freq) {
+      const int64_t linear = freq * ey_numel + index;
+      ey_real[linear] += value * weighted_cos[freq];
+      ey_imag[linear] += value * weighted_sin[freq];
+    }
+  }
+  if (index < ez_numel) {
+    const float value = ez[index];
+    for (int64_t freq = 0; freq < frequency_count; ++freq) {
+      const int64_t linear = freq * ez_numel + index;
+      ez_real[linear] += value * weighted_cos[freq];
+      ez_imag[linear] += value * weighted_sin[freq];
+    }
+  }
+}
+
+__global__ void accumulate_dft_component_kernel(
+    int64_t field_numel,
+    const float* __restrict__ field,
+    const float* __restrict__ weighted_cos,
+    const float* __restrict__ 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 >= field_numel) {
+    return;
+  }
+  const int64_t freq = static_cast<int64_t>(blockIdx.y);
+  const int64_t linear = freq * field_numel + index;
+  const float value = field[index];
+  real_accum[linear] += value * weighted_cos[freq];
+  imag_accum[linear] += value * weighted_sin[freq];
+}
+
+void check_dft_component(
+    const torch::stable::Tensor& field,
+    const torch::stable::Tensor& real_accum,
+    const torch::stable::Tensor& imag_accum,
+    int64_t frequency_count,
+    const char* name) {
+  check_float32_tensor(field, name);
+  check_float32_tensor(real_accum, "real_accum");
+  check_float32_tensor(imag_accum, "imag_accum");
+  check_contiguous_tensor(field, name);
+  check_contiguous_tensor(real_accum, "real_accum");
+  check_contiguous_tensor(imag_accum, "imag_accum");
+  check_same_cuda_device(field, real_accum, "real_accum");
+  check_same_cuda_device(field, imag_accum, "imag_accum");
+  STD_TORCH_CHECK(real_accum.dim() == field.dim() + 1, "real_accum rank must be field rank + 1");
+  STD_TORCH_CHECK(imag_accum.sizes().equals(real_accum.sizes()), "imag_accum must match real_accum shape");
+  STD_TORCH_CHECK(real_accum.size(0) == frequency_count, "accum frequency dimension mismatch");
+  for (int64_t dim = 0; dim < field.dim(); ++dim) {
+    STD_TORCH_CHECK(real_accum.size(dim + 1) == field.size(dim), "accum spatial shape must match field");
+  }
+}
+
+}  // namespace
+
+void accumulate_dft_batched_cuda(
+    const torch::stable::Tensor& ex,
+    const torch::stable::Tensor& ey,
+    const torch::stable::Tensor& ez,
+    torch::stable::Tensor ex_real,
+    torch::stable::Tensor ex_imag,
+    torch::stable::Tensor ey_real,
+    torch::stable::Tensor ey_imag,
+    torch::stable::Tensor ez_real,
+    torch::stable::Tensor ez_imag,
+    const torch::stable::Tensor& weighted_cos,
+    const torch::stable::Tensor& weighted_sin) {
+  check_float32_tensor(weighted_cos, "weighted_cos");
+  check_float32_tensor(weighted_sin, "weighted_sin");
+  check_contiguous_tensor(weighted_cos, "weighted_cos");
+  check_contiguous_tensor(weighted_sin, "weighted_sin");
+  STD_TORCH_CHECK(weighted_cos.dim() == 1, "weighted_cos must be rank 1");
+  STD_TORCH_CHECK(weighted_sin.sizes().equals(weighted_cos.sizes()), "weighted_sin must match weighted_cos");
+  const int64_t frequency_count = weighted_cos.size(0);
+  check_dft_component(ex, ex_real, ex_imag, frequency_count, "ex");
+  check_dft_component(ey, ey_real, ey_imag, frequency_count, "ey");
+  check_dft_component(ez, ez_real, ez_imag, frequency_count, "ez");
+  check_same_cuda_device(ex, ey, "ey");
+  check_same_cuda_device(ex, ez, "ez");
+  check_same_cuda_device(ex, weighted_cos, "weighted_cos");
+  check_same_cuda_device(ex, weighted_sin, "weighted_sin");
+  check_same_cuda_device(ex, ey_real, "ey_real");
+  check_same_cuda_device(ex, ey_imag, "ey_imag");
+  check_same_cuda_device(ex, ez_real, "ez_real");
+  check_same_cuda_device(ex, ez_imag, "ez_imag");
+  torch::stable::accelerator::DeviceGuard guard(ex.get_device_index());
+  const int64_t max_field_numel = std::max(ex.numel(), std::max(ey.numel(), ez.numel()));
+  constexpr int block_size = 256;
+  if (frequency_count == 1) {
+    accumulate_dft_single_frequency_kernel<<<linear_grid(max_field_numel, block_size), block_size, 0, current_cuda_stream()>>>(
+        ex.numel(),
+        ey.numel(),
+        ez.numel(),
+        ex.mutable_data_ptr<float>(),
+        ey.mutable_data_ptr<float>(),
+        ez.mutable_data_ptr<float>(),
+        weighted_cos.mutable_data_ptr<float>(),
+        weighted_sin.mutable_data_ptr<float>(),
+        ex_real.mutable_data_ptr<float>(),
+        ex_imag.mutable_data_ptr<float>(),
+        ey_real.mutable_data_ptr<float>(),
+        ey_imag.mutable_data_ptr<float>(),
+        ez_real.mutable_data_ptr<float>(),
+        ez_imag.mutable_data_ptr<float>());
+  } else if (
+      (frequency_count == 4 && max_field_numel < 1000000) ||
+      (frequency_count == 3 && max_field_numel >= 400000 && max_field_numel < 1000000)) {
+    accumulate_dft_frequency_coarsened_kernel<<<linear_grid(max_field_numel, block_size), block_size, 0, current_cuda_stream()>>>(
+        max_field_numel,
+        frequency_count,
+        ex.numel(),
+        ey.numel(),
+        ez.numel(),
+        ex.mutable_data_ptr<float>(),
+        ey.mutable_data_ptr<float>(),
+        ez.mutable_data_ptr<float>(),
+        weighted_cos.mutable_data_ptr<float>(),
+        weighted_sin.mutable_data_ptr<float>(),
+        ex_real.mutable_data_ptr<float>(),
+        ex_imag.mutable_data_ptr<float>(),
+        ey_real.mutable_data_ptr<float>(),
+        ey_imag.mutable_data_ptr<float>(),
+        ez_real.mutable_data_ptr<float>(),
+        ez_imag.mutable_data_ptr<float>());
+  } else if (frequency_count >= 4 && max_field_numel >= 1000000) {
+    const dim3 ex_grid(
+        static_cast<unsigned int>((ex.numel() + block_size - 1) / block_size),
+        static_cast<unsigned int>(frequency_count),
+        1);
+    const dim3 ey_grid(
+        static_cast<unsigned int>((ey.numel() + block_size - 1) / block_size),
+        static_cast<unsigned int>(frequency_count),
+        1);
+    const dim3 ez_grid(
+        static_cast<unsigned int>((ez.numel() + block_size - 1) / block_size),
+        static_cast<unsigned int>(frequency_count),
+        1);
+    accumulate_dft_component_kernel<<<ex_grid, block_size, 0, current_cuda_stream()>>>(
+        ex.numel(),
+        ex.mutable_data_ptr<float>(),
+        weighted_cos.mutable_data_ptr<float>(),
+        weighted_sin.mutable_data_ptr<float>(),
+        ex_real.mutable_data_ptr<float>(),
+        ex_imag.mutable_data_ptr<float>());
+    accumulate_dft_component_kernel<<<ey_grid, block_size, 0, current_cuda_stream()>>>(
+        ey.numel(),
+        ey.mutable_data_ptr<float>(),
+        weighted_cos.mutable_data_ptr<float>(),
+        weighted_sin.mutable_data_ptr<float>(),
+        ey_real.mutable_data_ptr<float>(),
+        ey_imag.mutable_data_ptr<float>());
+    accumulate_dft_component_kernel<<<ez_grid, block_size, 0, current_cuda_stream()>>>(
+        ez.numel(),
+        ez.mutable_data_ptr<float>(),
+        weighted_cos.mutable_data_ptr<float>(),
+        weighted_sin.mutable_data_ptr<float>(),
+        ez_real.mutable_data_ptr<float>(),
+        ez_imag.mutable_data_ptr<float>());
+  } else {
+    const dim3 grid(
+        static_cast<unsigned int>((max_field_numel + block_size - 1) / block_size),
+        static_cast<unsigned int>(frequency_count),
+        1);
+    accumulate_dft_batched_kernel<<<grid, block_size, 0, current_cuda_stream()>>>(
+        ex.numel(),
+        ey.numel(),
+        ez.numel(),
+        ex.mutable_data_ptr<float>(),
+        ey.mutable_data_ptr<float>(),
+        ez.mutable_data_ptr<float>(),
+        weighted_cos.mutable_data_ptr<float>(),
+        weighted_sin.mutable_data_ptr<float>(),
+        ex_real.mutable_data_ptr<float>(),
+        ex_imag.mutable_data_ptr<float>(),
+        ey_real.mutable_data_ptr<float>(),
+        ey_imag.mutable_data_ptr<float>(),
+        ez_real.mutable_data_ptr<float>(),
+        ez_imag.mutable_data_ptr<float>());
+  }
+  WITWIN_CUDA_CHECK();
+}

witwin/maxwell/fdtd/cuda/launch.h

@@ -0,0 +1,23 @@
+#pragma once
+
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/macros.h>
+
+#include <cuda_runtime.h>
+
+inline dim3 linear_grid(long long elements, int block_size = 256) {
+  const auto blocks = (elements + block_size - 1) / block_size;
+  return dim3(static_cast<unsigned int>(blocks), 1, 1);
+}
+
+inline cudaStream_t current_cuda_stream() {
+  void* stream_ptr = nullptr;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(
+      torch::stable::accelerator::getCurrentDeviceIndex(),
+      &stream_ptr));
+  return static_cast<cudaStream_t>(stream_ptr);
+}
+
+#define WITWIN_CUDA_CHECK() STD_CUDA_KERNEL_LAUNCH_CHECK()
+

witwin/maxwell/fdtd/cuda/runtime/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from .dispatcher import NativeFDTDModule, get_native_fdtd_module
+
+__all__ = ["NativeFDTDModule", "get_native_fdtd_module"]
+

witwin/maxwell/fdtd/cuda/runtime/dispatcher.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from ..backend import NativeFDTDModule, get_native_fdtd_module
+
+__all__ = ["NativeFDTDModule", "get_native_fdtd_module"]
+

witwin/maxwell/fdtd/cuda/runtime/graph.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+
+class CudaGraphRunner:
+    def __init__(self, *, enabled: bool = True, warmup_steps: int = 1):
+        self.enabled = bool(enabled)
+        self.warmup_steps = max(0, int(warmup_steps))
+        self.graph: torch.cuda.CUDAGraph | None = None
+
+    def capture(self, fn: Callable[[], None]) -> Callable[[], None]:
+        if not self.enabled:
+            return fn
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA Graph capture requires CUDA.")
+
+        for _ in range(self.warmup_steps):
+            fn()
+        torch.cuda.synchronize()
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            fn()
+        torch.cuda.synchronize()
+        self.graph = graph
+
+        def replay() -> None:
+            graph.replay()
+
+        return replay

witwin/maxwell/fdtd/cuda/tensors.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+
+#include <optional>
+
+inline void check_cuda_tensor(const torch::stable::Tensor& tensor, const char* name) {
+  STD_TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+}
+
+inline void check_same_cuda_device(
+    const torch::stable::Tensor& reference,
+    const torch::stable::Tensor& tensor,
+    const char* name) {
+  STD_TORCH_CHECK(
+      tensor.device() == reference.device(),
+      name,
+      " must be on the same CUDA device as the reference tensor");
+}
+
+inline void check_float32_tensor(const torch::stable::Tensor& tensor, const char* name) {
+  check_cuda_tensor(tensor, name);
+  STD_TORCH_CHECK(
+      tensor.scalar_type() == torch::headeronly::ScalarType::Float,
+      name,
+      " must be float32");
+}
+
+inline void check_contiguous_tensor(const torch::stable::Tensor& tensor, const char* name) {
+  STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+

witwin/maxwell/fdtd/dispersion.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import math
+
+
+def solve_numerical_wavenumber(solver, direction, deltas) -> float:
+    """Solve the discrete dispersion relation for |k|.
+
+    ``deltas`` maps axis name ("x"/"y"/"z") to the grid spacing (float) to use
+    along that axis — the locally uniform spacing of the injection region.
+    """
+    omega = float(solver.source_omega)
+    if omega <= 0.0:
+        return 0.0
+
+    target = math.sin(0.5 * omega * float(solver.dt)) ** 2 / (float(solver.c) ** 2 * float(solver.dt) ** 2)
+    if target <= 0.0:
+        return omega / float(solver.c)
+
+    axis_limits = []
+    for axis, component in zip("xyz", direction):
+        abs_component = abs(float(component))
+        if abs_component <= 1e-12:
+            continue
+        delta = float(deltas[axis])
+        axis_limits.append(2.0 * math.pi / max(abs_component * delta, 1e-30))
+    if not axis_limits:
+        return omega / float(solver.c)
+
+    def dispersion_residual(k_mag):
+        total = 0.0
+        for axis, component in zip("xyz", direction):
+            abs_component = abs(float(component))
+            if abs_component <= 1e-12:
+                continue
+            delta = float(deltas[axis])
+            total += math.sin(0.5 * k_mag * abs_component * delta) ** 2 / (delta * delta)
+        return total - target
+
+    low = 0.0
+    high = min(axis_limits)
+    for _ in range(80):
+        mid = 0.5 * (low + high)
+        if dispersion_residual(mid) >= 0.0:
+            high = mid
+        else:
+            low = mid
+    return 0.5 * (low + high)