diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index 053e95d..50f602f 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -18,8 +18,8 @@ from ..cuda import compiler_v2 as compiler from .blocks import BlockLayout2D from ..body.jets import jet_link_velocity from .block_stepper import MAX_GRID_DIM -from .refinement import (ACTIVE, HALO, FixedPatch2D, FixedPatchSet2D, - nest_patch, nest_patch_set, prepare_fine_boundary) +from .hierarchy import build_hierarchy, is_rectangle +from .refinement import ACTIVE, HALO, FixedPatch2D, FixedPatchSet2D, prepare_fine_boundary from .refinement_tables import (blocked_base_tables, blocked_fine_tables, # noqa: F401 (re-export) blocked_kind, _prepare_blocked_boundary) from .descriptors import SOLID @@ -55,6 +55,65 @@ class _Level: boundary: object = None +def _level_name(level): + return ("base", "fine", "child")[level] if level < 3 else f"level {level}" + + +def _legacy_levels(bounds, *, fine_block, base_block, child, child_block, child_rectangles_in_parent, + fine_objects, jets, child_objects, child_jets): + """Normalise the original two/three-level keywords to per-level lists (original errors).""" + if child is not None and child_rectangles_in_parent is not None: + raise ValueError("child and child_rectangles_in_parent are mutually exclusive") + nested = child is not None or child_rectangles_in_parent is not None + bounds = tuple(bounds) + if not is_rectangle(bounds) and (fine_block is None or base_block is None): + raise ValueError("several patches require fine_block and base_block") + fine_objects, jets = tuple(fine_objects), tuple(jets) + child_objects, child_jets = tuple(child_objects), tuple(child_jets) + if not nested and (child_objects or child_jets): + raise ValueError("child_objects/child_jets require a nested child") + if base_block is not None and fine_block is None: + raise ValueError("base_block requires fine_block") + refinements, blocks = [bounds], [base_block, fine_block] + objects, level_jets = [(), fine_objects], [(), jets] + if nested: + if child is not None and not is_rectangle(bounds): + raise ValueError("scalar child requires one parent patch; use child_rectangles_in_parent") + if fine_block is None or base_block is None or child_block is None: + raise ValueError("a nested child needs fine_block, base_block and child_block") + if fine_objects or jets: + raise ValueError("a nested child supports bodies/Jets only on L2; use child_objects/child_jets") + refinements.append(tuple(child) if child is not None else tuple(child_rectangles_in_parent)) + blocks.append(child_block) + objects.append(child_objects) + level_jets.append(child_jets) + return refinements, blocks, objects, level_jets + + +def _check_levels(refinements, blocks, objects, jets): + """Candidate limits shared by the list and the legacy entry points.""" + depth = len(refinements) + 1 + if depth < 2: + raise ValueError("at least one refinement level is required") + for name, values in (("level_blocks", blocks), ("level_objects", objects), ("level_jets", jets)): + if len(values) != depth: + raise ValueError(f"{name} needs one entry per level ({depth})") + for level, block in enumerate(blocks): + if block is not None and len(block) != 2: + raise ValueError(f"{_level_name(level)}_block must be (by, bx)") + if objects[0] or jets[0]: + raise ValueError("base-level bodies/Jets are not supported") + if any(blocks[level - 1] is not None and blocks[level] is None for level in range(1, depth)): + raise ValueError("a blocked level requires blocked refined levels below it") + if not is_rectangle(refinements[0]) and (blocks[0] is None or blocks[1] is None): + raise ValueError("several patches require fine_block and base_block") + if depth > 2: + if any(block is None for block in blocks): + raise ValueError("three or more levels must all be blocked") + if any(objects[level] or jets[level] for level in range(1, depth - 1)): + raise ValueError("bodies/Jets are supported only on the finest level of a nested hierarchy") + + class FixedPatchStepper: """D2Q9/SRT/FP32, one unforced internal 2:1 patch in an x-periodic box. @@ -118,233 +177,149 @@ class FixedPatchStepper: stencil must fit strictly away from the L2 transfer footprint. L2 runs halo -> BFL -> step four times per base step. Runtime child Jet changes require ``set_jet_amplitude(index, amplitude, level=2)``. + + Level lists (any depth): ``refinements[l]`` gives level l+1 in level-l + integer cells (scalar window or rectangle list, see lbm.hierarchy); + ``level_blocks``/``level_objects``/``level_jets`` hold one entry per level, + L0 included. They are mutually exclusive with the legacy keywords above, + which are normalised to the same lists. One recursive hop advances every + level: a child explodes from its parent's current hop source, takes two + hops, then the parent gathers halos, stages restriction, applies BFL and + steps. Depth >= 3 must be fully blocked with bodies/Jets on the finest level. """ - def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(), + def __init__(self, cfg, bounds=None, *, context, coarse_boundary="periodic", fine_objects=(), fine_block=None, base_block=None, jets=(), child=None, child_block=None, - child_objects=(), child_jets=(), child_rectangles_in_parent=None): + child_objects=(), child_jets=(), child_rectangles_in_parent=None, + refinements=None, level_blocks=None, level_objects=None, level_jets=None): if HALO != 4: raise RuntimeError("refinement.cu requires a four-cell fine halo") if coarse_boundary not in ("periodic", "slot_reservoir"): raise ValueError("coarse_boundary must be periodic or slot_reservoir") - if child is not None and child_rectangles_in_parent is not None: - raise ValueError("child and child_rectangles_in_parent are mutually exclusive") - nested = child is not None or child_rectangles_in_parent is not None - bounds = tuple(bounds) - multipatch = len(bounds) > 0 and not isinstance(bounds[0], (Integral, np.integer)) - if multipatch: - if fine_block is None or base_block is None: - raise ValueError("several patches require fine_block and base_block") - self.patch = FixedPatchSet2D.from_config(cfg, bounds) + if refinements is None: + if bounds is None: + raise ValueError("bounds or refinements is required") + if level_blocks is not None or level_objects is not None or level_jets is not None: + raise ValueError("level_blocks/level_objects/level_jets require refinements") + refinements, level_blocks, level_objects, level_jets = _legacy_levels( + bounds, fine_block=fine_block, base_block=base_block, child=child, + child_block=child_block, child_rectangles_in_parent=child_rectangles_in_parent, + fine_objects=fine_objects, jets=jets, child_objects=child_objects, + child_jets=child_jets) else: - self.patch = FixedPatch2D.from_config(cfg, bounds) + if (bounds is not None or child is not None or child_rectangles_in_parent is not None + or fine_block is not None or base_block is not None or child_block is not None + or tuple(fine_objects) or tuple(jets) or tuple(child_objects) or tuple(child_jets)): + raise ValueError("refinements and the legacy patch keywords are mutually exclusive") + depth = len(refinements) + 1 + level_blocks = [None] * depth if level_blocks is None else list(level_blocks) + level_objects = [()] * depth if level_objects is None else list(level_objects) + level_jets = [()] * depth if level_jets is None else list(level_jets) + level_objects = [tuple(o) for o in level_objects] + level_jets = [tuple(j) for j in level_jets] + level_blocks = [None if b is None else tuple(b) for b in level_blocks] + _check_levels(refinements, level_blocks, level_objects, level_jets) + depth = len(refinements) + 1 self._wall_aware = coarse_boundary == "slot_reservoir" + if self._wall_aware and (cfg.y_wall_bc != "on_site_velocity" or cfg.top_wall_bc != "reservoir" + or cfg.velocity != 0 or not 2 <= cfg.wall_slot_width < cfg.nx): + raise ValueError("slot_reservoir requires on_site_velocity, reservoir, " + "velocity=0 and 2 <= wall_slot_width < nx") + self.hierarchy = build_hierarchy(cfg, refinements, wall_aware=self._wall_aware) + pairs = self.hierarchy.pairs if self._wall_aware: - if (cfg.y_wall_bc != "on_site_velocity" or cfg.top_wall_bc != "reservoir" - or cfg.velocity != 0 or not 2 <= cfg.wall_slot_width < cfg.nx): - raise ValueError("slot_reservoir requires on_site_velocity, reservoir, " - "velocity=0 and 2 <= wall_slot_width < nx") - rects = self.patch.bounds if multipatch else (self.patch.bounds,) + rects = (pairs[0].bounds,) if isinstance(pairs[0], FixedPatch2D) else pairs[0].bounds if any(r[1] < 4 or r[3] > cfg.ny - 4 for r in rects): raise ValueError("slot_reservoir requires y0 >= 4 and y1 <= ny-4") - fine_objects, jets = tuple(fine_objects), tuple(jets) - child_objects, child_jets = tuple(child_objects), tuple(child_jets) - if not nested and (child_objects or child_jets): - raise ValueError("child_objects/child_jets require a nested child") - if base_block is not None and fine_block is None: - raise ValueError("base_block requires fine_block") - self.base_layout = None - if base_block is not None: - if len(base_block) != 2: - raise ValueError("base_block must be (by, bx)") - self.base_layout = BlockLayout2D.build( - cfg.ny, cfg.nx, *base_block, periodic=(True, not self._wall_aware)) - if (self.base_layout.nblocks > MAX_GRID_DIM - or self.base_layout.local_shape[0] > MAX_GRID_DIM): - raise ValueError("base blocks exceed launch limits") - self.fine_layout = None - if fine_block is not None: - if len(fine_block) != 2: - raise ValueError("fine_block must be (by, bx)") - fy, fx = self.patch.fine_shape - self.fine_layout = (self.patch.fine_layout(fine_block) if multipatch else - BlockLayout2D.build(fy, fx, *fine_block, periodic=(False, False))) - if (self.fine_layout.nblocks > MAX_GRID_DIM - or self.fine_layout.local_shape[0] > MAX_GRID_DIM): - raise ValueError("fine blocks exceed launch limits") - self.child_patch = self.child_layout = None - self._middle_kind = None - if nested: - if child is not None and multipatch: - raise ValueError("scalar child requires one parent patch; use child_rectangles_in_parent") - if fine_block is None or base_block is None or child_block is None: - raise ValueError("a nested child needs fine_block, base_block and child_block") - if tuple(fine_objects) or tuple(jets): - raise ValueError("a nested child supports bodies/Jets only on L2; use child_objects/child_jets") - if len(child_block) != 2: - raise ValueError("child_block must be (by, bx)") - if child_rectangles_in_parent is not None: - self.child_patch, self._middle_kind = nest_patch_set( - cfg, self.patch, child_rectangles_in_parent) - self.child_layout = self.child_patch.fine_layout(child_block) - else: - self.child_patch, self._middle_kind = nest_patch(cfg, self.patch, child) - self.child_layout = BlockLayout2D.build(*self.child_patch.fine_shape, *child_block, - periodic=(False, False)) - if (self.child_layout.nblocks > MAX_GRID_DIM - or self.child_layout.local_shape[0] > MAX_GRID_DIM): - raise ValueError("child blocks exceed launch limits") - self._jets = tuple(jets) - self.fine_boundary = prepare_fine_boundary( - cfg, self.patch, fine_objects, wall_aware=self._wall_aware, jets=self._jets) - fine_cfg = replace(cfg, nx=self.patch.fine_shape[1], ny=self.patch.fine_shape[0], + self._layouts = [] + for level, block in enumerate(level_blocks): + layout = None + if block is not None: + if level == 0: + layout = BlockLayout2D.build(cfg.ny, cfg.nx, *block, + periodic=(True, not self._wall_aware)) + elif isinstance(pairs[level - 1], FixedPatchSet2D): + layout = pairs[level - 1].fine_layout(block) + else: + layout = BlockLayout2D.build(*pairs[level - 1].fine_shape, *block, + periodic=(False, False)) + if layout.nblocks > MAX_GRID_DIM or layout.local_shape[0] > MAX_GRID_DIM: + raise ValueError(f"{_level_name(level)} blocks exceed launch limits") + self._layouts.append(layout) + # Kernel configurations per level (geometry configs live in self.hierarchy). + fine_cfg = replace(cfg, nx=pairs[0].fine_shape[1], ny=pairs[0].fine_shape[0], viscosity=2 * cfg.viscosity) if self._wall_aware: fine_cfg = replace(fine_cfg, y_wall_bc="bounce_back", top_wall_bc="velocity", wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0) - self.child_boundary = (None if self.child_patch is None else prepare_fine_boundary( - fine_cfg, self.child_patch, child_objects, jets=child_jets, - parent_origin=self.patch.fine_origin, parent_spacing=.5)) - child_cfg = (None if self.child_patch is None else replace( - fine_cfg, nx=self.child_patch.fine_shape[1], ny=self.child_patch.fine_shape[0], - viscosity=4 * cfg.viscosity)) - self._boundaries = { - 1: _prepare_blocked_boundary(self.fine_boundary, self.fine_layout, fine_cfg, self._jets), - 2: _prepare_blocked_boundary(self.child_boundary, self.child_layout, child_cfg, child_jets), - } - # Compatibility aliases for the existing L1 inspection API. - b = self._boundaries[1] - self._fine_tables = None if b is None else b.tables - self._actuated = False if b is None else b.actuated - self._jet_cache = None if b is None else b.jet_cache - self._fine_leaf = (self.patch.fine_kind if self._middle_kind is None - else self._middle_kind) == ACTIVE - self._child_leaf = (None if self.child_patch is None - else self.child_patch.fine_kind == ACTIVE) - if self.child_boundary is not None: - self._child_leaf &= (self.child_boundary.flags.reshape( - self.child_patch.fine_shape) & SOLID) == 0 - self._child = None - if self.fine_boundary is not None: - self._fine_leaf &= (self.fine_boundary.flags.reshape(self.patch.fine_shape) & SOLID) == 0 + self._level_cfgs = [replace(cfg), fine_cfg] + for pair in pairs[1:]: + parent = self._level_cfgs[-1] + self._level_cfgs.append(replace(parent, nx=pair.fine_shape[1], ny=pair.fine_shape[0], + viscosity=2 * parent.viscosity)) + # Per-level boundary geometry in the root edge frame (level l spacing 2**-l). + self._boundary_hosts = [None] + origin, spacing = (0.0, 0.0), 1.0 + for level in range(1, depth): + pair = pairs[level - 1] + if level == 1: + host = prepare_fine_boundary(cfg, pair, level_objects[1], + wall_aware=self._wall_aware, jets=level_jets[1]) + else: + host = prepare_fine_boundary(self._level_cfgs[level - 1], pair, level_objects[level], + jets=level_jets[level], parent_origin=origin, + parent_spacing=spacing) + self._boundary_hosts.append(host) + origin = tuple(o + spacing * f for o, f in zip(origin, pair.fine_origin)) + spacing /= 2 + self._level_jets = level_jets + self._level_boundaries = [None] + [ + _prepare_blocked_boundary(self._boundary_hosts[level], self._layouts[level], + self._level_cfgs[level], level_jets[level]) + for level in range(1, depth)] + self._leaf = [kind == ACTIVE for kind in self.hierarchy.kinds] + for level in range(1, depth): + host = self._boundary_hosts[level] + if host is not None: + self._leaf[level] &= (host.flags.reshape(pairs[level - 1].fine_shape) & SOLID) == 0 # Private constructor inputs for the bounded manual L2 continuation. self._remap_cfg = replace(cfg) - self._remap_blocks = tuple(None if b is None else tuple(b) - for b in (base_block, fine_block, child_block)) + self._level_blocks = level_blocks self.ctx = context self._lock = RLock() self._status = "new" self._step_count = 0 self._allocations = [] - self._levels = [] + self._levels, self._pairs = [], [] self._stream = None self._fine_flags = self._fine_curved = self._curved = None - self._coalesce = None self._link_tables = {} self._block = (cfg.threads_per_block, 1, 1) with self.ctx.activate(): try: self._stream = cuda.Stream() - specs = [(replace(cfg), self.base_layout, self.patch.tau_coarse, "PatchCoarseStep"), - (fine_cfg, self.fine_layout, self.patch.tau_fine, "PatchFineStep")] - if self.child_patch is not None: - specs.append((child_cfg, self.child_layout, - self.child_patch.tau_fine, "PatchFineStep")) - for index, (config, layout, tau, name) in enumerate(specs): - source = "refinement.cu" - if layout is not None: - config = replace(config, nx=layout.local_shape[1], ny=layout.local_shape[0], - y_wall_bc="bounce_back", top_wall_bc="velocity", - wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0) - name, source = "BlockStep", "blocks.cu" - config.validate() - if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0): - raise RuntimeError("level omega/tau mismatch") - n_objects = (len(self.fine_boundary.objects) - if name == "PatchFineStep" and self.fine_boundary is not None else 0) - build = compiler.generate_config(config, n_objects=n_objects) - try: - ptx = compiler.compile_kernel(self.ctx.sm_arch, source=source, - build_dir=build) - module = compiler.load_module(ptx) - finally: - compiler.cleanup_build_dir(build) - ptr, size = module.get_global("d_params") - if size != struct.calcsize("=ffffI"): - raise RuntimeError("LBMParams host/device layout mismatch") - cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, n_objects)) - shape = (9, *(self.patch.coarse_shape, self.patch.fine_shape, - self.child_patch and self.child_patch.fine_shape)[index]) - size = int(np.prod(shape if layout is None else layout.blocked_shape)) * 4 - grid = ((config.nx + self._block[0] - 1) // self._block[0], config.ny, - 1 if layout is None else layout.nblocks) - level = _Level(shape, module, self._allocate(size), self._allocate(size), - module.get_function(name), grid, layout) - level.boundary = self._boundaries.get(index) - self._levels.append(level) - if layout is not None: - level.halo_copy = module.get_function("BlockHaloCopy") - if len(layout.halo_dst): - level.halo_dst = self._upload(layout.halo_dst) - level.halo_src = self._upload(layout.halo_src) - if index == 0 and self._wall_aware: - ptr, size = module.get_global("d_block_wall") - if size != struct.calcsize("=IIIf"): - raise RuntimeError("Block wall host/device layout mismatch") - cuda.memcpy_htod(ptr, struct.pack("=IIIf", cfg.nx, cfg.ny, - cfg.wall_slot_width, cfg.wall_normal_velocity)) - p = self.patch - c, f = self._levels[:2] - order = np.argsort(p.coalesce_coarse) - fine_tables = (None if f.layout is None else - blocked_fine_tables(p, f.layout, order)) - base_tables = (None if c.layout is None else - blocked_base_tables(p, c.layout, order, self._wall_aware)) - self._ghosts = self._upload(p.ghost_fine if fine_tables is None else fine_tables[0]) - donors = p.ghost_coarse if base_tables is None else base_tables[0] - self._donors = self._upload(donors if fine_tables is None else donors.astype(np.uint64)) - self._targets = self._upload(p.coalesce_coarse[order] if base_tables is None - else base_tables[1]) - self._children = self._upload(p.coalesce_fine[order] if fine_tables is None - else fine_tables[1]) - self._explode = f.module.get_function("PatchExplode" if fine_tables is None - else "BlockExplode") - if fine_tables is not None: - f.kind = self._upload(fine_tables[2] if self._middle_kind is None - else blocked_kind(f.layout, self._middle_kind)) - if base_tables is not None: - c.kind = self._upload(base_tables[2]) - self._coalesce = c.module.get_function("BlockCoalesce") - if self.child_patch is not None: - g, cp = self._levels[2], self.child_patch - order = np.argsort(cp.coalesce_coarse) - ghosts, children, g_kind = blocked_fine_tables(cp, g.layout, order) - donors, targets, _ = blocked_base_tables( - cp, f.layout, order, composed_kind=self._middle_kind) - g.kind = self._upload(g_kind) - g.halo_copy = g.module.get_function("BlockHaloCopy") - self._child = _Pair(g.module.get_function("BlockExplode"), - f.module.get_function("BlockCoalesce"), - self._upload(ghosts), self._upload(donors.astype(np.uint64)), - self._upload(targets), self._upload(children), - len(ghosts), len(targets), int(np.prod(f.layout.local_shape))) + for level in range(depth): + self._levels.append(self._build_level(level)) + for index, pair in enumerate(pairs): + self._pairs.append(self._build_pair(index, pair)) for level in self._levels: b = level.boundary if b is not None: b.flags = self._upload(b.tables["flags"]) b.device = {k: self._upload(v) for k, v in b.tables.items() if k != "flags"} b.kernel = level.module.get_function("BlockCurvedBoundary") - if self._boundaries[1] is not None: - b = self._boundaries[1] + if self._level_boundaries[1] is not None: + b = self._level_boundaries[1] self._fine_flags, self._link_tables, self._curved = b.flags, b.device, b.kernel - elif self.fine_boundary is not None: + elif self._boundary_hosts[1] is not None: from .curved_links import CurvedLinkSoA - boundary = self.fine_boundary + boundary = self._boundary_hosts[1] self._fine_flags = self._upload(boundary.flags) # SoA owns its allocations; never also register them in _allocations. self._fine_curved = CurvedLinkSoA() self._fine_curved.assign_host(*boundary.curved_host) self._fine_curved.assign_surface(boundary.objects, fine_cfg, boundary.flags, - jets=self._jets, hits=boundary.curved_hits) + jets=self._level_jets[1], + hits=boundary.curved_hits) self._fine_curved.upload() self._curved = self._levels[1].module.get_function("CurvedBoundaryKernel") except BaseException: @@ -354,6 +329,84 @@ class FixedPatchStepper: pass # Preserve the original construction failure. raise + def _build_level(self, index): + """Compile one level's step module and allocate its two G buffers.""" + layout, config = self._layouts[index], self._level_cfgs[index] + pairs = self.hierarchy.pairs + tau = pairs[0].tau_coarse if index == 0 else pairs[index - 1].tau_fine + name, source = ("PatchCoarseStep" if index == 0 else "PatchFineStep"), "refinement.cu" + if layout is not None: + config = replace(config, nx=layout.local_shape[1], ny=layout.local_shape[0], + y_wall_bc="bounce_back", top_wall_bc="velocity", + wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0) + name, source = "BlockStep", "blocks.cu" + config.validate() + if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0): + raise RuntimeError("level omega/tau mismatch") + host = self._boundary_hosts[index] + n_objects = len(host.objects) if name == "PatchFineStep" and host is not None else 0 + build = compiler.generate_config(config, n_objects=n_objects) + try: + ptx = compiler.compile_kernel(self.ctx.sm_arch, source=source, build_dir=build) + module = compiler.load_module(ptx) + finally: + compiler.cleanup_build_dir(build) + ptr, size = module.get_global("d_params") + if size != struct.calcsize("=ffffI"): + raise RuntimeError("LBMParams host/device layout mismatch") + cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, n_objects)) + shape = (9, *self.hierarchy.kinds[index].shape) + size = int(np.prod(shape if layout is None else layout.blocked_shape)) * 4 + grid = ((config.nx + self._block[0] - 1) // self._block[0], config.ny, + 1 if layout is None else layout.nblocks) + level = _Level(shape, module, self._allocate(size), self._allocate(size), + module.get_function(name), grid, layout) + level.boundary = self._level_boundaries[index] + if layout is not None: + level.halo_copy = module.get_function("BlockHaloCopy") + if len(layout.halo_dst): + level.halo_dst = self._upload(layout.halo_dst) + level.halo_src = self._upload(layout.halo_src) + if index == 0 and self._wall_aware: + cfg = self._remap_cfg + ptr, size = module.get_global("d_block_wall") + if size != struct.calcsize("=IIIf"): + raise RuntimeError("Block wall host/device layout mismatch") + cuda.memcpy_htod(ptr, struct.pack("=IIIf", cfg.nx, cfg.ny, + cfg.wall_slot_width, cfg.wall_normal_velocity)) + return level + + def _build_pair(self, index, pair): + """Upload the transfer tables of levels (index, index+1) and the child's block kinds.""" + parent, child = self._levels[index], self._levels[index + 1] + order = np.argsort(pair.coalesce_coarse) + fine = None if child.layout is None else blocked_fine_tables(pair, child.layout, order) + base = None + if parent.layout is not None: + if index == 0: + base = blocked_base_tables(pair, parent.layout, order, self._wall_aware) + else: + # Nested pairs keep inherited INNER/OUTER and sparse holes of the parent. + base = blocked_base_tables(pair, parent.layout, order, + composed_kind=self.hierarchy.kinds[index]) + ghosts = pair.ghost_fine if fine is None else fine[0] + donors = pair.ghost_coarse if base is None else base[0] + if fine is not None: + donors = donors.astype(np.uint64) + targets = pair.coalesce_coarse[order] if base is None else base[1] + children = pair.coalesce_fine[order] if fine is None else fine[1] + result = _Pair(child.module.get_function("PatchExplode" if fine is None else "BlockExplode"), + None if parent.layout is None else parent.module.get_function("BlockCoalesce"), + self._upload(ghosts), self._upload(donors), self._upload(targets), + self._upload(children), len(ghosts), len(targets), + int(np.prod(pair.coarse_shape if parent.layout is None + else parent.layout.local_shape))) + if child.layout is not None: + child.kind = self._upload(blocked_kind(child.layout, self.hierarchy.kinds[index + 1])) + if index == 0 and base is not None: + parent.kind = self._upload(base[2]) + return result + def _allocate(self, size): pointer = cuda.mem_alloc(size) self._allocations.append(pointer) @@ -373,28 +426,21 @@ class FixedPatchStepper: def step_count(self): return self._step_count - def _mask_nonleaves(self, coarse, fine, child=None): - coarse[:, self.patch.covered] = np.nan - if self._wall_aware: - coarse[:, (0, -1), :] = np.nan - fine[:, ~self._fine_leaf] = np.nan - if child is not None: - child[:, ~self._child_leaf] = np.nan + def _mask_nonleaves(self, *arrays): + for array, leaf in zip(arrays, self._leaf): + array[:, ~leaf] = np.nan - def initialize(self, coarse_g, fine_g, child_g=None): - """Upload an explicit G_0 (child_g only with a nested child); non-leaves are discarded.""" + def initialize(self, *level_g, child_g=None): + """Upload an explicit G_0, one Q-major array per level; non-leaves are discarded.""" with self._lock, self.ctx.activate(): self._require("new", "ready") - given = (coarse_g, fine_g) + (() if child_g is None else (child_g,)) + given = level_g + (() if child_g is None else (child_g,)) if len(given) != len(self._levels): raise ValueError("Expected one Q-major array per level") arrays = [np.array(a, dtype=np.float32, order="C", copy=True) for a in given] if any(a.shape != level.shape for a, level in zip(arrays, self._levels)): raise ValueError("Expected Q-major (9,ny,nx) arrays for every level") - leaf = ~self.patch.covered - if self._wall_aware: - leaf[[0, -1]] = False - for array, mask in zip(arrays, (leaf, self._fine_leaf, self._child_leaf)): + for array, mask in zip(arrays, self._leaf): if not np.isfinite(array[:, mask]).all() or (array[:, mask].sum(axis=0) <= 0).any(): raise ValueError("Leaf populations must be finite with positive density") self._mask_nonleaves(*arrays) @@ -433,34 +479,35 @@ class FixedPatchStepper: with self._lock: self._require("ready") - if (self._wall_aware or len(self._levels) != 3 - or not isinstance(self.child_patch, FixedPatchSet2D) - or self.fine_boundary is not None - or any(level.boundary is not None for level in self._levels[:2])): + pairs = self.hierarchy.pairs + if (self._wall_aware or self.hierarchy.depth != 3 + or not isinstance(pairs[1], FixedPatchSet2D) + or self._boundary_hosts[1] is not None + or any(b is not None for b in self._level_boundaries[:2])): raise ValueError("remap_child requires periodic three-level L2 sets with no L0/L1 bodies") cfg = self._remap_cfg child_objects = tuple(child_objects) - actual = self._boundaries[2] + actual = self._level_boundaries[2] jets = () if actual is None else actual.jets - old = RemapTopology.from_config(cfg, self.patch, self.child_patch.bounds, + old = RemapTopology.from_config(cfg, pairs[0], pairs[1].bounds, child_objects=child_objects, child_jets=jets) - if (not same_physical_boundary(old.boundary, self.child_boundary) - or (self.child_boundary is not None and actual is None)): + if (not same_physical_boundary(old.boundary, self._boundary_hosts[2]) + or (self._boundary_hosts[2] is not None and actual is None)): raise ValueError("child_objects must reproduce the original root-edge-frame boundary") if actual is not None: velocity = actual.tables.get("surface_velocity") if velocity is None: - velocity = boundary_velocity(cfg, self.child_patch, actual.host, jets) + velocity = boundary_velocity(cfg, pairs[1], actual.host, jets) if (not same_physical_boundary(old.boundary, actual.host) or not np.array_equal(old.surface_velocity, velocity)): raise ValueError("child_objects changed boundary ownership or surface velocity") - new = RemapTopology.from_config(cfg, self.patch, rectangles, + new = RemapTopology.from_config(cfg, pairs[0], rectangles, child_objects=child_objects, child_jets=jets) if (old.pinned_rectangles != new.pinned_rectangles or not same_physical_boundary(old.boundary, new.boundary) or not np.array_equal(old.surface_velocity, new.surface_velocity)): raise ValueError("remap requires the identical pinned body rectangles and boundary") - base_block, fine_block, child_block = self._remap_blocks + base_block, fine_block, child_block = self._level_blocks # Check sparse block alignment/launch size before any device read/allocation. layout = new.child.fine_layout(child_block) if layout.nblocks > MAX_GRID_DIM or layout.local_shape[0] > MAX_GRID_DIM: @@ -477,7 +524,7 @@ class FixedPatchStepper: old, new, source, base_count=count, level_steps=(count, 2*count, 4*count), phase="postcollision_g") candidate = type(self)( - cfg, self.patch.bounds, context=self.ctx, base_block=base_block, + cfg, pairs[0].bounds, context=self.ctx, base_block=base_block, fine_block=fine_block, child_block=child_block, child_rectangles_in_parent=new.child.bounds, child_objects=child_objects, child_jets=jets) @@ -496,13 +543,13 @@ class FixedPatchStepper: """Stage solid source slots; exchange buffer holds this fine substep only.""" cl = self._fine_curved if cl is None: - if self.fine_boundary is not None: + if self._boundary_hosts[1] is not None: raise RuntimeError("Fine boundary device buffers are missing") return if self._curved is None or cl.fluid_idx_gpu is None: raise RuntimeError("Fine boundary device buffers are missing") - driven = bool(self._jets) or any( - obj.surface_actuation is not None for obj in self.fine_boundary.objects) + driven = bool(self._level_jets[1]) or any( + obj.surface_actuation is not None for obj in self._boundary_hosts[1].objects) if driven and (cl.surface_velocity_gpu is None or cl.link_exchange_gpu is None): raise RuntimeError("Surface boundary device buffers are missing") self._curved(buffer, cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu, @@ -523,10 +570,10 @@ class FixedPatchStepper: with self._lock, self.ctx.activate(): self._require("new", "ready") if (isinstance(level, bool) or not isinstance(level, Integral) - or level not in (1, 2) or level >= len(self._levels)): - raise ValueError("Jet level must be 1 or an existing child level 2") - boundary = self._boundaries[level] - jets = list(self._jets if level == 1 else (() if boundary is None else boundary.jets)) + or not 1 <= level < self.hierarchy.depth): + raise ValueError("Jet level must name an existing refined level (1 = first)") + boundary = self._level_boundaries[level] + jets = list(self._level_jets[level]) jets[index] = jets[index].with_amplitude(amplitude) # validates first velocity = None if boundary is not None and boundary.jet_cache is not None: @@ -547,11 +594,7 @@ class FixedPatchStepper: boundary.tables["surface_velocity"] = velocity if boundary is not None: boundary.jets = tuple(jets) - if level == 1: - self._jets = tuple(jets) - - def _blocked_fine_boundary(self, buffer, **kw): - self._blocked_boundary(self._levels[1], buffer, **kw) + self._level_jets[level] = tuple(jets) def _blocked_boundary(self, level, buffer, **kw): b = level.boundary @@ -567,17 +610,45 @@ class FixedPatchStepper: np.uint32(count), velocity if velocity is not None else np.uintp(0), grid=((count + self._block[0] - 1) // self._block[0], 1), **kw) - def _advance_child(self, parent_buffer, **kw): - """Explode from the given L1 buffer, then two L2 hops (halo -> BFL -> step).""" - pair, g = self._child, self._levels[2] - pair.explode(parent_buffer, g.current, pair.ghosts, pair.donors, np.uint32(pair.n_ghosts), + def _advance(self, index, parent_buffer, **kw): + """Explode from the given parent buffer, then two hops of level ``index``.""" + pair, level = self._pairs[index - 1], self._levels[index] + pair.explode(parent_buffer, level.current, pair.ghosts, pair.donors, np.uint32(pair.n_ghosts), np.uint64(pair.parent_cells), grid=((pair.n_ghosts + self._block[0] - 1) // self._block[0], 1), **kw) - for src, dst in ((g.current, g.spare), (g.spare, g.current)): - self._gather_halo(g, src, **kw) - self._blocked_boundary(g, src, **kw) - flags = np.uintp(0) if g.boundary is None else g.boundary.flags - g.kernel(src, dst, np.uint32(g.layout.halo), flags, g.kind, grid=g.grid, **kw) + self._hop(index, level.current, level.spare, **kw) + self._hop(index, level.spare, level.current, **kw) + + def _hop(self, index, src, dst, **kw): + """One interval of level ``index``: a child reads ``src`` (fresh at the midpoint) and + advances twice; then halo -> receiver-local restriction -> BFL -> step on ``src``.""" + level = self._levels[index] + nested = index + 1 < len(self._levels) + if nested: + self._advance(index + 1, src, **kw) + if level.layout is None: + if index == 0: # dense root: PatchCoarseStep fuses restriction and collision + pair = self._pairs[0] + level.kernel(src, dst, self._levels[1].current, pair.targets, pair.children, + np.uint32(pair.n_rows), *(np.uint32(v) for v in self.hierarchy.pairs[0].bounds), + np.uint32(self._wall_aware), grid=level.grid, **kw) + return + if self._boundary_hosts[index] is not None and self._fine_flags is None: + raise RuntimeError("Fine boundary flags are missing") + self._apply_fine_boundary(src, **kw) + flags = self._fine_flags if self._fine_flags is not None else np.uintp(0) + level.kernel(src, dst, flags, grid=level.grid, **kw) + return + # Child donors were read above; no halo gather may overwrite staged arrivals below. + self._gather_halo(level, src, **kw) + if nested: + pair = self._pairs[index] + pair.coalesce(self._levels[index + 1].current, src, pair.targets, pair.children, + np.uint32(pair.n_rows), + grid=((pair.n_rows + self._block[0] - 1) // self._block[0], 1), **kw) + self._blocked_boundary(level, src, **kw) + flags = np.uintp(0) if level.boundary is None else level.boundary.flags + level.kernel(src, dst, np.uint32(level.layout.halo), flags, level.kind, grid=level.grid, **kw) def _gather_halo(self, level, buffer, **kw): pairs = len(level.layout.halo_dst) @@ -586,67 +657,19 @@ class FixedPatchStepper: grid=((pairs + self._block[0] - 1) // self._block[0], 1), **kw) def step(self, n=1): - """Advance G→G: one explosion, two fine collisions, one coarse collision.""" + """Advance G→G by n base intervals: level l takes 2**l hops per interval.""" if isinstance(n, bool) or not isinstance(n, Integral) or n < 0: raise ValueError("n must be a nonnegative integer") with self._lock, self.ctx.activate(): self._require("ready") - c, f = self._levels[:2] # a nested child (third level) runs inside the fine hops - p = self.patch + root = self._levels[0] kw = dict(block=self._block, stream=self._stream) try: self._status = "advancing" for _ in range(n): - self._explode(c.current, f.current, self._ghosts, self._donors, - np.uint32(len(p.ghost_fine)), - np.uint64(np.prod(p.coarse_shape if c.layout is None - else c.layout.local_shape)), - grid=((len(p.ghost_fine) + self._block[0] - 1) // self._block[0], 1), - **kw) - if self.fine_layout is not None: - # Gather one-cell block halos before EACH fine hop. - # Same buffer per hop: halo -> receiver-local BFL -> step. - flags = self._fine_flags if self._fine_flags is not None else np.uintp(0) - for src, dst in ((f.current, f.spare), (f.spare, f.current)): - if self._child is not None: - # L1 feeds L2 from its CURRENT G (fresh at the midpoint), - # L2 takes two hops, then is restricted into src. - self._advance_child(src, **kw) - self._gather_halo(f, src, **kw) - if self._child is not None: - pair = self._child - pair.coalesce(self._levels[2].current, src, pair.targets, - pair.children, np.uint32(pair.n_rows), - grid=((pair.n_rows + self._block[0] - 1) - // self._block[0], 1), **kw) - self._blocked_fine_boundary(src, **kw) - f.kernel(src, dst, np.uint32(self.fine_layout.halo), flags, - f.kind, grid=f.grid, **kw) - else: - flags = self._fine_flags if self._fine_flags is not None else np.uintp(0) - if self.fine_boundary is not None and self._fine_flags is None: - raise RuntimeError("Fine boundary flags are missing") - self._apply_fine_boundary(f.current, **kw) - f.kernel(f.current, f.spare, flags, grid=f.grid, **kw) - self._apply_fine_boundary(f.spare, **kw) - f.kernel(f.spare, f.current, flags, grid=f.grid, **kw) - if c.layout is None: - c.kernel(c.current, c.spare, f.current, self._targets, self._children, - np.uint32(len(p.coalesce_coarse)), - *(np.uint32(v) for v in p.bounds), np.uint32(self._wall_aware), - grid=c.grid, **kw) - else: - # Child donors must be read before staging mutates parent G. - # No halo gather may overwrite these receiver-local arrivals. - self._gather_halo(c, c.current, **kw) - self._coalesce(f.current, c.current, self._targets, self._children, - np.uint32(len(p.coalesce_coarse)), - grid=((len(p.coalesce_coarse) + self._block[0] - 1) - // self._block[0], 1), **kw) - c.kernel(c.current, c.spare, np.uint32(c.layout.halo), np.uintp(0), - c.kind, grid=c.grid, **kw) + self._hop(0, root.current, root.spare, **kw) self._stream.synchronize() - c.current, c.spare = c.spare, c.current + root.current, root.spare = root.spare, root.current self._step_count += 1 self._status = "ready" except BaseException: @@ -684,11 +707,11 @@ class FixedPatchStepper: self._stream.synchronize() except Exception as exc: error = exc - for b in getattr(self, "_boundaries", {}).values(): + for b in getattr(self, "_level_boundaries", ()): if b is not None: b.flags = b.device = b.kernel = None self._levels.clear() - self._explode = self._coalesce = self._child = None + self._pairs = [] if getattr(self, "_fine_curved", None) is not None: try: self._fine_curved.free() @@ -717,6 +740,28 @@ class FixedPatchStepper: if exc[0] is None: raise + # Read-only compatibility views over the level lists. Internal code must not use them; + # REMOVE AFTER the existing tests/docs migrate to hierarchy/_layouts/_level_* lists. + patch = property(lambda self: self.hierarchy.pairs[0]) + child_patch = property(lambda self: self.hierarchy.pairs[1] if self.hierarchy.depth > 2 else None) + base_layout = property(lambda self: self._layouts[0]) + fine_layout = property(lambda self: self._layouts[1]) + child_layout = property(lambda self: self._layouts[2] if len(self._layouts) > 2 else None) + fine_boundary = property(lambda self: self._boundary_hosts[1]) + child_boundary = property(lambda self: self._boundary_hosts[2] if len(self._boundary_hosts) > 2 + else None) + _jets = property(lambda self: self._level_jets[1]) + _boundaries = property(lambda self: {i: b for i, b in enumerate(self._level_boundaries) if i}) + _fine_tables = property(lambda self: None if self._level_boundaries[1] is None + else self._level_boundaries[1].tables) + _actuated = property(lambda self: self._level_boundaries[1] is not None + and self._level_boundaries[1].actuated) + _jet_cache = property(lambda self: None if self._level_boundaries[1] is None + else self._level_boundaries[1].jet_cache) + _fine_leaf = property(lambda self: self._leaf[1]) + _child_leaf = property(lambda self: self._leaf[2] if len(self._leaf) > 2 else None) + _middle_kind = property(lambda self: self.hierarchy.kinds[1] if self.hierarchy.depth > 2 else None) + # Preserve the original entry point and its periodic default. PeriodicPatchStepper = FixedPatchStepper