ClaudeandClaude Opus 5.5
2b3c2a7529
kernel_v2.cu -> uniform.cu (whole-domain NX/NY, Simulation) and
blocks.cu -> blocked.cu (one block window, every blocked level), with
InitTubeFlow_v2 -> InitTubeFlow. All live callers, include comments,
CONFIG.md, lifecycle doc and current-tree tests move together; no alias.
refinement.cu keeps its name until the dense-route decision.
Historical-revision tooling that names the old files (LES source pins,
timing/J campaign runners, TG run_case, supported_settings citations at
888c8038, development notes) is unchanged on purpose.
No arithmetic change; the PTX symbol InitTubeFlow_v2 is renamed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
207 lines
7.5 KiB
Python
207 lines
7.5 KiB
Python
# CelerisLab/lbm/stepper.py
|
|
"""
|
|
LBMStepper — time-advance driver for LBM kernels.
|
|
|
|
Owns kernel function handles and manages the double-buffer swap.
|
|
Does NOT return observations by default — callers use field.get_macroscopic()
|
|
or body manager methods to pull data when they need it.
|
|
"""
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
|
|
class LBMStepper:
|
|
"""Drive the LBM kernel forward in time."""
|
|
|
|
def __init__(self, field, module: cuda.Module, cfg):
|
|
self.field = field
|
|
self.module = module
|
|
self.cfg = cfg
|
|
|
|
self._esopull = (cfg.streaming == "esopull")
|
|
|
|
if self._esopull:
|
|
self.step_fn = module.get_function("EsoPullStep")
|
|
self.init_fn = module.get_function("InitEsoPull")
|
|
self.prepare_boundaries_fn = module.get_function("PrepareEsoPullBoundaries")
|
|
self.boundary_values_gpu = cuda.mem_alloc((6 * cfg.nx + 2 * cfg.ny * cfg.nq) * 4)
|
|
else:
|
|
self.step_fn = module.get_function("OneStep")
|
|
self.init_fn = module.get_function("InitTubeFlow")
|
|
|
|
self.curved_fn = module.get_function("CurvedBoundaryKernel")
|
|
self.prepare_esopull_curved_fn = module.get_function("PrepareEsoPullCurvedKernel")
|
|
self.apply_esopull_curved_fn = module.get_function("ApplyEsoPullCurvedKernel")
|
|
self.sensor_fn = module.get_function("SensorKernel")
|
|
self.force_region_fn = module.get_function("ForceRegionKernel")
|
|
|
|
tpb = cfg.threads_per_block
|
|
self.block = (tpb, 1, 1)
|
|
if cfg.is_d2q9:
|
|
self.grid = ((cfg.nx + tpb - 1) // tpb, cfg.ny, 1)
|
|
else:
|
|
self.grid = ((cfg.nx + tpb - 1) // tpb, cfg.ny, cfg.nz)
|
|
|
|
self.boundary_grid = ((max(cfg.nx, cfg.ny) + tpb - 1) // tpb, 2, 1)
|
|
self._step_count = 0
|
|
|
|
def initialize(self):
|
|
"""Run the init kernel to set up channel flow + flags."""
|
|
f = self.field
|
|
self.init_fn(
|
|
f.flag_gpu, f.ddf_gpu,
|
|
block=self.block, grid=self.grid,
|
|
)
|
|
if not self._esopull:
|
|
_gpu_bytes = f.n * f.nq * f.store_bytes
|
|
cuda.memcpy_dtod(f.temp_gpu, f.ddf_gpu, _gpu_bytes)
|
|
cuda.memcpy_dtoh(f.flag, f.flag_gpu)
|
|
# Initialization publishes a new t=0 state, including EsoPull parity.
|
|
self._step_count = 0
|
|
f.download_ddf(step_id=0, force=True)
|
|
|
|
def step(
|
|
self,
|
|
n: int = 1,
|
|
*,
|
|
action_gpu: cuda.DeviceAllocation,
|
|
obs_gpu: cuda.DeviceAllocation,
|
|
stream: cuda.Stream | None = None,
|
|
):
|
|
"""Advance *n* time steps."""
|
|
f = self.field
|
|
launch_kw = {}
|
|
if stream is not None:
|
|
launch_kw["stream"] = stream
|
|
|
|
for _ in range(n):
|
|
if self._esopull:
|
|
self.prepare_boundaries_fn(
|
|
f.ddf_gpu, self.boundary_values_gpu, np.uint64(self._step_count),
|
|
block=self.block, grid=self.boundary_grid, **launch_kw,
|
|
)
|
|
self._launch_curved(action_gpu, obs_gpu, **launch_kw)
|
|
|
|
if self._esopull:
|
|
self.step_fn(
|
|
f.ddf_gpu, f.flag_gpu,
|
|
action_gpu, obs_gpu, self.boundary_values_gpu,
|
|
np.uint64(self._step_count),
|
|
block=self.block, grid=self.grid,
|
|
**launch_kw,
|
|
)
|
|
else:
|
|
self.step_fn(
|
|
f.flag_gpu, f.ddf_gpu, f.temp_gpu,
|
|
block=self.block, grid=self.grid,
|
|
**launch_kw,
|
|
)
|
|
f.ddf_gpu, f.temp_gpu = f.temp_gpu, f.ddf_gpu
|
|
|
|
self._launch_force_region(action_gpu, **launch_kw)
|
|
self._launch_sensor(obs_gpu, **launch_kw)
|
|
self._step_count += 1
|
|
if n > 0:
|
|
# Any completed GPU step invalidates host DDF mirror.
|
|
f.invalidate_host_ddf_cache()
|
|
|
|
def _launch_curved(self, action_gpu, obs_gpu, **launch_kw):
|
|
f = self.field
|
|
if f.n_curved == 0:
|
|
return
|
|
tpb = self.cfg.threads_per_block
|
|
grid_c = ((f.n_curved + tpb - 1) // tpb, 1, 1)
|
|
cl = f.curved
|
|
assert cl.fluid_idx_gpu is not None
|
|
assert cl.dir_gpu is not None
|
|
assert cl.q_gpu is not None
|
|
assert cl.rx_gpu is not None
|
|
assert cl.ry_gpu is not None
|
|
assert cl.rz_gpu is not None
|
|
assert cl.fallback_class_gpu is not None
|
|
assert cl.body_id_gpu is not None
|
|
|
|
if self._esopull:
|
|
eb = f.esopull_curved
|
|
assert eb.corr_value_gpu is not None
|
|
assert eb.target_cell_gpu is not None
|
|
assert eb.target_slot_gpu is not None
|
|
self.prepare_esopull_curved_fn(
|
|
f.ddf_gpu,
|
|
cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
|
|
cl.rx_gpu, cl.ry_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
|
|
action_gpu, obs_gpu,
|
|
eb.corr_value_gpu, eb.target_cell_gpu, eb.target_slot_gpu,
|
|
np.uint32(f.n_curved),
|
|
np.uint64(self._step_count),
|
|
block=(tpb, 1, 1), grid=grid_c,
|
|
**launch_kw,
|
|
)
|
|
self.apply_esopull_curved_fn(
|
|
f.ddf_gpu,
|
|
eb.corr_value_gpu, eb.target_cell_gpu, eb.target_slot_gpu,
|
|
np.uint32(f.n_curved),
|
|
block=(tpb, 1, 1), grid=grid_c,
|
|
**launch_kw,
|
|
)
|
|
else:
|
|
assert cl.fluid_idx_gpu is not None
|
|
if cl.surface_velocity is not None and (
|
|
cl.surface_velocity_gpu is None or cl.link_exchange_gpu is None):
|
|
raise RuntimeError("Surface boundary device buffers are missing")
|
|
self.curved_fn(
|
|
f.ddf_gpu,
|
|
cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
|
|
cl.rx_gpu, cl.ry_gpu, cl.rz_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
|
|
action_gpu, obs_gpu,
|
|
np.uint32(f.n_curved),
|
|
cl.surface_velocity_gpu if cl.surface_velocity_gpu is not None else np.uintp(0),
|
|
cl.link_exchange_gpu if cl.link_exchange_gpu is not None else np.uintp(0),
|
|
block=(tpb, 1, 1), grid=grid_c,
|
|
**launch_kw,
|
|
)
|
|
|
|
def _launch_sensor(self, obs_gpu, **launch_kw):
|
|
f = self.field
|
|
if f.n_sensor == 0:
|
|
return
|
|
tpb = self.cfg.threads_per_block
|
|
grid_s = ((f.n_sensor + tpb - 1) // tpb, 1, 1)
|
|
fi_in = f.ddf_gpu
|
|
se = f.sensors
|
|
assert se.cells_gpu is not None and se.obj_id_gpu is not None
|
|
t = np.uint64(self._step_count + 1 if self._esopull else 0)
|
|
self.sensor_fn(
|
|
fi_in, f.flag_gpu,
|
|
se.cells_gpu, se.obj_id_gpu,
|
|
obs_gpu,
|
|
np.uint32(f.n_sensor),
|
|
t,
|
|
block=(tpb, 1, 1), grid=grid_s,
|
|
**launch_kw,
|
|
)
|
|
|
|
def _launch_force_region(self, action_gpu, **launch_kw):
|
|
f = self.field
|
|
if f.n_force_region == 0:
|
|
return
|
|
tpb = self.cfg.threads_per_block
|
|
grid_f = ((f.n_force_region + tpb - 1) // tpb, 1, 1)
|
|
fr = f.force_regions
|
|
assert fr.cells_gpu is not None and fr.obj_id_gpu is not None
|
|
t = np.uint64(self._step_count + 1 if self._esopull else 0)
|
|
self.force_region_fn(
|
|
f.ddf_gpu, f.flag_gpu,
|
|
fr.cells_gpu, fr.obj_id_gpu,
|
|
action_gpu,
|
|
np.uint32(f.n_force_region),
|
|
t,
|
|
block=(tpb, 1, 1), grid=grid_f,
|
|
**launch_kw,
|
|
)
|
|
|
|
@property
|
|
def step_count(self) -> int:
|
|
return self._step_count
|