From 5ff176fd2db6159fc38127db289b0c7da25da1c5 Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Wed, 15 Jul 2026 05:01:43 -0700 Subject: [PATCH] feat(fdtd): add multi-GPU joint solve --- FEATURE_LIST.md | 3 + docs/README.md | 26 + docs/reference/fdtd-multi-gpu-joint-solve.md | 315 +++++ scripts/build_fdtd_cuda_prebuilt.py | 5 +- scripts/dev/fdtd/multi_gpu/bench_joint.py | 939 ++++++++++++++ tests/fdtd/cuda/test_cuda_prebuilt_abi.py | 57 + tests/fdtd/multi_gpu/conftest.py | 37 + tests/fdtd/multi_gpu/test_bench_joint_cli.py | 282 ++++ .../test_bounded_standard_dispatch.py | 242 ++++ .../multi_gpu/test_capacity_cpml_stats.py | 237 ++++ .../fdtd/multi_gpu/test_guard_regressions.py | 202 +++ .../multi_gpu/test_monitor_merge_ownership.py | 351 +++++ .../fdtd/multi_gpu/test_monitor_numerical.py | 245 ++++ tests/fdtd/multi_gpu/test_numerical_matrix.py | 418 ++++++ tests/fdtd/multi_gpu/test_partition.py | 282 ++++ tests/fdtd/multi_gpu/test_public_api.py | 219 ++++ .../fdtd/multi_gpu/test_result_persistence.py | 131 ++ .../fdtd/multi_gpu/test_review_regressions.py | 240 ++++ tests/fdtd/multi_gpu/test_sharded_export.py | 184 +++ .../multi_gpu/test_sharded_persistence.py | 278 ++++ .../fdtd/multi_gpu/test_solver_acceptance.py | 390 ++++++ tests/fdtd/multi_gpu/test_transport.py | 231 ++++ witwin/maxwell/__init__.py | 2 + witwin/maxwell/fdtd/cuda/backend.py | 72 ++ witwin/maxwell/fdtd/cuda/build.py | 35 +- witwin/maxwell/fdtd/cuda/extension.cpp | 90 ++ witwin/maxwell/fdtd/cuda/kernels/electric.cu | 160 ++- witwin/maxwell/fdtd/cuda/kernels/magnetic.cu | 106 +- witwin/maxwell/fdtd/cuda/tensors.h | 20 + witwin/maxwell/fdtd/distributed/__init__.py | 10 + witwin/maxwell/fdtd/distributed/capacity.py | 87 ++ witwin/maxwell/fdtd/distributed/counts.py | 25 + .../fdtd/distributed/frequency_counts.py | 24 + .../maxwell/fdtd/distributed/interpolation.py | 35 + .../maxwell/fdtd/distributed/monitor_merge.py | 416 ++++++ witwin/maxwell/fdtd/distributed/output.py | 49 + .../maxwell/fdtd/distributed/persistence.py | 879 +++++++++++++ witwin/maxwell/fdtd/distributed/solver.py | 1137 +++++++++++++++++ .../fdtd/distributed/source_corrections.py | 51 + .../fdtd/distributed/source_rebuild.py | 76 ++ witwin/maxwell/fdtd/distributed/sources.py | 77 ++ witwin/maxwell/fdtd/distributed/transport.py | 161 +++ witwin/maxwell/fdtd/excitation/injection.py | 15 +- witwin/maxwell/fdtd/excitation/temporal.py | 8 + witwin/maxwell/fdtd_parallel.py | 689 ++++++++++ witwin/maxwell/result.py | 141 +- witwin/maxwell/simulation.py | 59 +- 47 files changed, 9687 insertions(+), 51 deletions(-) create mode 100644 docs/README.md create mode 100644 docs/reference/fdtd-multi-gpu-joint-solve.md create mode 100644 scripts/dev/fdtd/multi_gpu/bench_joint.py create mode 100644 tests/fdtd/cuda/test_cuda_prebuilt_abi.py create mode 100644 tests/fdtd/multi_gpu/conftest.py create mode 100644 tests/fdtd/multi_gpu/test_bench_joint_cli.py create mode 100644 tests/fdtd/multi_gpu/test_bounded_standard_dispatch.py create mode 100644 tests/fdtd/multi_gpu/test_capacity_cpml_stats.py create mode 100644 tests/fdtd/multi_gpu/test_guard_regressions.py create mode 100644 tests/fdtd/multi_gpu/test_monitor_merge_ownership.py create mode 100644 tests/fdtd/multi_gpu/test_monitor_numerical.py create mode 100644 tests/fdtd/multi_gpu/test_numerical_matrix.py create mode 100644 tests/fdtd/multi_gpu/test_partition.py create mode 100644 tests/fdtd/multi_gpu/test_public_api.py create mode 100644 tests/fdtd/multi_gpu/test_result_persistence.py create mode 100644 tests/fdtd/multi_gpu/test_review_regressions.py create mode 100644 tests/fdtd/multi_gpu/test_sharded_export.py create mode 100644 tests/fdtd/multi_gpu/test_sharded_persistence.py create mode 100644 tests/fdtd/multi_gpu/test_solver_acceptance.py create mode 100644 tests/fdtd/multi_gpu/test_transport.py create mode 100644 witwin/maxwell/fdtd/distributed/__init__.py create mode 100644 witwin/maxwell/fdtd/distributed/capacity.py create mode 100644 witwin/maxwell/fdtd/distributed/counts.py create mode 100644 witwin/maxwell/fdtd/distributed/frequency_counts.py create mode 100644 witwin/maxwell/fdtd/distributed/interpolation.py create mode 100644 witwin/maxwell/fdtd/distributed/monitor_merge.py create mode 100644 witwin/maxwell/fdtd/distributed/output.py create mode 100644 witwin/maxwell/fdtd/distributed/persistence.py create mode 100644 witwin/maxwell/fdtd/distributed/solver.py create mode 100644 witwin/maxwell/fdtd/distributed/source_corrections.py create mode 100644 witwin/maxwell/fdtd/distributed/source_rebuild.py create mode 100644 witwin/maxwell/fdtd/distributed/sources.py create mode 100644 witwin/maxwell/fdtd/distributed/transport.py create mode 100644 witwin/maxwell/fdtd_parallel.py diff --git a/FEATURE_LIST.md b/FEATURE_LIST.md index 3c0a026..467a9ab 100644 --- a/FEATURE_LIST.md +++ b/FEATURE_LIST.md @@ -233,6 +233,9 @@ result = mw.Simulation.fdtd(scene, frequencies=[200e12]).run() - Time-domain simulation through `Simulation.fdtd(...)` - CUDA-only solver execution; `Simulation.fdtd(...)` requires `Scene(device="cuda")` - Compiled native CUDA kernels are the only FDTD runtime backend; the Torch reference path has been retired from the runtime and relocated to the CUDA parity tests +- Engineering-preview multi-GPU joint solves through public `FDTDParallelConfig` and `Simulation.fdtd(..., parallel=...)`: Python remains the control plane, and `DistributedFDTD` owns an independent distributed coordinator/time loop for x-slab ownership, P2P halos, streams, and events. Single- and multi-GPU execution share the same native CUDA numerical core; the six bounded standard real-field operations and legacy full-domain wrappers invoke the same implementations with different x ranges +- Monitor-first multi-GPU result semantics: `gather_fields=False` returns supported point/time and assembled plane/flux/mode payloads on `result_device` without a global field, while `gather_fields=True` preflights and gathers global electric last-step or DFT fields. y/z-normal `PlaneMonitor`/`FinitePlaneMonitor`/`FluxMonitor`/`ModeMonitor` payloads are tiled across owned x intervals; x-normal planes have one owner, with an explicit guard for ordinary `Ex` planes exactly on an internal split. Gathered `Result.save/load` and manifest-based `save_sharded/load_sharded` provide eager or lazy/gathered persistence without restoring live transport state +- Qualification remains limited to two RTX A6000 GPUs connected by NV4 and a guarded real-valued forward subset: no three/four-GPU or PCIe-only qualification, Nsight trace, NCCL/multi-node execution, distributed adjoint, peer-aware CUDA Graph, advanced source families, x periodic/Bloch/symmetry, nonlinear/full-off-diagonal media, or SIBC claim is made. The final CUDA 13 acceptance observed 1.97385x strong speedup for the `257^3` vacuum case and exact single/two-GPU diagnostics, while the `129 x 65 x 65` case remained below break-even; see the [FDTD multi-GPU joint-solve guide](docs/reference/fdtd-multi-gpu-joint-solve.md) for the complete support matrix, persistence contract, test checkpoints, and benchmark table - Single- or multi-frequency DFT extraction through `frequency=` or `frequencies=[...]` - Source temporal frequency (`source_time.frequency`) remains distinct from simulation / monitor extraction frequencies; Maxwell does not infer extraction frequencies implicitly - ADE-based electric and magnetic dispersive-material updates for Debye, Drude, and Lorentz media diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..8c16da5 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,26 @@ +# Documentation + +The documentation tree is organized by document purpose. Solver and feature names stay in filenames rather than creating another directory layer. + +## Active documentation + +- [`plans/`](plans/): implementation plans and functional roadmaps that still have open work. +- [`memory/`](memory/): current implementation notes and durable engineering context. +- [`assessments/`](assessments/): current audits, gap analyses, and capability evaluations. +- [`reference/`](reference/): maintained technical facts, capability matrices, and code-linked inventories, including the [FDTD multi-GPU joint-solve guide](reference/fdtd-multi-gpu-joint-solve.md). +- [`releases/`](releases/): release notes. +- [`validation/`](validation/): stable release-validation records. + +The active functional roadmap is indexed in [`plans/next-functional-2026-07/README.md`](plans/next-functional-2026-07/README.md). + +## Archive + +[`archive/`](archive/) contains completed or superseded plans, historical assessments, implementation snapshots, and design records. Archived material preserves engineering history but is not the current source of truth. + +## Placement rules + +1. Choose the directory from the document's purpose, not its solver backend. +2. Use a descriptive filename such as `fdtd-cuda-graph-option-b.md`; do not add solver-specific directory layers. +3. Move a document to `archive/` when its work is complete, its implementation model no longer exists, or a newer document supersedes its facts. +4. Keep repository-wide user documents such as `README.md`, `CHANGELOG.md`, and `FEATURE_LIST.md` at the repository root. +5. Keep executable development utilities under `scripts/`, never under `docs/`. diff --git a/docs/reference/fdtd-multi-gpu-joint-solve.md b/docs/reference/fdtd-multi-gpu-joint-solve.md new file mode 100644 index 0000000..d2f179e --- /dev/null +++ b/docs/reference/fdtd-multi-gpu-joint-solve.md @@ -0,0 +1,315 @@ +# FDTD Multi-GPU Joint Solve + +- Status: engineering preview +- Execution model: one Python process, homogeneous NVIDIA GPUs, CUDA peer-to-peer +- Validated hardware scope: two NVIDIA RTX A6000 GPUs only + +This document is the maintained user and acceptance guide for one FDTD solve that +spans multiple GPUs. It implements the joint-solve direction from the +[detailed implementation plan](../plans/fdtd-multi-gpu-implementation-plan.md) and +the [multi-GPU execution roadmap](../plans/next-functional-2026-07/02-multi-gpu-execution.md). +It is not the independent-simulation ensemble scheduler described in the roadmap. + +## Quick start + +The public entrypoint remains `Scene -> Simulation -> Result`. Add an +`FDTDParallelConfig` to an ordinary FDTD simulation: + +```python +import witwin.maxwell as mw + +# Construct the scene on a participating CUDA device. PointDipole sources must use +# profile="ideal" in the current multi-GPU runtime. +scene = mw.Scene( + domain=mw.Domain(bounds=((-1.0, 1.0), (-0.5, 0.5), (-0.5, 0.5))), + grid=mw.GridSpec.uniform(0.02), + boundary=mw.BoundarySpec.pml(num_layers=8), + device="cuda:0", +) +scene.add_source( + mw.PointDipole( + position=(0.0, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=mw.CW(frequency=1.0e9), + name="source", + ) +) +scene.add_monitor( + mw.PointMonitor("probe", (0.2, 0.0, 0.0), fields=("Ez",)) +) + +parallel = mw.FDTDParallelConfig( + devices=("cuda:0", "cuda:1"), + decomposition_axis="x", + transport="auto", # "auto" and "cuda_p2p" use the same-process P2P path + overlap=True, + gather_fields=False, # monitor-first; do not rebuild a global field by default + result_device="cuda:0", +) + +simulation = mw.Simulation.fdtd( + scene, + frequencies=(1.0e9,), + run_time=mw.TimeConfig(time_steps=1000), + parallel=parallel, + cuda_graph=False, +) + +# prepare() performs topology, homogeneity, and partition checks. run() performs the +# DFT/output-capacity preflight before advancing the first time step. +prepared = simulation.prepare() +result = prepared.run() + +probe = result.monitor("probe") +parallel_stats = result.solver_stats["parallel_stats"] +print(parallel_stats["partitions"]) +print(parallel_stats["peak_memory_bytes"]) +``` + +`devices` is ordered and must contain at least two unique, explicitly indexed CUDA +devices. `result_device` defaults to the first device and must be one of the +participants. Only `decomposition_axis="x"` is accepted. Passing `parallel=None` +continues to use the unchanged single-GPU solver. + +`transport="auto"` currently selects CUDA P2P. `transport="nccl"` is reserved but +raises `RuntimeError`; it does not silently switch to P2P or stage through host +memory. + +## Result placement and gathering + +The default is monitor-first because gathering a large global field can erase the +memory-capacity benefit of domain decomposition. + +| Setting | Public result behavior | +| --- | --- | +| `gather_fields=False` | Supported point/time and assembled plane/flux/mode monitor payloads and metadata are moved to `result_device`; `result.fields` is empty. Rank-local mutable fields remain solver internals and are not a public distributed-field object. | +| `gather_fields=True`, one frequency, `full_field_dft=False` | The last-step global `Ex`, `Ey`, and `Ez` fields are assembled on `result_device`. | +| `gather_fields=True`, full-field DFT or multiple frequencies | Global electric DFT fields are assembled on `result_device`; multi-frequency tensors retain the leading frequency axis. | + +Magnetic full fields are not emitted by the public `Result` field contract in this +preview. Request magnetic components through supported point monitors when they are +needed. + +Before a gathered run starts, the solver checks the destination's free memory for the +global electric output plus pending local DFT work. It raises `MemoryError` rather than +starting a solve that cannot complete the requested gather. The check is reported in +`parallel_stats["gather_preflight"]`. + +### Result persistence + +`Result.save(path)` and `Result.load(path, scene=...)` keep their ordinary gathered +snapshot semantics: fields, monitors, metadata, and solver statistics are serialized, +but live solvers, CUDA streams, peer transport state, and prepared-scene runtime state +are not. + +For a distributed result, `Result.save_sharded(directory)` asks the coordinator for +owned electric field shards and writes a manifest plus one detached CPU tensor payload +per rank. It does not assemble a global field first. The manifest records component +global x intervals, tensor shapes/dtypes, frequency metadata, and rank files; monitor +payloads, result metadata, and solver statistics remain in the result-level metadata. + +`Result.load_sharded(directory, scene=..., gather_fields=False)` is lazy with respect +to rank tensor deserialization. It validates the manifest and required rank paths, +returns `fields == {}`, and exposes the sharded manifest and rank paths. Passing +`gather_fields=True` loads and validates the rank payloads, stitches each component's +contiguous owned intervals in global x order, and returns conventional gathered +fields on `map_location`. These APIs restore detached inference results, not a +resumable distributed solve. + +## Architecture and ownership + +This is spatial domain decomposition, not replicated data parallelism: + +1. `FDTDPartitionPlan` balances the global physical x cells into contiguous half-open + slabs. PML cells at the global x faces stay with the first and last slabs. +2. `FDTDShardLayout` gives every staggered component explicit owned, allocated, halo, + and global-coordinate slices. `Ex`, `Hy`, and `Hz` follow cell ownership; + `Ey`, `Ez`, and `Hx` follow low-node ownership, with the terminal node assigned to + the final shard. An interface value therefore has exactly one owner. +3. Python is the control plane in both modes, but the time-loop ownership differs: + the single-GPU solver owns its ordinary loop, while `DistributedFDTD` is an + independent distributed coordinator with its own shard orchestration, halo exchange, + event ordering, and distributed time loop. This is not a claim that the Python loops + are the same. +4. The numerical core is shared. Every shard calls the same native CUDA Yee, boundary, + CPML, conductivity, and ADE kernels as the single-GPU solver. The standard real-field + path exposes six bounded native operations for `Hx/Hy/Hz` and `Ex/Ey/Ez`; the + legacy full-domain standard wrappers invoke those same implementations with the full + x range. There is no second Python or Torch field-update implementation. +5. One-plane asymmetric halos carry only the tangential components required by the + second-order stencil. `Ey/Ez` move from the right shard to the left shard before + the magnetic update. `Hy/Hz` move from the left shard to the right shard before + the electric update. +6. A dedicated communication stream and CUDA events order each peer copy against both + producer and destination use. With `overlap=True`, eligible standard real-field + runs compute interior regions while halos are in flight and wait only for the + boundary strip. CPML and modulation currently use the serialized-safe schedule; + inspect `overlap_active` instead of assuming a request was activated. +7. Only the first and last slabs apply physical x-face boundary behavior. Internal + interfaces never apply PML, PEC, PMC, Mur, periodic, or symmetry boundary rules. + Every shard still owns the global y/z faces. + +The steady-state fields and field-sized CPML/ADE/DFT state are shard local. The +coordinator prepares global grid metadata and a common time step, but does not +initialize a complete duplicate six-field simulation on every GPU. + +## Current support matrix + +“Accepted” means the current public runtime has an implementation path. “A6000 parity” +means a two-GPU case is in the hardware acceptance suite and is compared with the +single-GPU native solver. Features accepted by the runtime but not named in that suite +must not be described as independently hardware-qualified. + +| Area | Current joint-solve status | Evidence / notes | +| --- | --- | --- | +| Devices and topology | Accepted for homogeneous, explicitly indexed NVIDIA CUDA devices with bidirectional direct peer access | Two RTX A6000 GPUs validated. No CPU/host-staged fallback. | +| Decomposition | x slabs, one-cell halo, balanced or uneven physical-cell counts | A6000 parity covers uniform and uneven nonuniform x grids, including an interface source. | +| Grid | Uniform, anisotropic-spacing, custom/nonuniform, and prepared auto grids through shard-local custom coordinates | Uniform and uneven custom-grid parity are explicit acceptance cases. | +| Linear materials | Vacuum, scalar dielectric, diagonal electric/magnetic tensors, scalar/diagonal conductivity | Scalar dielectric plus CPML and `sigma_e` are in A6000 parity. Full off-diagonal electric tensors fail fast. | +| Dispersion | Electric and magnetic Debye/Drude/Lorentz ADE state is shard local | Conductive electric Lorentz and an interface-crossing magnetic Lorentz case have A6000 parity. Other pole combinations retain runtime support but are not separately hardware-qualified. | +| Modulation | Forward serialized-safe path | Communication/compute overlap is disabled for this path; no final A6000 acceptance claim yet. | +| Nonlinearity | Rejected | Additional collocation halos and bounded nonlinear kernels are required. | +| Absorbers and boundaries | Global-face CPML/PML/StablePML/absorber behavior and mixed `none`/PEC/PMC/Mur; y/z periodic and y/z symmetry are accepted | XYZ CPML plus dielectric, mixed PEC/PMC/Mur, y periodic, and y symmetry have two-GPU parity coverage. x periodic, all Bloch, and x symmetry fail fast. | +| Sources | `PointDipole(profile="ideal")` and `UniformCurrentSource` | Point sources on/near interfaces and a current volume crossing the interface have A6000 parity coverage. Plane-wave/TFSF, beam, mode, and other surface sources fail fast. | +| Monitors | Point spectral monitor, point `FieldTimeMonitor`, a valid `DipoleEmissionMonitor`, and spectral `PlaneMonitor`/`FinitePlaneMonitor`/`FluxMonitor`/`ModeMonitor` | y/z-normal planes are tiled across x and stitched from owned component intervals on `result_device`; x-normal planes have exactly one shard owner. Five Plane/Flux/Mode A6000 numerical cases observed exact single/two-GPU parity. Closed-surface, diffraction, flux-time, non-point field-time, and material monitors remain rejected. | +| Spectral output | One or many requested frequencies; supported point/plane monitor assembly; optional gathered electric fields | Multiple frequencies preserve frequency metadata and leading frequency dimension. | +| Persistence | Gathered `save`/`load` and distributed `save_sharded`/`load_sharded` | Lazy sharded load leaves `fields` empty and rank tensors unopened; explicit gather validates and stitches owned intervals. Persistence does not restore live solver/transport state. | +| Source normalization | Accepted only when exactly one logical source is present | Multiple logical sources fail fast when normalization is requested. | +| Auto shutoff | Global owned-electric-energy reduction on `result_device` | `steps_run`, halo totals, and normalization reflect early termination. | +| CUDA Graph | Rejected | Peer communication is not graph captured. | +| Plotting | Rejected during solve | Run with `gather_fields=True`, then consume the gathered result explicitly. | +| Trainable scenes / adjoint | Rejected before distributed allocation | The existing differentiable path remains single GPU. No multi-GPU backward or distributed checkpoint/replay claim. | +| NCCL / multi-process / multi-node | Reserved, not implemented | `transport="nccl"` raises explicitly. | +| Three or four GPUs | Structurally representable by the partition and neighbor transport | Not qualified in the current hardware acceptance record; no validation claim is made here. | + +## Fail-fast conditions + +Configuration, preparation, or the solve preflight raises a specific error for the +following cases rather than silently changing the physics or execution model: + +- fewer than two devices, duplicate/unindexed/non-CUDA devices, a non-x axis, or a + `result_device` outside the participant list; +- unavailable devices, heterogeneous device model/compute capability, any missing + bidirectional neighbor P2P link, or missing bidirectional P2P between + `result_device` and a shard; +- too few physical x cells for the shard count/halo, or insufficient device memory + for local DFT state or an explicitly requested gather; +- trainable scene parameters, `MaterialRegion` density designs, nonlinear media, + full off-diagonal electric anisotropy, or lossy-metal/SIBC ownership; +- Bloch boundaries, periodic x, or x-axis symmetry; +- ports, unsupported source classes, non-ideal point-dipole profiles, or an invalid + dipole-emission source reference; +- `ClosedSurfaceMonitor`, `DiffractionMonitor`, `FluxTimeMonitor`, non-point + `FieldTimeMonitor`, `PermittivityMonitor`, and `MediumMonitor`; +- an ordinary non-flux x-normal `PlaneMonitor` or `FinitePlaneMonitor` that requests + `Ex` exactly on an internal x partition node. The owner lacks a current + monitor-only low-`Ex` halo; move the plane off the split, remove `Ex`, or use a + tangential-field `FluxMonitor`/`ModeMonitor`; +- `cuda_graph=True`, solve-time plotting, NCCL transport, or source normalization with + more than one logical source. + +This list is intentionally narrower than the single-GPU FDTD feature list. Removing a +guard requires interface-aware numerical parity and hardware evidence; accepting a +scene locally is not enough. + +## Diagnostics and solver statistics + +Use either `result.solver_stats["parallel_stats"]` or +`result.stats()["parallel_stats"]`. The nested record includes: + +- `devices`, `decomposition_axis`, `transport`, and the peer `topology` snapshot; +- `overlap_requested` and the truthful `overlap_active` result; +- `gather_fields`, `result_device`, `gather_preflight`, and per-device `dft_preflight`; +- per-rank physical/global cell and node extents in `partitions`; +- `halo_bytes_per_step` and `halo_bytes_total`; +- per-device `peak_memory_bytes` before gathering and + `peak_memory_bytes_including_gather` after output collection; +- `wall_time_s`. + +`compute_time_s`, `communication_time_s`, and +`exposed_communication_time_s` are currently `None`. Per-phase timing would require +synchronization that is deliberately absent from the production loop; use a CUDA +profiler and retain the accompanying `timing_note` instead of interpreting these keys +as zero communication cost. + +## Dual-A6000 acceptance + +Run the tests and benchmark from the repository root with both A6000 devices visible. +The examples below use the repository's `maxwell` environment; an equivalent Python +environment with the current native extension is acceptable. + +```bash +nvidia-smi topo -m + +/home/xingyu/miniconda3/envs/maxwell/bin/python -m pytest \ + tests/fdtd/multi_gpu -q --basetemp=/tmp/witwin-mgpu-pytest + +/home/xingyu/miniconda3/envs/maxwell/bin/python \ + scripts/dev/fdtd/multi_gpu/bench_joint.py \ + --devices cuda:0 cuda:1 \ + --nodes-x 257 --nodes-y 257 --nodes-z 257 \ + --steps 100 --warmups 1 --repeats 5 \ + --weak-scaling --p2p-size-mib 256 \ + --assert-gates \ + --json /tmp/fdtd-multi-gpu-a6000.json +``` + +Run `--assert-gates` only on an otherwise idle, thermally stable host. The harness +checks a six-component one-versus-two-GPU diagnostic, both P2P directions, strong and +weak scaling, and peak memory. Its default gates are: + +- diagnostic maximum absolute error at most `2e-6`, maximum relative error at most + `2e-5`, finite fields, and nonzero reference signal; +- median P2P bandwidth at least 40 GB/s; +- strong-scaling speedup at least 1.0; +- weak-scaling efficiency at least 0.70; +- maximum per-GPU peak at most 1.15 times the ideal half-single-GPU allocation, plus + any explicitly supplied fixed-overhead allowance. + +The benchmark's default `gather_fields=False` measures the scalable monitor-first +solve. Run an additional `--gather-fields` pass when validating result assembly, but +do not mix gather allocation/time into the core scaling number. + +### Clean-build test record + +The final bounded-operator source was rebuilt cleanly with the CUDA 13 toolchain before +this acceptance record. Scoped checkpoints from that build are listed separately +because their membership can evolve; no aggregate total is implied: + +- the multi-GPU suite passed 109 tests, then the five dual-A6000 numerical + Plane/Flux/Mode cases raised that suite checkpoint to 114 passed; +- the CUDA FDTD suite passed 65 tests; +- the public API suite passed 31 tests; +- the focused numerical/persistence matrix passed 7 tests. + +### Final benchmark measurements + +These are final measurements for this engineering-preview implementation on the named +two-A6000 host, not estimates from the pre-bounded-kernel prototype. The memory factor +is the reported maximum per-GPU peak divided by ideal half of the one-GPU peak. + +| Workload | One GPU | Two GPUs | Strong speedup | Weak efficiency | Memory factor | Numerical observation | +| --- | ---: | ---: | ---: | ---: | ---: | --- | +| Vacuum `257 x 257 x 257`, 100 steps | 3.4659 ms/step | 1.7559 ms/step | 1.97385 | 0.99274 | 1.01058 | Six-field diagnostic observed exact single/two-GPU agreement | +| CPML dielectric `257 x 129 x 129`, slab CPML | 1.0040 ms/step | 0.8116 ms/step | 1.23701 | — | 1.03698 | Diagnostic observed exact agreement | +| Two-frequency full-field DFT `257 x 129 x 129` | 1.7001 ms/step | 1.1110 ms/step | 1.53026 | — | 0.76273 | Multi-frequency DFT numerical acceptance passed | +| Two-frequency full-field DFT `129 x 65 x 65` | 0.3012 ms/step | 1.1284 ms/step | 0.26692 | — | passed | Numerically exact and memory gate passed; below the two-GPU break-even size | + +The same host measured 52.65-52.69 GB/s direct P2P bandwidth across the two directions. +The small case is intentionally retained: it shows that correct domain decomposition +does not imply a speedup when launch and halo overhead dominate. + +### Evidence boundary + +The acceptance host contained only two RTX A6000 GPUs connected through NV4. No +three- or four-GPU system and no PCIe-only peer topology were available, so this record +does not qualify either case. Nsight Systems and Nsight Compute were not available on +the host; the benchmark therefore supplies no profiler trace or per-kernel roofline +claim. + +This remains an engineering preview. The measurements and tests above do not qualify +NCCL, multi-process or multi-node execution, distributed adjoint/checkpoint replay, +peer-aware CUDA Graph capture, advanced plane/beam/mode/TFSF sources, ports, x-periodic +or Bloch decomposition, x symmetry, nonlinear media, full off-diagonal media, SIBC, or +any other guarded feature. diff --git a/scripts/build_fdtd_cuda_prebuilt.py b/scripts/build_fdtd_cuda_prebuilt.py index b35ef80..684d99b 100644 --- a/scripts/build_fdtd_cuda_prebuilt.py +++ b/scripts/build_fdtd_cuda_prebuilt.py @@ -43,7 +43,7 @@ def main() -> None: build_dir = Path( os.environ.get( "WITWIN_MAXWELL_FDTD_CUDA_BUILD_DIR", - Path(tempfile.gettempdir()) / "witwin_maxwell_fdtd_cuda_wheel" / "stable_abi_v1", + Path(tempfile.gettempdir()) / "witwin_maxwell_fdtd_cuda_wheel" / "stable_abi_v2", ) ) os.environ["WITWIN_MAXWELL_FDTD_CUDA_BUILD_DIR"] = str(build_dir) @@ -61,6 +61,9 @@ def main() -> None: existing.unlink() target = cuda_build.prebuilt_extension_path() shutil.copy2(module_file, target) + marker = cuda_build.prebuilt_abi_marker_path(target) + marker.write_text(f"{cuda_build.STABLE_ABI_VERSION}\n", encoding="utf-8") + print(f"Built prebuilt FDTD CUDA extension: {target}") diff --git a/scripts/dev/fdtd/multi_gpu/bench_joint.py b/scripts/dev/fdtd/multi_gpu/bench_joint.py new file mode 100644 index 0000000..520c0c5 --- /dev/null +++ b/scripts/dev/fdtd/multi_gpu/bench_joint.py @@ -0,0 +1,939 @@ +"""Hardware-qualified performance and memory harness for joint-solve FDTD. + +The harness intentionally stays outside pytest: wall-clock assertions are noisy on a +shared GPU host. It reports synchronized solve timing, per-device allocator peaks, +strong/optional weak scaling, and a direct CUDA P2P bandwidth health check. Example:: + + python scripts/dev/fdtd/multi_gpu/bench_joint.py \ + --workload cpml_dielectric \ + --nodes-x 513 --nodes-y 257 --nodes-z 257 --steps 500 --repeats 5 \ + --weak-scaling --json /tmp/fdtd-multi-gpu-a6000.json + +Pass ``--assert-gates`` only on an otherwise idle, thermally stable machine. The +memory gate accepts an explicitly measured fixed per-device overhead instead of +hiding it inside the 15% local-volume allowance from the implementation plan. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import platform +import statistics +import subprocess +import sys +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable + +import numpy as np +import torch + +import witwin.maxwell as mw + + +_FREQUENCY = 1.0e9 +_MIB = 1024**2 +_FIELD_NAMES = ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") +_MAX_FIELD_ABS = 2.0e-6 +_MAX_FIELD_REL = 2.0e-5 + + +@dataclass(frozen=True) +class _WorkloadSpec: + name: str + frequencies: tuple[float, ...] + full_field_dft: bool + cpml_mode: str + dielectric: bool = False + broadband_source: bool = False + graph_mode: str = "disabled" + + +_WORKLOADS = { + "vacuum": _WorkloadSpec( + name="vacuum", + frequencies=(_FREQUENCY,), + full_field_dft=False, + cpml_mode="disabled", + ), + "cpml_dielectric": _WorkloadSpec( + name="cpml_dielectric", + frequencies=(_FREQUENCY,), + full_field_dft=False, + cpml_mode="slab", + dielectric=True, + ), + "multifrequency_dft": _WorkloadSpec( + name="multifrequency_dft", + frequencies=(0.8e9, 1.2e9), + full_field_dft=True, + cpml_mode="disabled", + broadband_source=True, + ), +} + + +def _workload_spec(name: str) -> _WorkloadSpec: + try: + return _WORKLOADS[str(name)] + except KeyError as exc: + choices = ", ".join(_WORKLOADS) + raise ValueError(f"Unknown workload {name!r}; expected one of: {choices}.") from exc + + +def _require_devices(values: tuple[str, ...]) -> tuple[torch.device, torch.device]: + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + raise SystemExit("This benchmark requires at least two visible CUDA devices.") + if len(values) != 2: + raise SystemExit("The current x-slab benchmark requires exactly two devices.") + devices = tuple(torch.device(value) for value in values) + for device in devices: + if device.type != "cuda" or device.index is None: + raise SystemExit(f"Expected indexed CUDA devices, got {device}.") + if device.index >= torch.cuda.device_count(): + raise SystemExit(f"{device} is not visible to this process.") + if not torch.cuda.can_device_access_peer(devices[0].index, devices[1].index): + raise SystemExit(f"No direct peer access from {devices[0]} to {devices[1]}.") + if not torch.cuda.can_device_access_peer(devices[1].index, devices[0].index): + raise SystemExit(f"No direct peer access from {devices[1]} to {devices[0]}.") + properties = tuple(torch.cuda.get_device_properties(device) for device in devices) + if len({(item.name, item.major, item.minor) for item in properties}) != 1: + raise SystemExit("Joint-solve acceptance requires homogeneous CUDA devices.") + return devices + + +def _axis_nodes(count: int, spacing: float) -> np.ndarray: + if count < 4: + raise ValueError("Every benchmark node count must be >= 4.") + centered = np.arange(count, dtype=np.float64) - 0.5 * (count - 1) + return centered * float(spacing) + + +def _vacuum_scene( + shape: tuple[int, int, int], + *, + spacing: float, + device: torch.device, +) -> mw.Scene: + x, y, z = (_axis_nodes(count, spacing) for count in shape) + source_position = (float(x[len(x) // 2]), 0.0, 0.0) + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.none(), + device=str(device), + ) + scene.add_source( + mw.PointDipole( + position=source_position, + polarization="Ez", + profile="ideal", + source_time=mw.CW( + frequency=_FREQUENCY, + amplitude=1.0, + ), + name="source", + ) + ) + scene.add_monitor(mw.PointMonitor("probe", source_position, fields=("Ez",))) + return scene + + +def _workload_scene( + workload: str, + shape: tuple[int, int, int], + *, + spacing: float, + device: torch.device, +) -> mw.Scene: + spec = _workload_spec(workload) + if spec.name == "vacuum": + return _vacuum_scene(shape, spacing=spacing, device=device) + + x, y, z = (_axis_nodes(count, spacing) for count in shape) + source_x_index = max(1, len(x) // 3) + probe_x_index = min(len(x) - 2, 2 * len(x) // 3) + source_position = (float(x[source_x_index]), 0.0, 0.0) + probe_position = (float(x[probe_x_index]), 0.0, 0.0) + boundary = ( + mw.BoundarySpec.pml(num_layers=4, strength=1.0e6) + if spec.cpml_mode != "disabled" + else mw.BoundarySpec.none() + ) + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=boundary, + device=str(device), + ) + if spec.dielectric: + spans = tuple(float(axis[-1] - axis[0]) for axis in (x, y, z)) + scene.add_structure( + mw.Structure( + name="dielectric_scatterer", + geometry=mw.Box( + position=(0.1 * spans[0], 0.0, 0.0), + size=tuple(0.25 * span for span in spans), + ), + material=mw.Material(eps_r=4.0), + ) + ) + source_time = ( + mw.GaussianPulse(frequency=_FREQUENCY, fwidth=0.5e9, amplitude=1.0) + if spec.broadband_source + else mw.CW(frequency=_FREQUENCY, amplitude=1.0) + ) + scene.add_source( + mw.PointDipole( + position=source_position, + polarization="Ez", + profile="ideal", + source_time=source_time, + name="source", + ) + ) + scene.add_monitor(mw.PointMonitor("probe", probe_position, fields=("Ez",))) + return scene + + +def _simulation( + workload: str, + shape: tuple[int, int, int], + *, + spacing: float, + steps: int, + device: torch.device, + parallel: mw.FDTDParallelConfig | None = None, +): + spec = _workload_spec(workload) + cpml_options = ( + {"cpml_config": {"memory_mode": spec.cpml_mode}} + if spec.cpml_mode != "disabled" + else {} + ) + return mw.Simulation.fdtd( + _workload_scene(workload, shape, spacing=spacing, device=device), + frequencies=spec.frequencies, + run_time=mw.TimeConfig(time_steps=steps), + spectral_sampler=mw.SpectralSampler(window="none"), + full_field_dft=spec.full_field_dft, + cuda_graph=False, + parallel=parallel, + **cpml_options, + ) + + +def _release_cuda(devices: tuple[torch.device, ...]) -> None: + gc.collect() + for device in devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + +def _baseline_allocated(devices: tuple[torch.device, ...]) -> dict[str, int]: + baseline = {} + for device in devices: + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + baseline[str(device)] = int(torch.cuda.memory_allocated(device)) + return baseline + + +def _single_sample( + shape: tuple[int, int, int], + *, + workload: str = "vacuum", + spacing: float, + steps: int, + device: torch.device, +) -> dict[str, object]: + _release_cuda((device,)) + baseline = _baseline_allocated((device,))[str(device)] + result = _simulation( + workload, + shape, + spacing=spacing, + steps=steps, + device=device, + ).run() + torch.cuda.synchronize(device) + stats = result.stats() + sample = { + "elapsed_s": float(stats["elapsed_s"]), + "ms_per_step": float(stats["elapsed_s"]) * 1.0e3 / steps, + "peak_memory_bytes": max( + int(torch.cuda.max_memory_allocated(device)) - baseline, + 0, + ), + "local_node_shapes": {str(device): shape}, + } + for key in ("cpml_requested_memory_mode", "cpml_memory_mode"): + if key in stats: + sample[key] = stats[key] + del result + _release_cuda((device,)) + return sample + + +def _multi_sample( + shape: tuple[int, int, int], + *, + workload: str = "vacuum", + spacing: float, + steps: int, + devices: tuple[torch.device, torch.device], + overlap: bool, + gather_fields: bool, +) -> dict[str, object]: + _release_cuda(devices) + baseline = _baseline_allocated(devices) + parallel = mw.FDTDParallelConfig( + devices=devices, + transport="cuda_p2p", + overlap=overlap, + gather_fields=gather_fields, + result_device=devices[0], + ) + result = _simulation( + workload, + shape, + spacing=spacing, + steps=steps, + device=devices[0], + parallel=parallel, + ).run() + for device in devices: + torch.cuda.synchronize(device) + stats = result.stats() + parallel_stats = stats["parallel_stats"] + peaks = { + name: max(int(value) - baseline[name], 0) + for name, value in parallel_stats["peak_memory_bytes"].items() + } + sample = { + "elapsed_s": float(stats["elapsed_s"]), + "ms_per_step": float(stats["elapsed_s"]) * 1.0e3 / steps, + "peak_memory_bytes": peaks, + "halo_bytes_per_step": int(parallel_stats["halo_bytes_per_step"]), + "halo_bytes_total": int(parallel_stats["halo_bytes_total"]), + "overlap_active": bool(parallel_stats["overlap_active"]), + "transport": parallel_stats["transport"], + "partitions": parallel_stats["partitions"], + "topology": parallel_stats["topology"], + "local_node_shapes": { + str(shard.device): ( + int(shard.solver.Nx), + int(shard.solver.Ny), + int(shard.solver.Nz), + ) + for shard in result.solver.shards + }, + } + for key in ("cpml_requested_memory_mode", "cpml_memory_mode"): + if key in stats: + sample[key] = stats[key] + del result + _release_cuda(devices) + return sample + + +def _field_error( + actual: torch.Tensor, + reference: torch.Tensor, +) -> dict[str, float | bool]: + actual = actual.detach().to(device="cpu") + reference = reference.detach().to(device="cpu") + absolute = torch.abs(actual - reference) + max_abs = float(absolute.max().item()) + reference_scale = float(torch.abs(reference).max().item()) + significant = torch.abs(reference) >= max(1.0e-6, 1.0e-4 * reference_scale) + max_rel = 0.0 + if bool(significant.any()): + relative = absolute[significant] / torch.abs(reference[significant]) + max_rel = float(relative.max().item()) + return { + "max_abs": max_abs, + "max_rel": max_rel, + "actual_peak": float(torch.abs(actual).max().item()), + "reference_peak": reference_scale, + "finite": bool(torch.isfinite(actual).all() and torch.isfinite(reference).all()), + } + + +def _to_cpu_data(value): + if isinstance(value, torch.Tensor): + return value.detach().to(device="cpu") + if isinstance(value, Mapping): + return {key: _to_cpu_data(item) for key, item in value.items()} + if isinstance(value, list): + return [_to_cpu_data(item) for item in value] + if isinstance(value, tuple): + return tuple(_to_cpu_data(item) for item in value) + return value + + +def _numeric_leaves(value, path: tuple[str, ...] = ()) -> dict[str, torch.Tensor]: + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return {} + return {".".join(path): value} + if isinstance(value, np.ndarray): + if value.size == 0 or not np.issubdtype(value.dtype, np.number): + return {} + return {".".join(path): torch.as_tensor(value)} + if isinstance(value, Mapping): + leaves = {} + for key, item in value.items(): + leaves.update(_numeric_leaves(item, (*path, str(key)))) + return leaves + if isinstance(value, (list, tuple)): + leaves = {} + for index, item in enumerate(value): + leaves.update(_numeric_leaves(item, (*path, str(index)))) + return leaves + return {} + + +def _monitor_error( + actual: Mapping[str, object], + reference: Mapping[str, object], +) -> dict[str, object]: + actual_names = tuple(actual) + reference_names = tuple(reference) + missing_monitors = tuple(name for name in reference_names if name not in actual) + unexpected_monitors = tuple(name for name in actual_names if name not in reference) + per_monitor = {} + all_errors = [] + finite = not missing_monitors and not unexpected_monitors + for name in reference_names: + if name not in actual: + continue + actual_leaves = _numeric_leaves(actual[name]) + reference_leaves = _numeric_leaves(reference[name]) + missing_paths = tuple(path for path in reference_leaves if path not in actual_leaves) + unexpected_paths = tuple(path for path in actual_leaves if path not in reference_leaves) + common_paths = tuple(path for path in reference_leaves if path in actual_leaves) + path_errors = {} + shapes_match = True + for path in common_paths: + actual_value = actual_leaves[path] + reference_value = reference_leaves[path] + if actual_value.shape != reference_value.shape: + shapes_match = False + path_errors[path] = { + "actual_shape": tuple(actual_value.shape), + "reference_shape": tuple(reference_value.shape), + "shape_match": False, + } + continue + row = _field_error(actual_value, reference_value) + row["shape_match"] = True + path_errors[path] = row + all_errors.append(row) + monitor_finite = ( + bool(common_paths) + and shapes_match + and not missing_paths + and not unexpected_paths + and all(bool(row.get("finite", False)) for row in path_errors.values()) + ) + finite = finite and monitor_finite + per_monitor[name] = { + "paths": path_errors, + "missing_numeric_paths": missing_paths, + "unexpected_numeric_paths": unexpected_paths, + "finite": monitor_finite, + } + + max_abs = max((float(row["max_abs"]) for row in all_errors), default=0.0) + max_rel = max((float(row["max_rel"]) for row in all_errors), default=0.0) + return { + "per_monitor": per_monitor, + "missing_monitors": missing_monitors, + "unexpected_monitors": unexpected_monitors, + "max_abs": max_abs, + "max_rel": max_rel, + "finite": finite, + "abs_tolerance": _MAX_FIELD_ABS, + "rel_tolerance": _MAX_FIELD_REL, + "passed": finite and max_abs <= _MAX_FIELD_ABS and max_rel <= _MAX_FIELD_REL, + } + + +def _parity_diagnostics( + shape: tuple[int, int, int], + *, + workload: str = "vacuum", + spacing: float, + steps: int, + devices: tuple[torch.device, torch.device], +) -> dict[str, object]: + """Run a cheap public-API parity check and report real six-field errors.""" + + _release_cuda(devices) + single = _simulation( + workload, + shape, + spacing=spacing, + steps=steps, + device=devices[0], + ).run() + single_raw = single.raw_output if isinstance(single.raw_output, dict) else {} + references = { + name: single_raw.get(name, getattr(single.solver, name)).detach().to(device="cpu") + for name in _FIELD_NAMES + } + reference_monitors = _to_cpu_data(single.monitors) + del single + _release_cuda(devices) + + parallel = mw.FDTDParallelConfig( + devices=devices, + transport="cuda_p2p", + overlap=True, + gather_fields=True, + result_device=devices[0], + ) + distributed = _simulation( + workload, + shape, + spacing=spacing, + steps=steps, + device=devices[0], + parallel=parallel, + ).run() + solver = distributed.solver + per_component = {} + for name in _FIELD_NAMES: + if name in distributed.raw_output: + actual = distributed.raw_output[name] + else: + local_values = tuple(getattr(shard.solver, name) for shard in solver.shards) + actual = solver._gather_component(name, local_values) + per_component[name] = _field_error(actual, references[name]) + + max_abs = max(float(row["max_abs"]) for row in per_component.values()) + max_rel = max(float(row["max_rel"]) for row in per_component.values()) + finite = all(bool(row["finite"]) for row in per_component.values()) + reference_peak = max(float(row["reference_peak"]) for row in per_component.values()) + monitor_error = _monitor_error(distributed.monitors, reference_monitors) + fields_passed = ( + finite + and reference_peak > 0.0 + and max_abs <= _MAX_FIELD_ABS + and max_rel <= _MAX_FIELD_REL + ) + report = { + "workload": workload, + "node_shape": shape, + "steps": steps, + "per_component": per_component, + "max_abs": max_abs, + "max_rel": max_rel, + "field_max_abs": max_abs, + "field_max_rel": max_rel, + "monitor_max_abs": float(monitor_error["max_abs"]), + "monitor_max_rel": float(monitor_error["max_rel"]), + "monitor_error": monitor_error, + "reference_peak": reference_peak, + "finite": finite, + "abs_tolerance": _MAX_FIELD_ABS, + "rel_tolerance": _MAX_FIELD_REL, + "fields_passed": fields_passed, + "monitors_passed": bool(monitor_error["passed"]), + "passed": fields_passed and bool(monitor_error["passed"]), + } + del distributed, solver, references, reference_monitors + _release_cuda(devices) + return report + + +def _run_samples( + function: Callable[[], dict[str, object]], + *, + warmups: int, + repeats: int, +) -> list[dict[str, object]]: + for _ in range(warmups): + function() + return [function() for _ in range(repeats)] + + +def _median_summary(samples: list[dict[str, object]]) -> dict[str, object]: + elapsed = [float(sample["elapsed_s"]) for sample in samples] + ms_per_step = [float(sample["ms_per_step"]) for sample in samples] + summary = { + "elapsed_s_median": statistics.median(elapsed), + "elapsed_s_runs": elapsed, + "ms_per_step_median": statistics.median(ms_per_step), + "ms_per_step_runs": ms_per_step, + } + peak = samples[0]["peak_memory_bytes"] + if isinstance(peak, dict): + summary["peak_memory_bytes_median"] = { + device: int( + statistics.median( + [int(sample["peak_memory_bytes"][device]) for sample in samples] + ) + ) + for device in peak + } + else: + summary["peak_memory_bytes_median"] = int( + statistics.median([int(sample["peak_memory_bytes"]) for sample in samples]) + ) + for key in ( + "halo_bytes_per_step", + "halo_bytes_total", + "overlap_active", + "transport", + "partitions", + "topology", + "local_node_shapes", + "cpml_requested_memory_mode", + "cpml_memory_mode", + ): + if key in samples[-1]: + summary[key] = samples[-1][key] + return summary + + +def _p2p_bandwidth( + devices: tuple[torch.device, torch.device], + *, + size_bytes: int, + trials: int = 5, +) -> dict[str, object]: + count = max(1, int(math.ceil(size_bytes / 4))) + actual_bytes = count * 4 + repeats = max(4, int(math.ceil((1024**3) / actual_bytes))) + rows = {} + for source_device, destination_device in (devices, tuple(reversed(devices))): + with torch.cuda.device(source_device): + source = torch.empty(count, device=source_device, dtype=torch.float32).normal_() + with torch.cuda.device(destination_device): + destination = torch.empty(count, device=destination_device, dtype=torch.float32) + stream = torch.cuda.Stream(device=destination_device) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(stream): + for _ in range(4): + destination.copy_(source, non_blocking=True) + stream.synchronize() + bandwidths = [] + latency_us = [] + for _ in range(trials): + with torch.cuda.stream(stream): + start.record(stream) + for _ in range(repeats): + destination.copy_(source, non_blocking=True) + end.record(stream) + end.synchronize() + elapsed_ms = float(start.elapsed_time(end)) / repeats + latency_us.append(elapsed_ms * 1.0e3) + bandwidths.append(actual_bytes / (elapsed_ms / 1.0e3) / 1.0e9) + label = f"{source_device}->{destination_device}" + rows[label] = { + "size_bytes": actual_bytes, + "repeats": repeats, + "bandwidth_gbps_median": statistics.median(bandwidths), + "bandwidth_gbps_runs": bandwidths, + "latency_us_median": statistics.median(latency_us), + } + del source, destination + _release_cuda(devices) + return rows + + +def _command_output(command: list[str]) -> str | None: + try: + return subprocess.check_output( + command, + text=True, + stderr=subprocess.STDOUT, + timeout=10, + ).strip() + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired): + return None + + +def _hardware(devices: tuple[torch.device, torch.device]) -> dict[str, object]: + return { + "platform": platform.platform(), + "python_version": sys.version, + "nvidia_smi_query": _command_output( + [ + "nvidia-smi", + "--query-gpu=index,name,driver_version,pstate," + "clocks.current.sm,clocks.current.memory", + "--format=csv,noheader", + ] + ), + "nvidia_smi_topology": _command_output(["nvidia-smi", "topo", "-m"]), + "torch_version": torch.__version__, + "torch_cuda_version": torch.version.cuda, + "devices": tuple( + { + "device": str(device), + "name": torch.cuda.get_device_name(device), + "compute_capability": torch.cuda.get_device_capability(device), + "total_memory_bytes": torch.cuda.get_device_properties(device).total_memory, + } + for device in devices + ), + "peer_access": { + f"{devices[0]}->{devices[1]}": torch.cuda.can_device_access_peer( + devices[0].index, devices[1].index + ), + f"{devices[1]}->{devices[0]}": torch.cuda.can_device_access_peer( + devices[1].index, devices[0].index + ), + }, + } + + +def _argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--workload", choices=tuple(_WORKLOADS), default="vacuum") + parser.add_argument("--devices", nargs=2, default=("cuda:0", "cuda:1")) + parser.add_argument("--nodes-x", type=int, default=257) + parser.add_argument("--nodes-y", type=int, default=129) + parser.add_argument("--nodes-z", type=int, default=129) + parser.add_argument("--spacing", type=float, default=0.01) + parser.add_argument("--steps", type=int, default=200) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--no-overlap", action="store_true") + parser.add_argument("--gather-fields", action="store_true") + parser.add_argument("--weak-scaling", action="store_true") + parser.add_argument("--p2p-size-mib", type=float, default=256.0) + parser.add_argument("--fixed-overhead-mib", type=float, default=0.0) + parser.add_argument("--diagnostic-nodes-x", type=int, default=33) + parser.add_argument("--diagnostic-nodes-y", type=int, default=17) + parser.add_argument("--diagnostic-nodes-z", type=int, default=17) + parser.add_argument("--diagnostic-steps", type=int, default=24) + parser.add_argument("--skip-diagnostics", action="store_true") + parser.add_argument("--min-p2p-gbps", type=float, default=40.0) + parser.add_argument("--min-strong-speedup", type=float, default=1.0) + parser.add_argument("--min-weak-efficiency", type=float, default=0.70) + parser.add_argument("--max-memory-factor", type=float, default=1.15) + parser.add_argument("--assert-gates", action="store_true") + parser.add_argument("--json", type=Path) + return parser + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + return _argument_parser().parse_args(argv) + + +def main(argv: list[str] | None = None) -> dict[str, object]: + args = _parse_args(argv) + + if args.steps <= 0 or args.diagnostic_steps <= 0 or args.warmups < 0 or args.repeats <= 0: + raise SystemExit("steps/repeats must be positive and warmups must be nonnegative.") + devices = _require_devices(tuple(args.devices)) + shape = (args.nodes_x, args.nodes_y, args.nodes_z) + workload = _workload_spec(args.workload) + p2p = _p2p_bandwidth( + devices, + size_bytes=max(4, int(args.p2p_size_mib * _MIB)), + ) + diagnostics = None + if not args.skip_diagnostics: + diagnostics = _parity_diagnostics( + ( + args.diagnostic_nodes_x, + args.diagnostic_nodes_y, + args.diagnostic_nodes_z, + ), + workload=workload.name, + spacing=args.spacing, + steps=args.diagnostic_steps, + devices=devices, + ) + + single_samples = _run_samples( + lambda: _single_sample( + shape, + workload=workload.name, + spacing=args.spacing, + steps=args.steps, + device=devices[0], + ), + warmups=args.warmups, + repeats=args.repeats, + ) + multi_samples = _run_samples( + lambda: _multi_sample( + shape, + workload=workload.name, + spacing=args.spacing, + steps=args.steps, + devices=devices, + overlap=not args.no_overlap, + gather_fields=args.gather_fields, + ), + warmups=args.warmups, + repeats=args.repeats, + ) + single = _median_summary(single_samples) + multi = _median_summary(multi_samples) + strong_speedup = float(single["ms_per_step_median"]) / float(multi["ms_per_step_median"]) + + single_peak = int(single["peak_memory_bytes_median"]) + multi_peaks = {key: int(value) for key, value in multi["peak_memory_bytes_median"].items()} + fixed_overhead = int(args.fixed_overhead_mib * _MIB) + memory_limit = args.max_memory_factor * single_peak / 2.0 + fixed_overhead + memory_factor_to_ideal = max(multi_peaks.values()) / max(single_peak / 2.0, 1.0) + + weak = None + if args.weak_scaling: + local_x_cells = int(math.ceil((args.nodes_x - 1) / 2)) + weak_shape = (local_x_cells + 1, args.nodes_y, args.nodes_z) + weak_samples = _run_samples( + lambda: _single_sample( + weak_shape, + workload=workload.name, + spacing=args.spacing, + steps=args.steps, + device=devices[0], + ), + warmups=args.warmups, + repeats=args.repeats, + ) + weak_single = _median_summary(weak_samples) + weak = { + "single_local_shape": weak_shape, + "single": weak_single, + "efficiency": float(weak_single["ms_per_step_median"]) + / float(multi["ms_per_step_median"]), + } + + local_shapes = multi.get("local_node_shapes", {}) + cpml_actual = ( + "disabled" + if workload.cpml_mode == "disabled" + else multi.get("cpml_memory_mode", workload.cpml_mode) + ) + numerical_error = { + "diagnostics_run": diagnostics is not None, + "field_max_abs": ( + float(diagnostics["field_max_abs"]) if diagnostics is not None else None + ), + "field_max_rel": ( + float(diagnostics["field_max_rel"]) if diagnostics is not None else None + ), + "monitor_max_abs": ( + float(diagnostics["monitor_max_abs"]) if diagnostics is not None else None + ), + "monitor_max_rel": ( + float(diagnostics["monitor_max_rel"]) if diagnostics is not None else None + ), + } + report = { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "workload": workload.name, + "hardware": _hardware(devices), + "shapes": { + "global_node_shape": shape, + "local_node_shapes": local_shapes, + }, + "execution": { + "steps": args.steps, + "frequencies": workload.frequencies, + "full_field_dft": workload.full_field_dft, + "graph_mode": workload.graph_mode, + "cpml_mode_requested": workload.cpml_mode, + "cpml_mode_actual": cpml_actual, + "gather_fields": args.gather_fields, + "overlap_requested": not args.no_overlap, + }, + "communication": { + "transport": multi.get("transport"), + "halo_bytes_per_step": multi.get("halo_bytes_per_step"), + "halo_bytes_total": multi.get("halo_bytes_total"), + "topology": multi.get("topology"), + }, + "numerical_error": numerical_error, + "protocol": { + "workload": workload.name, + "global_node_shape": shape, + "local_node_shapes": local_shapes, + "steps": args.steps, + "warmups": args.warmups, + "repeats": args.repeats, + "spacing": args.spacing, + "overlap_requested": not args.no_overlap, + "gather_fields": args.gather_fields, + "frequencies": workload.frequencies, + "full_field_dft": workload.full_field_dft, + "graph_mode": workload.graph_mode, + "cpml_mode": cpml_actual, + }, + "p2p": p2p, + "numerical_parity": diagnostics, + "single_gpu": single, + "two_gpu": multi, + "strong_speedup": strong_speedup, + "memory": { + "single_peak_bytes": single_peak, + "two_gpu_peak_bytes": multi_peaks, + "fixed_overhead_bytes": fixed_overhead, + "allowed_peak_per_gpu_bytes": memory_limit, + "factor_to_ideal_half_single": memory_factor_to_ideal, + }, + "weak_scaling": weak, + } + + failures = [] + if diagnostics is not None and not bool(diagnostics["passed"]): + failures.append( + "field/monitor parity " + f"field_max_abs={float(diagnostics['field_max_abs']):.3e}, " + f"field_max_rel={float(diagnostics['field_max_rel']):.3e}, " + f"monitor_max_abs={float(diagnostics['monitor_max_abs']):.3e}, " + f"monitor_max_rel={float(diagnostics['monitor_max_rel']):.3e}" + ) + if min(float(row["bandwidth_gbps_median"]) for row in p2p.values()) < args.min_p2p_gbps: + failures.append(f"P2P bandwidth < {args.min_p2p_gbps:.1f} GB/s") + if strong_speedup < args.min_strong_speedup: + failures.append(f"strong speedup {strong_speedup:.3f} < {args.min_strong_speedup:.3f}") + if max(multi_peaks.values()) > memory_limit: + failures.append( + f"per-GPU peak {max(multi_peaks.values())} > memory allowance {memory_limit:.0f}" + ) + if weak is not None and float(weak["efficiency"]) < args.min_weak_efficiency: + failures.append( + f"weak efficiency {float(weak['efficiency']):.3f} < {args.min_weak_efficiency:.3f}" + ) + report["gates"] = {"passed": not failures, "failures": failures} + + rendered = json.dumps(report, indent=2, sort_keys=True) + print(rendered) + if args.json is not None: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(rendered + "\n", encoding="utf-8") + if args.assert_gates and failures: + raise SystemExit("; ".join(failures)) + return report + + +if __name__ == "__main__": + main() diff --git a/tests/fdtd/cuda/test_cuda_prebuilt_abi.py b/tests/fdtd/cuda/test_cuda_prebuilt_abi.py new file mode 100644 index 0000000..1a737fb --- /dev/null +++ b/tests/fdtd/cuda/test_cuda_prebuilt_abi.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import pytest + +from witwin.maxwell.fdtd.cuda import build as cuda_build + + +@pytest.mark.parametrize("marker_contents", [None, "stable_abi_v1\n", "garbage\n"]) +def test_stale_or_missing_packaged_prebuilt_is_rejected_before_loading( + monkeypatch, + tmp_path, + marker_contents, +): + library = tmp_path / f"witwin_maxwell_fdtd_cuda{cuda_build.extension_suffix()}" + library.write_bytes(b"not-loaded") + marker = cuda_build.prebuilt_abi_marker_path(library) + if marker_contents is not None: + marker.write_text(marker_contents, encoding="utf-8") + + load_calls = [] + monkeypatch.setattr(cuda_build, "prebuilt_extension_path", lambda: library) + monkeypatch.setattr( + cuda_build, + "_load_extension_file", + lambda path: load_calls.append(path), + ) + + assert cuda_build._load_packaged_prebuilt_extension() is None + assert load_calls == [] + + +def test_matching_packaged_prebuilt_load_failure_is_fail_fast_without_jit( + monkeypatch, + tmp_path, +): + library = tmp_path / f"witwin_maxwell_fdtd_cuda{cuda_build.extension_suffix()}" + library.write_bytes(b"invalid-v2-binary") + cuda_build.prebuilt_abi_marker_path(library).write_text( + f"{cuda_build.STABLE_ABI_VERSION}\n", + encoding="utf-8", + ) + monkeypatch.setattr(cuda_build, "prebuilt_extension_path", lambda: library) + monkeypatch.delenv("WITWIN_MAXWELL_FDTD_CUDA_SKIP_PREBUILT", raising=False) + monkeypatch.delenv("WITWIN_MAXWELL_FDTD_CUDA_PREBUILT", raising=False) + + jit_calls = [] + + def fail_load(path): + raise RuntimeError(f"failed to load {path}") + + monkeypatch.setattr(cuda_build, "_load_extension_file", fail_load) + monkeypatch.setattr(cuda_build, "load", lambda **kwargs: jit_calls.append(kwargs)) + + with pytest.raises(RuntimeError, match="failed to load"): + cuda_build.build_extension() + + assert jit_calls == [] diff --git a/tests/fdtd/multi_gpu/conftest.py b/tests/fdtd/multi_gpu/conftest.py new file mode 100644 index 0000000..f34df58 --- /dev/null +++ b/tests/fdtd/multi_gpu/conftest.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import pytest +import torch + + +@pytest.fixture(scope="session") +def cuda_p2p_devices() -> tuple[torch.device, torch.device]: + """Return two mutually peer-accessible CUDA devices or skip the test session.""" + + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("requires at least two CUDA devices") + + devices = (torch.device("cuda:0"), torch.device("cuda:1")) + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("cuda:0 cannot directly access cuda:1") + if not torch.cuda.can_device_access_peer(1, 0): + pytest.skip("cuda:1 cannot directly access cuda:0") + + properties = tuple(torch.cuda.get_device_properties(device) for device in devices) + signatures = {(item.name, item.major, item.minor) for item in properties} + if len(signatures) != 1: + pytest.skip("multi-GPU FDTD acceptance requires homogeneous CUDA devices") + return devices + + +@pytest.fixture +def cuda_memory_cleanup(cuda_p2p_devices): + """Synchronize and release allocator caches around an isolated GPU test.""" + + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + yield + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() diff --git a/tests/fdtd/multi_gpu/test_bench_joint_cli.py b/tests/fdtd/multi_gpu/test_bench_joint_cli.py new file mode 100644 index 0000000..884d802 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_bench_joint_cli.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + + +@pytest.fixture(scope="module") +def bench_joint(): + root = Path(__file__).resolve().parents[3] + path = root / "scripts" / "dev" / "fdtd" / "multi_gpu" / "bench_joint.py" + module_name = "_witwin_test_bench_joint" + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_existing_cli_defaults_remain_vacuum_compatible(bench_joint): + args = bench_joint._parse_args([]) + + assert args.workload == "vacuum" + assert tuple(args.devices) == ("cuda:0", "cuda:1") + assert (args.nodes_x, args.nodes_y, args.nodes_z) == (257, 129, 129) + assert args.steps == 200 + assert args.warmups == 1 + assert args.repeats == 5 + assert args.gather_fields is False + assert args.weak_scaling is False + + +def test_monitor_error_reports_maximum_numeric_payload_difference(bench_joint): + reference = { + "probe": { + "data": torch.tensor((1.0, 2.0)), + "coords": (np.asarray((0.0, 0.1)),), + } + } + actual = { + "probe": { + "data": torch.tensor((1.0, 2.25)), + "coords": (np.asarray((0.0, 0.1)),), + } + } + + error = bench_joint._monitor_error(actual, reference) + + assert error["max_abs"] == pytest.approx(0.25) + assert error["max_rel"] == pytest.approx(0.125) + assert error["finite"] is True + assert error["passed"] is False + + +@pytest.mark.parametrize( + ("name", "frequencies", "full_field_dft", "cpml_mode", "structure_count"), + ( + ("vacuum", (1.0e9,), False, "disabled", 0), + ("cpml_dielectric", (1.0e9,), False, "slab", 1), + ("multifrequency_dft", (0.8e9, 1.2e9), True, "disabled", 0), + ), +) +def test_workload_specs_build_expected_declarative_scene( + bench_joint, + name, + frequencies, + full_field_dft, + cpml_mode, + structure_count, +): + spec = bench_joint._workload_spec(name) + scene = bench_joint._workload_scene( + name, + (17, 9, 9), + spacing=0.01, + device=torch.device("cpu"), + ) + + assert spec.frequencies == frequencies + assert spec.full_field_dft is full_field_dft + assert spec.cpml_mode == cpml_mode + assert len(scene.structures) == structure_count + assert len(scene.sources) == 1 + assert tuple(monitor.name for monitor in scene.monitors) == ("probe",) + assert scene.boundary.uses_kind("pml") is (cpml_mode != "disabled") + + +@pytest.mark.parametrize( + ("workload", "expected_cpml_config"), + ( + ("vacuum", None), + ("cpml_dielectric", {"memory_mode": "slab"}), + ("multifrequency_dft", None), + ), +) +def test_simulation_routes_cpml_memory_mode_only_when_enabled( + bench_joint, + monkeypatch, + workload, + expected_cpml_config, +): + calls = [] + monkeypatch.setattr( + bench_joint.mw.Simulation, + "fdtd", + lambda scene, **kwargs: calls.append((scene, kwargs)) or object(), + ) + + bench_joint._simulation( + workload, + (17, 9, 9), + spacing=0.01, + steps=5, + device=torch.device("cpu"), + ) + + assert len(calls) == 1 + kwargs = calls[0][1] + if expected_cpml_config is None: + assert "cpml_config" not in kwargs + else: + assert kwargs["cpml_config"] == expected_cpml_config + + +@pytest.mark.parametrize( + ("workload", "expected_cpml", "expected_dft"), + ( + ("vacuum", "disabled", False), + ("cpml_dielectric", "slab", False), + ("multifrequency_dft", "disabled", True), + ), +) +def test_report_schema_records_workload_shapes_runtime_and_errors( + bench_joint, + monkeypatch, + tmp_path, + capsys, + workload, + expected_cpml, + expected_dft, +): + devices = (torch.device("cuda:0"), torch.device("cuda:1")) + calls = [] + topology = { + "kind": "cuda_p2p", + "neighbor_pairs": [{"left": "cuda:0", "right": "cuda:1"}], + } + local_shapes = {"cuda:0": (6, 7, 7), "cuda:1": (7, 7, 7)} + + monkeypatch.setattr(bench_joint, "_require_devices", lambda _values: devices) + monkeypatch.setattr( + bench_joint, + "_p2p_bandwidth", + lambda *_args, **_kwargs: { + "cuda:0->cuda:1": {"bandwidth_gbps_median": 100.0}, + "cuda:1->cuda:0": {"bandwidth_gbps_median": 100.0}, + }, + ) + monkeypatch.setattr( + bench_joint, + "_hardware", + lambda _devices: { + "torch_version": "2.test", + "torch_cuda_version": "12.test", + "devices": ({"device": "cuda:0"}, {"device": "cuda:1"}), + }, + ) + + def fake_diagnostics(shape, *, workload, spacing, steps, devices): + calls.append(("diagnostics", workload, shape, steps)) + return { + "passed": True, + "field_max_abs": 1.0e-7, + "field_max_rel": 2.0e-7, + "monitor_max_abs": 3.0e-7, + "monitor_max_rel": 4.0e-7, + } + + def fake_single(shape, *, workload, spacing, steps, device): + calls.append(("single", workload, shape, steps)) + return { + "elapsed_s": 2.0, + "ms_per_step": 200.0, + "peak_memory_bytes": 2000, + "local_node_shapes": {str(device): shape}, + "cpml_memory_mode": expected_cpml, + } + + def fake_multi( + shape, + *, + workload, + spacing, + steps, + devices, + overlap, + gather_fields, + ): + calls.append(("multi", workload, shape, steps)) + return { + "elapsed_s": 1.0, + "ms_per_step": 100.0, + "peak_memory_bytes": {"cuda:0": 800, "cuda:1": 850}, + "halo_bytes_per_step": 128, + "halo_bytes_total": 1280, + "overlap_active": overlap, + "transport": "cuda_p2p", + "partitions": ( + {"rank": 0, "device": "cuda:0"}, + {"rank": 1, "device": "cuda:1"}, + ), + "topology": topology, + "local_node_shapes": local_shapes, + "cpml_memory_mode": expected_cpml, + } + + monkeypatch.setattr(bench_joint, "_parity_diagnostics", fake_diagnostics) + monkeypatch.setattr(bench_joint, "_single_sample", fake_single) + monkeypatch.setattr(bench_joint, "_multi_sample", fake_multi) + + output_path = tmp_path / f"{workload}.json" + returned = bench_joint.main( + [ + "--workload", + workload, + "--nodes-x", + "11", + "--nodes-y", + "7", + "--nodes-z", + "7", + "--steps", + "10", + "--warmups", + "0", + "--repeats", + "1", + "--json", + str(output_path), + ] + ) + rendered = json.loads(output_path.read_text(encoding="utf-8")) + stdout = json.loads(capsys.readouterr().out) + + assert returned["workload"] == workload + assert stdout == rendered + assert rendered["workload"] == workload + assert rendered["shapes"] == { + "global_node_shape": [11, 7, 7], + "local_node_shapes": { + "cuda:0": [6, 7, 7], + "cuda:1": [7, 7, 7], + }, + } + assert rendered["execution"]["steps"] == 10 + assert rendered["execution"]["full_field_dft"] is expected_dft + assert rendered["execution"]["graph_mode"] == "disabled" + assert rendered["execution"]["cpml_mode_actual"] == expected_cpml + assert rendered["communication"]["halo_bytes_per_step"] == 128 + assert rendered["communication"]["halo_bytes_total"] == 1280 + assert rendered["communication"]["topology"] == topology + assert rendered["hardware"]["torch_version"] == "2.test" + assert rendered["hardware"]["torch_cuda_version"] == "12.test" + assert rendered["numerical_error"] == { + "diagnostics_run": True, + "field_max_abs": 1.0e-7, + "field_max_rel": 2.0e-7, + "monitor_max_abs": 3.0e-7, + "monitor_max_rel": 4.0e-7, + } + assert rendered["memory"]["two_gpu_peak_bytes"] == { + "cuda:0": 800, + "cuda:1": 850, + } + assert {row[0] for row in calls} == {"diagnostics", "single", "multi"} + assert {row[1] for row in calls} == {workload} diff --git a/tests/fdtd/multi_gpu/test_bounded_standard_dispatch.py b/tests/fdtd/multi_gpu/test_bounded_standard_dispatch.py new file mode 100644 index 0000000..2890137 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_bounded_standard_dispatch.py @@ -0,0 +1,242 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +from witwin.maxwell.fdtd.cuda import backend +from witwin.maxwell.fdtd.distributed import solver as distributed_solver +from witwin.maxwell.fdtd_parallel import FDTDPartitionPlan + + +class _RecordedLaunch: + def __init__(self, calls, name: str, kwargs: dict) -> None: + self._calls = calls + self._name = name + self._kwargs = kwargs + + def launchRaw(self): # noqa: N802 - mirrors the CUDA backend launch API + self._calls.append((self._name, self._kwargs)) + + +class _RecordingModule: + def __init__(self) -> None: + self.calls: list[tuple[str, dict]] = [] + + def __getattr__(self, name: str): + def bind(**kwargs): + return _RecordedLaunch(self.calls, name, kwargs) + + return bind + + +def _sentinel_solver(module: _RecordingModule): + names = ( + "Ex", "Ey", "Ez", "Hx", "Hy", "Hz", + "chx_decay", "chx_curl", "chy_decay", "chy_curl", "chz_decay", "chz_curl", + "inv_dx_e", "inv_dy_e", "inv_dz_e", "inv_dx_h", "inv_dy_h", "inv_dz_h", + ) + return SimpleNamespace( + fdtd_module=module, + boundary_y_low_code=11, + boundary_y_high_code=12, + boundary_z_low_code=13, + boundary_z_high_code=14, + **{name: object() for name in names}, + ) + + +def test_bounded_helpers_pass_full_allocations_and_component_coordinates(monkeypatch): + plan = FDTDPartitionPlan( + global_shape=(13, 7, 6), + devices=("cuda:0", "cuda:1", "cuda:2"), + ) + layout = plan.layout(1) + module = _RecordingModule() + solver = _sentinel_solver(module) + electric_decay = (object(), object(), object()) + electric_curl = (object(), object(), object()) + monkeypatch.setattr( + distributed_solver, + "_electric_coefficients", + lambda _solver: (electric_decay, electric_curl), + ) + + node_owned = layout.storage_node_owned + cell_owned = layout.storage_cell_owned + distributed_solver._launch_magnetic_hx(solver, layout, node_owned) + distributed_solver._launch_magnetic_hy_hz(solver, layout, cell_owned) + distributed_solver._launch_electric_ex(solver, layout, cell_owned) + distributed_solver._launch_electric_ey_ez( + solver, + layout, + node_owned, + x_low_mode=21, + x_high_mode=22, + ) + + expected = ( + ("updateMagneticFieldHxStandardBounded3D", "Hx", node_owned), + ("updateMagneticFieldHyStandardBounded3D", "Hy", cell_owned), + ("updateMagneticFieldHzStandardBounded3D", "Hz", cell_owned), + ("updateElectricFieldExStandardBounded3D", "Ex", cell_owned), + ("updateElectricFieldEyStandardBounded3D", "Ey", node_owned), + ("updateElectricFieldEzStandardBounded3D", "Ez", node_owned), + ) + assert tuple(name for name, _ in module.calls) == tuple(item[0] for item in expected) + for (name, field_name, x_slice), (actual_name, kwargs) in zip(expected, module.calls): + assert actual_name == name + assert kwargs[field_name] is getattr(solver, field_name) + component_layout = layout.component(field_name) + assert kwargs["localXBegin"] == x_slice.start + assert kwargs["localXEnd"] == x_slice.stop + assert kwargs["globalXOffset"] == component_layout.global_origin[0] + assert kwargs["globalXExtent"] == component_layout.global_shape[0] + + assert module.calls[3][1]["ExDecay"] is electric_decay[0] + assert module.calls[4][1]["EyCurl"] is electric_curl[1] + assert module.calls[5][1]["EzCurl"] is electric_curl[2] + assert module.calls[4][1]["xLowBoundaryMode"] == 21 + assert module.calls[4][1]["xHighBoundaryMode"] == 22 + + before = len(module.calls) + distributed_solver._launch_magnetic_hx(solver, layout, slice(2, 2)) + distributed_solver._launch_electric_ex(solver, layout, slice(3, 3)) + assert len(module.calls) == before + + +def test_bounded_operators_have_schemas_and_backend_bindings(): + op_names = ( + "update_magnetic_hx_standard_bounded", + "update_magnetic_hy_standard_bounded", + "update_magnetic_hz_standard_bounded", + "update_electric_ex_standard_bounded", + "update_electric_ey_standard_bounded", + "update_electric_ez_standard_bounded", + ) + public_names = ( + "updateMagneticFieldHxStandardBounded3D", + "updateMagneticFieldHyStandardBounded3D", + "updateMagneticFieldHzStandardBounded3D", + "updateElectricFieldExStandardBounded3D", + "updateElectricFieldEyStandardBounded3D", + "updateElectricFieldEzStandardBounded3D", + ) + source = Path(backend.__file__).with_name("extension.cpp").read_text(encoding="utf-8") + for name in op_names: + schema = source.split(f'_({name}, "', maxsplit=1)[1].split('"', maxsplit=1)[0] + assert "int local_x_begin, int local_x_end" in schema + assert "int global_x_offset, int global_x_extent" in schema + assert all(name in backend._KERNELS for name in public_names) + + +class _RecordingStream: + def __init__(self, rank: int, events: list[tuple]) -> None: + self.rank = rank + self.events = events + + def wait_event(self, event) -> None: + self.events.append(("wait", self.rank, event)) + + +class _RecordingTransport: + def __init__(self, events: list[tuple]) -> None: + self.events = events + + def exchange_electric(self, shards) -> None: + self.events.append(("exchange_electric",)) + + def exchange_magnetic(self, shards) -> None: + self.events.append(("exchange_magnetic",)) + + +def _assert_exact_owned_partition(boxes: list[slice], owned: slice) -> None: + assert boxes + assert boxes[0].start == owned.start + assert boxes[-1].stop == owned.stop + assert all(box.start >= owned.start and box.stop <= owned.stop for box in boxes) + assert all(left.stop == right.start for left, right in zip(boxes, boxes[1:])) + + +def test_overlap_schedules_wait_between_disjoint_owned_launch_boxes(monkeypatch): + from contextlib import nullcontext + + plan = FDTDPartitionPlan(global_shape=(13, 7, 6), devices=("cuda:0", "cuda:1")) + events: list[tuple] = [] + shards = [] + for rank, layout in enumerate(plan.shard_layouts): + solver = SimpleNamespace(boundary_x_low_code=31, boundary_x_high_code=32) + shards.append( + SimpleNamespace( + rank=rank, + layout=layout, + solver=solver, + device=layout.device, + compute_stream=_RecordingStream(rank, events), + electric_received=f"electric-{rank}", + magnetic_received=f"magnetic-{rank}", + is_first=rank == 0, + is_last=rank == len(plan.shard_layouts) - 1, + ) + ) + coordinator = distributed_solver.DistributedFDTD.__new__( + distributed_solver.DistributedFDTD + ) + coordinator.shards = tuple(shards) + coordinator.transport = _RecordingTransport(events) + + monkeypatch.setattr(distributed_solver.torch.cuda, "device", lambda *_: nullcontext()) + monkeypatch.setattr(distributed_solver.torch.cuda, "stream", lambda *_: nullcontext()) + monkeypatch.setattr( + distributed_solver, + "_launch_magnetic_hx", + lambda solver, layout, box: events.append(("hx", layout.rank, box)), + ) + monkeypatch.setattr( + distributed_solver, + "_launch_magnetic_hy_hz", + lambda solver, layout, box: events.append(("hyhz", layout.rank, box)), + ) + monkeypatch.setattr( + distributed_solver, + "_launch_electric_ex", + lambda solver, layout, box: events.append(("ex", layout.rank, box)), + ) + monkeypatch.setattr( + distributed_solver, + "_launch_electric_ey_ez", + lambda solver, layout, box, **_: events.append(("eyez", layout.rank, box)), + ) + + coordinator._advance_magnetic_overlapped() + nonlast = plan.layout(0) + h_owned = nonlast.storage_cell_owned + h_boxes = [event[2] for event in events if event[:2] == ("hyhz", 0)] + assert h_boxes == [ + slice(h_owned.start, h_owned.stop - 1), + slice(h_owned.stop - 1, h_owned.stop), + ] + _assert_exact_owned_partition(h_boxes, h_owned) + h_interior = events.index(("hyhz", 0, h_boxes[0])) + h_wait = events.index(("wait", 0, "electric-0")) + h_boundary = events.index(("hyhz", 0, h_boxes[1])) + assert h_interior < h_wait < h_boundary + assert ("hx", 0, nonlast.storage_node_owned) in events + + events.clear() + coordinator._advance_electric_overlapped() + nonfirst = plan.layout(1) + e_owned = nonfirst.storage_node_owned + e_boxes = [event[2] for event in events if event[:2] == ("eyez", 1)] + assert e_boxes == [ + slice(e_owned.start + 1, e_owned.stop), + slice(e_owned.start, e_owned.start + 1), + ] + ordered_boxes = sorted(e_boxes, key=lambda box: box.start) + _assert_exact_owned_partition(ordered_boxes, e_owned) + e_interior = events.index(("eyez", 1, e_boxes[0])) + e_wait = events.index(("wait", 1, "magnetic-1")) + e_boundary = events.index(("eyez", 1, e_boxes[1])) + assert e_interior < e_wait < e_boundary + assert e_owned.start > 0 + assert all(box.start >= e_owned.start for box in e_boxes) + assert ("ex", 1, nonfirst.storage_cell_owned) in events diff --git a/tests/fdtd/multi_gpu/test_capacity_cpml_stats.py b/tests/fdtd/multi_gpu/test_capacity_cpml_stats.py new file mode 100644 index 0000000..7fb0103 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_capacity_cpml_stats.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed import capacity +from witwin.maxwell.fdtd.distributed.capacity import ( + local_dft_working_set_bytes, + require_gather_capacity, + require_local_dft_capacity, +) +from witwin.maxwell.fdtd.distributed.output import electric_field_output_bytes +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig + + +_FREQUENCY = 1.0e9 +_FIELD_NAMES = ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") + + +def test_local_dft_capacity_counts_accumulators_and_postprocessed_fields_exactly( + monkeypatch, +): + shape = (5, 4, 3) + frequencies = (0.8e9, 1.2e9) + complex_fields = electric_field_output_bytes( + shape, + frequency_count=len(frequencies), + complex_output=True, + ) + # One complex-field equivalent is the real+imag float32 accumulators; the + # other is the postprocessed complex64 local result. + expected = 2 * complex_fields + + assert expected == 4256 + assert local_dft_working_set_bytes( + shape, + dft_frequency=frequencies, + full_field_dft=True, + ) == expected + + monkeypatch.setattr(capacity, "effective_cuda_free_bytes", lambda _device: 1 << 40) + preflight = require_local_dft_capacity( + torch.device("cuda:0"), + shape, + dft_frequency=frequencies, + full_field_dft=True, + ) + assert preflight["required_bytes"] == expected + assert preflight["pending_local_dft_bytes"] == expected + + +def test_result_gather_capacity_oom_uses_combined_global_and_local_bytes(monkeypatch): + shape = (9, 7, 5) + output_bytes = electric_field_output_bytes( + shape, + frequency_count=2, + complex_output=True, + ) + pending_local = output_bytes + combined = output_bytes + pending_local + + # Enough for the global result alone, but not for result + resident local + # DFT accumulators/postprocess. The preflight must reject the combination. + monkeypatch.setattr( + capacity, + "effective_cuda_free_bytes", + lambda _device: int(output_bytes * 1.10), + ) + with pytest.raises(MemoryError, match=rf"requires {combined} bytes"): + require_gather_capacity( + torch.device("cuda:0"), + shape, + dft_frequency=(0.8e9, 1.2e9), + full_field_dft=True, + pending_local_dft_bytes=pending_local, + ) + + monkeypatch.setattr( + capacity, + "effective_cuda_free_bytes", + lambda _device: 2 * combined, + ) + preflight = require_gather_capacity( + torch.device("cuda:0"), + shape, + dft_frequency=(0.8e9, 1.2e9), + full_field_dft=True, + pending_local_dft_bytes=pending_local, + ) + assert preflight["output_bytes"] == output_bytes + assert preflight["pending_local_dft_bytes"] == pending_local + assert preflight["required_bytes"] == combined + + +def _cpml_scene() -> mw.Scene: + x = np.linspace(-0.6, 0.6, 13, dtype=np.float64) + y = np.linspace(-0.4, 0.4, 9, dtype=np.float64) + z = np.linspace(-0.4, 0.4, 9, dtype=np.float64) + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.pml(num_layers=2, strength=1.0e6), + device="cuda:0", + ) + scene.add_source( + mw.PointDipole( + position=(-0.2, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=mw.CW(frequency=_FREQUENCY, amplitude=1.0), + name="source", + ) + ) + scene.add_monitor( + mw.PointMonitor("probe", (0.0, 0.0, 0.0), fields=_FIELD_NAMES) + ) + return scene + + +def _run_cpml(*, parallel=None): + return mw.Simulation.fdtd( + _cpml_scene(), + frequencies=(_FREQUENCY,), + run_time=mw.TimeConfig(time_steps=24), + spectral_sampler=mw.SpectralSampler(window="none"), + absorber="cpml", + cpml_config={"memory_mode": "slab"}, + full_field_dft=False, + cuda_graph=False, + parallel=parallel, + ).run() + + +def _multi_component(result, name: str) -> torch.Tensor: + if name in result.raw_output: + return result.raw_output[name] + solver = result.solver + local = tuple(getattr(shard.solver, name) for shard in solver.shards) + return solver._gather_component(name, local) + + +@pytest.fixture(scope="module") +def slab_cpml_results(cuda_p2p_devices): + single = _run_cpml() + multi = _run_cpml( + parallel=FDTDParallelConfig( + devices=cuda_p2p_devices, + transport="cuda_p2p", + overlap=True, + gather_fields=True, + result_device=cuda_p2p_devices[0], + ) + ) + yield single, multi + del single, multi + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + +def test_multi_gpu_slab_cpml_preserves_config_and_six_field_parity(slab_cpml_results): + single, multi = slab_cpml_results + distributed = multi.solver + + assert distributed.cpml_config["memory_mode"] == "slab" + assert distributed._cpml_memory_mode_requested == "slab" + assert distributed._cpml_memory_mode == "slab" + assert all( + shard.solver._cpml_memory_mode_requested == "slab" + and shard.solver._cpml_memory_mode == "slab" + for shard in distributed.shards + ) + + for name in _FIELD_NAMES: + torch.testing.assert_close( + _multi_component(multi, name).to("cuda:0"), + getattr(single.solver, name).to("cuda:0"), + rtol=2.0e-5, + atol=2.0e-6, + ) + + +def test_public_solver_stats_aggregate_slab_cpml_memory(slab_cpml_results): + _single, multi = slab_cpml_results + solver = multi.solver + stats = multi.stats() + + allocated = sum( + int(shard.solver._cpml_allocated_memory_bytes) for shard in solver.shards + ) + dense = sum(int(shard.solver._cpml_dense_memory_bytes) for shard in solver.shards) + slab = sum(int(shard.solver._cpml_slab_memory_bytes) for shard in solver.shards) + + assert stats["cpml_requested_memory_mode"] == "slab" + assert stats["cpml_memory_mode"] == "slab" + assert stats["cpml_allocated_memory_bytes"] == allocated == slab + assert stats["cpml_dense_memory_bytes"] == dense + assert stats["cpml_slab_memory_bytes"] == slab + assert 0 < slab < dense + + +def test_parallel_stats_report_local_and_including_gather_peaks(slab_cpml_results): + _single, multi = slab_cpml_results + parallel_stats = multi.stats()["parallel_stats"] + local = parallel_stats["peak_memory_bytes"] + including_gather = parallel_stats["peak_memory_bytes_including_gather"] + + assert set(local) == {"cuda:0", "cuda:1"} + assert set(including_gather) == set(local) + assert all(local[device] > 0 for device in local) + assert all(including_gather[device] >= local[device] for device in local) + assert { + partition["device"]: partition["peak_memory_bytes"] + for partition in parallel_stats["partitions"] + } == local + + +def test_parallel_phase_timings_are_explicitly_unavailable_with_explanation( + slab_cpml_results, +): + _single, multi = slab_cpml_results + parallel_stats = multi.stats()["parallel_stats"] + + assert parallel_stats["wall_time_s"] > 0.0 + assert parallel_stats["compute_time_s"] is None + assert parallel_stats["communication_time_s"] is None + assert parallel_stats["exposed_communication_time_s"] is None + assert "external CUDA profiler" in parallel_stats["timing_note"] + assert "not inserted" in parallel_stats["timing_note"] diff --git a/tests/fdtd/multi_gpu/test_guard_regressions.py b/tests/fdtd/multi_gpu/test_guard_regressions.py new file mode 100644 index 0000000..2a0780c --- /dev/null +++ b/tests/fdtd/multi_gpu/test_guard_regressions.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed import DistributedFDTD +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig + + +_FREQUENCY = 1.0e9 + + +def _parallel(*, devices=("cuda:0", "cuda:1"), result_device="cuda:0"): + return FDTDParallelConfig( + devices=devices, + transport="cuda_p2p", + gather_fields=False, + result_device=result_device, + ) + + +def _scene(*, device="cpu") -> mw.Scene: + return mw.Scene( + domain=mw.Domain(bounds=((-0.4, 0.4), (-0.3, 0.3), (-0.3, 0.3))), + grid=mw.GridSpec.uniform(0.1), + boundary=mw.BoundarySpec.none(), + device=device, + ) + + +@pytest.mark.parametrize( + "material, message", + ( + ( + mw.Material( + eps_r=2.0, + nonlinearity=mw.NonlinearSusceptibility(chi3=1.0e-10), + ), + "nonlinear media", + ), + ( + mw.Material( + epsilon_tensor=mw.Tensor3x3( + ((2.0, 0.2, 0.1), (0.2, 2.5, 0.3), (0.1, 0.3, 3.0)) + ) + ), + "off-diagonal anisotropy", + ), + ), + ids=("nonlinear", "full_offdiagonal_anisotropy"), +) +def test_public_prepare_rejects_materials_without_complete_interface_halos( + material, + message, + cuda_p2p_devices, + cuda_memory_cleanup, +): + scene = _scene(device="cuda:0") + scene.add_structure( + mw.Structure( + name="cross_interface_material", + geometry=mw.Box(position=(0.0, 0.0, 0.0), size=(0.4, 0.4, 0.4)), + material=material, + ) + ) + simulation = mw.Simulation.fdtd( + scene, + frequencies=(_FREQUENCY,), + parallel=_parallel(devices=cuda_p2p_devices), + ) + + with pytest.raises(ValueError, match=message): + simulation.prepare() + + +def test_invalid_dipole_emission_source_is_rejected_before_hardware_prepare(): + scene = _scene() + scene.add_monitor( + mw.DipoleEmissionMonitor( + "emission", + source_name="missing_dipole", + frequencies=(_FREQUENCY,), + ) + ) + + with pytest.raises(ValueError, match="missing_dipole"): + DistributedFDTD( + scene, + frequency=_FREQUENCY, + parallel=_parallel(), + ) + + +@pytest.mark.parametrize( + "monitor", + ( + mw.PermittivityMonitor( + "permittivity", + position=(0.0, 0.0, 0.0), + size=(0.2, 0.2, 0.2), + ), + mw.MediumMonitor( + "medium", + position=(0.0, 0.0, 0.0), + size=(0.2, 0.2, 0.2), + ), + ), + ids=("permittivity", "medium"), +) +def test_material_monitors_are_explicitly_rejected_before_hardware_prepare(monitor): + scene = _scene() + scene.add_monitor(monitor) + + with pytest.raises(ValueError, match="material monitors"): + DistributedFDTD( + scene, + frequency=_FREQUENCY, + parallel=_parallel(), + ) + + +def _count_shard(frequencies=(), counts=()): + return SimpleNamespace( + solver=SimpleNamespace( + observer_frequencies=tuple(frequencies), + observer_sample_counts=tuple(counts), + dft_sample_counts=tuple(counts), + ) + ) + + +def test_observer_counts_merge_by_frequency_across_different_shard_subsets(): + solver = object.__new__(DistributedFDTD) + solver.shards = ( + _count_shard((_FREQUENCY,), (7,)), + _count_shard((1.5e9,), (11,)), + ) + solver._observer_frequencies = (_FREQUENCY, 1.5e9) + + assert solver.observer_frequencies == (_FREQUENCY, 1.5e9) + assert solver.observer_sample_counts == (7, 11) + + +def test_empty_distributed_sample_counts_use_empty_tuple_contract(): + solver = object.__new__(DistributedFDTD) + solver.shards = () + solver._observer_frequencies = () + + assert solver.dft_sample_counts == () + assert solver.observer_sample_counts == () + + +def test_active_absorber_type_aggregates_x_high_only_pml_from_last_shard(): + solver = object.__new__(DistributedFDTD) + solver.shards = ( + SimpleNamespace(solver=SimpleNamespace(active_absorber_type="none")), + SimpleNamespace(solver=SimpleNamespace(active_absorber_type="cpml")), + ) + + assert solver.active_absorber_type == "cpml" + + +def _hardware_probe(monkeypatch, *, missing_pair=None): + solver = object.__new__(DistributedFDTD) + solver.devices = tuple(torch.device(f"cuda:{index}") for index in range(3)) + solver.device = torch.device("cuda:0") + solver.transport = SimpleNamespace(preflight=Mock()) + calls = [] + + def can_access(source, destination): + pair = (int(source), int(destination)) + calls.append(pair) + return pair != missing_pair + + properties = SimpleNamespace(name="mock-gpu", major=8, minor=6) + monkeypatch.setattr(torch.cuda, "can_device_access_peer", can_access) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _device: properties) + return solver, calls + + +@pytest.mark.parametrize("missing_pair", ((0, 2), (2, 0))) +def test_result_device_preflight_rejects_missing_bidirectional_peer_link( + monkeypatch, + missing_pair, +): + solver, _calls = _hardware_probe(monkeypatch, missing_pair=missing_pair) + + with pytest.raises(RuntimeError, match="result gathering/reduction"): + solver._validate_hardware() + + +def test_result_device_preflight_checks_every_shard_in_both_directions(monkeypatch): + solver, calls = _hardware_probe(monkeypatch) + + solver._validate_hardware() + + assert {(0, 1), (1, 0), (0, 2), (2, 0)} <= set(calls) + solver.transport.preflight.assert_called_once_with() diff --git a/tests/fdtd/multi_gpu/test_monitor_merge_ownership.py b/tests/fdtd/multi_gpu/test_monitor_merge_ownership.py new file mode 100644 index 0000000..ff27a74 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_monitor_merge_ownership.py @@ -0,0 +1,351 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from witwin.maxwell.fdtd.distributed.monitor_merge import ( + _crop_aligned_to_physical_bounds, + merge_sharded_monitor_payloads, +) +from witwin.maxwell.fdtd_parallel import FDTDPartitionPlan, FDTDShardLayout + + +_FREQUENCY = 1.0e9 +_X_NODES = np.linspace(-0.5, 0.5, 11, dtype=np.float64) +_Y_NODES = np.linspace(-0.2, 0.2, 5, dtype=np.float64) +_Z_NODES = np.linspace(-0.3, 0.3, 7, dtype=np.float64) +_PHYSICAL_BOUNDS = ((-0.4, 0.4), (-0.2, 0.2), (-0.2, 0.2)) +_NODE_X_COMPONENTS = frozenset(("Ey", "Ez", "Hx")) + + +def _plan() -> FDTDPartitionPlan: + return FDTDPartitionPlan( + global_shape=(_X_NODES.size, _Y_NODES.size, _Z_NODES.size), + devices=("cuda:0", "cuda:1"), + low_pml_cells=1, + high_pml_cells=1, + ) + + +def _layout_map(plan: FDTDPartitionPlan) -> dict[int, FDTDShardLayout]: + return {layout.rank: layout for layout in plan.shard_layouts} + + +def _global_component_coords(component: str, axis: str) -> np.ndarray: + if axis == "x": + nodes = _X_NODES + is_node = component in _NODE_X_COMPONENTS + elif axis == "z": + nodes = _Z_NODES + is_node = component in {"Ex", "Ey", "Hz"} + else: + raise ValueError(f"Unsupported test axis {axis!r}.") + return nodes if is_node else 0.5 * (nodes[:-1] + nodes[1:]) + + +def _local_coords( + layout: FDTDShardLayout, + component: str, +) -> tuple[np.ndarray, np.ndarray]: + component_layout = layout.component(component) + x = _global_component_coords(component, "x") + allocation = component_layout.allocation_global_slice[0] + return np.array(x[allocation], copy=True), _global_component_coords(component, "z") + + +def _tagged_component( + layout: FDTDShardLayout, + component: str, + *, + device: str | torch.device = "cpu", + offset: float = 0.0, +) -> dict[str, object]: + component_layout = layout.component(component) + x, z = _local_coords(layout, component) + data = torch.full( + (x.size, z.size), + complex(9000.0 + layout.rank, 0.0), + device=device, + dtype=torch.complex64, + ) + owned_local = component_layout.owned_local_slice[0] + owned_global = component_layout.owned_global_slice[0] + values = torch.arange( + int(owned_global.start), + int(owned_global.stop), + device=device, + dtype=torch.float32, + ) + data[owned_local] = (values + float(offset))[:, None].to(torch.complex64) + return {"data": data, "coords": (x, z)} + + +def _tagged_plane(layout: FDTDShardLayout) -> dict[str, object]: + fields = ("Ex", "Ez") + return { + "kind": "plane", + "monitor_type": "plane", + "fields": fields, + "components": { + "Ex": _tagged_component(layout, "Ex"), + "Ez": _tagged_component(layout, "Ez", offset=100.0), + }, + "samples": 8, + "frequency": _FREQUENCY, + "frequencies": (_FREQUENCY,), + "axis": "y", + "position": 0.0, + "compute_flux": False, + "normal_direction": "+", + "Ex": torch.full((1, 1), complex(7000.0 + layout.rank, 0.0)), + "Ez": torch.full((1, 1), complex(7000.0 + layout.rank, 0.0)), + } + + +def _constant_vector_plane( + layout: FDTDShardLayout, + *, + device: str | torch.device, + monitor_type: str, + compute_flux: bool, +) -> dict[str, object]: + constants = {"Ex": 0.0, "Ez": 1.0, "Hx": 2.0, "Hz": 0.0} + components = {} + for component, value in constants.items(): + component_layout = layout.component(component) + x, z = _local_coords(layout, component) + data = torch.full( + (x.size, z.size), + complex(5000.0 + layout.rank, 0.0), + device=device, + dtype=torch.complex64, + ) + data[component_layout.owned_local_slice[0]] = complex(value, 0.0) + components[component] = {"data": data, "coords": (x, z)} + + payload: dict[str, object] = { + "kind": "plane", + "monitor_type": monitor_type, + "fields": tuple(constants), + "components": components, + "samples": 8, + "frequency": _FREQUENCY, + "frequencies": (_FREQUENCY,), + "axis": "y", + "position": 0.0, + "compute_flux": compute_flux, + "normal_direction": "+", + "mode_spec": {"mode_index": 0} if monitor_type == "mode" else None, + "flux": torch.tensor(1234.0, device=device), + "power": torch.tensor(1234.0, device=device), + } + for component in constants: + payload[component] = torch.full( + (1, 1), + complex(6000.0 + layout.rank, 0.0), + device=device, + dtype=torch.complex64, + ) + return payload + + +def test_tiled_plane_discards_node_high_ghost_and_cell_low_ghost(): + plan = _plan() + layouts = _layout_map(plan) + payloads = tuple( + (layout.rank, {"plane": _tagged_plane(layout)}) + for layout in reversed(plan.shard_layouts) + ) + + output = merge_sharded_monitor_payloads( + ("plane",), + payloads, + shard_layouts=layouts, + physical_bounds=_PHYSICAL_BOUNDS, + result_device="cpu", + )["plane"] + + ex = output["components"]["Ex"] + ez = output["components"]["Ez"] + np.testing.assert_array_equal( + ex["coords"][0], + 0.5 * (_X_NODES[:-1] + _X_NODES[1:]), + ) + np.testing.assert_array_equal(ez["coords"][0], _X_NODES) + torch.testing.assert_close( + ex["data"].real[:, 0], + torch.arange(0, 10, dtype=torch.float32), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + ez["data"].real[:, 0], + torch.arange(0, 11, dtype=torch.float32) + 100.0, + rtol=0.0, + atol=0.0, + ) + assert float(ez["data"][5, 0].real) == 105.0 + + +def test_flux_recomputes_from_global_owned_components_on_result_gpu( + cuda_p2p_devices, + cuda_memory_cleanup, +): + plan = _plan() + layouts = _layout_map(plan) + x_plane = { + "kind": "plane", + "monitor_type": "plane", + "fields": ("Ez",), + "axis": "x", + "position": 0.1, + "compute_flux": False, + "normal_direction": "+", + "data": torch.ones((2, 2), device=cuda_p2p_devices[1]), + } + payloads = tuple( + ( + layout.rank, + { + "flux": _constant_vector_plane( + layout, + device=cuda_p2p_devices[layout.rank], + monitor_type="plane", + compute_flux=True, + ) + } + | ({"x_plane": x_plane} if layout.rank == 1 else {}), + ) + for layout in reversed(plan.shard_layouts) + ) + + merged = merge_sharded_monitor_payloads( + ("x_plane", "flux"), + payloads, + shard_layouts=layouts, + physical_bounds=_PHYSICAL_BOUNDS, + result_device=cuda_p2p_devices[0], + ) + assert tuple(merged) == ("x_plane", "flux") + assert merged["x_plane"]["data"].device == cuda_p2p_devices[0] + output = merged["flux"] + + assert output["flux"].device == cuda_p2p_devices[0] + assert output["power"].device == cuda_p2p_devices[0] + torch.testing.assert_close( + output["flux"], + torch.tensor(0.32, device=cuda_p2p_devices[0]), + rtol=1.0e-6, + atol=1.0e-6, + ) + torch.testing.assert_close(output["power"], output["flux"]) + np.testing.assert_allclose(output["x"], np.linspace(-0.35, 0.35, 8)) + np.testing.assert_allclose(output["z"], np.linspace(-0.15, 0.15, 4)) + assert output["Ez"].shape == (8, 4) + assert torch.all(output["Ez"] == 1.0) + assert torch.all(output["Hx"] == 2.0) + + +def test_tiled_mode_realigns_global_raw_components_after_owned_stitch(): + plan = _plan() + layouts = _layout_map(plan) + payloads = tuple( + ( + layout.rank, + { + "mode": _constant_vector_plane( + layout, + device="cpu", + monitor_type="mode", + compute_flux=False, + ) + }, + ) + for layout in reversed(plan.shard_layouts) + ) + + output = merge_sharded_monitor_payloads( + ("mode",), + payloads, + shard_layouts=layouts, + physical_bounds=_PHYSICAL_BOUNDS, + result_device="cpu", + )["mode"] + + assert output["monitor_type"] == "mode" + assert output["mode_spec"] == {"mode_index": 0} + assert output["Ez"].shape == (10, 6) + assert torch.all(output["Ez"] == 1.0) + assert torch.all(output["Hx"] == 2.0) + assert float(output["Ez"].real.max()) < 10.0 + + +def test_tiled_plane_requires_layout_for_every_contributing_rank(): + plan = _plan() + left, right = plan.shard_layouts + payloads = ( + (left.rank, {"plane": _tagged_plane(left)}), + (right.rank, {"plane": _tagged_plane(right)}), + ) + with pytest.raises(ValueError, match="missing shard layouts.*rank 1"): + merge_sharded_monitor_payloads( + ("plane",), + payloads, + shard_layouts={0: left}, + physical_bounds=_PHYSICAL_BOUNDS, + result_device="cpu", + ) + + +@pytest.mark.parametrize( + ("monitor_type", "compute_flux"), + (("plane", True), ("mode", False)), +) +def test_x_normal_flux_and_mode_require_exactly_one_owner(monitor_type, compute_flux): + payload = { + "kind": "plane", + "monitor_type": monitor_type, + "fields": ("Ez",), + "axis": "x", + "position": 0.0, + "compute_flux": compute_flux, + "normal_direction": "+", + "data": torch.ones((2, 2)), + } + with pytest.raises(RuntimeError, match="x-normal plane.*more than one shard owner"): + merge_sharded_monitor_payloads( + ("plane",), + ((0, {"plane": payload}), (1, {"plane": payload})), + shard_layouts=_layout_map(_plan()), + physical_bounds=_PHYSICAL_BOUNDS, + result_device="cpu", + ) + + + +def test_nm_scale_physical_crop_does_not_admit_pml_samples(): + nm = 1.0e-9 + coords = np.asarray((-150.0, -50.0, 50.0, 150.0), dtype=np.float64) * nm + aligned = { + "x": coords, + "z": coords, + "fields": { + "Ez": torch.arange(16, dtype=torch.float32).reshape(4, 4), + }, + } + + cropped = _crop_aligned_to_physical_bounds( + "y", + aligned, + ((-100.0 * nm, 100.0 * nm), (0.0, 0.0), (-100.0 * nm, 100.0 * nm)), + ) + + np.testing.assert_array_equal(cropped["x"], coords[1:3]) + np.testing.assert_array_equal(cropped["z"], coords[1:3]) + torch.testing.assert_close( + cropped["fields"]["Ez"], + torch.tensor(((5.0, 6.0), (9.0, 10.0))), + rtol=0.0, + atol=0.0, + ) diff --git a/tests/fdtd/multi_gpu/test_monitor_numerical.py b/tests/fdtd/multi_gpu/test_monitor_numerical.py new file mode 100644 index 0000000..bc85e81 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_monitor_numerical.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed import DistributedFDTD +from witwin.maxwell.fdtd.solver import FDTD +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig +from witwin.maxwell.result import Result +from witwin.maxwell.scene import prepare_scene + + +_FREQUENCY = 1.0e9 +_TIME_STEPS = 48 +_MONITOR_RTOL = 5.0e-5 +_MONITOR_ATOL = 5.0e-6 +_FLUX_REL_TOL = 1.0e-3 + + +def _monitor_scene() -> mw.Scene: + x = np.linspace(-0.5, 0.5, 11, dtype=np.float64) + y = np.linspace(-0.2, 0.2, 9, dtype=np.float64) + z = np.linspace(-0.2, 0.2, 9, dtype=np.float64) + scene = mw.Scene( + domain=mw.Domain( + bounds=((float(x[0]), float(x[-1])), (float(y[0]), float(y[-1])), (float(z[0]), float(z[-1]))) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + scene.add_source( + mw.PointDipole( + position=(-0.2, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=mw.CW(frequency=_FREQUENCY, amplitude=1.0), + name="source", + ) + ) + scene.add_monitor( + mw.PlaneMonitor( + "y_plane", + axis="y", + position=0.05, + fields=("Ez",), + frequencies=(_FREQUENCY,), + ) + ) + scene.add_monitor( + mw.FluxMonitor( + "z_flux", + axis="z", + position=0.1, + frequencies=(_FREQUENCY,), + ) + ) + scene.add_monitor( + mw.PlaneMonitor( + "x_plane", + axis="x", + position=0.1, + fields=("Ez",), + frequencies=(_FREQUENCY,), + ) + ) + scene.add_monitor( + mw.FluxMonitor( + "x_flux", + axis="x", + position=0.0, + frequencies=(_FREQUENCY,), + ) + ) + scene.add_monitor( + mw.ModeMonitor( + "z_mode", + position=(0.0, 0.0, 0.1), + size=(0.8, 0.3, 0.0), + mode_index=0, + polarization="Ex", + frequencies=(_FREQUENCY,), + ) + ) + return scene + + +def _solve(solver): + return solver.solve( + time_steps=_TIME_STEPS, + dft_frequency=_FREQUENCY, + dft_window="none", + full_field_dft=False, + normalize_source=False, + use_cuda_graph=False, + ) + + +@pytest.fixture(scope="module") +def monitor_pair(cuda_p2p_devices): + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + scene = _monitor_scene() + single_solver = FDTD(prepare_scene(scene), frequency=_FREQUENCY, absorber_type="cpml") + single_solver.init_field() + single_output = _solve(single_solver) + + distributed_solver = DistributedFDTD( + scene, + frequency=_FREQUENCY, + parallel=FDTDParallelConfig( + devices=cuda_p2p_devices, + transport="cuda_p2p", + overlap=True, + gather_fields=False, + result_device=cuda_p2p_devices[0], + ), + absorber_type="cpml", + ) + distributed_solver.init_field() + distributed_output = _solve(distributed_solver) + + single_result = Result( + method="fdtd", + scene=scene, + prepared_scene=single_solver.scene, + frequency=_FREQUENCY, + solver=single_solver, + monitors=single_output["observers"], + raw_output=single_output, + ) + distributed_result = Result( + method="fdtd", + scene=scene, + prepared_scene=distributed_solver.scene, + frequency=_FREQUENCY, + solver=distributed_solver, + monitors=distributed_output["observers"], + raw_output=distributed_output, + ) + yield SimpleNamespace( + single=single_output["observers"], + distributed=distributed_output["observers"], + single_result=single_result, + distributed_result=distributed_result, + ) + + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + +def _assert_coords_close(actual, expected) -> None: + np.testing.assert_allclose( + np.asarray(actual, dtype=np.float64), + np.asarray(expected, dtype=np.float64), + rtol=0.0, + atol=1.0e-15, + ) + + +def _assert_tensor_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close( + actual.to(expected.device), + expected, + rtol=_MONITOR_RTOL, + atol=_MONITOR_ATOL, + ) + + +def _assert_raw_fields_close(actual, expected, fields) -> None: + for component in fields: + actual_component = actual["components"][component] + expected_component = expected["components"][component] + _assert_tensor_close(actual_component["data"], expected_component["data"]) + _assert_coords_close(actual_component["coords"][0], expected_component["coords"][0]) + _assert_coords_close(actual_component["coords"][1], expected_component["coords"][1]) + _assert_tensor_close(actual[component], expected[component]) + + +def _assert_flux_close(actual, expected) -> None: + for key in ("flux", "power"): + actual_value = actual[key].to(expected[key].device) + expected_value = expected[key] + denominator = max(float(torch.abs(expected_value).item()), 1.0e-30) + relative_error = float(torch.abs(actual_value - expected_value).item()) / denominator + assert denominator > 1.0e-12 + assert relative_error <= _FLUX_REL_TOL + + +def test_y_normal_single_component_plane_monitor_matches_one_gpu(monitor_pair): + expected = monitor_pair.single["y_plane"] + actual = monitor_pair.distributed["y_plane"] + _assert_raw_fields_close(actual, expected, ("Ez",)) + _assert_coords_close(actual["coords"][0], expected["coords"][0]) + _assert_coords_close(actual["coords"][1], expected["coords"][1]) + assert float(torch.abs(expected["Ez"]).max().item()) > _MONITOR_ATOL + + +def test_z_normal_flux_monitor_matches_fields_and_integrated_flux(monitor_pair): + expected = monitor_pair.single["z_flux"] + actual = monitor_pair.distributed["z_flux"] + _assert_raw_fields_close(actual, expected, ("Ex", "Ey", "Hx", "Hy")) + _assert_flux_close(actual, expected) + + +def test_off_split_x_normal_plane_monitor_matches_one_gpu(monitor_pair): + expected = monitor_pair.single["x_plane"] + actual = monitor_pair.distributed["x_plane"] + _assert_raw_fields_close(actual, expected, ("Ez",)) + assert float(torch.abs(expected["Ez"]).max().item()) > _MONITOR_ATOL + + +def test_split_x_normal_flux_monitor_matches_fields_and_integrated_flux(monitor_pair): + expected = monitor_pair.single["x_flux"] + actual = monitor_pair.distributed["x_flux"] + _assert_raw_fields_close(actual, expected, ("Ey", "Ez", "Hy", "Hz")) + _assert_flux_close(actual, expected) + + +def test_mode_monitor_matches_raw_power_flux_and_resolved_mode(monitor_pair): + expected_raw = monitor_pair.single["z_mode"] + actual_raw = monitor_pair.distributed["z_mode"] + _assert_raw_fields_close(actual_raw, expected_raw, ("Ex", "Ey", "Hx", "Hy")) + _assert_flux_close(actual_raw, expected_raw) + + expected = monitor_pair.single_result.monitor("z_mode") + actual = monitor_pair.distributed_result.monitor("z_mode") + assert actual["effective_index"] == pytest.approx(expected["effective_index"], rel=1.0e-9) + assert actual["beta"] == pytest.approx(expected["beta"], rel=1.0e-9) + for key in ( + "mode_power", + "total_power", + "amplitude_forward", + "amplitude_backward", + "power_forward", + "power_backward", + ): + _assert_tensor_close(actual[key], expected[key]) diff --git a/tests/fdtd/multi_gpu/test_numerical_matrix.py b/tests/fdtd/multi_gpu/test_numerical_matrix.py new file mode 100644 index 0000000..a0e5b54 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_numerical_matrix.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed import DistributedFDTD +from witwin.maxwell.fdtd.solver import FDTD +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig +from witwin.maxwell.result import Result +from witwin.maxwell.scene import prepare_scene + + +_FREQUENCY = 1.0e9 +_FIELD_NAMES = ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") +_DFT_FREQUENCIES = (0.8e9, 1.0e9, 1.2e9) + + +@dataclass +class _SolvePair: + single: FDTD + single_output: dict[str, Any] | None + distributed: DistributedFDTD + distributed_output: dict[str, Any] | None + + +def _parallel(devices, *, gather_fields: bool = True) -> FDTDParallelConfig: + return FDTDParallelConfig( + devices=devices, + transport="cuda_p2p", + overlap=True, + gather_fields=gather_fields, + result_device=devices[0], + ) + + +def _run_pair( + scene: mw.Scene, + devices, + *, + time_steps: int, + dft_frequency: float | tuple[float, ...] = _FREQUENCY, + full_field_dft: bool = False, + shutoff: float = 0.0, + shutoff_check_interval: int = 10, +) -> _SolvePair: + single = FDTD(prepare_scene(scene), frequency=_FREQUENCY, absorber_type="cpml") + single.init_field() + single_output = single.solve( + time_steps=time_steps, + dft_frequency=dft_frequency, + dft_window="none", + full_field_dft=full_field_dft, + normalize_source=False, + shutoff=shutoff, + shutoff_check_interval=shutoff_check_interval, + use_cuda_graph=False, + ) + + distributed = DistributedFDTD( + scene, + frequency=_FREQUENCY, + parallel=_parallel(devices), + absorber_type="cpml", + ) + distributed.init_field() + distributed_output = distributed.solve( + time_steps=time_steps, + dft_frequency=dft_frequency, + dft_window="none", + full_field_dft=full_field_dft, + normalize_source=False, + shutoff=shutoff, + shutoff_check_interval=shutoff_check_interval, + use_cuda_graph=False, + ) + return _SolvePair(single, single_output, distributed, distributed_output) + + +def _small_scene(*, source_time=None, time_monitor: bool = False) -> mw.Scene: + x = np.linspace(-0.4, 0.4, 9, dtype=np.float64) + y = np.linspace(-0.2, 0.2, 5, dtype=np.float64) + z = np.linspace(-0.2, 0.2, 5, dtype=np.float64) + scene = mw.Scene( + domain=mw.Domain( + bounds=((float(x[0]), float(x[-1])), (float(y[0]), float(y[-1])), (float(z[0]), float(z[-1]))) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + if source_time is not None: + scene.add_source( + mw.PointDipole( + position=(0.0, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=source_time, + name="interface_pulse", + ) + ) + if time_monitor: + scene.add_monitor( + mw.FieldTimeMonitor( + "six_field_trace", + components=_FIELD_NAMES, + position=(0.1, 0.0, 0.0), + start=2, + stop=46, + interval=3, + ) + ) + return scene + + +def _distributed_field(pair: _SolvePair, name: str) -> torch.Tensor: + local = tuple(getattr(shard.solver, name) for shard in pair.distributed.shards) + return pair.distributed._gather_component(name, local) + + +def _assert_step_field_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + actual = actual.to(expected.device) + assert torch.isfinite(actual).all() + assert torch.isfinite(expected).all() + error = torch.abs(actual - expected) + assert float(error.max().item()) <= 2.0e-6 + + scale = float(torch.abs(expected).max().item()) + significant = torch.abs(expected) >= max(1.0e-6, 1.0e-4 * scale) + if bool(significant.any()): + relative = error[significant] / torch.abs(expected[significant]) + assert float(relative.max().item()) <= 2.0e-5 + + +@pytest.fixture(scope="module") +def pulse_matrix_pair(cuda_p2p_devices): + scene = _small_scene( + source_time=mw.GaussianPulse( + frequency=_FREQUENCY, + fwidth=0.4e9, + amplitude=10.0, + ), + time_monitor=True, + ) + pair = _run_pair( + scene, + cuda_p2p_devices, + time_steps=48, + dft_frequency=_DFT_FREQUENCIES, + full_field_dft=True, + ) + yield pair + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + +def test_zero_field_remains_zero_on_one_and_two_gpus( + cuda_p2p_devices, + cuda_memory_cleanup, +): + pair = _run_pair( + _small_scene(), + cuda_p2p_devices, + time_steps=8, + full_field_dft=False, + ) + for name in _FIELD_NAMES: + single = getattr(pair.single, name) + distributed = _distributed_field(pair, name) + assert torch.count_nonzero(single).item() == 0 + assert torch.count_nonzero(distributed).item() == 0 + torch.testing.assert_close(distributed, single, rtol=0.0, atol=0.0) + + +def test_gaussian_impulse_six_step_fields_match_one_gpu(pulse_matrix_pair): + pair = pulse_matrix_pair + assert float(torch.abs(pair.single.Ez).max().item()) > 0.0 + for name in _FIELD_NAMES: + _assert_step_field_close( + _distributed_field(pair, name), + getattr(pair.single, name), + ) + + +def test_actual_six_component_field_time_monitor_matches_one_gpu(pulse_matrix_pair): + pair = pulse_matrix_pair + single = pair.single_output["observers"]["six_field_trace"] + distributed = pair.distributed_output["observers"]["six_field_trace"] + + assert distributed["kind"] == single["kind"] == "field_time" + assert distributed["fields"] == single["fields"] == _FIELD_NAMES + assert distributed["start"] == single["start"] == 2 + assert distributed["stop"] == single["stop"] == 46 + assert distributed["interval"] == single["interval"] == 3 + torch.testing.assert_close(distributed["t"], single["t"], rtol=0.0, atol=0.0) + assert tuple(distributed["t"].shape) == (15,) + for name in _FIELD_NAMES: + torch.testing.assert_close( + distributed["components"][name], + single["components"][name], + rtol=5.0e-5, + atol=5.0e-6, + ) + assert float(torch.abs(single["components"]["Ez"]).max().item()) > 0.0 + + +def test_multifrequency_full_field_dft_values_match_one_gpu(pulse_matrix_pair): + pair = pulse_matrix_pair + assert pair.single_output["frequencies"] == _DFT_FREQUENCIES + assert pair.distributed_output["frequencies"] == _DFT_FREQUENCIES + for name in ("Ex", "Ey", "Ez"): + single = pair.single_output[name] + distributed = pair.distributed_output[name] + assert single.shape[0] == distributed.shape[0] == len(_DFT_FREQUENCIES) + assert torch.is_complex(single) + assert torch.is_complex(distributed) + assert float(torch.abs(single).max().item()) > 0.0 + torch.testing.assert_close( + distributed, + single, + rtol=5.0e-5, + atol=5.0e-6, + ) + + +def test_two_gpu_sharded_result_save_lazy_and_gather_round_trip( + pulse_matrix_pair, + tmp_path, +): + pair = pulse_matrix_pair + solver = pair.distributed + output = pair.distributed_output + references = {name: output[name] for name in ("Ex", "Ey", "Ez")} + result = Result( + method="fdtd", + scene=solver.logical_scene, + prepared_scene=solver.scene, + frequencies=_DFT_FREQUENCIES, + solver=solver, + fields={name.upper(): tensor for name, tensor in references.items()}, + monitors=output.get("observers", {}), + solver_stats={"parallel_stats": solver.parallel_stats}, + ) + + directory = tmp_path / "two-gpu-sharded-result" + manifest = result.save_sharded(directory) + + assert tuple(path.name for path in sorted(directory.iterdir())) == ( + "manifest.json", + "rank-0000.pt", + "rank-0001.pt", + "result.pt", + ) + assert manifest.frequencies == _DFT_FREQUENCIES + assert tuple(shard.rank for shard in manifest.shards) == (0, 1) + manifest_components = {component.name: component for component in manifest.components} + for name, reference in references.items(): + global_component = manifest_components[name] + assert global_component.shape == tuple(reference.shape) + assert global_component.x_axis == reference.ndim - 3 == 1 + + cursor = 0 + persisted_x = 0 + for shard, path in zip(manifest.shards, manifest.shard_paths(directory)): + local = next(component for component in shard.components if component.name == name) + assert local.global_x_slice[0] == cursor + cursor = local.global_x_slice[1] + persisted_x += local.shape[local.x_axis] + payload = torch.load(path, map_location="cpu", weights_only=True) + tensor = payload["components"][name]["tensor"] + assert tensor.shape[local.x_axis] == ( + local.global_x_slice[1] - local.global_x_slice[0] + ) + assert cursor == reference.shape[global_component.x_axis] + assert persisted_x == reference.shape[global_component.x_axis] + + lazy = Result.load_sharded( + directory, + scene=solver.logical_scene, + prepared_scene=solver.scene, + gather_fields=False, + map_location="cpu", + ) + assert lazy.is_sharded is True + assert lazy.fields == {} + assert tuple(path.name for path in lazy.shard_paths) == ( + "rank-0000.pt", + "rank-0001.pt", + ) + + gathered = Result.load_sharded( + directory, + scene=solver.logical_scene, + prepared_scene=solver.scene, + gather_fields=True, + map_location="cpu", + ) + assert tuple(gathered.fields) == ("EX", "EY", "EZ") + for name, reference in references.items(): + torch.testing.assert_close( + gathered.fields[name.upper()], + reference.detach().cpu(), + rtol=0.0, + atol=0.0, + ) + + +def _magnetic_ade_scene() -> mw.Scene: + scene = _small_scene( + source_time=mw.GaussianPulse( + frequency=_FREQUENCY, + fwidth=0.4e9, + amplitude=10.0, + ) + ) + scene.add_structure( + mw.Structure( + name="interface_magnetic_lorentz", + geometry=mw.Box( + position=(0.0, 0.0, 0.0), + size=(0.4, 0.3, 0.3), + ), + material=mw.Material( + mu_lorentz_poles=( + mw.LorentzPole( + delta_eps=0.5, + resonance_frequency=1.2e9, + gamma=1.0e8, + ), + ), + ), + ) + ) + return scene + + +def test_magnetic_lorentz_ade_six_fields_match_one_gpu( + cuda_p2p_devices, + cuda_memory_cleanup, +): + pair = _run_pair( + _magnetic_ade_scene(), + cuda_p2p_devices, + time_steps=48, + full_field_dft=False, + ) + assert pair.single.magnetic_dispersive_enabled is True + assert all( + shard.solver.magnetic_dispersive_enabled + for shard in pair.distributed.shards + ) + for name in _FIELD_NAMES: + _assert_step_field_close( + _distributed_field(pair, name), + getattr(pair.single, name), + ) + + +def _shutoff_scene() -> mw.Scene: + nodes = np.linspace(-0.5, 0.5, 26, dtype=np.float64) + scene = mw.Scene( + domain=mw.Domain(bounds=((-0.5, 0.5), (-0.5, 0.5), (-0.5, 0.5))), + grid=mw.GridSpec.custom(nodes, nodes, nodes), + boundary=mw.BoundarySpec.pml(num_layers=6, strength=1.0), + device="cuda:0", + ) + scene.add_source( + mw.UniformCurrentSource( + size=(0.08, 0.08, 0.08), + center=(0.0, 0.0, 0.0), + polarization="Ez", + source_time=mw.GaussianPulse( + frequency=_FREQUENCY, + fwidth=0.4e9, + amplitude=100.0, + ), + name="ringdown_pulse", + ) + ) + return scene + + +def test_early_shutoff_step_and_six_fields_match_one_gpu( + cuda_p2p_devices, + cuda_memory_cleanup, +): + planned_steps = 1500 + pair = _run_pair( + _shutoff_scene(), + cuda_p2p_devices, + time_steps=planned_steps, + full_field_dft=False, + shutoff=5.0e-2, + shutoff_check_interval=25, + ) + + assert pair.single._shutoff_triggered is True + assert pair.distributed._shutoff_triggered is True + assert pair.single._shutoff_step == pair.distributed._shutoff_step + assert pair.single._shutoff_step is not None + assert pair.single._shutoff_step < planned_steps - 1 + executed_steps = pair.single._shutoff_step + 1 + assert all( + shard.solver._shutoff_step == pair.single._shutoff_step + for shard in pair.distributed.shards + ) + stats = pair.distributed.parallel_stats + assert stats["halo_bytes_total"] == stats["halo_bytes_per_step"] * executed_steps + for name in _FIELD_NAMES: + _assert_step_field_close( + _distributed_field(pair, name), + getattr(pair.single, name), + ) diff --git a/tests/fdtd/multi_gpu/test_partition.py b/tests/fdtd/multi_gpu/test_partition.py new file mode 100644 index 0000000..9c60fb3 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_partition.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest +import torch + +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig, FDTDPartitionPlan + + +def _plan(*, shape=(11, 6, 5), devices=("cuda:0", "cuda:1", "cuda:2")): + return FDTDPartitionPlan(global_shape=shape, devices=devices) + + +def _x_extent(index3): + return index3[0].start, index3[0].stop + + +def test_parallel_config_normalizes_without_touching_cuda_runtime(): + config = FDTDParallelConfig( + devices=(torch.device("cuda:1"), "cuda:0"), + decomposition_axis=" X ", + transport=" CUDA_P2P ", + gather_fields=True, + ) + + assert config.devices == ("cuda:1", "cuda:0") + assert config.decomposition_axis == "x" + assert config.transport == "cuda_p2p" + assert config.result_device == "cuda:1" + assert config.gather_fields is True + with pytest.raises(FrozenInstanceError): + config.transport = "auto" + + +@pytest.mark.parametrize( + ("kwargs", "error", "message"), + [ + ({"devices": ("cuda:0",)}, ValueError, "at least 2"), + ({"devices": ("cuda:0", "cuda:0")}, ValueError, "unique"), + ({"devices": ("cuda:0", "cpu")}, ValueError, "CUDA"), + ({"devices": ("cuda", "cuda:1")}, ValueError, "explicit CUDA index"), + ({"devices": ("cuda:0", "cuda:1"), "decomposition_axis": "y"}, ValueError, "must be 'x'"), + ({"devices": ("cuda:0", "cuda:1"), "transport": "gloo"}, ValueError, "transport"), + ( + {"devices": ("cuda:0", "cuda:1"), "result_device": "cuda:2"}, + ValueError, + "participating", + ), + ({"devices": ("cuda:0", "cuda:1"), "overlap": 1}, TypeError, "overlap"), + ], +) +def test_parallel_config_rejects_invalid_structure(kwargs, error, message): + with pytest.raises(error, match=message): + FDTDParallelConfig(**kwargs) + + +def test_balanced_partition_puts_remainder_on_low_ranks(): + plan = _plan() + + assert plan.cell_count == 10 + assert plan.cell_intervals == ((0, 4), (4, 7), (7, 10)) + assert tuple(layout.owned_cell_count for layout in plan.shard_layouts) == (4, 3, 3) + assert tuple(layout.device for layout in plan.shard_layouts) == ( + "cuda:0", + "cuda:1", + "cuda:2", + ) + + +def test_padded_storage_slices_keep_low_ghost_and_terminal_node_explicit(): + plan = _plan() + + first, middle, last = plan.shard_layouts + assert first.global_cell_owned == slice(0, 4) + assert first.global_node_owned == slice(0, 4) + assert first.storage_cell_owned == slice(0, 4) + assert first.storage_node_owned == slice(0, 4) + + assert middle.global_cell_owned == slice(4, 7) + assert middle.global_node_owned == slice(4, 7) + assert middle.storage_cell_owned == slice(1, 4) + assert middle.storage_node_owned == slice(1, 4) + + assert last.global_cell_owned == slice(7, 10) + assert last.global_node_owned == slice(7, 11) + assert last.storage_cell_owned == slice(1, 4) + assert last.storage_node_owned == slice(1, 5) + + assert _x_extent(first.component("Ex").allocation_global_slice) == (0, 4) + assert _x_extent(first.component("Hx").allocation_global_slice) == (0, 5) + assert _x_extent(middle.component("Ex").allocation_global_slice) == (3, 7) + assert _x_extent(middle.component("Hx").allocation_global_slice) == (3, 8) + assert _x_extent(last.component("Ex").allocation_global_slice) == (6, 10) + assert _x_extent(last.component("Hx").allocation_global_slice) == (6, 11) + + +def test_outer_shards_include_pml_while_balancing_only_physical_cells(): + plan = FDTDPartitionPlan( + global_shape=(15, 6, 5), + devices=("cuda:0", "cuda:1", "cuda:2"), + low_pml_cells=2, + high_pml_cells=2, + ) + + assert plan.physical_cell_count == 10 + assert plan.cell_intervals == ((0, 4), (4, 7), (7, 10)) + assert plan.global_cell_intervals == ((0, 6), (6, 9), (9, 14)) + + first, middle, last = plan.shard_layouts + assert (first.physical_cell_begin, first.physical_cell_end) == (0, 4) + assert first.global_cell_owned == slice(0, 6) + assert first.global_node_owned == slice(0, 6) + + assert (middle.physical_cell_begin, middle.physical_cell_end) == (4, 7) + assert middle.global_cell_owned == slice(6, 9) + assert middle.global_node_owned == slice(6, 9) + + assert (last.physical_cell_begin, last.physical_cell_end) == (7, 10) + assert last.global_cell_owned == slice(9, 14) + assert last.global_node_owned == slice(9, 15) + + assert plan.owner_of_cell(0) == 0 + assert plan.owner_of_cell(13) == 2 + assert _x_extent(first.component("Ex").owned_global_slice) == (0, 6) + assert _x_extent(last.component("Ey").owned_global_slice) == (9, 15) + + +def test_all_six_component_shapes_and_owned_extents_follow_yee_staggering(): + plan = _plan(shape=(11, 6, 5)) + assert dict(plan.component_shapes) == { + "Ex": (10, 6, 5), + "Ey": (11, 5, 5), + "Ez": (11, 6, 4), + "Hx": (11, 5, 4), + "Hy": (10, 6, 4), + "Hz": (10, 5, 5), + } + + first, middle, last = plan.shard_layouts + for component in ("Ex", "Hy", "Hz"): + assert _x_extent(first.component(component).owned_global_slice) == (0, 4) + assert _x_extent(middle.component(component).owned_global_slice) == (4, 7) + assert _x_extent(last.component(component).owned_global_slice) == (7, 10) + for component in ("Ey", "Ez", "Hx"): + assert _x_extent(first.component(component).owned_global_slice) == (0, 4) + assert _x_extent(middle.component(component).owned_global_slice) == (4, 7) + assert _x_extent(last.component(component).owned_global_slice) == (7, 11) + + +def test_interface_and_terminal_values_have_exactly_one_owner(): + plan = _plan() + + for component, x_extent in plan.component_shapes.items(): + for x_index in range(x_extent[0]): + owners = [ + shard.rank + for shard in plan.shard_layouts + if shard.component(component).owns((x_index, 0, 0)) + ] + assert owners == [plan.owner_of_component_x(component, x_index)] + + assert plan.owner_of_node(4) == 1 + assert plan.owner_of_node(7) == 2 + assert plan.owner_of_node(10) == 2 + assert plan.owner_of_cell(3) == 0 + assert plan.owner_of_cell(4) == 1 + + +def test_tangential_halos_and_send_regions_match_half_step_dependencies(): + first, middle, last = _plan().shard_layouts + + for component in ("Ey", "Ez"): + first_component = first.component(component) + middle_component = middle.component(component) + last_component = last.component(component) + assert _x_extent(first_component.high_halo.global_slice) == (4, 5) + assert first_component.high_halo.neighbor_rank == 1 + assert first_component.low_halo is None + assert _x_extent(middle_component.low_send.global_slice) == (4, 5) + assert middle_component.low_send.neighbor_rank == 0 + assert _x_extent(middle_component.high_halo.global_slice) == (7, 8) + assert last_component.high_halo is None + assert _x_extent(last_component.low_send.global_slice) == (7, 8) + + for component in ("Hy", "Hz"): + first_component = first.component(component) + middle_component = middle.component(component) + last_component = last.component(component) + assert first_component.low_halo is None + assert _x_extent(first_component.high_send.global_slice) == (3, 4) + assert first_component.high_send.neighbor_rank == 1 + assert _x_extent(middle_component.low_halo.global_slice) == (3, 4) + assert middle_component.low_halo.neighbor_rank == 0 + assert _x_extent(middle_component.high_send.global_slice) == (6, 7) + assert _x_extent(last_component.low_halo.global_slice) == (6, 7) + assert last_component.high_send is None + + for shard in (first, middle, last): + for component in ("Ex", "Hx"): + layout = shard.component(component) + assert layout.receive_halos == () + assert layout.send_regions == () + assert layout.local_shape[0] >= layout.owned_shape[0] + + +def test_global_local_round_trip_includes_read_only_halo_when_requested(): + middle = _plan().layout(1) + ey = middle.component("Ey") + hy = middle.component("Hy") + + assert ey.global_origin == (3, 0, 0) + assert ey.global_to_local((4, 2, 3)) == (1, 2, 3) + assert ey.local_to_global((2, 1, 1)) == (5, 1, 1) + assert ey.global_to_local((7, 1, 1), include_halo=True) == (4, 1, 1) + with pytest.raises(IndexError, match="owned extent"): + ey.global_to_local((7, 1, 1)) + + assert hy.global_origin == (3, 0, 0) + assert hy.global_to_local((4, 2, 2)) == (1, 2, 2) + assert hy.global_to_local((3, 2, 2), include_halo=True) == (0, 2, 2) + assert hy.local_to_global((0, 2, 2), include_halo=True) == (3, 2, 2) + + +def test_only_outer_x_shards_own_x_physical_faces(): + first, middle, last = _plan().shard_layouts + + assert first.owns_physical_face("x", "low") + assert not first.owns_physical_face("x_high") + assert not middle.owns_physical_face("x_low") + assert not middle.owns_physical_face("x_high") + assert last.owns_physical_face("x", "high") + for shard in (first, middle, last): + assert shard.owns_physical_face("y_low") + assert shard.owns_physical_face("y", "high") + assert shard.owns_physical_face("z_low") + assert shard.owns_physical_face("z", "high") + + +def test_single_shard_plan_owns_both_x_faces_and_has_no_halos(): + plan = _plan(devices=("cuda:4",)) + shard = plan.layout_for_device("cuda:4") + + assert shard.cell_interval == (0, 10) + assert shard.owns_physical_face("x_low") + assert shard.owns_physical_face("x_high") + assert all(not layout.receive_halos for layout in shard.component_layouts) + assert all(not layout.send_regions for layout in shard.component_layouts) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"global_shape": (1, 4, 4), "devices": ("cuda:0",)}, ">= 2"), + ({"global_shape": (4, 4), "devices": ("cuda:0",)}, "global_shape"), + ( + {"global_shape": (3, 4, 4), "devices": ("cuda:0", "cuda:1", "cuda:2")}, + "Cannot partition", + ), + ( + { + "global_shape": (6, 4, 4), + "devices": ("cuda:0", "cuda:1", "cuda:2"), + "halo_width": 2, + }, + "halo_width", + ), + ], +) +def test_partition_rejects_illegal_extents(kwargs, message): + with pytest.raises(ValueError, match=message): + FDTDPartitionPlan(**kwargs) + + +def test_partition_and_nested_layouts_are_frozen(): + plan = _plan() + with pytest.raises(FrozenInstanceError): + plan.halo_width = 2 + with pytest.raises(FrozenInstanceError): + plan.layout(0).cell_interval = (1, 4) + with pytest.raises(FrozenInstanceError): + plan.layout(0).component("Ex").global_origin = (1, 0, 0) diff --git a/tests/fdtd/multi_gpu/test_public_api.py b/tests/fdtd/multi_gpu/test_public_api.py new file mode 100644 index 0000000..b383798 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_public_api.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import witwin.maxwell as mw +import witwin.maxwell.simulation as simulation_module +from witwin.maxwell.result import Result + + +def _scene() -> mw.Scene: + return mw.Scene( + domain=mw.Domain(bounds=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0))), + grid=mw.GridSpec.uniform(0.25), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + + +def _parallel(*, gather_fields: bool = False, result_device: str = "cuda:0"): + return mw.FDTDParallelConfig( + devices=("cuda:0", "cuda:1"), + gather_fields=gather_fields, + result_device=result_device, + ) + + +def test_parallel_config_is_public_and_forwarded_by_fdtd_factory(): + parallel = _parallel() + simulation = mw.Simulation.fdtd(_scene(), frequency=1.0e9, parallel=parallel) + + assert mw.FDTDParallelConfig is simulation_module.FDTDParallelConfig + assert simulation.config.parallel is parallel + + +def test_parallel_none_keeps_single_gpu_backend_path(monkeypatch): + scene = _scene() + prepared_scene = object() + calls = [] + + class FakeFDTD: + def __init__(self, prepared, **kwargs): + calls.append((prepared, kwargs)) + self.initialized = False + + def init_field(self): + self.initialized = True + + monkeypatch.setattr(simulation_module, "prepare_scene", lambda value: prepared_scene) + monkeypatch.setattr(simulation_module, "_require_cuda_scene", lambda *args, **kwargs: None) + monkeypatch.setattr(simulation_module, "_resolve_fdtd_backend", lambda: (FakeFDTD, object())) + + simulation = mw.Simulation.fdtd(scene, frequency=1.0e9) + solver = simulation._build_fdtd_solver_for_scene(scene, initialize=True) + + assert solver.initialized is True + assert calls == [ + ( + prepared_scene, + { + "frequency": 1.0e9, + "absorber_type": mw.AbsorberKind.CPML, + "cpml_config": {}, + }, + ) + ] + + +def test_parallel_builds_distributed_solver_and_initializes(monkeypatch): + import witwin.maxwell.fdtd.distributed as distributed_module + + scene = _scene() + parallel = _parallel() + calls = [] + + class FakeDistributedFDTD: + def __init__(self, logical_scene, **kwargs): + calls.append((logical_scene, kwargs)) + self.initialized = False + + def init_field(self): + self.initialized = True + + monkeypatch.setattr(simulation_module, "prepare_scene", lambda value: value) + monkeypatch.setattr(simulation_module, "_require_cuda_scene", lambda *args, **kwargs: None) + monkeypatch.setattr(distributed_module, "DistributedFDTD", FakeDistributedFDTD) + + simulation = mw.Simulation.fdtd(scene, frequency=2.0e9, parallel=parallel) + solver = simulation._build_fdtd_solver_for_scene(scene, initialize=True) + + assert solver.initialized is True + assert calls == [ + ( + scene, + { + "frequency": 2.0e9, + "parallel": parallel, + "absorber_type": mw.AbsorberKind.CPML, + "cpml_config": {}, + }, + ) + ] + + +@pytest.mark.parametrize("entrypoint", ["prepare", "run"]) +def test_trainable_parallel_is_rejected_before_solver_allocation(monkeypatch, entrypoint): + simulation = mw.Simulation.fdtd( + _scene(), + frequency=1.0e9, + parallel=_parallel(), + ) + simulation.has_trainable_parameters = True + + def unexpected_allocation(*args, **kwargs): + raise AssertionError("solver allocation must not be reached") + + monkeypatch.setattr(simulation, "_build_fdtd_solver", unexpected_allocation) + + with pytest.raises(ValueError, match="does not support trainable"): + getattr(simulation, entrypoint)() + + +def _fake_solved_distributed_solver(): + return SimpleNamespace( + scene=object(), + dt=0.125, + last_solve_elapsed_s=1.0, + dft_sample_counts=(0,), + observer_sample_counts=(0,), + _shutoff_triggered=False, + _shutoff_step=None, + parallel_stats={"devices": ("cuda:0", "cuda:1"), "halo_bytes_total": 64}, + ) + + +def test_gather_fields_false_does_not_fallback_to_shard_fields(monkeypatch): + simulation = mw.Simulation.fdtd( + _scene(), + frequency=1.0e9, + run_time=mw.TimeConfig(time_steps=2), + parallel=_parallel(gather_fields=False), + ) + solver = _fake_solved_distributed_solver() + monkeypatch.setattr( + simulation, + "_execute_fdtd_solve", + lambda *args: (None, 2, False, simulation.config.spectral_sampler), + ) + monkeypatch.setattr( + simulation, + "_fdtd_last_step_field_payload", + lambda *args: (_ for _ in ()).throw(AssertionError("field fallback must not run")), + ) + + result = simulation._run_fdtd_from_solver(solver) + + assert result.fields == {} + assert result.raw_output == {} + assert result.solver_stats["parallel_stats"]["halo_bytes_total"] == 64 + + +def test_gather_fields_true_requires_distributed_field_output(monkeypatch): + simulation = mw.Simulation.fdtd( + _scene(), + frequency=1.0e9, + parallel=_parallel(gather_fields=True), + ) + solver = _fake_solved_distributed_solver() + monkeypatch.setattr( + simulation, + "_execute_fdtd_solve", + lambda *args: (None, 1, False, simulation.config.spectral_sampler), + ) + + with pytest.raises(RuntimeError, match="did not return any output"): + simulation._run_fdtd_from_solver(solver) + + +def test_gathered_fields_are_normalized_on_explicit_result_device(monkeypatch): + parallel = _parallel(gather_fields=True, result_device="cuda:1") + simulation = mw.Simulation.fdtd(_scene(), frequency=1.0e9, parallel=parallel) + solver = _fake_solved_distributed_solver() + raw_fields = {"Ex": torch.tensor([1.0]), "Ey": torch.tensor([2.0]), "Ez": torch.tensor([3.0])} + requested_devices = [] + + monkeypatch.setattr( + simulation, + "_execute_fdtd_solve", + lambda *args: (raw_fields, 1, False, simulation.config.spectral_sampler), + ) + + def capture_device(fields, device): + requested_devices.append(str(device)) + return {name.upper(): value for name, value in fields.items()} + + monkeypatch.setattr(simulation_module, "_to_tensor_fields", capture_device) + + result = simulation._run_fdtd_from_solver(solver) + + assert requested_devices == ["cuda:1"] + assert set(result.fields) == {"EX", "EY", "EZ"} + + +def test_result_solver_stats_is_a_read_only_copy_property(): + result = Result( + method="fdtd", + scene=_scene(), + frequency=1.0e9, + solver_stats={"steps_run": 3}, + ) + + stats = result.solver_stats + stats["steps_run"] = 99 + + assert result.solver_stats == {"steps_run": 3} + with pytest.raises(AttributeError): + result.solver_stats = {} diff --git a/tests/fdtd/multi_gpu/test_result_persistence.py b/tests/fdtd/multi_gpu/test_result_persistence.py new file mode 100644 index 0000000..122c6fa --- /dev/null +++ b/tests/fdtd/multi_gpu/test_result_persistence.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import numpy as np +import torch + +import witwin.maxwell as mw +from witwin.maxwell.result import Result + + +_FREQUENCIES = (0.8e9, 1.2e9) + + +def _scene() -> mw.Scene: + return mw.Scene( + domain=mw.Domain(bounds=((-0.4, 0.4), (-0.2, 0.2), (-0.2, 0.2))), + grid=mw.GridSpec.uniform(0.1), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + + +def _parallel_stats(device: torch.device) -> dict: + return { + "parallel_stats": { + "devices": ("cuda:0", "cuda:1"), + "partition_extents": ((0, 4), (4, 8)), + "transport": "cuda_p2p", + "halo_bytes_total": 4096, + "communication_time_s": 0.01, + "compute_time_s": 0.04, + "peak_memory_bytes": {"cuda:0": 1024, "cuda:1": 1088}, + "device_counter": torch.tensor((7, 9), device=device), + } + } + + +def test_monitor_first_result_round_trip_preserves_order_and_stats_without_runtime_state( + tmp_path, + cuda_p2p_devices, + cuda_memory_cleanup, +): + coords = np.asarray((-0.3, -0.1, 0.1, 0.3), dtype=np.float64) + left = torch.tensor((1.0 + 2.0j, 3.0 + 4.0j), device=cuda_p2p_devices[0]) + right = torch.tensor((5.0 + 6.0j, 7.0 + 8.0j), device=cuda_p2p_devices[1]) + monitors = { + "right_declared_first": { + "kind": "point", + "monitor_type": "point", + "fields": ("Ez",), + "frequencies": _FREQUENCIES, + "data": right, + "components": {"Ez": right}, + }, + "plane_declared_second": { + "kind": "plane", + "monitor_type": "plane", + "fields": ("Ez",), + "frequencies": _FREQUENCIES, + "axis": "y", + "x": coords, + "z": np.asarray((-0.1, 0.1), dtype=np.float64), + "data": left[:, None, None].expand(-1, coords.size, 2), + }, + } + result = Result( + method="fdtd", + scene=_scene(), + frequencies=_FREQUENCIES, + fields={}, + monitors=monitors, + metadata={"rank_metric": right.real}, + solver_stats=_parallel_stats(cuda_p2p_devices[1]), + solver=lambda: None, + raw_output={"transport_handle": lambda: None}, + ) + + path = tmp_path / "monitor-first.pt" + result.save(path) + loaded = Result.load(path, scene=_scene()) + + assert loaded.fields == {} + assert tuple(loaded.monitors) == ( + "right_declared_first", + "plane_declared_second", + ) + torch.testing.assert_close(loaded.monitors["right_declared_first"]["data"], right.cpu()) + np.testing.assert_array_equal(loaded.monitors["plane_declared_second"]["x"], coords) + assert loaded.monitors["plane_declared_second"]["data"].device.type == "cpu" + stats = loaded.solver_stats["parallel_stats"] + assert stats["devices"] == ("cuda:0", "cuda:1") + assert stats["partition_extents"] == ((0, 4), (4, 8)) + assert stats["halo_bytes_total"] == 4096 + assert stats["device_counter"].device.type == "cpu" + torch.testing.assert_close(stats["device_counter"], torch.tensor((7, 9))) + assert loaded._metadata["rank_metric"].device.type == "cpu" + assert loaded.solver is None + assert loaded.raw_output is None + + +def test_gathered_field_result_round_trip_preserves_frequency_selection_on_cpu( + tmp_path, + cuda_p2p_devices, + cuda_memory_cleanup, +): + ex = torch.arange(48, device=cuda_p2p_devices[1], dtype=torch.float32).reshape( + 2, + 4, + 3, + 2, + ) + ez = (ex + 1.0).to(torch.complex64) + result = Result( + method="fdtd", + scene=_scene(), + frequencies=_FREQUENCIES, + fields={"EX": ex, "EZ": ez}, + monitors={}, + solver_stats=_parallel_stats(cuda_p2p_devices[0]), + ) + + path = tmp_path / "gathered-fields.pt" + result.save(path) + loaded = Result.load(path, scene=_scene()) + + assert loaded.frequencies == _FREQUENCIES + assert tuple(loaded.fields) == ("EX", "EZ") + assert all(tensor.device.type == "cpu" for tensor in loaded.fields.values()) + torch.testing.assert_close(loaded.tensor("EX", freq_index=1), ex[1].cpu()) + torch.testing.assert_close(loaded.tensor("EZ", frequency=_FREQUENCIES[0]), ez[0].cpu()) + assert loaded.stats()["num_fields"] == 2 + assert loaded.stats()["num_monitors"] == 0 diff --git a/tests/fdtd/multi_gpu/test_review_regressions.py b/tests/fdtd/multi_gpu/test_review_regressions.py new file mode 100644 index 0000000..114822b --- /dev/null +++ b/tests/fdtd/multi_gpu/test_review_regressions.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed import DistributedFDTD +from witwin.maxwell.fdtd.solver import FDTD +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig +from witwin.maxwell.scene import prepare_scene + + +_FREQUENCY = 1.0e9 +_FIELD_NAMES = ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") + + +def _parallel(devices, *, gather_fields: bool) -> FDTDParallelConfig: + return FDTDParallelConfig( + devices=devices, + transport="cuda_p2p", + overlap=True, + gather_fields=gather_fields, + result_device=devices[0], + ) + + +def _direct_distributed(scene, devices, *, gather_fields: bool) -> DistributedFDTD: + solver = DistributedFDTD( + scene, + frequency=_FREQUENCY, + parallel=_parallel(devices, gather_fields=gather_fields), + absorber_type="cpml", + ) + solver.init_field() + return solver + + +def _regular_scene() -> mw.Scene: + x = np.linspace(-0.4, 0.4, 9, dtype=np.float64) + y = np.linspace(-0.2, 0.2, 5, dtype=np.float64) + z = np.linspace(-0.2, 0.2, 5, dtype=np.float64) + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + scene.add_source( + mw.PointDipole( + position=(-0.2, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=mw.CW(frequency=_FREQUENCY, amplitude=1.0), + name="source", + ) + ) + return scene + + +def test_nanometer_scale_monitors_on_interface_sides_have_unique_owner( + cuda_p2p_devices, + cuda_memory_cleanup, +): + x = np.linspace(-4.0e-9, 4.0e-9, 9, dtype=np.float64) + y = np.linspace(-2.0e-9, 2.0e-9, 5, dtype=np.float64) + z = np.linspace(-2.0e-9, 2.0e-9, 5, dtype=np.float64) + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + scene.add_monitor( + mw.PointMonitor("left_of_interface", (-0.25e-9, 0.0, 0.0), fields=("Ez",)) + ) + scene.add_monitor( + mw.PointMonitor("right_of_interface", (0.25e-9, 0.0, 0.0), fields=("Ez",)) + ) + + solver = _direct_distributed(scene, cuda_p2p_devices, gather_fields=False) + owners = {monitor.name: [] for monitor in scene.monitors} + for shard in solver.shards: + for monitor in shard.solver.scene.monitors: + owners[monitor.name].append(shard.rank) + + assert owners == { + "left_of_interface": [0], + "right_of_interface": [1], + } + + +def test_monitor_output_keeps_scene_declaration_order_across_shards( + cuda_p2p_devices, + cuda_memory_cleanup, +): + scene = _regular_scene() + declared = ( + ("right_first", 0.25), + ("left_second", -0.25), + ("right_third", 0.35), + ("left_fourth", -0.35), + ) + for name, x in declared: + scene.add_monitor(mw.PointMonitor(name, (x, 0.0, 0.0), fields=("Ez",))) + + solver = _direct_distributed(scene, cuda_p2p_devices, gather_fields=False) + output = solver.solve( + time_steps=8, + dft_frequency=_FREQUENCY, + dft_window="none", + full_field_dft=False, + normalize_source=False, + use_cuda_graph=False, + ) + + assert tuple(output["observers"]) == tuple(name for name, _ in declared) + + +def test_direct_distributed_multifrequency_output_preserves_frequency_metadata( + cuda_p2p_devices, + cuda_memory_cleanup, +): + frequencies = (0.8e9, 1.2e9) + solver = _direct_distributed( + _regular_scene(), + cuda_p2p_devices, + gather_fields=True, + ) + output = solver.solve( + time_steps=12, + dft_frequency=frequencies, + dft_window="none", + full_field_dft=True, + normalize_source=False, + use_cuda_graph=False, + ) + + assert output["frequencies"] == frequencies + for name in ("Ex", "Ey", "Ez"): + assert output[name].shape[0] == len(frequencies) + + +def _ex_interface_scene(image_case: str) -> mw.Scene: + x = np.linspace(-0.5, 0.5, 11, dtype=np.float64) + z = np.linspace(-0.2, 0.2, 5, dtype=np.float64) + if image_case == "y_symmetry": + y = np.linspace(0.0, 0.4, 5, dtype=np.float64) + boundary = mw.BoundarySpec.none() + symmetry = (None, ("PMC", "low"), None) + source_y = 0.0 + elif image_case == "y_periodic": + y = np.linspace(-0.2, 0.2, 5, dtype=np.float64) + boundary = mw.BoundarySpec.faces(default="none", y="periodic") + symmetry = None + source_y = -0.19 + else: + raise ValueError(f"Unknown image case {image_case!r}.") + + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=boundary, + symmetry=symmetry, + device="cuda:0", + ) + scene.add_source( + mw.PointDipole( + # The first slab owns x in [-0.5, 0); its final Ex sample is -0.05. + position=(-0.05, source_y, 0.0), + polarization="Ex", + profile="ideal", + source_time=mw.CW(frequency=_FREQUENCY, amplitude=1.0), + name="interface_adjacent_ex", + ) + ) + return scene + + +def _distributed_component(solver, output, name: str) -> torch.Tensor: + if name in output: + return output[name] + local = tuple(getattr(shard.solver, name) for shard in solver.shards) + return solver._gather_component(name, local) + + +@pytest.mark.parametrize("image_case", ("y_symmetry", "y_periodic")) +def test_interface_adjacent_ideal_ex_source_image_parity( + image_case, + cuda_p2p_devices, + cuda_memory_cleanup, +): + scene = _ex_interface_scene(image_case) + single = FDTD(prepare_scene(scene), frequency=_FREQUENCY, absorber_type="cpml") + single.init_field() + single.solve( + time_steps=24, + dft_frequency=_FREQUENCY, + dft_window="none", + full_field_dft=False, + normalize_source=False, + use_cuda_graph=False, + ) + + distributed = _direct_distributed(scene, cuda_p2p_devices, gather_fields=True) + output = distributed.solve( + time_steps=24, + dft_frequency=_FREQUENCY, + dft_window="none", + full_field_dft=False, + normalize_source=False, + use_cuda_graph=False, + ) + + assert float(torch.abs(single.Ex).max().item()) > 0.0 + for name in _FIELD_NAMES: + torch.testing.assert_close( + _distributed_component(distributed, output, name).to("cuda:0"), + getattr(single, name).to("cuda:0"), + rtol=2.0e-5, + atol=2.0e-6, + ) diff --git a/tests/fdtd/multi_gpu/test_sharded_export.py b/tests/fdtd/multi_gpu/test_sharded_export.py new file mode 100644 index 0000000..498b3d9 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_sharded_export.py @@ -0,0 +1,184 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from witwin.maxwell.fdtd.distributed.persistence import ( + export_distributed_field_shards, +) + + +_SENTINEL = -98765.0 +_FREQUENCIES = (0.8e9, 1.2e9) +_LAYOUTS = ( + SimpleNamespace( + storage_cell_owned=slice(0, 3), + global_cell_owned=slice(0, 3), + storage_node_owned=slice(0, 3), + global_node_owned=slice(0, 3), + ), + SimpleNamespace( + storage_cell_owned=slice(1, 3), + global_cell_owned=slice(3, 5), + storage_node_owned=slice(1, 4), + global_node_owned=slice(3, 6), + ), +) + + +class _DftSolver: + dft_enabled = True + + def __init__(self, fields, frequencies): + self._fields = fields + self._frequencies = frequencies + self.calls = [] + + def get_frequency_solution(self, *, all_frequencies=False): + self.calls.append(all_frequencies) + return {**self._fields, "frequencies": self._frequencies} + + +def _global_fields(*, frequency_count=None): + prefix = () if frequency_count is None else (frequency_count,) + return { + "Ex": torch.arange(5 * 2 * 2 * (frequency_count or 1), dtype=torch.float32) + .reshape(prefix + (5, 2, 2)), + "Ey": ( + 1000 + + torch.arange(6 * 2 * 2 * (frequency_count or 1), dtype=torch.float32) + ).reshape(prefix + (6, 2, 2)), + "Ez": ( + 2000 + + torch.arange(6 * 2 * 2 * (frequency_count or 1), dtype=torch.float32) + ).reshape(prefix + (6, 2, 2)), + } + + +def _with_ghosts(field, *, storage_slice, global_slice, local_x): + shape = list(field.shape) + x_axis = field.ndim - 3 + shape[x_axis] = local_x + local = torch.full(shape, _SENTINEL, dtype=field.dtype) + source = [slice(None)] * field.ndim + destination = [slice(None)] * field.ndim + source[x_axis] = global_slice + destination[x_axis] = storage_slice + local[tuple(destination)] = field[tuple(source)] + return local + + +def _local_fields(global_fields, rank): + layout = _LAYOUTS[rank] + result = {} + for name, field in global_fields.items(): + is_cell = name == "Ex" + storage = layout.storage_cell_owned if is_cell else layout.storage_node_owned + global_slice = layout.global_cell_owned if is_cell else layout.global_node_owned + local_x = 4 if rank == 0 or not is_cell else 3 + result[name] = _with_ghosts( + field, + storage_slice=storage, + global_slice=global_slice, + local_x=local_x, + ) + return result + + +def _distributed_solver(global_fields, *, dft=True, frequencies=_FREQUENCIES): + shards = [] + local_solvers = [] + for rank in range(2): + fields = _local_fields(global_fields, rank) + if dft: + metadata = ( + torch.tensor(frequencies, dtype=torch.float64) + if rank == 0 + else frequencies + ) + local_solver = _DftSolver(fields, metadata) + else: + local_solver = SimpleNamespace(dft_enabled=False, **fields) + local_solvers.append(local_solver) + shards.append( + SimpleNamespace( + rank=rank, + device=f"cuda:{rank}", + layout=_LAYOUTS[rank], + solver=local_solver, + ) + ) + root = SimpleNamespace(frequency=1.5e9, shards=tuple(reversed(shards))) + return root, tuple(local_solvers) + + +def _assert_owned_fields(artifacts, global_fields, frequencies): + assert tuple(artifact.rank for artifact in artifacts) == (0, 1) + assert all(artifact.frequencies == frequencies for artifact in artifacts) + for name, expected in global_fields.items(): + components = [ + next(component for component in artifact.components if component.name == name) + for artifact in artifacts + ] + x_axis = expected.ndim - 3 + gathered = torch.cat([component.tensor for component in components], dim=x_axis) + torch.testing.assert_close(gathered, expected) + assert all(component.tensor.is_contiguous() for component in components) + assert all(not torch.any(component.tensor == _SENTINEL) for component in components) + + +def test_export_last_step_fields_crops_halos_and_uses_yee_owned_slices(): + fields = _global_fields() + solver, _locals = _distributed_solver(fields, dft=False) + + artifacts = export_distributed_field_shards(solver) + + _assert_owned_fields(artifacts, fields, (solver.frequency,)) + assert tuple(artifact.components[0].global_x_slice for artifact in artifacts) == ( + (0, 3), + (3, 5), + ) + assert tuple(artifact.components[1].global_x_slice for artifact in artifacts) == ( + (0, 3), + (3, 6), + ) + assert tuple(artifact.components[2].global_x_slice for artifact in artifacts) == ( + (0, 3), + (3, 6), + ) + + +def test_export_multifrequency_dft_crops_x_axis_after_frequency_axis(): + fields = _global_fields(frequency_count=len(_FREQUENCIES)) + solver, local_solvers = _distributed_solver(fields) + + artifacts = export_distributed_field_shards(solver) + + _assert_owned_fields(artifacts, fields, _FREQUENCIES) + assert all( + component.x_axis == 1 + for artifact in artifacts + for component in artifact.components + ) + assert all(local_solver.calls == [True] for local_solver in local_solvers) + + +def test_export_rejects_mixed_dft_and_last_step_shards_before_reading_fields(): + fields = _global_fields(frequency_count=len(_FREQUENCIES)) + solver, local_solvers = _distributed_solver(fields) + local_solvers[1].dft_enabled = False + + with pytest.raises(RuntimeError, match="mixes DFT and last-step"): + export_distributed_field_shards(solver) + assert all(local_solver.calls == [] for local_solver in local_solvers) + + +def test_export_rejects_inconsistent_dft_frequencies(): + fields = _global_fields(frequency_count=len(_FREQUENCIES)) + solver, local_solvers = _distributed_solver(fields) + local_solvers[1]._frequencies = (0.8e9, 1.3e9) + + with pytest.raises(RuntimeError, match=r"rank 1 DFT frequencies.*do not match"): + export_distributed_field_shards(solver) diff --git a/tests/fdtd/multi_gpu/test_sharded_persistence.py b/tests/fdtd/multi_gpu/test_sharded_persistence.py new file mode 100644 index 0000000..9b298f7 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_sharded_persistence.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed.persistence import ( + FieldComponentArtifact, + FieldShardArtifact, +) +from witwin.maxwell.result import Result + + +_FREQUENCIES = (0.8e9, 1.2e9) + + +def _scene() -> mw.Scene: + return mw.Scene( + domain=mw.Domain(bounds=((-0.5, 0.5), (-0.2, 0.2), (-0.2, 0.2))), + grid=mw.GridSpec.uniform(0.1), + boundary=mw.BoundarySpec.none(), + device="cpu", + ) + + +def _global_fields() -> dict[str, torch.Tensor]: + return { + "Ex": torch.arange(2 * 9 * 4 * 5, dtype=torch.float32).reshape(2, 9, 4, 5), + "Ey": ( + 1000 + + torch.arange(2 * 10 * 3 * 5, dtype=torch.float32).reshape(2, 10, 3, 5) + ), + "Ez": ( + 2000 + + torch.arange(2 * 10 * 4 * 4, dtype=torch.float32).reshape(2, 10, 4, 4) + ), + } + + +def _artifacts( + fields: dict[str, torch.Tensor], + *, + order=(2, 0, 1), +) -> tuple[FieldShardArtifact, ...]: + cell_slices = ((0, 4), (4, 7), (7, 9)) + node_slices = ((0, 4), (4, 7), (7, 10)) + by_rank = [] + for rank in range(3): + components = [] + for name in ("Ex", "Ey", "Ez"): + extent = cell_slices[rank] if name == "Ex" else node_slices[rank] + components.append( + FieldComponentArtifact( + name=name, + tensor=fields[name][:, extent[0] : extent[1]], + global_x_slice=extent, + ) + ) + by_rank.append( + FieldShardArtifact( + rank=rank, + device=f"cuda:{rank % 2}", + frequencies=_FREQUENCIES, + components=tuple(components), + ) + ) + return tuple(by_rank[rank] for rank in order) + + +def _result(artifacts, *, fields=None) -> Result: + solver = SimpleNamespace(export_field_shards=lambda: artifacts) + monitor_data = torch.tensor((1.0 + 2.0j, 3.0 + 4.0j)) + return Result( + method="fdtd", + scene=_scene(), + frequencies=_FREQUENCIES, + solver=solver, + fields=fields or {}, + monitors={ + "probe": { + "kind": "point", + "fields": ("Ez",), + "frequencies": _FREQUENCIES, + "data": monitor_data, + } + }, + metadata={"run": "uneven-multifrequency"}, + solver_stats={"parallel_stats": {"devices": ("cuda:0", "cuda:1")}}, + ) + + +def _save_fixture(tmp_path): + fields = _global_fields() + directory = tmp_path / "sharded-result" + manifest = _result( + _artifacts(fields), + fields={name.upper(): tensor for name, tensor in fields.items()}, + ).save_sharded(directory) + return directory, manifest, fields + + +def test_sharded_round_trip_sorts_ranks_and_gathers_uneven_multifrequency_fields( + tmp_path, +): + directory, manifest, fields = _save_fixture(tmp_path) + + assert tuple(path.name for path in sorted(directory.iterdir())) == ( + "manifest.json", + "rank-0000.pt", + "rank-0001.pt", + "rank-0002.pt", + "result.pt", + ) + assert tuple(shard.rank for shard in manifest.shards) == (0, 1, 2) + assert tuple(component.shape for component in manifest.components) == ( + (2, 9, 4, 5), + (2, 10, 3, 5), + (2, 10, 4, 4), + ) + assert all(component.x_axis == 1 for component in manifest.components) + + encoded = json.loads((directory / "manifest.json").read_text(encoding="utf-8")) + assert [shard["rank"] for shard in encoded["shards"]] == [0, 1, 2] + assert [ + shard["components"][0]["global_x_slice"] for shard in encoded["shards"] + ] == [[0, 4], [4, 7], [7, 9]] + assert [ + shard["components"][1]["global_x_slice"] for shard in encoded["shards"] + ] == [[0, 4], [4, 7], [7, 10]] + + rank = torch.load(directory / "rank-0001.pt", weights_only=True) + ex = rank["components"]["Ex"] + assert ex["tensor"].device.type == "cpu" + assert ex["dtype"] == "torch.float32" + assert ex["shape"] == (2, 3, 4, 5) + assert ex["frequencies"] == _FREQUENCIES + assert ex["x_axis"] == 1 + + lazy = Result.load_sharded(directory, scene=_scene()) + assert lazy.is_sharded is True + assert lazy.fields == {} + assert tuple(path.name for path in lazy.shard_paths) == ( + "rank-0000.pt", + "rank-0001.pt", + "rank-0002.pt", + ) + assert tuple(shard.rank for shard in lazy.sharded_manifest.shards) == (0, 1, 2) + assert tuple(lazy.monitors) == ("probe",) + + gathered = Result.load_sharded( + directory, + scene=_scene(), + gather_fields=True, + map_location="cpu", + ) + assert tuple(gathered.fields) == ("EX", "EY", "EZ") + for name, reference in fields.items(): + torch.testing.assert_close(gathered.fields[name.upper()], reference) + torch.testing.assert_close(gathered.tensor("EX", freq_index=1), fields["Ex"][1]) + assert gathered.solver_stats["parallel_stats"]["devices"] == ( + "cuda:0", + "cuda:1", + ) + + +def test_lazy_load_does_not_deserialize_existing_corrupt_rank_until_gather(tmp_path): + directory, _manifest, _fields = _save_fixture(tmp_path) + corrupt = directory / "rank-0001.pt" + corrupt.write_bytes(b"not a torch shard") + + lazy = Result.load_sharded(directory, scene=_scene(), gather_fields=False) + assert lazy.fields == {} + assert lazy.shard_paths[1] == corrupt + + with pytest.raises(ValueError, match=r"rank 1 file.*Failed|rank 1 file"): + Result.load_sharded(directory, scene=_scene(), gather_fields=True) + + +def test_missing_rank_is_reported_before_lazy_result_is_returned(tmp_path): + directory, _manifest, _fields = _save_fixture(tmp_path) + missing = directory / "rank-0002.pt" + missing.unlink() + + with pytest.raises(FileNotFoundError, match=r"rank 2 file is missing"): + Result.load_sharded(directory, scene=_scene()) + + +def test_corrupt_manifest_is_rejected(tmp_path): + directory, _manifest, _fields = _save_fixture(tmp_path) + (directory / "manifest.json").write_text("{not-json", encoding="utf-8") + + with pytest.raises(ValueError, match="Invalid sharded Result manifest"): + Result.load_sharded(directory, scene=_scene()) + + +def test_save_sharded_requires_export_field_shards_provider(tmp_path): + result = Result( + method="fdtd", + scene=_scene(), + frequency=1.0e9, + solver=object(), + ) + + with pytest.raises(RuntimeError, match=r"requires solver\.export_field_shards"): + result.save_sharded(tmp_path / "missing-exporter") + + +def test_noncontiguous_owned_intervals_are_rejected_before_publication(tmp_path): + fields = _global_fields() + artifacts = list(_artifacts(fields, order=(0, 1, 2))) + broken = artifacts[1] + components = list(broken.components) + components[0] = FieldComponentArtifact( + name="Ex", + tensor=fields["Ex"][:, 5:8], + global_x_slice=(5, 8), + ) + artifacts[1] = FieldShardArtifact( + rank=broken.rank, + device=broken.device, + frequencies=broken.frequencies, + components=tuple(components), + ) + + directory = tmp_path / "gap" + with pytest.raises(ValueError, match="leaves a gap"): + _result(tuple(artifacts)).save_sharded(directory) + assert directory.exists() is False + + +def test_rank_tensors_are_detached_and_written_on_cpu( + tmp_path, + cuda_p2p_devices, + cuda_memory_cleanup, +): + global_ex = torch.arange( + 2 * 6 * 2 * 2, + dtype=torch.float32, + device=cuda_p2p_devices[0], + ).reshape(2, 6, 2, 2) + artifacts = ( + { + "rank": 1, + "device": "cuda:1", + "frequencies": _FREQUENCIES, + "components": { + "Ex": { + "tensor": global_ex[:, 3:].to(cuda_p2p_devices[1]), + "global_x_slice": (3, 6), + } + }, + }, + { + "rank": 0, + "device": "cuda:0", + "frequencies": _FREQUENCIES, + "components": { + "Ex": { + "tensor": global_ex[:, :3].requires_grad_(), + "global_x_slice": (0, 3), + } + }, + }, + ) + directory = tmp_path / "cuda-shards" + + _result(artifacts).save_sharded(directory) + + for rank in range(2): + payload = torch.load(directory / f"rank-{rank:04d}.pt", weights_only=True) + tensor = payload["components"]["Ex"]["tensor"] + assert tensor.device.type == "cpu" + assert tensor.requires_grad is False + loaded = Result.load_sharded(directory, scene=_scene(), gather_fields=True) + torch.testing.assert_close(loaded.fields["EX"], global_ex.cpu()) diff --git a/tests/fdtd/multi_gpu/test_solver_acceptance.py b/tests/fdtd/multi_gpu/test_solver_acceptance.py new file mode 100644 index 0000000..c1f6651 --- /dev/null +++ b/tests/fdtd/multi_gpu/test_solver_acceptance.py @@ -0,0 +1,390 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pytest +import torch + +import witwin.maxwell as mw +from witwin.maxwell.fdtd.distributed import DistributedFDTD +from witwin.maxwell.fdtd.solver import FDTD +from witwin.maxwell.fdtd_parallel import FDTDParallelConfig +from witwin.maxwell.scene import prepare_scene + + +_FREQUENCY = 1.0e9 +_TIME_STEPS = 24 +_FIELD_NAMES = ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") + + +@dataclass +class _SolvePair: + single_solver: FDTD + single_output: dict + distributed_solver: DistributedFDTD + distributed_output: dict + + +def _grid_nodes(*, nonuniform_x: bool): + if nonuniform_x: + # Eleven uneven x cells give a 6/5 partition. The finer, replicated y/z + # spacing keeps the shard-local CFL time steps exactly equal. + x = np.asarray( + [-0.55, -0.44, -0.33, -0.21, -0.10, 0.00, 0.09, 0.19, 0.30, 0.42, 0.53, 0.65], + dtype=np.float64, + ) + else: + x = np.linspace(-0.5, 0.5, 11, dtype=np.float64) + y = np.linspace(-0.2, 0.2, 9, dtype=np.float64) + z = np.linspace(-0.2, 0.2, 9, dtype=np.float64) + return x, y, z + + +def _vacuum_scene(*, nonuniform_x: bool, source_on_interface: bool) -> mw.Scene: + x, y, z = _grid_nodes(nonuniform_x=nonuniform_x) + interface_node = (len(x) - 1 + 1) // 2 + source_x = float(x[interface_node] if source_on_interface else x[2]) + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=mw.BoundarySpec.none(), + device="cuda:0", + ) + scene.add_source( + mw.PointDipole( + position=(source_x, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=mw.CW( + frequency=_FREQUENCY, + amplitude=1.0, + ), + name="interface_source" if source_on_interface else "source", + ) + ) + scene.add_monitor( + mw.PointMonitor( + "probe", + (source_x, 0.0, 0.0), + fields=("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"), + ) + ) + return scene + + +def _run_single(scene: mw.Scene): + solver = FDTD(prepare_scene(scene), frequency=_FREQUENCY, absorber_type="cpml") + solver.init_field() + output = solver.solve( + time_steps=_TIME_STEPS, + dft_frequency=_FREQUENCY, + dft_window="none", + full_field_dft=False, + normalize_source=False, + use_cuda_graph=False, + ) + return solver, output + + +def _run_distributed( + scene: mw.Scene, + devices: tuple[torch.device, torch.device], + *, + gather_fields: bool, +): + config = FDTDParallelConfig( + devices=devices, + transport="cuda_p2p", + overlap=True, + gather_fields=gather_fields, + result_device=devices[0], + ) + solver = DistributedFDTD( + scene, + frequency=_FREQUENCY, + parallel=config, + absorber_type="cpml", + ) + solver.init_field() + output = solver.solve( + time_steps=_TIME_STEPS, + dft_frequency=_FREQUENCY, + dft_window="none", + full_field_dft=False, + normalize_source=False, + use_cuda_graph=False, + ) + return solver, output + + +def _assert_step_field_close(actual: torch.Tensor, reference: torch.Tensor) -> None: + actual = actual.to(reference.device) + assert torch.isfinite(actual).all() + assert torch.isfinite(reference).all() + absolute = torch.abs(actual - reference) + assert float(absolute.max().item()) <= 2.0e-6 + + reference_scale = float(torch.abs(reference).max().item()) + significant = torch.abs(reference) >= max(1.0e-6, 1.0e-4 * reference_scale) + if bool(significant.any()): + relative = absolute[significant] / torch.abs(reference[significant]) + assert float(relative.max().item()) <= 2.0e-5 + + +def _distributed_step_field(pair: _SolvePair, name: str) -> torch.Tensor: + if name in pair.distributed_output: + return pair.distributed_output[name] + local_values = tuple( + getattr(shard.solver, name) for shard in pair.distributed_solver.shards + ) + return pair.distributed_solver._gather_component(name, local_values) + + +def _monitor_component(payload: dict, name: str) -> torch.Tensor: + value = payload["components"][name] + if isinstance(value, dict): + value = value["data"] + return torch.as_tensor(value) + + +@pytest.fixture(scope="module") +def uniform_vacuum_pair(cuda_p2p_devices): + scene = _vacuum_scene(nonuniform_x=False, source_on_interface=False) + single_solver, single_output = _run_single(scene) + distributed_solver, distributed_output = _run_distributed( + scene, cuda_p2p_devices, gather_fields=True + ) + yield _SolvePair(single_solver, single_output, distributed_solver, distributed_output) + del single_solver, single_output, distributed_solver, distributed_output + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + +@pytest.fixture(scope="module") +def uneven_interface_pair(cuda_p2p_devices): + scene = _vacuum_scene(nonuniform_x=True, source_on_interface=True) + single_solver, single_output = _run_single(scene) + distributed_solver, distributed_output = _run_distributed( + scene, cuda_p2p_devices, gather_fields=True + ) + yield _SolvePair(single_solver, single_output, distributed_solver, distributed_output) + del single_solver, single_output, distributed_solver, distributed_output + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("fixture_name", ["uniform_vacuum_pair", "uneven_interface_pair"]) +def test_two_gpu_step_fields_match_one_gpu(request, fixture_name): + pair: _SolvePair = request.getfixturevalue(fixture_name) + for name in _FIELD_NAMES: + _assert_step_field_close( + _distributed_step_field(pair, name), + getattr(pair.single_solver, name), + ) + + +@pytest.mark.parametrize("fixture_name", ["uniform_vacuum_pair", "uneven_interface_pair"]) +def test_two_gpu_point_monitor_matches_one_gpu(request, fixture_name): + pair: _SolvePair = request.getfixturevalue(fixture_name) + single = pair.single_output["observers"]["probe"] + distributed = pair.distributed_output["observers"]["probe"] + for name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"): + torch.testing.assert_close( + _monitor_component(distributed, name).to("cuda:0"), + _monitor_component(single, name).to("cuda:0"), + rtol=5.0e-5, + atol=5.0e-6, + ) + + +def test_gather_false_returns_monitors_without_global_fields_and_reports_stats( + cuda_p2p_devices, + cuda_memory_cleanup, +): + scene = _vacuum_scene(nonuniform_x=False, source_on_interface=True) + solver, output = _run_distributed(scene, cuda_p2p_devices, gather_fields=False) + + assert output is not None + assert set(output) == {"observers"} + assert "probe" in output["observers"] + stats = solver.parallel_stats + assert stats["devices"] == ("cuda:0", "cuda:1") + assert stats["transport"] == "cuda_p2p" + assert stats["gather_fields"] is False + assert stats["result_device"] == "cuda:0" + assert stats["overlap_requested"] is True + assert stats["overlap_active"] is True + assert stats["halo_bytes_per_step"] > 0 + assert stats["halo_bytes_total"] == stats["halo_bytes_per_step"] * _TIME_STEPS + assert set(stats["peak_memory_bytes"]) == {"cuda:0", "cuda:1"} + assert all(value > 0 for value in stats["peak_memory_bytes"].values()) + assert stats["topology"]["neighbor_pairs"][0]["peer_left_to_right"] is True + assert stats["topology"]["neighbor_pairs"][0]["peer_right_to_left"] is True + + +def test_gather_true_places_owned_global_fields_on_requested_result_device( + uniform_vacuum_pair, +): + pair = uniform_vacuum_pair + output = pair.distributed_output + solver = pair.distributed_solver + + expected_shapes = { + "Ex": (solver.Nx - 1, solver.Ny, solver.Nz), + "Ey": (solver.Nx, solver.Ny - 1, solver.Nz), + "Ez": (solver.Nx, solver.Ny, solver.Nz - 1), + } + for name, shape in expected_shapes.items(): + assert tuple(output[name].shape) == shape + assert output[name].device == torch.device("cuda:0") + assert solver.parallel_stats["gather_fields"] is True + assert len(solver.parallel_stats["partitions"]) == 2 + assert solver.parallel_stats["partitions"][0]["physical_cells"][0] == 0 + assert ( + solver.parallel_stats["partitions"][-1]["physical_cells"][1] + == solver.partition_plan.physical_cell_count + ) + + +def _feature_scene(case: str) -> mw.Scene: + x = np.linspace(-0.6, 0.6, 13, dtype=np.float64) + y = np.linspace(-0.4, 0.4, 9, dtype=np.float64) + z = np.linspace(-0.4, 0.4, 9, dtype=np.float64) + if case == "xyz_cpml_dielectric": + boundary = mw.BoundarySpec.pml(num_layers=2, strength=1.0e6) + elif case == "mixed_pec_pmc_mur": + boundary = mw.BoundarySpec.faces( + default="none", + x=("pec", "pmc"), + y=("mur", "mur"), + z=("pec", "mur"), + ) + else: + boundary = mw.BoundarySpec.none() + + scene = mw.Scene( + domain=mw.Domain( + bounds=( + (float(x[0]), float(x[-1])), + (float(y[0]), float(y[-1])), + (float(z[0]), float(z[-1])), + ) + ), + grid=mw.GridSpec.custom(x, y, z), + boundary=boundary, + device="cuda:0", + ) + if case == "xyz_cpml_dielectric": + scene.add_structure( + mw.Structure( + name="interface_dielectric", + geometry=mw.Box(position=(0.0, 0.0, 0.0), size=(0.5, 0.4, 0.4)), + material=mw.Material(eps_r=3.0), + ) + ) + elif case == "conductive_electric_dispersion": + scene.add_structure( + mw.Structure( + name="lossy_lorentz_interface", + geometry=mw.Box(position=(0.0, 0.0, 0.0), size=(0.5, 0.4, 0.4)), + material=mw.Material( + eps_r=2.0, + sigma_e=0.02, + lorentz_poles=( + mw.LorentzPole( + delta_eps=0.5, + resonance_frequency=2.0e9, + gamma=2.0e8, + ), + ), + ), + ) + ) + + if case == "uniform_current_interface": + scene.add_source( + mw.UniformCurrentSource( + size=(0.5, 0.3, 0.3), + center=(0.0, 0.0, 0.0), + polarization="Ez", + source_time=mw.CW( + frequency=_FREQUENCY, + amplitude=1.0, + ), + name="cross_interface_current", + ) + ) + else: + scene.add_source( + mw.PointDipole( + position=(-0.3, 0.0, 0.0), + polarization="Ez", + profile="ideal", + source_time=mw.CW( + frequency=_FREQUENCY, + amplitude=1.0, + ), + name="source", + ) + ) + scene.add_monitor( + mw.PointMonitor("probe", (0.0, 0.0, 0.0), fields=_FIELD_NAMES) + ) + return scene + + +@pytest.mark.parametrize( + "case", + ( + "xyz_cpml_dielectric", + "conductive_electric_dispersion", + "mixed_pec_pmc_mur", + "uniform_current_interface", + ), +) +def test_two_gpu_feature_matrix_matches_one_gpu( + case, + cuda_p2p_devices, + cuda_memory_cleanup, +): + scene = _feature_scene(case) + single_solver, single_output = _run_single(scene) + distributed_solver, distributed_output = _run_distributed( + scene, + cuda_p2p_devices, + gather_fields=True, + ) + pair = _SolvePair( + single_solver, + single_output, + distributed_solver, + distributed_output, + ) + + for name in _FIELD_NAMES: + _assert_step_field_close( + _distributed_step_field(pair, name), + getattr(single_solver, name), + ) + torch.testing.assert_close( + _monitor_component(distributed_output["observers"]["probe"], name).to("cuda:0"), + _monitor_component(single_output["observers"]["probe"], name).to("cuda:0"), + rtol=5.0e-5, + atol=5.0e-6, + ) + + if case == "xyz_cpml_dielectric": + assert all(shard.solver.uses_cpml for shard in distributed_solver.shards) + for shard in distributed_solver.shards: + assert shard.solver.scene.boundary.uses_kind("pml") + if case == "uniform_current_interface": + assert all(shard.solver._source_terms for shard in distributed_solver.shards) diff --git a/tests/fdtd/multi_gpu/test_transport.py b/tests/fdtd/multi_gpu/test_transport.py new file mode 100644 index 0000000..4a999af --- /dev/null +++ b/tests/fdtd/multi_gpu/test_transport.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from witwin.maxwell.fdtd.distributed.transport import CudaP2PHaloTransport +from witwin.maxwell.fdtd_parallel import FDTDPartitionPlan + + +_SENTINEL = -91.0 + + +def _make_transport_shards(devices: tuple[torch.device, torch.device]): + plan = FDTDPartitionPlan(global_shape=(10, 7, 6), devices=devices) + shards = [] + for layout in plan.shard_layouts: + device = torch.device(layout.device) + with torch.cuda.device(device): + solver = SimpleNamespace( + Ey=torch.full( + layout.component("Ey").local_shape, + _SENTINEL, + device=device, + dtype=torch.float32, + ), + Ez=torch.full( + layout.component("Ez").local_shape, + _SENTINEL, + device=device, + dtype=torch.float32, + ), + Hy=torch.full( + layout.component("Hy").local_shape, + _SENTINEL, + device=device, + dtype=torch.float32, + ), + Hz=torch.full( + layout.component("Hz").local_shape, + _SENTINEL, + device=device, + dtype=torch.float32, + ), + ) + compute_stream = torch.cuda.Stream(device=device) + communication_stream = torch.cuda.Stream(device=device, priority=-1) + shards.append( + SimpleNamespace( + rank=layout.rank, + device=device, + layout=layout, + solver=solver, + compute_stream=compute_stream, + communication_stream=communication_stream, + electric_ready=torch.cuda.Event(), + electric_received=torch.cuda.Event(), + magnetic_ready=torch.cuda.Event(), + magnetic_received=torch.cuda.Event(), + halo_hy_low=solver.Hy[0] if layout.rank > 0 else None, + halo_hz_low=solver.Hz[0] if layout.rank > 0 else None, + ) + ) + return plan, tuple(shards) + + +def _tag_plane(plane: torch.Tensor, *, offset: float) -> torch.Tensor: + values = torch.arange(plane.numel(), device=plane.device, dtype=plane.dtype) + return (values.reshape_as(plane) + offset).contiguous() + + +def test_cuda_p2p_transport_round_trips_tagged_asymmetric_y_halos( + cuda_p2p_devices, + cuda_memory_cleanup, +): + transport = CudaP2PHaloTransport(cuda_p2p_devices) + transport.preflight() + plan, shards = _make_transport_shards(cuda_p2p_devices) + left, right = shards + + right_node = right.layout.storage_node_owned.start + left_ghost = left.layout.storage_node_owned.stop + ey_tag = _tag_plane(right.solver.Ey[right_node], offset=1_000.0) + ez_tag = _tag_plane(right.solver.Ez[right_node], offset=2_000.0) + right.solver.Ey[right_node].copy_(ey_tag) + right.solver.Ez[right_node].copy_(ez_tag) + untouched_left_ey = left.solver.Ey[left_ghost - 1].clone() + + transport.exchange_electric(shards) + left.electric_received.synchronize() + + torch.testing.assert_close( + left.solver.Ey[left_ghost], ey_tag.to(left.device), rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + left.solver.Ez[left_ghost], ez_tag.to(left.device), rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + left.solver.Ey[left_ghost - 1], untouched_left_ey, rtol=0.0, atol=0.0 + ) + + left_last = left.layout.storage_cell_owned.stop - 1 + hy_tag = _tag_plane(left.solver.Hy[left_last], offset=3_000.0) + hz_tag = _tag_plane(left.solver.Hz[left_last], offset=4_000.0) + left.solver.Hy[left_last].copy_(hy_tag) + left.solver.Hz[left_last].copy_(hz_tag) + untouched_right_hy = right.solver.Hy[1].clone() + + transport.exchange_magnetic(shards) + right.magnetic_received.synchronize() + + torch.testing.assert_close(right.halo_hy_low, hy_tag.to(right.device), rtol=0.0, atol=0.0) + torch.testing.assert_close(right.halo_hz_low, hz_tag.to(right.device), rtol=0.0, atol=0.0) + torch.testing.assert_close(right.solver.Hy[1], untouched_right_hy, rtol=0.0, atol=0.0) + + assert plan.layout(0).component("Ey").high_halo is not None + assert plan.layout(1).component("Hy").low_halo is not None + + +def test_cuda_p2p_transport_reuses_storage_without_torch_allocator_growth( + cuda_p2p_devices, + cuda_memory_cleanup, +): + transport = CudaP2PHaloTransport(cuda_p2p_devices) + transport.preflight() + _, shards = _make_transport_shards(cuda_p2p_devices) + left, right = shards + + # Warm up peer mappings and event dependencies before taking allocator snapshots. + transport.exchange_electric(shards) + transport.exchange_magnetic(shards) + left.electric_received.synchronize() + right.magnetic_received.synchronize() + + pointers = ( + left.solver.Ey.data_ptr(), + left.solver.Ez.data_ptr(), + right.solver.Hy.data_ptr(), + right.solver.Hz.data_ptr(), + ) + allocated_before = tuple(torch.cuda.memory_allocated(device) for device in cuda_p2p_devices) + for _ in range(32): + transport.exchange_electric(shards) + transport.exchange_magnetic(shards) + left.electric_received.synchronize() + right.magnetic_received.synchronize() + allocated_after = tuple(torch.cuda.memory_allocated(device) for device in cuda_p2p_devices) + + assert allocated_after == allocated_before + assert pointers == ( + left.solver.Ey.data_ptr(), + left.solver.Ez.data_ptr(), + right.solver.Hy.data_ptr(), + right.solver.Hz.data_ptr(), + ) + + +def test_cuda_p2p_transport_waits_for_pending_destination_writes( + cuda_p2p_devices, + cuda_memory_cleanup, +): + """A delayed compute write must finish before the receive copy overwrites it.""" + + transport = CudaP2PHaloTransport(cuda_p2p_devices) + transport.preflight() + _, shards = _make_transport_shards(cuda_p2p_devices) + left, right = shards + + right_node = right.layout.storage_node_owned.start + left_ghost = left.layout.storage_node_owned.stop + ey_tag = _tag_plane(right.solver.Ey[right_node], offset=11_000.0) + ez_tag = _tag_plane(right.solver.Ez[right_node], offset=12_000.0) + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + + with torch.cuda.device(right.device), torch.cuda.stream(right.compute_stream): + right.solver.Ey[right_node].copy_(ey_tag) + right.solver.Ez[right_node].copy_(ez_tag) + with torch.cuda.device(left.device), torch.cuda.stream(left.compute_stream): + torch.cuda._sleep(20_000_000) + left.solver.Ey[left_ghost].fill_(-17.0) + left.solver.Ez[left_ghost].fill_(-17.0) + + transport.exchange_electric(shards) + left.electric_received.synchronize() + torch.testing.assert_close( + left.solver.Ey[left_ghost], ey_tag.to(left.device), rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + left.solver.Ez[left_ghost], ez_tag.to(left.device), rtol=0.0, atol=0.0 + ) + + left_last = left.layout.storage_cell_owned.stop - 1 + hy_tag = _tag_plane(left.solver.Hy[left_last], offset=13_000.0) + hz_tag = _tag_plane(left.solver.Hz[left_last], offset=14_000.0) + for device in cuda_p2p_devices: + torch.cuda.synchronize(device) + + with torch.cuda.device(left.device), torch.cuda.stream(left.compute_stream): + left.solver.Hy[left_last].copy_(hy_tag) + left.solver.Hz[left_last].copy_(hz_tag) + with torch.cuda.device(right.device), torch.cuda.stream(right.compute_stream): + torch.cuda._sleep(20_000_000) + right.halo_hy_low.fill_(-23.0) + right.halo_hz_low.fill_(-23.0) + + transport.exchange_magnetic(shards) + right.magnetic_received.synchronize() + torch.testing.assert_close(right.halo_hy_low, hy_tag.to(right.device), rtol=0.0, atol=0.0) + torch.testing.assert_close(right.halo_hz_low, hz_tag.to(right.device), rtol=0.0, atol=0.0) + + +def test_cuda_p2p_transport_reports_bidirectional_neighbor_capability(cuda_p2p_devices): + transport = CudaP2PHaloTransport(cuda_p2p_devices) + transport.preflight() + + assert transport.name == "cuda_p2p" + assert len(transport.links) == 2 + assert {link.direction for link in transport.links} == { + "electric_right_to_left", + "magnetic_left_to_right", + } + topology = transport.topology + assert topology["kind"] == "cuda_p2p" + assert topology["neighbor_pairs"] == ( + { + "devices": ("cuda:0", "cuda:1"), + "peer_left_to_right": True, + "peer_right_to_left": True, + }, + ) diff --git a/witwin/maxwell/__init__.py b/witwin/maxwell/__init__.py index e3b6719..4900d7c 100644 --- a/witwin/maxwell/__init__.py +++ b/witwin/maxwell/__init__.py @@ -76,6 +76,7 @@ TimeConfig, run, ) +from .fdtd_parallel import FDTDParallelConfig __all__ = [ "AstigmaticGaussianBeam", @@ -106,6 +107,7 @@ "FluxTimeMonitor", "FDFDConfig", "FDTDConfig", + "FDTDParallelConfig", "FinitePlaneMonitor", "GaussianBeam", "GaussianPulse", diff --git a/witwin/maxwell/fdtd/cuda/backend.py b/witwin/maxwell/fdtd/cuda/backend.py index 5de82d2..719614c 100644 --- a/witwin/maxwell/fdtd/cuda/backend.py +++ b/witwin/maxwell/fdtd/cuda/backend.py @@ -127,6 +127,36 @@ def _magnetic_hz_standard(*, Hz, Ex, Ey, HzDecay, HzCurl, invDx, invDy): _COMPILED_EXTENSION.update_magnetic_hz_standard(Hz, Ex, Ey, HzDecay, HzCurl, invDx, invDy) +def _magnetic_hx_standard_bounded( + *, Hx, Ey, Ez, HxDecay, HxCurl, invDy, invDz, + localXBegin, localXEnd, globalXOffset, globalXExtent, +): + _COMPILED_EXTENSION.update_magnetic_hx_standard_bounded( + Hx, Ey, Ez, HxDecay, HxCurl, invDy, invDz, + int(localXBegin), int(localXEnd), int(globalXOffset), int(globalXExtent), + ) + + +def _magnetic_hy_standard_bounded( + *, Hy, Ex, Ez, HyDecay, HyCurl, invDx, invDz, + localXBegin, localXEnd, globalXOffset, globalXExtent, +): + _COMPILED_EXTENSION.update_magnetic_hy_standard_bounded( + Hy, Ex, Ez, HyDecay, HyCurl, invDx, invDz, + int(localXBegin), int(localXEnd), int(globalXOffset), int(globalXExtent), + ) + + +def _magnetic_hz_standard_bounded( + *, Hz, Ex, Ey, HzDecay, HzCurl, invDx, invDy, + localXBegin, localXEnd, globalXOffset, globalXExtent, +): + _COMPILED_EXTENSION.update_magnetic_hz_standard_bounded( + Hz, Ex, Ey, HzDecay, HzCurl, invDx, invDy, + int(localXBegin), int(localXEnd), int(globalXOffset), int(globalXExtent), + ) + + def _magnetic_hx_cpml( *, Hx, @@ -484,6 +514,42 @@ def _electric_ez_standard( ) + + +def _electric_ex_standard_bounded( + *, Ex, Hy, Hz, ExDecay, ExCurl, invDy, invDz, + yLowBoundaryMode, yHighBoundaryMode, zLowBoundaryMode, zHighBoundaryMode, + localXBegin, localXEnd, globalXOffset, globalXExtent, +): + _COMPILED_EXTENSION.update_electric_ex_standard_bounded( + Ex, Hy, Hz, ExDecay, ExCurl, invDy, invDz, + int(yLowBoundaryMode), int(yHighBoundaryMode), int(zLowBoundaryMode), int(zHighBoundaryMode), + int(localXBegin), int(localXEnd), int(globalXOffset), int(globalXExtent), + ) + + +def _electric_ey_standard_bounded( + *, Ey, Hx, Hz, EyDecay, EyCurl, invDx, invDz, + xLowBoundaryMode, xHighBoundaryMode, zLowBoundaryMode, zHighBoundaryMode, + localXBegin, localXEnd, globalXOffset, globalXExtent, +): + _COMPILED_EXTENSION.update_electric_ey_standard_bounded( + Ey, Hx, Hz, EyDecay, EyCurl, invDx, invDz, + int(xLowBoundaryMode), int(xHighBoundaryMode), int(zLowBoundaryMode), int(zHighBoundaryMode), + int(localXBegin), int(localXEnd), int(globalXOffset), int(globalXExtent), + ) + + +def _electric_ez_standard_bounded( + *, Ez, Hx, Hy, EzDecay, EzCurl, invDx, invDy, + xLowBoundaryMode, xHighBoundaryMode, yLowBoundaryMode, yHighBoundaryMode, + localXBegin, localXEnd, globalXOffset, globalXExtent, +): + _COMPILED_EXTENSION.update_electric_ez_standard_bounded( + Ez, Hx, Hy, EzDecay, EzCurl, invDx, invDy, + int(xLowBoundaryMode), int(xHighBoundaryMode), int(yLowBoundaryMode), int(yHighBoundaryMode), + int(localXBegin), int(localXEnd), int(globalXOffset), int(globalXExtent), + ) def _electric_ex_modulated( *, Ex, @@ -3398,6 +3464,9 @@ def _project_bloch_boundary(*, fieldReal, fieldImag, axis, phaseCos, phaseSin): "updateMagneticFieldHxStandard3D": _magnetic_hx_standard, "updateMagneticFieldHyStandard3D": _magnetic_hy_standard, "updateMagneticFieldHzStandard3D": _magnetic_hz_standard, + "updateMagneticFieldHxStandardBounded3D": _magnetic_hx_standard_bounded, + "updateMagneticFieldHyStandardBounded3D": _magnetic_hy_standard_bounded, + "updateMagneticFieldHzStandardBounded3D": _magnetic_hz_standard_bounded, "updateMagneticFieldHx3D": _magnetic_hx_cpml, "updateMagneticFieldHy3D": _magnetic_hy_cpml, "updateMagneticFieldHz3D": _magnetic_hz_cpml, @@ -3405,6 +3474,7 @@ def _project_bloch_boundary(*, fieldReal, fieldImag, axis, phaseCos, phaseSin): "updateMagneticFieldHyCpmlCompressed3D": _magnetic_hy_cpml_compressed, "updateMagneticFieldHzCpmlCompressed3D": _magnetic_hz_cpml_compressed, "updateElectricFieldExStandard3D": _electric_ex_standard, + "updateElectricFieldExStandardBounded3D": _electric_ex_standard_bounded, "advanceModulationTime3D": _advance_modulation_time, "updateElectricFieldExModulated3D": _electric_ex_modulated, "updateElectricFieldEyModulated3D": _electric_ey_modulated, @@ -3414,6 +3484,8 @@ def _project_bloch_boundary(*, fieldReal, fieldImag, axis, phaseCos, phaseSin): "updateElectricFieldEzCpmlModulated3D": _electric_ez_cpml_modulated, "updateElectricFieldEyStandard3D": _electric_ey_standard, "updateElectricFieldEzStandard3D": _electric_ez_standard, + "updateElectricFieldEyStandardBounded3D": _electric_ey_standard_bounded, + "updateElectricFieldEzStandardBounded3D": _electric_ez_standard_bounded, "updateElectricFieldExCpml3D": _electric_ex_cpml, "updateElectricFieldEyCpml3D": _electric_ey_cpml, "updateElectricFieldEzCpml3D": _electric_ez_cpml, diff --git a/witwin/maxwell/fdtd/cuda/build.py b/witwin/maxwell/fdtd/cuda/build.py index 02e5e7a..6eb345d 100644 --- a/witwin/maxwell/fdtd/cuda/build.py +++ b/witwin/maxwell/fdtd/cuda/build.py @@ -10,6 +10,17 @@ import torch from torch.utils.cpp_extension import load +STABLE_ABI_VERSION = "stable_abi_v2" +_REQUIRED_STABLE_OPS = ( + "update_magnetic_hx_standard_bounded", + "update_magnetic_hy_standard_bounded", + "update_magnetic_hz_standard_bounded", + "update_electric_ex_standard_bounded", + "update_electric_ey_standard_bounded", + "update_electric_ez_standard_bounded", +) + + def _candidate_vcvars64_paths() -> list[Path]: paths: list[Path] = [] @@ -159,6 +170,11 @@ def prebuilt_extension_path() -> Path: return prebuilt_root() / f"witwin_maxwell_fdtd_cuda{extension_suffix()}" +def prebuilt_abi_marker_path(library_path: Path | None = None) -> Path: + selected = prebuilt_extension_path() if library_path is None else Path(library_path) + return selected.with_suffix(f"{selected.suffix}.abi") + + def extension_sources() -> list[Path]: root = source_root() return [ @@ -227,7 +243,8 @@ def __getattr__(self, name: str): def _load_extension_file(library_path: Path) -> _StableOpsModule: torch.ops.load_library(str(library_path)) - if not hasattr(torch.ops.witwin_maxwell_fdtd_cuda, "update_magnetic_hx_standard"): + namespace = torch.ops.witwin_maxwell_fdtd_cuda + if any(not hasattr(namespace, name) for name in _REQUIRED_STABLE_OPS): raise ImportError(f"{library_path} does not register the Stable ABI Maxwell operators.") return _StableOpsModule(library_path) @@ -236,6 +253,11 @@ def _load_packaged_prebuilt_extension(): module_path = prebuilt_extension_path() if not module_path.exists(): return None + marker_path = prebuilt_abi_marker_path(module_path) + if not marker_path.exists(): + return None + if marker_path.read_text(encoding="utf-8").strip() != STABLE_ABI_VERSION: + return None return _load_extension_file(module_path) @@ -255,13 +277,14 @@ def _load_prebuilt_extension(build_directory: Path): def build_extension(*, verbose: bool = False): root = source_root() - default_build_directory = Path(tempfile.gettempdir()) / "witwin_maxwell_fdtd_cuda" / "stable_abi_v1" + default_build_directory = Path(tempfile.gettempdir()) / "witwin_maxwell_fdtd_cuda" / STABLE_ABI_VERSION build_directory = Path(os.environ.get("WITWIN_MAXWELL_FDTD_CUDA_BUILD_DIR", default_build_directory)) if os.environ.get("WITWIN_MAXWELL_FDTD_CUDA_SKIP_PREBUILT") != "1": - try: - module = _load_packaged_prebuilt_extension() - except Exception: # noqa: BLE001 - stale/ABI-mismatched prebuilt, rebuild instead - module = None + # Missing or stale markers are rejected before loading the library so a + # JIT rebuild remains safe. A matching marker is an ABI promise: if + # loading or probing that binary fails, propagate the error instead of + # trying to register the same dispatcher namespace again in-process. + module = _load_packaged_prebuilt_extension() if module is not None: return module if os.environ.get("WITWIN_MAXWELL_FDTD_CUDA_PREBUILT") == "1": diff --git a/witwin/maxwell/fdtd/cuda/extension.cpp b/witwin/maxwell/fdtd/cuda/extension.cpp index 8915217..2f589ce 100644 --- a/witwin/maxwell/fdtd/cuda/extension.cpp +++ b/witwin/maxwell/fdtd/cuda/extension.cpp @@ -36,6 +36,42 @@ void update_magnetic_hz_standard_cuda( const torch::stable::Tensor& curl, const torch::stable::Tensor& inv_dx, const torch::stable::Tensor& inv_dy); +void update_magnetic_hx_standard_bounded_cuda( + torch::stable::Tensor hx, + const torch::stable::Tensor& ey, + const torch::stable::Tensor& ez, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dy, + const torch::stable::Tensor& inv_dz, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent); +void update_magnetic_hy_standard_bounded_cuda( + torch::stable::Tensor hy, + const torch::stable::Tensor& ex, + const torch::stable::Tensor& ez, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dz, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent); +void update_magnetic_hz_standard_bounded_cuda( + torch::stable::Tensor hz, + const torch::stable::Tensor& ex, + const torch::stable::Tensor& ey, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dy, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent); void update_magnetic_hx_cpml_cuda( torch::stable::Tensor hx, const torch::stable::Tensor& ey, @@ -192,6 +228,54 @@ void update_electric_ez_standard_cuda( int64_t x_high_mode, int64_t y_low_mode, int64_t y_high_mode); +void update_electric_ex_standard_bounded_cuda( + torch::stable::Tensor ex, + const torch::stable::Tensor& hy, + const torch::stable::Tensor& hz, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dy, + const torch::stable::Tensor& inv_dz, + int64_t y_low_mode, + int64_t y_high_mode, + int64_t z_low_mode, + int64_t z_high_mode, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent); +void update_electric_ey_standard_bounded_cuda( + torch::stable::Tensor ey, + const torch::stable::Tensor& hx, + const torch::stable::Tensor& hz, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dz, + int64_t x_low_mode, + int64_t x_high_mode, + int64_t z_low_mode, + int64_t z_high_mode, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent); +void update_electric_ez_standard_bounded_cuda( + torch::stable::Tensor ez, + const torch::stable::Tensor& hx, + const torch::stable::Tensor& hy, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dy, + int64_t x_low_mode, + int64_t x_high_mode, + int64_t y_low_mode, + int64_t y_high_mode, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent); void update_electric_ex_modulated_cuda( torch::stable::Tensor ex, const torch::stable::Tensor& hy, @@ -1910,6 +1994,9 @@ void add_source_patch_ez_periodic_cuda( _(update_magnetic_hx_standard, "update_magnetic_hx_standard(Tensor(a!) hx, Tensor ey, Tensor ez, Tensor decay, Tensor curl, Tensor inv_dy, Tensor inv_dz) -> ()", update_magnetic_hx_standard_cuda) \ _(update_magnetic_hy_standard, "update_magnetic_hy_standard(Tensor(a!) hy, Tensor ex, Tensor ez, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dz) -> ()", update_magnetic_hy_standard_cuda) \ _(update_magnetic_hz_standard, "update_magnetic_hz_standard(Tensor(a!) hz, Tensor ex, Tensor ey, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dy) -> ()", update_magnetic_hz_standard_cuda) \ + _(update_magnetic_hx_standard_bounded, "update_magnetic_hx_standard_bounded(Tensor(a!) hx, Tensor ey, Tensor ez, Tensor decay, Tensor curl, Tensor inv_dy, Tensor inv_dz, int local_x_begin, int local_x_end, int global_x_offset, int global_x_extent) -> ()", update_magnetic_hx_standard_bounded_cuda) \ + _(update_magnetic_hy_standard_bounded, "update_magnetic_hy_standard_bounded(Tensor(a!) hy, Tensor ex, Tensor ez, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dz, int local_x_begin, int local_x_end, int global_x_offset, int global_x_extent) -> ()", update_magnetic_hy_standard_bounded_cuda) \ + _(update_magnetic_hz_standard_bounded, "update_magnetic_hz_standard_bounded(Tensor(a!) hz, Tensor ex, Tensor ey, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dy, int local_x_begin, int local_x_end, int global_x_offset, int global_x_extent) -> ()", update_magnetic_hz_standard_bounded_cuda) \ _(update_magnetic_hx_cpml, "update_magnetic_hx_cpml(Tensor(a!) hx, Tensor ey, Tensor ez, Tensor decay, Tensor curl, Tensor(b!) psi_y, Tensor(c!) psi_z, Tensor inv_kappa_y, Tensor b_y, Tensor c_y, Tensor inv_kappa_z, Tensor b_z, Tensor c_z, Tensor inv_dy, Tensor inv_dz) -> ()", update_magnetic_hx_cpml_cuda) \ _(update_magnetic_hy_cpml, "update_magnetic_hy_cpml(Tensor(a!) hy, Tensor ex, Tensor ez, Tensor decay, Tensor curl, Tensor(b!) psi_x, Tensor(c!) psi_z, Tensor inv_kappa_x, Tensor b_x, Tensor c_x, Tensor inv_kappa_z, Tensor b_z, Tensor c_z, Tensor inv_dx, Tensor inv_dz) -> ()", update_magnetic_hy_cpml_cuda) \ _(update_magnetic_hz_cpml, "update_magnetic_hz_cpml(Tensor(a!) hz, Tensor ex, Tensor ey, Tensor decay, Tensor curl, Tensor(b!) psi_x, Tensor(c!) psi_y, Tensor inv_kappa_x, Tensor b_x, Tensor c_x, Tensor inv_kappa_y, Tensor b_y, Tensor c_y, Tensor inv_dx, Tensor inv_dy) -> ()", update_magnetic_hz_cpml_cuda) \ @@ -1919,6 +2006,9 @@ void add_source_patch_ez_periodic_cuda( _(update_electric_ex_standard, "update_electric_ex_standard(Tensor(a!) ex, Tensor hy, Tensor hz, Tensor decay, Tensor curl, Tensor inv_dy, Tensor inv_dz, int y_low_mode, int y_high_mode, int z_low_mode, int z_high_mode) -> ()", update_electric_ex_standard_cuda) \ _(update_electric_ey_standard, "update_electric_ey_standard(Tensor(a!) ey, Tensor hx, Tensor hz, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dz, int x_low_mode, int x_high_mode, int z_low_mode, int z_high_mode) -> ()", update_electric_ey_standard_cuda) \ _(update_electric_ez_standard, "update_electric_ez_standard(Tensor(a!) ez, Tensor hx, Tensor hy, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dy, int x_low_mode, int x_high_mode, int y_low_mode, int y_high_mode) -> ()", update_electric_ez_standard_cuda) \ + _(update_electric_ex_standard_bounded, "update_electric_ex_standard_bounded(Tensor(a!) ex, Tensor hy, Tensor hz, Tensor decay, Tensor curl, Tensor inv_dy, Tensor inv_dz, int y_low_mode, int y_high_mode, int z_low_mode, int z_high_mode, int local_x_begin, int local_x_end, int global_x_offset, int global_x_extent) -> ()", update_electric_ex_standard_bounded_cuda) \ + _(update_electric_ey_standard_bounded, "update_electric_ey_standard_bounded(Tensor(a!) ey, Tensor hx, Tensor hz, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dz, int x_low_mode, int x_high_mode, int z_low_mode, int z_high_mode, int local_x_begin, int local_x_end, int global_x_offset, int global_x_extent) -> ()", update_electric_ey_standard_bounded_cuda) \ + _(update_electric_ez_standard_bounded, "update_electric_ez_standard_bounded(Tensor(a!) ez, Tensor hx, Tensor hy, Tensor decay, Tensor curl, Tensor inv_dx, Tensor inv_dy, int x_low_mode, int x_high_mode, int y_low_mode, int y_high_mode, int local_x_begin, int local_x_end, int global_x_offset, int global_x_extent) -> ()", update_electric_ez_standard_bounded_cuda) \ _(advance_modulation_time, "advance_modulation_time(Tensor(a!) modulation_time, float dt) -> ()", advance_modulation_time_cuda) \ _(update_electric_ex_modulated, "update_electric_ex_modulated(Tensor(a!) ex, Tensor hy, Tensor hz, Tensor decay, Tensor curl, Tensor mod_cos, Tensor mod_sin, Tensor mod_omega, Tensor modulation_time, Tensor inv_dy, Tensor inv_dz, int y_low_mode, int y_high_mode, int z_low_mode, int z_high_mode) -> ()", update_electric_ex_modulated_cuda) \ _(update_electric_ey_modulated, "update_electric_ey_modulated(Tensor(a!) ey, Tensor hx, Tensor hz, Tensor decay, Tensor curl, Tensor mod_cos, Tensor mod_sin, Tensor mod_omega, Tensor modulation_time, Tensor inv_dx, Tensor inv_dz, int x_low_mode, int x_high_mode, int z_low_mode, int z_high_mode) -> ()", update_electric_ey_modulated_cuda) \ diff --git a/witwin/maxwell/fdtd/cuda/kernels/electric.cu b/witwin/maxwell/fdtd/cuda/kernels/electric.cu index f00934f..ae16e0e 100644 --- a/witwin/maxwell/fdtd/cuda/kernels/electric.cu +++ b/witwin/maxwell/fdtd/cuda/kernels/electric.cu @@ -209,6 +209,8 @@ __global__ void update_electric_standard_kernel( unsigned int nx, unsigned int ny, unsigned int nz, + unsigned int local_x_begin, + unsigned int local_x_end, const float* __restrict__ first, const float* __restrict__ second, const float* __restrict__ decay, @@ -222,8 +224,8 @@ __global__ void update_electric_standard_kernel( float* __restrict__ field) { const unsigned int k = blockIdx.x * blockDim.x + threadIdx.x; const unsigned int j = blockIdx.y * blockDim.y + threadIdx.y; - const unsigned int i = blockIdx.z * blockDim.z + threadIdx.z; - if (i >= nx || j >= ny || k >= nz) { + const unsigned int i = local_x_begin + blockIdx.z * blockDim.z + threadIdx.z; + if (i >= local_x_end || i >= nx || j >= ny || k >= nz) { return; } const long long linear = offset3d(i, j, k, ny, nz); @@ -260,6 +262,8 @@ __global__ void update_electric_standard_interior_kernel( unsigned int nx, unsigned int ny, unsigned int nz, + unsigned int local_x_begin, + unsigned int local_x_end, const float* __restrict__ first, const float* __restrict__ second, const float* __restrict__ decay, @@ -269,8 +273,8 @@ __global__ void update_electric_standard_interior_kernel( float* __restrict__ field) { const unsigned int k = blockIdx.x * blockDim.x + threadIdx.x + (Component == 2 ? 0 : 1); const unsigned int j = blockIdx.y * blockDim.y + threadIdx.y + (Component == 1 ? 0 : 1); - const unsigned int i = blockIdx.z * blockDim.z + threadIdx.z + (Component == 0 ? 0 : 1); - if (i + (Component == 0 ? 0 : 1) >= nx || j + (Component == 1 ? 0 : 1) >= ny || + const unsigned int i = local_x_begin + blockIdx.z * blockDim.z + threadIdx.z; + if (i >= local_x_end || i >= nx || j + (Component == 1 ? 0 : 1) >= ny || k + (Component == 2 ? 0 : 1) >= nz) { return; } @@ -1862,7 +1866,7 @@ void advance_modulation_time_cuda(torch::stable::Tensor modulation_time, double WITWIN_CUDA_CHECK(); } -void update_electric_ex_standard_cuda( +void update_electric_ex_standard_bounded_cuda( torch::stable::Tensor ex, const torch::stable::Tensor& hy, const torch::stable::Tensor& hz, @@ -1873,21 +1877,32 @@ void update_electric_ex_standard_cuda( int64_t y_low_mode, int64_t y_high_mode, int64_t z_low_mode, - int64_t z_high_mode) { + int64_t z_high_mode, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent) { check_electric_inputs(ex, hy, hz, decay, curl, "ex"); check_rank3_shape(hy, "hy", ex.size(0), ex.size(1), ex.size(2) - 1); check_rank3_shape(hz, "hz", ex.size(0), ex.size(1) - 1, ex.size(2)); check_spacing_vector(ex, inv_dy, 1, "inv_dy"); check_spacing_vector(ex, inv_dz, 2, "inv_dz"); + check_bounded_x_launch( + ex, local_x_begin, local_x_end, global_x_offset, global_x_extent, "ex"); + if (local_x_begin == local_x_end) { + return; + } torch::stable::accelerator::DeviceGuard guard(ex.get_device_index()); const auto sizes = ex.sizes(); const dim3 block = field_block3d(); if (inactive_boundary_pair(y_low_mode, y_high_mode) && inactive_boundary_pair(z_low_mode, z_high_mode)) { if (sizes[1] > 2 && sizes[2] > 2) { - update_electric_standard_interior_kernel<0><<>>( + update_electric_standard_interior_kernel<0><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), hy.mutable_data_ptr(), hz.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -1897,10 +1912,12 @@ void update_electric_ex_standard_cuda( ex.mutable_data_ptr()); } } else { - update_electric_standard_kernel<0><<>>( + update_electric_standard_kernel<0><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), hy.mutable_data_ptr(), hz.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -1916,7 +1933,26 @@ void update_electric_ex_standard_cuda( WITWIN_CUDA_CHECK(); } -void update_electric_ey_standard_cuda( +void update_electric_ex_standard_cuda( + torch::stable::Tensor ex, + const torch::stable::Tensor& hy, + const torch::stable::Tensor& hz, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dy, + const torch::stable::Tensor& inv_dz, + int64_t y_low_mode, + int64_t y_high_mode, + int64_t z_low_mode, + int64_t z_high_mode) { + const int64_t x_extent = ex.size(0); + update_electric_ex_standard_bounded_cuda( + ex, hy, hz, decay, curl, inv_dy, inv_dz, + y_low_mode, y_high_mode, z_low_mode, z_high_mode, + 0, x_extent, 0, x_extent); +} + +void update_electric_ey_standard_bounded_cuda( torch::stable::Tensor ey, const torch::stable::Tensor& hx, const torch::stable::Tensor& hz, @@ -1927,21 +1963,38 @@ void update_electric_ey_standard_cuda( int64_t x_low_mode, int64_t x_high_mode, int64_t z_low_mode, - int64_t z_high_mode) { + int64_t z_high_mode, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent) { check_electric_inputs(ey, hx, hz, decay, curl, "ey"); check_rank3_shape(hx, "hx", ey.size(0), ey.size(1), ey.size(2) - 1); check_rank3_shape(hz, "hz", ey.size(0) - 1, ey.size(1), ey.size(2)); check_spacing_vector(ey, inv_dx, 0, "inv_dx"); check_spacing_vector(ey, inv_dz, 2, "inv_dz"); + check_bounded_x_launch( + ey, local_x_begin, local_x_end, global_x_offset, global_x_extent, "ey"); + if (local_x_begin == local_x_end) { + return; + } + const int64_t effective_x_low_mode = + global_x_offset == 0 ? x_low_mode : BOUNDARY_NONE; + const int64_t effective_x_high_mode = + global_x_offset + ey.size(0) == global_x_extent ? x_high_mode : BOUNDARY_NONE; torch::stable::accelerator::DeviceGuard guard(ey.get_device_index()); const auto sizes = ey.sizes(); const dim3 block = field_block3d(); - if (inactive_boundary_pair(x_low_mode, x_high_mode) && inactive_boundary_pair(z_low_mode, z_high_mode)) { - if (sizes[0] > 2 && sizes[2] > 2) { - update_electric_standard_interior_kernel<1><<>>( + if (inactive_boundary_pair(effective_x_low_mode, effective_x_high_mode) && inactive_boundary_pair(z_low_mode, z_high_mode)) { + const int64_t interior_x_begin = local_x_begin < 1 ? 1 : local_x_begin; + const int64_t interior_x_end = local_x_end > sizes[0] - 1 ? sizes[0] - 1 : local_x_end; + if (interior_x_end > interior_x_begin && sizes[2] > 2) { + update_electric_standard_interior_kernel<1><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(interior_x_begin), + static_cast(interior_x_end), hx.mutable_data_ptr(), hz.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -1951,18 +2004,20 @@ void update_electric_ey_standard_cuda( ey.mutable_data_ptr()); } } else { - update_electric_standard_kernel<1><<>>( + update_electric_standard_kernel<1><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), hx.mutable_data_ptr(), hz.mutable_data_ptr(), decay.mutable_data_ptr(), curl.mutable_data_ptr(), inv_dx.mutable_data_ptr(), inv_dz.mutable_data_ptr(), - static_cast(x_low_mode), - static_cast(x_high_mode), + static_cast(effective_x_low_mode), + static_cast(effective_x_high_mode), static_cast(z_low_mode), static_cast(z_high_mode), ey.mutable_data_ptr()); @@ -1970,7 +2025,26 @@ void update_electric_ey_standard_cuda( WITWIN_CUDA_CHECK(); } -void update_electric_ez_standard_cuda( +void update_electric_ey_standard_cuda( + torch::stable::Tensor ey, + const torch::stable::Tensor& hx, + const torch::stable::Tensor& hz, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dz, + int64_t x_low_mode, + int64_t x_high_mode, + int64_t z_low_mode, + int64_t z_high_mode) { + const int64_t x_extent = ey.size(0); + update_electric_ey_standard_bounded_cuda( + ey, hx, hz, decay, curl, inv_dx, inv_dz, + x_low_mode, x_high_mode, z_low_mode, z_high_mode, + 0, x_extent, 0, x_extent); +} + +void update_electric_ez_standard_bounded_cuda( torch::stable::Tensor ez, const torch::stable::Tensor& hx, const torch::stable::Tensor& hy, @@ -1981,21 +2055,38 @@ void update_electric_ez_standard_cuda( int64_t x_low_mode, int64_t x_high_mode, int64_t y_low_mode, - int64_t y_high_mode) { + int64_t y_high_mode, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent) { check_electric_inputs(ez, hx, hy, decay, curl, "ez"); check_rank3_shape(hx, "hx", ez.size(0), ez.size(1) - 1, ez.size(2)); check_rank3_shape(hy, "hy", ez.size(0) - 1, ez.size(1), ez.size(2)); check_spacing_vector(ez, inv_dx, 0, "inv_dx"); check_spacing_vector(ez, inv_dy, 1, "inv_dy"); + check_bounded_x_launch( + ez, local_x_begin, local_x_end, global_x_offset, global_x_extent, "ez"); + if (local_x_begin == local_x_end) { + return; + } + const int64_t effective_x_low_mode = + global_x_offset == 0 ? x_low_mode : BOUNDARY_NONE; + const int64_t effective_x_high_mode = + global_x_offset + ez.size(0) == global_x_extent ? x_high_mode : BOUNDARY_NONE; torch::stable::accelerator::DeviceGuard guard(ez.get_device_index()); const auto sizes = ez.sizes(); const dim3 block = field_block3d(); - if (inactive_boundary_pair(x_low_mode, x_high_mode) && inactive_boundary_pair(y_low_mode, y_high_mode)) { - if (sizes[0] > 2 && sizes[1] > 2) { - update_electric_standard_interior_kernel<2><<>>( + if (inactive_boundary_pair(effective_x_low_mode, effective_x_high_mode) && inactive_boundary_pair(y_low_mode, y_high_mode)) { + const int64_t interior_x_begin = local_x_begin < 1 ? 1 : local_x_begin; + const int64_t interior_x_end = local_x_end > sizes[0] - 1 ? sizes[0] - 1 : local_x_end; + if (interior_x_end > interior_x_begin && sizes[1] > 2) { + update_electric_standard_interior_kernel<2><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(interior_x_begin), + static_cast(interior_x_end), hx.mutable_data_ptr(), hy.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -2005,18 +2096,20 @@ void update_electric_ez_standard_cuda( ez.mutable_data_ptr()); } } else { - update_electric_standard_kernel<2><<>>( + update_electric_standard_kernel<2><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), hx.mutable_data_ptr(), hy.mutable_data_ptr(), decay.mutable_data_ptr(), curl.mutable_data_ptr(), inv_dx.mutable_data_ptr(), inv_dy.mutable_data_ptr(), - static_cast(x_low_mode), - static_cast(x_high_mode), + static_cast(effective_x_low_mode), + static_cast(effective_x_high_mode), static_cast(y_low_mode), static_cast(y_high_mode), ez.mutable_data_ptr()); @@ -2024,6 +2117,25 @@ void update_electric_ez_standard_cuda( WITWIN_CUDA_CHECK(); } +void update_electric_ez_standard_cuda( + torch::stable::Tensor ez, + const torch::stable::Tensor& hx, + const torch::stable::Tensor& hy, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dy, + int64_t x_low_mode, + int64_t x_high_mode, + int64_t y_low_mode, + int64_t y_high_mode) { + const int64_t x_extent = ez.size(0); + update_electric_ez_standard_bounded_cuda( + ez, hx, hy, decay, curl, inv_dx, inv_dy, + x_low_mode, x_high_mode, y_low_mode, y_high_mode, + 0, x_extent, 0, x_extent); +} + void update_electric_ex_bloch_cuda( torch::stable::Tensor ex_real, torch::stable::Tensor ex_imag, diff --git a/witwin/maxwell/fdtd/cuda/kernels/magnetic.cu b/witwin/maxwell/fdtd/cuda/kernels/magnetic.cu index 5849e8c..9f5a5b2 100644 --- a/witwin/maxwell/fdtd/cuda/kernels/magnetic.cu +++ b/witwin/maxwell/fdtd/cuda/kernels/magnetic.cu @@ -22,6 +22,8 @@ __global__ void update_magnetic_standard_kernel( unsigned int nx, unsigned int ny, unsigned int nz, + unsigned int local_x_begin, + unsigned int local_x_end, const float* __restrict__ first, const float* __restrict__ second, const float* __restrict__ decay, @@ -31,8 +33,8 @@ __global__ void update_magnetic_standard_kernel( float* __restrict__ field) { const unsigned int k = blockIdx.x * blockDim.x + threadIdx.x; const unsigned int j = blockIdx.y * blockDim.y + threadIdx.y; - const unsigned int i = blockIdx.z * blockDim.z + threadIdx.z; - if (i >= nx || j >= ny || k >= nz) { + const unsigned int i = local_x_begin + blockIdx.z * blockDim.z + threadIdx.z; + if (i >= local_x_end || i >= nx || j >= ny || k >= nz) { return; } const long long linear = offset3d(i, j, k, ny, nz); @@ -351,26 +353,37 @@ void check_magnetic_cpml_compressed_inputs( } // namespace -void update_magnetic_hx_standard_cuda( +void update_magnetic_hx_standard_bounded_cuda( torch::stable::Tensor hx, const torch::stable::Tensor& ey, const torch::stable::Tensor& ez, const torch::stable::Tensor& decay, const torch::stable::Tensor& curl, const torch::stable::Tensor& inv_dy, - const torch::stable::Tensor& inv_dz) { + const torch::stable::Tensor& inv_dz, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent) { check_magnetic_inputs(hx, ey, ez, decay, curl, "hx"); check_rank3_shape(ey, "ey", hx.size(0), hx.size(1), hx.size(2) + 1); check_rank3_shape(ez, "ez", hx.size(0), hx.size(1) + 1, hx.size(2)); check_spacing_vector(hx, inv_dy, 1, "inv_dy"); check_spacing_vector(hx, inv_dz, 2, "inv_dz"); + check_bounded_x_launch( + hx, local_x_begin, local_x_end, global_x_offset, global_x_extent, "hx"); + if (local_x_begin == local_x_end) { + return; + } torch::stable::accelerator::DeviceGuard guard(hx.get_device_index()); const auto sizes = hx.sizes(); const dim3 block = field_block3d(); - update_magnetic_standard_kernel<0><<>>( + update_magnetic_standard_kernel<0><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), ey.mutable_data_ptr(), ez.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -381,26 +394,60 @@ void update_magnetic_hx_standard_cuda( WITWIN_CUDA_CHECK(); } -void update_magnetic_hy_standard_cuda( +void update_magnetic_hx_standard_cuda( + torch::stable::Tensor hx, + const torch::stable::Tensor& ey, + const torch::stable::Tensor& ez, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dy, + const torch::stable::Tensor& inv_dz) { + const int64_t x_extent = hx.size(0); + update_magnetic_hx_standard_bounded_cuda( + hx, + ey, + ez, + decay, + curl, + inv_dy, + inv_dz, + 0, + x_extent, + 0, + x_extent); +} + +void update_magnetic_hy_standard_bounded_cuda( torch::stable::Tensor hy, const torch::stable::Tensor& ex, const torch::stable::Tensor& ez, const torch::stable::Tensor& decay, const torch::stable::Tensor& curl, const torch::stable::Tensor& inv_dx, - const torch::stable::Tensor& inv_dz) { + const torch::stable::Tensor& inv_dz, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent) { check_magnetic_inputs(hy, ex, ez, decay, curl, "hy"); check_rank3_shape(ex, "ex", hy.size(0), hy.size(1), hy.size(2) + 1); check_rank3_shape(ez, "ez", hy.size(0) + 1, hy.size(1), hy.size(2)); check_spacing_vector(hy, inv_dx, 0, "inv_dx"); check_spacing_vector(hy, inv_dz, 2, "inv_dz"); + check_bounded_x_launch( + hy, local_x_begin, local_x_end, global_x_offset, global_x_extent, "hy"); + if (local_x_begin == local_x_end) { + return; + } torch::stable::accelerator::DeviceGuard guard(hy.get_device_index()); const auto sizes = hy.sizes(); const dim3 block = field_block3d(); - update_magnetic_standard_kernel<1><<>>( + update_magnetic_standard_kernel<1><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), ex.mutable_data_ptr(), ez.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -411,26 +458,50 @@ void update_magnetic_hy_standard_cuda( WITWIN_CUDA_CHECK(); } -void update_magnetic_hz_standard_cuda( +void update_magnetic_hy_standard_cuda( + torch::stable::Tensor hy, + const torch::stable::Tensor& ex, + const torch::stable::Tensor& ez, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dz) { + const int64_t x_extent = hy.size(0); + update_magnetic_hy_standard_bounded_cuda( + hy, ex, ez, decay, curl, inv_dx, inv_dz, 0, x_extent, 0, x_extent); +} + +void update_magnetic_hz_standard_bounded_cuda( torch::stable::Tensor hz, const torch::stable::Tensor& ex, const torch::stable::Tensor& ey, const torch::stable::Tensor& decay, const torch::stable::Tensor& curl, const torch::stable::Tensor& inv_dx, - const torch::stable::Tensor& inv_dy) { + const torch::stable::Tensor& inv_dy, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent) { check_magnetic_inputs(hz, ex, ey, decay, curl, "hz"); check_rank3_shape(ex, "ex", hz.size(0), hz.size(1) + 1, hz.size(2)); check_rank3_shape(ey, "ey", hz.size(0) + 1, hz.size(1), hz.size(2)); check_spacing_vector(hz, inv_dx, 0, "inv_dx"); check_spacing_vector(hz, inv_dy, 1, "inv_dy"); + check_bounded_x_launch( + hz, local_x_begin, local_x_end, global_x_offset, global_x_extent, "hz"); + if (local_x_begin == local_x_end) { + return; + } torch::stable::accelerator::DeviceGuard guard(hz.get_device_index()); const auto sizes = hz.sizes(); const dim3 block = field_block3d(); - update_magnetic_standard_kernel<2><<>>( + update_magnetic_standard_kernel<2><<>>( static_cast(sizes[0]), static_cast(sizes[1]), static_cast(sizes[2]), + static_cast(local_x_begin), + static_cast(local_x_end), ex.mutable_data_ptr(), ey.mutable_data_ptr(), decay.mutable_data_ptr(), @@ -441,6 +512,19 @@ void update_magnetic_hz_standard_cuda( WITWIN_CUDA_CHECK(); } +void update_magnetic_hz_standard_cuda( + torch::stable::Tensor hz, + const torch::stable::Tensor& ex, + const torch::stable::Tensor& ey, + const torch::stable::Tensor& decay, + const torch::stable::Tensor& curl, + const torch::stable::Tensor& inv_dx, + const torch::stable::Tensor& inv_dy) { + const int64_t x_extent = hz.size(0); + update_magnetic_hz_standard_bounded_cuda( + hz, ex, ey, decay, curl, inv_dx, inv_dy, 0, x_extent, 0, x_extent); +} + void update_magnetic_hx_cpml_cuda( torch::stable::Tensor hx, const torch::stable::Tensor& ey, diff --git a/witwin/maxwell/fdtd/cuda/tensors.h b/witwin/maxwell/fdtd/cuda/tensors.h index 4562b2d..78d2186 100644 --- a/witwin/maxwell/fdtd/cuda/tensors.h +++ b/witwin/maxwell/fdtd/cuda/tensors.h @@ -33,3 +33,23 @@ inline void check_contiguous_tensor(const torch::stable::Tensor& tensor, const c STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); } +inline void check_bounded_x_launch( + const torch::stable::Tensor& field, + int64_t local_x_begin, + int64_t local_x_end, + int64_t global_x_offset, + int64_t global_x_extent, + const char* name) { + STD_TORCH_CHECK( + local_x_begin >= 0 && local_x_begin <= local_x_end && + local_x_end <= field.size(0), + name, + " local x launch interval must satisfy 0 <= begin <= end <= field.size(0)"); + STD_TORCH_CHECK(global_x_offset >= 0, name, " global x offset must be nonnegative"); + STD_TORCH_CHECK(global_x_extent > 0, name, " global x extent must be positive"); + STD_TORCH_CHECK( + global_x_offset <= global_x_extent && + field.size(0) <= global_x_extent - global_x_offset, + name, + " padded local x allocation must fit within the component global x extent"); +} diff --git a/witwin/maxwell/fdtd/distributed/__init__.py b/witwin/maxwell/fdtd/distributed/__init__.py new file mode 100644 index 0000000..5277a4a --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/__init__.py @@ -0,0 +1,10 @@ +"""Internal single-node distributed FDTD runtime. + +Only :class:`~witwin.maxwell.fdtd_parallel.FDTDParallelConfig` is public. The +objects in this package deliberately remain implementation details so transport +handles and mutable rank-local state never become part of the Result contract. +""" + +from .solver import DistributedFDTD + +__all__ = ["DistributedFDTD"] diff --git a/witwin/maxwell/fdtd/distributed/capacity.py b/witwin/maxwell/fdtd/distributed/capacity.py new file mode 100644 index 0000000..75019c8 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/capacity.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import torch + +from .output import effective_cuda_free_bytes, electric_field_output_bytes + + +def _frequency_count(dft_frequency) -> int: + if isinstance(dft_frequency, (tuple, list)): + return max(len(dft_frequency), 1) + return 1 + + +def local_dft_working_set_bytes( + local_shape: tuple[int, int, int], + *, + dft_frequency, + full_field_dft: bool, +) -> int: + """Budget real/imag accumulators plus the postprocessed complex local fields.""" + + if dft_frequency is None or not full_field_dft: + return 0 + complex_field_bytes = electric_field_output_bytes( + local_shape, + frequency_count=_frequency_count(dft_frequency), + complex_output=True, + ) + return 2 * complex_field_bytes + + +def _require_capacity(device: torch.device, required: int, *, operation: str) -> dict[str, int]: + available = effective_cuda_free_bytes(device) + if required * 1.05 > available: + raise MemoryError( + f"{operation} requires {required} bytes on {device}, but only {available} " + "effective bytes are free. Disable full-field DFT/gathering or choose a device " + "with more free memory." + ) + return {"required_bytes": int(required), "available_bytes": int(available)} + + +def require_local_dft_capacity( + device: torch.device, + local_shape: tuple[int, int, int], + *, + dft_frequency, + full_field_dft: bool, +) -> dict[str, int]: + pending = local_dft_working_set_bytes( + local_shape, + dft_frequency=dft_frequency, + full_field_dft=full_field_dft, + ) + result = _require_capacity(device, pending, operation="Local multi-GPU DFT working set") + result["pending_local_dft_bytes"] = int(pending) + return result + + +def require_gather_capacity( + device: torch.device, + global_shape: tuple[int, int, int], + *, + dft_frequency, + full_field_dft: bool, + pending_local_dft_bytes: int = 0, +) -> dict[str, int]: + """Fail before stepping when the result device's combined working set cannot fit.""" + + complex_output = bool(dft_frequency is not None and full_field_dft) + output_bytes = electric_field_output_bytes( + global_shape, + frequency_count=_frequency_count(dft_frequency), + complex_output=complex_output, + ) + required = output_bytes + max(int(pending_local_dft_bytes), 0) + result = _require_capacity(device, required, operation="Explicit multi-GPU field gathering") + result["output_bytes"] = int(output_bytes) + result["pending_local_dft_bytes"] = int(pending_local_dft_bytes) + return result + + +__all__ = [ + "local_dft_working_set_bytes", + "require_gather_capacity", + "require_local_dft_capacity", +] diff --git a/witwin/maxwell/fdtd/distributed/counts.py b/witwin/maxwell/fdtd/distributed/counts.py new file mode 100644 index 0000000..6ace8fa --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/counts.py @@ -0,0 +1,25 @@ +from __future__ import annotations + + +def reduce_sample_counts(shards, attribute: str) -> tuple[int, ...]: + """Reduce per-shard spectral counters without double-counting spatial tiles.""" + + if not shards: + return () + sequences = tuple( + tuple(int(value) for value in (getattr(shard.solver, attribute, ()) or ())) + for shard in shards + ) + width = max((len(sequence) for sequence in sequences), default=0) + if width == 0: + return () + return tuple( + max( + (sequence[index] if index < len(sequence) else 0) + for sequence in sequences + ) + for index in range(width) + ) + + +__all__ = ["reduce_sample_counts"] diff --git a/witwin/maxwell/fdtd/distributed/frequency_counts.py b/witwin/maxwell/fdtd/distributed/frequency_counts.py new file mode 100644 index 0000000..58d0a75 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/frequency_counts.py @@ -0,0 +1,24 @@ +from __future__ import annotations + + +def reduce_frequency_sample_counts(shards, frequencies) -> tuple[int, ...]: + """Reduce observer counters by frequency rather than shard-local entry index.""" + + requested = tuple(float(frequency) for frequency in frequencies) + if not requested: + return () + merged = {frequency: 0 for frequency in requested} + for shard in shards: + local_frequencies = tuple( + float(frequency) for frequency in (shard.solver.observer_frequencies or ()) + ) + local_counts = tuple( + int(count) for count in (shard.solver.observer_sample_counts or ()) + ) + for frequency, count in zip(local_frequencies, local_counts): + if frequency in merged: + merged[frequency] = max(merged[frequency], count) + return tuple(merged[frequency] for frequency in requested) + + +__all__ = ["reduce_frequency_sample_counts"] diff --git a/witwin/maxwell/fdtd/distributed/interpolation.py b/witwin/maxwell/fdtd/distributed/interpolation.py new file mode 100644 index 0000000..a86f2b9 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/interpolation.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import numpy as np + + +def ideal_axis_weight(coords64: np.ndarray, position: float, index: int) -> float: + """Match the native point-source compiler's float32 linear-axis weight.""" + + coords = np.asarray(coords64, dtype=np.float32) + probe = np.float32(position) + upper = int(np.searchsorted(coords, probe, side="left")) + if upper <= 0: + return 1.0 if index == 0 else 0.0 + if upper >= coords.size: + return 1.0 if index == coords.size - 1 else 0.0 + + lower = upper - 1 + lower_coord = float(coords[lower]) + upper_coord = float(coords[upper]) + span = upper_coord - lower_coord + if span <= np.finfo(np.float32).eps: + return 1.0 if index == lower else 0.0 + fraction = (float(position) - lower_coord) / span + if fraction <= 1.0e-12: + return 1.0 if index == lower else 0.0 + if fraction >= 1.0 - 1.0e-12: + return 1.0 if index == upper else 0.0 + if index == lower: + return 1.0 - fraction + if index == upper: + return fraction + return 0.0 + + +__all__ = ["ideal_axis_weight"] diff --git a/witwin/maxwell/fdtd/distributed/monitor_merge.py b/witwin/maxwell/fdtd/distributed/monitor_merge.py new file mode 100644 index 0000000..6905634 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/monitor_merge.py @@ -0,0 +1,416 @@ +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +import numpy as np +import torch + +from ...fdtd_parallel import FDTDShardLayout +from ..observers import _align_plane_monitor_payload, _compute_plane_flux +from .output import move_tensors_to_device + + +_PLANE_COORD_NAMES = { + "x": ("y", "z"), + "y": ("x", "z"), + "z": ("x", "y"), +} + + +def _validate_matching_metadata(entries: list[tuple[int, Mapping[str, Any]]]) -> None: + first = entries[0][1] + for rank, payload in entries[1:]: + for key in ( + "kind", + "monitor_type", + "fields", + "frequencies", + "axis", + "position", + "compute_flux", + "normal_direction", + ): + if payload.get(key) != first.get(key): + raise ValueError( + f"Shard {rank} monitor metadata {key!r} does not match rank " + f"{entries[0][0]}." + ) + + +def _normalize_shard_layouts( + shard_layouts: Mapping[int, FDTDShardLayout] | Iterable[FDTDShardLayout] | None, +) -> dict[int, FDTDShardLayout]: + if shard_layouts is None: + return {} + values = shard_layouts.values() if isinstance(shard_layouts, Mapping) else shard_layouts + normalized: dict[int, FDTDShardLayout] = {} + for layout in values: + if not isinstance(layout, FDTDShardLayout): + raise TypeError("shard_layouts must contain FDTDShardLayout instances.") + if layout.rank in normalized: + raise ValueError(f"shard_layouts contains duplicate rank {layout.rank}.") + normalized[layout.rank] = layout + return normalized + + +def _slice_tensor_x(tensor: torch.Tensor, x_slice: slice) -> torch.Tensor: + index = [slice(None)] * tensor.ndim + index[-2] = x_slice + return tensor[tuple(index)] + + +def _stitch_owned_component( + entries: list[tuple[int, Mapping[str, Any]]], + component: str, + *, + shard_layouts: Mapping[int, FDTDShardLayout], + result_device: torch.device, +) -> dict[str, Any]: + owned_tiles = [] + expected_transverse = None + expected_global_extent = None + for rank, monitor_payload in entries: + component_payloads = monitor_payload.get("components", {}) + if component not in component_payloads: + raise ValueError(f"Shard {rank} is missing plane component {component!r}.") + component_payload = component_payloads[component] + data = component_payload.get("data") + if not isinstance(data, torch.Tensor): + raise TypeError(f"Shard {rank} plane component {component!r} data must be a tensor.") + if data.ndim < 2: + raise ValueError( + f"Shard {rank} plane component {component!r} data must have at least two dimensions." + ) + coords = component_payload.get("coords", ()) + if len(coords) != 2: + raise ValueError( + f"Shard {rank} component {component!r} requires two plane coordinate arrays." + ) + x_values = np.asarray(coords[0], dtype=np.float64) + transverse = np.asarray(coords[1], dtype=np.float64) + + component_layout = shard_layouts[rank].component(component) + local_x_extent = int(component_layout.local_shape[0]) + if x_values.size != local_x_extent or int(data.shape[-2]) != local_x_extent: + raise ValueError( + f"Shard {rank} component {component!r} local x extent does not match " + f"its layout: coords={x_values.size}, data={int(data.shape[-2])}, " + f"layout={local_x_extent}." + ) + if int(data.shape[-1]) != transverse.size: + raise ValueError( + f"Shard {rank} component {component!r} transverse coordinates do not " + f"match data shape {tuple(data.shape)}." + ) + if expected_transverse is None: + expected_transverse = np.array(transverse, copy=True) + elif not np.array_equal(transverse, expected_transverse): + raise ValueError( + f"Shard {rank} component {component!r} transverse coordinates do not match." + ) + + owned_local = component_layout.owned_local_slice[0] + owned_global = component_layout.owned_global_slice[0] + owned_count = int(owned_global.stop) - int(owned_global.start) + cropped_x = np.array(x_values[owned_local], copy=True) + cropped_data = _slice_tensor_x(data, owned_local) + if cropped_x.size != owned_count or int(cropped_data.shape[-2]) != owned_count: + raise RuntimeError( + f"Shard {rank} component {component!r} owned local/global x extents disagree." + ) + global_extent = int(component_layout.global_shape[0]) + if expected_global_extent is None: + expected_global_extent = global_extent + elif expected_global_extent != global_extent: + raise ValueError( + f"Shard {rank} component {component!r} global x extent is inconsistent." + ) + owned_tiles.append( + ( + int(owned_global.start), + int(owned_global.stop), + rank, + cropped_data.to(device=result_device, non_blocking=True), + cropped_x, + component_payload, + ) + ) + + owned_tiles.sort(key=lambda item: (item[0], item[2])) + cursor = 0 + for begin, end, rank, _data, _x, _payload in owned_tiles: + if begin != cursor: + relation = "overlap" if begin < cursor else "gap" + raise ValueError( + f"Component {component!r} owned global x slices contain a {relation} " + f"before shard {rank}: expected {cursor}, got {begin}." + ) + cursor = end + if expected_global_extent is None or cursor != expected_global_extent: + raise ValueError( + f"Component {component!r} owned global x coverage ends at {cursor}, " + f"expected {expected_global_extent}." + ) + + merged_data = torch.cat([item[3] for item in owned_tiles], dim=-2) + merged_x = np.concatenate([item[4] for item in owned_tiles]) + if merged_x.size > 1 and not np.all(np.diff(merged_x) > 0.0): + raise ValueError( + f"Component {component!r} owned global x coordinates are not strictly increasing." + ) + template = { + key: value + for key, value in owned_tiles[0][5].items() + if key not in {"data", "coords"} + } + merged = move_tensors_to_device(template, result_device) + merged["data"] = merged_data + merged["coords"] = (merged_x, expected_transverse) + return merged + + +def _merge_owned_components( + entries: list[tuple[int, Mapping[str, Any]]], + *, + shard_layouts: Mapping[int, FDTDShardLayout], + result_device: torch.device, +) -> dict[str, Any]: + first_components = entries[0][1].get("components", {}) + if not isinstance(first_components, Mapping) or not first_components: + raise ValueError("A tiled plane monitor requires raw Yee component payloads.") + expected = tuple(first_components) + for rank, payload in entries[1:]: + if tuple(payload.get("components", {})) != expected: + raise ValueError( + f"Shard {rank} plane component order/set does not match rank {entries[0][0]}." + ) + return { + component: _stitch_owned_component( + entries, + component, + shard_layouts=shard_layouts, + result_device=result_device, + ) + for component in expected + } + + +def _plane_metadata_template(payload: Mapping[str, Any]) -> dict[str, Any]: + excluded = { + "components", + "component", + "data", + "coords", + "x", + "y", + "z", + "flux", + "power", + "Ex", + "Ey", + "Ez", + "Hx", + "Hy", + "Hz", + } + return {key: value for key, value in payload.items() if key not in excluded} + + +def _normalize_physical_bounds(physical_bounds): + if physical_bounds is None or len(physical_bounds) != 3: + raise ValueError( + "physical_bounds must contain the three logical-domain (low, high) pairs." + ) + normalized = [] + for bounds in physical_bounds: + if len(bounds) != 2: + raise ValueError("Each physical_bounds entry must be a (low, high) pair.") + low, high = float(bounds[0]), float(bounds[1]) + if not np.isfinite(low) or not np.isfinite(high) or high < low: + raise ValueError("physical_bounds entries must be finite and ordered.") + normalized.append((low, high)) + return tuple(normalized) + + +def _physical_coord_indices(values, bounds: tuple[float, float]) -> np.ndarray: + coords = np.asarray(values, dtype=np.float64) + low, high = bounds + coordinate_magnitude = max( + abs(low), + abs(high), + float(np.max(np.abs(coords))) if coords.size else 0.0, + ) + domain_span = abs(high - low) + unique = np.unique(coords) + positive_spacing = np.diff(unique) + positive_spacing = positive_spacing[positive_spacing > 0.0] + local_spacing = ( + float(np.min(positive_spacing)) if positive_spacing.size else domain_span + ) + scale = max( + coordinate_magnitude, + domain_span, + local_spacing, + np.finfo(np.float64).tiny, + ) + tolerance = 64.0 * np.finfo(np.float64).eps * scale + indices = np.flatnonzero( + (coords >= low - tolerance) & (coords <= high + tolerance) + ) + if indices.size == 0: + raise ValueError("FluxMonitor plane has no samples inside the physical domain.") + return indices + + +def _index_plane_values(values, first: np.ndarray, second: np.ndarray): + if isinstance(values, torch.Tensor): + first_index = torch.as_tensor(first, device=values.device, dtype=torch.long) + second_index = torch.as_tensor(second, device=values.device, dtype=torch.long) + selected = torch.index_select(values, values.ndim - 2, first_index) + return torch.index_select(selected, selected.ndim - 1, second_index) + selected = np.take(np.asarray(values), first, axis=np.ndim(values) - 2) + return np.take(selected, second, axis=np.ndim(selected) - 1) + + +def _crop_aligned_to_physical_bounds(axis: str, aligned, physical_bounds): + bounds = _normalize_physical_bounds(physical_bounds) + coord_names = _PLANE_COORD_NAMES[axis] + axis_index = {"x": 0, "y": 1, "z": 2} + first_coord = np.asarray(aligned[coord_names[0]], dtype=np.float64) + second_coord = np.asarray(aligned[coord_names[1]], dtype=np.float64) + first_indices = _physical_coord_indices( + first_coord, + bounds[axis_index[coord_names[0]]], + ) + second_indices = _physical_coord_indices( + second_coord, + bounds[axis_index[coord_names[1]]], + ) + return { + coord_names[0]: np.array(first_coord[first_indices], copy=True), + coord_names[1]: np.array(second_coord[second_indices], copy=True), + "fields": { + component: _index_plane_values(values, first_indices, second_indices) + for component, values in aligned["fields"].items() + }, + } + + +def _merge_tiled_plane( + entries: list[tuple[int, Mapping[str, Any]]], + *, + shard_layouts: Mapping[int, FDTDShardLayout], + physical_bounds, + result_device: torch.device, +) -> dict[str, Any]: + _validate_matching_metadata(entries) + first = entries[0][1] + axis = str(first.get("axis", "")).lower() + if axis not in {"y", "z"}: + raise ValueError(f"Only y/z-normal planes are tiled across x, got axis={axis!r}.") + + missing_layouts = sorted(rank for rank, _payload in entries if rank not in shard_layouts) + if missing_layouts: + ranks = ", ".join(str(rank) for rank in missing_layouts) + raise ValueError(f"Tiled plane merge is missing shard layouts for rank {ranks}.") + + merged = move_tensors_to_device(_plane_metadata_template(first), result_device) + components = _merge_owned_components( + entries, + shard_layouts=shard_layouts, + result_device=result_device, + ) + merged["components"] = components + + fields = tuple(first.get("fields", ())) + coord_names = _PLANE_COORD_NAMES[axis] + if len(fields) == 1: + if bool(first.get("compute_flux", False)): + raise ValueError("FluxMonitor requires multiple tangential E/H components.") + component = fields[0] + component_payload = components[component] + merged["component"] = component.lower() + merged["data"] = component_payload["data"] + merged[component] = component_payload["data"] + merged["coords"] = component_payload["coords"] + merged[coord_names[0]] = component_payload["coords"][0] + merged[coord_names[1]] = component_payload["coords"][1] + return merged + + aligned = _align_plane_monitor_payload(axis, components) + if aligned is None: + if bool(first.get("compute_flux", False)): + raise ValueError( + "FluxMonitor is missing aligned tangential fields required for integration." + ) + return merged + if bool(first.get("compute_flux", False)): + aligned = _crop_aligned_to_physical_bounds(axis, aligned, physical_bounds) + + merged[coord_names[0]] = aligned[coord_names[0]] + merged[coord_names[1]] = aligned[coord_names[1]] + merged["coords"] = (aligned[coord_names[0]], aligned[coord_names[1]]) + for component, values in aligned["fields"].items(): + merged[component] = values + + if bool(first.get("compute_flux", False)): + flux = _compute_plane_flux(merged) + merged["flux"] = flux + merged["power"] = flux + return merged + + +def merge_sharded_monitor_payloads( + monitor_order: Iterable[str], + shard_payloads: Iterable[tuple[int, Mapping[str, Mapping[str, Any]]]], + *, + result_device: str | torch.device, + shard_layouts: ( + Mapping[int, FDTDShardLayout] | Iterable[FDTDShardLayout] | None + ) = None, + physical_bounds=None, +) -> dict[str, Any]: + """Merge rank-local monitor payloads in scene declaration order. + + Point monitors and x-normal planes have exactly one owner. For y/z-normal + planes, every raw Yee component is first cropped by its component-specific + owned-local x slice and concatenated by owned-global x slice. Multi-component + planes are aligned only after global assembly. Flux is recomputed on the + result device after cropping to the logical physical bounds. + """ + + device = torch.device(result_device) + layouts = _normalize_shard_layouts(shard_layouts) + by_name: dict[str, list[tuple[int, Mapping[str, Any]]]] = {} + for rank, payloads in sorted(shard_payloads, key=lambda item: int(item[0])): + for name, payload in payloads.items(): + by_name.setdefault(str(name), []).append((int(rank), payload)) + + merged = {} + for name in monitor_order: + entries = by_name.pop(str(name), []) + if not entries: + continue + first = entries[0][1] + if first.get("kind") == "plane" and first.get("axis") in {"y", "z"}: + merged[str(name)] = _merge_tiled_plane( + entries, + shard_layouts=layouts, + physical_bounds=physical_bounds, + result_device=device, + ) + else: + if len(entries) != 1: + axis = first.get("axis") + descriptor = "x-normal plane" if axis == "x" else "monitor" + raise RuntimeError(f"{descriptor} {name!r} has more than one shard owner.") + merged[str(name)] = move_tensors_to_device(entries[0][1], device) + + if by_name: + undeclared = ", ".join(sorted(by_name)) + raise ValueError(f"Shard payloads contain undeclared monitors: {undeclared}.") + return merged + + +__all__ = ["merge_sharded_monitor_payloads"] diff --git a/witwin/maxwell/fdtd/distributed/output.py b/witwin/maxwell/fdtd/distributed/output.py new file mode 100644 index 0000000..6f94048 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/output.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from typing import Any + +import torch + + +def move_tensors_to_device(value: Any, device: torch.device) -> Any: + """Move a nested, small result payload directly between CUDA devices.""" + + if isinstance(value, torch.Tensor): + return value.to(device=device, non_blocking=True) + if isinstance(value, dict): + return {key: move_tensors_to_device(item, device) for key, item in value.items()} + if isinstance(value, list): + return [move_tensors_to_device(item, device) for item in value] + if isinstance(value, tuple): + return tuple(move_tensors_to_device(item, device) for item in value) + return value + + +def electric_field_output_bytes( + global_shape: tuple[int, int, int], + *, + frequency_count: int = 1, + complex_output: bool = False, +) -> int: + nx, ny, nz = (int(value) for value in global_shape) + elements = ( + (nx - 1) * ny * nz + + nx * (ny - 1) * nz + + nx * ny * (nz - 1) + ) + element_size = 8 if complex_output else 4 + return elements * max(int(frequency_count), 1) * element_size + + +def effective_cuda_free_bytes(device: torch.device) -> int: + with torch.cuda.device(device): + driver_free, _total = torch.cuda.mem_get_info(device) + reusable_cache = torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device) + return int(driver_free + max(int(reusable_cache), 0)) + + +__all__ = [ + "effective_cuda_free_bytes", + "electric_field_output_bytes", + "move_tensors_to_device", +] diff --git a/witwin/maxwell/fdtd/distributed/persistence.py b/witwin/maxwell/fdtd/distributed/persistence.py new file mode 100644 index 0000000..d08b8ad --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/persistence.py @@ -0,0 +1,879 @@ +from __future__ import annotations + +import json +import os +import shutil +import tempfile +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + + +_SCHEMA = "witwin.maxwell.fdtd.sharded-result" +_FORMAT_VERSION = 1 +_ELECTRIC_COMPONENTS = ("Ex", "Ey", "Ez") + + +def _component_name(value: str) -> str: + text = str(value).strip() + normalized = text[0].upper() + text[1:].lower() if text else text + if normalized not in _ELECTRIC_COMPONENTS: + choices = ", ".join(_ELECTRIC_COMPONENTS) + raise ValueError( + f"Sharded Result component must be one of {choices}, got {value!r}." + ) + return normalized + + +def _half_open_slice(value, *, label: str) -> tuple[int, int]: + if isinstance(value, slice): + if value.step not in (None, 1): + raise ValueError(f"{label} must have unit stride.") + start, stop = value.start, value.stop + else: + try: + start, stop = value + except (TypeError, ValueError) as exc: + raise TypeError(f"{label} must be a (start, stop) pair or slice.") from exc + if isinstance(start, bool) or isinstance(stop, bool): + raise TypeError(f"{label} endpoints must be integers.") + start, stop = int(start), int(stop) + if start < 0 or stop <= start: + raise ValueError( + f"{label} must be a nonempty half-open interval, got ({start}, {stop})." + ) + return start, stop + + +def _frequencies(values) -> tuple[float, ...]: + result = tuple(float(value) for value in values) + if not result: + raise ValueError("Sharded Result frequencies must not be empty.") + return result + + +def _cpu_data(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, Mapping): + return {key: _cpu_data(item) for key, item in value.items()} + if isinstance(value, list): + return [_cpu_data(item) for item in value] + if isinstance(value, tuple): + return tuple(_cpu_data(item) for item in value) + return value + + +@dataclass(frozen=True) +class FieldComponentArtifact: + """Owned-only tensor and global x interval for one Yee E component.""" + + name: str + tensor: torch.Tensor + global_x_slice: tuple[int, int] + + def __post_init__(self) -> None: + object.__setattr__(self, "name", _component_name(self.name)) + if not isinstance(self.tensor, torch.Tensor): + raise TypeError(f"{self.name} shard tensor must be a torch.Tensor.") + if self.tensor.ndim < 3: + raise ValueError( + f"{self.name} shard tensor must have at least three field dimensions." + ) + extent = _half_open_slice( + self.global_x_slice, + label=f"{self.name} global_x_slice", + ) + object.__setattr__(self, "global_x_slice", extent) + x_axis = self.tensor.ndim - 3 + if int(self.tensor.shape[x_axis]) != extent[1] - extent[0]: + raise ValueError( + f"{self.name} tensor x size {self.tensor.shape[x_axis]} does not match " + f"owned interval {extent}." + ) + + @property + def x_axis(self) -> int: + return self.tensor.ndim - 3 + + +@dataclass(frozen=True) +class FieldShardArtifact: + """Pure-data export produced by ``solver.export_field_shards()``.""" + + rank: int + components: tuple[FieldComponentArtifact, ...] + frequencies: tuple[float, ...] + device: str | None = None + + def __post_init__(self) -> None: + if isinstance(self.rank, bool) or int(self.rank) < 0: + raise ValueError(f"Shard rank must be a nonnegative integer, got {self.rank!r}.") + object.__setattr__(self, "rank", int(self.rank)) + components = tuple(self.components) + if not components: + raise ValueError(f"Shard rank {self.rank} has no field components.") + if len({component.name for component in components}) != len(components): + raise ValueError(f"Shard rank {self.rank} contains duplicate field components.") + object.__setattr__(self, "components", components) + object.__setattr__(self, "frequencies", _frequencies(self.frequencies)) + if self.device is not None: + object.__setattr__(self, "device", str(self.device)) + + +def _export_frequency_values(values: Any, *, label: str) -> tuple[float, ...]: + if isinstance(values, torch.Tensor): + values = values.detach().cpu().reshape(-1).tolist() + elif isinstance(values, (int, float)): + values = (values,) + try: + return _frequencies(values) + except (TypeError, ValueError) as exc: + raise ValueError(f"{label} must contain one or more numeric values.") from exc + + +def _dft_frequencies( + local_solver: Any, + output: Mapping[str, Any], + *, + rank: int, +) -> tuple[float, ...]: + if "frequencies" in output: + values = output["frequencies"] + else: + values = getattr(local_solver, "dft_frequencies", ()) + if not values: + frequency = getattr(local_solver, "dft_frequency", None) + values = () if frequency is None else (frequency,) + return _export_frequency_values(values, label=f"Shard rank {rank} DFT frequencies") + + +def _owned_component( + *, + name: str, + tensor: Any, + layout: Any, +) -> FieldComponentArtifact: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"Shard-local {name} field must be a torch.Tensor.") + if tensor.ndim < 3: + raise ValueError(f"Shard-local {name} field must have at least three dimensions.") + is_cell = name == "Ex" + storage_name = "storage_cell_owned" if is_cell else "storage_node_owned" + global_name = "global_cell_owned" if is_cell else "global_node_owned" + storage = _half_open_slice( + getattr(layout, storage_name), + label=f"{name} {storage_name}", + ) + global_extent = _half_open_slice( + getattr(layout, global_name), + label=f"{name} {global_name}", + ) + if storage[1] - storage[0] != global_extent[1] - global_extent[0]: + raise ValueError( + f"{name} storage owned interval {storage} and global owned interval " + f"{global_extent} have different widths." + ) + x_axis = tensor.ndim - 3 + if storage[1] > int(tensor.shape[x_axis]): + raise ValueError( + f"{name} storage owned interval {storage} exceeds local x size " + f"{tensor.shape[x_axis]}." + ) + index = [slice(None)] * tensor.ndim + index[x_axis] = slice(*storage) + return FieldComponentArtifact( + name=name, + tensor=tensor[tuple(index)].contiguous(), + global_x_slice=global_extent, + ) + + +def export_distributed_field_shards(solver: Any) -> tuple[FieldShardArtifact, ...]: + """Export owned Ex/Ey/Ez regions from a distributed solver without gathering.""" + + shards = tuple(getattr(solver, "shards", ())) + if not shards: + raise ValueError("Distributed solver has no shards to export.") + ordered = tuple(sorted(shards, key=lambda shard: int(shard.rank))) + modes = tuple(bool(shard.solver.dft_enabled) for shard in ordered) + if len(set(modes)) != 1: + raise RuntimeError("Shard-local field export mixes DFT and last-step modes.") + + expected_frequencies: tuple[float, ...] | None = None + if not modes[0]: + expected_frequencies = _export_frequency_values( + (solver.frequency,), + label="Distributed solver frequency", + ) + + artifacts = [] + for shard in ordered: + rank = int(shard.rank) + local_solver = shard.solver + if modes[0]: + output = local_solver.get_frequency_solution(all_frequencies=True) + if not isinstance(output, Mapping): + raise TypeError(f"Shard rank {rank} DFT solution must be a mapping.") + frequencies = _dft_frequencies(local_solver, output, rank=rank) + if expected_frequencies is None: + expected_frequencies = frequencies + elif frequencies != expected_frequencies: + raise RuntimeError( + f"Shard rank {rank} DFT frequencies {frequencies} do not match " + f"{expected_frequencies}." + ) + fields = output + else: + frequencies = expected_frequencies + fields = { + name: getattr(local_solver, name) for name in _ELECTRIC_COMPONENTS + } + + components = [] + for name in _ELECTRIC_COMPONENTS: + if name not in fields: + raise KeyError(f"Shard rank {rank} field export is missing {name}.") + component = _owned_component( + name=name, + tensor=fields[name], + layout=shard.layout, + ) + _validate_frequency_axis(component, frequencies) + components.append(component) + artifacts.append( + FieldShardArtifact( + rank=rank, + components=tuple(components), + frequencies=frequencies, + device=getattr(shard, "device", None), + ) + ) + return tuple(artifacts) + + +@dataclass(frozen=True) +class FieldComponentManifest: + name: str + global_x_slice: tuple[int, int] + shape: tuple[int, ...] + dtype: str + x_axis: int + frequencies: tuple[float, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "global_x_slice": self.global_x_slice, + "shape": self.shape, + "dtype": self.dtype, + "x_axis": self.x_axis, + "frequencies": self.frequencies, + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "FieldComponentManifest": + return cls( + name=_component_name(payload["name"]), + global_x_slice=_half_open_slice( + payload["global_x_slice"], + label=f"{payload['name']} manifest global_x_slice", + ), + shape=tuple(int(value) for value in payload["shape"]), + dtype=str(payload["dtype"]), + x_axis=int(payload["x_axis"]), + frequencies=_frequencies(payload["frequencies"]), + ) + + +@dataclass(frozen=True) +class FieldShardManifest: + rank: int + file: str + device: str + components: tuple[FieldComponentManifest, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "rank": self.rank, + "file": self.file, + "device": self.device, + "components": tuple(component.to_dict() for component in self.components), + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "FieldShardManifest": + rank = int(payload["rank"]) + expected_file = f"rank-{rank:04d}.pt" + file_name = str(payload["file"]) + if file_name != expected_file: + raise ValueError( + f"Shard rank {rank} file must be {expected_file!r}, got {file_name!r}." + ) + return cls( + rank=rank, + file=file_name, + device=str(payload["device"]), + components=tuple( + FieldComponentManifest.from_dict(component) + for component in payload["components"] + ), + ) + + +@dataclass(frozen=True) +class GlobalFieldManifest: + name: str + shape: tuple[int, ...] + dtype: str + x_axis: int + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "shape": self.shape, + "dtype": self.dtype, + "x_axis": self.x_axis, + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "GlobalFieldManifest": + return cls( + name=_component_name(payload["name"]), + shape=tuple(int(value) for value in payload["shape"]), + dtype=str(payload["dtype"]), + x_axis=int(payload["x_axis"]), + ) + + +@dataclass(frozen=True) +class DistributedResultManifest: + schema: str + format_version: int + result_file: str + frequencies: tuple[float, ...] + components: tuple[GlobalFieldManifest, ...] + shards: tuple[FieldShardManifest, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "schema": self.schema, + "format_version": self.format_version, + "result_file": self.result_file, + "frequencies": self.frequencies, + "components": tuple(component.to_dict() for component in self.components), + "shards": tuple(shard.to_dict() for shard in self.shards), + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "DistributedResultManifest": + schema = str(payload["schema"]) + version = int(payload["format_version"]) + result_file = str(payload["result_file"]) + if schema != _SCHEMA: + raise ValueError(f"Unsupported sharded Result schema {schema!r}.") + if version != _FORMAT_VERSION: + raise ValueError(f"Unsupported sharded Result format_version={version}.") + if result_file != "result.pt": + raise ValueError("Sharded Result manifest result_file must be 'result.pt'.") + manifest = cls( + schema=schema, + format_version=version, + result_file=result_file, + frequencies=_frequencies(payload["frequencies"]), + components=tuple( + GlobalFieldManifest.from_dict(component) + for component in payload["components"] + ), + shards=tuple( + FieldShardManifest.from_dict(shard) for shard in payload["shards"] + ), + ) + _validate_manifest_structure(manifest) + return manifest + + @classmethod + def read(cls, directory: str | Path) -> "DistributedResultManifest": + path = Path(directory) / "manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"Sharded Result manifest is missing: {path}.") + try: + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise TypeError("manifest root is not a mapping") + return cls.from_dict(payload) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + raise ValueError(f"Invalid sharded Result manifest {path}: {exc}") from exc + + def shard_paths(self, directory: str | Path) -> tuple[Path, ...]: + root = Path(directory) + return tuple(root / shard.file for shard in self.shards) + + +@dataclass(frozen=True) +class LoadedShardedResult: + manifest: DistributedResultManifest + result_payload: dict[str, Any] + shard_paths: tuple[Path, ...] + fields: dict[str, torch.Tensor] + + +def _coerce_component(name: str, value: Any) -> FieldComponentArtifact: + if isinstance(value, FieldComponentArtifact): + if _component_name(name) != value.name: + raise ValueError( + f"Component mapping key {name!r} does not match artifact {value.name!r}." + ) + return value + if not isinstance(value, Mapping): + raise TypeError(f"Shard component {name!r} must be a mapping or artifact.") + return FieldComponentArtifact( + name=name, + tensor=value["tensor"], + global_x_slice=value["global_x_slice"], + ) + + +def _coerce_shard(value: Any) -> FieldShardArtifact: + if isinstance(value, FieldShardArtifact): + return value + if not isinstance(value, Mapping): + raise TypeError("export_field_shards() entries must be mappings or FieldShardArtifact.") + components_value = value["components"] + if isinstance(components_value, Mapping): + components = tuple( + _coerce_component(name, component) + for name, component in components_value.items() + ) + else: + components = tuple(components_value) + return FieldShardArtifact( + rank=value["rank"], + device=value.get("device"), + frequencies=value["frequencies"], + components=components, + ) + + +def _validate_frequency_axis( + component: FieldComponentArtifact, + frequencies: tuple[float, ...], +) -> None: + if len(frequencies) <= 1: + return + if component.tensor.ndim < 4 or int(component.tensor.shape[0]) != len(frequencies): + raise ValueError( + f"{component.name} multi-frequency tensor shape {tuple(component.tensor.shape)} " + f"does not start with frequency count {len(frequencies)}." + ) + + +def _build_manifest( + shards: Iterable[FieldShardArtifact | Mapping[str, Any]], + *, + frequencies: tuple[float, ...], +) -> tuple[DistributedResultManifest, tuple[FieldShardArtifact, ...]]: + expected_frequencies = _frequencies(frequencies) + ordered = tuple( + sorted( + (_coerce_shard(shard) for shard in shards), + key=lambda item: item.rank, + ) + ) + if not ordered: + raise ValueError("export_field_shards() returned no shard artifacts.") + expected_ranks = tuple(range(len(ordered))) + ranks = tuple(shard.rank for shard in ordered) + if ranks != expected_ranks: + raise ValueError(f"Shard ranks must be contiguous from zero, got {ranks}.") + + component_names = tuple(component.name for component in ordered[0].components) + for shard in ordered: + if shard.frequencies != expected_frequencies: + raise ValueError( + f"Shard rank {shard.rank} frequencies {shard.frequencies} do not match " + f"Result frequencies {expected_frequencies}." + ) + names = tuple(component.name for component in shard.components) + if names != component_names: + raise ValueError( + f"Shard rank {shard.rank} component order {names} does not match " + f"rank 0 order {component_names}." + ) + + global_components = [] + for component_index, name in enumerate(component_names): + entries = tuple(shard.components[component_index] for shard in ordered) + first = entries[0] + _validate_frequency_axis(first, expected_frequencies) + expected_start = 0 + for shard, component in zip(ordered, entries): + _validate_frequency_axis(component, expected_frequencies) + start, stop = component.global_x_slice + if start != expected_start: + relation = "overlaps" if start < expected_start else "leaves a gap before" + raise ValueError( + f"Shard rank {shard.rank} {name} interval {(start, stop)} {relation} " + f"global x={expected_start}." + ) + if component.x_axis != first.x_axis: + raise ValueError(f"Shard rank {shard.rank} {name} x axis is inconsistent.") + if component.tensor.dtype != first.tensor.dtype: + raise ValueError(f"Shard rank {shard.rank} {name} dtype is inconsistent.") + for axis, (actual, reference) in enumerate( + zip(component.tensor.shape, first.tensor.shape) + ): + if axis != first.x_axis and int(actual) != int(reference): + raise ValueError( + f"Shard rank {shard.rank} {name} non-x shape is inconsistent." + ) + expected_start = stop + global_shape = list(first.tensor.shape) + global_shape[first.x_axis] = expected_start + global_components.append( + GlobalFieldManifest( + name=name, + shape=tuple(int(value) for value in global_shape), + dtype=str(first.tensor.dtype), + x_axis=first.x_axis, + ) + ) + + shard_manifests = [] + for shard in ordered: + source_device = shard.device or str(shard.components[0].tensor.device) + shard_manifests.append( + FieldShardManifest( + rank=shard.rank, + file=f"rank-{shard.rank:04d}.pt", + device=source_device, + components=tuple( + FieldComponentManifest( + name=component.name, + global_x_slice=component.global_x_slice, + shape=tuple(int(value) for value in component.tensor.shape), + dtype=str(component.tensor.dtype), + x_axis=component.x_axis, + frequencies=expected_frequencies, + ) + for component in shard.components + ), + ) + ) + + manifest = DistributedResultManifest( + schema=_SCHEMA, + format_version=_FORMAT_VERSION, + result_file="result.pt", + frequencies=expected_frequencies, + components=tuple(global_components), + shards=tuple(shard_manifests), + ) + _validate_manifest_structure(manifest) + return manifest, ordered + + +def _validate_manifest_structure(manifest: DistributedResultManifest) -> None: + if not manifest.components or not manifest.shards: + raise ValueError("Sharded Result manifest must contain components and shards.") + ranks = tuple(shard.rank for shard in manifest.shards) + if ranks != tuple(range(len(manifest.shards))): + raise ValueError(f"Manifest shard ranks must be contiguous from zero, got {ranks}.") + component_names = tuple(component.name for component in manifest.components) + if len(set(component_names)) != len(component_names): + raise ValueError("Manifest contains duplicate global components.") + for shard in manifest.shards: + if tuple(component.name for component in shard.components) != component_names: + raise ValueError( + f"Manifest shard rank {shard.rank} component order is inconsistent." + ) + for component_index, global_component in enumerate(manifest.components): + if len(global_component.shape) < 3 or any( + int(value) <= 0 for value in global_component.shape + ): + raise ValueError( + f"Manifest global {global_component.name} has invalid shape " + f"{global_component.shape}." + ) + if global_component.x_axis != len(global_component.shape) - 3: + raise ValueError( + f"Manifest global {global_component.name} x_axis is inconsistent with ndim." + ) + expected_start = 0 + for shard in manifest.shards: + local = shard.components[component_index] + if local.frequencies != manifest.frequencies: + raise ValueError( + f"Manifest shard rank {shard.rank} frequencies are inconsistent." + ) + if local.dtype != global_component.dtype or local.x_axis != global_component.x_axis: + raise ValueError( + f"Manifest shard rank {shard.rank} {local.name} metadata is inconsistent." + ) + if ( + len(local.shape) != len(global_component.shape) + or local.x_axis != len(local.shape) - 3 + ): + raise ValueError( + f"Manifest shard rank {shard.rank} {local.name} shape/x_axis is invalid." + ) + if local.shape[local.x_axis] != ( + local.global_x_slice[1] - local.global_x_slice[0] + ): + raise ValueError( + f"Manifest shard rank {shard.rank} {local.name} x size does not " + "match its owned interval." + ) + for axis, (local_size, global_size) in enumerate( + zip(local.shape, global_component.shape) + ): + if axis != local.x_axis and local_size != global_size: + raise ValueError( + f"Manifest shard rank {shard.rank} {local.name} non-x shape " + "does not match the global field." + ) + if local.global_x_slice[0] != expected_start: + raise ValueError( + f"Manifest shard rank {shard.rank} {local.name} intervals are not contiguous." + ) + expected_start = local.global_x_slice[1] + if expected_start != global_component.shape[global_component.x_axis]: + raise ValueError( + f"Manifest global {global_component.name} x size does not match shard intervals." + ) + + +def _rank_payload( + artifact: FieldShardArtifact, + shard_manifest: FieldShardManifest, +) -> dict[str, Any]: + components = {} + for component, component_manifest in zip(artifact.components, shard_manifest.components): + components[component.name] = { + "tensor": component.tensor, + "global_x_slice": component_manifest.global_x_slice, + "shape": component_manifest.shape, + "dtype": component_manifest.dtype, + "x_axis": component_manifest.x_axis, + "frequencies": component_manifest.frequencies, + } + return _cpu_data( + { + "format_version": _FORMAT_VERSION, + "rank": artifact.rank, + "device": shard_manifest.device, + "frequencies": artifact.frequencies, + "components": components, + } + ) + + +def _sync_file(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def write_sharded_result( + directory: str | Path, + *, + result_payload: Mapping[str, Any], + shard_artifacts: Iterable[FieldShardArtifact | Mapping[str, Any]], + frequencies: tuple[float, ...], +) -> DistributedResultManifest: + """Atomically publish ``manifest.json``, ``result.pt``, and rank tensor files.""" + + destination = Path(directory) + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists(): + raise FileExistsError( + "Sharded Result directory already exists; refusing non-atomic overwrite: " + f"{destination}." + ) + manifest, artifacts = _build_manifest(shard_artifacts, frequencies=frequencies) + staging = Path( + tempfile.mkdtemp( + prefix=f".{destination.name}.tmp-", + dir=destination.parent, + ) + ) + try: + result_path = staging / manifest.result_file + torch.save(_cpu_data(dict(result_payload)), result_path) + _sync_file(result_path) + for artifact, shard_manifest in zip(artifacts, manifest.shards): + shard_path = staging / shard_manifest.file + torch.save(_rank_payload(artifact, shard_manifest), shard_path) + _sync_file(shard_path) + manifest_path = staging / "manifest.json" + manifest_path.write_text( + json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + _sync_file(manifest_path) + os.rename(staging, destination) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + return manifest + + +def _load_result_payload( + directory: Path, + manifest: DistributedResultManifest, + *, + map_location: Any, +) -> dict[str, Any]: + path = directory / manifest.result_file + if not path.is_file(): + raise FileNotFoundError(f"Sharded Result metadata file is missing: {path}.") + try: + payload = torch.load(path, map_location=map_location, weights_only=False) + except Exception as exc: + raise ValueError(f"Failed to load sharded Result metadata {path}: {exc}") from exc + if not isinstance(payload, dict): + raise ValueError(f"Sharded Result metadata {path} must contain a mapping.") + payload_frequencies = _frequencies( + payload.get("frequencies", (payload.get("frequency"),)) + ) + if payload_frequencies != manifest.frequencies: + raise ValueError("Sharded Result metadata frequencies do not match manifest.") + return payload + + +def _load_rank_payload( + path: Path, + shard_manifest: FieldShardManifest, + manifest: DistributedResultManifest, + *, + map_location: Any, +) -> dict[str, Any]: + try: + payload = torch.load(path, map_location=map_location, weights_only=True) + except Exception as exc: + raise ValueError( + f"Failed to load sharded Result rank {shard_manifest.rank} file {path}: {exc}" + ) from exc + if not isinstance(payload, dict): + raise ValueError(f"Shard rank {shard_manifest.rank} file must contain a mapping.") + if int(payload.get("format_version", -1)) != _FORMAT_VERSION: + raise ValueError(f"Shard rank {shard_manifest.rank} has unsupported format_version.") + if int(payload.get("rank", -1)) != shard_manifest.rank: + raise ValueError(f"Shard file {path} rank does not match manifest.") + if _frequencies(payload.get("frequencies", ())) != manifest.frequencies: + raise ValueError(f"Shard rank {shard_manifest.rank} frequencies do not match manifest.") + components = payload.get("components") + if not isinstance(components, dict): + raise ValueError(f"Shard rank {shard_manifest.rank} components must be a mapping.") + expected_names = tuple(component.name for component in shard_manifest.components) + if tuple(components) != expected_names: + raise ValueError(f"Shard rank {shard_manifest.rank} component order is corrupt.") + for component_manifest in shard_manifest.components: + component = components[component_manifest.name] + if not isinstance(component, dict): + raise ValueError( + f"Shard rank {shard_manifest.rank} {component_manifest.name} is corrupt." + ) + tensor = component.get("tensor") + if not isinstance(tensor, torch.Tensor): + raise ValueError( + f"Shard rank {shard_manifest.rank} {component_manifest.name} tensor is missing." + ) + checks = { + "global_x_slice": tuple(component.get("global_x_slice", ())) + == component_manifest.global_x_slice, + "shape": tuple(component.get("shape", ())) == component_manifest.shape, + "tensor_shape": tuple(tensor.shape) == component_manifest.shape, + "dtype": str(component.get("dtype")) == component_manifest.dtype, + "tensor_dtype": str(tensor.dtype) == component_manifest.dtype, + "x_axis": int(component.get("x_axis", -1)) == component_manifest.x_axis, + "frequencies": _frequencies(component.get("frequencies", ())) + == component_manifest.frequencies, + } + failed = tuple(name for name, passed in checks.items() if not passed) + if failed: + raise ValueError( + f"Shard rank {shard_manifest.rank} {component_manifest.name} metadata " + f"is corrupt: {', '.join(failed)}." + ) + return payload + + +def _gather_fields( + manifest: DistributedResultManifest, + shard_payloads: tuple[dict[str, Any], ...], +) -> dict[str, torch.Tensor]: + fields = {} + for component_index, global_component in enumerate(manifest.components): + first = shard_payloads[0]["components"][global_component.name]["tensor"] + destination = torch.empty( + global_component.shape, + device=first.device, + dtype=first.dtype, + ) + for shard_manifest, payload in zip(manifest.shards, shard_payloads): + component_manifest = shard_manifest.components[component_index] + source = payload["components"][global_component.name]["tensor"] + destination_index = [slice(None)] * destination.ndim + destination_index[global_component.x_axis] = slice( + *component_manifest.global_x_slice + ) + destination[tuple(destination_index)].copy_(source) + fields[global_component.name.upper()] = destination + return fields + + +def load_sharded_result( + directory: str | Path, + *, + gather_fields: bool = False, + map_location: Any = "cpu", +) -> LoadedShardedResult: + """Load metadata lazily, reading rank tensors only for an explicit gather.""" + + root = Path(directory) + manifest = DistributedResultManifest.read(root) + paths = manifest.shard_paths(root) + for shard, path in zip(manifest.shards, paths): + if not path.is_file(): + raise FileNotFoundError( + f"Sharded Result rank {shard.rank} file is missing: {path}." + ) + result_payload = _load_result_payload(root, manifest, map_location=map_location) + fields = {} + if gather_fields: + shard_payloads = tuple( + _load_rank_payload( + path, + shard, + manifest, + map_location=map_location, + ) + for shard, path in zip(manifest.shards, paths) + ) + fields = _gather_fields(manifest, shard_payloads) + return LoadedShardedResult( + manifest=manifest, + result_payload=result_payload, + shard_paths=paths, + fields=fields, + ) + + +__all__ = [ + "DistributedResultManifest", + "FieldComponentArtifact", + "FieldComponentManifest", + "FieldShardArtifact", + "FieldShardManifest", + "GlobalFieldManifest", + "LoadedShardedResult", + "export_distributed_field_shards", + "load_sharded_result", + "write_sharded_result", +] diff --git a/witwin/maxwell/fdtd/distributed/solver.py b/witwin/maxwell/fdtd/distributed/solver.py new file mode 100644 index 0000000..8464bac --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/solver.py @@ -0,0 +1,1137 @@ +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from ...monitors import ( + ClosedSurfaceMonitor, + DiffractionMonitor, + DipoleEmissionMonitor, + FieldTimeMonitor, + FinitePlaneMonitor, + FluxTimeMonitor, + MediumMonitor, + ModeMonitor, + PermittivityMonitor, + PlaneMonitor, + PointMonitor, +) +from ...scene import Domain, GridSpec, Scene, prepare_scene +from ...sources import PointDipole, UniformCurrentSource +from ...fdtd_parallel import FDTDParallelConfig, FDTDPartitionPlan, FDTDShardLayout +from ..excitation import ( + inject_electric_surface_source_terms, + inject_magnetic_surface_source_terms, +) +from ..observers import _merge_frequency_lists +from ..solver import FDTD +from ..runtime import stepping +from .counts import reduce_sample_counts +from .capacity import ( + local_dft_working_set_bytes, + require_gather_capacity, + require_local_dft_capacity, +) +from .frequency_counts import reduce_frequency_sample_counts +from .monitor_merge import merge_sharded_monitor_payloads +from .persistence import export_distributed_field_shards +from .sources import crop_solver_source_terms_to_owned_x +from .source_corrections import correct_ideal_point_ex_control_volume +from .transport import CudaP2PHaloTransport + + +_CELL_COMPONENTS = frozenset(("Ex", "Hy", "Hz")) +_NODE_COMPONENTS = frozenset(("Ey", "Ez", "Hx")) + + +@dataclass +class FDTDShard: + rank: int + device: torch.device + layout: FDTDShardLayout + solver: FDTD + compute_stream: torch.cuda.Stream + communication_stream: torch.cuda.Stream + electric_ready: torch.cuda.Event + electric_received: torch.cuda.Event + magnetic_ready: torch.cuda.Event + magnetic_received: torch.cuda.Event + halo_hy_low: torch.Tensor | None + halo_hz_low: torch.Tensor | None + peak_memory_bytes: int = 0 + + @property + def is_first(self) -> bool: + return self.rank == 0 + + @property + def is_last(self) -> bool: + return self.layout.owns_physical_face("x", "high") + + +def _physical_axis_nodes(prepared_scene, axis: str) -> np.ndarray: + nodes = np.asarray(getattr(prepared_scene, f"{axis}_nodes64"), dtype=np.float64) + low = int(prepared_scene.pml_thickness_for_face(axis, "low")) + high = int(prepared_scene.pml_thickness_for_face(axis, "high")) + stop = nodes.size - high if high else nodes.size + return np.array(nodes[low:stop], copy=True) + + +def _owns_global_x(layout: FDTDShardLayout, prepared_scene, x: float) -> bool: + nodes = np.asarray(prepared_scene.x_nodes64, dtype=np.float64) + value = float(x) + if value <= float(nodes[0]): + cell = 0 + elif value >= float(nodes[-1]): + cell = nodes.size - 2 + else: + cell = int(np.searchsorted(nodes, value, side="right")) - 1 + return int(layout.global_cell_owned.start) <= cell < int(layout.global_cell_owned.stop) + + +def _point_source_position(scene: Scene, source_name: str): + for source in scene.sources: + if getattr(source, "name", None) == source_name and hasattr(source, "position"): + return source.position + return None + + +def _plane_position(monitor: PlaneMonitor | FinitePlaneMonitor | ModeMonitor) -> float: + if isinstance(monitor, PlaneMonitor): + return float(monitor.position) + return float(monitor.plane_position) + + +def _matches_x_node(nodes: np.ndarray, node_index: int, value: float) -> bool: + coordinate = float(nodes[node_index]) + spacings = [] + if node_index > 0: + spacings.append(abs(coordinate - float(nodes[node_index - 1]))) + if node_index + 1 < nodes.size: + spacings.append(abs(float(nodes[node_index + 1]) - coordinate)) + scale = max( + abs(coordinate), + abs(float(value)), + max(spacings, default=0.0), + np.finfo(np.float64).tiny, + ) + tolerance = 64.0 * np.finfo(np.float64).eps * scale + return abs(float(value) - coordinate) <= tolerance + + +def _local_monitors(scene: Scene, layout: FDTDShardLayout, global_prepared): + selected = [] + for monitor in scene.monitors: + if isinstance(monitor, (PermittivityMonitor, MediumMonitor)): + # Material monitors are resolved from compiled tensors, not accumulated. + selected.append(monitor) + continue + if isinstance(monitor, PointMonitor): + if _owns_global_x(layout, global_prepared, monitor.position[0]): + selected.append(monitor) + continue + if isinstance(monitor, FieldTimeMonitor) and monitor.region_kind == "point": + if _owns_global_x(layout, global_prepared, monitor.position[0]): + selected.append(monitor) + continue + if isinstance(monitor, DipoleEmissionMonitor): + position = _point_source_position(scene, monitor.source_name) + if position is not None and _owns_global_x(layout, global_prepared, position[0]): + selected.append(monitor) + continue + if isinstance(monitor, (PlaneMonitor, FinitePlaneMonitor, ModeMonitor)): + if monitor.axis != "x" or _owns_global_x( + layout, + global_prepared, + _plane_position(monitor), + ): + selected.append(monitor) + continue + raise ValueError( + f"Multi-GPU FDTD does not support {type(monitor).__name__}." + ) + return selected + + +def _build_local_scene( + logical_scene: Scene, + global_prepared, + layout: FDTDShardLayout, + physical_x_nodes: np.ndarray, +) -> Scene: + begin = int(layout.physical_cell_begin) + end = int(layout.physical_cell_end) + storage_begin = begin - (1 if layout.rank > 0 else 0) + local_x = np.array(physical_x_nodes[storage_begin : end + 1], copy=True) + local_y = _physical_axis_nodes(global_prepared, "y") + local_z = _physical_axis_nodes(global_prepared, "z") + + global_boundary = logical_scene.boundary + local_boundary = global_boundary.with_faces( + x_low=global_boundary.face_kind("x", "low") if layout.rank == 0 else "none", + x_high=( + global_boundary.face_kind("x", "high") + if layout.owns_physical_face("x", "high") + else "none" + ), + ) + local_domain = Domain( + bounds=( + (float(local_x[0]), float(local_x[-1])), + (float(local_y[0]), float(local_y[-1])), + (float(local_z[0]), float(local_z[-1])), + ) + ) + local_grid = GridSpec.custom(local_x, local_y, local_z) + return logical_scene.clone( + domain=local_domain, + grid=local_grid, + boundary=local_boundary, + monitors=_local_monitors(logical_scene, layout, global_prepared), + ports=(), + device=str(layout.device), + symmetry=(None, logical_scene.symmetry[1], logical_scene.symmetry[2]), + ) + + +def _bounded_x_kwargs(layout: FDTDShardLayout, component: str, x_slice: slice) -> dict[str, int]: + component_layout = layout.component(component) + return { + "localXBegin": int(x_slice.start), + "localXEnd": int(x_slice.stop), + "globalXOffset": int(component_layout.global_origin[0]), + "globalXExtent": int(component_layout.global_shape[0]), + } + + +def _launch_magnetic_hx(solver, layout: FDTDShardLayout, x_slice: slice) -> None: + if x_slice.stop <= x_slice.start: + return + solver.fdtd_module.updateMagneticFieldHxStandardBounded3D( + Hx=solver.Hx, + Ey=solver.Ey, + Ez=solver.Ez, + HxDecay=solver.chx_decay, + HxCurl=solver.chx_curl, + invDy=solver.inv_dy_h, + invDz=solver.inv_dz_h, + **_bounded_x_kwargs(layout, "Hx", x_slice), + ).launchRaw() + + +def _launch_magnetic_hy_hz(solver, layout: FDTDShardLayout, x_slice: slice) -> None: + if x_slice.stop <= x_slice.start: + return + solver.fdtd_module.updateMagneticFieldHyStandardBounded3D( + Hy=solver.Hy, + Ex=solver.Ex, + Ez=solver.Ez, + HyDecay=solver.chy_decay, + HyCurl=solver.chy_curl, + invDx=solver.inv_dx_h, + invDz=solver.inv_dz_h, + **_bounded_x_kwargs(layout, "Hy", x_slice), + ).launchRaw() + solver.fdtd_module.updateMagneticFieldHzStandardBounded3D( + Hz=solver.Hz, + Ex=solver.Ex, + Ey=solver.Ey, + HzDecay=solver.chz_decay, + HzCurl=solver.chz_curl, + invDx=solver.inv_dx_h, + invDy=solver.inv_dy_h, + **_bounded_x_kwargs(layout, "Hz", x_slice), + ).launchRaw() + + +def _electric_coefficients(solver): + return stepping._electric_decay_tensors(solver), stepping._electric_curl_tensors(solver) + + +def _launch_electric_ex(solver, layout: FDTDShardLayout, x_slice: slice) -> None: + if x_slice.stop <= x_slice.start: + return + decay, curl = _electric_coefficients(solver) + solver.fdtd_module.updateElectricFieldExStandardBounded3D( + Ex=solver.Ex, + Hy=solver.Hy, + Hz=solver.Hz, + ExDecay=decay[0], + ExCurl=curl[0], + invDy=solver.inv_dy_e, + invDz=solver.inv_dz_e, + yLowBoundaryMode=solver.boundary_y_low_code, + yHighBoundaryMode=solver.boundary_y_high_code, + zLowBoundaryMode=solver.boundary_z_low_code, + zHighBoundaryMode=solver.boundary_z_high_code, + **_bounded_x_kwargs(layout, "Ex", x_slice), + ).launchRaw() + + +def _launch_electric_ey_ez( + solver, + layout: FDTDShardLayout, + node_slice: slice, + *, + x_low_mode: int, + x_high_mode: int, +) -> None: + if node_slice.stop <= node_slice.start: + return + decay, curl = _electric_coefficients(solver) + solver.fdtd_module.updateElectricFieldEyStandardBounded3D( + Ey=solver.Ey, + Hx=solver.Hx, + Hz=solver.Hz, + EyDecay=decay[1], + EyCurl=curl[1], + invDx=solver.inv_dx_e, + invDz=solver.inv_dz_e, + xLowBoundaryMode=x_low_mode, + xHighBoundaryMode=x_high_mode, + zLowBoundaryMode=solver.boundary_z_low_code, + zHighBoundaryMode=solver.boundary_z_high_code, + **_bounded_x_kwargs(layout, "Ey", node_slice), + ).launchRaw() + solver.fdtd_module.updateElectricFieldEzStandardBounded3D( + Ez=solver.Ez, + Hx=solver.Hx, + Hy=solver.Hy, + EzDecay=decay[2], + EzCurl=curl[2], + invDx=solver.inv_dx_e, + invDy=solver.inv_dy_e, + xLowBoundaryMode=x_low_mode, + xHighBoundaryMode=x_high_mode, + yLowBoundaryMode=solver.boundary_y_low_code, + yHighBoundaryMode=solver.boundary_y_high_code, + **_bounded_x_kwargs(layout, "Ez", node_slice), + ).launchRaw() + + +class DistributedFDTD: + """Single-process x-slab coordinator that duck-types the prepared FDTD solver.""" + + def __init__( + self, + scene: Scene, + *, + frequency: float, + parallel: FDTDParallelConfig, + absorber_type: str = "cpml", + cpml_config: dict[str, Any] | None = None, + ): + if not isinstance(parallel, FDTDParallelConfig): + raise TypeError("parallel must be an FDTDParallelConfig instance.") + if parallel.transport == "nccl": + raise RuntimeError( + "NCCL transport is reserved by the API but is not available in the same-process runtime." + ) + self.logical_scene = scene + self.scene = prepare_scene(scene) + self.frequency = float(frequency) + self.parallel = parallel + self.devices = tuple(torch.device(device) for device in parallel.devices) + self.device = torch.device(parallel.result_device) + self.absorber_type = str(absorber_type) + self.cpml_config = dict(cpml_config or {}) + requested_cpml_mode = str(self.cpml_config.get("memory_mode", "auto")).strip().lower() + if requested_cpml_mode not in {"auto", "dense", "slab"}: + raise ValueError( + "cpml_config['memory_mode'] must be one of 'auto', 'dense', or 'slab'." + ) + self._cpml_memory_mode_requested = requested_cpml_mode + self._cpml_memory_mode = "none" + self._cpml_allocated_memory_bytes = 0 + self._cpml_dense_memory_bytes = 0 + self._cpml_slab_memory_bytes = 0 + self.c = 299792458.0 + self.Nx, self.Ny, self.Nz = self.scene.Nx, self.scene.Ny, self.scene.Nz + self.dt = None + self.shards: tuple[FDTDShard, ...] = () + self.partition_plan: FDTDPartitionPlan | None = None + self.shard_layouts: tuple[FDTDShardLayout, ...] = () + self.transport = CudaP2PHaloTransport(self.devices) + self.last_solve_elapsed_s = None + self._cuda_graph_active = False + self._tail_graph_active = False + self._gather_preflight: dict[str, int] = {} + self._dft_preflight: dict[str, dict[str, int]] = {} + self._peak_memory_including_gather: dict[str, int] = {} + self._shutoff_triggered = False + self._shutoff_step = None + self._shutoff_peak = None + self._parallel_stats: dict[str, Any] = {} + self._observer_frequencies: tuple[float, ...] = () + self._initialized = False + self._validate_static_capabilities() + + def _validate_static_capabilities(self) -> None: + boundary = self.logical_scene.boundary + if boundary.uses_kind("bloch"): + raise ValueError( + "Multi-GPU Bloch fields require complex halo exchange, which is not enabled yet." + ) + if boundary.axis_kind("x") == "periodic": + raise ValueError( + "Multi-GPU x decomposition does not yet support periodic x ring exchange." + ) + if self.logical_scene.symmetry[0] is not None: + raise ValueError("Multi-GPU x decomposition does not support x-axis symmetry yet.") + if self.logical_scene.material_regions: + raise ValueError( + "Multi-GPU material density regions require distributed density slicing and are " + "currently limited to the single-GPU adjoint path." + ) + if self.logical_scene.ports: + raise ValueError("Multi-GPU FDTD mode ports require distributed modal plane tiling.") + for monitor in self.logical_scene.monitors: + if isinstance(monitor, ClosedSurfaceMonitor): + raise ValueError( + "Multi-GPU FDTD does not yet support closed-surface monitor assembly." + ) + if isinstance(monitor, DiffractionMonitor): + raise ValueError( + "Multi-GPU FDTD does not yet support distributed diffraction monitors." + ) + if isinstance(monitor, FluxTimeMonitor): + raise ValueError( + "Multi-GPU FDTD does not yet support time-domain flux plane monitors." + ) + if isinstance(monitor, FieldTimeMonitor) and monitor.region_kind != "point": + raise ValueError( + "Multi-GPU FDTD currently supports FieldTimeMonitor only for point regions." + ) + + prospective_plan = FDTDPartitionPlan( + global_shape=(self.Nx, self.Ny, self.Nz), + devices=self.devices, + decomposition_axis="x", + halo_width=1, + low_pml_cells=int(self.scene.pml_thickness_for_face("x", "low")), + high_pml_cells=int(self.scene.pml_thickness_for_face("x", "high")), + ) + internal_x_nodes = tuple( + int(layout.global_cell_owned.stop) + for layout in prospective_plan.shard_layouts[:-1] + ) + x_nodes = np.asarray(self.scene.x_nodes64, dtype=np.float64) + for monitor in self.logical_scene.monitors: + if ( + isinstance(monitor, (PlaneMonitor, FinitePlaneMonitor)) + and monitor.axis == "x" + and "Ex" in monitor.fields + and not bool(monitor.compute_flux) + and any( + _matches_x_node(x_nodes, node_index, _plane_position(monitor)) + for node_index in internal_x_nodes + ) + ): + raise ValueError( + f"Multi-GPU x-normal monitor {monitor.name!r} requests Ex exactly on an " + "internal x partition node. The owning shard does not have a current " + "monitor-only low Ex halo. Move the plane off the split, remove Ex, or " + "use a tangential-field FluxMonitor/ModeMonitor." + ) + if any( + bool(getattr(getattr(structure, "material", None), "is_lossy_metal", False)) + for structure in self.logical_scene.structures + ): + raise ValueError("Multi-GPU FDTD does not yet support distributed SIBC ownership.") + for monitor in self.logical_scene.resolved_monitors(): + if isinstance(monitor, (PermittivityMonitor, MediumMonitor)): + raise ValueError( + "Multi-GPU material monitors require sharded material compilation; " + "full-domain material gathering is intentionally disabled." + ) + if isinstance(monitor, DipoleEmissionMonitor): + matches = tuple( + source + for source in self.logical_scene.sources + if getattr(source, "name", None) == monitor.source_name + ) + if len(matches) != 1 or not isinstance(matches[0], PointDipole): + raise ValueError( + f"DipoleEmissionMonitor {monitor.name!r} requires exactly one named " + f"PointDipole source {monitor.source_name!r}." + ) + for source in self.logical_scene.sources: + if not isinstance(source, (PointDipole, UniformCurrentSource)): + raise ValueError( + f"Multi-GPU FDTD currently supports PointDipole and UniformCurrentSource; " + f"got {type(source).__name__}." + ) + if isinstance(source, PointDipole) and source.profile != "ideal": + raise ValueError( + "Multi-GPU PointDipole currently requires profile='ideal' so source " + "normalization has exactly one interface owner." + ) + + def _validate_hardware(self) -> None: + self.transport.preflight() + if self.device.index is None: + raise RuntimeError("Multi-GPU result_device must be an indexed CUDA device.") + for device in self.devices: + if device == self.device: + continue + if device.index is None or not torch.cuda.can_device_access_peer( + self.device.index, device.index + ): + raise RuntimeError( + f"Direct CUDA peer access is unavailable from {self.device} to {device}; " + "result gathering/reduction cannot silently stage through host memory." + ) + if not torch.cuda.can_device_access_peer(device.index, self.device.index): + raise RuntimeError( + f"Direct CUDA peer access is unavailable from {device} to {self.device}; " + "result gathering/reduction cannot silently stage through host memory." + ) + properties = [torch.cuda.get_device_properties(device) for device in self.devices] + signatures = {(prop.name, prop.major, prop.minor) for prop in properties} + if len(signatures) != 1: + raise RuntimeError( + "Multi-GPU FDTD requires homogeneous devices with the same model and compute capability." + ) + + def init_field(self) -> None: + if self._initialized: + return + self._validate_hardware() + reference_solver = FDTD( + self.scene, + frequency=self.frequency, + absorber_type=self.absorber_type, + cpml_config=self.cpml_config, + ) + self.dt = float(reference_solver.dt) + del reference_solver + low_pml = int(self.scene.pml_thickness_for_face("x", "low")) + high_pml = int(self.scene.pml_thickness_for_face("x", "high")) + self.partition_plan = FDTDPartitionPlan( + global_shape=(self.Nx, self.Ny, self.Nz), + devices=self.devices, + decomposition_axis="x", + halo_width=1, + low_pml_cells=low_pml, + high_pml_cells=high_pml, + ) + self.shard_layouts = self.partition_plan.shard_layouts + physical_x = _physical_axis_nodes(self.scene, "x") + + shards = [] + for layout in self.shard_layouts: + device = torch.device(layout.device) + with torch.cuda.device(device): + torch.cuda.reset_peak_memory_stats(device) + local_scene = _build_local_scene( + self.logical_scene, + self.scene, + layout, + physical_x, + ) + local_solver = FDTD( + local_scene, + frequency=self.frequency, + absorber_type=self.absorber_type, + cpml_config=self.cpml_config, + ) + local_solver.dt = self.dt + local_solver.init_field() + if local_solver.nonlinear_enabled: + raise ValueError( + "Multi-GPU nonlinear media require additional collocation halos and " + "bounded nonlinear kernels." + ) + if local_solver.full_aniso_enabled: + raise ValueError( + "Multi-GPU full off-diagonal anisotropy requires additional H/curl halos." + ) + crop_solver_source_terms_to_owned_x(local_solver, layout) + correct_ideal_point_ex_control_volume(local_solver, layout, self.scene) + compute_stream = torch.cuda.Stream(device=device) + communication_stream = torch.cuda.Stream(device=device, priority=-1) + halo_hy = local_solver.Hy[0] if layout.rank > 0 else None + compute_stream.wait_stream(torch.cuda.current_stream(device)) + halo_hz = local_solver.Hz[0] if layout.rank > 0 else None + shard = FDTDShard( + rank=layout.rank, + device=device, + layout=layout, + solver=local_solver, + compute_stream=compute_stream, + communication_stream=communication_stream, + electric_ready=torch.cuda.Event(), + electric_received=torch.cuda.Event(), + magnetic_ready=torch.cuda.Event(), + magnetic_received=torch.cuda.Event(), + halo_hy_low=halo_hy, + halo_hz_low=halo_hz, + ) + self._validate_local_layout(shard) + shards.append(shard) + + self.shards = tuple(shards) + dts = tuple(float(shard.solver.dt) for shard in self.shards) + if not all(np.isclose(self.dt, dt, rtol=0.0, atol=1e-18) for dt in dts): + raise RuntimeError(f"Shard-local time steps disagree: {dts}.") + modes = tuple( + getattr(shard.solver, "_cpml_memory_mode", "none") + for shard in self.shards + ) + active_modes = {mode for mode in modes if mode != "none"} + self._cpml_memory_mode = ( + "none" + if not active_modes + else next(iter(active_modes)) if len(active_modes) == 1 else "mixed" + ) + + def _sum_shard_bytes(attribute: str) -> int: + return sum( + int(getattr(shard.solver, attribute, 0)) + for shard in self.shards + ) + + self._cpml_allocated_memory_bytes = _sum_shard_bytes("_cpml_allocated_memory_bytes") + self._cpml_dense_memory_bytes = _sum_shard_bytes("_cpml_dense_memory_bytes") + self._cpml_slab_memory_bytes = _sum_shard_bytes("_cpml_slab_memory_bytes") + self._initialized = True + + @staticmethod + def _validate_local_layout(shard: FDTDShard) -> None: + layout = shard.layout + solver = shard.solver + cell_owned = layout.storage_cell_owned + node_owned = layout.storage_node_owned + if cell_owned.stop > solver.Nx - 1 or node_owned.stop > solver.Nx: + raise RuntimeError( + f"Rank {shard.rank} padded storage is smaller than its declared owned slices." + ) + expected_low_pad = 0 if shard.rank == 0 else 1 + if cell_owned.start != expected_low_pad or node_owned.start != expected_low_pad: + raise RuntimeError( + f"Rank {shard.rank} has invalid padded low ownership: " + f"cell={cell_owned}, node={node_owned}." + ) + + def _prepare_outputs(self, time_steps, dft_frequency, dft_window, full_field_dft): + spectral_enabled = any(shard.solver.observers for shard in self.shards) + default_frequencies = dft_frequency if dft_frequency is not None else (self.frequency,) + monitor_frequencies = tuple( + getattr(monitor, "frequencies", None) + for monitor in self.logical_scene.resolved_monitors() + if not isinstance(monitor, FieldTimeMonitor) + ) + self._observer_frequencies = ( + _merge_frequency_lists(default_frequencies, *monitor_frequencies) if spectral_enabled else () + ) + for shard in self.shards: + solver = shard.solver + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + if dft_frequency is not None and full_field_dft: + solver.enable_dft(dft_frequency, window_type=dft_window, end_step=time_steps) + else: + solver.dft_enabled = False + solver._dft_entries = [] + solver._sync_dft_legacy_state() + observer_frequency = dft_frequency if dft_frequency is not None else solver.source_frequency + if solver.observers: + solver._prepare_observers(observer_frequency, dft_window, time_steps) + if solver.time_observers: + solver._prepare_time_observers(time_steps) + solver._shutoff_triggered = False + solver._shutoff_step = None + + def _overlap_active(self) -> bool: + if not self.parallel.overlap: + return False + return all( + not shard.solver.uses_cpml + and not shard.solver.complex_fields_enabled + and not getattr(shard.solver, "modulation_enabled", False) + for shard in self.shards + ) + + def _advance_magnetic_overlapped(self) -> None: + self.transport.exchange_electric(self.shards) + for shard in self.shards: + solver = shard.solver + cs = shard.layout.storage_cell_owned + ns = shard.layout.storage_node_owned + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + _launch_magnetic_hx(solver, shard.layout, ns) + if shard.is_last: + _launch_magnetic_hy_hz(solver, shard.layout, cs) + continue + interior = slice(cs.start, max(cs.start, cs.stop - 1)) + _launch_magnetic_hy_hz(solver, shard.layout, interior) + shard.compute_stream.wait_event(shard.electric_received) + _launch_magnetic_hy_hz( + solver, shard.layout, slice(cs.stop - 1, cs.stop) + ) + + def _advance_electric_overlapped(self) -> None: + self.transport.exchange_magnetic(self.shards) + for shard in self.shards: + solver = shard.solver + cs = shard.layout.storage_cell_owned + ns = shard.layout.storage_node_owned + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + _launch_electric_ex(solver, shard.layout, cs) + if shard.is_first: + _launch_electric_ey_ez( + solver, + shard.layout, + ns, + x_low_mode=solver.boundary_x_low_code, + x_high_mode=solver.boundary_x_high_code, + ) + continue + interior = slice(min(ns.start + 1, ns.stop), ns.stop) + _launch_electric_ey_ez( + solver, + shard.layout, + interior, + x_low_mode=solver.boundary_x_low_code, + x_high_mode=solver.boundary_x_high_code, + ) + shard.compute_stream.wait_event(shard.magnetic_received) + boundary = slice(ns.start, min(ns.start + 1, ns.stop)) + _launch_electric_ey_ez( + solver, + shard.layout, + boundary, + x_low_mode=solver.boundary_x_low_code, + x_high_mode=solver.boundary_x_high_code, + ) + + def _advance_magnetic_serialized(self) -> None: + self.transport.exchange_electric(self.shards) + for shard in self.shards: + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + if not shard.is_last: + shard.compute_stream.wait_event(shard.electric_received) + solver = shard.solver + if not solver.uses_cpml and not solver.complex_fields_enabled: + _launch_magnetic_hx( + solver, shard.layout, shard.layout.storage_node_owned + ) + _launch_magnetic_hy_hz( + solver, shard.layout, shard.layout.storage_cell_owned + ) + else: + stepping.update_magnetic_fields( + solver, + solver.Hx, solver.Hy, solver.Hz, + solver.Ex, solver.Ey, solver.Ez, + ) + + def _advance_electric_serialized(self, time_value: float) -> None: + self.transport.exchange_magnetic(self.shards) + for shard in self.shards: + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + if not shard.is_first: + shard.compute_stream.wait_event(shard.magnetic_received) + solver = shard.solver + if ( + not solver.uses_cpml + and not solver.complex_fields_enabled + and not solver.modulation_enabled + ): + _launch_electric_ex( + solver, shard.layout, shard.layout.storage_cell_owned + ) + _launch_electric_ey_ez( + solver, + shard.layout, + shard.layout.storage_node_owned, + x_low_mode=solver.boundary_x_low_code, + x_high_mode=solver.boundary_x_high_code, + ) + else: + stepping.update_electric_fields( + solver, + solver.Ex, + solver.Ey, + solver.Ez, + solver.Hx, + solver.Hy, + solver.Hz, + time_value=time_value, + ) + + def _advance_one_step(self, n: int, *, overlap_active: bool) -> None: + time_value = n * self.dt + for shard in self.shards: + solver = shard.solver + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + if solver.modulation_enabled: + solver.fdtd_module.advanceModulationTime3D( + ModulationTime=solver._modulation_time, + dt=solver.dt, + ).launchRaw() + solver._advance_magnetic_dispersive_state() + + if overlap_active: + self._advance_magnetic_overlapped() + else: + self._advance_magnetic_serialized() + + for shard in self.shards: + solver = shard.solver + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + if solver._magnetic_source_terms: + inject_magnetic_surface_source_terms(solver, time_value=time_value) + solver._apply_magnetic_dispersive_corrections() + solver._advance_dispersive_state() + if solver.nonlinear_enabled: + solver._update_nonlinear_electric_coefficients() + stepping.capture_aniso_conduction_currents(solver) + + if overlap_active: + self._advance_electric_overlapped() + else: + self._advance_electric_serialized(time_value) + + for shard in self.shards: + solver = shard.solver + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + if solver.full_aniso_enabled: + stepping.apply_full_aniso_corrections(solver) + stepping.apply_full_aniso_conduction(solver) + if getattr(solver, "_sibc", None) is not None: + raise RuntimeError("Distributed SIBC surface ownership is not enabled.") + if solver._electric_source_terms: + inject_electric_surface_source_terms( + solver, time_value=time_value + 0.5 * float(solver.dt) + ) + if solver._source_terms: + solver.add_source(time_value=time_value) + solver._apply_dispersive_corrections() + stepping.enforce_pec_boundaries(solver) + stepping.apply_mur_boundaries(solver) + solver.accumulate_dft(n) + solver.accumulate_observers(n) + solver.accumulate_time_observers(n) + + def _owned_electric_energy(self, shard: FDTDShard) -> torch.Tensor: + solver = shard.solver + cs = shard.layout.storage_cell_owned + ns = shard.layout.storage_node_owned + return ( + (solver.eps_Ex[cs] * solver.Ex[cs] * solver.Ex[cs]).sum() + + (solver.eps_Ey[ns] * solver.Ey[ns] * solver.Ey[ns]).sum() + + (solver.eps_Ez[ns] * solver.Ez[ns] * solver.Ez[ns]).sum() + ) + + def _global_shutoff_energy(self) -> torch.Tensor: + local_energies = [] + for shard in self.shards: + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + local = self._owned_electric_energy(shard) + shard.electric_ready.record(shard.compute_stream) + local_energies.append(local) + with torch.cuda.device(self.device): + result_stream = torch.cuda.current_stream(self.device) + total = torch.zeros((), device=self.device, dtype=torch.float32) + for shard, local in zip(self.shards, local_energies): + result_stream.wait_event(shard.electric_ready) + total.add_(local.to(self.device, non_blocking=True)) + return total + + def _synchronize_all(self) -> None: + for shard in self.shards: + shard.compute_stream.synchronize() + shard.communication_stream.synchronize() + + def _gather_component(self, component: str, local_values: tuple[torch.Tensor, ...]) -> torch.Tensor: + is_cell = component.capitalize() in _CELL_COMPONENTS + global_x = self.Nx - 1 if is_cell else self.Nx + sample = local_values[0] + x_axis = sample.ndim - 3 + shape = list(sample.shape) + shape[x_axis] = global_x + destination = torch.empty(tuple(shape), device=self.device, dtype=sample.dtype) + for shard, value in zip(self.shards, local_values): + local_slice = ( + shard.layout.storage_cell_owned if is_cell else shard.layout.storage_node_owned + ) + global_slice = ( + shard.layout.global_cell_owned if is_cell else shard.layout.global_node_owned + ) + src_index = [slice(None)] * value.ndim + dst_index = [slice(None)] * destination.ndim + src_index[x_axis] = local_slice + dst_index[x_axis] = global_slice + destination[tuple(dst_index)].copy_(value[tuple(src_index)], non_blocking=True) + return destination + + def _collect_output(self) -> dict[str, Any] | None: + local_fields: dict[str, list[torch.Tensor]] = {} + shard_monitor_payloads: list[tuple[int, dict[str, Any]]] = [] + frequency_metadata: tuple[float, ...] | None = None + for shard in self.shards: + solver = shard.solver + shard_monitors: dict[str, Any] = {} + if solver.dft_enabled: + local = solver.get_frequency_solution(all_frequencies=True) + metadata = local.get("frequencies") + if metadata is not None: + if isinstance(metadata, torch.Tensor): + values = tuple( + float(value) for value in metadata.detach().cpu().tolist() + ) + else: + values = tuple(float(value) for value in metadata) + if frequency_metadata is None: + frequency_metadata = values + elif frequency_metadata != values: + raise RuntimeError( + "Shard-local DFT frequency metadata is inconsistent." + ) + for name, tensor in local.items(): + if name not in {"Ex", "Ey", "Ez"}: + continue + local_fields.setdefault(name, []).append(tensor) + for enabled, getter in ( + (solver.observers_enabled, solver.get_observer_results), + (solver.time_observers_enabled, solver.get_time_observer_results), + ): + if not enabled: + continue + for name, payload in getter().items(): + if name in shard_monitors: + raise RuntimeError( + f"Monitor {name!r} appears in multiple observer groups on shard " + f"{shard.rank}." + ) + shard_monitors[name] = payload + shard_monitor_payloads.append((shard.rank, shard_monitors)) + + monitors = merge_sharded_monitor_payloads( + (monitor.name for monitor in self.logical_scene.resolved_monitors()), + shard_monitor_payloads, + result_device=self.device, + shard_layouts=self.shard_layouts, + physical_bounds=self.logical_scene.domain.bounds, + ) + + output: dict[str, Any] = {} + if self.parallel.gather_fields: + if local_fields: + for name, values in local_fields.items(): + output[name] = self._gather_component(name, tuple(values)) + else: + for name in ("Ex", "Ey", "Ez"): + values = tuple(getattr(shard.solver, name) for shard in self.shards) + output[name] = self._gather_component(name, values) + if frequency_metadata is not None: + output["frequencies"] = frequency_metadata + if monitors: + output["observers"] = monitors + return output or None + + def solve( + self, + time_steps: int, + dft_frequency: float = None, + enable_plot: bool = False, + dft_window: str = "hanning", + full_field_dft: bool = True, + normalize_source: bool = False, + shutoff: float = 0.0, + shutoff_check_interval: int = 100, + use_cuda_graph: bool = False, + ): + if not self._initialized: + self.init_field() + self._gather_preflight = {} + self._dft_preflight = {} + for shard in self.shards: + local_shape = ( + shard.solver.Nx, + shard.solver.Ny, + shard.solver.Nz, + ) + pending_dft = local_dft_working_set_bytes( + local_shape, + dft_frequency=dft_frequency, + full_field_dft=full_field_dft, + ) + if shard.device == self.device and self.parallel.gather_fields: + preflight = require_gather_capacity( + self.device, + (self.Nx, self.Ny, self.Nz), + dft_frequency=dft_frequency, + full_field_dft=full_field_dft, + pending_local_dft_bytes=pending_dft, + ) + self._gather_preflight = dict(preflight) + else: + preflight = require_local_dft_capacity( + shard.device, + local_shape, + dft_frequency=dft_frequency, + full_field_dft=full_field_dft, + ) + self._dft_preflight[str(shard.device)] = dict(preflight) + if enable_plot: + raise ValueError("Multi-GPU FDTD plotting requires an explicit gathered result slice.") + if use_cuda_graph: + raise ValueError("Multi-GPU FDTD does not capture peer communication in CUDA Graphs.") + if normalize_source and len(self.logical_scene.sources) != 1: + raise ValueError( + "Multi-GPU source normalization requires exactly one logical source." + ) + for shard in self.shards: + shard.solver._normalize_source = bool(normalize_source) + + self._prepare_outputs(time_steps, dft_frequency, dft_window, full_field_dft) + shutoff_min_step = max( + stepping._compute_shutoff_min_step(shard.solver, int(shutoff_check_interval)) + for shard in self.shards + ) + overlap_active = self._overlap_active() + self._shutoff_triggered = False + self._shutoff_step = None + self._shutoff_peak = torch.zeros((), device=self.device, dtype=torch.float32) + self._synchronize_all() + start = time.perf_counter() + + for n in range(int(time_steps)): + self._advance_one_step(n, overlap_active=overlap_active) + if shutoff > 0.0 and (n + 1) % int(shutoff_check_interval) == 0: + energy = self._global_shutoff_energy() + self._shutoff_peak = torch.maximum(self._shutoff_peak, energy) + if bool( + n >= shutoff_min_step + and ((self._shutoff_peak > 0.0) & (energy < float(shutoff) * self._shutoff_peak)).item() + ): + self._shutoff_triggered = True + self._shutoff_step = n + break + + self._synchronize_all() + self.last_solve_elapsed_s = time.perf_counter() - start + for shard in self.shards: + solver = shard.solver + solver.last_solve_elapsed_s = self.last_solve_elapsed_s + solver._shutoff_triggered = self._shutoff_triggered + solver._shutoff_step = self._shutoff_step + if self._shutoff_triggered: + with torch.cuda.device(shard.device): + stepping._complete_spectral_normalization(solver, int(time_steps)) + if solver.dft_enabled: + solver._sync_dft_legacy_state() + if solver.observers_enabled: + solver._sync_observer_legacy_state() + shard.peak_memory_bytes = int(torch.cuda.max_memory_allocated(shard.device)) + + output = self._collect_output() + for device in self.devices: + torch.cuda.synchronize(device) + self._peak_memory_including_gather = { + str(device): int(torch.cuda.max_memory_allocated(device)) + for device in self.devices + } + self._parallel_stats = self._build_parallel_stats(time_steps, overlap_active) + return output + + def _build_parallel_stats(self, time_steps: int, overlap_active: bool) -> dict[str, Any]: + steps_run = (self._shutoff_step + 1) if self._shutoff_triggered else int(time_steps) + halo_bytes_per_step = 0 + for left, right in zip(self.shards[:-1], self.shards[1:]): + halo_bytes_per_step += ( + left.solver.Ey[-1].numel() + + left.solver.Ez[-1].numel() + + right.solver.Hy[0].numel() + + right.solver.Hz[0].numel() + ) * left.solver.Ex.element_size() + partitions = tuple( + { + "rank": shard.rank, + "device": str(shard.device), + "physical_cells": ( + shard.layout.physical_cell_begin, + shard.layout.physical_cell_end, + ), + "global_cells": ( + shard.layout.global_cell_owned.start, + shard.layout.global_cell_owned.stop, + ), + "global_nodes": ( + shard.layout.global_node_owned.start, + shard.layout.global_node_owned.stop, + ), + "peak_memory_bytes": shard.peak_memory_bytes, + } + for shard in self.shards + ) + return { + "devices": tuple(str(device) for device in self.devices), + "decomposition_axis": "x", + "transport": self.transport.name, + "topology": self.transport.topology, + "overlap_requested": bool(self.parallel.overlap), + "overlap_active": bool(overlap_active), + "gather_fields": bool(self.parallel.gather_fields), + "result_device": str(self.device), + "gather_preflight": dict(self._gather_preflight), + "dft_preflight": {device: dict(stats) for device, stats in self._dft_preflight.items()}, + "partitions": partitions, + "halo_bytes_per_step": int(halo_bytes_per_step), + "halo_bytes_total": int(halo_bytes_per_step * steps_run), + "peak_memory_bytes": { + str(shard.device): shard.peak_memory_bytes for shard in self.shards + }, + "peak_memory_bytes_including_gather": dict(self._peak_memory_including_gather), + "wall_time_s": self.last_solve_elapsed_s, + "compute_time_s": None, + "communication_time_s": None, + "exposed_communication_time_s": None, + "timing_note": ( + "Phase timings require an external CUDA profiler; per-step synchronization " + "is intentionally not inserted into the production solve loop." + ), + } + + @property + def parallel_stats(self) -> dict[str, Any]: + return dict(self._parallel_stats) + + def export_field_shards(self): + """Export owned electric-field shards without a global gather.""" + + return export_distributed_field_shards(self) + + @property + def dft_sample_counts(self): + return reduce_sample_counts(self.shards, "dft_sample_counts") + + @property + def observer_frequencies(self): + return tuple(self._observer_frequencies) + + @property + def observer_sample_counts(self): + return reduce_frequency_sample_counts(self.shards, self._observer_frequencies) + + @property + def active_absorber_type(self): + if not self.shards: + return "none" + active = { + shard.solver.active_absorber_type for shard in self.shards + } - {"none"} + if not active: + return "none" + if len(active) != 1: + raise RuntimeError(f"Shard-local absorber types disagree: {sorted(active)}.") + return next(iter(active)) + + @property + def boundary_kind(self): + return self.logical_scene.boundary.kind diff --git a/witwin/maxwell/fdtd/distributed/source_corrections.py b/witwin/maxwell/fdtd/distributed/source_corrections.py new file mode 100644 index 0000000..df7f70c --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/source_corrections.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import numpy as np + +from ...fdtd_parallel import FDTDShardLayout +from .source_rebuild import rebuild_ideal_ex_interface_plane + + +def correct_ideal_point_ex_control_volume(solver, layout: FDTDShardLayout, global_scene) -> None: + """Restore the global Ex control volume at an artificial shard-high edge. + + The local point-source compiler derives sample control widths from its local + ``x_half`` array. On a non-final shard the last owned Ex sample is an + artificial endpoint, so that generic endpoint rule would give it half a + control volume. The Yee owner is still this shard; rescaling the compiled + source patch by local/global width makes interface-adjacent ideal dipoles + identical to global compilation on uniform and nonuniform grids. + """ + + if layout.owns_physical_face("x", "high"): + return + global_x_half = np.asarray(global_scene.x_half64, dtype=np.float64) + local_index = int(layout.storage_cell_owned.stop) - 1 + global_index = int(layout.global_cell_owned.stop) - 1 + if local_index <= 0 or global_index <= 0 or global_index + 1 >= global_x_half.size: + return + + compiled_sources = tuple(getattr(solver, "_compiled_sources", ()) or ()) + for term in tuple(getattr(solver, "_source_terms", ()) or ()): + if term.get("field_name") != "Ex" or term.get("patch") is None: + continue + source_index = term.get("source_index") + if source_index is None or not (0 <= int(source_index) < len(compiled_sources)): + continue + source = compiled_sources[int(source_index)] + if source.get("kind") != "point_dipole" or source.get("profile") != "ideal": + continue + offset = int(term["offsets"][0]) + relative_index = local_index - offset + if 0 <= relative_index < int(term["patch"].shape[0]): + rebuild_ideal_ex_interface_plane( + term, + source, + solver, + local_index=local_index, + global_index=global_index, + global_x_half64=global_x_half, + ) + + +__all__ = ["correct_ideal_point_ex_control_volume"] diff --git a/witwin/maxwell/fdtd/distributed/source_rebuild.py b/witwin/maxwell/fdtd/distributed/source_rebuild.py new file mode 100644 index 0000000..9acc2b9 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/source_rebuild.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import numpy as np +import torch + +from .interpolation import ideal_axis_weight + + +def _sample_control_widths(coords: torch.Tensor) -> torch.Tensor: + count = int(coords.numel()) + if count <= 1: + return torch.ones_like(coords) + diffs = coords[1:] - coords[:-1] + widths = torch.empty_like(coords) + widths[0] = 0.5 * diffs[0] + widths[-1] = 0.5 * diffs[-1] + if count > 2: + widths[1:-1] = 0.5 * (diffs[:-1] + diffs[1:]) + return widths + + +def rebuild_ideal_ex_interface_plane( + term: dict, + source: dict, + solver, + *, + local_index: int, + global_index: int, + global_x_half64: np.ndarray, +) -> None: + """Rebuild one artificial endpoint plane using global interpolation arithmetic.""" + + offset_x, offset_y, offset_z = (int(value) for value in term["offsets"]) + relative_x = local_index - offset_x + patch = term["patch"] + if not (0 <= relative_x < int(patch.shape[0])): + return + + coords = np.asarray(global_x_half64, dtype=np.float32) + left = np.float32(coords[global_index] - coords[global_index - 1]) + right = np.float32(coords[global_index + 1] - coords[global_index]) + global_x_width = np.float32(0.5) * np.float32(left + right) + position = term.get("source_position") + if position is None: + raise RuntimeError("Compiled ideal point-source term is missing its image position.") + x_weight = ideal_axis_weight(coords, position[0], global_index) + + dtype = solver.Ex.dtype + device = solver.device + y_widths = _sample_control_widths(solver.scene.y.to(device=device, dtype=dtype)) + z_widths = _sample_control_widths(solver.scene.z.to(device=device, dtype=dtype)) + y_stop = offset_y + int(patch.shape[1]) + z_stop = offset_z + int(patch.shape[2]) + x_width = torch.as_tensor(global_x_width, device=device, dtype=dtype) + volume = x_width * y_widths[offset_y:y_stop, None] * z_widths[None, offset_z:z_stop] + eps = solver.eps_Ex[local_index, offset_y:y_stop, offset_z:z_stop] + denominator = eps * volume + symmetry_scale = term.get("source_symmetry_scale") + if symmetry_scale is None: + raise RuntimeError("Compiled ideal point-source term is missing its symmetry scale.") + source_scale = ( + -float(solver.dt) * float(source["polarization"][0]) * float(symmetry_scale) + ) + y_coords = np.asarray(solver.scene.y_nodes64, dtype=np.float64) + z_coords = np.asarray(solver.scene.z_nodes64, dtype=np.float64) + for local_y, global_y in enumerate(range(offset_y, y_stop)): + y_weight = ideal_axis_weight(y_coords, position[1], global_y) + for local_z, global_z in enumerate(range(offset_z, z_stop)): + z_weight = ideal_axis_weight(z_coords, position[2], global_z) + numerator = source_scale * x_weight * y_weight * z_weight + patch[relative_x, local_y, local_z].copy_( + numerator / denominator[local_y, local_z] + ) + + +__all__ = ["rebuild_ideal_ex_interface_plane"] diff --git a/witwin/maxwell/fdtd/distributed/sources.py b/witwin/maxwell/fdtd/distributed/sources.py new file mode 100644 index 0000000..cd6d76b --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/sources.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import torch + +from ...fdtd_parallel import FDTDShardLayout + + +_CELL_COMPONENTS = frozenset(("Ex", "Hy", "Hz")) +_PAYLOAD_KEYS = ( + "patch", + "delay_patch", + "activation_delay_patch", + "cw_cos_patch", + "cw_sin_patch", +) +_TERM_LIST_NAMES = ( + "_source_terms", + "_magnetic_source_terms", + "_electric_source_terms", +) + + +def _first_payload(term: dict) -> torch.Tensor: + for key in _PAYLOAD_KEYS: + value = term.get(key) + if isinstance(value, torch.Tensor): + return value + raise RuntimeError("A compiled FDTD source term has no tensor payload.") + + +def _crop_term_to_owned_x(term: dict, layout: FDTDShardLayout) -> dict | None: + """Crop a shard-local source term to component values owned by this rank. + + Local scenes intentionally contain one low ghost cell on every nonzero rank. + Their ordinary source compiler may therefore generate terms on receive ghosts or + on the high electric ghost plane. This final ownership crop makes all source + samples single-writer while retaining the existing source compiler and kernels. + """ + + field_name = str(term["field_name"]) + owned = ( + layout.storage_cell_owned + if field_name in _CELL_COMPONENTS + else layout.storage_node_owned + ) + payload = _first_payload(term) + offset_x = int(term["offsets"][0]) + term_stop_x = offset_x + int(payload.shape[0]) + owned_begin = max(offset_x, int(owned.start)) + owned_end = min(term_stop_x, int(owned.stop)) + if owned_end <= owned_begin: + return None + + local_x = slice(owned_begin - offset_x, owned_end - offset_x) + cropped = dict(term) + cropped["offsets"] = (owned_begin, *tuple(term["offsets"])[1:]) + for key in _PAYLOAD_KEYS: + value = term.get(key) + if isinstance(value, torch.Tensor): + cropped[key] = value[local_x].contiguous() + return cropped + + +def crop_solver_source_terms_to_owned_x(solver, layout: FDTDShardLayout) -> None: + """Apply the immutable Yee ownership contract to every compiled source list.""" + + for attribute in _TERM_LIST_NAMES: + terms = tuple(getattr(solver, attribute, ()) or ()) + cropped = [] + for term in terms: + owned_term = _crop_term_to_owned_x(term, layout) + if owned_term is not None: + cropped.append(owned_term) + setattr(solver, attribute, cropped) + + +__all__ = ["crop_solver_source_terms_to_owned_x"] diff --git a/witwin/maxwell/fdtd/distributed/transport.py b/witwin/maxwell/fdtd/distributed/transport.py new file mode 100644 index 0000000..d5f8525 --- /dev/null +++ b/witwin/maxwell/fdtd/distributed/transport.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Protocol + +import torch + + +class _TransportShard(Protocol): + rank: int + device: torch.device + layout: object + solver: object + compute_stream: torch.cuda.Stream + communication_stream: torch.cuda.Stream + electric_ready: torch.cuda.Event + electric_received: torch.cuda.Event + magnetic_ready: torch.cuda.Event + magnetic_received: torch.cuda.Event + halo_hy_low: torch.Tensor | None + halo_hz_low: torch.Tensor | None + + +@dataclass(frozen=True) +class PeerLink: + source_rank: int + destination_rank: int + source_device: str + destination_device: str + direction: str + + +class HaloTransport(ABC): + """Transport contract used by the distributed time-step coordinator.""" + + name = "abstract" + + @abstractmethod + def preflight(self) -> None: + ... + + @abstractmethod + def exchange_electric(self, shards: tuple[_TransportShard, ...]) -> None: + ... + + @abstractmethod + def exchange_magnetic(self, shards: tuple[_TransportShard, ...]) -> None: + ... + + def teardown(self) -> None: + return None + + +class CudaP2PHaloTransport(HaloTransport): + """One-process CUDA peer transport for asymmetric Yee x halos.""" + + name = "cuda_p2p" + + def __init__(self, devices: tuple[torch.device, ...]): + self.devices = tuple(torch.device(device) for device in devices) + self.links = tuple( + PeerLink( + source_rank=rank + 1, + destination_rank=rank, + source_device=str(self.devices[rank + 1]), + destination_device=str(self.devices[rank]), + direction="electric_right_to_left", + ) + for rank in range(len(self.devices) - 1) + ) + tuple( + PeerLink( + source_rank=rank, + destination_rank=rank + 1, + source_device=str(self.devices[rank]), + destination_device=str(self.devices[rank + 1]), + direction="magnetic_left_to_right", + ) + for rank in range(len(self.devices) - 1) + ) + + def preflight(self) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA P2P transport requires torch.cuda.is_available().") + count = torch.cuda.device_count() + for device in self.devices: + if device.type != "cuda" or device.index is None: + raise ValueError(f"CUDA P2P transport requires indexed CUDA devices, got {device}.") + if device.index < 0 or device.index >= count: + raise ValueError( + f"CUDA device {device} is unavailable; this process exposes {count} CUDA devices." + ) + for left, right in zip(self.devices, self.devices[1:]): + if not torch.cuda.can_device_access_peer(left.index, right.index): + raise RuntimeError( + f"Direct CUDA peer access is unavailable from {left} to {right}; " + "host-staged halo fallback is intentionally disabled." + ) + if not torch.cuda.can_device_access_peer(right.index, left.index): + raise RuntimeError( + f"Direct CUDA peer access is unavailable from {right} to {left}; " + "host-staged halo fallback is intentionally disabled." + ) + + def exchange_electric(self, shards: tuple[_TransportShard, ...]) -> None: + for shard in shards: + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + shard.electric_ready.record(shard.compute_stream) + + for destination, source in zip(shards[:-1], shards[1:]): + source_node = source.layout.storage_node_owned.start + destination_ghost = destination.layout.storage_node_owned.stop + stream = destination.communication_stream + with torch.cuda.device(destination.device), torch.cuda.stream(stream): + stream.wait_event(source.electric_ready) + stream.wait_event(destination.electric_ready) + destination.solver.Ey[destination_ghost].copy_( + source.solver.Ey[source_node], non_blocking=True + ) + destination.solver.Ez[destination_ghost].copy_( + source.solver.Ez[source_node], non_blocking=True + ) + destination.electric_received.record(stream) + + def exchange_magnetic(self, shards: tuple[_TransportShard, ...]) -> None: + for shard in shards: + with torch.cuda.device(shard.device), torch.cuda.stream(shard.compute_stream): + shard.magnetic_ready.record(shard.compute_stream) + + for source, destination in zip(shards[:-1], shards[1:]): + if destination.halo_hy_low is None or destination.halo_hz_low is None: + raise RuntimeError("Magnetic receive halos were not allocated during prepare().") + source_last = source.layout.storage_cell_owned.stop - 1 + stream = destination.communication_stream + with torch.cuda.device(destination.device), torch.cuda.stream(stream): + stream.wait_event(source.magnetic_ready) + stream.wait_event(destination.magnetic_ready) + destination.halo_hy_low.copy_( + source.solver.Hy[source_last], non_blocking=True + ) + destination.halo_hz_low.copy_( + source.solver.Hz[source_last], non_blocking=True + ) + destination.magnetic_received.record(stream) + + @property + def topology(self) -> dict[str, object]: + pairs = [] + for left, right in zip(self.devices, self.devices[1:]): + pairs.append( + { + "devices": (str(left), str(right)), + "peer_left_to_right": bool( + torch.cuda.can_device_access_peer(left.index, right.index) + ), + "peer_right_to_left": bool( + torch.cuda.can_device_access_peer(right.index, left.index) + ), + } + ) + return {"kind": self.name, "neighbor_pairs": tuple(pairs)} diff --git a/witwin/maxwell/fdtd/excitation/injection.py b/witwin/maxwell/fdtd/excitation/injection.py index 6da7678..0f1a1d0 100644 --- a/witwin/maxwell/fdtd/excitation/injection.py +++ b/witwin/maxwell/fdtd/excitation/injection.py @@ -271,7 +271,8 @@ def _prepare_point_dipole_source(solver, source, *, source_index): * profile ) patch_offsets = (ix_start, iy_start, iz_start) - source_patch *= _symmetry_plane_source_scale(scene, "Ez", image_position) + symmetry_scale = _symmetry_plane_source_scale(scene, "Ez", image_position) + source_patch *= symmetry_scale append_source_term( solver._source_terms, solver, @@ -281,6 +282,8 @@ def _prepare_point_dipole_source(solver, source, *, source_index): phase_real=phase_real, phase_imag=phase_imag, source_index=source_index, + source_position=image_position, + source_symmetry_scale=symmetry_scale, source_time=source_time, omega=source_omega, ) @@ -318,7 +321,8 @@ def _prepare_point_dipole_source(solver, source, *, source_index): * profile ) patch_offsets = (ix_start, iy_start, iz_start) - source_patch *= _symmetry_plane_source_scale(scene, "Ex", image_position) + symmetry_scale = _symmetry_plane_source_scale(scene, "Ex", image_position) + source_patch *= symmetry_scale append_source_term( solver._source_terms, solver, @@ -328,6 +332,8 @@ def _prepare_point_dipole_source(solver, source, *, source_index): phase_real=phase_real, phase_imag=phase_imag, source_index=source_index, + source_position=image_position, + source_symmetry_scale=symmetry_scale, source_time=source_time, omega=source_omega, ) @@ -365,7 +371,8 @@ def _prepare_point_dipole_source(solver, source, *, source_index): * profile ) patch_offsets = (ix_start, iy_start, iz_start) - source_patch *= _symmetry_plane_source_scale(scene, "Ey", image_position) + symmetry_scale = _symmetry_plane_source_scale(scene, "Ey", image_position) + source_patch *= symmetry_scale append_source_term( solver._source_terms, solver, @@ -375,6 +382,8 @@ def _prepare_point_dipole_source(solver, source, *, source_index): phase_real=phase_real, phase_imag=phase_imag, source_index=source_index, + source_position=image_position, + source_symmetry_scale=symmetry_scale, source_time=source_time, omega=source_omega, ) diff --git a/witwin/maxwell/fdtd/excitation/temporal.py b/witwin/maxwell/fdtd/excitation/temporal.py index e5763ec..c58fdef 100644 --- a/witwin/maxwell/fdtd/excitation/temporal.py +++ b/witwin/maxwell/fdtd/excitation/temporal.py @@ -24,6 +24,8 @@ def build_source_term( cw_sin_patch=None, source_index=None, source_time=None, + source_position=None, + source_symmetry_scale=None, omega=None, ): grid_tensor = patch @@ -49,6 +51,12 @@ def build_source_term( "cw_cos_patch": None if cw_cos_patch is None else cw_cos_patch.contiguous(), "cw_sin_patch": None if cw_sin_patch is None else cw_sin_patch.contiguous(), "source_index": None if source_index is None else int(source_index), + "source_position": ( + None if source_position is None else tuple(float(value) for value in source_position) + ), + "source_symmetry_scale": ( + None if source_symmetry_scale is None else float(source_symmetry_scale) + ), "source_time": source_time, "omega": None if omega is None else float(omega), } diff --git a/witwin/maxwell/fdtd_parallel.py b/witwin/maxwell/fdtd_parallel.py new file mode 100644 index 0000000..ea9e373 --- /dev/null +++ b/witwin/maxwell/fdtd_parallel.py @@ -0,0 +1,689 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Mapping + +import torch + + +_COMPONENT_SHAPE_OFFSETS = { + "Ex": (-1, 0, 0), + "Ey": (0, -1, 0), + "Ez": (0, 0, -1), + "Hx": (0, -1, -1), + "Hy": (-1, 0, -1), + "Hz": (-1, -1, 0), +} +_CELL_EXTENT_COMPONENTS = frozenset(("Ex", "Hy", "Hz")) +_NODE_EXTENT_COMPONENTS = frozenset(("Ey", "Ez", "Hx")) +_ELECTRIC_HALO_COMPONENTS = frozenset(("Ey", "Ez")) +_MAGNETIC_HALO_COMPONENTS = frozenset(("Hy", "Hz")) +_PHYSICAL_FACES = frozenset( + ("x_low", "x_high", "y_low", "y_high", "z_low", "z_high") +) + +Index3 = tuple[slice, slice, slice] + + +def _require_positive_int(value, *, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer, not {type(value).__name__}.") + if value <= 0: + raise ValueError(f"{name} must be > 0, got {value}.") + return value + + +def _require_nonnegative_int(value, *, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer, not {type(value).__name__}.") + if value < 0: + raise ValueError(f"{name} must be >= 0, got {value}.") + return value + + +def _normalize_cuda_device(value, *, name: str) -> str: + try: + device = torch.device(value) + except (TypeError, RuntimeError) as exc: + raise TypeError(f"{name} must be a CUDA device, got {value!r}.") from exc + if device.type != "cuda": + raise ValueError(f"{name} must be a CUDA device, got {device}.") + if device.index is None: + raise ValueError(f"{name} must include an explicit CUDA index, got {device}.") + return str(device) + + +def _normalize_devices(devices, *, minimum: int) -> tuple[str, ...]: + if isinstance(devices, (str, bytes, torch.device)): + raise TypeError("devices must be an ordered iterable of CUDA devices.") + try: + values = tuple(devices) + except TypeError as exc: + raise TypeError("devices must be an ordered iterable of CUDA devices.") from exc + if len(values) < minimum: + noun = "device" if minimum == 1 else "unique CUDA devices" + raise ValueError(f"devices must contain at least {minimum} {noun}.") + normalized = tuple( + _normalize_cuda_device(value, name=f"devices[{index}]") + for index, value in enumerate(values) + ) + if len(set(normalized)) != len(normalized): + raise ValueError("devices must contain unique CUDA devices.") + return normalized + + +def _normalize_axis(value: str) -> str: + axis = str(value).strip().lower() + if axis != "x": + raise ValueError("decomposition_axis must be 'x' for FDTD joint solves.") + return axis + + +@dataclass(frozen=True) +class FDTDParallelConfig: + """Public execution descriptor for one FDTD solve spanning CUDA devices. + + This type performs only structural validation. CUDA availability, peer access, + topology, homogeneity, and memory capacity are intentionally prepare-time + checks so the public configuration remains usable in CPU-only API tests. + """ + + devices: tuple[str | torch.device, ...] + decomposition_axis: str = "x" + transport: str = "auto" + overlap: bool = True + gather_fields: bool = False + result_device: str | torch.device | None = None + + def __post_init__(self) -> None: + devices = _normalize_devices(self.devices, minimum=2) + axis = _normalize_axis(self.decomposition_axis) + transport = str(self.transport).strip().lower() + if transport not in {"auto", "cuda_p2p", "nccl"}: + raise ValueError("transport must be one of: auto, cuda_p2p, nccl.") + if not isinstance(self.overlap, bool): + raise TypeError("overlap must be a bool.") + if not isinstance(self.gather_fields, bool): + raise TypeError("gather_fields must be a bool.") + + result_device = devices[0] + if self.result_device is not None: + result_device = _normalize_cuda_device(self.result_device, name="result_device") + if result_device not in devices: + raise ValueError("result_device must be one of the participating devices.") + + object.__setattr__(self, "devices", devices) + object.__setattr__(self, "decomposition_axis", axis) + object.__setattr__(self, "transport", transport) + object.__setattr__(self, "result_device", result_device) + + +@dataclass(frozen=True) +class FDTDHaloRegion: + """A persistent receive-halo or send region for one component and neighbor.""" + + side: str + neighbor_rank: int + global_slice: Index3 + local_slice: Index3 + + def __post_init__(self) -> None: + if self.side not in {"low", "high"}: + raise ValueError("side must be 'low' or 'high'.") + if isinstance(self.neighbor_rank, bool) or not isinstance(self.neighbor_rank, int): + raise TypeError("neighbor_rank must be an integer.") + + @property + def width(self) -> int: + axis_slice = self.global_slice[0] + return int(axis_slice.stop) - int(axis_slice.start) + + +@dataclass(frozen=True) +class FDTDComponentLayout: + """Owned and allocated extents of one staggered Yee component on one shard.""" + + component: str + global_shape: tuple[int, int, int] + owned_global_slice: Index3 + allocation_global_slice: Index3 + owned_local_slice: Index3 + global_origin: tuple[int, int, int] + receive_halos: tuple[FDTDHaloRegion, ...] = () + send_regions: tuple[FDTDHaloRegion, ...] = () + + def __post_init__(self) -> None: + if self.component not in _COMPONENT_SHAPE_OFFSETS: + raise ValueError(f"Unsupported Yee component {self.component!r}.") + + @property + def owned_shape(self) -> tuple[int, int, int]: + return tuple( + int(axis_slice.stop) - int(axis_slice.start) + for axis_slice in self.owned_global_slice + ) + + @property + def local_shape(self) -> tuple[int, int, int]: + return tuple( + int(axis_slice.stop) - int(axis_slice.start) + for axis_slice in self.allocation_global_slice + ) + + @property + def low_halo(self) -> FDTDHaloRegion | None: + return next((region for region in self.receive_halos if region.side == "low"), None) + + @property + def high_halo(self) -> FDTDHaloRegion | None: + return next((region for region in self.receive_halos if region.side == "high"), None) + + @property + def low_send(self) -> FDTDHaloRegion | None: + return next((region for region in self.send_regions if region.side == "low"), None) + + @property + def high_send(self) -> FDTDHaloRegion | None: + return next((region for region in self.send_regions if region.side == "high"), None) + + def owns(self, global_index: tuple[int, int, int]) -> bool: + return _index_in_slice(global_index, self.owned_global_slice) + + def contains(self, global_index: tuple[int, int, int]) -> bool: + """Whether the local allocation contains an owned or receive-halo value.""" + + return _index_in_slice(global_index, self.allocation_global_slice) + + def global_to_local( + self, + global_index: tuple[int, int, int], + *, + include_halo: bool = False, + ) -> tuple[int, int, int]: + selected = self.allocation_global_slice if include_halo else self.owned_global_slice + if not _index_in_slice(global_index, selected): + extent = "local allocation" if include_halo else "owned extent" + raise IndexError( + f"Global index {global_index!r} is outside {self.component} {extent}." + ) + return tuple( + int(index) - int(origin) + for index, origin in zip(global_index, self.global_origin) + ) + + def local_to_global( + self, + local_index: tuple[int, int, int], + *, + include_halo: bool = False, + ) -> tuple[int, int, int]: + if len(local_index) != 3: + raise ValueError("local_index must have three entries.") + global_index = tuple( + int(index) + int(origin) + for index, origin in zip(local_index, self.global_origin) + ) + selected = self.allocation_global_slice if include_halo else self.owned_global_slice + if not _index_in_slice(global_index, selected): + extent = "local allocation" if include_halo else "owned extent" + raise IndexError( + f"Local index {local_index!r} is outside {self.component} {extent}." + ) + return global_index + + +def _index_in_slice(index: tuple[int, int, int], extent: Index3) -> bool: + if len(index) != 3: + raise ValueError("A Yee-grid index must have three entries.") + return all( + int(axis_slice.start) <= int(value) < int(axis_slice.stop) + for value, axis_slice in zip(index, extent) + ) + + +@dataclass(frozen=True) +class FDTDShardLayout: + """Immutable ownership, halo, and physical-face metadata for one CUDA shard.""" + + rank: int + device: str + global_shape: tuple[int, int, int] + physical_cell_begin: int + physical_cell_end: int + global_cell_owned: slice + global_node_owned: slice + storage_cell_owned: slice + storage_node_owned: slice + halo_width: int + component_layouts: tuple[FDTDComponentLayout, ...] + physical_faces: frozenset[str] + + @property + def cell_begin(self) -> int: + return self.physical_cell_begin + + @property + def cell_end(self) -> int: + return self.physical_cell_end + + @property + def cell_interval(self) -> tuple[int, int]: + """Tuple form of the global physical-cell interval.""" + + return self.physical_cell_begin, self.physical_cell_end + + @property + def physical_owned_cell_count(self) -> int: + return self.physical_cell_end - self.physical_cell_begin + + @property + def owned_cell_count(self) -> int: + return int(self.global_cell_owned.stop) - int(self.global_cell_owned.start) + + @property + def owned_node_count(self) -> int: + return int(self.global_node_owned.stop) - int(self.global_node_owned.start) + + @property + def storage_cell_count(self) -> int: + return self.owned_cell_count + int(self.storage_cell_owned.start) + + @property + def storage_node_count(self) -> int: + return self.storage_cell_count + 1 + + @property + def components(self) -> Mapping[str, FDTDComponentLayout]: + return MappingProxyType( + {layout.component: layout for layout in self.component_layouts} + ) + + def component(self, name: str) -> FDTDComponentLayout: + normalized = str(name).strip() + normalized = normalized[0].upper() + normalized[1:].lower() if normalized else normalized + for layout in self.component_layouts: + if layout.component == normalized: + return layout + raise KeyError(f"Unknown Yee component {name!r}.") + + def owns_physical_face(self, axis: str, side: str | None = None) -> bool: + face = str(axis).strip().lower() + if side is not None: + face = f"{face}_{str(side).strip().lower()}" + if face not in _PHYSICAL_FACES: + choices = ", ".join(sorted(_PHYSICAL_FACES)) + raise ValueError(f"Physical face must be one of: {choices}.") + return face in self.physical_faces + + +@dataclass(frozen=True) +class FDTDPartitionPlan: + """Balanced x-cell domain decomposition for all six staggered components.""" + + global_shape: tuple[int, int, int] + devices: tuple[str | torch.device, ...] + decomposition_axis: str = "x" + halo_width: int = 1 + low_pml_cells: int = 0 + high_pml_cells: int = 0 + shard_layouts: tuple[FDTDShardLayout, ...] = field(init=False) + + def __post_init__(self) -> None: + if len(self.global_shape) != 3: + raise ValueError("global_shape must be an (Nx, Ny, Nz) tuple.") + shape = tuple( + _require_positive_int(value, name=f"global_shape[{axis}]") + for axis, value in enumerate(self.global_shape) + ) + if any(value < 2 for value in shape): + raise ValueError("FDTD global_shape dimensions must all be >= 2.") + devices = _normalize_devices(self.devices, minimum=1) + axis = _normalize_axis(self.decomposition_axis) + halo_width = _require_positive_int(self.halo_width, name="halo_width") + if halo_width != 1: + raise ValueError("halo_width must be 1 for the current second-order Yee stencil.") + low_pml_cells = _require_nonnegative_int(self.low_pml_cells, name="low_pml_cells") + high_pml_cells = _require_nonnegative_int(self.high_pml_cells, name="high_pml_cells") + + cell_count = shape[0] - 1 + physical_cell_count = cell_count - low_pml_cells - high_pml_cells + if physical_cell_count <= 0: + raise ValueError( + "low_pml_cells + high_pml_cells must leave at least one physical x cell." + ) + if len(devices) > physical_cell_count: + raise ValueError( + f"Cannot partition {physical_cell_count} physical x cells across " + f"{len(devices)} devices." + ) + intervals = _balanced_intervals(physical_cell_count, len(devices)) + if len(devices) > 1 and any(end - begin < halo_width for begin, end in intervals): + raise ValueError( + "Every x partition must own at least halo_width cells; " + f"got intervals={intervals} and halo_width={halo_width}." + ) + + layouts = tuple( + _build_shard_layout( + rank=rank, + device=device, + global_shape=shape, + physical_cell_interval=intervals[rank], + shard_count=len(devices), + halo_width=halo_width, + low_pml_cells=low_pml_cells, + high_pml_cells=high_pml_cells, + ) + for rank, device in enumerate(devices) + ) + object.__setattr__(self, "global_shape", shape) + object.__setattr__(self, "devices", devices) + object.__setattr__(self, "decomposition_axis", axis) + object.__setattr__(self, "halo_width", halo_width) + object.__setattr__(self, "low_pml_cells", low_pml_cells) + object.__setattr__(self, "high_pml_cells", high_pml_cells) + object.__setattr__(self, "shard_layouts", layouts) + + @property + def num_shards(self) -> int: + return len(self.shard_layouts) + + @property + def cell_count(self) -> int: + return self.global_shape[0] - 1 + + @property + def physical_cell_count(self) -> int: + return self.cell_count - self.low_pml_cells - self.high_pml_cells + + @property + def cell_intervals(self) -> tuple[tuple[int, int], ...]: + return tuple(layout.cell_interval for layout in self.shard_layouts) + + @property + def global_cell_intervals(self) -> tuple[tuple[int, int], ...]: + return tuple( + ( + int(layout.global_cell_owned.start), + int(layout.global_cell_owned.stop), + ) + for layout in self.shard_layouts + ) + + @property + def component_shapes(self) -> Mapping[str, tuple[int, int, int]]: + return MappingProxyType(_component_shapes(self.global_shape)) + + def layout(self, rank: int) -> FDTDShardLayout: + if isinstance(rank, bool) or not isinstance(rank, int): + raise TypeError("rank must be an integer.") + if rank < 0 or rank >= self.num_shards: + raise IndexError(f"rank={rank} is out of range for {self.num_shards} shards.") + return self.shard_layouts[rank] + + def layout_for_device(self, device: str | torch.device) -> FDTDShardLayout: + normalized = _normalize_cuda_device(device, name="device") + try: + rank = self.devices.index(normalized) + except ValueError as exc: + raise KeyError(f"Device {normalized} does not participate in this plan.") from exc + return self.shard_layouts[rank] + + def owner_of_cell(self, cell_index: int) -> int: + if isinstance(cell_index, bool) or not isinstance(cell_index, int): + raise TypeError("cell_index must be an integer.") + if cell_index < 0 or cell_index >= self.cell_count: + raise IndexError( + f"cell_index={cell_index} is out of range for extent {self.cell_count}." + ) + for layout in self.shard_layouts: + if int(layout.global_cell_owned.start) <= cell_index < int(layout.global_cell_owned.stop): + return layout.rank + raise RuntimeError(f"No owner found for cell_index={cell_index}; invalid partition.") + + def owner_of_node(self, node_index: int) -> int: + if isinstance(node_index, bool) or not isinstance(node_index, int): + raise TypeError("node_index must be an integer.") + if node_index < 0 or node_index >= self.global_shape[0]: + raise IndexError( + f"node_index={node_index} is out of range for {self.global_shape[0]} x nodes." + ) + if node_index == self.global_shape[0] - 1: + return self.num_shards - 1 + return self.owner_of_cell(node_index) + + def owner_of_component_x(self, component: str, x_index: int) -> int: + normalized = str(component).strip() + normalized = normalized[0].upper() + normalized[1:].lower() if normalized else normalized + if normalized in _CELL_EXTENT_COMPONENTS: + return self.owner_of_cell(x_index) + if normalized in _NODE_EXTENT_COMPONENTS: + return self.owner_of_node(x_index) + raise KeyError(f"Unknown Yee component {component!r}.") + + + + +def _balanced_intervals(cell_count: int, shard_count: int) -> tuple[tuple[int, int], ...]: + quotient, remainder = divmod(cell_count, shard_count) + intervals = [] + begin = 0 + for rank in range(shard_count): + length = quotient + (1 if rank < remainder else 0) + end = begin + length + intervals.append((begin, end)) + begin = end + return tuple(intervals) + + +def _component_shapes(global_shape: tuple[int, int, int]) -> dict[str, tuple[int, int, int]]: + return { + component: tuple(size + offset for size, offset in zip(global_shape, offsets)) + for component, offsets in _COMPONENT_SHAPE_OFFSETS.items() + } + + +def _slice3(x_begin: int, x_end: int, shape: tuple[int, int, int]) -> Index3: + return (slice(x_begin, x_end), slice(0, shape[1]), slice(0, shape[2])) + + +def _build_component_layout( + *, + component: str, + component_shape: tuple[int, int, int], + rank: int, + shard_count: int, + global_cell_owned: slice, + global_node_owned: slice, + storage_cell_owned: slice, + storage_node_owned: slice, + halo_width: int, +) -> FDTDComponentLayout: + cell_extent = component in _CELL_EXTENT_COMPONENTS + global_owned = global_cell_owned if cell_extent else global_node_owned + storage_owned = storage_cell_owned if cell_extent else storage_node_owned + owned_begin = int(global_owned.start) + owned_end = int(global_owned.stop) + allocation_begin = owned_begin - int(storage_owned.start) + storage_cell_count = ( + int(global_cell_owned.stop) + - int(global_cell_owned.start) + + int(storage_cell_owned.start) + ) + storage_count = storage_cell_count if cell_extent else storage_cell_count + 1 + allocation_end = allocation_begin + storage_count + receive_halos: list[FDTDHaloRegion] = [] + send_regions: list[FDTDHaloRegion] = [] + + if component in _ELECTRIC_HALO_COMPONENTS: + if rank < shard_count - 1: + receive_halos.append( + FDTDHaloRegion( + side="high", + neighbor_rank=rank + 1, + global_slice=_slice3( + owned_end, + owned_end + halo_width, + component_shape, + ), + local_slice=_slice3( + int(storage_node_owned.stop), + int(storage_node_owned.stop) + halo_width, + component_shape, + ), + ) + ) + if rank > 0: + send_regions.append( + FDTDHaloRegion( + side="low", + neighbor_rank=rank - 1, + global_slice=_slice3( + owned_begin, + owned_begin + halo_width, + component_shape, + ), + local_slice=_slice3( + int(storage_node_owned.start), + int(storage_node_owned.start) + halo_width, + component_shape, + ), + ) + ) + elif component in _MAGNETIC_HALO_COMPONENTS: + if rank > 0: + receive_halos.append( + FDTDHaloRegion( + side="low", + neighbor_rank=rank - 1, + global_slice=_slice3( + owned_begin - halo_width, + owned_begin, + component_shape, + ), + local_slice=_slice3(0, halo_width, component_shape), + ) + ) + if rank < shard_count - 1: + send_regions.append( + FDTDHaloRegion( + side="high", + neighbor_rank=rank + 1, + global_slice=_slice3(owned_end - halo_width, owned_end, component_shape), + local_slice=_slice3( + int(storage_cell_owned.stop) - halo_width, + int(storage_cell_owned.stop), + component_shape, + ), + ) + ) + + allocation_global = _slice3(allocation_begin, allocation_end, component_shape) + owned_global = _slice3(owned_begin, owned_end, component_shape) + owned_local = _slice3( + int(storage_owned.start), + int(storage_owned.stop), + component_shape, + ) + return FDTDComponentLayout( + component=component, + global_shape=component_shape, + owned_global_slice=owned_global, + allocation_global_slice=allocation_global, + owned_local_slice=owned_local, + global_origin=(allocation_begin, 0, 0), + receive_halos=tuple(receive_halos), + send_regions=tuple(send_regions), + ) + + +def _build_shard_layout( + *, + rank: int, + device: str, + global_shape: tuple[int, int, int], + physical_cell_interval: tuple[int, int], + shard_count: int, + halo_width: int, + low_pml_cells: int, + high_pml_cells: int, +) -> FDTDShardLayout: + shapes = _component_shapes(global_shape) + physical_cell_begin, physical_cell_end = physical_cell_interval + total_cell_count = global_shape[0] - 1 + if rank == 0 and physical_cell_begin != 0: + raise RuntimeError( + "The first physical partition must start at physical cell zero." + ) + if ( + rank == shard_count - 1 + and physical_cell_end + != total_cell_count - low_pml_cells - high_pml_cells + ): + raise RuntimeError( + "The final physical partition must end at the physical cell count." + ) + global_cell_begin = 0 if rank == 0 else physical_cell_begin + low_pml_cells + global_cell_end = ( + total_cell_count + if rank == shard_count - 1 + else physical_cell_end + low_pml_cells + ) + global_cell_owned = slice(global_cell_begin, global_cell_end) + terminal_node_count = 1 if rank == shard_count - 1 else 0 + global_node_owned = slice(global_cell_begin, global_cell_end + terminal_node_count) + storage_owned_start = 0 if rank == 0 else halo_width + owned_cell_count = ( + int(global_cell_owned.stop) - int(global_cell_owned.start) + ) + storage_cell_owned = slice( + storage_owned_start, storage_owned_start + owned_cell_count + ) + storage_node_owned = slice( + storage_owned_start, + storage_owned_start + owned_cell_count + terminal_node_count, + ) + + components = tuple( + _build_component_layout( + component=component, + component_shape=shapes[component], + rank=rank, + shard_count=shard_count, + global_cell_owned=global_cell_owned, + global_node_owned=global_node_owned, + storage_cell_owned=storage_cell_owned, + storage_node_owned=storage_node_owned, + halo_width=halo_width, + ) + for component in _COMPONENT_SHAPE_OFFSETS + ) + physical_faces = {"y_low", "y_high", "z_low", "z_high"} + if rank == 0: + physical_faces.add("x_low") + if rank == shard_count - 1: + physical_faces.add("x_high") + return FDTDShardLayout( + rank=rank, + device=device, + global_shape=global_shape, + physical_cell_begin=physical_cell_begin, + physical_cell_end=physical_cell_end, + global_cell_owned=global_cell_owned, + global_node_owned=global_node_owned, + storage_cell_owned=storage_cell_owned, + storage_node_owned=storage_node_owned, + halo_width=halo_width, + component_layouts=components, + physical_faces=frozenset(physical_faces), + ) + + +__all__ = [ + "FDTDComponentLayout", + "FDTDHaloRegion", + "FDTDParallelConfig", + "FDTDPartitionPlan", + "FDTDShardLayout", +] diff --git a/witwin/maxwell/result.py b/witwin/maxwell/result.py index 5b4ef53..7b7b25f 100644 --- a/witwin/maxwell/result.py +++ b/witwin/maxwell/result.py @@ -741,6 +741,8 @@ def __init__( self._monitors = _clone_mapping(monitors) self._metadata = _clone_mapping(metadata) self._solver_stats = _clone_mapping(solver_stats) + self._sharded_manifest = None + self._shard_paths: tuple[Path, ...] = () self.raw_output = raw_output self.plot = ResultPlotter(self) @@ -776,6 +778,22 @@ def fields(self) -> dict[str, torch.Tensor]: def monitors(self) -> dict[str, Any]: return dict(self._monitors) + @property + def solver_stats(self) -> dict[str, Any]: + return dict(self._solver_stats) + + @property + def is_sharded(self) -> bool: + return self._sharded_manifest is not None + + @property + def sharded_manifest(self): + return self._sharded_manifest + + @property + def shard_paths(self) -> tuple[Path, ...]: + return self._shard_paths + @property def E(self) -> ResultFieldAccessor: return self.at().E @@ -955,13 +973,132 @@ def save(self, path: str | Path): output_path.parent.mkdir(parents=True, exist_ok=True) torch.save( { + "format_version": 1, "method": self.method, "frequency": self.frequency, "frequencies": self.frequencies, "fields": {name: tensor.detach().cpu() for name, tensor in self._fields.items()}, "monitors": _cpu_serializable(self._monitors), - "metadata": self._metadata, - "solver_stats": self._solver_stats, + "metadata": _cpu_serializable(self._metadata), + "solver_stats": _cpu_serializable(self._solver_stats), }, output_path, ) + + def save_sharded(self, directory: str | Path): + """Persist owned field shards without assembling global field tensors.""" + + exporter = getattr(self.solver, "export_field_shards", None) + if not callable(exporter): + raise RuntimeError( + "Result.save_sharded() requires solver.export_field_shards(); " + "this Result has no shard field export provider." + ) + from .fdtd.distributed.persistence import write_sharded_result + + payload = { + "format_version": 1, + "method": self.method, + "frequency": self.frequency, + "frequencies": self.frequencies, + "fields": {}, + "monitors": _cpu_serializable(self._monitors), + "metadata": _cpu_serializable(self._metadata), + "solver_stats": _cpu_serializable(self._solver_stats), + } + return write_sharded_result( + directory, + result_payload=payload, + shard_artifacts=exporter(), + frequencies=self.frequencies, + ) + + @classmethod + def load( + cls, + path: str | Path, + *, + scene, + prepared_scene=None, + map_location: Any = "cpu", + ) -> "Result": + """Load detached inference data saved by :meth:`save`. + + The declarative scene is supplied by the caller because runtime solvers, + prepared scenes, and transport state are intentionally not persisted. + Result files use pickle-backed ``torch.save`` and must therefore only be + loaded from trusted sources. + """ + + payload = torch.load( + Path(path), + map_location=map_location, + weights_only=False, + ) + if not isinstance(payload, dict): + raise ValueError("Result checkpoint must contain a mapping payload.") + version = int(payload.get("format_version", 1)) + if version != 1: + raise ValueError(f"Unsupported Result checkpoint format_version={version}.") + missing = {"method", "fields", "monitors"}.difference(payload) + if missing: + names = ", ".join(sorted(missing)) + raise ValueError(f"Result checkpoint is missing required keys: {names}.") + + return cls( + method=payload["method"], + scene=scene, + prepared_scene=prepared_scene, + frequency=payload.get("frequency"), + frequencies=payload.get("frequencies"), + fields=payload["fields"], + monitors=payload["monitors"], + metadata=payload.get("metadata"), + solver_stats=payload.get("solver_stats"), + ) + + @classmethod + def load_sharded( + cls, + directory: str | Path, + *, + scene, + prepared_scene=None, + gather_fields: bool = False, + map_location: Any = "cpu", + ) -> "Result": + """Load sharded metadata lazily, gathering owned fields only on request.""" + + if not isinstance(gather_fields, bool): + raise TypeError("gather_fields must be a bool.") + from .fdtd.distributed.persistence import load_sharded_result + + loaded = load_sharded_result( + directory, + gather_fields=gather_fields, + map_location=map_location, + ) + payload = loaded.result_payload + version = int(payload.get("format_version", 1)) + if version != 1: + raise ValueError( + f"Unsupported sharded Result metadata format_version={version}." + ) + missing = {"method", "monitors"}.difference(payload) + if missing: + names = ", ".join(sorted(missing)) + raise ValueError(f"Sharded Result metadata is missing required keys: {names}.") + result = cls( + method=payload["method"], + scene=scene, + prepared_scene=prepared_scene, + frequency=payload.get("frequency"), + frequencies=payload.get("frequencies"), + fields=loaded.fields, + monitors=payload["monitors"], + metadata=payload.get("metadata"), + solver_stats=payload.get("solver_stats"), + ) + result._sharded_manifest = loaded.manifest + result._shard_paths = loaded.shard_paths + return result diff --git a/witwin/maxwell/simulation.py b/witwin/maxwell/simulation.py index 885da7d..ea16e6e 100644 --- a/witwin/maxwell/simulation.py +++ b/witwin/maxwell/simulation.py @@ -6,6 +6,8 @@ import torch +from .fdtd_parallel import FDTDParallelConfig + from .monitors import MediumMonitor, PermittivityMonitor from .result import Result from .scene import Scene, SceneModule, prepare_scene @@ -147,6 +149,7 @@ class FDTDConfig: shutoff: float = 0.0 # relative E-energy threshold for auto-shutoff; 0 disables (opt-in) shutoff_check_interval: int = 100 cuda_graph: bool = False # capture the field-update core into a CUDA graph (opt-in) + parallel: FDTDParallelConfig | None = None def __post_init__(self): if not isinstance(self.run_time, TimeConfig): @@ -163,7 +166,15 @@ def __post_init__(self): self.shutoff_check_interval = int(self.shutoff_check_interval) if self.shutoff_check_interval <= 0: raise ValueError("shutoff_check_interval must be > 0.") + if self.parallel is not None and not isinstance(self.parallel, FDTDParallelConfig): + raise TypeError("parallel must be an FDTDParallelConfig instance or None.") self.cuda_graph = bool(self.cuda_graph) + if self.parallel is not None and self.cuda_graph: + raise ValueError("Multi-GPU FDTD does not support CUDA Graph capture.") + if self.parallel is not None and self.enable_plot: + raise ValueError( + "Multi-GPU FDTD plotting requires running first and requesting gathered fields." + ) if self.adjoint_checkpoint_stride is not None: self.adjoint_checkpoint_stride = int(self.adjoint_checkpoint_stride) if self.adjoint_checkpoint_stride <= 0: @@ -323,6 +334,7 @@ def fdtd( shutoff: float = 0.0, shutoff_check_interval: int = 100, cuda_graph: bool = False, + parallel: FDTDParallelConfig | None = None, ) -> "Simulation": return cls( scene=scene, @@ -338,6 +350,7 @@ def fdtd( shutoff=shutoff, shutoff_check_interval=shutoff_check_interval, cuda_graph=cuda_graph, + parallel=parallel, ), ) @@ -346,6 +359,7 @@ def prepare(self): if self.method == SimulationMethod.FDFD: solver = self._build_fdfd_solver() elif self.method == SimulationMethod.FDTD: + self._reject_trainable_parallel_fdtd() solver = self._build_fdtd_solver(initialize=True) if self.has_trainable_parameters: from .fdtd.adjoint.dispatch import validate_native_adjoint_preparation @@ -443,6 +457,7 @@ def _build_fdfd_result(self, solver, solver_cfg: GMRES) -> Result: ) def _run_fdtd(self) -> Result: + self._reject_trainable_parallel_fdtd() if self.has_trainable_parameters: return self._run_fdtd_with_gradient_bridge() solver = self._build_fdtd_solver(initialize=True) @@ -451,7 +466,29 @@ def _run_fdtd(self) -> Result: def _build_fdtd_solver(self, *, initialize: bool): return self._build_fdtd_solver_for_scene(self.scene, initialize=initialize) + def _reject_trainable_parallel_fdtd(self) -> None: + if self.config.parallel is not None and self.has_trainable_parameters: + raise ValueError( + "Multi-GPU FDTD does not support trainable scene parameters; " + "use the single-GPU adjoint path by omitting parallel." + ) + def _build_fdtd_solver_for_scene(self, scene, *, initialize: bool): + if self.config.parallel is not None: + _require_cuda_scene(scene, method="fdtd") + from .fdtd.distributed import DistributedFDTD + + solver = DistributedFDTD( + scene, + frequency=self.frequency, + parallel=self.config.parallel, + absorber_type=self.config.absorber, + cpml_config=self.config.cpml_config, + ) + if initialize: + solver.init_field() + return solver + prepared_scene = prepare_scene(scene) _require_cuda_scene(prepared_scene, method="fdtd") fdtd_backend, _ = _resolve_fdtd_backend() @@ -551,7 +588,9 @@ def _field(name: str): def _run_fdtd_from_solver(self, solver) -> Result: raw_output, time_steps, use_full_field_dft, dft_cfg = self._execute_fdtd_solve(solver, self.scene) if raw_output is None: - raise RuntimeError("FDTD solve did not return any output.") + if self.config.parallel is None or self.config.parallel.gather_fields: + raise RuntimeError("FDTD solve did not return any output.") + raw_output = {} monitors = raw_output.get("observers", {}) if isinstance(raw_output, dict) else {} field_payload = { @@ -559,10 +598,19 @@ def _run_fdtd_from_solver(self, solver) -> Result: for key, value in (raw_output.items() if isinstance(raw_output, dict) else []) if key in {"Ex", "Ey", "Ez"} } - if not field_payload and len(self.frequencies) == 1: + if self.config.parallel is not None and self.config.parallel.gather_fields and not field_payload: + raise RuntimeError( + "Multi-GPU FDTD was configured with gather_fields=True but returned no fields." + ) + if not field_payload and len(self.frequencies) == 1 and self.config.parallel is None: field_payload = self._fdtd_last_step_field_payload(solver) - fields = _to_tensor_fields(field_payload, self.scene.device) + result_device = ( + self.scene.device + if self.config.parallel is None + else self.config.parallel.result_device + ) + fields = _to_tensor_fields(field_payload, result_device) solver_stats = self._build_fdtd_solver_stats( solver, time_steps=time_steps, @@ -602,7 +650,7 @@ def _build_fdtd_solver_stats( ) shutoff_triggered = bool(getattr(solver, "_shutoff_triggered", False)) shutoff_step = getattr(solver, "_shutoff_step", None) - return { + stats = { "time_steps": time_steps, "shutoff": self.config.shutoff, "shutoff_check_interval": self.config.shutoff_check_interval, @@ -640,6 +688,9 @@ def _build_fdtd_solver_stats( else None ), } + if self.config.parallel is not None: + stats["parallel_stats"] = dict(getattr(solver, "parallel_stats", {})) + return stats def _run_fdtd_with_gradient_bridge(self) -> Result: from .fdtd.adjoint import run_fdtd_with_gradient_bridge