ClaudeandClaude Opus 5.5
a454f02355
surface_link_velocity already runs jet_link_shapes; CurvedLinkSoA. assign_surface and _prepare_blocked_boundary then ran it again with the same jets and hits to fill the host jet cache. A private _surface_links returns (velocity, aperture, shape, jet_id) from the one call; both callers use it. surface_link_velocity keeps its two-value public contract (now a thin wrapper). Same values, same validation order, same cache contents; host geometry preparation only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
453 lines
17 KiB
Python
453 lines
17 KiB
Python
# CelerisLab/lbm/curved_links.py
|
|
"""Structure-of-arrays helpers for curved Bouzidi links and sensor probe lists.
|
|
|
|
Responsibilities:
|
|
Hold host columns and lazy GPU mirrors for compact boundary kernels.
|
|
|
|
Design::
|
|
Device columns are reused by capacity to avoid reallocating every upload.
|
|
``count`` tracks active rows while ``capacity`` tracks allocated rows.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
from ..body.jets import surface_link_velocity # host-only; re-exported for callers
|
|
|
|
|
|
def _free_alloc(ptr: Optional[cuda.DeviceAllocation]) -> None:
|
|
if ptr is not None:
|
|
ptr.free()
|
|
|
|
|
|
def _upload_columns(columns, stream=None):
|
|
"""Copy low-frequency topology data; keep pinned staging alive to completion."""
|
|
if stream is None:
|
|
for device, host in columns:
|
|
cuda.memcpy_htod(device, host)
|
|
return
|
|
staging = []
|
|
try:
|
|
for device, host in columns:
|
|
pinned = cuda.pagelocked_empty(host.shape, host.dtype)
|
|
pinned[:] = host
|
|
staging.append(pinned)
|
|
cuda.memcpy_htod_async(device, pinned, stream)
|
|
finally:
|
|
# Returning pageable/temporary storage while DMA is active is invalid.
|
|
# Topology uploads are infrequent and intentionally complete on return.
|
|
stream.synchronize()
|
|
|
|
|
|
@dataclass
|
|
class CurvedLinkSoA:
|
|
"""Cut-link columns on host with matching ``*_gpu`` device buffers."""
|
|
|
|
fluid_idx: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint32))
|
|
dir: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint8))
|
|
q: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.float32))
|
|
body_id: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.int32))
|
|
rx: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.float32))
|
|
ry: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.float32))
|
|
rz: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.float32))
|
|
fallback_class: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint8))
|
|
|
|
fluid_idx_gpu: Optional[cuda.DeviceAllocation] = None
|
|
dir_gpu: Optional[cuda.DeviceAllocation] = None
|
|
q_gpu: Optional[cuda.DeviceAllocation] = None
|
|
body_id_gpu: Optional[cuda.DeviceAllocation] = None
|
|
rx_gpu: Optional[cuda.DeviceAllocation] = None
|
|
ry_gpu: Optional[cuda.DeviceAllocation] = None
|
|
rz_gpu: Optional[cuda.DeviceAllocation] = None
|
|
fallback_class_gpu: Optional[cuda.DeviceAllocation] = None
|
|
|
|
count: int = 0
|
|
capacity: int = 0
|
|
|
|
# Optional fixed surface velocity derived from true hit-center columns.
|
|
surface_velocity: Optional[np.ndarray] = None
|
|
aperture: Optional[np.ndarray] = None
|
|
surface_velocity_gpu: Optional[cuda.DeviceAllocation] = None
|
|
link_exchange_gpu: Optional[cuda.DeviceAllocation] = None
|
|
# Host-only jet cache (never uploaded or checkpointed).
|
|
jets: tuple = ()
|
|
jet_shape: Optional[np.ndarray] = None
|
|
jet_id: Optional[np.ndarray] = None
|
|
|
|
def assign_surface(self, objects, cfg, flags, jets=(), hits=None):
|
|
"""Keep the original eight-column geometry/checkpoint contract unchanged.
|
|
|
|
Jet shapes are cached host-side so amplitude changes skip geometry work.
|
|
"""
|
|
self.surface_velocity = self.aperture = None
|
|
self.jets, self.jet_shape, self.jet_id = tuple(jets), None, None
|
|
from ..body.jets import _surface_links
|
|
self.surface_velocity, self.aperture, self.jet_shape, self.jet_id = _surface_links(
|
|
objects, cfg, flags, self.fluid_idx, self.dir, self.body_id, self.rx, self.ry,
|
|
jets=self.jets, hits=hits)
|
|
|
|
def set_jet_amplitude(self, index, amplitude, stream=None):
|
|
"""Rewrite jet rows from cached shapes; re-upload only the velocity column.
|
|
|
|
Host state is committed only after the upload returns; on a device error
|
|
the host keeps the previous amplitude (device contents are then unknown).
|
|
"""
|
|
from ..body.jets import jet_link_velocity
|
|
jets = list(self.jets)
|
|
jets[index] = jets[index].with_amplitude(amplitude) # validates before any change
|
|
velocity = self.surface_velocity.copy()
|
|
driven = self.jet_id >= 0
|
|
velocity[driven] = jet_link_velocity(jets, self.jet_shape, self.jet_id)[driven]
|
|
if self.surface_velocity_gpu is not None:
|
|
_upload_columns([(self.surface_velocity_gpu, velocity)], stream)
|
|
self.jets, self.surface_velocity = tuple(jets), velocity
|
|
|
|
def surface_exchange(self, body_id):
|
|
"""Latest boundary step's (Fx,Fy,Tz), split by aperture; not nozzle thrust.
|
|
|
|
Caller synchronizes the CUDA stream/context. Rebuild/restart clears this
|
|
diagnostic; it is not the time-accumulated rigid-body observation.
|
|
"""
|
|
if self.link_exchange_gpu is None:
|
|
raise RuntimeError("No surface actuation diagnostic is allocated")
|
|
values = np.empty((self.count, 3), np.float32)
|
|
cuda.memcpy_dtoh(values, self.link_exchange_gpu)
|
|
body = self.body_id == body_id
|
|
return {name: values[body & mask].sum(axis=0, dtype=np.float64)
|
|
for name, mask in (("wall", ~self.aperture), ("aperture", self.aperture))}
|
|
|
|
def assign_host(
|
|
self,
|
|
fluid_idx: np.ndarray,
|
|
dir: np.ndarray,
|
|
q: np.ndarray,
|
|
body_id: np.ndarray,
|
|
rx: np.ndarray,
|
|
ry: np.ndarray,
|
|
rz: np.ndarray,
|
|
fallback_class: np.ndarray,
|
|
) -> None:
|
|
"""Store and validate host SoA columns for curved-link kernel inputs."""
|
|
expected_dtypes = {
|
|
"fluid_idx": np.uint32,
|
|
"dir": np.uint8,
|
|
"q": np.float32,
|
|
"body_id": np.int32,
|
|
"rx": np.float32,
|
|
"ry": np.float32,
|
|
"rz": np.float32,
|
|
"fallback_class": np.uint8,
|
|
}
|
|
raw_columns = {
|
|
"fluid_idx": fluid_idx,
|
|
"dir": dir,
|
|
"q": q,
|
|
"body_id": body_id,
|
|
"rx": rx,
|
|
"ry": ry,
|
|
"rz": rz,
|
|
"fallback_class": fallback_class,
|
|
}
|
|
|
|
columns = {}
|
|
for name, value in raw_columns.items():
|
|
arr = np.asarray(value)
|
|
if not np.can_cast(arr.dtype, expected_dtypes[name], casting="same_kind"):
|
|
raise TypeError(
|
|
f"CurvedLinkSoA column '{name}' dtype {arr.dtype} "
|
|
f"is incompatible with {np.dtype(expected_dtypes[name])}"
|
|
)
|
|
columns[name] = arr.astype(expected_dtypes[name], copy=False)
|
|
|
|
n = len(columns["fluid_idx"])
|
|
for name, col in columns.items():
|
|
if len(col) != n:
|
|
raise ValueError(
|
|
f"CurvedLinkSoA column '{name}' has length {len(col)} != {n}"
|
|
)
|
|
if not col.flags["C_CONTIGUOUS"]:
|
|
columns[name] = np.ascontiguousarray(col)
|
|
|
|
self.fluid_idx = columns["fluid_idx"]
|
|
self.dir = columns["dir"]
|
|
self.q = columns["q"]
|
|
self.body_id = columns["body_id"]
|
|
self.rx = columns["rx"]
|
|
self.ry = columns["ry"]
|
|
self.rz = columns["rz"]
|
|
self.fallback_class = columns["fallback_class"]
|
|
|
|
if self.q.size > 0:
|
|
if np.any((self.q <= 0.0) | (self.q > 1.0)):
|
|
raise ValueError("CurvedLinkSoA column 'q' must satisfy 0 < q <= 1")
|
|
if self.fallback_class.size > 0:
|
|
if not np.all(np.isin(self.fallback_class, (0, 1))):
|
|
raise ValueError(
|
|
"CurvedLinkSoA column 'fallback_class' must use classes {0, 1}"
|
|
)
|
|
self.surface_velocity = self.aperture = None
|
|
self.jets, self.jet_shape, self.jet_id = (), None, None
|
|
|
|
def upload(self, stream: Optional[cuda.Stream] = None) -> None:
|
|
"""Upload host columns to GPU with capacity-based allocation reuse.
|
|
|
|
Args:
|
|
stream: If set, use pinned copies ordered on this stream and wait
|
|
for completion before releasing staging memory.
|
|
"""
|
|
n = int(len(self.fluid_idx))
|
|
self.count = n
|
|
if n == 0:
|
|
_free_alloc(self.surface_velocity_gpu)
|
|
_free_alloc(self.link_exchange_gpu)
|
|
self.surface_velocity_gpu = self.link_exchange_gpu = None
|
|
return
|
|
|
|
if n > self.capacity or self.fluid_idx_gpu is None:
|
|
self.free_gpu_columns()
|
|
self.fluid_idx_gpu = cuda.mem_alloc(int(n * self.fluid_idx.dtype.itemsize))
|
|
self.dir_gpu = cuda.mem_alloc(int(n * self.dir.dtype.itemsize))
|
|
self.q_gpu = cuda.mem_alloc(int(n * self.q.dtype.itemsize))
|
|
self.body_id_gpu = cuda.mem_alloc(int(n * self.body_id.dtype.itemsize))
|
|
self.rx_gpu = cuda.mem_alloc(int(n * self.rx.dtype.itemsize))
|
|
self.ry_gpu = cuda.mem_alloc(int(n * self.ry.dtype.itemsize))
|
|
self.rz_gpu = cuda.mem_alloc(int(n * self.rz.dtype.itemsize))
|
|
self.fallback_class_gpu = cuda.mem_alloc(
|
|
int(n * self.fallback_class.dtype.itemsize)
|
|
)
|
|
self.capacity = n
|
|
|
|
assert self.fluid_idx_gpu is not None
|
|
assert self.dir_gpu is not None
|
|
assert self.q_gpu is not None
|
|
assert self.body_id_gpu is not None
|
|
assert self.rx_gpu is not None
|
|
assert self.ry_gpu is not None
|
|
assert self.rz_gpu is not None
|
|
assert self.fallback_class_gpu is not None
|
|
|
|
_upload_columns([
|
|
(self.fluid_idx_gpu, self.fluid_idx),
|
|
(self.dir_gpu, self.dir),
|
|
(self.q_gpu, self.q),
|
|
(self.body_id_gpu, self.body_id),
|
|
(self.rx_gpu, self.rx),
|
|
(self.ry_gpu, self.ry),
|
|
(self.rz_gpu, self.rz),
|
|
(self.fallback_class_gpu, self.fallback_class),
|
|
], stream)
|
|
_free_alloc(self.surface_velocity_gpu)
|
|
_free_alloc(self.link_exchange_gpu)
|
|
self.surface_velocity_gpu = self.link_exchange_gpu = None
|
|
if self.surface_velocity is not None:
|
|
self.surface_velocity_gpu = cuda.mem_alloc(self.surface_velocity.nbytes)
|
|
self.link_exchange_gpu = cuda.mem_alloc(n*3*4)
|
|
_upload_columns([(self.surface_velocity_gpu, self.surface_velocity),
|
|
(self.link_exchange_gpu, np.zeros((n, 3), np.float32))], stream)
|
|
|
|
def free_gpu_columns(self) -> None:
|
|
"""Release device columns only (host arrays unchanged)."""
|
|
_free_alloc(self.fluid_idx_gpu)
|
|
_free_alloc(self.dir_gpu)
|
|
_free_alloc(self.q_gpu)
|
|
_free_alloc(self.body_id_gpu)
|
|
_free_alloc(self.rx_gpu)
|
|
_free_alloc(self.ry_gpu)
|
|
_free_alloc(self.rz_gpu)
|
|
_free_alloc(self.fallback_class_gpu)
|
|
_free_alloc(self.surface_velocity_gpu)
|
|
_free_alloc(self.link_exchange_gpu)
|
|
self.surface_velocity_gpu = self.link_exchange_gpu = None
|
|
self.fluid_idx_gpu = None
|
|
self.dir_gpu = None
|
|
self.q_gpu = None
|
|
self.body_id_gpu = None
|
|
self.rx_gpu = None
|
|
self.ry_gpu = None
|
|
self.rz_gpu = None
|
|
self.fallback_class_gpu = None
|
|
self.capacity = 0
|
|
|
|
def free(self) -> None:
|
|
"""Free GPU storage and reset ``count``."""
|
|
self.free_gpu_columns()
|
|
self.count = 0
|
|
|
|
|
|
@dataclass
|
|
class SensorSoA:
|
|
"""Sensor cell list with GPU mirrors for ``SensorKernel``."""
|
|
|
|
cells: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint32))
|
|
obj_id: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.int32))
|
|
|
|
cells_gpu: Optional[cuda.DeviceAllocation] = None
|
|
obj_id_gpu: Optional[cuda.DeviceAllocation] = None
|
|
|
|
count: int = 0
|
|
|
|
def assign_host(self, cells: np.ndarray, obj_id: np.ndarray) -> None:
|
|
"""Store host sensor columns."""
|
|
self.cells = cells
|
|
self.obj_id = obj_id
|
|
|
|
def upload(self, stream: Optional[cuda.Stream] = None) -> None:
|
|
"""Reallocate GPU columns and upload.
|
|
|
|
Args:
|
|
stream: If set, use pinned copies ordered on this stream and wait
|
|
for completion before releasing staging memory.
|
|
"""
|
|
self.free_gpu_columns()
|
|
n = int(len(self.cells))
|
|
self.count = n
|
|
if n == 0:
|
|
return
|
|
|
|
self.cells_gpu = cuda.mem_alloc(int(self.cells.nbytes))
|
|
self.obj_id_gpu = cuda.mem_alloc(int(self.obj_id.nbytes))
|
|
_upload_columns([
|
|
(self.cells_gpu, self.cells),
|
|
(self.obj_id_gpu, self.obj_id),
|
|
], stream)
|
|
|
|
def free_gpu_columns(self) -> None:
|
|
_free_alloc(self.cells_gpu)
|
|
_free_alloc(self.obj_id_gpu)
|
|
self.cells_gpu = None
|
|
self.obj_id_gpu = None
|
|
|
|
def free(self) -> None:
|
|
"""Free GPU storage and reset ``count``."""
|
|
self.free_gpu_columns()
|
|
self.count = 0
|
|
|
|
|
|
@dataclass
|
|
class ForceRegionSoA:
|
|
"""Force-region cell list with GPU mirrors for ForceRegionKernel."""
|
|
|
|
cells: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint32))
|
|
obj_id: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.int32))
|
|
|
|
cells_gpu: Optional[cuda.DeviceAllocation] = None
|
|
obj_id_gpu: Optional[cuda.DeviceAllocation] = None
|
|
|
|
count: int = 0
|
|
|
|
def assign_host(self, cells: np.ndarray, obj_id: np.ndarray) -> None:
|
|
"""Store host force-region columns."""
|
|
self.cells = cells
|
|
self.obj_id = obj_id
|
|
|
|
def upload(self, stream: Optional[cuda.Stream] = None) -> None:
|
|
"""Reallocate GPU columns and upload.
|
|
|
|
Args:
|
|
stream: If set, use pinned copies ordered on this stream and wait
|
|
for completion before releasing staging memory.
|
|
"""
|
|
self.free_gpu_columns()
|
|
n = int(len(self.cells))
|
|
self.count = n
|
|
if n == 0:
|
|
return
|
|
|
|
self.cells_gpu = cuda.mem_alloc(int(self.cells.nbytes))
|
|
self.obj_id_gpu = cuda.mem_alloc(int(self.obj_id.nbytes))
|
|
_upload_columns([
|
|
(self.cells_gpu, self.cells),
|
|
(self.obj_id_gpu, self.obj_id),
|
|
], stream)
|
|
|
|
def free_gpu_columns(self) -> None:
|
|
_free_alloc(self.cells_gpu)
|
|
_free_alloc(self.obj_id_gpu)
|
|
self.cells_gpu = None
|
|
self.obj_id_gpu = None
|
|
|
|
def free(self) -> None:
|
|
"""Free GPU storage and reset ``count``."""
|
|
self.free_gpu_columns()
|
|
self.count = 0
|
|
|
|
|
|
@dataclass
|
|
class EsoPullCurvedBuffer:
|
|
"""Intermediate buffers for EsoPull curved boundary (Prepare -> Apply).
|
|
|
|
Three arrays per link: correction value, target cell index, target slot.
|
|
Allocated once per n_curved, reused across steps (layout is static).
|
|
"""
|
|
|
|
corr_value: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.float32))
|
|
target_cell: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint32))
|
|
target_slot: np.ndarray = field(
|
|
default_factory=lambda: np.zeros(0, dtype=np.uint8))
|
|
|
|
corr_value_gpu: Optional[cuda.DeviceAllocation] = None
|
|
target_cell_gpu: Optional[cuda.DeviceAllocation] = None
|
|
target_slot_gpu: Optional[cuda.DeviceAllocation] = None
|
|
|
|
count: int = 0
|
|
|
|
def assign_host(self, corr_value: np.ndarray, target_cell: np.ndarray, target_slot: np.ndarray) -> None:
|
|
"""Store host-side intermediate buffer columns."""
|
|
self.corr_value = corr_value
|
|
self.target_cell = target_cell
|
|
self.target_slot = target_slot
|
|
|
|
def upload(self, stream: Optional[cuda.Stream] = None) -> None:
|
|
"""Reallocate GPU columns and upload.
|
|
|
|
Args:
|
|
stream: If set, use pinned copies ordered on this stream and wait
|
|
for completion before releasing staging memory.
|
|
"""
|
|
self.free_gpu_columns()
|
|
n = int(len(self.corr_value))
|
|
self.count = n
|
|
if n == 0:
|
|
return
|
|
|
|
self.corr_value_gpu = cuda.mem_alloc(int(self.corr_value.nbytes))
|
|
self.target_cell_gpu = cuda.mem_alloc(int(self.target_cell.nbytes))
|
|
self.target_slot_gpu = cuda.mem_alloc(int(self.target_slot.nbytes))
|
|
_upload_columns([
|
|
(self.corr_value_gpu, self.corr_value),
|
|
(self.target_cell_gpu, self.target_cell),
|
|
(self.target_slot_gpu, self.target_slot),
|
|
], stream)
|
|
|
|
def free_gpu_columns(self) -> None:
|
|
_free_alloc(self.corr_value_gpu)
|
|
_free_alloc(self.target_cell_gpu)
|
|
_free_alloc(self.target_slot_gpu)
|
|
self.corr_value_gpu = None
|
|
self.target_cell_gpu = None
|
|
self.target_slot_gpu = None
|
|
|
|
def free(self) -> None:
|
|
"""Free GPU storage and reset ``count``."""
|
|
self.free_gpu_columns()
|
|
self.count = 0
|