diff --git a/README.md b/README.md index b312953..7642f11 100644 --- a/README.md +++ b/README.md @@ -1,125 +1,63 @@ # PyFastFlow -**First full-GPU geomorphological and hydrodynamic toolbox powered by Taichi-lang** +> ⚠️ **Experimental.** This describes the in-progress v1 core. The API is +> settling and will change. -[![Python](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/) -[![Taichi](https://img.shields.io/badge/taichi-≥1.6.0-orange.svg)](https://github.com/taichi-dev/taichi) -[![License](https://img.shields.io/badge/license-Custom-red.svg)](./LICENSE) - -## Overview - -**A lot of work will be put into finalising v1.0 in Q1 2026 + paper** +**GPU routines for Earth-surface processes — portable across backends.** -**02/03/2026: I am currently in a hackathlon to refactor the interface to fix it before submission, stay tuned** +Two things in one: +- **A portable GPU-routine engine.** Compose parameters, helpers and + multi-kernel *routines* once, and run them on **Taichi**, **Quadrants**, + or **CuPy** — same composition model, backend-native kernels. +- **A geomorphology toolbox built on it.** Flow routing, flooding, and + landscape evolution, engineered for grids of hundreds of millions of nodes. - ---- +CeCILL v2.1 — Boris Gailleton (Géosciences Rennes) · Guillaume Cordonnier (INRIA). diff --git a/examples/core/depressions/depressions_cupy.py b/examples/core/depressions/depressions_cupy.py new file mode 100644 index 0000000..5de2041 --- /dev/null +++ b/examples/core/depressions/depressions_cupy.py @@ -0,0 +1,139 @@ +""" +Perlin terrain -> receivers -> carve depressions -> drainage area, on cupy. + +The shortest path through the flow stack, on the new builder/frozen/bound +stack (pyfastflow/experimental/core/context/builder.py, frozen.py, bound.py): +noise fills z inline in the init kernel (make_noise_group composes an `at(i)` +device helper, never a field), one make_receivers pass builds the D8 +receiver graph, make_depression_solver resolves every pit by carving, and +make_accumulation sums a unit source over the resolved graph so `q` is +drainage area in cells. + +Every buffer the depression solver touches is allocated here: the flow +factories take no pool and allocate nothing, so scratch is the caller's +throughout. Two things differ from the Taichi/Quadrants files: every launch +carries its own grid/block, and the atomic accumulation is two kernels +rather than one, since a single `__global__` has no grid-wide barrier +between initializing q and accumulating into it. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.cupy_backend import CupyParameter +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool +from pyfastflow.experimental.flow import ( + make_accumulation, + make_depression_solver, + make_depressions, + make_receivers, +) +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +N = 2048 +DX = 50.0 +n_flat = N * N +BLOCK = 256 +LAUNCH = {"grid": ((n_flat + BLOCK - 1) // BLOCK,), "block": (BLOCK,)} + +pool = CupyPool() +grid_group = make_grid_group("cupy", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("cupy", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("cupy", kind="perlin") +noise_params = make_noise_parameters("cupy", pool, kind="perlin", amplitude=300.0, frequency=6.0, octaves=6) + +# z plus every scratch buffer the carve solver and the accumulation need +z = pool.get_data(np.float32, (n_flat,)) +z_prime = pool.get_data(np.float32, (n_flat,)) +q = pool.get_data(np.float32, (n_flat,)) +rec = pool.get_data(np.int32, (n_flat,)) +rec_jump = pool.get_data(np.int32, (n_flat,)) +rec_scratch = pool.get_data(np.int32, (n_flat,)) +bid = pool.get_data(np.int32, (n_flat,)) +basin_saddlenode = pool.get_data(np.int32, (n_flat,)) +basin_saddle = pool.get_data(np.int64, (n_flat,)) +outlet = pool.get_data(np.int64, (n_flat,)) +is_border = pool.get_data(np.uint8, (n_flat,)) +tag = pool.get_data(np.uint8, (n_flat,)) +tag_alt = pool.get_data(np.uint8, (n_flat,)) +rerouted = pool.get_data(np.uint8, (n_flat,)) + +init_bound = ( + KernelBuilder().compose("noise", noise_group).wire_data("z").ingest( + f""" +extern "C" __global__ void init_z(float* z) {{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= {n_flat}) return; + z[i] = $ctx.noise.at(i)$; +}} +""" + ).build() +) +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("cupy", **LAUNCH) + +recv = make_receivers("cupy", grid_group, mode="steepest") +recv_bound = recv["receivers"].build() +recv_bound.bind_leaf(grid_params) +recv_bound.bind("z", z.data) +recv_bound.bind("rec", rec.data) +receivers_kernel = recv_bound.compile("cupy", **LAUNCH) + +ndep = CupyParameter("NDEP", dtype=np.int32, mode="scalar", value=0, pool=pool) +deps = make_depressions("cupy", grid_group, ndep, method="vanilla", reroute="carve", n_flat=n_flat) +solver = make_depression_solver( + "cupy", deps, grid_params, method="vanilla", reroute="carve", + rec=rec.data, z=z.data, bid=bid.data, rec_jump=rec_jump.data, z_prime=z_prime.data, + is_border=is_border.data, basin_saddle=basin_saddle.data, basin_saddlenode=basin_saddlenode.data, + outlet=outlet.data, rerouted=rerouted.data, tag=tag.data, tag_alt=tag_alt.data, + rec_scratch=rec_scratch.data, n_flat=n_flat, block_size=BLOCK, +) + +source = CupyParameter("SRC", dtype=np.float32, mode="const", value=1.0, pool=pool) +accumulation = make_accumulation("cupy", grid_group, method="atomic", n_flat=n_flat) +q_init_bound = accumulation["q_init"].build() +q_init_bound.bind("SOURCE", source) +q_init_bound.bind("q", q.data) +q_init = q_init_bound.compile("cupy", **LAUNCH) +accum_bound = accumulation["accum"].build() +accum_bound.bind("SOURCE", source) +accum_bound.bind("rec", rec.data) +accum_bound.bind("q", q.data) +accum = accum_bound.compile("cupy", **LAUNCH) + +init_kernel() +receivers_kernel() +solver() +q_init() +accum() + +print(f"depressions left: {ndep.read()}, passes taken: {solver.last_trip_counts}") + +zz = z.data.get().reshape(N, N) +qq = q.data.get().reshape(N, N) + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zz, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zz, cmap="terrain", alpha=0.6) +axes[0].set_title("terrain (m)") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/depressions_quadrants.py b/examples/core/depressions/depressions_quadrants.py new file mode 100644 index 0000000..22f5538 --- /dev/null +++ b/examples/core/depressions/depressions_quadrants.py @@ -0,0 +1,127 @@ +""" +Perlin terrain -> receivers -> carve depressions -> drainage area, on Quadrants. + +The shortest path through the flow stack, on the new builder/frozen/bound +stack (pyfastflow/experimental/core/context/builder.py, frozen.py, bound.py): +noise fills z inline in the init kernel (make_noise_group composes an `at(i)` +device helper, never a field), one make_receivers pass builds the D8 +receiver graph, make_depression_solver resolves every pit by carving, and +make_accumulation sums a unit source over the resolved graph so `q` is +drainage area in cells. + +Every buffer the depression solver touches is allocated here: the flow +factories take no pool and allocate nothing, so scratch is the caller's +throughout. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.quadrants_backend import QuadrantsParameter +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool +from pyfastflow.experimental.flow import ( + make_accumulation, + make_depression_solver, + make_depressions, + make_receivers, +) +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +qd.init(arch=qd.gpu) + +N = 1024 +DX = 50.0 +n_flat = N * N + +pool = QuadrantsPool() +grid_group = make_grid_group("quadrants", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("quadrants", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("quadrants", kind="perlin") +noise_params = make_noise_parameters("quadrants", pool, kind="perlin", amplitude=300.0, frequency=6.0, octaves=6) + +# z plus every scratch buffer the carve solver and the accumulation need +z = pool.get_data(qd.f32, (n_flat,)) +z_prime = pool.get_data(qd.f32, (n_flat,)) +q = pool.get_data(qd.f32, (n_flat,)) +rec = pool.get_data(qd.i32, (n_flat,)) +rec_jump = pool.get_data(qd.i32, (n_flat,)) +rec_scratch = pool.get_data(qd.i32, (n_flat,)) +bid = pool.get_data(qd.i32, (n_flat,)) +basin_saddlenode = pool.get_data(qd.i32, (n_flat,)) +basin_saddle = pool.get_data(qd.i64, (n_flat,)) +outlet = pool.get_data(qd.i64, (n_flat,)) +is_border = pool.get_data(qd.u8, (n_flat,)) +tag = pool.get_data(qd.u8, (n_flat,)) +tag_alt = pool.get_data(qd.u8, (n_flat,)) +rerouted = pool.get_data(qd.u8, (n_flat,)) + + +def init_z_tmpl(ctx, z: qd.Tensor): + for i in z: + z[i] = ctx.noise.at(i) + + +init_bound = KernelBuilder().compose("noise", noise_group).wire_data("z").ingest(init_z_tmpl).build() +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("quadrants") + +recv = make_receivers("quadrants", grid_group, mode="steepest") +recv_bound = recv["receivers"].build() +recv_bound.bind_leaf(grid_params) +recv_bound.bind("z", z.data) +recv_bound.bind("rec", rec.data) +receivers_kernel = recv_bound.compile("quadrants") + +ndep = QuadrantsParameter("NDEP", dtype=qd.i32, mode="scalar", value=0, pool=pool) +deps = make_depressions("quadrants", grid_group, ndep, method="optimized", reroute="carve", n_flat=n_flat) +solver = make_depression_solver( + "quadrants", deps, grid_params, method="optimized", reroute="carve", + rec=rec.data, z=z.data, bid=bid.data, rec_jump=rec_jump.data, z_prime=z_prime.data, + is_border=is_border.data, basin_saddle=basin_saddle.data, basin_saddlenode=basin_saddlenode.data, + outlet=outlet.data, rerouted=rerouted.data, tag=tag.data, tag_alt=tag_alt.data, + rec_scratch=rec_scratch.data, n_flat=n_flat, +) + +source = QuadrantsParameter("SRC", dtype=qd.f32, mode="const", value=1.0, pool=pool) +accum_bound = make_accumulation("quadrants", grid_group, method="atomic", n_flat=n_flat)["accum"].build() +accum_bound.bind("SOURCE", source) +accum_bound.bind("rec", rec.data) +accum_bound.bind("q", q.data) +accum_kernel = accum_bound.compile("quadrants") + +init_kernel() +receivers_kernel() +solver() +accum_kernel() + +print(f"depressions left: {ndep.read()}, passes taken: {solver.last_trip_counts}") + +zz = z.data.to_numpy().reshape(N, N) +qq = q.data.to_numpy().reshape(N, N) + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zz, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zz, cmap="terrain", alpha=0.6) +axes[0].set_title("terrain (m)") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/depressions_taichi.py b/examples/core/depressions/depressions_taichi.py new file mode 100644 index 0000000..578ec11 --- /dev/null +++ b/examples/core/depressions/depressions_taichi.py @@ -0,0 +1,127 @@ +""" +Perlin terrain -> receivers -> carve depressions -> drainage area, on Taichi. + +The shortest path through the flow stack, on the new builder/frozen/bound +stack (pyfastflow/experimental/core/context/builder.py, frozen.py, bound.py): +noise fills z inline in the init kernel (make_noise_group composes an `at(i)` +device helper, never a field), one make_receivers pass builds the D8 +receiver graph, make_depression_solver resolves every pit by carving, and +make_accumulation sums a unit source over the resolved graph so `q` is +drainage area in cells. + +Every buffer the depression solver touches is allocated here: the flow +factories take no pool and allocate nothing, so scratch is the caller's +throughout. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.taichi_backend import TaichiParameter +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool +from pyfastflow.experimental.flow import ( + make_accumulation, + make_depression_solver, + make_depressions, + make_receivers, +) +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +ti.init(arch=ti.gpu) + +N = 512 +DX = 50.0 +n_flat = N * N + +pool = TaichiPool() +grid_group = make_grid_group("taichi", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("taichi", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("taichi", kind="perlin") +noise_params = make_noise_parameters("taichi", pool, kind="perlin", amplitude=300.0, frequency=6.0, octaves=6) + +# z plus every scratch buffer the carve solver and the accumulation need +z = pool.get_data(ti.f32, (n_flat,)) +z_prime = pool.get_data(ti.f32, (n_flat,)) +q = pool.get_data(ti.f32, (n_flat,)) +rec = pool.get_data(ti.i32, (n_flat,)) +rec_jump = pool.get_data(ti.i32, (n_flat,)) +rec_scratch = pool.get_data(ti.i32, (n_flat,)) +bid = pool.get_data(ti.i32, (n_flat,)) +basin_saddlenode = pool.get_data(ti.i32, (n_flat,)) +basin_saddle = pool.get_data(ti.i64, (n_flat,)) +outlet = pool.get_data(ti.i64, (n_flat,)) +is_border = pool.get_data(ti.u8, (n_flat,)) +tag = pool.get_data(ti.u8, (n_flat,)) +tag_alt = pool.get_data(ti.u8, (n_flat,)) +rerouted = pool.get_data(ti.u8, (n_flat,)) + + +def init_z_tmpl(ctx, z: ti.template()): + for i in z: + z[i] = ctx.noise.at(i) + + +init_bound = KernelBuilder().compose("noise", noise_group).wire_data("z").ingest(init_z_tmpl).build() +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("taichi") + +recv = make_receivers("taichi", grid_group, mode="steepest") +recv_bound = recv["receivers"].build() +recv_bound.bind_leaf(grid_params) +recv_bound.bind("z", z.data) +recv_bound.bind("rec", rec.data) +receivers_kernel = recv_bound.compile("taichi") + +ndep = TaichiParameter("NDEP", dtype=ti.i32, mode="scalar", value=0, pool=pool) +deps = make_depressions("taichi", grid_group, ndep, method="vanilla", reroute="carve", n_flat=n_flat) +solver = make_depression_solver( + "taichi", deps, grid_params, method="vanilla", reroute="carve", + rec=rec.data, z=z.data, bid=bid.data, rec_jump=rec_jump.data, z_prime=z_prime.data, + is_border=is_border.data, basin_saddle=basin_saddle.data, basin_saddlenode=basin_saddlenode.data, + outlet=outlet.data, rerouted=rerouted.data, tag=tag.data, tag_alt=tag_alt.data, + rec_scratch=rec_scratch.data, n_flat=n_flat, +) + +source = TaichiParameter("SRC", dtype=ti.f32, mode="const", value=1.0, pool=pool) +accum_bound = make_accumulation("taichi", grid_group, method="atomic", n_flat=n_flat)["accum"].build() +accum_bound.bind("SOURCE", source) +accum_bound.bind("rec", rec.data) +accum_bound.bind("q", q.data) +accum_kernel = accum_bound.compile("taichi") + +init_kernel() +receivers_kernel() +solver() +accum_kernel() + +print(f"depressions left: {ndep.read()}, passes taken: {solver.last_trip_counts}") + +zz = z.data.to_numpy().reshape(N, N) +qq = q.data.to_numpy().reshape(N, N) + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zz, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zz, cmap="terrain", alpha=0.6) +axes[0].set_title("terrain (m)") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/fill_reconstruct_cupy.py b/examples/core/depressions/fill_reconstruct_cupy.py new file mode 100644 index 0000000..ec6744d --- /dev/null +++ b/examples/core/depressions/fill_reconstruct_cupy.py @@ -0,0 +1,128 @@ +""" +Perlin terrain -> fill by grayscale morphological reconstruction -> +drainage area, on cupy. + +The reconstruction alternative to depressions_cupy.py's carve/label/saddle +loop, on the new builder/frozen/bound stack (pyfastflow/experimental/core/ +context/builder.py, frozen.py, bound.py): make_fill_reconstruct/ +make_fill_reconstruct_solver converge `filled`/`parent` (the receiver graph) +directly to a fixed point - no basin ids, no saddle search, no outlet +routing. See pyfastflow/experimental/flow/__init__.py's module docstring and +experimental/LM/fill_reconstruct_optimised.py for the algorithm. + +Every buffer the solver touches is allocated here: the factory takes no pool +and allocates nothing, so scratch is the caller's throughout - including the +two buffers this algorithm needs that make_depressions' solver does not: +`frontier` (2*n_flat, the ping-ponged active-cell list - see +make_fill_reconstruct's module note for why it is one combined buffer here) +and `counters` (per-pass frontier sizes), plus `queued_gen`, which - like +`counters` - needs a one-time init before the first call (`-1` and `0` +respectively) since the solver never resets them itself. `active_p` is the +solver's early-stop scalar Parameter (make_fill_reconstruct_solver's own +docstring) - no caller-side init needed, the solver zeroes it every pass. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.cupy_backend import CupyParameter +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool +from pyfastflow.experimental.flow import make_accumulation, make_fill_reconstruct, make_fill_reconstruct_solver +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +N = 2048 +DX = 50.0 +n_flat = N * N +BLOCK = 256 +LAUNCH = {"grid": ((n_flat + BLOCK - 1) // BLOCK,), "block": (BLOCK,)} +MAX_PASSES = 4 * N + +pool = CupyPool() +grid_group = make_grid_group("cupy", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("cupy", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("cupy", kind="perlin") +noise_params = make_noise_parameters("cupy", pool, kind="perlin", amplitude=300.0, frequency=6.0, octaves=6) + +z = pool.get_data(np.float32, (n_flat,)) +filled = pool.get_data(np.float32, (n_flat,)) +parent = pool.get_data(np.int32, (n_flat,)) +frontier = pool.get_data(np.int32, (2 * n_flat,)) +counters = pool.get_data(np.int32, (MAX_PASSES + 2,)) +queued_gen = pool.get_data(np.int32, (n_flat,)) +q = pool.get_data(np.float32, (n_flat,)) + +init_bound = ( + KernelBuilder().compose("noise", noise_group).wire_data("z").ingest( + f""" +extern "C" __global__ void init_z(float* z) {{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= {n_flat}) return; + z[i] = $ctx.noise.at(i)$; +}} +""" + ).build() +) +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("cupy", **LAUNCH) + +pass_p = CupyParameter("PASS", dtype=np.int32, mode="scalar", value=0, pool=pool) +active_p = CupyParameter("ACTIVE", dtype=np.int32, mode="scalar", value=0, pool=pool) +deps = make_fill_reconstruct("cupy", grid_group, nx=N, ny=N) +solver = make_fill_reconstruct_solver( + "cupy", deps, grid_params, + z=z.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + n_flat=n_flat, nx=N, ny=N, block_size=BLOCK, max_passes=MAX_PASSES, +) + +source = CupyParameter("SRC", dtype=np.float32, mode="const", value=1.0, pool=pool) +accumulation = make_accumulation("cupy", grid_group, method="atomic", n_flat=n_flat) +q_init_bound = accumulation["q_init"].build() +q_init_bound.bind("SOURCE", source) +q_init_bound.bind("q", q.data) +q_init = q_init_bound.compile("cupy", **LAUNCH) +accum_bound = accumulation["accum"].build() +accum_bound.bind("SOURCE", source) +accum_bound.bind("rec", parent.data) +accum_bound.bind("q", q.data) +accum = accum_bound.compile("cupy", **LAUNCH) + +init_kernel() +counters.data.fill(0) +queued_gen.data.fill(-1) +solver() +q_init() +accum() + +print(f"reconstruction fill: passes taken = {solver.last_trip_counts}") + +zz = z.data.get().reshape(N, N) +zf = filled.data.get().reshape(N, N) +qq = q.data.get().reshape(N, N) +print(f"cells raised: {int(np.count_nonzero(zf > zz))}/{n_flat}, max raise = {float((zf - zz).max()):.4f} m") + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zf, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zf, cmap="terrain", alpha=0.6) +axes[0].set_title("filled DEM (m), reconstruction") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/fill_reconstruct_epsilon_mfd_cupy.py b/examples/core/depressions/fill_reconstruct_epsilon_mfd_cupy.py new file mode 100644 index 0000000..7cb53fa --- /dev/null +++ b/examples/core/depressions/fill_reconstruct_epsilon_mfd_cupy.py @@ -0,0 +1,239 @@ +""" +Perlin terrain -> fill by grayscale morphological reconstruction -> +"reconstruct_epsilon" (a parent-hop-distance-scaled epsilon added on top of +`filled`) -> persistent-kernel MFD accumulation on that surface, on cupy. + +Mirrors fill_reconstruct_cupy.py, swapping its single-flow-direction +`make_accumulation(method="atomic")` step for multiple-flow-direction +accumulation over the SAME fill_reconstruct result - see +pyfastflow/experimental/graphflood/__init__.py's module docstring, +"kind='vanilla_mfd'" section, for the full algorithm this example runs by +hand (make_graphflood's own kind="vanilla_mfd" wraps exactly this pipeline +plus the graphflood-specific friction/divergence steps this example has no +use for, computing drainage area rather than water depth). + +Why "reconstruct_epsilon" rather than plain `filled`: MFD topology +(pyfastflow/experimental/graphflood/_cupy_mfd_topology.py) derives +`dirs`/`mfd_w` from `slope(filled[i], filled[j]) > 0` between neighbours - +exactly 0 for every pair inside a resolved depression's flat lake bottom, +which gives every cell in there zero outgoing MFD edges and stalls +accumulation at the flat's boundary. SFD accumulation (fill_reconstruct_ +cupy.py's own `accum_bound.bind("rec", parent.data)`) never has this +problem because it walks `parent` directly rather than re-deriving edges +from `filled`'s elevation values. "reconstruct_epsilon" +(pyfastflow/experimental/graphflood/_cupy_reconstruct_epsilon.py) fixes +this without touching _cupy_mfd_topology.py's own slope-based logic at +all: `filled_eps[i] = filled[i] + MFD_EPSILON * hops[i]`, where `hops[i]` +is i's distance to the outlet along `parent` (pointer-jumping, double- +buffered - see build_hops_jump's own docstring for why an earlier, +in-place version of this raced and gave wrong distances). Real slopes are +unaffected; a flat gets a small but strictly monotonic, acyclic synthetic +gradient along the direction `parent` already established. + +Every buffer every step touches is allocated here - none of these +factories take a pool or allocate anything themselves. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.cupy_backend import CupyParameter +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool +from pyfastflow.experimental.flow import make_fill_reconstruct, make_fill_reconstruct_solver +from pyfastflow.experimental.flow._cupy_mfd_accum import build_persistent_mfd, init_frontier_mfd, persistent_grid_block +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.graphflood._cupy_mfd_topology import build_mfd_topology +from pyfastflow.experimental.graphflood._cupy_reconstruct_epsilon import build_apply_epsilon, build_hops_init, build_hops_jump +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +N = 2048 +DX = 50.0 +n_flat = N * N +N_NEIGHBOURS = 8 # D8 +BLOCK = 256 +LAUNCH = {"grid": ((n_flat + BLOCK - 1) // BLOCK,), "block": (BLOCK,)} +MAX_PASSES = 4 * N + +pool = CupyPool() +grid_group = make_grid_group("cupy", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("cupy", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("cupy", kind="perlin") +noise_params = make_noise_parameters("cupy", pool, kind="perlin", amplitude=0.001, frequency=6.0, octaves=6) + +z = pool.get_data(np.float32, (n_flat,)) +filled = pool.get_data(np.float32, (n_flat,)) +parent = pool.get_data(np.int32, (n_flat,)) +frontier = pool.get_data(np.int32, (2 * n_flat,)) +counters = pool.get_data(np.int32, (MAX_PASSES + 2,)) +queued_gen = pool.get_data(np.int32, (n_flat,)) + +dist = pool.get_data(np.float32, (n_flat,)) +anc = pool.get_data(np.int32, (n_flat,)) +dist2 = pool.get_data(np.float32, (n_flat,)) +anc2 = pool.get_data(np.int32, (n_flat,)) +filled_eps = pool.get_data(np.float32, (n_flat,)) + +dirs = pool.get_data(np.uint8, (n_flat,)) +mfd_w = pool.get_data(np.float32, (n_flat * N_NEIGHBOURS,)) +indegree = pool.get_data(np.int32, (n_flat,)) +frontier0 = pool.get_data(np.int32, (n_flat,)) +frontier1 = pool.get_data(np.int32, (n_flat,)) +count = pool.get_data(np.int32, (2,)) +barrier = pool.get_data(np.uint32, (1,)) +q = pool.get_data(np.float32, (n_flat,)) + +init_bound = ( + KernelBuilder().compose("noise", noise_group).wire_data("z").ingest( + f""" +extern "C" __global__ void init_z(float* z) {{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= {n_flat}) return; + z[i] = $ctx.noise.at(i)$; +}} +""" + ).build() +) +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("cupy", **LAUNCH) + +# --- fill by reconstruction -------------------------------------------------- +pass_p = CupyParameter("PASS", dtype=np.int32, mode="scalar", value=0, pool=pool) +active_p = CupyParameter("ACTIVE", dtype=np.int32, mode="scalar", value=0, pool=pool) +deps = make_fill_reconstruct("cupy", grid_group, nx=N, ny=N) +solver = make_fill_reconstruct_solver( + "cupy", deps, grid_params, + z=z.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + n_flat=n_flat, nx=N, ny=N, block_size=BLOCK, max_passes=MAX_PASSES, +) + +# --- reconstruct_epsilon: hops-to-outlet along parent, then filled_eps ----- +hops_init_bound = build_hops_init(n_flat=n_flat).build() +hops_init_bound.bind("parent", parent.data) +hops_init_bound.bind("filled", filled.data) +hops_init_bound.bind("dist", dist.data) +hops_init_bound.bind("anc", anc.data) +hops_init = hops_init_bound.compile("cupy", **LAUNCH) + +hops_jump_fk = build_hops_jump(n_flat=n_flat) +hops_jump_fwd_bound = hops_jump_fk.build() +hops_jump_fwd_bound.bind("dist_in", dist.data) +hops_jump_fwd_bound.bind("anc_in", anc.data) +hops_jump_fwd_bound.bind("dist_out", dist2.data) +hops_jump_fwd_bound.bind("anc_out", anc2.data) +hops_jump_fwd = hops_jump_fwd_bound.compile("cupy", **LAUNCH) + +hops_jump_bwd_bound = hops_jump_fk.build() +hops_jump_bwd_bound.bind("dist_in", dist2.data) +hops_jump_bwd_bound.bind("anc_in", anc2.data) +hops_jump_bwd_bound.bind("dist_out", dist.data) +hops_jump_bwd_bound.bind("anc_out", anc.data) +hops_jump_bwd = hops_jump_bwd_bound.compile("cupy", **LAUNCH) + +# rounded up to even so alternating fwd/bwd always ends back in dist/anc +HOPS_ROUNDS = int(np.ceil(np.log2(max(2, n_flat)))) + 1 +if HOPS_ROUNDS % 2 != 0: + HOPS_ROUNDS += 1 + +apply_epsilon_bound = build_apply_epsilon(n_flat=n_flat).build() +apply_epsilon_bound.bind("filled", filled.data) +apply_epsilon_bound.bind("dist", dist.data) +apply_epsilon_bound.bind("filled_eps", filled_eps.data) +apply_epsilon = apply_epsilon_bound.compile("cupy", **LAUNCH) + +# --- MFD topology on filled_eps, then persistent-kernel accumulation ------- +topo = build_mfd_topology( + grid=grid_group, n_flat=n_flat, topology="D8", diagonal_partition_correction=True, +) +dirs_weights_bound = topo["dirs_weights"].build() +dirs_weights_bound.bind("filled", filled_eps.data) +dirs_weights_bound.bind("dirs", dirs.data) +dirs_weights_bound.bind("mfd_w", mfd_w.data) +dirs_weights_bound.bind_leaf(grid_params) +dirs_weights = dirs_weights_bound.compile("cupy", **LAUNCH) + +indegree_reset_bound = topo["indegree_reset"].build() +indegree_reset_bound.bind("indegree", indegree.data) +indegree_reset = indegree_reset_bound.compile("cupy", **LAUNCH) + +indegree_count_bound = topo["indegree_count"].build() +indegree_count_bound.bind("dirs", dirs.data) +indegree_count_bound.bind("indegree", indegree.data) +indegree_count_bound.bind_leaf(grid_params) +indegree_count = indegree_count_bound.compile("cupy", **LAUNCH) + +source = CupyParameter("SOURCE", dtype=np.float32, mode="const", value=1.0, pool=pool) +persistent = build_persistent_mfd(grid=grid_group, n_flat=n_flat, n_neighbours=N_NEIGHBOURS) +q_init_bound = persistent["q_init"].build() +q_init_bound.bind("SOURCE", source) +q_init_bound.bind("accum", q.data) +q_init = q_init_bound.compile("cupy", **LAUNCH) + +accum_bound = persistent["accum"].build() +accum_bound.bind("frontier0", frontier0.data) +accum_bound.bind("frontier1", frontier1.data) +accum_bound.bind("count", count.data) +accum_bound.bind("barrier", barrier.data) +accum_bound.bind("dirs", dirs.data) +accum_bound.bind("mfd_w", mfd_w.data) +accum_bound.bind("accum", q.data) +accum_bound.bind("indegree", indegree.data) +accum_bound.bind_leaf(grid_params) +persistent_grid, persistent_block = persistent_grid_block() +accum = accum_bound.compile("cupy", grid=persistent_grid, block=persistent_block) + +# --- run --------------------------------------------------------------------- +init_kernel() +counters.data.fill(0) +queued_gen.data.fill(-1) +solver() + +hops_init() +for _ in range(HOPS_ROUNDS // 2): + hops_jump_fwd() + hops_jump_bwd() +apply_epsilon() + +indegree_reset() +dirs_weights() +indegree_count() + +q_init() +n0 = init_frontier_mfd(indegree.data, frontier0.data) +count.data[0:1] = n0 +count.data[1:2] = 0 +barrier.data[0:1] = 0 +accum() + +print(f"reconstruction fill: passes taken = {solver.last_trip_counts}") +print(f"MFD: n0 (ready cells) = {n0}, indegree stuck > 0 after run = {int((indegree.data.get() > 0).sum())}") + +zz = z.data.get().reshape(N, N) +zf = filled.data.get().reshape(N, N) +qq = q.data.get().reshape(N, N) +print(f"cells raised: {int(np.count_nonzero(zf > zz))}/{n_flat}, max raise = {float((zf - zz).max()):.4f} m") + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zf, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zf, cmap="terrain", alpha=0.6) +axes[0].set_title("filled DEM (m), reconstruction") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 MFD drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/fill_reconstruct_epsilon_mfd_greenriver_cupy.py b/examples/core/depressions/fill_reconstruct_epsilon_mfd_greenriver_cupy.py new file mode 100644 index 0000000..9208995 --- /dev/null +++ b/examples/core/depressions/fill_reconstruct_epsilon_mfd_greenriver_cupy.py @@ -0,0 +1,197 @@ +""" +Real DEM (topotoolbox's "greenriver") -> fill by grayscale morphological +reconstruction -> "reconstruct_epsilon" (a parent-hop-distance-scaled +epsilon added on top of `filled`) -> persistent-kernel MFD accumulation on +that surface, on cupy. + +Duplicate of fill_reconstruct_epsilon_mfd_cupy.py with `ttb.load_dem` +swapped in for the Perlin terrain init - see that file's own module +docstring for the full algorithm and for why "reconstruct_epsilon" (not +plain `filled`) is what MFD topology needs fed to it. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +import topotoolbox as ttb +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.cupy_backend import CupyParameter +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool +from pyfastflow.experimental.flow import make_fill_reconstruct, make_fill_reconstruct_solver +from pyfastflow.experimental.flow._cupy_mfd_accum import build_persistent_mfd, init_frontier_mfd, persistent_grid_block +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.graphflood._cupy_mfd_topology import build_mfd_topology +from pyfastflow.experimental.graphflood._cupy_reconstruct_epsilon import build_apply_epsilon, build_hops_init, build_hops_jump + +dem = ttb.load_dem("greenriver") +N_NEIGHBOURS = 8 # D8 +NX, NY, DX = dem.columns, dem.rows, dem.cellsize +n_flat = NX * NY +BLOCK = 256 +LAUNCH = {"grid": ((n_flat + BLOCK - 1) // BLOCK,), "block": (BLOCK,)} +MAX_PASSES = 4 * max(NX, NY) + +pool = CupyPool() +grid_group = make_grid_group("cupy", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("cupy", pool, NX, NY, DX, topology="D8", outlet="edge") + +z = pool.get_data(np.float32, (n_flat,)) +filled = pool.get_data(np.float32, (n_flat,)) +parent = pool.get_data(np.int32, (n_flat,)) +frontier = pool.get_data(np.int32, (2 * n_flat,)) +counters = pool.get_data(np.int32, (MAX_PASSES + 2,)) +queued_gen = pool.get_data(np.int32, (n_flat,)) + +dist = pool.get_data(np.float32, (n_flat,)) +anc = pool.get_data(np.int32, (n_flat,)) +dist2 = pool.get_data(np.float32, (n_flat,)) +anc2 = pool.get_data(np.int32, (n_flat,)) +filled_eps = pool.get_data(np.float32, (n_flat,)) + +dirs = pool.get_data(np.uint8, (n_flat,)) +mfd_w = pool.get_data(np.float32, (n_flat * N_NEIGHBOURS,)) +indegree = pool.get_data(np.int32, (n_flat,)) +frontier0 = pool.get_data(np.int32, (n_flat,)) +frontier1 = pool.get_data(np.int32, (n_flat,)) +count = pool.get_data(np.int32, (2,)) +barrier = pool.get_data(np.uint32, (1,)) +q = pool.get_data(np.float32, (n_flat,)) + +z.from_numpy(dem.z.ravel().astype(np.float32)) + +# --- fill by reconstruction -------------------------------------------------- +pass_p = CupyParameter("PASS", dtype=np.int32, mode="scalar", value=0, pool=pool) +active_p = CupyParameter("ACTIVE", dtype=np.int32, mode="scalar", value=0, pool=pool) +deps = make_fill_reconstruct("cupy", grid_group, nx=NX, ny=NY) +solver = make_fill_reconstruct_solver( + "cupy", deps, grid_params, + z=z.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + n_flat=n_flat, nx=NX, ny=NY, block_size=BLOCK, max_passes=MAX_PASSES, +) + +# --- reconstruct_epsilon: hops-to-outlet along parent, then filled_eps ----- +hops_init_bound = build_hops_init(n_flat=n_flat).build() +hops_init_bound.bind("parent", parent.data) +hops_init_bound.bind("filled", filled.data) +hops_init_bound.bind("dist", dist.data) +hops_init_bound.bind("anc", anc.data) +hops_init = hops_init_bound.compile("cupy", **LAUNCH) + +hops_jump_fk = build_hops_jump(n_flat=n_flat) +hops_jump_fwd_bound = hops_jump_fk.build() +hops_jump_fwd_bound.bind("dist_in", dist.data) +hops_jump_fwd_bound.bind("anc_in", anc.data) +hops_jump_fwd_bound.bind("dist_out", dist2.data) +hops_jump_fwd_bound.bind("anc_out", anc2.data) +hops_jump_fwd = hops_jump_fwd_bound.compile("cupy", **LAUNCH) + +hops_jump_bwd_bound = hops_jump_fk.build() +hops_jump_bwd_bound.bind("dist_in", dist2.data) +hops_jump_bwd_bound.bind("anc_in", anc2.data) +hops_jump_bwd_bound.bind("dist_out", dist.data) +hops_jump_bwd_bound.bind("anc_out", anc.data) +hops_jump_bwd = hops_jump_bwd_bound.compile("cupy", **LAUNCH) + +# rounded up to even so alternating fwd/bwd always ends back in dist/anc +HOPS_ROUNDS = int(np.ceil(np.log2(max(2, n_flat)))) + 1 +if HOPS_ROUNDS % 2 != 0: + HOPS_ROUNDS += 1 + +apply_epsilon_bound = build_apply_epsilon(n_flat=n_flat).build() +apply_epsilon_bound.bind("filled", filled.data) +apply_epsilon_bound.bind("dist", dist.data) +apply_epsilon_bound.bind("filled_eps", filled_eps.data) +apply_epsilon = apply_epsilon_bound.compile("cupy", **LAUNCH) + +# --- MFD topology on filled_eps, then persistent-kernel accumulation ------- +topo = build_mfd_topology( + grid=grid_group, n_flat=n_flat, topology="D8", diagonal_partition_correction=True, +) +dirs_weights_bound = topo["dirs_weights"].build() +dirs_weights_bound.bind("filled", filled_eps.data) +dirs_weights_bound.bind("dirs", dirs.data) +dirs_weights_bound.bind("mfd_w", mfd_w.data) +dirs_weights_bound.bind_leaf(grid_params) +dirs_weights = dirs_weights_bound.compile("cupy", **LAUNCH) + +indegree_reset_bound = topo["indegree_reset"].build() +indegree_reset_bound.bind("indegree", indegree.data) +indegree_reset = indegree_reset_bound.compile("cupy", **LAUNCH) + +indegree_count_bound = topo["indegree_count"].build() +indegree_count_bound.bind("dirs", dirs.data) +indegree_count_bound.bind("indegree", indegree.data) +indegree_count_bound.bind_leaf(grid_params) +indegree_count = indegree_count_bound.compile("cupy", **LAUNCH) + +source = CupyParameter("SOURCE", dtype=np.float32, mode="const", value=1.0, pool=pool) +persistent = build_persistent_mfd(grid=grid_group, n_flat=n_flat, n_neighbours=N_NEIGHBOURS) +q_init_bound = persistent["q_init"].build() +q_init_bound.bind("SOURCE", source) +q_init_bound.bind("accum", q.data) +q_init = q_init_bound.compile("cupy", **LAUNCH) + +accum_bound = persistent["accum"].build() +accum_bound.bind("frontier0", frontier0.data) +accum_bound.bind("frontier1", frontier1.data) +accum_bound.bind("count", count.data) +accum_bound.bind("barrier", barrier.data) +accum_bound.bind("dirs", dirs.data) +accum_bound.bind("mfd_w", mfd_w.data) +accum_bound.bind("accum", q.data) +accum_bound.bind("indegree", indegree.data) +accum_bound.bind_leaf(grid_params) +persistent_grid, persistent_block = persistent_grid_block() +accum = accum_bound.compile("cupy", grid=persistent_grid, block=persistent_block) + +# --- run --------------------------------------------------------------------- +counters.data.fill(0) +queued_gen.data.fill(-1) +solver() + +hops_init() +for _ in range(HOPS_ROUNDS // 2): + hops_jump_fwd() + hops_jump_bwd() +apply_epsilon() + +indegree_reset() +dirs_weights() +indegree_count() + +q_init() +n0 = init_frontier_mfd(indegree.data, frontier0.data) +count.data[0:1] = n0 +count.data[1:2] = 0 +barrier.data[0:1] = 0 +accum() + +print(f"reconstruction fill: passes taken = {solver.last_trip_counts}") +print(f"MFD: n0 (ready cells) = {n0}, indegree stuck > 0 after run = {int((indegree.data.get() > 0).sum())}") + +zz = z.data.get().reshape(NY, NX) +zf = filled.data.get().reshape(NY, NX) +qq = q.data.get().reshape(NY, NX) +print(f"cells raised: {int(np.count_nonzero(zf > zz))}/{n_flat}, max raise = {float((zf - zz).max()):.4f} m") + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zf, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zf, cmap="terrain", alpha=0.6) +axes[0].set_title("filled DEM (m), reconstruction") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 MFD drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/fill_reconstruct_quadrants.py b/examples/core/depressions/fill_reconstruct_quadrants.py new file mode 100644 index 0000000..b7b733b --- /dev/null +++ b/examples/core/depressions/fill_reconstruct_quadrants.py @@ -0,0 +1,119 @@ +""" +Perlin terrain -> fill by grayscale morphological reconstruction -> +drainage area, on Quadrants. + +The reconstruction alternative to depressions_quadrants.py's carve/label/ +saddle loop, on the new builder/frozen/bound stack (pyfastflow/experimental/ +core/context/builder.py, frozen.py, bound.py): make_fill_reconstruct/ +make_fill_reconstruct_solver converge `filled`/`parent` (the receiver graph) +directly to a fixed point - no basin ids, no saddle search, no outlet +routing. See pyfastflow/experimental/flow/__init__.py's module docstring and +experimental/LM/fill_reconstruct_optimised.py for the algorithm. + +Every buffer the solver touches is allocated here: the factory takes no pool +and allocates nothing, so scratch is the caller's throughout - including the +two buffers this algorithm needs that make_depressions' solver does not: +`frontier` (2*n_flat, the ping-ponged active-cell list - see +make_fill_reconstruct's module note for why it is one combined buffer here) +and `counters` (per-pass frontier sizes), plus `queued_gen`, which - like +`counters` - needs a one-time init before the first call (`-1` and `0` +respectively) since the solver never resets them itself. `active_p` is the +solver's early-stop scalar Parameter (make_fill_reconstruct_solver's own +docstring) - no caller-side init needed, the solver zeroes it every pass. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.quadrants_backend import QuadrantsParameter +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool +from pyfastflow.experimental.flow import make_accumulation, make_fill_reconstruct, make_fill_reconstruct_solver +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +qd.init(arch=qd.gpu) + +N = 1024 +DX = 50.0 +n_flat = N * N +MAX_PASSES = 4 * N + +pool = QuadrantsPool() +grid_group = make_grid_group("quadrants", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("quadrants", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("quadrants", kind="perlin") +noise_params = make_noise_parameters("quadrants", pool, kind="perlin", amplitude=300.0, frequency=6.0, octaves=6) + +z = pool.get_data(qd.f32, (n_flat,)) +filled = pool.get_data(qd.f32, (n_flat,)) +parent = pool.get_data(qd.i32, (n_flat,)) +frontier = pool.get_data(qd.i32, (2 * n_flat,)) +counters = pool.get_data(qd.i32, (MAX_PASSES + 2,)) +queued_gen = pool.get_data(qd.i32, (n_flat,)) +q = pool.get_data(qd.f32, (n_flat,)) + + +def init_z_tmpl(ctx, z: qd.Tensor): + for i in z: + z[i] = ctx.noise.at(i) + + +init_bound = KernelBuilder().compose("noise", noise_group).wire_data("z").ingest(init_z_tmpl).build() +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("quadrants") + +pass_p = QuadrantsParameter("PASS", dtype=qd.i32, mode="scalar", value=0, pool=pool) +active_p = QuadrantsParameter("ACTIVE", dtype=qd.i32, mode="scalar", value=0, pool=pool) +deps = make_fill_reconstruct("quadrants", grid_group, nx=N, ny=N) +solver = make_fill_reconstruct_solver( + "quadrants", deps, grid_params, + z=z.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + n_flat=n_flat, nx=N, ny=N, max_passes=MAX_PASSES, +) + +source = QuadrantsParameter("SRC", dtype=qd.f32, mode="const", value=1.0, pool=pool) +accum_bound = make_accumulation("quadrants", grid_group, method="atomic", n_flat=n_flat)["accum"].build() +accum_bound.bind("SOURCE", source) +accum_bound.bind("rec", parent.data) +accum_bound.bind("q", q.data) +accum_kernel = accum_bound.compile("quadrants") + +init_kernel() +counters.data.fill(0) +queued_gen.data.fill(-1) +solver() +accum_kernel() + +print(f"reconstruction fill: passes taken = {solver.last_trip_counts}") + +zz = z.data.to_numpy().reshape(N, N) +zf = filled.data.to_numpy().reshape(N, N) +qq = q.data.to_numpy().reshape(N, N) +print(f"cells raised: {int(np.count_nonzero(zf > zz))}/{n_flat}, max raise = {float((zf - zz).max()):.4f} m") + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zf, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zf, cmap="terrain", alpha=0.6) +axes[0].set_title("filled DEM (m), reconstruction") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/depressions/fill_reconstruct_taichi.py b/examples/core/depressions/fill_reconstruct_taichi.py new file mode 100644 index 0000000..012e613 --- /dev/null +++ b/examples/core/depressions/fill_reconstruct_taichi.py @@ -0,0 +1,119 @@ +""" +Perlin terrain -> fill by grayscale morphological reconstruction -> +drainage area, on Taichi. + +The reconstruction alternative to depressions_taichi.py's carve/label/saddle +loop, on the new builder/frozen/bound stack (pyfastflow/experimental/core/ +context/builder.py, frozen.py, bound.py): make_fill_reconstruct/ +make_fill_reconstruct_solver converge `filled`/`parent` (the receiver graph) +directly to a fixed point - no basin ids, no saddle search, no outlet +routing. See pyfastflow/experimental/flow/__init__.py's module docstring and +experimental/LM/fill_reconstruct_optimised.py for the algorithm. + +Every buffer the solver touches is allocated here: the factory takes no pool +and allocates nothing, so scratch is the caller's throughout - including the +two buffers this algorithm needs that make_depressions' solver does not: +`frontier` (2*n_flat, the ping-ponged active-cell list - see +make_fill_reconstruct's module note for why it is one combined buffer here) +and `counters` (per-pass frontier sizes), plus `queued_gen`, which - like +`counters` - needs a one-time init before the first call (`-1` and `0` +respectively) since the solver never resets them itself. `active_p` is the +solver's early-stop scalar Parameter (make_fill_reconstruct_solver's own +docstring) - no caller-side init needed, the solver zeroes it every pass. + +Author: B.G (08/2026) +""" + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.builder import KernelBuilder +from pyfastflow.experimental.core.context.taichi_backend import TaichiParameter +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool +from pyfastflow.experimental.flow import make_accumulation, make_fill_reconstruct, make_fill_reconstruct_solver +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.noise import make_noise_group, make_noise_parameters + +ti.init(arch=ti.gpu) + +N = 512 +DX = 50.0 +n_flat = N * N +MAX_PASSES = 4 * N + +pool = TaichiPool() +grid_group = make_grid_group("taichi", topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters("taichi", pool, N, N, DX, topology="D8", outlet="edge") +noise_group = make_noise_group("taichi", kind="perlin") +noise_params = make_noise_parameters("taichi", pool, kind="perlin", amplitude=300.0, frequency=6.0, octaves=6) + +z = pool.get_data(ti.f32, (n_flat,)) +filled = pool.get_data(ti.f32, (n_flat,)) +parent = pool.get_data(ti.i32, (n_flat,)) +frontier = pool.get_data(ti.i32, (2 * n_flat,)) +counters = pool.get_data(ti.i32, (MAX_PASSES + 2,)) +queued_gen = pool.get_data(ti.i32, (n_flat,)) +q = pool.get_data(ti.f32, (n_flat,)) + + +def init_z_tmpl(ctx, z: ti.template()): + for i in z: + z[i] = ctx.noise.at(i) + + +init_bound = KernelBuilder().compose("noise", noise_group).wire_data("z").ingest(init_z_tmpl).build() +init_bound.bind_leaf(grid_params, prefix=("noise",)) +init_bound.bind_leaf(noise_params, prefix=("noise",)) +init_bound.bind("z", z.data) +init_kernel = init_bound.compile("taichi") + +pass_p = TaichiParameter("PASS", dtype=ti.i32, mode="scalar", value=0, pool=pool) +active_p = TaichiParameter("ACTIVE", dtype=ti.i32, mode="scalar", value=0, pool=pool) +deps = make_fill_reconstruct("taichi", grid_group, nx=N, ny=N) +solver = make_fill_reconstruct_solver( + "taichi", deps, grid_params, + z=z.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + n_flat=n_flat, nx=N, ny=N, max_passes=MAX_PASSES, +) + +source = TaichiParameter("SRC", dtype=ti.f32, mode="const", value=1.0, pool=pool) +accum_bound = make_accumulation("taichi", grid_group, method="atomic", n_flat=n_flat)["accum"].build() +accum_bound.bind("SOURCE", source) +accum_bound.bind("rec", parent.data) +accum_bound.bind("q", q.data) +accum_kernel = accum_bound.compile("taichi") + +init_kernel() +counters.data.fill(0) +queued_gen.data.fill(-1) +solver() +accum_kernel() + +print(f"reconstruction fill: passes taken = {solver.last_trip_counts}") + +zz = z.data.to_numpy().reshape(N, N) +zf = filled.data.to_numpy().reshape(N, N) +qq = q.data.to_numpy().reshape(N, N) +print(f"cells raised: {int(np.count_nonzero(zf > zz))}/{n_flat}, max raise = {float((zf - zz).max()):.4f} m") + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zf, vert_exag=2.0, dx=DX, dy=DX) + +fig, axes = plt.subplots(1, 2, figsize=(13, 6), constrained_layout=True) +axes[0].imshow(hs, cmap="gray") +im0 = axes[0].imshow(zf, cmap="terrain", alpha=0.6) +axes[0].set_title("filled DEM (m), reconstruction") +fig.colorbar(im0, ax=axes[0], shrink=0.8) + +axes[1].imshow(hs, cmap="gray") +im1 = axes[1].imshow(np.log10(qq), cmap="Blues", alpha=0.7) +axes[1].set_title("log10 drainage area (cells)") +fig.colorbar(im1, ax=axes[1], shrink=0.8) + +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +plt.show() diff --git a/examples/core/graphflood/graphflood_cli.py b/examples/core/graphflood/graphflood_cli.py new file mode 100644 index 0000000..ec376fc --- /dev/null +++ b/examples/core/graphflood/graphflood_cli.py @@ -0,0 +1,321 @@ +""" +GraphFlood CLI: run make_graphflood to steady state (dh/dt convergence) or +n_max on a real DEM (any file ttb.read_tif() accepts), reporting progress +and saving results. + +NoData handling +------------------ +If the DEM has NaNs, the grid is built with `nodata=True, outlet="mask"` +(never the plain `outlet="edge"` this package's other examples use) - +NODATA_MASK marks every NaN cell; OUTLET_MASK marks every DEM-edge cell AND +every cell that neighbours a NaN cell (not the NaN cells themselves - see +below). `z`'s NaNs are replaced with a large finite sentinel for +computation (kept as real NaN only in the array used for the hillshade +plot) - the grid's own `_move_allowed`/`_valid` machinery (../grid/ +_closure_blocks.py) already makes `neighbour()` return -1 for any lookup +that touches a NoData cell on either end, so nothing ever dereferences that +sentinel's neighbours; it exists only so a NoData cell's own z value is +never a bare NaN sitting in the buffer. + +A NoData cell is deliberately NOT put in OUTLET_MASK - only cells +neighbouring one are, per spec. That leaves a NoData cell with no valid +downslope neighbour at all (can never route out) and not can_out either - +harmless for h and Qo (both stay exactly 0 there, since compute_qo's own +neighbour loop only ever executes when the target list is nonempty), but +Q_in would otherwise accumulate that cell's own rain contribution forever +with nowhere for it to go, drifting h upward without bound over many steps. +SOURCE is switched from a uniform const to a field masked to 0 at every +NoData cell to prevent exactly that - a direct consequence of "NoData means +excluded from the simulation", not a change to the physics anywhere real +data exists. + +Convergence metric +--------------------- +The `CONVERGENCE_PERCENTILE`th-percentile (90th, not 99th - see below) +|dh/dt| check at every `--n_check` steps is computed only over currently +"wet" cells (`h > WET_H`, 1cm) - not the whole domain. Most of a real DEM +never floods, or floods far more slowly than the active front; a global +percentile is dominated by that static dry background the moment the wet +fraction is small, reporting a flat, falsely tiny value (essentially +float32 noise around 0) that never reflects whether the actually-flooding +part of the domain is still changing. Before any cell is wet, the check +reports `n_wet=0` and does not count towards convergence. + +The wet-cell set itself is typically a small fraction of the domain, so its +own 99th percentile sits close to the noisy tail of a comparatively small +sample (whichever handful of wet cells happen to be most active right now, +not the bulk). 90th tracks the bulk of the wet region instead - still high +enough to be a "is the flood still changing" check, not a mean. + +Run: + python graphflood_cli.py DEM.tif [--kind mfd|sfd|unstable] [--backend cupy|taichi|quadrants] + [--manning 0.033] [--dt 5e-3] [--rain 50.0] [--n_check 10] [--threshold 1e-5] + [--n_max 5000] [--prefix NAME_out_] + +Author: B.G (08/2026) +""" + +import argparse +import os + +import matplotlib.pyplot as plt +import numpy as np +import topotoolbox as ttb +from matplotlib.colors import LightSource +from scipy.ndimage import binary_dilation + +from pyfastflow.experimental.core.context.backends import backend_classes +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.graphflood import make_graphflood + +FRICTION_EXPONENT = 2.0 / 3.0 +N_NEIGHBOURS = 8 # D8 +NODATA_Z_SENTINEL = 1.0e8 +WET_H = 1e-2 # depth above which a cell counts as "wet" - convergence metric and plot both use this +CONVERGENCE_PERCENTILE = 95 # of |dh/dt| over wet cells - wet fraction is typically small, 99th was too + # close to the noisy tail of a small sample; 90th tracks the bulk instead + +_KIND_MAP = {"sfd": "vanilla_sfd", "mfd": "vanilla_mfd", "unstable": "unstable"} + + +def parse_args(): + p = argparse.ArgumentParser(description="GraphFlood CLI runner") + p.add_argument("dem", type=str, help="path to a DEM readable by topotoolbox.read_tif()") + p.add_argument("--kind", choices=sorted(_KIND_MAP), default="mfd") + p.add_argument("--backend", choices=["taichi", "quadrants", "cupy"], default="cupy") + p.add_argument("--manning", type=float, default=0.033) + p.add_argument("--dt", type=float, default=5e-3, help="timestep, seconds") + p.add_argument("--rain", type=float, default=50.0, help="uniform rain rate, mm/h (converted to m/s internally)") + p.add_argument("--n_check", type=int, default=10) + p.add_argument("--threshold", type=float, default=1e-5) + p.add_argument("--n_max", type=int, default=5000) + p.add_argument("--prefix", type=str, default=None) + return p.parse_args() + + +def main(): + args = parse_args() + kind = _KIND_MAP[args.kind] + backend = args.backend + if kind == "vanilla_mfd" and backend != "cupy": + raise ValueError("--kind mfd is cupy-only") + + prefix = args.prefix + if prefix is None: + prefix = os.path.splitext(os.path.basename(args.dem))[0] + "_out_" + + if backend == "taichi": + import taichi as ti + ti.init(arch=ti.gpu) + from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool as PoolCls + elif backend == "quadrants": + import quadrants as qd + qd.init(arch=qd.gpu) + from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool as PoolCls + else: + from pyfastflow.experimental.core.pool.cupy_pool import CupyPool as PoolCls + + dem = ttb.read_tif(args.dem) + NX, NY, DX = dem.columns, dem.rows, dem.cellsize + n_flat = NX * NY + + z_display = dem.z.astype(np.float32) + nodata_np = ~np.isfinite(z_display) + has_nodata = bool(nodata_np.any()) + + z_np = z_display.copy() + z_np[nodata_np] = NODATA_Z_SENTINEL + + outlet_np = np.zeros((NY, NX), dtype=bool) + outlet_np[0, :] = True + outlet_np[-1, :] = True + outlet_np[:, 0] = True + outlet_np[:, -1] = True + if has_nodata: + touches_nodata = binary_dilation(nodata_np, structure=np.ones((3, 3), dtype=bool)) & ~nodata_np + outlet_np |= touches_nodata + + outlet_mode = "mask" if has_nodata else "edge" + + _, ParamCls, _, dtypes = backend_classes(backend) + i32, i64, f32, u8 = dtypes["i32"], dtypes["i64"], dtypes["f32"], dtypes["u8"] + pool = PoolCls() + + grid_group = make_grid_group(backend, topology="D8", boundary="normal", nodata=has_nodata, outlet=outlet_mode) + grid_params = make_grid_parameters( + backend, pool, NX, NY, DX, topology="D8", nodata=has_nodata, outlet=outlet_mode, + ) + if has_nodata: + grid_params["NODATA_MASK"].get().from_numpy(nodata_np.ravel().astype(np.uint8)) + if outlet_mode == "mask": + grid_params["OUTLET_MASK"].get().from_numpy(outlet_np.ravel().astype(np.uint8)) + + z = pool.get_data(f32, (n_flat,)) + h = pool.get_data(f32, (n_flat,)) + Q_in = pool.get_data(f32, (n_flat,)) + Qo = pool.get_data(f32, (n_flat,)) + z.from_numpy(z_np.ravel()) + h.from_numpy(np.zeros(n_flat, dtype=np.float32)) + + rain_m_s = args.rain * 1e-3 / 3600.0 # mm/h -> m/s + # SOURCE is Q (m^3/s per cell), not a bare rate - apply_divergence computes + # (Q_in - Qo)/area*dt against Qo (m^3/s, from the friction law), so Q_in + # must be in the same units: rain rate * cell area. + rain_q = rain_m_s * DX * DX + source_np = np.where(nodata_np.ravel(), 0.0, rain_q).astype(np.float32) + source_p = ParamCls("SOURCE", dtype=f32, mode="field", value=source_np, pool=pool, n_flat=n_flat) + manning_p = ParamCls("MANNING", dtype=f32, mode="const", value=args.manning, pool=pool) + expo_p = ParamCls("EXPO", dtype=f32, mode="const", value=FRICTION_EXPONENT, pool=pool) + dt_p = ParamCls("DT", dtype=f32, mode="const", value=args.dt, pool=pool) + gf_min_increment_p = ParamCls("GF_MIN_INCREMENT", dtype=f32, mode="const", value=0.0, pool=pool) + boundary_h_p = ParamCls("BOUNDARY_H", dtype=f32, mode="const", value=0.0, pool=pool) + + kwargs = dict( + n_flat=n_flat, nx=NX, ny=NY, z=z.data, h=h.data, Q_in=Q_in.data, Qo=Qo.data, + source_p=source_p, manning_p=manning_p, friction_exponent_p=expo_p, dt_p=dt_p, + gf_min_increment_p=gf_min_increment_p, boundary_h_p=boundary_h_p, + outlet_behavior="fixed_h", kind=kind, + ) + + if kind == "unstable": + Q_next = pool.get_data(f32, (n_flat,)) + kwargs["Q_next"] = Q_next.data + + elif kind == "vanilla_mfd": + surface = pool.get_data(f32, (n_flat,)) + filled = pool.get_data(f32, (n_flat,)) + parent = pool.get_data(i32, (n_flat,)) + frontier = pool.get_data(i32, (2 * n_flat,)) + max_passes = 4 * max(NX, NY) + counters = pool.get_data(i32, (max_passes + 2,)) + queued_gen = pool.get_data(i32, (n_flat,)) + pass_p = ParamCls("P", dtype=i32, mode="scalar", value=0, pool=pool) + active_p = ParamCls("ACTIVE", dtype=i32, mode="scalar", value=0, pool=pool) + dirs = pool.get_data(u8, (n_flat,)) + mfd_w = pool.get_data(f32, (n_flat * N_NEIGHBOURS,)) + indegree = pool.get_data(i32, (n_flat,)) + frontier0 = pool.get_data(i32, (n_flat,)) + frontier1 = pool.get_data(i32, (n_flat,)) + count = pool.get_data(i32, (2,)) + barrier = pool.get_data(dtypes.get("u32", i32), (1,)) + dist = pool.get_data(f32, (n_flat,)) + anc = pool.get_data(i32, (n_flat,)) + dist2 = pool.get_data(f32, (n_flat,)) + anc2 = pool.get_data(i32, (n_flat,)) + filled_eps = pool.get_data(f32, (n_flat,)) + kwargs.update( + surface=surface.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + max_passes=max_passes, dirs=dirs.data, mfd_w=mfd_w.data, indegree=indegree.data, + frontier0=frontier0.data, frontier1=frontier1.data, count=count.data, barrier=barrier.data, + dist=dist.data, anc=anc.data, dist2=dist2.data, anc2=anc2.data, filled_eps=filled_eps.data, + ) + + else: # kind == "vanilla_sfd" + # depression_method="vanilla", not the default "optimized" - see + # this package's memory note depression_optimized_carve_hang.md: + # "optimized"'s carve kernel can hang forever on real DEM data. + kwargs["fill_method"] = "jump" + kwargs["depression_method"] = "vanilla" + rec = pool.get_data(i32, (n_flat,)) + bid = pool.get_data(i32, (n_flat,)) + rec_jump = pool.get_data(i32, (n_flat,)) + z_prime = pool.get_data(f32, (n_flat,)) + is_border = pool.get_data(i32, (n_flat,)) + basin_saddle = pool.get_data(i64, (n_flat,)) + basin_saddlenode = pool.get_data(i32, (n_flat,)) + outlet_h = pool.get_data(i64, (n_flat,)) + rerouted = pool.get_data(i32, (n_flat,)) + tag = pool.get_data(i32, (n_flat,)) + tag_alt = pool.get_data(i32, (n_flat,)) + rec_scratch = pool.get_data(i32, (n_flat,)) + ndep_p = ParamCls("NDEP", dtype=i32, mode="scalar", value=0, pool=pool) + kwargs.update( + rec=rec.data, ndep_p=ndep_p, bid=bid.data, rec_jump=rec_jump.data, z_prime=z_prime.data, + is_border=is_border.data, basin_saddle=basin_saddle.data, basin_saddlenode=basin_saddlenode.data, + outlet=outlet_h.data, rerouted=rerouted.data, tag=tag.data, tag_alt=tag_alt.data, + rec_scratch=rec_scratch.data, + ) + + gf = make_graphflood(backend, grid_group, grid_params, **kwargs) + + h_prev = h.to_numpy() + below_streak = 0 + converged = False + step = 0 + for step in range(1, args.n_max + 1): + gf.step() + if step % args.n_check == 0: + h_now = h.to_numpy() + wet = h_now > WET_H + n_wet = int(wet.sum()) + if n_wet == 0: + # nothing has flooded yet - nothing to declare converged + print(f"step {step}/{args.n_max} n_wet=0 (nothing flooded yet)", flush=True) + h_prev = h_now + below_streak = 0 + continue + dhdt = np.abs(h_now[wet] - h_prev[wet]) / (args.n_check * args.dt) + metric = float(np.percentile(dhdt, CONVERGENCE_PERCENTILE)) + q_in_now = Q_in.to_numpy() + q_out_now = Qo.to_numpy() + print( + f"step {step}/{args.n_max} n_wet={n_wet} {CONVERGENCE_PERCENTILE}th pct |dh/dt| (wet only) = {metric:.6g} " + f"h_max={h_now.max():.6g} Qin_max={q_in_now.max():.6g} Qout_max={q_out_now.max():.6g}", + flush=True, + ) + h_prev = h_now + if metric < args.threshold: + below_streak += 1 + if below_streak >= 2: + converged = True + break + else: + below_streak = 0 + + if converged: + print( + f"converged at step {step} ({CONVERGENCE_PERCENTILE}th pct |dh/dt| over wet cells < " + f"{args.threshold} for 2 consecutive checks)" + ) + else: + print(f"n_max={args.n_max} reached without converging") + + h_np = h.to_numpy().reshape(NY, NX) + q_in_np = Q_in.to_numpy().reshape(NY, NX) + q_out_np = Qo.to_numpy().reshape(NY, NX) + np.save(prefix + "h.npy", h_np) + np.save(prefix + "Qin.npy", q_in_np) + np.save(prefix + "Qout.npy", q_out_np) + + ls = LightSource(azdeg=315, altdeg=45) + hs = ls.hillshade(z_display, vert_exag=2.0, dx=DX, dy=DX) + wet = h_np > WET_H + vmax = float(np.percentile(h_np[wet], 90)) if wet.any() else 1.0 + status = " (converged)" if converged else " (n_max)" + + fig, axes = plt.subplots(1, 3, figsize=(20, 7), constrained_layout=True) + + ax = axes[0] + ax.imshow(hs, cmap="gray") + im = ax.imshow(np.where(wet, h_np, np.nan), cmap="Blues", vmin=0.0, vmax=vmax, alpha=0.8) + fig.colorbar(im, ax=ax, shrink=0.8, label="water depth h (m)") + ax.set_title(f"GraphFlood {args.kind}, {backend}, step {step}{status}") + + for ax, data, name in ((axes[1], q_in_np, "Qin"), (axes[2], q_out_np, "Qout")): + ax.imshow(hs, cmap="gray") + log_data = np.full_like(data, np.nan) + np.log10(data, out=log_data, where=data > 0.0) + im = ax.imshow(log_data, cmap="Blues", alpha=0.8) + fig.colorbar(im, ax=ax, shrink=0.8, label=f"log10 {name} (m3/s)") + ax.set_title(f"{name}, step {step}{status}") + + for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) + fig.savefig(prefix + "hillshade_h_Qin_Qout.png", dpi=150) + print(f"saved {prefix}h.npy, {prefix}Qin.npy, {prefix}Qout.npy, {prefix}hillshade_h_Qin_Qout.png") + + +if __name__ == "__main__": + main() diff --git a/examples/core/graphflood/graphflood_vanilla_sfd_taichi.py b/examples/core/graphflood/graphflood_vanilla_sfd_taichi.py new file mode 100644 index 0000000..f48326a --- /dev/null +++ b/examples/core/graphflood/graphflood_vanilla_sfd_taichi.py @@ -0,0 +1,207 @@ +""" +GraphFlood on a real DEM (topotoolbox's "greenriver") - any backend, any +`kind` (pyfastflow.experimental.graphflood.make_graphflood's own dispatch: +"vanilla_sfd" with fill_method="jump"|"reconstruct", "unstable", or the +cupy-only "vanilla_mfd" - see make_graphflood's own module docstring, +pyfastflow/experimental/graphflood/__init__.py, for what each does). + +Only the buffers the selected `kind`/`fill_method` combination actually +needs are allocated - make_graphflood takes no pool and allocates nothing +itself (every array argument is caller-supplied), so each branch below +allocates exactly its own combination's own required set, per +make_graphflood's own docstring. + +fill_method="jump" pins `depression_method="vanilla"` explicitly, not the +factory's own default "optimized": that method's carve reroute kernel has +an unbounded on-device loop that hangs forever on real DEM data (confirmed +on this exact DEM, all three backends - not GraphFlood-specific, reproduces +calling ../../flow's make_depression_solver directly). "vanilla" solves +this same DEM in under a second. + +Run: + python graphflood_vanilla_sfd_taichi.py [backend] [kind] [fill_method] + backend: taichi (default) | quadrants | cupy + kind: vanilla_sfd (default) | unstable | vanilla_mfd (cupy-only) + fill_method: jump (default) | reconstruct - only used by vanilla_sfd + +Author: B.G (08/2026) +""" + +import sys + +import matplotlib.pyplot as plt +import numpy as np +import topotoolbox as ttb +from matplotlib.colors import LightSource + +from pyfastflow.experimental.core.context.backends import backend_classes +from pyfastflow.experimental.grid import make_grid_group, make_grid_parameters +from pyfastflow.experimental.graphflood import make_graphflood + +BACKEND = sys.argv[1] if len(sys.argv) > 1 else "taichi" +KIND = sys.argv[2] if len(sys.argv) > 2 else "vanilla_sfd" +FILL_METHOD = sys.argv[3] if len(sys.argv) > 3 else "jump" + +if KIND == "vanilla_mfd" and BACKEND != "cupy": + raise ValueError("kind='vanilla_mfd' is cupy-only") + +if BACKEND == "taichi": + import taichi as ti + ti.init(arch=ti.gpu) + from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool as PoolCls +elif BACKEND == "quadrants": + import quadrants as qd + qd.init(arch=qd.gpu) + from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool as PoolCls +elif BACKEND == "cupy": + from pyfastflow.experimental.core.pool.cupy_pool import CupyPool as PoolCls +else: + raise ValueError(f"unknown backend {BACKEND!r}, expected 'taichi', 'quadrants' or 'cupy'") + +N_STEPS = 100 +RAIN = 100e-3 / 3600.0 # 50 mm/hr, in m/s +DT = 1e-2 +MANNING = 0.033 +FRICTION_EXPONENT = 2.0 / 3.0 + +dem = ttb.load_dem("greenriver") +NX, NY, DX = dem.columns, dem.rows, dem.cellsize +n_flat = NX * NY +N_NEIGHBOURS = 8 # D8 + +_, ParamCls, _, dtypes = backend_classes(BACKEND) +i32, i64, f32, u8 = dtypes["i32"], dtypes["i64"], dtypes["f32"], dtypes["u8"] +pool = PoolCls() + +grid_group = make_grid_group(BACKEND, topology="D8", boundary="normal", outlet="edge") +grid_params = make_grid_parameters(BACKEND, pool, NX, NY, DX, topology="D8", outlet="edge") + +# --- buffers/params every kind needs --------------------------------------- +z = pool.get_data(f32, (n_flat,)) +h = pool.get_data(f32, (n_flat,)) +Q_in = pool.get_data(f32, (n_flat,)) +Qo = pool.get_data(f32, (n_flat,)) +z.from_numpy(dem.z.ravel().astype(np.float32)) +h.from_numpy(np.zeros(n_flat, dtype=np.float32)) + +# SOURCE is Q (m^3/s per cell), not a bare rate - apply_divergence computes +# (Q_in - Qo)/area*dt against Qo (m^3/s, from the friction law), so Q_in +# must be in the same units: rain rate * cell area. +source_p = ParamCls("SOURCE", dtype=f32, mode="const", value=RAIN * DX * DX, pool=pool) +manning_p = ParamCls("MANNING", dtype=f32, mode="const", value=MANNING, pool=pool) +expo_p = ParamCls("EXPO", dtype=f32, mode="const", value=FRICTION_EXPONENT, pool=pool) +dt_p = ParamCls("DT", dtype=f32, mode="const", value=DT, pool=pool) +gf_min_increment_p = ParamCls("GF_MIN_INCREMENT", dtype=f32, mode="const", value=0.0, pool=pool) +boundary_h_p = ParamCls("BOUNDARY_H", dtype=f32, mode="const", value=0.0, pool=pool) + +kwargs = dict( + n_flat=n_flat, nx=NX, ny=NY, z=z.data, h=h.data, Q_in=Q_in.data, Qo=Qo.data, + source_p=source_p, manning_p=manning_p, friction_exponent_p=expo_p, dt_p=dt_p, + gf_min_increment_p=gf_min_increment_p, boundary_h_p=boundary_h_p, + outlet_behavior="fixed_h", kind=KIND, +) + +# --- buffers only this kind/fill_method combination needs ------------------- +if KIND == "unstable": + Q_next = pool.get_data(f32, (n_flat,)) + kwargs["Q_next"] = Q_next.data + +elif KIND == "vanilla_mfd": + surface = pool.get_data(f32, (n_flat,)) + filled = pool.get_data(f32, (n_flat,)) + parent = pool.get_data(i32, (n_flat,)) + frontier = pool.get_data(i32, (2 * n_flat,)) + max_passes = 4 * max(NX, NY) + counters = pool.get_data(i32, (max_passes + 2,)) + queued_gen = pool.get_data(i32, (n_flat,)) + pass_p = ParamCls("P", dtype=i32, mode="scalar", value=0, pool=pool) + active_p = ParamCls("ACTIVE", dtype=i32, mode="scalar", value=0, pool=pool) + dirs = pool.get_data(u8, (n_flat,)) + mfd_w = pool.get_data(f32, (n_flat * N_NEIGHBOURS,)) + indegree = pool.get_data(i32, (n_flat,)) + frontier0 = pool.get_data(i32, (n_flat,)) + frontier1 = pool.get_data(i32, (n_flat,)) + count = pool.get_data(i32, (2,)) + barrier = pool.get_data(dtypes.get("u32", i32), (1,)) + dist = pool.get_data(f32, (n_flat,)) + anc = pool.get_data(i32, (n_flat,)) + dist2 = pool.get_data(f32, (n_flat,)) + anc2 = pool.get_data(i32, (n_flat,)) + filled_eps = pool.get_data(f32, (n_flat,)) + kwargs.update( + surface=surface.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + max_passes=max_passes, dirs=dirs.data, mfd_w=mfd_w.data, indegree=indegree.data, + frontier0=frontier0.data, frontier1=frontier1.data, count=count.data, barrier=barrier.data, + dist=dist.data, anc=anc.data, dist2=dist2.data, anc2=anc2.data, filled_eps=filled_eps.data, + ) + +else: # kind == "vanilla_sfd" + kwargs["fill_method"] = FILL_METHOD + if FILL_METHOD == "jump": + # depression_method="vanilla", not the default "optimized": that + # method's own carve reroute kernel (carve_basins_serial, + # _closure_depressions.py) has an unbounded on-device while loop + # that hangs forever on real DEM data (confirmed on greenriver, + # all three backends) - see the memory note + # depression_optimized_carve_hang.md. "vanilla"'s own carve is a + # bounded pointer-jump sweep and solves this same DEM in <1s. + kwargs["depression_method"] = "vanilla" + rec = pool.get_data(i32, (n_flat,)) + bid = pool.get_data(i32, (n_flat,)) + rec_jump = pool.get_data(i32, (n_flat,)) + z_prime = pool.get_data(f32, (n_flat,)) + is_border = pool.get_data(i32, (n_flat,)) + basin_saddle = pool.get_data(i64, (n_flat,)) + basin_saddlenode = pool.get_data(i32, (n_flat,)) + outlet_h = pool.get_data(i64, (n_flat,)) + rerouted = pool.get_data(i32, (n_flat,)) + tag = pool.get_data(i32, (n_flat,)) + tag_alt = pool.get_data(i32, (n_flat,)) + rec_scratch = pool.get_data(i32, (n_flat,)) + ndep_p = ParamCls("NDEP", dtype=i32, mode="scalar", value=0, pool=pool) + kwargs.update( + rec=rec.data, ndep_p=ndep_p, bid=bid.data, rec_jump=rec_jump.data, z_prime=z_prime.data, + is_border=is_border.data, basin_saddle=basin_saddle.data, basin_saddlenode=basin_saddlenode.data, + outlet=outlet_h.data, rerouted=rerouted.data, tag=tag.data, tag_alt=tag_alt.data, + rec_scratch=rec_scratch.data, + ) + else: # "reconstruct" + surface = pool.get_data(f32, (n_flat,)) + filled = pool.get_data(f32, (n_flat,)) + parent = pool.get_data(i32, (n_flat,)) + frontier = pool.get_data(i32, (2 * n_flat,)) + max_passes = 4 * max(NX, NY) + counters = pool.get_data(i32, (max_passes + 2,)) + queued_gen = pool.get_data(i32, (n_flat,)) + pass_p = ParamCls("P", dtype=i32, mode="scalar", value=0, pool=pool) + active_p = ParamCls("ACTIVE", dtype=i32, mode="scalar", value=0, pool=pool) + kwargs.update( + surface=surface.data, filled=filled.data, parent=parent.data, frontier=frontier.data, + counters=counters.data, queued_gen=queued_gen.data, pass_p=pass_p, active_p=active_p, + max_passes=max_passes, + ) + +gf = make_graphflood(BACKEND, grid_group, grid_params, **kwargs) + +for step in range(N_STEPS): + gf.step() + if step % 10 == 0: + h_np = h.to_numpy() + print(f"step {step}/{N_STEPS} h_max={h_np.max():.4g} h_mean={h_np.mean():.4g}") + +zz = z.to_numpy().reshape(NY, NX) +hh = h.to_numpy().reshape(NY, NX) + +ls = LightSource(azdeg=315, altdeg=45) +hs = ls.hillshade(zz, vert_exag=2.0, dx=DX, dy=DX) + +fig, ax = plt.subplots(figsize=(9, 8), constrained_layout=True) +ax.imshow(hs, cmap="gray") +im = ax.imshow(np.where(hh > 1e-3, hh, np.nan), cmap="Blues", vmin=0.0, vmax=0.5, alpha=0.8) +fig.colorbar(im, ax=ax, shrink=0.8, label="water depth h (m)") +title = f"GraphFlood {KIND}" + (f" ({FILL_METHOD})" if KIND == "vanilla_sfd" else "") +ax.set_title(f"{title}, {BACKEND}, {N_STEPS} steps, dt={DT}s") +ax.set_xticks([]) +ax.set_yticks([]) +plt.show() diff --git a/examples/core/heat_diffusion/heat_diffusion_cupy.py b/examples/core/heat_diffusion/heat_diffusion_cupy.py new file mode 100644 index 0000000..fc8d06e --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_cupy.py @@ -0,0 +1,342 @@ +""" +Heat diffusion through a procedurally-generated floor plan (air + walls), +heated by a single stove, built on pyfastflow's backend-agnostic core +(Parameter/Helper/Kernel/Pool), Cupy (cp.RawKernel) backend. + +Same model as heat_diffusion_taichi.py, authored as CUDA source strings. The +grid is stored flat (N*N); kernels launch one thread per cell. Params/helpers +are referenced through `$...$` spans, uniform with the closure backends: + + $wall.get(idx)$ read a field param + $alpha.set_node(idx, v)$ device-side write of a field param + $stove.temp.get(0)$ read a scalar param through a bound Bag + $whash(a, b)$ call a bound __device__ helper (source auto-spliced) + +For scalar/field params the parser auto-generates the matching pointer +argument into the __global__ signature and appends the launch array - the +source never declares them. Top-level const params (N, ROOM, ...) become +#defines and are written bare; a const inside a Bag does not, and is reached +through a span like any other member. Spans do not nest, so a helper's param argument is read into a temp +first (see diffuse: laplacian takes the T_in pointer directly as a data arg). + +Binding styles, all three visible in one file: + - flat, one bind() per object (most kernels here); + - a Bag bound whole and reached by dotted path - `stove` in apply_source, + which nests a sub-bag for the position, and `heat` in diffuse, which mixes + a Parameter, a device helper and two consts under one name; + - bind_bag(), merging a bag's members in flat under their own names, so the + kernel still sees plain names - `alpha_seeds` in set_alpha. + +Author: B.G (07/2026) +""" + +import math +import time + +import cupy as cp +import matplotlib.pyplot as plt +import numpy as np + +from pyfastflow.experimental.core.context.bag import Bag +from pyfastflow.experimental.core.context.cupy_backend import ( + CupyHelperBuilder, + CupyKernelBuilder, + CupyParameter, +) +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool + +# --------------------------------------------------------------------------- +# host-side constants +# --------------------------------------------------------------------------- +GRID_N = 512 +NN = GRID_N * GRID_N +STEPS_PER_FRAME = 10000 +PULSE_FREQ = 0.0 # stove temperature oscillation speed, rad/s (0 = steady stove) +BLOCK = 256 +GRID = (NN + BLOCK - 1) // BLOCK + +# Physical grounding (see the taichi demo for the reasoning). +ROOM_M = 3.0 +DX_M = ROOM_M / (GRID_N // 4) +ALPHA_AIR_VAL = 0.015 # m^2/s, effective convective air diffusivity +ALPHA_WALL_VAL = 1.0e-6 # m^2/s, real solid diffusivity +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) + +pool = CupyPool() + +# Structural constants -> #define, used bare in the CUDA source. +n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool) +room_p = CupyParameter("ROOM", dtype=np.int32, mode="const", value=GRID_N // 4, pool=pool) +wall_thick_p = CupyParameter("WALL_THICK", dtype=np.int32, mode="const", value=8, pool=pool) +door_p = CupyParameter("DOOR", dtype=np.int32, mode="const", value=6, pool=pool) +seed_p = CupyParameter("SEED", dtype=np.float32, mode="const", value=17.0, pool=pool) + +dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool) +dx2_p = CupyParameter("DX2", dtype=np.float32, mode="const", value=DX_M**2, pool=pool) + +alpha_air_seed_p = CupyParameter("ALPHA_AIR_SEED", dtype=np.float32, mode="const", value=ALPHA_AIR_VAL, pool=pool) +alpha_wall_seed_p = CupyParameter("ALPHA_WALL_SEED", dtype=np.float32, mode="const", value=ALPHA_WALL_VAL, pool=pool) +t_bg_p = CupyParameter("T_BG", dtype=np.float32, mode="const", value=15.0, pool=pool) + +src_i_p = CupyParameter("SRC_I", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_j_p = CupyParameter("SRC_J", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_r_p = CupyParameter("SRC_R", dtype=np.int32, mode="const", value=10, pool=pool) + +# scalar mode: host-set each frame -> pulsing stove temperature +OG_stove = 70 +stove_p = CupyParameter("STOVE_T", dtype=np.float32, mode="scalar", value=OG_stove, pool=pool) + +# field mode: per-cell wall/air mask and per-cell diffusivity +wall_p = CupyParameter("WALL", dtype=np.int32, mode="field", value=np.zeros(NN), pool=pool, n_flat=NN) +alpha_p = CupyParameter("ALPHA", dtype=np.float32, mode="field", value=np.zeros(NN), pool=pool, n_flat=NN) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- +clamp_fn = ( + CupyHelperBuilder() + .bind("N", n_p) + .ingest("__device__ int clampi(int i) { return i < 0 ? 0 : (i >= N ? N - 1 : i); }") +) + +laplacian_fn = ( + CupyHelperBuilder() + .bind("N", n_p) + .bind("clampi", clamp_fn) + .ingest( + r""" +__device__ float laplacian(const float* f, int i, int j) { + int ip = $clampi(i + 1)$; + int im = $clampi(i - 1)$; + int jp = $clampi(j + 1)$; + int jm = $clampi(j - 1)$; + return f[ip * N + j] + f[im * N + j] + f[i * N + jp] + f[i * N + jm] - 4.0f * f[i * N + j]; +} +""" + ) +) + +whash_fn = ( + CupyHelperBuilder() + .bind("SEED", seed_p) + .ingest( + r""" +__device__ float whash(int a, int b) { + float x = (float)a * 12.9898f + (float)b * 78.233f + SEED; + float s = sinf(x) * 43758.5453f; + return s - floorf(s); +} +""" + ) +) + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- +generate_walls_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("ROOM", room_p) + .bind("WALL_THICK", wall_thick_p) + .bind("DOOR", door_p) + .bind("wall", wall_p) + .bind("whash", whash_fn) + .ingest( + r""" +__global__ void generate_walls() { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N; + int j = idx % N; + int is_wall = 0; + if (i < WALL_THICK || i >= N - WALL_THICK || j < WALL_THICK || j >= N - WALL_THICK) { + is_wall = 1; + } else if (i % ROOM < WALL_THICK) { + int door = (int)($whash(i / ROOM, j / ROOM)$ * ROOM); + int r = j % ROOM; + if (!(r >= door && r < door + DOOR)) is_wall = 1; + } else if (j % ROOM < WALL_THICK) { + int door = (int)($whash(j / ROOM + 7919, i / ROOM)$ * ROOM); + int r = i % ROOM; + if (!(r >= door && r < door + DOOR)) is_wall = 1; + } + $wall.set_node(idx, is_wall)$; +} +""" + ) + .compile() +) + +# The two seed values are grouped on the host for tidiness, then merged in with +# bind_bag() - which binds each member flat, under its own name. The source +# below is unaware: the seeds stay top-level consts, so they still arrive as +# #defines and are written bare inside the spans. bind() the bag whole instead +# when you want a dotted path. +alpha_seeds = Bag({"ALPHA_WALL_SEED": alpha_wall_seed_p, "ALPHA_AIR_SEED": alpha_air_seed_p}) + +set_alpha_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind_bag(alpha_seeds) + .bind("wall", wall_p) + .bind("alpha", alpha_p) + .ingest( + r""" +__global__ void set_alpha() { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + if ($wall.get(idx)$ == 1) { + $alpha.set_node(idx, ALPHA_WALL_SEED)$; + } else { + $alpha.set_node(idx, ALPHA_AIR_SEED)$; + } +} +""" + ) + .compile() +) + +init_temperature_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("T_BG", t_bg_p) + .ingest( + r""" +__global__ void init_temperature(float* T) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + T[idx] = T_BG; +} +""" + ) + .compile() +) + +# The stove travels as ONE nested Bag rather than four flat binds: its position +# is grouped into an `at` sub-bag, and the span parser walks the dotted path +# through both levels - const members expand to CUDA literals, the scalar +# Parameter to its generated pointer arg. Everything else here still binds +# flat, so the two styles sit side by side in one file. +stove = Bag( + { + "at": Bag({"i": src_i_p, "j": src_j_p}), + "r": src_r_p, + "temp": stove_p, + } +) + +apply_source_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("stove", stove) + .ingest( + r""" +__global__ void apply_source(float* T) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int dx = idx / N - $stove.at.i.get(0)$; + int dy = idx % N - $stove.at.j.get(0)$; + if (dx * dx + dy * dy <= $stove.r.get(0)$ * $stove.r.get(0)$) { + T[idx] = $stove.temp.get(0)$; + } +} +""" + ) + .compile() +) + +# A MIXED Bag: everything the diffusion step needs, whatever kind it is - a +# field Parameter, a device helper, two const Parameters - under one name. A +# bag has no member type; each member is resolved on its own when the spans +# expand. +# Note the consts are reached through spans here rather than written bare: only +# top-level const params become #defines, members of a bag do not. +heat = Bag({"alpha": alpha_p, "lap": laplacian_fn, "dt": dt_p, "dx2": dx2_p}) + +diffuse_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("heat", heat) + .ingest( + r""" +__global__ void diffuse(float* T_out, const float* T_in) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N; + int j = idx % N; + float a = $heat.alpha.get(idx)$; + float lap = $heat.lap(T_in, i, j)$ / $heat.dx2.get(0)$; + T_out[idx] = T_in[idx] + $heat.dt.get(0)$ * a * lap; +} +""" + ) + .compile() +) + +# --------------------------------------------------------------------------- +# fields (pooled - two flat buffers for ping-pong) +# --------------------------------------------------------------------------- +T0 = pool.get_data(np.float32, (NN,)) +T1 = pool.get_data(np.float32, (NN,)) + +generate_walls_kernel(grid=GRID, block=BLOCK) +set_alpha_kernel(grid=GRID, block=BLOCK) +init_temperature_kernel(T0.data, grid=GRID, block=BLOCK) +apply_source_kernel(T0.data, grid=GRID, block=BLOCK) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy().reshape(GRID_N, GRID_N), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall_p.get().to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (Cupy backend)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + clock += PULSE_FREQ * DT_VAL + stove_p.set(OG_stove + 20.0 * math.sin(clock)) + + diffuse_kernel(T1.data, T0.data, grid=GRID, block=BLOCK) + apply_source_kernel(T1.data, grid=GRID, block=BLOCK) + T0, T1 = T1, T0 + sim_time += DT_VAL + + cp.cuda.Device().synchronize() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy().reshape(GRID_N, GRID_N)) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +# destroy() hands a Parameter's storage back to the pool; it is a no-op on a +# const, which owns none. Safe only because nothing will launch again - the +# pool may reissue these buffers, while the compiled kernels above still point +# at them (see parameter.py, "Lifetime of a compiled object"). +for param in (stove_p, wall_p, alpha_p): + param.destroy() +pool.release_data(T0) +pool.release_data(T1) +print("pooled storage released") diff --git a/examples/core/heat_diffusion/heat_diffusion_pure_taichi.py b/examples/core/heat_diffusion/heat_diffusion_pure_taichi.py new file mode 100644 index 0000000..af2ed19 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_pure_taichi.py @@ -0,0 +1,194 @@ +""" +Heat diffusion through a procedurally-generated floor plan (air + walls), +heated by a single stove - PURE Taichi, no pyfastflow framework. + +Byte-for-byte the same model as heat_diffusion_taichi.py, written with plain +ti.field / ti.func / ti.kernel and module-global constants, so it can be timed +against the framework version as a zero-abstraction baseline. Constants are +captured as compile-time literals by Taichi directly (the framework's const +mode does the same via bindings); wall/alpha are flat fields; the stove +temperature is a 0-d field host-set each substep. + +Author: B.G (07/2026) +""" + +import math +import time + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti + +ti.init(arch=ti.gpu) + +# --------------------------------------------------------------------------- +# constants (baked into kernels as literals) +# --------------------------------------------------------------------------- +GRID_N = 512 +STEPS_PER_FRAME = 10000 +PULSE_FREQ = 0.0 + +ROOM_M = 3.0 +DX_M = ROOM_M / (GRID_N // 4) +ALPHA_AIR_VAL = 0.015 +ALPHA_WALL_VAL = 1.0e-6 +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) +DX2_VAL = DX_M**2 + +N = GRID_N +ROOM = GRID_N // 4 +WALL_THICK = 8 +DOOR = 6 +SEED = 17.0 +T_BG = 15.0 +SRC_I = GRID_N // 4 + GRID_N // 8 +SRC_J = GRID_N // 4 + GRID_N // 8 +SRC_R = 10 +OG_stove = 70 + +# --------------------------------------------------------------------------- +# fields +# --------------------------------------------------------------------------- +T0 = ti.field(ti.f32, shape=(GRID_N, GRID_N)) +T1 = ti.field(ti.f32, shape=(GRID_N, GRID_N)) +wall = ti.field(ti.i32, shape=(GRID_N * GRID_N,)) +alpha = ti.field(ti.f32, shape=(GRID_N * GRID_N,)) +stove_t = ti.field(ti.f32, shape=()) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +@ti.func +def clamp(i): + return min(max(i, 0), N - 1) + + +@ti.func +def laplacian(field_, i, j): + ip = clamp(i + 1) + im = clamp(i - 1) + jp = clamp(j + 1) + jm = clamp(j - 1) + return field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + + +@ti.func +def whash(a, b): + x = ti.cast(a, ti.f32) * 12.9898 + ti.cast(b, ti.f32) * 78.233 + SEED + s = ti.sin(x) * 43758.5453 + return s - ti.floor(s) + + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- + + +@ti.kernel +def generate_walls(): + for i, j in ti.ndrange(N, N): + is_wall = 0 + if i < WALL_THICK or i >= N - WALL_THICK or j < WALL_THICK or j >= N - WALL_THICK: + is_wall = 1 + elif i % ROOM < WALL_THICK: + vline = i // ROOM + seg = j // ROOM + door = ti.cast(whash(vline, seg) * ROOM, ti.i32) + gap = (j % ROOM) >= door and (j % ROOM) < door + DOOR + if not gap: + is_wall = 1 + elif j % ROOM < WALL_THICK: + hline = j // ROOM + seg = i // ROOM + door = ti.cast(whash(hline + 7919, seg) * ROOM, ti.i32) + gap = (i % ROOM) >= door and (i % ROOM) < door + DOOR + if not gap: + is_wall = 1 + wall[i * N + j] = is_wall + + +@ti.kernel +def set_alpha(): + for i, j in ti.ndrange(N, N): + idx = i * N + j + if wall[idx] == 1: + alpha[idx] = ALPHA_WALL_VAL + else: + alpha[idx] = ALPHA_AIR_VAL + + +@ti.kernel +def init_temperature(T: ti.template()): + for i, j in T: + T[i, j] = T_BG + + +@ti.kernel +def apply_source(T: ti.template()): + for i, j in T: + dx = i - SRC_I + dy = j - SRC_J + if dx * dx + dy * dy <= SRC_R * SRC_R: + T[i, j] = stove_t[None] + + +@ti.kernel +def diffuse(T_out: ti.template(), T_in: ti.template()): + for i, j in T_in: + idx = i * N + j + a = alpha[idx] + lap = laplacian(T_in, i, j) / DX2_VAL + T_out[i, j] = T_in[i, j] + DT_VAL * a * lap + + +# --------------------------------------------------------------------------- +# setup +# --------------------------------------------------------------------------- +stove_t[None] = OG_stove +generate_walls() +set_alpha() +init_temperature(T0) +apply_source(T0) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy(), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall.to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (pure Taichi)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + clock += PULSE_FREQ * DT_VAL + stove_t[None] = OG_stove + 20.0 * math.sin(clock) + + diffuse(T1, T0) + apply_source(T1) + T0, T1 = T1, T0 + sim_time += DT_VAL + + ti.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) diff --git a/examples/core/heat_diffusion/heat_diffusion_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_quadrants.py new file mode 100644 index 0000000..0c97d39 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_quadrants.py @@ -0,0 +1,336 @@ +""" +Heat diffusion through a procedurally-generated floor plan (air + walls), +heated by a single stove, built on pyfastflow's backend-agnostic core +(Parameter/Helper/Kernel/Pool), Quadrants backend. + +Pipeline: + - generate_walls: pointwise kernel, carves a grid of rooms with doors into + a wall mask (field-mode Parameter `wall`, written device-side via + wall.set_node), using a deterministic hash device helper instead of + per-cell RNG so wall/door layout is reproducible from SEED. + - set_alpha: seeds a per-cell diffusivity field (`alpha`, field mode) from + `wall` - air diffuses fast, walls slow. + - init_temperature: fills T with the background temperature. + - apply_source: clamps a disc of cells around the stove to `stove.temp` + (scalar-mode Parameter, updated from the host every substep -> a gently + pulsing stove). + - diffuse: explicit FTCS heat equation dT/dt = alpha(i,j) * lap(T), with a + clamped (Neumann / no-flux) boundary Laplacian. + +Uniform device surface: every Parameter is read with p.get(node) and written +with p.set_node(node, val) regardless of const/scalar/field mode - the +kernels never branch on mode, so re-declaring `alpha` as a single const +(uniform room, no walls) needs no kernel-body change. + +Binding styles, all three visible in one file: + - flat, one bind() per object (most kernels here); + - a Bag bound whole and reached by dotted path - `stove` in apply_source, + which nests a sub-bag for the position, and `heat` in diffuse, which mixes + a Parameter, a device helper and two consts under one name; + - bind_bag(), merging a bag's members in flat under their own names, so the + kernel still sees plain names - `alpha_seeds` in set_alpha. + +Compilation is the two-layer builder: QuadrantsKernelBuilder / +QuadrantsHelperBuilder collect bind()ed params + helper builders and one +ingest()ed template. A HelperBuilder (clamp_fn, laplacian_fn, whash_fn below) +is a recipe, not a compiled object - it has no compile() of its own. Binding +one into a kernel, flat or through a Bag, is what specializes it: +QuadrantsKernelBuilder.compile() specializes every HelperBuilder the kernel +reaches, against that kernel's own bindings, before compiling the kernel +body. Recompiling the kernel after rebinding a const the helper reads picks +up the new value without touching the helper builder itself. + +Author: B.G (07/2026) +""" + +import math +import time + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd + +from pyfastflow.experimental.core.context.bag import Bag +from pyfastflow.experimental.core.context.quadrants_backend import ( + QuadrantsHelperBuilder, + QuadrantsKernelBuilder, + QuadrantsParameter, +) +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool + +qd.init(arch=qd.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, loop/timing counts - never used as kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 512 +STEPS_PER_FRAME = 10000 +PULSE_FREQ = 0.0 # stove temperature oscillation speed, rad/s (0 = steady stove) + +# Physical grounding: without a cell size, DT/ALPHA are just numbers tuned by +# feel - here they're derived from a real room size and real diffusivities so +# "seconds" and "m^2/s" mean what they say. +ROOM_M = 3.0 # room span, meters (rooms are GRID_N//4 cells across) +DX_M = ROOM_M / (GRID_N // 4) # meters per cell + +# Air's real molecular thermal diffusivity (~2.2e-5 m^2/s) would take DAYS to +# spread heat by pure conduction - rooms actually heat by convective mixing. +# ALPHA_AIR below is an effective/turbulent diffusivity standing in for that +# mixing, not molecular diffusion - otherwise a stove would need real hours. +ALPHA_AIR_VAL = 0.015 # m^2/s, effective convective air diffusivity +ALPHA_WALL_VAL = 1.0e-6 # m^2/s, real solid (drywall/brick-like) diffusivity + +# Explicit FTCS stability limit is dt <= dx^2 / (4*alpha); stay well under it. +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) # seconds + +pool = QuadrantsPool() + +# Structural constants: const mode, bake to compile-time literals in generated +# code even though the kernel body still reads them via .get(0). +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool) +room_p = QuadrantsParameter("ROOM", dtype=qd.i32, mode="const", value=GRID_N // 4, pool=pool) +wall_thick_p = QuadrantsParameter("WALL_THICK", dtype=qd.i32, mode="const", value=8, pool=pool) +door_p = QuadrantsParameter("DOOR", dtype=qd.i32, mode="const", value=6, pool=pool) +seed_p = QuadrantsParameter("SEED", dtype=qd.f32, mode="const", value=17.0, pool=pool) + +dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool) # seconds +dx2_p = QuadrantsParameter("DX2", dtype=qd.f32, mode="const", value=DX_M**2, pool=pool) # meters^2 + +# Seed values for the alpha field - read via .get(0) inside set_alpha. +alpha_air_seed_p = QuadrantsParameter("ALPHA_AIR_SEED", dtype=qd.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool) +alpha_wall_seed_p = QuadrantsParameter("ALPHA_WALL_SEED", dtype=qd.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool) +t_bg_p = QuadrantsParameter("T_BG", dtype=qd.f32, mode="const", value=15.0, pool=pool) + +src_i_p = QuadrantsParameter("SRC_I", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_j_p = QuadrantsParameter("SRC_J", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_r_p = QuadrantsParameter("SRC_R", dtype=qd.i32, mode="const", value=10, pool=pool) # stove radius, cells + +# scalar mode: a 0-d field, host-settable every frame -> a pulsing stove +# temperature. Reached in-kernel as stove.temp.get(0) (see the stove Bag). +OG_stove = 70 +stove_p = QuadrantsParameter("STOVE_T", dtype=qd.f32, mode="scalar", value=OG_stove, pool=pool) + +# field mode: per-cell wall/air mask, written device-side via wall.set_node, +# read via wall.get. +wall_p = QuadrantsParameter("WALL", dtype=qd.i32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) + +# field mode: per-cell thermal diffusivity, read in diffuse via alpha.get - so +# switching this Parameter to const/scalar mode later needs no kernel edits. +alpha_p = QuadrantsParameter("ALPHA", dtype=qd.f32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +def clamp(i): + return min(max(i, 0), N.get(0) - 1) + + +clamp_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clamp) + + +def laplacian(field_, i, j): + ip = clamp(i + 1) + im = clamp(i - 1) + jp = clamp(j + 1) + jm = clamp(j - 1) + return field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + + +laplacian_fn = QuadrantsHelperBuilder().bind("clamp", clamp_fn).ingest(laplacian) + + +def whash(a, b): + """Deterministic pseudo-random value in [0, 1) for two integer indices.""" + x = qd.cast(a, qd.f32) * 12.9898 + qd.cast(b, qd.f32) * 78.233 + SEED.get(0) + s = qd.sin(x) * 43758.5453 + return s - qd.floor(s) + + +whash_fn = QuadrantsHelperBuilder().bind("SEED", seed_p).ingest(whash) + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- + + +def generate_walls_template(): + for i, j in qd.ndrange(N.get(0), N.get(0)): + is_wall = 0 + if i < WALL_THICK.get(0) or i >= N.get(0) - WALL_THICK.get(0) or j < WALL_THICK.get(0) or j >= N.get(0) - WALL_THICK.get(0): + is_wall = 1 + elif i % ROOM.get(0) < WALL_THICK.get(0): + vline = i // ROOM.get(0) + seg = j // ROOM.get(0) + door = qd.cast(whash(vline, seg) * ROOM.get(0), qd.i32) + gap = (j % ROOM.get(0)) >= door and (j % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + elif j % ROOM.get(0) < WALL_THICK.get(0): + hline = j // ROOM.get(0) + seg = i // ROOM.get(0) + door = qd.cast(whash(hline + 7919, seg) * ROOM.get(0), qd.i32) + gap = (i % ROOM.get(0)) >= door and (i % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + wall.set_node(i * N.get(0) + j, is_wall) + + +generate_walls_kernel = ( + QuadrantsKernelBuilder() + .bind("N", n_p) + .bind("ROOM", room_p) + .bind("WALL_THICK", wall_thick_p) + .bind("DOOR", door_p) + .bind("wall", wall_p) + .bind("whash", whash_fn) + .ingest(generate_walls_template) + .compile() +) + + +def set_alpha_template(): + for i, j in qd.ndrange(N.get(0), N.get(0)): + idx = i * N.get(0) + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) + else: + alpha.set_node(idx, ALPHA_AIR_SEED.get(0)) + + +# The two seed values are grouped on the host for tidiness, then merged in with +# bind_bag() - which binds each member flat, under its own name. The template +# above is unaware: it still reads ALPHA_WALL_SEED / ALPHA_AIR_SEED bare. Use +# this when a bag is a convenient way to carry things around but the kernel +# wants plain names; bind() the bag whole instead when you want a dotted path. +alpha_seeds = Bag({"ALPHA_WALL_SEED": alpha_wall_seed_p, "ALPHA_AIR_SEED": alpha_air_seed_p}) + +set_alpha_kernel = ( + QuadrantsKernelBuilder() + .bind("N", n_p) + .bind("wall", wall_p) + .bind("alpha", alpha_p) + .bind_bag(alpha_seeds) + .ingest(set_alpha_template) + .compile() +) + + +def init_temperature_template(T: qd.Tensor): + for i, j in T: + T[i, j] = T_BG.get(0) + + +init_temperature_kernel = QuadrantsKernelBuilder().bind("T_BG", t_bg_p).ingest(init_temperature_template).compile() + + +# The stove travels as ONE nested Bag rather than four flat binds: its position +# is grouped into an `at` sub-bag. Every member, whatever mode, is reached the +# same way - .get(0) - so const and scalar Parameters sit side by side under +# one name. Everything else here still binds flat, so the two styles sit side +# by side in one file. +stove = Bag( + { + "at": Bag({"i": src_i_p, "j": src_j_p}), + "r": src_r_p, + "temp": stove_p, + } +) + + +def apply_source_template(T: qd.Tensor): + for i, j in T: + dx = i - stove.at.i.get(0) + dy = j - stove.at.j.get(0) + if dx * dx + dy * dy <= stove.r.get(0) * stove.r.get(0): + T[i, j] = stove.temp.get(0) + + +apply_source_kernel = QuadrantsKernelBuilder().bind("stove", stove).ingest(apply_source_template).compile() + + +# A MIXED Bag: everything the diffusion step needs, whatever kind it is - a +# field Parameter, a device helper, two const Parameters - under one name. A +# bag has no member type; each is resolved on its own at compile time, so +# `heat.alpha` becomes a device accessor, `heat.lap` a compiled func, and +# `heat.dx2` a device accessor whose .get(0) bakes to a literal. +heat = Bag({"alpha": alpha_p, "lap": laplacian_fn, "dt": dt_p, "dx2": dx2_p}) + + +def diffuse_template(T_out: qd.Tensor, T_in: qd.Tensor): + for i, j in T_in: + idx = i * N.get(0) + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2.get(0) + T_out[i, j] = T_in[i, j] + heat.dt.get(0) * a * lap + +diffuse_kernel = QuadrantsKernelBuilder().bind("N", n_p).bind("heat", heat).ingest(diffuse_template).compile() + +# --------------------------------------------------------------------------- +# fields (pooled - two buffers for ping-pong) +# --------------------------------------------------------------------------- +T0 = pool.get_data(qd.f32, (GRID_N, GRID_N)) +T1 = pool.get_data(qd.f32, (GRID_N, GRID_N)) + +generate_walls_kernel() +set_alpha_kernel() +init_temperature_kernel(T0.data) +apply_source_kernel(T0.data) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy(), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall_p.get().to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (Quadrants backend)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + clock += PULSE_FREQ * DT_VAL + stove_p.set(OG_stove + 20.0 * math.sin(clock)) + + diffuse_kernel(T1.data, T0.data) + apply_source_kernel(T1.data) + T0, T1 = T1, T0 + sim_time += dt_p.get() + + qd.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +# destroy() hands a Parameter's storage back to the pool; it is a no-op on a +# const, which owns none. Safe only because nothing will launch again - the +# pool may reissue these buffers, while the compiled kernels above still point +# at them (see parameter.py, "Lifetime of a compiled object"). +for param in (stove_p, wall_p, alpha_p): + param.destroy() +pool.release_data(T0) +pool.release_data(T1) +print("pooled storage released") diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py new file mode 100644 index 0000000..3467de6 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py @@ -0,0 +1,358 @@ +""" +Same model and setup as heat_diffusion_cupy.py, with the per-substep +ping-pong expressed as a Routine instead of a hand-written python loop. + +heat_diffusion_cupy.py alternates `diffuse_kernel(T1, T0, ...)`, +`apply_source_kernel(T1, ...)`, then swaps the T0/T1 python names each +iteration. A Routine has no python between its steps to do that swap in, so +the two iterations that one swap-pair covers are unrolled into one routine +with a repeat block, add_swap standing in for the python-level +`T0, T1 = T1, T0`: + + begin_repeat(times=2) + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + end_repeat() + +which records the body once and replays it twice, giving the same six-step +sequence as writing it out by hand: + + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + +Two swaps compose to the identity, which is exactly what compile() checks +for - so the compiled routine can be called over and over, each call +advancing the simulation by two substeps, and the result always ends up back +in the T0 buffer, matching two iterations of the manual loop. + +diffuse_builder and apply_source_builder are ordinary KernelBuilders, built +exactly as in heat_diffusion_cupy.py; apply_source_builder is also compiled +once on its own to seed T0 before the loop starts, same as that file does - +compile() does not consume a builder, so the same builder is later handed to +add_kernel() unchanged. The routine's one shared bag is the merge of what +each builder already binds, so nothing about either CUDA source template +changes. + +cupy has no auto-ranging launch the way Taichi/Quadrants derive one from the +template, so grid/block are set once on CupyRoutineBuilder's constructor and +apply to every step that does not override them - see cupy_backend.py, +CupyRoutineBuilder. + +The stove's pulse is only updated between routine() calls, not between the +two substeps a single call unrolls: set() on the stove's scalar Parameter is +safe between calls (see routine.py, "Contract: no set()/destroy() +mid-routine"), but doing it *inside* a routine's steps is exactly what that +contract forbids, since there is no python between steps to run it in. With +PULSE_FREQ=0.0 by default the stove is steady anyway and this has no visible +effect. + +Author: B.G (07/2026) +""" + +import math +import time + +import cupy as cp +import matplotlib.pyplot as plt +import numpy as np + +from pyfastflow.experimental.core.context.bag import Bag, merge +from pyfastflow.experimental.core.context.cupy_backend import ( + CupyHelperBuilder, + CupyKernelBuilder, + CupyParameter, + CupyRoutineBuilder, +) +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool + +# --------------------------------------------------------------------------- +# host-side constants +# --------------------------------------------------------------------------- +GRID_N = 512 +NN = GRID_N * GRID_N +STEPS_PER_FRAME = 10000 # two routine substeps per call - see the loop below +PULSE_FREQ = 0.0 # stove temperature oscillation speed, rad/s (0 = steady stove) +BLOCK = 256 +GRID = (NN + BLOCK - 1) // BLOCK + +ROOM_M = 3.0 +DX_M = ROOM_M / (GRID_N // 4) +ALPHA_AIR_VAL = 0.015 +ALPHA_WALL_VAL = 1.0e-6 +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) + +pool = CupyPool() + +n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool) +room_p = CupyParameter("ROOM", dtype=np.int32, mode="const", value=GRID_N // 4, pool=pool) +wall_thick_p = CupyParameter("WALL_THICK", dtype=np.int32, mode="const", value=8, pool=pool) +door_p = CupyParameter("DOOR", dtype=np.int32, mode="const", value=6, pool=pool) +seed_p = CupyParameter("SEED", dtype=np.float32, mode="const", value=17.0, pool=pool) + +dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool) +dx2_p = CupyParameter("DX2", dtype=np.float32, mode="const", value=DX_M**2, pool=pool) + +alpha_air_seed_p = CupyParameter("ALPHA_AIR_SEED", dtype=np.float32, mode="const", value=ALPHA_AIR_VAL, pool=pool) +alpha_wall_seed_p = CupyParameter("ALPHA_WALL_SEED", dtype=np.float32, mode="const", value=ALPHA_WALL_VAL, pool=pool) +t_bg_p = CupyParameter("T_BG", dtype=np.float32, mode="const", value=15.0, pool=pool) + +src_i_p = CupyParameter("SRC_I", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_j_p = CupyParameter("SRC_J", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_r_p = CupyParameter("SRC_R", dtype=np.int32, mode="const", value=10, pool=pool) + +OG_stove = 70 +stove_p = CupyParameter("STOVE_T", dtype=np.float32, mode="scalar", value=OG_stove, pool=pool) + +wall_p = CupyParameter("WALL", dtype=np.int32, mode="field", value=np.zeros(NN), pool=pool, n_flat=NN) +alpha_p = CupyParameter("ALPHA", dtype=np.float32, mode="field", value=np.zeros(NN), pool=pool, n_flat=NN) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- +clamp_fn = ( + CupyHelperBuilder() + .bind("N", n_p) + .ingest("__device__ int clampi(int i) { return i < 0 ? 0 : (i >= N ? N - 1 : i); }") +) + +laplacian_fn = ( + CupyHelperBuilder() + .bind("N", n_p) + .bind("clampi", clamp_fn) + .ingest( + r""" +__device__ float laplacian(const float* f, int i, int j) { + int ip = $clampi(i + 1)$; + int im = $clampi(i - 1)$; + int jp = $clampi(j + 1)$; + int jm = $clampi(j - 1)$; + return f[ip * N + j] + f[im * N + j] + f[i * N + jp] + f[i * N + jm] - 4.0f * f[i * N + j]; +} +""" + ) +) + +whash_fn = ( + CupyHelperBuilder() + .bind("SEED", seed_p) + .ingest( + r""" +__device__ float whash(int a, int b) { + float x = (float)a * 12.9898f + (float)b * 78.233f + SEED; + float s = sinf(x) * 43758.5453f; + return s - floorf(s); +} +""" + ) +) + +# --------------------------------------------------------------------------- +# one-shot setup kernels (run once, outside the routine, exactly as in the +# manual-loop example) +# --------------------------------------------------------------------------- +generate_walls_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("ROOM", room_p) + .bind("WALL_THICK", wall_thick_p) + .bind("DOOR", door_p) + .bind("wall", wall_p) + .bind("whash", whash_fn) + .ingest( + r""" +__global__ void generate_walls() { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N; + int j = idx % N; + int is_wall = 0; + if (i < WALL_THICK || i >= N - WALL_THICK || j < WALL_THICK || j >= N - WALL_THICK) { + is_wall = 1; + } else if (i % ROOM < WALL_THICK) { + int door = (int)($whash(i / ROOM, j / ROOM)$ * ROOM); + int r = j % ROOM; + if (!(r >= door && r < door + DOOR)) is_wall = 1; + } else if (j % ROOM < WALL_THICK) { + int door = (int)($whash(j / ROOM + 7919, i / ROOM)$ * ROOM); + int r = i % ROOM; + if (!(r >= door && r < door + DOOR)) is_wall = 1; + } + $wall.set_node(idx, is_wall)$; +} +""" + ) + .compile() +) + +alpha_seeds = Bag({"ALPHA_WALL_SEED": alpha_wall_seed_p, "ALPHA_AIR_SEED": alpha_air_seed_p}) + +set_alpha_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind_bag(alpha_seeds) + .bind("wall", wall_p) + .bind("alpha", alpha_p) + .ingest( + r""" +__global__ void set_alpha() { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + if ($wall.get(idx)$ == 1) { + $alpha.set_node(idx, ALPHA_WALL_SEED)$; + } else { + $alpha.set_node(idx, ALPHA_AIR_SEED)$; + } +} +""" + ) + .compile() +) + +init_temperature_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("T_BG", t_bg_p) + .ingest( + r""" +__global__ void init_temperature(float* T) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + T[idx] = T_BG; +} +""" + ) + .compile() +) + +stove = Bag( + { + "at": Bag({"i": src_i_p, "j": src_j_p}), + "r": src_r_p, + "temp": stove_p, + } +) + +# Kept as a builder, not just a compiled Kernel: compile() below seeds T0 +# once, standalone, and the very same builder is later handed to the +# routine's add_kernel() - compile() does not consume it (see compile.py). +apply_source_builder = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("stove", stove) + .ingest( + r""" +__global__ void apply_source(float* T) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int dx = idx / N - $stove.at.i.get(0)$; + int dy = idx % N - $stove.at.j.get(0)$; + if (dx * dx + dy * dy <= $stove.r.get(0)$ * $stove.r.get(0)$) { + T[idx] = $stove.temp.get(0)$; + } +} +""" + ) +) +apply_source_kernel = apply_source_builder.compile() + +heat = Bag({"alpha": alpha_p, "lap": laplacian_fn, "dt": dt_p, "dx2": dx2_p}) + +diffuse_builder = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("heat", heat) + .ingest( + r""" +__global__ void diffuse(float* T_out, const float* T_in) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N; + int j = idx % N; + float a = $heat.alpha.get(idx)$; + float lap = $heat.lap(T_in, i, j)$ / $heat.dx2.get(0)$; + T_out[idx] = T_in[idx] + $heat.dt.get(0)$ * a * lap; +} +""" + ) +) + +# --------------------------------------------------------------------------- +# fields (pooled - two flat buffers for ping-pong) +# --------------------------------------------------------------------------- +T0 = pool.get_data(np.float32, (NN,)) +T1 = pool.get_data(np.float32, (NN,)) + +generate_walls_kernel(grid=GRID, block=BLOCK) +set_alpha_kernel(grid=GRID, block=BLOCK) +init_temperature_kernel(T0.data, grid=GRID, block=BLOCK) +apply_source_kernel(T0.data, grid=GRID, block=BLOCK) + +# --------------------------------------------------------------------------- +# the routine: two unrolled substeps, T0/T1 swapped back to their starting +# roles by the end, so it can be called over and over. grid/block are set +# once on the builder and apply to both steps. +# --------------------------------------------------------------------------- +routine_bag = merge(diffuse_builder.as_bag(), apply_source_builder.as_bag()) + +diffusion_routine = ( + CupyRoutineBuilder(grid=GRID, block=BLOCK) + .add_data("T0", T0.data) + .add_data("T1", T1.data) + .bind_bag(routine_bag) + .begin_repeat(times=2) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .end_repeat() + .compile() +) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy().reshape(GRID_N, GRID_N), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall_p.get().to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (Cupy backend, Routine)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME // 2): + clock += 2.0 * PULSE_FREQ * DT_VAL + stove_p.set(OG_stove + 20.0 * math.sin(clock)) + + diffusion_routine() # two substeps, result lands back in T0 + sim_time += 2.0 * DT_VAL + + cp.cuda.Device().synchronize() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy().reshape(GRID_N, GRID_N)) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +for param in (stove_p, wall_p, alpha_p): + param.destroy() +pool.release_data(T0) +pool.release_data(T1) +print("pooled storage released") diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py new file mode 100644 index 0000000..e6b13cc --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py @@ -0,0 +1,318 @@ +""" +Same model and setup as heat_diffusion_quadrants.py, with the per-substep +ping-pong expressed as a Routine instead of a hand-written python loop. + +heat_diffusion_quadrants.py alternates `diffuse_kernel(T1, T0)`, +`apply_source_kernel(T1)`, then swaps the T0/T1 python names each iteration. +A Routine has no python between its steps to do that swap in, so the two +iterations that one swap-pair covers are unrolled into one routine with a +repeat block, add_swap standing in for the python-level `T0, T1 = T1, T0`: + + begin_repeat(times=2) + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + end_repeat() + +which records the body once and replays it twice, giving the same six-step +sequence as writing it out by hand: + + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + +Two swaps compose to the identity, which is exactly what compile() checks +for - so the compiled routine can be called over and over, each call +advancing the simulation by two substeps, and the result always ends up back +in the T0 buffer, matching two iterations of the manual loop. + +diffuse_builder and apply_source_builder are ordinary KernelBuilders, built +exactly as in heat_diffusion_quadrants.py; apply_source_builder is also +compiled once on its own to seed T0 before the loop starts, same as that file +does - compile() does not consume a builder, so the same builder is later +handed to add_kernel() unchanged. The routine's one shared bag is the merge +of what each builder already binds, so nothing about diffuse_template or +apply_source_template's own bodies changes. + +The stove's pulse is only updated between routine() calls, not between the +two substeps a single call unrolls: set() on the stove's scalar Parameter is +safe between calls (see routine.py, "Contract: no set()/destroy() +mid-routine"), but doing it *inside* a routine's steps is exactly what that +contract forbids, since there is no python between steps to run it in. With +PULSE_FREQ=0.0 by default the stove is steady anyway and this has no visible +effect. + +Author: B.G (07/2026) +""" + +import math +import time + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd + +from pyfastflow.experimental.core.context.bag import Bag, merge +from pyfastflow.experimental.core.context.quadrants_backend import ( + QuadrantsHelperBuilder, + QuadrantsKernelBuilder, + QuadrantsParameter, + QuadrantsRoutineBuilder, +) +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool + +qd.init(arch=qd.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, loop/timing counts - never used as kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 512 +STEPS_PER_FRAME = 10000 # two routine substeps per call - see the loop below +PULSE_FREQ = 0.0 # stove temperature oscillation speed, rad/s (0 = steady stove) + +ROOM_M = 3.0 +DX_M = ROOM_M / (GRID_N // 4) +ALPHA_AIR_VAL = 0.015 +ALPHA_WALL_VAL = 1.0e-6 +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) + +pool = QuadrantsPool() + +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool) +room_p = QuadrantsParameter("ROOM", dtype=qd.i32, mode="const", value=GRID_N // 4, pool=pool) +wall_thick_p = QuadrantsParameter("WALL_THICK", dtype=qd.i32, mode="const", value=8, pool=pool) +door_p = QuadrantsParameter("DOOR", dtype=qd.i32, mode="const", value=6, pool=pool) +seed_p = QuadrantsParameter("SEED", dtype=qd.f32, mode="const", value=17.0, pool=pool) + +dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool) +dx2_p = QuadrantsParameter("DX2", dtype=qd.f32, mode="const", value=DX_M**2, pool=pool) + +alpha_air_seed_p = QuadrantsParameter("ALPHA_AIR_SEED", dtype=qd.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool) +alpha_wall_seed_p = QuadrantsParameter("ALPHA_WALL_SEED", dtype=qd.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool) +t_bg_p = QuadrantsParameter("T_BG", dtype=qd.f32, mode="const", value=15.0, pool=pool) + +src_i_p = QuadrantsParameter("SRC_I", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_j_p = QuadrantsParameter("SRC_J", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_r_p = QuadrantsParameter("SRC_R", dtype=qd.i32, mode="const", value=10, pool=pool) + +OG_stove = 70 +stove_p = QuadrantsParameter("STOVE_T", dtype=qd.f32, mode="scalar", value=OG_stove, pool=pool) + +wall_p = QuadrantsParameter("WALL", dtype=qd.i32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) +alpha_p = QuadrantsParameter("ALPHA", dtype=qd.f32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +def clamp(i): + return min(max(i, 0), N.get(0) - 1) + + +clamp_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clamp) + + +def laplacian(field_, i, j): + ip = clamp(i + 1) + im = clamp(i - 1) + jp = clamp(j + 1) + jm = clamp(j - 1) + return field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + + +laplacian_fn = QuadrantsHelperBuilder().bind("clamp", clamp_fn).ingest(laplacian) + + +def whash(a, b): + """Deterministic pseudo-random value in [0, 1) for two integer indices.""" + x = qd.cast(a, qd.f32) * 12.9898 + qd.cast(b, qd.f32) * 78.233 + SEED.get(0) + s = qd.sin(x) * 43758.5453 + return s - qd.floor(s) + + +whash_fn = QuadrantsHelperBuilder().bind("SEED", seed_p).ingest(whash) + +# --------------------------------------------------------------------------- +# one-shot setup kernels (run once, outside the routine, exactly as in the +# manual-loop example) +# --------------------------------------------------------------------------- + + +def generate_walls_template(): + for i, j in qd.ndrange(N.get(0), N.get(0)): + is_wall = 0 + if i < WALL_THICK.get(0) or i >= N.get(0) - WALL_THICK.get(0) or j < WALL_THICK.get(0) or j >= N.get(0) - WALL_THICK.get(0): + is_wall = 1 + elif i % ROOM.get(0) < WALL_THICK.get(0): + vline = i // ROOM.get(0) + seg = j // ROOM.get(0) + door = qd.cast(whash(vline, seg) * ROOM.get(0), qd.i32) + gap = (j % ROOM.get(0)) >= door and (j % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + elif j % ROOM.get(0) < WALL_THICK.get(0): + hline = j // ROOM.get(0) + seg = i // ROOM.get(0) + door = qd.cast(whash(hline + 7919, seg) * ROOM.get(0), qd.i32) + gap = (i % ROOM.get(0)) >= door and (i % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + wall.set_node(i * N.get(0) + j, is_wall) + + +generate_walls_kernel = ( + QuadrantsKernelBuilder() + .bind("N", n_p) + .bind("ROOM", room_p) + .bind("WALL_THICK", wall_thick_p) + .bind("DOOR", door_p) + .bind("wall", wall_p) + .bind("whash", whash_fn) + .ingest(generate_walls_template) + .compile() +) + + +def set_alpha_template(): + for i, j in qd.ndrange(N.get(0), N.get(0)): + idx = i * N.get(0) + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) + else: + alpha.set_node(idx, ALPHA_AIR_SEED.get(0)) + + +alpha_seeds = Bag({"ALPHA_WALL_SEED": alpha_wall_seed_p, "ALPHA_AIR_SEED": alpha_air_seed_p}) + +set_alpha_kernel = ( + QuadrantsKernelBuilder() + .bind("N", n_p) + .bind("wall", wall_p) + .bind("alpha", alpha_p) + .bind_bag(alpha_seeds) + .ingest(set_alpha_template) + .compile() +) + + +def init_temperature_template(T: qd.Tensor): + for i, j in T: + T[i, j] = T_BG.get(0) + + +init_temperature_kernel = QuadrantsKernelBuilder().bind("T_BG", t_bg_p).ingest(init_temperature_template).compile() + +stove = Bag( + { + "at": Bag({"i": src_i_p, "j": src_j_p}), + "r": src_r_p, + "temp": stove_p, + } +) + + +def apply_source_template(T: qd.Tensor): + for i, j in T: + dx = i - stove.at.i.get(0) + dy = j - stove.at.j.get(0) + if dx * dx + dy * dy <= stove.r.get(0) * stove.r.get(0): + T[i, j] = stove.temp.get(0) + + +# Kept as a builder, not just a compiled Kernel: compile() below seeds T0 +# once, standalone, and the very same builder is later handed to the +# routine's add_kernel() - compile() does not consume it (see compile.py). +apply_source_builder = QuadrantsKernelBuilder().bind("stove", stove).ingest(apply_source_template) +apply_source_kernel = apply_source_builder.compile() + +heat = Bag({"alpha": alpha_p, "lap": laplacian_fn, "dt": dt_p, "dx2": dx2_p}) + + +def diffuse_template(T_out: qd.Tensor, T_in: qd.Tensor): + for i, j in T_in: + idx = i * N.get(0) + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2.get(0) + T_out[i, j] = T_in[i, j] + heat.dt.get(0) * a * lap + + +diffuse_builder = QuadrantsKernelBuilder().bind("N", n_p).bind("heat", heat).ingest(diffuse_template) + +# --------------------------------------------------------------------------- +# fields (pooled - two buffers for ping-pong) +# --------------------------------------------------------------------------- +T0 = pool.get_data(qd.f32, (GRID_N, GRID_N)) +T1 = pool.get_data(qd.f32, (GRID_N, GRID_N)) + +generate_walls_kernel() +set_alpha_kernel() +init_temperature_kernel(T0.data) +apply_source_kernel(T0.data) + +# --------------------------------------------------------------------------- +# the routine: two unrolled substeps, T0/T1 swapped back to their starting +# roles by the end, so it can be called over and over. +# --------------------------------------------------------------------------- +routine_bag = merge(diffuse_builder.as_bag(), apply_source_builder.as_bag()) + +diffusion_routine = ( + QuadrantsRoutineBuilder() + .add_data("T0", T0.data) + .add_data("T1", T1.data) + .bind_bag(routine_bag) + .begin_repeat(times=2) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .end_repeat() + .compile() +) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy(), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall_p.get().to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (Quadrants backend, Routine)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME // 2): + clock += 2.0 * PULSE_FREQ * DT_VAL + stove_p.set(OG_stove + 20.0 * math.sin(clock)) + + diffusion_routine() # two substeps, result lands back in T0 + sim_time += 2.0 * dt_p.get() + + qd.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +for param in (stove_p, wall_p, alpha_p): + param.destroy() +pool.release_data(T0) +pool.release_data(T1) +print("pooled storage released") diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py new file mode 100644 index 0000000..e88aa0b --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py @@ -0,0 +1,318 @@ +""" +Same model and setup as heat_diffusion_taichi.py, with the per-substep +ping-pong expressed as a Routine instead of a hand-written python loop. + +heat_diffusion_taichi.py alternates `diffuse_kernel(T1, T0)`, +`apply_source_kernel(T1)`, then swaps the T0/T1 python names each iteration. +A Routine has no python between its steps to do that swap in, so the two +iterations that one swap-pair covers are unrolled into one routine with a +repeat block, add_swap standing in for the python-level `T0, T1 = T1, T0`: + + begin_repeat(times=2) + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + end_repeat() + +which records the body once and replays it twice, giving the same six-step +sequence as writing it out by hand: + + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + diffuse(T1, T0); apply_source(T1); swap(T0, T1) + +Two swaps compose to the identity, which is exactly what compile() checks +for - so the compiled routine can be called over and over, each call +advancing the simulation by two substeps, and the result always ends up back +in the T0 buffer, matching two iterations of the manual loop. + +diffuse_builder and apply_source_builder are ordinary KernelBuilders, built +exactly as in heat_diffusion_taichi.py; apply_source_builder is also compiled +once on its own to seed T0 before the loop starts, same as that file does - +compile() does not consume a builder, so the same builder is later handed to +add_kernel() unchanged. The routine's one shared bag is the merge of what +each builder already binds, so nothing about diffuse_template or +apply_source_template's own bodies changes. + +The stove's pulse is only updated between routine() calls, not between the +two substeps a single call unrolls: set() on the stove's scalar Parameter is +safe between calls (see routine.py, "Contract: no set()/destroy() +mid-routine"), but doing it *inside* a routine's steps is exactly what that +contract forbids, since there is no python between steps to run it in. With +PULSE_FREQ=0.0 by default the stove is steady anyway and this has no visible +effect. + +Author: B.G (07/2026) +""" + +import math +import time + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti + +from pyfastflow.experimental.core.context.bag import Bag, merge +from pyfastflow.experimental.core.context.taichi_backend import ( + TaichiHelperBuilder, + TaichiKernelBuilder, + TaichiParameter, + TaichiRoutineBuilder, +) +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool + +ti.init(arch=ti.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, loop/timing counts - never used as kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 512 +STEPS_PER_FRAME = 10000 # two routine substeps per call - see the loop below +PULSE_FREQ = 0.0 # stove temperature oscillation speed, rad/s (0 = steady stove) + +ROOM_M = 3.0 +DX_M = ROOM_M / (GRID_N // 4) +ALPHA_AIR_VAL = 0.015 +ALPHA_WALL_VAL = 1.0e-6 +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) + +pool = TaichiPool() + +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool) +room_p = TaichiParameter("ROOM", dtype=ti.i32, mode="const", value=GRID_N // 4, pool=pool) +wall_thick_p = TaichiParameter("WALL_THICK", dtype=ti.i32, mode="const", value=8, pool=pool) +door_p = TaichiParameter("DOOR", dtype=ti.i32, mode="const", value=6, pool=pool) +seed_p = TaichiParameter("SEED", dtype=ti.f32, mode="const", value=17.0, pool=pool) + +dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool) +dx2_p = TaichiParameter("DX2", dtype=ti.f32, mode="const", value=DX_M**2, pool=pool) + +alpha_air_seed_p = TaichiParameter("ALPHA_AIR_SEED", dtype=ti.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool) +alpha_wall_seed_p = TaichiParameter("ALPHA_WALL_SEED", dtype=ti.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool) +t_bg_p = TaichiParameter("T_BG", dtype=ti.f32, mode="const", value=15.0, pool=pool) + +src_i_p = TaichiParameter("SRC_I", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_j_p = TaichiParameter("SRC_J", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_r_p = TaichiParameter("SRC_R", dtype=ti.i32, mode="const", value=10, pool=pool) + +OG_stove = 70 +stove_p = TaichiParameter("STOVE_T", dtype=ti.f32, mode="scalar", value=OG_stove, pool=pool) + +wall_p = TaichiParameter("WALL", dtype=ti.i32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) +alpha_p = TaichiParameter("ALPHA", dtype=ti.f32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +def clamp(i): + return min(max(i, 0), N.get(0) - 1) + + +clamp_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clamp) + + +def laplacian(field_, i, j): + ip = clamp(i + 1) + im = clamp(i - 1) + jp = clamp(j + 1) + jm = clamp(j - 1) + return field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + + +laplacian_fn = TaichiHelperBuilder().bind("clamp", clamp_fn).ingest(laplacian) + + +def whash(a, b): + """Deterministic pseudo-random value in [0, 1) for two integer indices.""" + x = ti.cast(a, ti.f32) * 12.9898 + ti.cast(b, ti.f32) * 78.233 + SEED.get(0) + s = ti.sin(x) * 43758.5453 + return s - ti.floor(s) + + +whash_fn = TaichiHelperBuilder().bind("SEED", seed_p).ingest(whash) + +# --------------------------------------------------------------------------- +# one-shot setup kernels (run once, outside the routine, exactly as in the +# manual-loop example) +# --------------------------------------------------------------------------- + + +def generate_walls_template(): + for i, j in ti.ndrange(N.get(0), N.get(0)): + is_wall = 0 + if i < WALL_THICK.get(0) or i >= N.get(0) - WALL_THICK.get(0) or j < WALL_THICK.get(0) or j >= N.get(0) - WALL_THICK.get(0): + is_wall = 1 + elif i % ROOM.get(0) < WALL_THICK.get(0): + vline = i // ROOM.get(0) + seg = j // ROOM.get(0) + door = ti.cast(whash(vline, seg) * ROOM.get(0), ti.i32) + gap = (j % ROOM.get(0)) >= door and (j % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + elif j % ROOM.get(0) < WALL_THICK.get(0): + hline = j // ROOM.get(0) + seg = i // ROOM.get(0) + door = ti.cast(whash(hline + 7919, seg) * ROOM.get(0), ti.i32) + gap = (i % ROOM.get(0)) >= door and (i % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + wall.set_node(i * N.get(0) + j, is_wall) + + +generate_walls_kernel = ( + TaichiKernelBuilder() + .bind("N", n_p) + .bind("ROOM", room_p) + .bind("WALL_THICK", wall_thick_p) + .bind("DOOR", door_p) + .bind("wall", wall_p) + .bind("whash", whash_fn) + .ingest(generate_walls_template) + .compile() +) + + +def set_alpha_template(): + for i, j in ti.ndrange(N.get(0), N.get(0)): + idx = i * N.get(0) + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) + else: + alpha.set_node(idx, ALPHA_AIR_SEED.get(0)) + + +alpha_seeds = Bag({"ALPHA_WALL_SEED": alpha_wall_seed_p, "ALPHA_AIR_SEED": alpha_air_seed_p}) + +set_alpha_kernel = ( + TaichiKernelBuilder() + .bind("N", n_p) + .bind("wall", wall_p) + .bind("alpha", alpha_p) + .bind_bag(alpha_seeds) + .ingest(set_alpha_template) + .compile() +) + + +def init_temperature_template(T: ti.template()): + for i, j in T: + T[i, j] = T_BG.get(0) + + +init_temperature_kernel = TaichiKernelBuilder().bind("T_BG", t_bg_p).ingest(init_temperature_template).compile() + +stove = Bag( + { + "at": Bag({"i": src_i_p, "j": src_j_p}), + "r": src_r_p, + "temp": stove_p, + } +) + + +def apply_source_template(T: ti.template()): + for i, j in T: + dx = i - stove.at.i.get(0) + dy = j - stove.at.j.get(0) + if dx * dx + dy * dy <= stove.r.get(0) * stove.r.get(0): + T[i, j] = stove.temp.get(0) + + +# Kept as a builder, not just a compiled Kernel: compile() below seeds T0 +# once, standalone, and the very same builder is later handed to the +# routine's add_kernel() - compile() does not consume it (see compile.py). +apply_source_builder = TaichiKernelBuilder().bind("stove", stove).ingest(apply_source_template) +apply_source_kernel = apply_source_builder.compile() + +heat = Bag({"alpha": alpha_p, "lap": laplacian_fn, "dt": dt_p, "dx2": dx2_p}) + + +def diffuse_template(T_out: ti.template(), T_in: ti.template()): + for i, j in T_in: + idx = i * N.get(0) + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2.get(0) + T_out[i, j] = T_in[i, j] + heat.dt.get(0) * a * lap + + +diffuse_builder = TaichiKernelBuilder().bind("N", n_p).bind("heat", heat).ingest(diffuse_template) + +# --------------------------------------------------------------------------- +# fields (pooled - two buffers for ping-pong) +# --------------------------------------------------------------------------- +T0 = pool.get_data(ti.f32, (GRID_N, GRID_N)) +T1 = pool.get_data(ti.f32, (GRID_N, GRID_N)) + +generate_walls_kernel() +set_alpha_kernel() +init_temperature_kernel(T0.data) +apply_source_kernel(T0.data) + +# --------------------------------------------------------------------------- +# the routine: two unrolled substeps, T0/T1 swapped back to their starting +# roles by the end, so it can be called over and over. +# --------------------------------------------------------------------------- +routine_bag = merge(diffuse_builder.as_bag(), apply_source_builder.as_bag()) + +diffusion_routine = ( + TaichiRoutineBuilder() + .add_data("T0", T0.data) + .add_data("T1", T1.data) + .bind_bag(routine_bag) + .begin_repeat(times=2) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .end_repeat() + .compile() +) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy(), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall_p.get().to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (Taichi backend, Routine)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME // 2): + clock += 2.0 * PULSE_FREQ * DT_VAL + stove_p.set(OG_stove + 20.0 * math.sin(clock)) + + diffusion_routine() # two substeps, result lands back in T0 + sim_time += 2.0 * dt_p.get() + + ti.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +for param in (stove_p, wall_p, alpha_p): + param.destroy() +pool.release_data(T0) +pool.release_data(T1) +print("pooled storage released") diff --git a/examples/core/heat_diffusion/heat_diffusion_taichi.py b/examples/core/heat_diffusion/heat_diffusion_taichi.py new file mode 100644 index 0000000..4f24567 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_taichi.py @@ -0,0 +1,336 @@ +""" +Heat diffusion through a procedurally-generated floor plan (air + walls), +heated by a single stove, built on pyfastflow's backend-agnostic core +(Parameter/Helper/Kernel/Pool), Taichi backend. + +Pipeline: + - generate_walls: pointwise kernel, carves a grid of rooms with doors into + a wall mask (field-mode Parameter `wall`, written device-side via + wall.set_node), using a deterministic hash device helper instead of + per-cell RNG so wall/door layout is reproducible from SEED. + - set_alpha: seeds a per-cell diffusivity field (`alpha`, field mode) from + `wall` - air diffuses fast, walls slow. + - init_temperature: fills T with the background temperature. + - apply_source: clamps a disc of cells around the stove to `stove.temp` + (scalar-mode Parameter, updated from the host every substep -> a gently + pulsing stove). + - diffuse: explicit FTCS heat equation dT/dt = alpha(i,j) * lap(T), with a + clamped (Neumann / no-flux) boundary Laplacian. + +Uniform device surface: every Parameter is read with p.get(node) and written +with p.set_node(node, val) regardless of const/scalar/field mode - the +kernels never branch on mode, so re-declaring `alpha` as a single const +(uniform room, no walls) needs no kernel-body change. + +Binding styles, all three visible in one file: + - flat, one bind() per object (most kernels here); + - a Bag bound whole and reached by dotted path - `stove` in apply_source, + which nests a sub-bag for the position, and `heat` in diffuse, which mixes + a Parameter, a device helper and two consts under one name; + - bind_bag(), merging a bag's members in flat under their own names, so the + kernel still sees plain names - `alpha_seeds` in set_alpha. + +Compilation is the two-layer builder: TaichiKernelBuilder / TaichiHelperBuilder +collect bind()ed params + helper builders and one ingest()ed template. A +HelperBuilder (clamp_fn, laplacian_fn, whash_fn below) is a recipe, not a +compiled object - it has no compile() of its own. Binding one into a kernel, +flat or through a Bag, is what specializes it: TaichiKernelBuilder.compile() +specializes every HelperBuilder the kernel reaches, against that kernel's own +bindings, before compiling the kernel body. Recompiling the kernel after +rebinding a const the helper reads picks up the new value without touching +the helper builder itself. + +Author: B.G (07/2026) +""" + +import math +import time + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti + +from pyfastflow.experimental.core.context.bag import Bag +from pyfastflow.experimental.core.context.taichi_backend import ( + TaichiHelperBuilder, + TaichiKernelBuilder, + TaichiParameter, +) +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool + +ti.init(arch=ti.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, loop/timing counts - never used as kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 512 +STEPS_PER_FRAME = 10000 +PULSE_FREQ = 0.0 # stove temperature oscillation speed, rad/s (0 = steady stove) + +# Physical grounding: without a cell size, DT/ALPHA are just numbers tuned by +# feel - here they're derived from a real room size and real diffusivities so +# "seconds" and "m^2/s" mean what they say. +ROOM_M = 3.0 # room span, meters (rooms are GRID_N//4 cells across) +DX_M = ROOM_M / (GRID_N // 4) # meters per cell + +# Air's real molecular thermal diffusivity (~2.2e-5 m^2/s) would take DAYS to +# spread heat by pure conduction - rooms actually heat by convective mixing. +# ALPHA_AIR below is an effective/turbulent diffusivity standing in for that +# mixing, not molecular diffusion - otherwise a stove would need real hours. +ALPHA_AIR_VAL = 0.015 # m^2/s, effective convective air diffusivity +ALPHA_WALL_VAL = 1.0e-6 # m^2/s, real solid (drywall/brick-like) diffusivity + +# Explicit FTCS stability limit is dt <= dx^2 / (4*alpha); stay well under it. +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * ALPHA_AIR_VAL) # seconds + +pool = TaichiPool() + +# Structural constants: const mode, bake to compile-time literals in generated +# code even though the kernel body still reads them via .get(0). +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool) +room_p = TaichiParameter("ROOM", dtype=ti.i32, mode="const", value=GRID_N // 4, pool=pool) +wall_thick_p = TaichiParameter("WALL_THICK", dtype=ti.i32, mode="const", value=8, pool=pool) +door_p = TaichiParameter("DOOR", dtype=ti.i32, mode="const", value=6, pool=pool) +seed_p = TaichiParameter("SEED", dtype=ti.f32, mode="const", value=17.0, pool=pool) + +dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool) # seconds +dx2_p = TaichiParameter("DX2", dtype=ti.f32, mode="const", value=DX_M**2, pool=pool) # meters^2 + +# Seed values for the alpha field - read via .get(0) inside set_alpha. +alpha_air_seed_p = TaichiParameter("ALPHA_AIR_SEED", dtype=ti.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool) +alpha_wall_seed_p = TaichiParameter("ALPHA_WALL_SEED", dtype=ti.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool) +t_bg_p = TaichiParameter("T_BG", dtype=ti.f32, mode="const", value=15.0, pool=pool) + +src_i_p = TaichiParameter("SRC_I", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_j_p = TaichiParameter("SRC_J", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool) +src_r_p = TaichiParameter("SRC_R", dtype=ti.i32, mode="const", value=10, pool=pool) # stove radius, cells + +# scalar mode: a 0-d field, host-settable every frame -> a pulsing stove +# temperature. Reached in-kernel as stove.temp.get(0) (see the stove Bag). +OG_stove = 70 +stove_p = TaichiParameter("STOVE_T", dtype=ti.f32, mode="scalar", value=OG_stove, pool=pool) + +# field mode: per-cell wall/air mask, written device-side via wall.set_node, +# read via wall.get. +wall_p = TaichiParameter("WALL", dtype=ti.i32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) + +# field mode: per-cell thermal diffusivity, read in diffuse via alpha.get - so +# switching this Parameter to const/scalar mode later needs no kernel edits. +alpha_p = TaichiParameter("ALPHA", dtype=ti.f32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +def clamp(i): + return min(max(i, 0), N.get(0) - 1) + + +clamp_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clamp) + + +def laplacian(field_, i, j): + ip = clamp(i + 1) + im = clamp(i - 1) + jp = clamp(j + 1) + jm = clamp(j - 1) + return field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + + +laplacian_fn = TaichiHelperBuilder().bind("clamp", clamp_fn).ingest(laplacian) + + +def whash(a, b): + """Deterministic pseudo-random value in [0, 1) for two integer indices.""" + x = ti.cast(a, ti.f32) * 12.9898 + ti.cast(b, ti.f32) * 78.233 + SEED.get(0) + s = ti.sin(x) * 43758.5453 + return s - ti.floor(s) + + +whash_fn = TaichiHelperBuilder().bind("SEED", seed_p).ingest(whash) + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- + + +def generate_walls_template(): + for i, j in ti.ndrange(N.get(0), N.get(0)): + is_wall = 0 + if i < WALL_THICK.get(0) or i >= N.get(0) - WALL_THICK.get(0) or j < WALL_THICK.get(0) or j >= N.get(0) - WALL_THICK.get(0): + is_wall = 1 + elif i % ROOM.get(0) < WALL_THICK.get(0): + vline = i // ROOM.get(0) + seg = j // ROOM.get(0) + door = ti.cast(whash(vline, seg) * ROOM.get(0), ti.i32) + gap = (j % ROOM.get(0)) >= door and (j % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + elif j % ROOM.get(0) < WALL_THICK.get(0): + hline = j // ROOM.get(0) + seg = i // ROOM.get(0) + door = ti.cast(whash(hline + 7919, seg) * ROOM.get(0), ti.i32) + gap = (i % ROOM.get(0)) >= door and (i % ROOM.get(0)) < door + DOOR.get(0) + if not gap: + is_wall = 1 + wall.set_node(i * N.get(0) + j, is_wall) + + +generate_walls_kernel = ( + TaichiKernelBuilder() + .bind("N", n_p) + .bind("ROOM", room_p) + .bind("WALL_THICK", wall_thick_p) + .bind("DOOR", door_p) + .bind("wall", wall_p) + .bind("whash", whash_fn) + .ingest(generate_walls_template) + .compile() +) + + +def set_alpha_template(): + for i, j in ti.ndrange(N.get(0), N.get(0)): + idx = i * N.get(0) + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) + else: + alpha.set_node(idx, ALPHA_AIR_SEED.get(0)) + + +# The two seed values are grouped on the host for tidiness, then merged in with +# bind_bag() - which binds each member flat, under its own name. The template +# above is unaware: it still reads ALPHA_WALL_SEED / ALPHA_AIR_SEED bare. Use +# this when a bag is a convenient way to carry things around but the kernel +# wants plain names; bind() the bag whole instead when you want a dotted path. +alpha_seeds = Bag({"ALPHA_WALL_SEED": alpha_wall_seed_p, "ALPHA_AIR_SEED": alpha_air_seed_p}) + +set_alpha_kernel = ( + TaichiKernelBuilder() + .bind("N", n_p) + .bind("wall", wall_p) + .bind("alpha", alpha_p) + .bind_bag(alpha_seeds) + .ingest(set_alpha_template) + .compile() +) + + +def init_temperature_template(T: ti.template()): + for i, j in T: + T[i, j] = T_BG.get(0) + + +init_temperature_kernel = TaichiKernelBuilder().bind("T_BG", t_bg_p).ingest(init_temperature_template).compile() + + +# The stove travels as ONE nested Bag rather than four flat binds: its position +# is grouped into an `at` sub-bag. Every member, whatever mode, is reached the +# same way - .get(0) - so const and scalar Parameters sit side by side under +# one name. Everything else here still binds flat, so the two styles sit side +# by side in one file. +stove = Bag( + { + "at": Bag({"i": src_i_p, "j": src_j_p}), + "r": src_r_p, + "temp": stove_p, + } +) + + +def apply_source_template(T: ti.template()): + for i, j in T: + dx = i - stove.at.i.get(0) + dy = j - stove.at.j.get(0) + if dx * dx + dy * dy <= stove.r.get(0) * stove.r.get(0): + T[i, j] = stove.temp.get(0) + + +apply_source_kernel = TaichiKernelBuilder().bind("stove", stove).ingest(apply_source_template).compile() + + +# A MIXED Bag: everything the diffusion step needs, whatever kind it is - a +# field Parameter, a device helper, two const Parameters - under one name. A +# bag has no member type; each is resolved on its own at compile time, so +# `heat.alpha` becomes a device accessor, `heat.lap` a compiled func, and +# `heat.dx2` a device accessor whose .get(0) bakes to a literal. +heat = Bag({"alpha": alpha_p, "lap": laplacian_fn, "dt": dt_p, "dx2": dx2_p}) + + +def diffuse_template(T_out: ti.template(), T_in: ti.template()): + for i, j in T_in: + idx = i * N.get(0) + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2.get(0) + T_out[i, j] = T_in[i, j] + heat.dt.get(0) * a * lap + +diffuse_kernel = TaichiKernelBuilder().bind("N", n_p).bind("heat", heat).ingest(diffuse_template).compile() + +# --------------------------------------------------------------------------- +# fields (pooled - two buffers for ping-pong) +# --------------------------------------------------------------------------- +T0 = pool.get_data(ti.f32, (GRID_N, GRID_N)) +T1 = pool.get_data(ti.f32, (GRID_N, GRID_N)) + +generate_walls_kernel() +set_alpha_kernel() +init_temperature_kernel(T0.data) +apply_source_kernel(T0.data) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(T0.to_numpy(), cmap="inferno", vmin=20.0, vmax=OG_stove) +fig.colorbar(im, ax=ax, label="Temperature (deg C)") + +wall_mask = wall_p.get().to_numpy().reshape(GRID_N, GRID_N) +wall_overlay = np.where(wall_mask == 1, 1.0, np.nan) +ax.imshow(wall_overlay, cmap="gray", vmin=0.0, vmax=1.0, alpha=0.35) + +ax.set_title("Heat diffusion in a floor plan (Taichi backend)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +clock = 0.0 +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + clock += PULSE_FREQ * DT_VAL + stove_p.set(OG_stove + 20.0 * math.sin(clock)) + + diffuse_kernel(T1.data, T0.data) + apply_source_kernel(T1.data) + T0, T1 = T1, T0 + sim_time += dt_p.get() + + ti.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.0f} s") + im.set_data(T0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +# destroy() hands a Parameter's storage back to the pool; it is a no-op on a +# const, which owns none. Safe only because nothing will launch again - the +# pool may reissue these buffers, while the compiled kernels above still point +# at them (see parameter.py, "Lifetime of a compiled object"). +for param in (stove_p, wall_p, alpha_p): + param.destroy() +pool.release_data(T0) +pool.release_data(T1) +print("pooled storage released") diff --git a/examples/core/lem/lem_routine_cupy.py b/examples/core/lem/lem_routine_cupy.py new file mode 100644 index 0000000..cf5b38c --- /dev/null +++ b/examples/core/lem/lem_routine_cupy.py @@ -0,0 +1,305 @@ +""" +Hillslope landscape evolution as one Routine, exercising every part of the +core in a single model. Cupy backend; same model as lem_routine_taichi.py. + +The physics is deliberately small - linear hillslope diffusion against a +spatially variable uplift field, with the domain edges pinned to base level: + + dz/dt = D * laplacian(z) + U(x, y) + +What the file is here to show is how the pieces fit together when a model +needs all of them at once. The other examples each isolate one thing; this +one carries the lot: + + Parameter modes N and DT are const Parameters bound flat, arriving as + #defines and read bare in the source. DX is a const bound + through a Bag, read as $grid.dx.get(0)$ - a const reached + through a span, rather than a top-level #define, always + goes through .get(0). D and SEA_LEVEL are scalars the + host retunes between frames. UPLIFT is a field, one rate + per node. + Helpers clampi binds a const; laplacian binds a bag and calls + clampi, so a helper reaches another helper; uplift_at + binds the UPLIFT *field* directly, which is what lets the + uplift kernel body stay a one-liner. Every scalar/field + parameter these reach lands in the module's __constant__ + block, so a helper reads one exactly as its caller does. + Bags grid is nested (grid.n, grid.dx), hill is mixed - a + scalar Parameter, a helper and a const under one name - + and the two noise seeds arrive flat through bind_bag. + Routine three kernels, two of them inside the routine, with the + z0/z1 ping-pong unrolled twice so the swaps compose to + the identity and the routine can be called repeatedly. + +Buffers are flat here rather than 2D, since a CUDA template indexes its own +data: a kernel takes one thread per node and recovers (i, j) itself. + +The step the routine runs is diffuse then uplift-and-clamp, so uplift is +applied to what diffusion just wrote. Two of those, plus the two swaps, make +one routine call - and the result always lands back in z0. + +The routine is captured into a CUDA graph (CupyRoutineBuilder.compile's +default), so a call replays recorded launches rather than re-issuing them. +D and SEA_LEVEL are retuned between calls, never between the steps inside +one: a write to a scalar Parameter goes through the same storage the graph +holds, so replay sees it, while there is no python between a routine's own +steps to run it in anyway (see routine.py, "Contract: no set()/destroy() +mid-routine"). + +Author: B.G (07/2026) +""" + +import time + +import cupy as cp +import matplotlib.pyplot as plt +import numpy as np + +from pyfastflow.experimental.core.context.bag import Bag, merge +from pyfastflow.experimental.core.context.cupy_backend import ( + CupyHelperBuilder, + CupyKernelBuilder, + CupyParameter, + CupyRoutineBuilder, +) +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool + +# --------------------------------------------------------------------------- +# host-side constants (grid size, launch config, timing) +# --------------------------------------------------------------------------- +GRID_N = 2048 +NN = GRID_N * GRID_N +DX_M = 100.0 +STEPS_PER_FRAME = 200 # two routine substeps per call - see the loop below + +BLOCK = 256 +GRID = (NN + BLOCK - 1) // BLOCK + +D_VAL = 1.0e-2 # hillslope diffusivity, m2/yr +UPLIFT_MAX = 1.0e-6 # m/yr at the range crest +CFL_SAFETY = 0.2 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * D_VAL) + +pool = CupyPool() + +# --------------------------------------------------------------------------- +# parameters - one of every mode +# --------------------------------------------------------------------------- +# const Parameters, bound flat at top level: emitted as #defines and read +# bare in the source. +n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool) +dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool) +seed_a_p = CupyParameter("SEED_A", dtype=np.float32, mode="const", value=12.9898, pool=pool) +seed_b_p = CupyParameter("SEED_B", dtype=np.float32, mode="const", value=78.233, pool=pool) + +# fixed at compile time like the ones above, but reached through a span +# (bound inside a Bag) rather than as a top-level #define +dx_p = CupyParameter("DX", dtype=np.float32, mode="const", value=DX_M, pool=pool) + +# scalars: one cell each, retuned from the host between routine calls +d_p = CupyParameter("D", dtype=np.float32, mode="scalar", value=D_VAL, pool=pool) +sea_p = CupyParameter("SEA_LEVEL", dtype=np.float32, mode="scalar", value=0.0, pool=pool) + +# field: one value per node, filled from the host below +uplift_p = CupyParameter("UPLIFT", dtype=np.float32, mode="field", value=np.zeros(NN), pool=pool, n_flat=NN) + +# a north-south uplift ridge, tapering to zero at the north and south edges +_yy = np.arange(GRID_N, dtype=np.float32)[:, None] * np.ones((1, GRID_N), np.float32) +_ridge = np.sin(np.pi * _yy / (GRID_N - 1)) ** 2 +uplift_p.set((UPLIFT_MAX * _ridge).ravel()) + +# --------------------------------------------------------------------------- +# bags +# --------------------------------------------------------------------------- +# nested: both grid.n and grid.dx are const Parameters, reached through a +# span (.get(0)) rather than a top-level #define - members resolve on their +# own type, not the bag's +grid = Bag({"n": n_p, "dx": dx_p}) + +# flat, for bind_bag: the kernel that uses these reads them as bare names +noise_seeds = Bag({"SEED_A": seed_a_p, "SEED_B": seed_b_p}) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- +clampi_fn = ( + CupyHelperBuilder() + .bind("N", n_p) + .ingest("__device__ int clampi(int i) { return i < 0 ? 0 : (i >= N ? N - 1 : i); }") +) + +laplacian_fn = ( + CupyHelperBuilder() + .bind("clampi", clampi_fn) + .bind("N", n_p) + .bind("grid", grid) + .ingest( + r""" +__device__ float laplacian(const float* f, int i, int j) { + // calls another helper, and reads a const Parameter out of a bound bag + int ip = $clampi(i + 1)$; + int im = $clampi(i - 1)$; + int jp = $clampi(j + 1)$; + int jm = $clampi(j - 1)$; + float acc = f[ip * N + j] + f[im * N + j] + f[i * N + jp] + f[i * N + jm] - 4.0f * f[i * N + j]; + float dx = $grid.dx.get(0)$; + return acc / (dx * dx); +} +""" + ) +) + +uplift_at_fn = ( + CupyHelperBuilder() + .bind("UPLIFT", uplift_p) + .ingest( + r""" +__device__ float uplift_at(int idx) { + // binds the UPLIFT *field* itself: a helper reads a non-const Parameter + // exactly the way a kernel does, so the caller passes only the index + return $UPLIFT.get(idx)$; +} +""" + ) +) + +# --------------------------------------------------------------------------- +# one-shot setup kernel (runs once, outside the routine) +# --------------------------------------------------------------------------- +init_topo_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind_bag(noise_seeds) + .ingest( + r""" +extern "C" __global__ void init_topo(float* z) { + // bind_bag put SEED_A / SEED_B in flat, so they read as bare names here + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N, j = idx % N; + float x = (float)i * SEED_A + (float)j * SEED_B; + float s = sinf(x) * 43758.5453f; + z[idx] = (s - floorf(s)) * 2.0f; +} +""" + ) + .compile() +) + +# --------------------------------------------------------------------------- +# routine kernels +# --------------------------------------------------------------------------- + +# mixed bag: a scalar Parameter, a helper and a const Parameter under one +# name. hill.d, hill.dt are span reads (.get(0)), hill.lap a spliced call. +hill = Bag({"d": d_p, "lap": laplacian_fn, "dt": dt_p}) + +diffuse_builder = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("hill", hill) + .ingest( + r""" +extern "C" __global__ void diffuse(float* z_out, const float* z_in) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N, j = idx % N; + z_out[idx] = z_in[idx] + $hill.dt.get(0)$ * $hill.d.get(0)$ * $hill.lap(z_in, i, j)$; +} +""" + ) +) + +uplift_builder = ( + CupyKernelBuilder() + .bind("up", uplift_at_fn) + .bind("DT", dt_p) + .bind("grid", grid) + .bind("SEA", sea_p) + .ingest( + r""" +extern "C" __global__ void uplift_bc(float* z) { + // grid.n is a const Parameter reached through the bag, so it splices in + // as a device accessor - .get(0) bakes to a literal just as a top-level + // #define would + int n = $grid.n.get(0)$; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n * n) return; + int j = idx % n; + z[idx] += $up(idx)$ * DT; + // base level: pin the east and west edges, so the ridge drains outward + if (j == 0 || j == n - 1) z[idx] = $SEA.get(0)$; +} +""" + ) +) + +# --------------------------------------------------------------------------- +# buffers (pooled - two for ping-pong) +# --------------------------------------------------------------------------- +z0 = pool.get_data(np.float32, (NN,)) +z1 = pool.get_data(np.float32, (NN,)) + +init_topo_kernel(z0.data, grid=GRID, block=BLOCK) + +# --------------------------------------------------------------------------- +# the routine +# --------------------------------------------------------------------------- +# One bag for the whole routine, merged from what each builder already binds. +# Both builders reach `grid` and `N` - the same objects, so the same uids, +# which is what lets merge() accept the collision instead of raising on it. +# grid/block are set once here and inherited by every step. +evolve = ( + CupyRoutineBuilder(grid=GRID, block=BLOCK) + .bind_bag(merge(diffuse_builder.as_bag(), uplift_builder.as_bag())) + .add_data("z0", z0.data) + .add_data("z1", z1.data) + .add_kernel(diffuse_builder, data_handle_ref=("z1", "z0")) + .add_kernel(uplift_builder, data_handle_ref=("z1",)) + .add_swap("z0", "z1") + .add_kernel(diffuse_builder, data_handle_ref=("z1", "z0")) + .add_kernel(uplift_builder, data_handle_ref=("z1",)) + .add_swap("z0", "z1") + .compile() +) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(z0.to_numpy().reshape(GRID_N, GRID_N), cmap="terrain", vmin=0.0, vmax=150.0) +fig.colorbar(im, ax=ax, label="Elevation (m)") +ax.set_title("Hillslope LEM (Cupy backend, Routine)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME // 2): + evolve() # two substeps, result lands back in z0 + sim_time += 2.0 * DT_VAL + + cp.cuda.Device().synchronize() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time / 1e6:.2f} Myr") + im.set_data(z0.to_numpy().reshape(GRID_N, GRID_N)) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +d_p.destroy() +sea_p.destroy() +uplift_p.destroy() +pool.release_data(z0) +pool.release_data(z1) diff --git a/examples/core/lem/lem_routine_quadrants.py b/examples/core/lem/lem_routine_quadrants.py new file mode 100644 index 0000000..906df3b --- /dev/null +++ b/examples/core/lem/lem_routine_quadrants.py @@ -0,0 +1,260 @@ +""" +Hillslope landscape evolution as one Routine, exercising every part of the +core in a single model. Quadrants backend; same model as lem_routine_taichi.py. + +The physics is deliberately small - linear hillslope diffusion against a +spatially variable uplift field, with the domain edges pinned to base level: + + dz/dt = D * laplacian(z) + U(x, y) + +What the file is here to show is how the pieces fit together when a model +needs all of them at once. The other examples each isolate one thing; this +one carries the lot: + + Parameter modes N, DT and DX are const Parameters, read uniformly via + .get(0) - the value still bakes to a compile-time literal + in generated code. D and SEA_LEVEL are scalars the host + retunes between frames. UPLIFT is a field, one rate per + node. + Helpers clampi binds a const; laplacian binds a bag and calls + clampi, so a helper reaches another helper; uplift_at + binds the UPLIFT *field* directly, which is what lets the + uplift kernel body stay a one-liner. + Bags grid is nested (grid.n, grid.dx), hill is mixed - a + scalar Parameter, a helper and a const under one name - + and the two noise seeds arrive flat through bind_bag. + Routine three kernels, two of them inside the routine, with the + z0/z1 ping-pong unrolled twice so the swaps compose to + the identity and the routine can be called repeatedly. + +The step the routine runs is diffuse then uplift-and-clamp, so uplift is +applied to what diffusion just wrote. Two of those, plus the two swaps, make +one routine call - and the result always lands back in z0. + +D and SEA_LEVEL are retuned between routine calls, never between the steps +inside one: set() on a scalar Parameter is what a routine expects between +calls (see routine.py, "Contract: no set()/destroy() mid-routine"), and there +is no python between a routine's own steps to run it in anyway. + +Author: B.G (07/2026) +""" + +import time + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd + +from pyfastflow.experimental.core.context.bag import Bag, merge +from pyfastflow.experimental.core.context.quadrants_backend import ( + QuadrantsHelperBuilder, + QuadrantsKernelBuilder, + QuadrantsParameter, + QuadrantsRoutineBuilder, +) +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool + +qd.init(arch=qd.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, timing - never used as kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 2048 +DX_M = 100.0 +STEPS_PER_FRAME = 200 # two routine substeps per call - see the loop below + +D_VAL = 1.0e-2 # hillslope diffusivity, m2/yr +UPLIFT_MAX = 1.0e-6 # m/yr at the range crest +CFL_SAFETY = 0.2 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * D_VAL) + +pool = QuadrantsPool() + +# --------------------------------------------------------------------------- +# parameters - one of every mode +# --------------------------------------------------------------------------- +# const Parameters: fixed at compile time and folded into the generated code +# as a literal, but read through .get(0) like any other mode, so a template +# can be written without knowing a given name is const. +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool) +dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool) +seed_a_p = QuadrantsParameter("SEED_A", dtype=qd.f32, mode="const", value=12.9898, pool=pool) +seed_b_p = QuadrantsParameter("SEED_B", dtype=qd.f32, mode="const", value=78.233, pool=pool) + +dx_p = QuadrantsParameter("DX", dtype=qd.f32, mode="const", value=DX_M, pool=pool) + +# scalars: one cell each, retuned from the host between routine calls +d_p = QuadrantsParameter("D", dtype=qd.f32, mode="scalar", value=D_VAL, pool=pool) +sea_p = QuadrantsParameter("SEA_LEVEL", dtype=qd.f32, mode="scalar", value=0.0, pool=pool) + +# field: one value per node, filled from the host below +uplift_p = QuadrantsParameter( + "UPLIFT", dtype=qd.f32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N +) + +# a north-south uplift ridge, tapering to zero at the north and south edges +_yy = np.arange(GRID_N, dtype=np.float32)[:, None] * np.ones((1, GRID_N), np.float32) +_ridge = np.sin(np.pi * _yy / (GRID_N - 1)) ** 2 +uplift_p.set((UPLIFT_MAX * _ridge).ravel()) + +# --------------------------------------------------------------------------- +# bags +# --------------------------------------------------------------------------- +# nested: both grid.n and grid.dx are const Parameters, read through .get(0) - +# members resolve on their own type, not the bag's +grid = Bag({"n": n_p, "dx": dx_p}) + +# flat, for bind_bag: the kernel that uses these reads them as bare names +noise_seeds = Bag({"SEED_A": seed_a_p, "SEED_B": seed_b_p}) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +def clampi(i): + return min(max(i, 0), N.get(0) - 1) + + +clampi_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clampi) + + +def laplacian(field_, i, j): + # calls another helper, and reads a const Parameter out of a bound bag + ip = clampi(i + 1) + im = clampi(i - 1) + jp = clampi(j + 1) + jm = clampi(j - 1) + acc = field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + return acc / (grid.dx.get(0) * grid.dx.get(0)) + + +laplacian_fn = QuadrantsHelperBuilder().bind("clampi", clampi_fn).bind("grid", grid).ingest(laplacian) + + +def uplift_at(i, j): + # binds the UPLIFT *field* itself: a helper reads a non-const Parameter + # exactly the way a kernel does, so the caller passes only the indices + return UPLIFT.get(i * N.get(0) + j) + + +uplift_at_fn = QuadrantsHelperBuilder().bind("UPLIFT", uplift_p).bind("N", n_p).ingest(uplift_at) + +# --------------------------------------------------------------------------- +# one-shot setup kernel (runs once, outside the routine) +# --------------------------------------------------------------------------- + + +def init_topo_template(z: qd.Tensor): + # bind_bag put SEED_A / SEED_B in flat, so they read as bare names here + for i, j in z: + x = qd.cast(i, qd.f32) * SEED_A.get(0) + qd.cast(j, qd.f32) * SEED_B.get(0) + s = qd.sin(x) * 43758.5453 + z[i, j] = (s - qd.floor(s)) * 2.0 + + +init_topo_kernel = QuadrantsKernelBuilder().bind_bag(noise_seeds).ingest(init_topo_template).compile() + +# --------------------------------------------------------------------------- +# routine kernels +# --------------------------------------------------------------------------- + +# mixed bag: a scalar Parameter, a helper and a const Parameter under one +# name. hill.d and hill.dt are both device accessors, hill.lap a specialized +# func. +hill = Bag({"d": d_p, "lap": laplacian_fn, "dt": dt_p}) + + +def diffuse_template(z_out: qd.Tensor, z_in: qd.Tensor): + for i, j in z_in: + z_out[i, j] = z_in[i, j] + hill.dt.get(0) * hill.d.get(0) * hill.lap(z_in, i, j) + + +diffuse_builder = QuadrantsKernelBuilder().bind("hill", hill).ingest(diffuse_template) + + +def uplift_template(z: qd.Tensor): + for i, j in z: + z[i, j] += up(i, j) * DT.get(0) + # base level: pin the east and west edges, so the ridge drains outward + if j == 0 or j == grid.n.get(0) - 1: + z[i, j] = SEA.get(0) + + +uplift_builder = ( + QuadrantsKernelBuilder() + .bind("up", uplift_at_fn) + .bind("DT", dt_p) + .bind("grid", grid) + .bind("SEA", sea_p) + .ingest(uplift_template) +) + +# --------------------------------------------------------------------------- +# fields (pooled - two buffers for ping-pong) +# --------------------------------------------------------------------------- +z0 = pool.get_data(qd.f32, (GRID_N, GRID_N)) +z1 = pool.get_data(qd.f32, (GRID_N, GRID_N)) + +init_topo_kernel(z0.data) + +# --------------------------------------------------------------------------- +# the routine +# --------------------------------------------------------------------------- +# One bag for the whole routine, merged from what each builder already binds. +# Both builders reach `grid` - the same Bag object, so the same uid, which is +# what lets merge() accept the collision instead of raising on it. +evolve = ( + QuadrantsRoutineBuilder() + .bind_bag(merge(diffuse_builder.as_bag(), uplift_builder.as_bag())) + .add_data("z0", z0.data) + .add_data("z1", z1.data) + .add_kernel(diffuse_builder, data_handle_ref=("z1", "z0")) + .add_kernel(uplift_builder, data_handle_ref=("z1",)) + .add_swap("z0", "z1") + .add_kernel(diffuse_builder, data_handle_ref=("z1", "z0")) + .add_kernel(uplift_builder, data_handle_ref=("z1",)) + .add_swap("z0", "z1") + .compile() +) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(z0.to_numpy(), cmap="terrain", vmin=0.0, vmax=150.0) +fig.colorbar(im, ax=ax, label="Elevation (m)") +ax.set_title("Hillslope LEM (Quadrants backend, Routine)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME // 2): + evolve() # two substeps, result lands back in z0 + sim_time += 2.0 * DT_VAL + + qd.sync() + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time / 1e6:.2f} Myr") + im.set_data(z0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +d_p.destroy() +sea_p.destroy() +uplift_p.destroy() +pool.release_data(z0) +pool.release_data(z1) diff --git a/examples/core/lem/lem_routine_taichi.py b/examples/core/lem/lem_routine_taichi.py new file mode 100644 index 0000000..223e3df --- /dev/null +++ b/examples/core/lem/lem_routine_taichi.py @@ -0,0 +1,260 @@ +""" +Hillslope landscape evolution as one Routine, exercising every part of the +core in a single model. + +The physics is deliberately small - linear hillslope diffusion against a +spatially variable uplift field, with the domain edges pinned to base level: + + dz/dt = D * laplacian(z) + U(x, y) + +What the file is here to show is how the pieces fit together when a model +needs all of them at once. The other examples each isolate one thing; this +one carries the lot: + + Parameter modes N, DT and DX are const Parameters, read uniformly via + .get(0) - the value still bakes to a compile-time literal + in generated code. D and SEA_LEVEL are scalars the host + retunes between frames. UPLIFT is a field, one rate per + node. + Helpers clampi binds a const; laplacian binds a bag and calls + clampi, so a helper reaches another helper; uplift_at + binds the UPLIFT *field* directly, which is what lets the + uplift kernel body stay a one-liner. + Bags grid is nested (grid.n, grid.dx), hill is mixed - a + scalar Parameter, a helper and a const under one name - + and the two noise seeds arrive flat through bind_bag. + Routine three kernels, two of them inside the routine, with the + z0/z1 ping-pong unrolled twice so the swaps compose to + the identity and the routine can be called repeatedly. + +The step the routine runs is diffuse then uplift-and-clamp, so uplift is +applied to what diffusion just wrote. Two of those, plus the two swaps, make +one routine call - and the result always lands back in z0. + +D and SEA_LEVEL are retuned between routine calls, never between the steps +inside one: set() on a scalar Parameter is what a routine expects between +calls (see routine.py, "Contract: no set()/destroy() mid-routine"), and there +is no python between a routine's own steps to run it in anyway. + +Author: B.G (07/2026) +""" + +import time + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti + +from pyfastflow.experimental.core.context.bag import Bag, merge +from pyfastflow.experimental.core.context.taichi_backend import ( + TaichiHelperBuilder, + TaichiKernelBuilder, + TaichiParameter, + TaichiRoutineBuilder, +) +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool + +ti.init(arch=ti.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, timing - never used as kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 2048 +DX_M = 100.0 +STEPS_PER_FRAME = 200 # two routine substeps per call - see the loop below + +D_VAL = 1.0e-2 # hillslope diffusivity, m2/yr +UPLIFT_MAX = 1.0e-6 # m/yr at the range crest +CFL_SAFETY = 0.2 +DT_VAL = CFL_SAFETY * DX_M**2 / (4.0 * D_VAL) + +pool = TaichiPool() + +# --------------------------------------------------------------------------- +# parameters - one of every mode +# --------------------------------------------------------------------------- +# const Parameters: fixed at compile time and folded into the generated code +# as a literal, but read through .get(0) like any other mode, so a template +# can be written without knowing a given name is const. +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool) +dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool) +seed_a_p = TaichiParameter("SEED_A", dtype=ti.f32, mode="const", value=12.9898, pool=pool) +seed_b_p = TaichiParameter("SEED_B", dtype=ti.f32, mode="const", value=78.233, pool=pool) + +dx_p = TaichiParameter("DX", dtype=ti.f32, mode="const", value=DX_M, pool=pool) + +# scalars: one cell each, retuned from the host between routine calls +d_p = TaichiParameter("D", dtype=ti.f32, mode="scalar", value=D_VAL, pool=pool) +sea_p = TaichiParameter("SEA_LEVEL", dtype=ti.f32, mode="scalar", value=0.0, pool=pool) + +# field: one value per node, filled from the host below +uplift_p = TaichiParameter( + "UPLIFT", dtype=ti.f32, mode="field", value=np.zeros(GRID_N * GRID_N), pool=pool, n_flat=GRID_N * GRID_N +) + +# a north-south uplift ridge, tapering to zero at the north and south edges +_yy = np.arange(GRID_N, dtype=np.float32)[:, None] * np.ones((1, GRID_N), np.float32) +_ridge = np.sin(np.pi * _yy / (GRID_N - 1)) ** 2 +uplift_p.set((UPLIFT_MAX * _ridge).ravel()) + +# --------------------------------------------------------------------------- +# bags +# --------------------------------------------------------------------------- +# nested: both grid.n and grid.dx are const Parameters, read through .get(0) - +# members resolve on their own type, not the bag's +grid = Bag({"n": n_p, "dx": dx_p}) + +# flat, for bind_bag: the kernel that uses these reads them as bare names +noise_seeds = Bag({"SEED_A": seed_a_p, "SEED_B": seed_b_p}) + +# --------------------------------------------------------------------------- +# device helpers +# --------------------------------------------------------------------------- + + +def clampi(i): + return min(max(i, 0), N.get(0) - 1) + + +clampi_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clampi) + + +def laplacian(field_, i, j): + # calls another helper, and reads a const Parameter out of a bound bag + ip = clampi(i + 1) + im = clampi(i - 1) + jp = clampi(j + 1) + jm = clampi(j - 1) + acc = field_[ip, j] + field_[im, j] + field_[i, jp] + field_[i, jm] - 4.0 * field_[i, j] + return acc / (grid.dx.get(0) * grid.dx.get(0)) + + +laplacian_fn = TaichiHelperBuilder().bind("clampi", clampi_fn).bind("grid", grid).ingest(laplacian) + + +def uplift_at(i, j): + # binds the UPLIFT *field* itself: a helper reads a non-const Parameter + # exactly the way a kernel does, so the caller passes only the indices + return UPLIFT.get(i * N.get(0) + j) + + +uplift_at_fn = TaichiHelperBuilder().bind("UPLIFT", uplift_p).bind("N", n_p).ingest(uplift_at) + +# --------------------------------------------------------------------------- +# one-shot setup kernel (runs once, outside the routine) +# --------------------------------------------------------------------------- + + +def init_topo_template(z: ti.template()): + # bind_bag put SEED_A / SEED_B in flat, so they read as bare names here + for i, j in z: + x = ti.cast(i, ti.f32) * SEED_A.get(0) + ti.cast(j, ti.f32) * SEED_B.get(0) + s = ti.sin(x) * 43758.5453 + z[i, j] = (s - ti.floor(s)) * 2.0 + + +init_topo_kernel = TaichiKernelBuilder().bind_bag(noise_seeds).ingest(init_topo_template).compile() + +# --------------------------------------------------------------------------- +# routine kernels +# --------------------------------------------------------------------------- + +# mixed bag: a scalar Parameter, a helper and a const Parameter under one +# name. hill.d and hill.dt are both device accessors, hill.lap a specialized +# func. +hill = Bag({"d": d_p, "lap": laplacian_fn, "dt": dt_p}) + + +def diffuse_template(z_out: ti.template(), z_in: ti.template()): + for i, j in z_in: + z_out[i, j] = z_in[i, j] + hill.dt.get(0) * hill.d.get(0) * hill.lap(z_in, i, j) + + +diffuse_builder = TaichiKernelBuilder().bind("hill", hill).ingest(diffuse_template) + + +def uplift_template(z: ti.template()): + for i, j in z: + z[i, j] += up(i, j) * DT.get(0) + # base level: pin the east and west edges, so the ridge drains outward + if j == 0 or j == grid.n.get(0) - 1: + z[i, j] = SEA.get(0) + + +uplift_builder = ( + TaichiKernelBuilder() + .bind("up", uplift_at_fn) + .bind("DT", dt_p) + .bind("grid", grid) + .bind("SEA", sea_p) + .ingest(uplift_template) +) + +# --------------------------------------------------------------------------- +# fields (pooled - two buffers for ping-pong) +# --------------------------------------------------------------------------- +z0 = pool.get_data(ti.f32, (GRID_N, GRID_N)) +z1 = pool.get_data(ti.f32, (GRID_N, GRID_N)) + +init_topo_kernel(z0.data) + +# --------------------------------------------------------------------------- +# the routine +# --------------------------------------------------------------------------- +# One bag for the whole routine, merged from what each builder already binds. +# Both builders reach `grid` - the same Bag object, so the same uid, which is +# what lets merge() accept the collision instead of raising on it. +evolve = ( + TaichiRoutineBuilder() + .bind_bag(merge(diffuse_builder.as_bag(), uplift_builder.as_bag())) + .add_data("z0", z0.data) + .add_data("z1", z1.data) + .add_kernel(diffuse_builder, data_handle_ref=("z1", "z0")) + .add_kernel(uplift_builder, data_handle_ref=("z1",)) + .add_swap("z0", "z1") + .add_kernel(diffuse_builder, data_handle_ref=("z1", "z0")) + .add_kernel(uplift_builder, data_handle_ref=("z1",)) + .add_swap("z0", "z1") + .compile() +) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +fig, ax = plt.subplots() +im = ax.imshow(z0.to_numpy(), cmap="terrain", vmin=0.0, vmax=150.0) +fig.colorbar(im, ax=ax, label="Elevation (m)") +ax.set_title("Hillslope LEM (Taichi backend, Routine)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="white", fontsize=9, bbox=dict(facecolor="black", alpha=0.4, pad=2), +) +fig.show() + +sim_time = 0.0 +try: + while True: + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME // 2): + evolve() # two substeps, result lands back in z0 + sim_time += 2.0 * DT_VAL + + ti.sync() + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time / 1e6:.2f} Myr") + im.set_data(z0.to_numpy()) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +d_p.destroy() +sea_p.destroy() +uplift_p.destroy() +pool.release_data(z0) +pool.release_data(z1) diff --git a/examples/core/shallow_water/shallow_water_cupy.py b/examples/core/shallow_water/shallow_water_cupy.py new file mode 100644 index 0000000..48e7f4c --- /dev/null +++ b/examples/core/shallow_water/shallow_water_cupy.py @@ -0,0 +1,308 @@ +""" +Shallow-water waves in a square tank, built on pyfastflow's backend-agnostic +core (Parameter/Helper/Kernel/Pool + Bags), Cupy (cp.RawKernel) backend. + +Same model as shallow_water_taichi.py (Kass & Miller 1990 on an Arakawa-C +staggered grid), authored as CUDA source strings. The grid is stored flat +(N*N, row-major: idx = i*N + j); kernels launch one thread per cell. + +Bag showcase: bags reach the CUDA source through the SAME `$...$` spans as +flat params, just with a dotted head - the span parser walks Bag members: + + $phys.dx.get(0)$ const bag member -> baked CUDA literal + $phys.g.get(0)$ scalar bag member -> auto-generated `phys_g` pointer arg + $drop.cx.get(0)$ host-set splash site, same mechanism + $ops.clamp(i + 1)$ helper from a Bag -> its __device__ source spliced + +so one .bind("phys", phys) / .bind("ops", ops) carries the whole group, and +the source never declares the generated pointer arguments. The three bags are +split by role, not by kind: a Bag has no member type, so one could equally hold +`phys` and `ops` together (see heat_diffusion's `heat`). Top-level const +params (N, REST_DEPTH, DROP_R) become #defines, used bare. Spans do not nest, +so span results are read into temps before being passed to a helper span. + +Author: B.G (07/2026) +""" + +import random +import time + +import cupy as cp +import matplotlib.pyplot as plt +import numpy as np + +from pyfastflow.experimental.core.context.bag import Bag +from pyfastflow.experimental.core.context.cupy_backend import ( + CupyHelperBuilder, + CupyKernelBuilder, + CupyParameter, +) +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool + +# --------------------------------------------------------------------------- +# host-side constants +# --------------------------------------------------------------------------- +GRID_N = 640 +NN = GRID_N * GRID_N +STEPS_PER_FRAME = 40 +DROP_EVERY = 25 # frames between automatic stone drops +BLOCK = 256 +GRID = (NN + BLOCK - 1) // BLOCK + +# Physical grounding (see the taichi demo for the reasoning): a 4 m x 4 m tank +# holding a thin (5 cm) sheet of water; c = sqrt(g*H), dt <= dx / (c*sqrt(2)). +WORLD_M = 4.0 +DX_M = WORLD_M / GRID_N +G_VAL = 9.81 +REST_DEPTH_VAL = 0.05 +WAVE_C = (G_VAL * REST_DEPTH_VAL) ** 0.5 +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M / (WAVE_C * 2.0**0.5) + +DAMP_RATE = 0.3 # 1/s +DAMP_VAL = 1.0 - DAMP_RATE * DT_VAL + +DROP_R_VAL = 12 # splash radius, cells +DROP_AMP_VAL = 0.02 # m, height a stone adds at impact + +pool = CupyPool() + +# --------------------------------------------------------------------------- +# parameters +# --------------------------------------------------------------------------- +# Structural constants -> #define, used bare in the CUDA source. +n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool) +rest_depth_p = CupyParameter("REST_DEPTH", dtype=np.float32, mode="const", value=REST_DEPTH_VAL, pool=pool) +drop_r_p = CupyParameter("DROP_R", dtype=np.int32, mode="const", value=DROP_R_VAL, pool=pool) + +# phys Bag: g is scalar (host-tunable live), dx/dt/damp are const - all +# written the same way in the source, $phys..get(0)$. +g_p = CupyParameter("g", dtype=np.float32, mode="scalar", value=G_VAL, pool=pool) +dx_p = CupyParameter("dx", dtype=np.float32, mode="const", value=DX_M, pool=pool) +dt_p = CupyParameter("dt", dtype=np.float32, mode="const", value=DT_VAL, pool=pool) +damp_p = CupyParameter("damp", dtype=np.float32, mode="const", value=DAMP_VAL, pool=pool) +phys = Bag({"g": g_p, "dx": dx_p, "dt": dt_p, "damp": damp_p}) + +# drop Bag: splash site + amplitude, host-set each time a stone falls. +drop_cx_p = CupyParameter("cx", dtype=np.int32, mode="scalar", value=GRID_N // 2, pool=pool) +drop_cy_p = CupyParameter("cy", dtype=np.int32, mode="scalar", value=GRID_N // 2, pool=pool) +drop_amp_p = CupyParameter("amp", dtype=np.float32, mode="scalar", value=0.0, pool=pool) +drop = Bag({"cx": drop_cx_p, "cy": drop_cy_p, "amp": drop_amp_p}) + +# --------------------------------------------------------------------------- +# device helpers -> ops Bag +# --------------------------------------------------------------------------- +clamp_fn = ( + CupyHelperBuilder() + .bind("N", n_p) + .ingest("__device__ int clampi(int i) { return i < 0 ? 0 : (i >= N ? N - 1 : i); }") +) + +face_depth_fn = ( + CupyHelperBuilder() + .ingest( + r""" +__device__ float face_depth(float up, float down, float vel) { + // upwind water depth at a face: the upstream cell when flow is outward + return vel > 0.0f ? up : down; +} +""" + ) +) + +ops = Bag({"clamp": clamp_fn, "face_depth": face_depth_fn}) + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- +init_height_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("REST_DEPTH", rest_depth_p) + .ingest( + r""" +__global__ void init_height(float* h) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + h[idx] = REST_DEPTH; +} +""" + ) + .compile() +) + +apply_drop_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("DROP_R", drop_r_p) + .bind("drop", drop) + .ingest( + r""" +__global__ void apply_drop(float* h) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int dxr = idx / N - $drop.cx.get(0)$; + int dyr = idx % N - $drop.cy.get(0)$; + if (dxr * dxr + dyr * dyr <= DROP_R * DROP_R) { + h[idx] += $drop.amp.get(0)$; + } +} +""" + ) + .compile() +) + +update_velocity_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("phys", phys) + .ingest( + r""" +__global__ void update_velocity(float* u, float* v, const float* h) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N; + int j = idx % N; + float gdtdx = $phys.g.get(0)$ * $phys.dt.get(0)$ / $phys.dx.get(0)$; + float damp = $phys.damp.get(0)$; + // u lives on the west face of cell (i,j); i==0 is the tank wall. + if (i > 0) { + u[idx] = (u[idx] + gdtdx * (h[idx - N] - h[idx])) * damp; + } else { + u[idx] = 0.0f; + } + // v lives on the south face of cell (i,j); j==0 is the tank wall. + if (j > 0) { + v[idx] = (v[idx] + gdtdx * (h[idx - 1] - h[idx])) * damp; + } else { + v[idx] = 0.0f; + } +} +""" + ) + .compile() +) + +update_height_kernel = ( + CupyKernelBuilder() + .bind("N", n_p) + .bind("phys", phys) + .bind("ops", ops) + .ingest( + r""" +__global__ void update_height(float* h_out, const float* h_in, const float* u, const float* v) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * N) return; + int i = idx / N; + int j = idx % N; + + // face velocities: u[i]=west face, u[i+1]=east face (0 at the wall). + float uw = u[idx]; + float ue = 0.0f; + if (i < N - 1) ue = u[idx + N]; + float vs = v[idx]; + float vn = 0.0f; + if (j < N - 1) vn = v[idx + 1]; + + int ip = $ops.clamp(i + 1)$; + int im = $ops.clamp(i - 1)$; + int jp = $ops.clamp(j + 1)$; + int jm = $ops.clamp(j - 1)$; + + float hc = h_in[idx]; + float hxm = h_in[im * N + j]; + float hxp = h_in[ip * N + j]; + float hym = h_in[i * N + jm]; + float hyp = h_in[i * N + jp]; + + float hw = $ops.face_depth(hxm, hc, uw)$; + float he = $ops.face_depth(hc, hxp, ue)$; + float hs = $ops.face_depth(hym, hc, vs)$; + float hn = $ops.face_depth(hc, hyp, vn)$; + + float flux = (he * ue - hw * uw) + (hn * vn - hs * vs); + h_out[idx] = hc - $phys.dt.get(0)$ / $phys.dx.get(0)$ * flux; +} +""" + ) + .compile() +) + +# --------------------------------------------------------------------------- +# fields (pooled flat buffers; h is ping-ponged, u/v updated in place) +# --------------------------------------------------------------------------- +h0 = pool.get_data(np.float32, (NN,)) +h1 = pool.get_data(np.float32, (NN,)) +u = pool.get_data(np.float32, (NN,)) +v = pool.get_data(np.float32, (NN,)) + +u.data[...] = 0.0 +v.data[...] = 0.0 +init_height_kernel(h0.data, grid=GRID, block=BLOCK) + +# first stone, dead center, so there is motion on frame 0 +drop_cx_p.set(GRID_N // 2) +drop_cy_p.set(GRID_N // 2) +drop_amp_p.set(DROP_AMP_VAL) +apply_drop_kernel(h0.data, grid=GRID, block=BLOCK) +drop_amp_p.set(0.0) + +# --------------------------------------------------------------------------- +# live view (surface elevation h - rest depth) +# --------------------------------------------------------------------------- +elev_lim = DROP_AMP_VAL * 0.35 +fig, ax = plt.subplots() +im = ax.imshow( + (h0.to_numpy().reshape(GRID_N, GRID_N) - REST_DEPTH_VAL).T, + cmap="RdBu_r", vmin=-elev_lim, vmax=elev_lim, origin="lower", +) +fig.colorbar(im, ax=ax, label="surface elevation (m)") +ax.set_title("Shallow-water waves in a tank (Cupy backend)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="black", fontsize=9, bbox=dict(facecolor="white", alpha=0.5, pad=2), +) +fig.show() + +sim_time = 0.0 +frame = 0 +try: + while True: + frame += 1 + if frame % DROP_EVERY == 0: + drop_cx_p.set(random.randint(DROP_R_VAL, GRID_N - 1 - DROP_R_VAL)) + drop_cy_p.set(random.randint(DROP_R_VAL, GRID_N - 1 - DROP_R_VAL)) + drop_amp_p.set(DROP_AMP_VAL) + apply_drop_kernel(h0.data, grid=GRID, block=BLOCK) + drop_amp_p.set(0.0) + + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + update_velocity_kernel(u.data, v.data, h0.data, grid=GRID, block=BLOCK) + update_height_kernel(h1.data, h0.data, u.data, v.data, grid=GRID, block=BLOCK) + h0, h1 = h1, h0 + sim_time += DT_VAL + + cp.cuda.Device().synchronize() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.1f} s") + im.set_data((h0.to_numpy().reshape(GRID_N, GRID_N) - REST_DEPTH_VAL).T) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +# destroy() hands a Parameter's storage back to the pool; it is a no-op on a +# const, which owns none. Safe only because nothing will launch again - the +# pool may reissue these buffers, while the compiled kernels above still point +# at them (see parameter.py, "Lifetime of a compiled object"). +for param in (g_p, drop_cx_p, drop_cy_p, drop_amp_p): + param.destroy() +for buf in (h0, h1, u, v): + pool.release_data(buf) +print("pooled storage released") diff --git a/examples/core/shallow_water/shallow_water_quadrants.py b/examples/core/shallow_water/shallow_water_quadrants.py new file mode 100644 index 0000000..33af678 --- /dev/null +++ b/examples/core/shallow_water/shallow_water_quadrants.py @@ -0,0 +1,281 @@ +""" +Shallow-water waves in a square tank, built on pyfastflow's backend-agnostic +core (Parameter/Helper/Kernel/Pool + Bags), Quadrants backend. + +Same model as shallow_water_taichi.py: the Kass & Miller (1990) stable +shallow-water update on an Arakawa-C staggered grid. Cell-centered water +column height h, x-velocity u on vertical faces, y-velocity v on horizontal +faces: + - update_velocity: u,v accelerate down the height gradient (g * grad h), + then light damping; domain-edge faces are pinned to 0 (reflective walls). + - update_height: h advected by the velocity divergence with upwind face + depths (mass-conserving, stable) -> waves, sloshing, reflection. + - apply_drop: a disc splash raises h wherever a "stone" lands; the landing + site + amplitude come from host-set scalar params, so stones drop live. + +Bag showcase (heat_diffusion mixes flat binds, bind_bag and a nested bag; here +everything goes through whole-bag binds): the physical constants travel +as ONE `phys` Bag (g/dx/dt/damp), read in-kernel as phys.g.get(0), +phys.dx.get(0), ...; the neighbour math travels as ONE `ops` Bag +(clamp, face_depth), called as ops.clamp(i), ops.face_depth(...); and the +splash controls travel as a `drop` Bag (cx/cy/amp), read drop.cx.get(0). +Bind the whole bag once (bind("phys", phys_bag)); dotted paths in the template +resolve to each member's device view - the kernel body never names them flat. +The three bags are split by role, not by kind: a Bag has no member type, so +one could equally hold `phys` and `ops` together (see heat_diffusion's `heat`). + +Structural constants (N, DROP_R, REST_DEPTH) are const Parameters, read +uniformly via .get(0) - the value still bakes to a compile-time literal in +generated code. + +Author: B.G (07/2026) +""" + +import random +import time + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd + +from pyfastflow.experimental.core.context.bag import Bag +from pyfastflow.experimental.core.context.quadrants_backend import ( + QuadrantsHelperBuilder, + QuadrantsKernelBuilder, + QuadrantsParameter, +) +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool + +qd.init(arch=qd.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, loop/timing counts - never kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 640 +STEPS_PER_FRAME = 40 +DROP_EVERY = 25 # frames between automatic stone drops + +# Physical grounding: a 4 m x 4 m tank holding a thin (5 cm) sheet of water. +# Shallow-water wave speed is c = sqrt(g*H); the explicit CFL limit is +# dt <= dx / (c*sqrt(2)), so dt is derived from the tank, not tuned by feel. +WORLD_M = 4.0 +DX_M = WORLD_M / GRID_N # meters per cell +G_VAL = 9.81 # m/s^2 +REST_DEPTH_VAL = 0.05 # m, still-water column height +WAVE_C = (G_VAL * REST_DEPTH_VAL) ** 0.5 # m/s +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M / (WAVE_C * 2.0**0.5) # seconds + +# Light drag so a tank eventually settles: per-step factor 1 - rate*dt. +DAMP_RATE = 0.3 # 1/s +DAMP_VAL = 1.0 - DAMP_RATE * DT_VAL + +DROP_R_VAL = 12 # splash radius, cells +DROP_AMP_VAL = 0.02 # m, height a stone adds at impact + +pool = QuadrantsPool() + +# --------------------------------------------------------------------------- +# parameters +# --------------------------------------------------------------------------- +# Structural constants: const mode, folded into the generated code as a +# literal but still read through .get(0). +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool) +rest_depth_p = QuadrantsParameter("REST_DEPTH", dtype=qd.f32, mode="const", value=REST_DEPTH_VAL, pool=pool) +drop_r_p = QuadrantsParameter("DROP_R", dtype=qd.i32, mode="const", value=DROP_R_VAL, pool=pool) + +# phys Bag: g is scalar (host-tunable live), dx/dt/damp are const - but all +# read uniformly as phys..get(0), so the kernels never branch on mode. +g_p = QuadrantsParameter("g", dtype=qd.f32, mode="scalar", value=G_VAL, pool=pool) +dx_p = QuadrantsParameter("dx", dtype=qd.f32, mode="const", value=DX_M, pool=pool) +dt_p = QuadrantsParameter("dt", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool) +damp_p = QuadrantsParameter("damp", dtype=qd.f32, mode="const", value=DAMP_VAL, pool=pool) +phys = Bag({"g": g_p, "dx": dx_p, "dt": dt_p, "damp": damp_p}) + +# drop Bag: splash site + amplitude, host-set each time a stone falls. +drop_cx_p = QuadrantsParameter("cx", dtype=qd.i32, mode="scalar", value=GRID_N // 2, pool=pool) +drop_cy_p = QuadrantsParameter("cy", dtype=qd.i32, mode="scalar", value=GRID_N // 2, pool=pool) +drop_amp_p = QuadrantsParameter("amp", dtype=qd.f32, mode="scalar", value=0.0, pool=pool) +drop = Bag({"cx": drop_cx_p, "cy": drop_cy_p, "amp": drop_amp_p}) + +# --------------------------------------------------------------------------- +# device helpers -> ops Bag +# --------------------------------------------------------------------------- + + +def clamp(i): + return min(max(i, 0), N.get(0) - 1) + + +clamp_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clamp) + + +def face_depth(up, down, vel): + """Upwind water depth at a face: the upstream cell when flow is outward.""" + d = down + if vel > 0.0: + d = up + return d + + +face_depth_fn = QuadrantsHelperBuilder().ingest(face_depth) + +ops = Bag({"clamp": clamp_fn, "face_depth": face_depth_fn}) + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- + + +def init_height_template(h: qd.Tensor): + for i, j in h: + h[i, j] = REST_DEPTH.get(0) + + +init_height_kernel = QuadrantsKernelBuilder().bind("REST_DEPTH", rest_depth_p).ingest(init_height_template).compile() + + +def apply_drop_template(h: qd.Tensor): + for i, j in h: + dxr = i - drop.cx.get(0) + dyr = j - drop.cy.get(0) + if dxr * dxr + dyr * dyr <= DROP_R.get(0) * DROP_R.get(0): + h[i, j] += drop.amp.get(0) + + +apply_drop_kernel = ( + QuadrantsKernelBuilder() + .bind("drop", drop) + .bind("DROP_R", drop_r_p) + .ingest(apply_drop_template) + .compile() +) + + +def update_velocity_template(u: qd.Tensor, v: qd.Tensor, h: qd.Tensor): + for i, j in h: + # u lives on the west face of cell (i,j); i==0 is the tank wall. + if i > 0: + acc = phys.g.get(0) * phys.dt.get(0) / phys.dx.get(0) * (h[i - 1, j] - h[i, j]) + u[i, j] = (u[i, j] + acc) * phys.damp.get(0) + else: + u[i, j] = 0.0 + # v lives on the south face of cell (i,j); j==0 is the tank wall. + if j > 0: + acc = phys.g.get(0) * phys.dt.get(0) / phys.dx.get(0) * (h[i, j - 1] - h[i, j]) + v[i, j] = (v[i, j] + acc) * phys.damp.get(0) + else: + v[i, j] = 0.0 + + +update_velocity_kernel = QuadrantsKernelBuilder().bind("phys", phys).ingest(update_velocity_template).compile() + + +def update_height_template(h_out: qd.Tensor, h_in: qd.Tensor, u: qd.Tensor, v: qd.Tensor): + for i, j in h_in: + # face velocities: u[i]=west face, u[i+1]=east face (0 at the wall). + uw = u[i, j] + ue = 0.0 + if i < N.get(0) - 1: + ue = u[i + 1, j] + vs = v[i, j] + vn = 0.0 + if j < N.get(0) - 1: + vn = v[i, j + 1] + + ip = ops.clamp(i + 1) + im = ops.clamp(i - 1) + jp = ops.clamp(j + 1) + jm = ops.clamp(j - 1) + + hw = ops.face_depth(h_in[im, j], h_in[i, j], uw) + he = ops.face_depth(h_in[i, j], h_in[ip, j], ue) + hs = ops.face_depth(h_in[i, jm], h_in[i, j], vs) + hn = ops.face_depth(h_in[i, j], h_in[i, jp], vn) + + flux = (he * ue - hw * uw) + (hn * vn - hs * vs) + h_out[i, j] = h_in[i, j] - phys.dt.get(0) / phys.dx.get(0) * flux + + +update_height_kernel = ( + QuadrantsKernelBuilder() + .bind("N", n_p) + .bind("phys", phys) + .bind("ops", ops) + .ingest(update_height_template) + .compile() +) + +# --------------------------------------------------------------------------- +# fields (pooled; h is ping-ponged, u/v updated in place - start at 0) +# --------------------------------------------------------------------------- +h0 = pool.get_data(qd.f32, (GRID_N, GRID_N)) +h1 = pool.get_data(qd.f32, (GRID_N, GRID_N)) +u = pool.get_data(qd.f32, (GRID_N, GRID_N)) +v = pool.get_data(qd.f32, (GRID_N, GRID_N)) + +init_height_kernel(h0.data) + +# first stone, dead center, so there is motion on frame 0 +drop_cx_p.set(GRID_N // 2) +drop_cy_p.set(GRID_N // 2) +drop_amp_p.set(DROP_AMP_VAL) +apply_drop_kernel(h0.data) +drop_amp_p.set(0.0) + +# --------------------------------------------------------------------------- +# live view (surface elevation h - rest depth) +# --------------------------------------------------------------------------- +elev_lim = DROP_AMP_VAL * 0.35 +fig, ax = plt.subplots() +im = ax.imshow((h0.to_numpy() - REST_DEPTH_VAL).T, cmap="RdBu_r", vmin=-elev_lim, vmax=elev_lim, origin="lower") +fig.colorbar(im, ax=ax, label="surface elevation (m)") +ax.set_title("Shallow-water waves in a tank (Quadrants backend)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="black", fontsize=9, bbox=dict(facecolor="white", alpha=0.5, pad=2), +) +fig.show() + +sim_time = 0.0 +frame = 0 +try: + while True: + frame += 1 + if frame % DROP_EVERY == 0: + drop_cx_p.set(random.randint(DROP_R_VAL, GRID_N - 1 - DROP_R_VAL)) + drop_cy_p.set(random.randint(DROP_R_VAL, GRID_N - 1 - DROP_R_VAL)) + drop_amp_p.set(DROP_AMP_VAL) + apply_drop_kernel(h0.data) + drop_amp_p.set(0.0) + + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + update_velocity_kernel(u.data, v.data, h0.data) + update_height_kernel(h1.data, h0.data, u.data, v.data) + h0, h1 = h1, h0 + sim_time += DT_VAL + + qd.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.1f} s") + im.set_data((h0.to_numpy() - REST_DEPTH_VAL).T) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +# destroy() hands a Parameter's storage back to the pool; it is a no-op on a +# const, which owns none. Safe only because nothing will launch again - the +# pool may reissue these buffers, while the compiled kernels above still point +# at them (see parameter.py, "Lifetime of a compiled object"). +for param in (g_p, drop_cx_p, drop_cy_p, drop_amp_p): + param.destroy() +for buf in (h0, h1, u, v): + pool.release_data(buf) +print("pooled storage released") diff --git a/examples/core/shallow_water/shallow_water_taichi.py b/examples/core/shallow_water/shallow_water_taichi.py new file mode 100644 index 0000000..cda70a2 --- /dev/null +++ b/examples/core/shallow_water/shallow_water_taichi.py @@ -0,0 +1,280 @@ +""" +Shallow-water waves in a square tank, built on pyfastflow's backend-agnostic +core (Parameter/Helper/Kernel/Pool + Bags), Taichi backend. + +Model: the Kass & Miller (1990) stable shallow-water update on an Arakawa-C +staggered grid. Cell-centered water column height h, x-velocity u on vertical +faces, y-velocity v on horizontal faces: + - update_velocity: u,v accelerate down the height gradient (g * grad h), + then light damping; domain-edge faces are pinned to 0 (reflective walls). + - update_height: h advected by the velocity divergence with upwind face + depths (mass-conserving, stable) -> waves, sloshing, reflection. + - apply_drop: a disc splash raises h wherever a "stone" lands; the landing + site + amplitude come from host-set scalar params, so stones drop live. + +Bag showcase (heat_diffusion mixes flat binds, bind_bag and a nested bag; here +everything goes through whole-bag binds): the physical constants travel +as ONE `phys` Bag (g/dx/dt/damp), read in-kernel as phys.g.get(0), +phys.dx.get(0), ...; the neighbour math travels as ONE `ops` Bag +(clamp, face_depth), called as ops.clamp(i), ops.face_depth(...); and the +splash controls travel as a `drop` Bag (cx/cy/amp), read drop.cx.get(0). +Bind the whole bag once (bind("phys", phys_bag)); dotted paths in the template +resolve to each member's device view - the kernel body never names them flat. +The three bags are split by role, not by kind: a Bag has no member type, so +one could equally hold `phys` and `ops` together (see heat_diffusion's `heat`). + +Structural constants (N, DROP_R, REST_DEPTH) are const Parameters, read +uniformly via .get(0) - the value still bakes to a compile-time literal in +generated code. + +Author: B.G (07/2026) +""" + +import random +import time + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti + +from pyfastflow.experimental.core.context.bag import Bag +from pyfastflow.experimental.core.context.taichi_backend import ( + TaichiHelperBuilder, + TaichiKernelBuilder, + TaichiParameter, +) +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool + +ti.init(arch=ti.gpu) + +# --------------------------------------------------------------------------- +# host-side constants (grid size, loop/timing counts - never kernel globals) +# --------------------------------------------------------------------------- +GRID_N = 640 +STEPS_PER_FRAME = 40 +DROP_EVERY = 25 # frames between automatic stone drops + +# Physical grounding: a 4 m x 4 m tank holding a thin (5 cm) sheet of water. +# Shallow-water wave speed is c = sqrt(g*H); the explicit CFL limit is +# dt <= dx / (c*sqrt(2)), so dt is derived from the tank, not tuned by feel. +WORLD_M = 4.0 +DX_M = WORLD_M / GRID_N # meters per cell +G_VAL = 9.81 # m/s^2 +REST_DEPTH_VAL = 0.05 # m, still-water column height +WAVE_C = (G_VAL * REST_DEPTH_VAL) ** 0.5 # m/s +CFL_SAFETY = 0.4 +DT_VAL = CFL_SAFETY * DX_M / (WAVE_C * 2.0**0.5) # seconds + +# Light drag so a tank eventually settles: per-step factor 1 - rate*dt. +DAMP_RATE = 0.3 # 1/s +DAMP_VAL = 1.0 - DAMP_RATE * DT_VAL + +DROP_R_VAL = 12 # splash radius, cells +DROP_AMP_VAL = 0.02 # m, height a stone adds at impact + +pool = TaichiPool() + +# --------------------------------------------------------------------------- +# parameters +# --------------------------------------------------------------------------- +# Structural constants: const mode, folded into the generated code as a +# literal but still read through .get(0). +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool) +rest_depth_p = TaichiParameter("REST_DEPTH", dtype=ti.f32, mode="const", value=REST_DEPTH_VAL, pool=pool) +drop_r_p = TaichiParameter("DROP_R", dtype=ti.i32, mode="const", value=DROP_R_VAL, pool=pool) + +# phys Bag: g is scalar (host-tunable live), dx/dt/damp are const - but all +# read uniformly as phys..get(0), so the kernels never branch on mode. +g_p = TaichiParameter("g", dtype=ti.f32, mode="scalar", value=G_VAL, pool=pool) +dx_p = TaichiParameter("dx", dtype=ti.f32, mode="const", value=DX_M, pool=pool) +dt_p = TaichiParameter("dt", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool) +damp_p = TaichiParameter("damp", dtype=ti.f32, mode="const", value=DAMP_VAL, pool=pool) +phys = Bag({"g": g_p, "dx": dx_p, "dt": dt_p, "damp": damp_p}) + +# drop Bag: splash site + amplitude, host-set each time a stone falls. +drop_cx_p = TaichiParameter("cx", dtype=ti.i32, mode="scalar", value=GRID_N // 2, pool=pool) +drop_cy_p = TaichiParameter("cy", dtype=ti.i32, mode="scalar", value=GRID_N // 2, pool=pool) +drop_amp_p = TaichiParameter("amp", dtype=ti.f32, mode="scalar", value=0.0, pool=pool) +drop = Bag({"cx": drop_cx_p, "cy": drop_cy_p, "amp": drop_amp_p}) + +# --------------------------------------------------------------------------- +# device helpers -> ops Bag +# --------------------------------------------------------------------------- + + +def clamp(i): + return min(max(i, 0), N.get(0) - 1) + + +clamp_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clamp) + + +def face_depth(up, down, vel): + """Upwind water depth at a face: the upstream cell when flow is outward.""" + d = down + if vel > 0.0: + d = up + return d + + +face_depth_fn = TaichiHelperBuilder().ingest(face_depth) + +ops = Bag({"clamp": clamp_fn, "face_depth": face_depth_fn}) + +# --------------------------------------------------------------------------- +# kernels +# --------------------------------------------------------------------------- + + +def init_height_template(h: ti.template()): + for i, j in h: + h[i, j] = REST_DEPTH.get(0) + + +init_height_kernel = TaichiKernelBuilder().bind("REST_DEPTH", rest_depth_p).ingest(init_height_template).compile() + + +def apply_drop_template(h: ti.template()): + for i, j in h: + dxr = i - drop.cx.get(0) + dyr = j - drop.cy.get(0) + if dxr * dxr + dyr * dyr <= DROP_R.get(0) * DROP_R.get(0): + h[i, j] += drop.amp.get(0) + + +apply_drop_kernel = ( + TaichiKernelBuilder() + .bind("drop", drop) + .bind("DROP_R", drop_r_p) + .ingest(apply_drop_template) + .compile() +) + + +def update_velocity_template(u: ti.template(), v: ti.template(), h: ti.template()): + for i, j in h: + # u lives on the west face of cell (i,j); i==0 is the tank wall. + if i > 0: + acc = phys.g.get(0) * phys.dt.get(0) / phys.dx.get(0) * (h[i - 1, j] - h[i, j]) + u[i, j] = (u[i, j] + acc) * phys.damp.get(0) + else: + u[i, j] = 0.0 + # v lives on the south face of cell (i,j); j==0 is the tank wall. + if j > 0: + acc = phys.g.get(0) * phys.dt.get(0) / phys.dx.get(0) * (h[i, j - 1] - h[i, j]) + v[i, j] = (v[i, j] + acc) * phys.damp.get(0) + else: + v[i, j] = 0.0 + + +update_velocity_kernel = TaichiKernelBuilder().bind("phys", phys).ingest(update_velocity_template).compile() + + +def update_height_template(h_out: ti.template(), h_in: ti.template(), u: ti.template(), v: ti.template()): + for i, j in h_in: + # face velocities: u[i]=west face, u[i+1]=east face (0 at the wall). + uw = u[i, j] + ue = 0.0 + if i < N.get(0) - 1: + ue = u[i + 1, j] + vs = v[i, j] + vn = 0.0 + if j < N.get(0) - 1: + vn = v[i, j + 1] + + ip = ops.clamp(i + 1) + im = ops.clamp(i - 1) + jp = ops.clamp(j + 1) + jm = ops.clamp(j - 1) + + hw = ops.face_depth(h_in[im, j], h_in[i, j], uw) + he = ops.face_depth(h_in[i, j], h_in[ip, j], ue) + hs = ops.face_depth(h_in[i, jm], h_in[i, j], vs) + hn = ops.face_depth(h_in[i, j], h_in[i, jp], vn) + + flux = (he * ue - hw * uw) + (hn * vn - hs * vs) + h_out[i, j] = h_in[i, j] - phys.dt.get(0) / phys.dx.get(0) * flux + + +update_height_kernel = ( + TaichiKernelBuilder() + .bind("N", n_p) + .bind("phys", phys) + .bind("ops", ops) + .ingest(update_height_template) + .compile() +) + +# --------------------------------------------------------------------------- +# fields (pooled; h is ping-ponged, u/v updated in place - start at 0) +# --------------------------------------------------------------------------- +h0 = pool.get_data(ti.f32, (GRID_N, GRID_N)) +h1 = pool.get_data(ti.f32, (GRID_N, GRID_N)) +u = pool.get_data(ti.f32, (GRID_N, GRID_N)) +v = pool.get_data(ti.f32, (GRID_N, GRID_N)) + +init_height_kernel(h0.data) + +# first stone, dead center, so there is motion on frame 0 +drop_cx_p.set(GRID_N // 2) +drop_cy_p.set(GRID_N // 2) +drop_amp_p.set(DROP_AMP_VAL) +apply_drop_kernel(h0.data) +drop_amp_p.set(0.0) + +# --------------------------------------------------------------------------- +# live view (surface elevation h - rest depth) +# --------------------------------------------------------------------------- +elev_lim = DROP_AMP_VAL * 0.35 +fig, ax = plt.subplots() +im = ax.imshow((h0.to_numpy() - REST_DEPTH_VAL).T, cmap="RdBu_r", vmin=-elev_lim, vmax=elev_lim, origin="lower") +fig.colorbar(im, ax=ax, label="surface elevation (m)") +ax.set_title("Shallow-water waves in a tank (Taichi backend)") +time_text = ax.text( + 0.02, 0.98, "", transform=ax.transAxes, va="top", ha="left", + color="black", fontsize=9, bbox=dict(facecolor="white", alpha=0.5, pad=2), +) +fig.show() + +sim_time = 0.0 +frame = 0 +try: + while True: + frame += 1 + if frame % DROP_EVERY == 0: + drop_cx_p.set(random.randint(DROP_R_VAL, GRID_N - 1 - DROP_R_VAL)) + drop_cy_p.set(random.randint(DROP_R_VAL, GRID_N - 1 - DROP_R_VAL)) + drop_amp_p.set(DROP_AMP_VAL) + apply_drop_kernel(h0.data) + drop_amp_p.set(0.0) + + t_start = time.perf_counter() + for _ in range(STEPS_PER_FRAME): + update_velocity_kernel(u.data, v.data, h0.data) + update_height_kernel(h1.data, h0.data, u.data, v.data) + h0, h1 = h1, h0 + sim_time += DT_VAL + + ti.sync() # GPU is async; sync before stopping the timer + frame_ms = (time.perf_counter() - t_start) * 1e3 + print(f"{STEPS_PER_FRAME} steps: {frame_ms:8.1f} ms ({frame_ms / STEPS_PER_FRAME * 1e3:6.1f} us/step)") + + time_text.set_text(f"t = {sim_time:.1f} s") + im.set_data((h0.to_numpy() - REST_DEPTH_VAL).T) + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.1) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +# destroy() hands a Parameter's storage back to the pool; it is a no-op on a +# const, which owns none. Safe only because nothing will launch again - the +# pool may reissue these buffers, while the compiled kernels above still point +# at them (see parameter.py, "Lifetime of a compiled object"). +for param in (g_p, drop_cx_p, drop_cy_p, drop_amp_p): + param.destroy() +for buf in (h0, h1, u, v): + pool.release_data(buf) +print("pooled storage released") diff --git a/examples/grid/ball_walk_cupy.py b/examples/grid/ball_walk_cupy.py new file mode 100644 index 0000000..e473ab5 --- /dev/null +++ b/examples/grid/ball_walk_cupy.py @@ -0,0 +1,250 @@ +""" +One kernel template, four grids: make_grid's boundary/nodata knobs made +visible, cupy backend. Same model as ball_walk_taichi.py; see that file's +docstring for the full explanation of the mechanism. This one differs only in +template syntax (CUDA source text with `$...$` spans, per cupy_backend.py) +and buffer shape (flat, since a CUDA thread indexes its own data). + +Two device templates are written ONCE, below, as CUDA source strings - `walk` +(moves a point one hop) and `spread` (one Jacobi sweep of a geodesic distance +field). Each is ingested by FOUR separate CupyKernelBuilders, one per grid +config, differing only in which grid Bag gets bound under the name `grid`; the +source text itself never changes. Every scalar/field Parameter a span reaches +- CENTRE, SEED, and whatever grid.neighbour_and_distance needs internally - +lands in that compile's own module-scope constant block (see cupy_backend.py), +so the same two device functions, compiled four times against four grids, read +four independent sets of pointers. + +The four grids, one per panel: + 1. boundary="normal", nodata=False - plain bounded grid + 2. boundary="periodic_EW", nodata=False - wraps east/west + 3. boundary="normal", nodata=True + island - an impassable disc + 4. boundary="periodic_EW", nodata=True + island - both at once + +`walk` xorshifts a SEED Parameter (stored as int32, reinterpreted as unsigned +inside the kernel body - cupy's Parameter has no unsigned dtype mapping, see +the module's own `_CTYPE` table, so the cast happens in the template instead +of the storage), turns the top byte into a direction k, and moves CENTRE to +whatever `grid.neighbour_and_distance` reports, only if that neighbour index +is not -1. `spread` does one Jacobi relaxation sweep of the geodesic distance +field: `d_out[i] = min(d_in[i], min_k(d_in[j] + w))` for every (j, w) pair +`neighbour_and_distance` returns through its out-pointers, skipping j == -1. +Reseed the field to 0 at CENTRE and +inf elsewhere every frame, run K sweeps, +and `d < R` is the disc shown per panel - wrapping across a periodic edge or +bending around the nodata island purely because the grid's own neighbour +lookup says so, never because the kernel text does. + +All four panels start from the same CENTRE but each gets its own SEED, so the +four balls are independent random walks rather than one walk replayed four +times. What the panel compares is how each configuration *confines* a ball - +edges that block, edges that wrap, an island that is never a valid target - +not four copies of one trajectory drifting apart and re-converging. + +Author: B.G (07/2026) +""" + +import time + +import cupy as cp +import matplotlib.pyplot as plt +import numpy as np + +from pyfastflow.experimental.core.context.cupy_backend import ( + CupyKernelBuilder, + CupyParameter, +) +from pyfastflow.experimental.grid import make_grid +from pyfastflow.experimental.core.pool.cupy_pool import CupyPool + +# --------------------------------------------------------------------------- +# host-side constants +# --------------------------------------------------------------------------- +NX, NY = 256, 256 +NN = NX * NY +DX = 1.0 + +INF = 1.0e6 # stand-in for +inf in the distance field (avoids float overflow) +R_BALL = 20.0 # display threshold: d < R_BALL is "inside the ball" +K_SWEEPS = 50 # Jacobi sweeps per frame - enough to converge for R_BALL=20 +WALK_STEPS_PER_FRAME = 5 + +BLOCK = 256 +GRID_DIM = (NN + BLOCK - 1) // BLOCK + +START_ROW, START_COL = 50, 50 +START_IDX = START_ROW * NX + START_COL +SEED_VALUE = 20260728 +SEED_STRIDE = 0x9E3779B9 # per-panel seed offset, so each ball walks on its own + +ISLAND_ROW, ISLAND_COL, ISLAND_R = NY // 2, NX // 2, 40 + +pool = CupyPool() + +# --------------------------------------------------------------------------- +# device templates - written once, ingested by four builders below +# --------------------------------------------------------------------------- +WALK_SRC = r""" +extern "C" __global__ void walk(void) { + // one hop of a random walk, entirely on-device. State (SEED, CENTRE) is + // bound per-panel as scalar Parameters; this source never changes. + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx != 0) return; + unsigned int s = (unsigned int)$SEED.get(0)$; + s ^= s << 13; + s ^= s >> 17; + s ^= s << 5; + $SEED.set_node(0, (int)s)$; + int k = (s >> 24) % $grid.n_neighbours.get(0)$; + int c = $CENTRE.get(0)$; + int n; + float w; + $grid.neighbour_and_distance(c, k, &n, &w)$; + if (n >= 0) { + $CENTRE.set_node(0, n)$; + } +} +""" + +SPREAD_SRC = r""" +extern "C" __global__ void spread(float* d_out, const float* d_in) { + // one Jacobi sweep of the geodesic distance field. + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= $grid.nx.get(0)$ * $grid.ny.get(0)$) return; + float best = d_in[i]; + for (int k = 0; k < $grid.n_neighbours.get(0)$; k++) { + int j; + float w; + $grid.neighbour_and_distance(i, k, &j, &w)$; + if (j != -1) { + float cand = d_in[j] + w; + if (cand < best) best = cand; + } + } + d_out[i] = best; +} +""" + +# --------------------------------------------------------------------------- +# nodata island mask (shared shape, independently allocated per grid) +# --------------------------------------------------------------------------- +_rr, _cc = np.mgrid[0:NY, 0:NX] +_island = ((_rr - ISLAND_ROW) ** 2 + (_cc - ISLAND_COL) ** 2) < ISLAND_R**2 +island_mask_flat = _island.astype(np.uint8).ravel() + +# --------------------------------------------------------------------------- +# build the four panels - same two templates, four different grid bindings +# --------------------------------------------------------------------------- +PANEL_CONFIGS = [ + ("normal, no nodata", "normal", False), + ("periodic_EW, no nodata", "periodic_EW", False), + ("normal, nodata island", "normal", True), + ("periodic_EW, nodata island", "periodic_EW", True), +] + +panels = [] +for panel_idx, (title, boundary, nodata) in enumerate(PANEL_CONFIGS): + grid_bag = make_grid( + "cupy", pool, NX, NY, DX, topology="D8", boundary=boundary, nodata=nodata + ) + if nodata: + grid_bag.nodata_mask.set(island_mask_flat) + + centre_p = CupyParameter("CENTRE", dtype=np.int32, mode="scalar", value=START_IDX, pool=pool) + panel_seed = (SEED_VALUE + panel_idx * SEED_STRIDE) & 0x7FFFFFFF + seed_p = CupyParameter("SEED", dtype=np.int32, mode="scalar", value=panel_seed, pool=pool) + + walk_kernel = ( + CupyKernelBuilder() + .bind("SEED", seed_p) + .bind("CENTRE", centre_p) + .bind("grid", grid_bag) + .ingest(WALK_SRC) + .compile() + ) + spread_kernel = CupyKernelBuilder().bind("grid", grid_bag).ingest(SPREAD_SRC).compile() + + d0 = pool.get_data(np.float32, (NN,)) + d1 = pool.get_data(np.float32, (NN,)) + + panels.append( + dict( + title=title, + nodata=nodata, + grid=grid_bag, + centre_p=centre_p, + seed_p=seed_p, + walk_kernel=walk_kernel, + spread_kernel=spread_kernel, + d0=d0, + d1=d1, + ) + ) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +cmap = plt.get_cmap("Blues").copy() +cmap.set_bad("dimgray") + +fig, axes = plt.subplots(2, 2, figsize=(9, 9)) +for panel, ax in zip(panels, axes.ravel()): + im = ax.imshow(np.zeros((NY, NX)), cmap=cmap, vmin=0.0, vmax=1.0) + ax.set_xticks([]) + ax.set_yticks([]) + panel["im"] = im + panel["ax"] = ax +fig.suptitle("Ball walk (Cupy) - one kernel template, four grids") +fig.tight_layout() +fig.show() + +frame = 0 +try: + while True: + t_start = time.perf_counter() + for panel in panels: + for _ in range(WALK_STEPS_PER_FRAME): + panel["walk_kernel"](grid=1, block=1) + + c = int(panel["centre_p"].get().to_numpy()) + seed_arr = np.full(NN, INF, dtype=np.float32) + seed_arr[c] = 0.0 + panel["d0"].from_numpy(seed_arr) + + d0, d1 = panel["d0"], panel["d1"] + for _ in range(K_SWEEPS): + panel["spread_kernel"](d1.data, d0.data, grid=GRID_DIM, block=BLOCK) + d0, d1 = d1, d0 + panel["d0"], panel["d1"] = d0, d1 + + dd = d0.to_numpy().reshape(NY, NX) + disp = (dd < R_BALL).astype(np.float32) + if panel["nodata"]: + disp[_island] = np.nan + panel["im"].set_data(disp) + panel["ax"].set_title(f"{panel['title']}\ncentre row={c // NX} col={c % NX}", fontsize=9) + + cp.cuda.Device().synchronize() # GPU is async; sync before stopping the timer + frame += 1 + frame_ms = (time.perf_counter() - t_start) * 1e3 + if frame % 20 == 0: + print(f"frame {frame}: {frame_ms:6.1f} ms") + + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.05) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +for panel in panels: + panel["centre_p"].destroy() + panel["seed_p"].destroy() + panel["grid"].nx.destroy() + panel["grid"].ny.destroy() + panel["grid"].dx.destroy() + panel["grid"].n_neighbours.destroy() + if panel["nodata"]: + panel["grid"].nodata_mask.destroy() + pool.release_data(panel["d0"]) + pool.release_data(panel["d1"]) diff --git a/examples/grid/ball_walk_quadrants.py b/examples/grid/ball_walk_quadrants.py new file mode 100644 index 0000000..51f6c6b --- /dev/null +++ b/examples/grid/ball_walk_quadrants.py @@ -0,0 +1,247 @@ +""" +One kernel template, four grids: make_grid's boundary/nodata knobs made +visible, Quadrants backend. Same model as ball_walk_taichi.py; see that +file's docstring for the full explanation of the mechanism. This one differs +only in which backend module TaichiKernelBuilder's counterpart wraps - +QuadrantsKernelBuilder compiles to qd.func/qd.kernel instead of ti.func/ +ti.kernel, through the same closure-splicing machinery +(context/_closure_backend.py). + +Two device templates are written ONCE, below, as plain python defs - `walk` +(moves a point one hop) and `spread` (one Jacobi sweep of a geodesic distance +field). Each is ingested by FOUR separate QuadrantsKernelBuilders, one per +grid config, differing only in which grid Bag gets bound under the name +`grid`. Nothing in either template body changes; make_grid's own block +substitution (see grid/_closure_blocks.py, shared verbatim with Taichi) is +what makes `grid.neighbour(i, k)` and `grid.neighbour_and_distance(i, k)` +mean something different per panel. + +The four grids, one per panel: + 1. boundary="normal", nodata=False - plain bounded grid + 2. boundary="periodic_EW", nodata=False - wraps east/west + 3. boundary="normal", nodata=True + island - an impassable disc + 4. boundary="periodic_EW", nodata=True + island - both at once + +The "ball" is a geodesic distance field, computed by Jacobi relaxation: +`spread(d_out, d_in)` sets each node's distance to +`min(d_in[i], min_k(d_in[neighbour(i,k)] + dist_from_k(k)))`, walking the +`neighbour_and_distance` helper's -1 sentinel to skip missing/blocked +neighbours. Reseed the field to 0 at one node and +inf elsewhere, run K +sweeps, and `d < R` is a disc that has flowed outward through the grid's own +notion of adjacency - wrapping across a periodic edge, or bending around a +nodata island, without the template knowing either is happening. + +The disc's centre does a random walk, also entirely on-device: `walk` xorshifts +a u32 SEED Parameter, turns the top byte into a direction k, and moves CENTRE +to `grid.neighbour(centre, k)` only if that is not -1. Since `neighbour()` +already folds in the edge gate and the nodata gate (see _valid_nodata_tmpl in +_closure_blocks.py), a lone `n >= 0` check is sufficient: normal boundaries +block the walk at the domain edge, periodic ones wrap it, and the nodata +island is simply never a valid target. CENTRE and SEED are scalar Parameters, +so the same kernel that mutates them on-device leaves the new value sitting in +their pooled storage for the host to read back (`.get().to_numpy()`) for the +panel title, no extra plumbing required. + +All four panels start from the same CENTRE but each gets its own SEED, so the +four balls are independent random walks rather than one walk replayed four +times. What the panel compares is how each configuration *confines* a ball - +edges that block, edges that wrap, an island that is never a valid target - +not four copies of one trajectory drifting apart and re-converging. + +Author: B.G (07/2026) +""" + +import time + +import matplotlib.pyplot as plt +import numpy as np +import quadrants as qd + +from pyfastflow.experimental.core.context.quadrants_backend import ( + QuadrantsKernelBuilder, + QuadrantsParameter, +) +from pyfastflow.experimental.grid import make_grid +from pyfastflow.experimental.core.pool.quadrants_pool import QuadrantsPool + +qd.init(arch=qd.gpu) + +# --------------------------------------------------------------------------- +# host-side constants +# --------------------------------------------------------------------------- +NX, NY = 256, 256 +NN = NX * NY +DX = 1.0 + +INF = 1.0e6 # stand-in for +inf in the distance field (avoids float overflow) +R_BALL = 20.0 # display threshold: d < R_BALL is "inside the ball" +K_SWEEPS = 50 # Jacobi sweeps per frame - enough to converge for R_BALL=20 +WALK_STEPS_PER_FRAME = 5 + +START_ROW, START_COL = 50, 50 +START_IDX = START_ROW * NX + START_COL +SEED_VALUE = 20260728 +SEED_STRIDE = 0x9E3779B9 # per-panel seed offset, so each ball walks on its own + +ISLAND_ROW, ISLAND_COL, ISLAND_R = NY // 2, NX // 2, 40 + +pool = QuadrantsPool() + +# --------------------------------------------------------------------------- +# device templates - written once, ingested by four builders below +# --------------------------------------------------------------------------- + + +def walk_template(): + # one hop of a random walk, entirely on-device. State (SEED, CENTRE) is + # bound per-panel as scalar Parameters; the body never changes. + for _dummy in range(1): + s = SEED.get(0) + s = s ^ (s << 13) + s = s ^ (s >> 17) + s = s ^ (s << 5) + SEED.set_node(0, s) + k = (s >> 24) % grid.n_neighbours.get(0) + c = CENTRE.get(0) + n = grid.neighbour(c, k) + if n >= 0: + CENTRE.set_node(0, n) + + +def spread_template(d_out: qd.template(), d_in: qd.template()): + # one Jacobi sweep of the geodesic distance field. + for i in d_in: + best = d_in[i] + for k in range(grid.n_neighbours.get(0)): + j, w = grid.neighbour_and_distance(i, k) + if j != -1: + cand = d_in[j] + w + if cand < best: + best = cand + d_out[i] = best + + +# --------------------------------------------------------------------------- +# nodata island mask (shared shape, independently allocated per grid) +# --------------------------------------------------------------------------- +_rr, _cc = np.mgrid[0:NY, 0:NX] +_island = ((_rr - ISLAND_ROW) ** 2 + (_cc - ISLAND_COL) ** 2) < ISLAND_R**2 +island_mask_flat = _island.astype(np.uint8).ravel() + +# --------------------------------------------------------------------------- +# build the four panels - same two templates, four different grid bindings +# --------------------------------------------------------------------------- +PANEL_CONFIGS = [ + ("normal, no nodata", "normal", False), + ("periodic_EW, no nodata", "periodic_EW", False), + ("normal, nodata island", "normal", True), + ("periodic_EW, nodata island", "periodic_EW", True), +] + +panels = [] +for panel_idx, (title, boundary, nodata) in enumerate(PANEL_CONFIGS): + grid_bag = make_grid( + "quadrants", pool, NX, NY, DX, topology="D8", boundary=boundary, nodata=nodata + ) + if nodata: + grid_bag.nodata_mask.set(island_mask_flat) + + centre_p = QuadrantsParameter("CENTRE", dtype=qd.i32, mode="scalar", value=START_IDX, pool=pool) + panel_seed = (SEED_VALUE + panel_idx * SEED_STRIDE) & 0x7FFFFFFF + seed_p = QuadrantsParameter("SEED", dtype=qd.u32, mode="scalar", value=panel_seed, pool=pool) + + walk_kernel = ( + QuadrantsKernelBuilder() + .bind("SEED", seed_p) + .bind("CENTRE", centre_p) + .bind("grid", grid_bag) + .ingest(walk_template) + .compile() + ) + spread_kernel = QuadrantsKernelBuilder().bind("grid", grid_bag).ingest(spread_template).compile() + + d0 = pool.get_data(qd.f32, (NN,)) + d1 = pool.get_data(qd.f32, (NN,)) + + panels.append( + dict( + title=title, + nodata=nodata, + grid=grid_bag, + centre_p=centre_p, + seed_p=seed_p, + walk_kernel=walk_kernel, + spread_kernel=spread_kernel, + d0=d0, + d1=d1, + ) + ) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +cmap = plt.get_cmap("Blues").copy() +cmap.set_bad("dimgray") + +fig, axes = plt.subplots(2, 2, figsize=(9, 9)) +for panel, ax in zip(panels, axes.ravel()): + im = ax.imshow(np.zeros((NY, NX)), cmap=cmap, vmin=0.0, vmax=1.0) + ax.set_xticks([]) + ax.set_yticks([]) + panel["im"] = im + panel["ax"] = ax +fig.suptitle("Ball walk (Quadrants) - one kernel template, four grids") +fig.tight_layout() +fig.show() + +frame = 0 +try: + while True: + t_start = time.perf_counter() + for panel in panels: + for _ in range(WALK_STEPS_PER_FRAME): + panel["walk_kernel"]() + + c = int(panel["centre_p"].get().to_numpy()) + seed_arr = np.full(NN, INF, dtype=np.float32) + seed_arr[c] = 0.0 + panel["d0"].from_numpy(seed_arr) + + d0, d1 = panel["d0"], panel["d1"] + for _ in range(K_SWEEPS): + panel["spread_kernel"](d1.data, d0.data) + d0, d1 = d1, d0 + panel["d0"], panel["d1"] = d0, d1 + + dd = d0.to_numpy().reshape(NY, NX) + disp = (dd < R_BALL).astype(np.float32) + if panel["nodata"]: + disp[_island] = np.nan + panel["im"].set_data(disp) + panel["ax"].set_title(f"{panel['title']}\ncentre row={c // NX} col={c % NX}", fontsize=9) + + qd.sync() + frame += 1 + frame_ms = (time.perf_counter() - t_start) * 1e3 + if frame % 20 == 0: + print(f"frame {frame}: {frame_ms:6.1f} ms") + + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.05) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +for panel in panels: + panel["centre_p"].destroy() + panel["seed_p"].destroy() + panel["grid"].nx.destroy() + panel["grid"].ny.destroy() + panel["grid"].dx.destroy() + panel["grid"].n_neighbours.destroy() + if panel["nodata"]: + panel["grid"].nodata_mask.destroy() + pool.release_data(panel["d0"]) + pool.release_data(panel["d1"]) diff --git a/examples/grid/ball_walk_taichi.py b/examples/grid/ball_walk_taichi.py new file mode 100644 index 0000000..a2b3125 --- /dev/null +++ b/examples/grid/ball_walk_taichi.py @@ -0,0 +1,243 @@ +""" +One kernel template, four grids: make_grid's boundary/nodata knobs made +visible, Taichi backend. + +Two device templates are written ONCE, below, as plain python defs - `walk` +(moves a point one hop) and `spread` (one Jacobi sweep of a geodesic distance +field). Each is ingested by FOUR separate TaichiKernelBuilders, one per grid +config, differing only in which grid Bag gets bound under the name `grid`. +Nothing in either template body changes; make_grid's own block substitution +(see grid/_closure_blocks.py) is what makes `grid.neighbour(i, k)` and +`grid.neighbour_and_distance(i, k)` mean something different per panel. + +The four grids, one per panel: + 1. boundary="normal", nodata=False - plain bounded grid + 2. boundary="periodic_EW", nodata=False - wraps east/west + 3. boundary="normal", nodata=True + island - an impassable disc + 4. boundary="periodic_EW", nodata=True + island - both at once + +The "ball" is a geodesic distance field, computed by Jacobi relaxation: +`spread(d_out, d_in)` sets each node's distance to +`min(d_in[i], min_k(d_in[neighbour(i,k)] + dist_from_k(k)))`, walking the +`neighbour_and_distance` helper's -1 sentinel to skip missing/blocked +neighbours. Reseed the field to 0 at one node and +inf elsewhere, run K +sweeps, and `d < R` is a disc that has flowed outward through the grid's own +notion of adjacency - wrapping across a periodic edge, or bending around a +nodata island, without the template knowing either is happening. + +The disc's centre does a random walk, also entirely on-device: `walk` xorshifts +a u32 SEED Parameter, turns the top byte into a direction k, and moves CENTRE +to `grid.neighbour(centre, k)` only if that is not -1. Since `neighbour()` +already folds in the edge gate and the nodata gate (see _valid_nodata_tmpl in +_closure_blocks.py), a lone `n >= 0` check is sufficient: normal boundaries +block the walk at the domain edge, periodic ones wrap it, and the nodata +island is simply never a valid target. CENTRE and SEED are scalar Parameters, +so the same kernel that mutates them on-device leaves the new value sitting in +their pooled storage for the host to read back (`.get().to_numpy()`) for the +panel title, no extra plumbing required. + +All four panels start from the same CENTRE but each gets its own SEED, so the +four balls are independent random walks rather than one walk replayed four +times. What the panel compares is how each configuration *confines* a ball - +edges that block, edges that wrap, an island that is never a valid target - +not four copies of one trajectory drifting apart and re-converging. + +Author: B.G (07/2026) +""" + +import time + +import matplotlib.pyplot as plt +import numpy as np +import taichi as ti + +from pyfastflow.experimental.core.context.taichi_backend import ( + TaichiKernelBuilder, + TaichiParameter, +) +from pyfastflow.experimental.grid import make_grid +from pyfastflow.experimental.core.pool.taichi_pool import TaichiPool + +ti.init(arch=ti.gpu) + +# --------------------------------------------------------------------------- +# host-side constants +# --------------------------------------------------------------------------- +NX, NY = 256, 256 +NN = NX * NY +DX = 1.0 + +INF = 1.0e6 # stand-in for +inf in the distance field (avoids float overflow) +R_BALL = 20.0 # display threshold: d < R_BALL is "inside the ball" +K_SWEEPS = 50 # Jacobi sweeps per frame - enough to converge for R_BALL=20 +WALK_STEPS_PER_FRAME = 50 + +START_ROW, START_COL = 50, 50 +START_IDX = START_ROW * NX + START_COL +SEED_VALUE = 20260728 +SEED_STRIDE = 0x9E3779B9 # per-panel seed offset, so each ball walks on its own + +ISLAND_ROW, ISLAND_COL, ISLAND_R = NY // 2, NX // 2, 40 + +pool = TaichiPool() + +# --------------------------------------------------------------------------- +# device templates - written once, ingested by four builders below +# --------------------------------------------------------------------------- + + +def walk_template(): + # one hop of a random walk, entirely on-device. State (SEED, CENTRE) is + # bound per-panel as scalar Parameters; the body never changes. + for _dummy in range(1): + s = SEED.get(0) + s = s ^ (s << 13) + s = s ^ (s >> 17) + s = s ^ (s << 5) + SEED.set_node(0, s) + k = (s >> 24) % grid.n_neighbours.get(0) + c = CENTRE.get(0) + n = grid.neighbour(c, k) + if n >= 0: + CENTRE.set_node(0, n) + + +def spread_template(d_out: ti.template(), d_in: ti.template()): + # one Jacobi sweep of the geodesic distance field. + for i in d_in: + best = d_in[i] + for k in range(grid.n_neighbours.get(0)): + j, w = grid.neighbour_and_distance(i, k) + if j != -1: + cand = d_in[j] + w + if cand < best: + best = cand + d_out[i] = best + + +# --------------------------------------------------------------------------- +# nodata island mask (shared shape, independently allocated per grid) +# --------------------------------------------------------------------------- +_rr, _cc = np.mgrid[0:NY, 0:NX] +_island = ((_rr - ISLAND_ROW) ** 2 + (_cc - ISLAND_COL) ** 2) < ISLAND_R**2 +island_mask_flat = _island.astype(np.uint8).ravel() + +# --------------------------------------------------------------------------- +# build the four panels - same two templates, four different grid bindings +# --------------------------------------------------------------------------- +PANEL_CONFIGS = [ + ("normal, no nodata", "normal", False), + ("periodic_EW, no nodata", "periodic_EW", False), + ("normal, nodata island", "normal", True), + ("periodic_EW, nodata island", "periodic_EW", True), +] + +panels = [] +for panel_idx, (title, boundary, nodata) in enumerate(PANEL_CONFIGS): + grid_bag = make_grid( + "taichi", pool, NX, NY, DX, topology="D8", boundary=boundary, nodata=nodata + ) + if nodata: + grid_bag.nodata_mask.set(island_mask_flat) + + centre_p = TaichiParameter("CENTRE", dtype=ti.i32, mode="scalar", value=START_IDX, pool=pool) + panel_seed = (SEED_VALUE + panel_idx * SEED_STRIDE) & 0x7FFFFFFF + seed_p = TaichiParameter("SEED", dtype=ti.u32, mode="scalar", value=panel_seed, pool=pool) + + walk_kernel = ( + TaichiKernelBuilder() + .bind("SEED", seed_p) + .bind("CENTRE", centre_p) + .bind("grid", grid_bag) + .ingest(walk_template) + .compile() + ) + spread_kernel = TaichiKernelBuilder().bind("grid", grid_bag).ingest(spread_template).compile() + + d0 = pool.get_data(ti.f32, (NN,)) + d1 = pool.get_data(ti.f32, (NN,)) + + panels.append( + dict( + title=title, + nodata=nodata, + grid=grid_bag, + centre_p=centre_p, + seed_p=seed_p, + walk_kernel=walk_kernel, + spread_kernel=spread_kernel, + d0=d0, + d1=d1, + ) + ) + +# --------------------------------------------------------------------------- +# live view +# --------------------------------------------------------------------------- +cmap = plt.get_cmap("Blues").copy() +cmap.set_bad("dimgray") + +fig, axes = plt.subplots(2, 2, figsize=(9, 9)) +ims = [] +for panel, ax in zip(panels, axes.ravel()): + im = ax.imshow(np.zeros((NY, NX)), cmap=cmap, vmin=0.0, vmax=1.0) + ax.set_xticks([]) + ax.set_yticks([]) + panel["im"] = im + panel["ax"] = ax + ims.append(im) +fig.suptitle("Ball walk (Taichi) - one kernel template, four grids") +fig.tight_layout() +fig.show() + +frame = 0 +try: + while True: + t_start = time.perf_counter() + for panel in panels: + for _ in range(WALK_STEPS_PER_FRAME): + panel["walk_kernel"]() + + c = int(panel["centre_p"].get().to_numpy()) + seed_arr = np.full(NN, INF, dtype=np.float32) + seed_arr[c] = 0.0 + panel["d0"].from_numpy(seed_arr) + + d0, d1 = panel["d0"], panel["d1"] + for _ in range(K_SWEEPS): + panel["spread_kernel"](d1.data, d0.data) + d0, d1 = d1, d0 + panel["d0"], panel["d1"] = d0, d1 + + dd = d0.to_numpy().reshape(NY, NX) + disp = (dd < R_BALL).astype(np.float32) + if panel["nodata"]: + disp[_island] = np.nan + panel["im"].set_data(disp) + panel["ax"].set_title(f"{panel['title']}\ncentre row={c // NX} col={c % NX}", fontsize=9) + + ti.sync() + frame += 1 + frame_ms = (time.perf_counter() - t_start) * 1e3 + if frame % 20 == 0: + print(f"frame {frame}: {frame_ms:6.1f} ms") + + fig.canvas.draw_idle() + fig.canvas.start_event_loop(0.05) +except KeyboardInterrupt: + pass + +# --------------------------------------------------------------------------- +# teardown +# --------------------------------------------------------------------------- +for panel in panels: + panel["centre_p"].destroy() + panel["seed_p"].destroy() + panel["grid"].nx.destroy() + panel["grid"].ny.destroy() + panel["grid"].dx.destroy() + panel["grid"].n_neighbours.destroy() + if panel["nodata"]: + panel["grid"].nodata_mask.destroy() + pool.release_data(panel["d0"]) + pool.release_data(panel["d1"]) diff --git a/pyfastflow/experimental/__init__.py b/pyfastflow/experimental/__init__.py new file mode 100644 index 0000000..a8f26f5 --- /dev/null +++ b/pyfastflow/experimental/__init__.py @@ -0,0 +1,5 @@ +""" +Namespace for the in-progress multi-backend refactor. + +Author: B.G (07/2026) +""" diff --git a/pyfastflow/experimental/core/__init__.py b/pyfastflow/experimental/core/__init__.py new file mode 100644 index 0000000..0a2cdcd --- /dev/null +++ b/pyfastflow/experimental/core/__init__.py @@ -0,0 +1,5 @@ +""" +New backend-agnostic core (Parameter/Helper/Kernel/Pool ABCs + backends). + +Author: B.G (07/2026) +""" diff --git a/pyfastflow/experimental/core/context/__init__.py b/pyfastflow/experimental/core/context/__init__.py new file mode 100644 index 0000000..0eca8ba --- /dev/null +++ b/pyfastflow/experimental/core/context/__init__.py @@ -0,0 +1,7 @@ +""" +New backend-agnostic context architecture (Parameter/Specializable ABCs + +backends, plus RoutineBuilder/Routine for a linear sequence of kernels +sharing one bag). + +Author: B.G (07/2026) +""" diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py new file mode 100644 index 0000000..5f58d0d --- /dev/null +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -0,0 +1,330 @@ +""" +Machinery shared by the two backends whose templates are python functions: +Taichi and Quadrants. + +Specialization works by rebuilding the template function around a globals dict +that carries the bound objects, so a name like `phys` in the template body +resolves to the bound object when the backend traces it. The rebuilt function +is then decorated with ti.func/qd.func or ti.kernel/qd.kernel. + +The two backends can share all of this because the pieces used here - func, +kernel, static, u8, i32, i64 - carry the same names and the same behaviour in +both modules. A backend subclass therefore only pins `_backend` to the ti or qd +module; nothing else varies. + +What lives here is only what a Parameter's device view (ClosureBackendParameter, +_build_device_view) needs to compile its own tiny get/set_node funcs - +specialize_closure and the two supporting classes. The kernel/helper/routine +compile path for Taichi/Quadrants is compile_closure.py, which composes a +BoundKernel's `ctx` tree instead of splicing bound objects into template +globals - see its own module docstring for why. + +cupy does not appear here: CUDA source text has no globals to patch, and that +backend substitutes into the source directly instead. + +Author: B.G (07/2026) +""" + +from types import FunctionType +from typing import Any, ClassVar + +import numpy as np + +from .parameter import MODES, Parameter + + +def specialize_closure(template, globals_: dict[str, Any]) -> FunctionType: + """ + Rebuild `template` as a new function whose globals carry `globals_`, + leaving the original untouched. + + The code object is reused as-is; only the globals differ, which is what + makes a name in the template body resolve to a bound object. Defaults, + annotations and the rest are copied over so the result still introspects + like the template it came from. + + Parameters + ---------- + template : FunctionType + Function to rebuild. + globals_ : dict[str, Any] + Names to inject into the rebuilt function's globals. + + Returns + ------- + FunctionType + A new function sharing `template`'s code object but with `globals_` + merged into its globals. + + Author: B.G (07/2026) + """ + source = getattr(template, "__wrapped__", template) + func_globals = dict(source.__globals__) + func_globals.update(globals_) + + specialised = FunctionType( + source.__code__, + func_globals, + source.__name__, + source.__defaults__, + source.__closure__, + ) + specialised.__kwdefaults__ = source.__kwdefaults__ + specialised.__annotations__ = dict(source.__annotations__) + specialised.__doc__ = source.__doc__ + specialised.__qualname__ = source.__qualname__ + return specialised + + +class ClosureParamDeviceView: + """ + What a Parameter looks like from inside device code. + + `.get` and `.set_node` are compiled device funcs, so a template body reads + `p.get(i)` and writes `p.set_node(i, v)` the same way whatever the + parameter's mode. A const parameter is read-only and carries no `.set_node` + at all, which turns a write to one into a trace-time error. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, get_fn, set_fn=None): + self._name = name + self.get = get_fn + if set_fn is not None: + self.set_node = set_fn + + +class ClosureBackendParameter(Parameter): + """ + Parameter backed by a const python value or by pooled device storage. + + Concrete backends subclass this and pin `_backend` to their module; the + dtype mapping and the device view are written once here against the names + both modules share. + + Author: B.G (07/2026) + """ + + _backend: ClassVar[Any] + + def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None): + """ + Declare one parameter and give it its initial value. + + scalar and field take pooled storage straight away; const stays a + plain python value. + + Parameters + ---------- + name : str + dtype : ti.* or qd.* dtype + mode : str + One of MODES ("const", "scalar", "field"). + value : Any + Initial value. + pool : Pool + Device-buffer pool backing scalar/field storage. + n_flat : int, optional + Node count, required for field mode. + + Raises + ------ + ValueError + If `mode` is not in MODES, or field mode is given without + `n_flat`. + + Author: B.G (07/2026) + """ + if mode not in MODES: + raise ValueError(f"{name}: mode must be one of {sorted(MODES)}, got {mode!r}") + + super().__init__() + self.name = name + self.dtype = dtype + self.mode = mode + self._pool = pool + self._const_value: Any = None + self._handle = None + self._device_view: "ClosureParamDeviceView | None" = None + + if mode == "scalar": + self._handle = pool.get_data(dtype, ()) + elif mode == "field": + if n_flat is None: + raise ValueError(f"{name}: field mode requires n_flat") + self._handle = pool.get_data(dtype, (n_flat,)) + + self._store(value) + + @classmethod + def _numpy_dtype(cls, dtype): + """ + Map a backend dtype (`ti.*`/`qd.*`) to the numpy dtype used for + host-side (de)serialization. + + Author: B.G (07/2026) + """ + backend = cls._backend + if dtype == backend.u8: + return np.uint8 + if dtype == backend.i32: + return np.int32 + if dtype == backend.i64: + return np.int64 + return np.float32 + + def get(self): + """ + The python value for const mode, the backing DataHandle otherwise. + + Author: B.G (07/2026) + """ + return self._const_value if self.mode == "const" else self._handle + + def set(self, value) -> None: + """ + Overwrite the whole value: a device write for scalar, a full + host->device copy for field. const is immutable - see Parameter.set. + + Raises + ------ + ValueError + If this parameter's mode is const. + + Author: B.G (07/2026) + """ + if self.mode == "const": + raise ValueError( + f"{self.name}: const parameter is immutable; build a new Parameter and " + f"replace() it into the bag, then recompile" + ) + self._store(value) + + def _store(self, value) -> None: + """ + Write `value` according to the mode, with no immutability check - the + one path that may set a const, used by __init__ to place its initial + value. + + Author: B.G (07/2026) + """ + if self.mode == "const": + self._const_value = self._numpy_dtype(self.dtype)(value).item() + elif self.mode == "scalar": + self._handle.data[None] = value + else: # field + arr = np.asarray(value, dtype=self._numpy_dtype(self.dtype)).reshape(-1) + self._handle.data.from_numpy(arr) + + def set_node(self, node, value) -> None: + """ + Host-side single-cell write. scalar ignores node; const is read-only. + + Raises + ------ + ValueError + If this parameter's mode is const. + + Author: B.G (07/2026) + """ + if self.mode == "const": + raise ValueError(f"{self.name}: const parameter is read-only") + if self.mode == "scalar": + self._handle.data[None] = value + else: # field + self._handle.data[node] = value + + def read(self): + """ + Host-side scalar read - see Parameter.read for the contract. + + Raises + ------ + ValueError + If this parameter's mode is field. + + Author: B.G (07/2026) + """ + if self.mode == "const": + return self._const_value + if self.mode == "field": + raise ValueError( + f"{self.name}: read() is for scalar/const only; a field is not meant to be " + f"read back to the host as a whole" + ) + return self._numpy_dtype(self.dtype)(self._handle.data.to_numpy()).item() + + def destroy(self) -> None: + """ + Return any pooled storage to the pool. const mode owns none, so this + is a no-op there. + + Author: B.G (07/2026) + """ + if self._handle is not None: + self._pool.release_data(self._handle) + self._handle = None + self._device_view = None # a cached view closes over the released handle + + def device_view(self) -> ClosureParamDeviceView: + """ + This parameter's device accessor, built on first use and kept. + + The compiled funcs come out identical every time, so one view serves + every kernel that binds this parameter, for the parameter's whole + life. A const's literal is fixed at construction, and a scalar or + field set() writes through the very storage the view already reads, so + neither can stale it. Only destroy() drops the view, having released + that storage - and that does not reach kernels compiled earlier, which + still hold it (see parameter.py, "Lifetime of a compiled object"). + + Author: B.G (07/2026) + """ + if self._device_view is None: + self._device_view = self._build_device_view() + return self._device_view + + def _build_device_view(self) -> ClosureParamDeviceView: + """ + Compile this parameter's device accessors as backend funcs. + + get(node) branches on the mode through `_backend.static`, which + resolves at trace time, so only one arm survives into the generated + code: a baked literal for const, HANDLE[None] for scalar, HANDLE[node] + for field. set_node is built for scalar and field only. MODE, VALUE and + HANDLE are ordinary python values spliced in as globals. + + Author: B.G (07/2026) + """ + backend = self._backend + mode = self.mode + value = self._const_value + handle = self._handle.data if self._handle is not None else None + + def get_template(node): + if STATIC(MODE == "const"): + return VALUE + elif STATIC(MODE == "scalar"): + return HANDLE[None] + else: + return HANDLE[node] + + get_fn = backend.func( + specialize_closure(get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle, "STATIC": backend.static}) + ) + + set_fn = None + if mode != "const": + + def set_node_template(node, val): + if STATIC(MODE == "scalar"): + HANDLE[None] = val + else: + HANDLE[node] = val + + set_fn = backend.func( + specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle, "STATIC": backend.static}) + ) + + return ClosureParamDeviceView(self.name, get_fn, set_fn) diff --git a/pyfastflow/experimental/core/context/backends.py b/pyfastflow/experimental/core/context/backends.py new file mode 100644 index 0000000..10e7e7f --- /dev/null +++ b/pyfastflow/experimental/core/context/backends.py @@ -0,0 +1,72 @@ +""" +Per-backend wiring shared by every factory (make_grid, make_noise, ...). + +A factory needs two things to build its Parameters against a chosen backend +name: the backend module itself (for a cupy-only helper that needs e.g. +`np.int32`, or a closure block that needs `ti`/`qd` directly), and the +backend's own dtype objects keyed by the short names factories write their +Parameters with ("i32", "i64", "f32", "u8", "u32"). backend_classes() is the +one place that knows the mapping from a backend name to those things, so a +factory does not carry its own copy of the same if-ladder. + +Picking which private block module implements a factory's device code +("_closure_blocks" vs "_cupy_blocks") stays the caller's job - a factory owns +its own blocks, this module does not know they exist. + +Author: B.G (07/2026) +""" + +import numpy as np + + +def backend_classes(backend: str): + """ + Look up the module, Parameter subclass and dtype table for one backend. + + Parameters + ---------- + backend : str + "taichi", "quadrants" or "cupy". + + Returns + ------- + module : module or None + `ti`/`qd` for the closure backends, `None` for cupy - cupy blocks + call plain C, never a bound backend module. + ParameterCls : type + The backend's Parameter subclass. + unused : None + Reserved, always None - kept so `_, ParamCls, _, dtypes = + backend_classes(backend)` call sites stay stable. + dtypes : dict + Maps "i32"/"i64"/"f32"/"u8"/"u32" to that backend's own dtype + objects (ti.*/qd.* for the closure backends, numpy dtypes for cupy). + + No blocks module is returned - each factory (grid, noise, ...) has its + own private block module and picks it itself. + + Author: B.G (07/2026) + """ + if backend == "taichi": + import taichi as ti + + from .taichi_backend import TaichiParameter + + return ti, TaichiParameter, None, { + "i32": ti.i32, "i64": ti.i64, "f32": ti.f32, "u8": ti.u8, "u32": ti.u32, + } + if backend == "quadrants": + import quadrants as qd + + from .quadrants_backend import QuadrantsParameter + + return qd, QuadrantsParameter, None, { + "i32": qd.i32, "i64": qd.i64, "f32": qd.f32, "u8": qd.u8, "u32": qd.u32, + } + if backend == "cupy": + from .cupy_backend import CupyParameter + + return None, CupyParameter, None, { + "i32": np.int32, "i64": np.int64, "f32": np.float32, "u8": np.uint8, "u32": np.uint32, + } + raise ValueError(f"unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") diff --git a/pyfastflow/experimental/core/context/bag.py b/pyfastflow/experimental/core/context/bag.py new file mode 100644 index 0000000..2912478 --- /dev/null +++ b/pyfastflow/experimental/core/context/bag.py @@ -0,0 +1,478 @@ +""" +Bag: a named collection of anything a template might bind, plus the operators +that reshape one. + +A Bag is a container and nothing more. It never inspects what it holds and has +no notion of backend, mode or compilation - each member is resolved on its own +type by whatever consumes the Bag. That is why this module stands on its own: +it depends on nothing else here beyond the shared uid counter. + +merge/extract/trim/replace all return a fresh Bag holding the very +same member objects - no device storage is ever copied, and the same Parameter +reachable from two bags is one Parameter. check_handles is the guard against +the one thing that aliasing can get wrong: a single name meaning two different +objects across the units of one compile. + +Author: B.G (07/2026) +""" + +from typing import Any + +from ..pool.base import new_uid + + +class Bag: + """ + A named collection that can be handed to a builder in one go. + + A bag holds whatever a template might want to reach under one name - + Parameters, Helpers, further Bags, plain python values - mixed + freely. Nothing dispatches on what a bag contains: each member is resolved + on its own type when the template is specialized, so a bag grouping a + quantity with the helpers that act on it works exactly like one holding + parameters alone. + + Bind it whole - bind("grid", bag) - and reach its members by dotted path + in the template body (grid.nx.get(i), grid.nbr(i)); or, at the + RoutineBuilder/SequenceBuilder layer (routine.py/sequence.py), + bind_bag(bag) sets the one bag every step/block is rebound against at + compile time. + + Build it, grow it, bind it. There is no removal or reassignment: to change + the contents, build another bag. + + Author: B.G (07/2026) + """ + + def __init__(self, items: dict[str, Any] | None = None): + self._uid = new_uid() + self._items: dict[str, Any] = {} + for name, item in (items or {}).items(): + self.add(name, item) + + @property + def uid(self) -> int: + """ + Process-wide identity assigned at construction, from the same counter + as Parameters, Helpers and pool data handles. See Parameter.uid. + + Author: B.G (07/2026) + """ + return self._uid + + def add(self, name: str, item: Any) -> None: + """ + Register `item` under `name`. + + Parameters + ---------- + name : str + Key to register `item` under. Must not already be taken. + item : Any + Value to store - a Parameter, Helper, nested Bag or plain value. + + Raises + ------ + KeyError + If `name` is already registered. + + Author: B.G (07/2026) + """ + if name in self._items: + raise KeyError(f"'{name}' is already registered in this bag") + self._items[name] = item + + def __getattr__(self, name: str) -> Any: + try: + return self._items[name] + except KeyError: + raise AttributeError(name) + + def __getitem__(self, name: str) -> Any: + return self._items[name] + + def __contains__(self, name: str) -> bool: + return name in self._items + + def __iter__(self): + return iter(self._items) + + def items(self): + return self._items.items() + + def __repr__(self) -> str: + """ + Every member on its own line at its dotted path, nested Bags shown as + the subtree they head rather than as an opaque entry. + + Bags are routinely built by merging several others, at which point the + only reliable way to see what one holds is to read it out; this is + that. Each leaf is labelled by what it is - a Parameter by mode and + dtype, anything else by its class - and by its uid, which is what + makes an alias visible: one object reached under two names shows the + same uid twice. + + Author: B.G (07/2026) + """ + lines = [f"Bag(uid={self._uid})"] + for handle, obj in self.walk(): + if isinstance(obj, Bag): + lines.append(f" {handle}/") + continue + mode = getattr(obj, "mode", None) + if mode is not None: + what = f"{mode} {getattr(obj, 'dtype', '?')}" + else: + what = type(obj).__name__ + uid = _uid_of(obj) + lines.append(f" {handle}: {what}" + (f" [uid {uid}]" if uid is not None else "")) + return "\n".join(lines) + + def walk(self, prefix: str = ""): + """ + Yield (dotted_handle, obj) for every member, descending into nested + Bags depth-first. + + A nested Bag produces two things: an entry for the Bag itself, at its + own dotted path, then one entry per member underneath it. So + `Bag({"at": Bag({"i": p1, "j": p2}), "r": p3})` walks as + `("at", )`, `("at.i", p1)`, `("at.j", p2)`, `("r", p3)` - the + parent Bag's entry always precedes its members'. + + Parameters + ---------- + prefix : str, optional + Dotted path prepended to every yielded handle. Used internally + for recursion; callers normally leave it at "". + + Returns + ------- + Iterator[tuple[str, Any]] + (dotted_handle, obj) pairs in depth-first order. + + Author: B.G (07/2026) + """ + for name, item in self._items.items(): + handle = f"{prefix}.{name}" if prefix else name + if isinstance(item, Bag): + yield handle, item + yield from item.walk(handle) + else: + yield handle, item + + +def _uid_of(obj: Any) -> int | None: + """ + An object's uid if it has one, else None. + + Handles bound without a uid (plain python values, unwrapped bindings) are + simply skipped by check_handles rather than treated as a conflict. + + Parameters + ---------- + obj : Any + Object to inspect. + + Returns + ------- + int or None + `obj.uid` if it is an int, else None. + + Author: B.G (07/2026) + """ + uid = getattr(obj, "uid", None) + return uid if isinstance(uid, int) else None + + +def check_handles(units: dict[str, dict[str, Any]]) -> None: + """ + Verify that a handle means the same object everywhere it is used. + + `units` maps a unit name (a kernel, a routine step - whatever the caller + is checking) to that unit's own {handle: obj} map, typically built from + Bag.walk(). Across every unit given, the same handle string must resolve + to objects sharing one uid; if two units bind the same handle to objects + with different uids, this raises naming the handle and both owning units. + + The converse is fine and common: two different handles pointing at the + same uid (an alias, or one Parameter reused under two names) is not a + conflict and is not reported. + + Objects with no `uid` attribute are ignored - there is nothing to compare. + + Parameters + ---------- + units : dict[str, dict[str, Any]] + Unit name -> {handle: obj} map, typically each built from a Bag's + `.walk()`. + + Raises + ------ + ValueError + If the same handle resolves to objects with different uids in two + units, naming the handle and both owning units. + + Author: B.G (07/2026) + """ + seen: dict[str, tuple[int, str]] = {} + for unit_name, handles in units.items(): + for handle, obj in handles.items(): + uid = _uid_of(obj) + if uid is None: + continue + prior = seen.get(handle) + if prior is None: + seen[handle] = (uid, unit_name) + elif prior[0] != uid: + raise ValueError( + f"handle '{handle}' is bound to different objects: " + f"uid {prior[0]} in '{prior[1]}' vs uid {uid} in '{unit_name}'" + ) + + +def _resolve_path(bag: "Bag", path: str) -> Any: + """ + Walk a dotted path through nested Bags and return what it names. + + Parameters + ---------- + bag : Bag + Bag to resolve `path` in. + path : str + Dotted path, e.g. "at.i". + + Returns + ------- + Any + The object named by `path`. + + Raises + ------ + KeyError + If any segment is missing or a non-terminal segment does not resolve + to a Bag, naming the exact prefix that failed. + + Author: B.G (07/2026) + """ + obj = bag + parts = path.split(".") + for depth, part in enumerate(parts): + if not isinstance(obj, Bag) or part not in obj: + failed = ".".join(parts[: depth + 1]) + raise KeyError(f"'{path}' not found in bag (no '{failed}')") + obj = obj[part] + return obj + + +def merge(*bags: "Bag") -> "Bag": + """ + Union of every member across `bags`, into one new Bag. + + Members are taken in argument order; nesting is kept rather than + flattened, so where two bags carry a Bag under the same name, those two + are merged recursively instead of one replacing the other. + + A same-name collision between two non-Bag members is allowed silently + when both share a uid - the same object reached through two bags - and + raises when they don't, naming the member and both uids. A collision + where either side has no uid (a plain python value) cannot be resolved + this way and always raises, since there is nothing to compare. + + No input bag is read from twice or mutated; the result is a fresh Bag. + + Parameters + ---------- + *bags : Bag + Bags to union, in precedence order. + + Returns + ------- + Bag + Fresh Bag holding the union of all members. + + Raises + ------ + ValueError + If a name collides between bags and cannot be proven to be the same + object. + + Author: B.G (07/2026) + """ + merged: dict[str, Any] = {} + for bag in bags: + for name, item in bag.items(): + if name not in merged: + merged[name] = item + continue + existing = merged[name] + if isinstance(existing, Bag) and isinstance(item, Bag): + merged[name] = merge(existing, item) + continue + euid, iuid = _uid_of(existing), _uid_of(item) + if euid is None or iuid is None: + raise ValueError( + f"merge: '{name}' collides between bags and at least one side has " + f"no uid to compare, so they cannot be proven to be the same object" + ) + if euid != iuid: + raise ValueError(f"merge: '{name}' collides between bags: uid {euid} vs uid {iuid}") + return Bag(merged) + + +def extract(bag: "Bag", names) -> "Bag": + """ + A new Bag holding just the named members of `bag`. + + Each entry in `names` may be a plain name or a dotted path + (`"stove.at.i"`); a dotted path is resolved through nested Bags and + reconstructed as nesting in the result, so extracting `"at.i"` and + `"at.j"` yields a result with an `at` sub-bag holding `i` and `j`, not + two flat members. + + Parameters + ---------- + bag : Bag + Bag to extract from. + names : Iterable[str] + Plain names or dotted paths to keep. + + Returns + ------- + Bag + Fresh Bag holding just the named members, with nesting rebuilt. + + Raises + ------ + KeyError + If any path does not resolve in `bag`. + + Author: B.G (07/2026) + """ + tree: dict[str, Any] = {} + for path in names: + resolved = _resolve_path(bag, path) + parts = path.split(".") + cursor = tree + for part in parts[:-1]: + cursor = cursor.setdefault(part, {}) + cursor[parts[-1]] = resolved + return _tree_to_bag(tree) + + +def _tree_to_bag(tree: dict[str, Any]) -> "Bag": + """ + Convert the nested-dict scaffolding built by extract()/trim() into + actual Bags, leaves left untouched. + + Author: B.G (07/2026) + """ + result = Bag() + for name, value in tree.items(): + result.add(name, _tree_to_bag(value) if isinstance(value, dict) else value) + return result + + +def trim(bag: "Bag", names) -> "Bag": + """ + `bag` minus the named members, as a new Bag. + + Accepts the same plain-name or dotted-path entries as extract(). Removing + `"at.i"` drops just that member, leaving `at` in the result with whatever + else it held; removing a bare name drops that member (and, if it names a + nested Bag, everything under it) whole. + + Parameters + ---------- + bag : Bag + Bag to trim. + names : Iterable[str] + Plain names or dotted paths to remove. + + Returns + ------- + Bag + Fresh Bag with the named members removed. + + Raises + ------ + KeyError + If any path does not resolve in `bag`. + + Author: B.G (07/2026) + """ + removal: dict[str, Any] = {} + for path in names: + _resolve_path(bag, path) # validates the path exists; raises otherwise + parts = path.split(".") + cursor = removal + for part in parts[:-1]: + cursor = cursor.setdefault(part, {}) + cursor[parts[-1]] = None + + def _copy_minus(b: "Bag", rem: dict[str, Any]) -> "Bag": + result = Bag() + for name, item in b.items(): + if name not in rem: + result.add(name, item) + continue + sub = rem[name] + if sub is None: + continue + if not isinstance(item, Bag): + raise KeyError(f"trim: cannot descend into '{name}': not a Bag") + result.add(name, _copy_minus(item, sub)) + return result + + return _copy_minus(bag, removal) + + +def replace(bag: "Bag", name: str, obj: Any) -> "Bag": + """ + `bag` with the member at `name` swapped for `obj`, as a new Bag. + + This is how anything fixed at a Parameter's construction is changed - its + mode (see Parameter.mode), or a const's value (see Parameter.set). Both + mean building a new Parameter, replacing it in here, and recompiling + whatever bound the old one. + + Parameters + ---------- + bag : Bag + Bag to modify. + name : str + Plain name or dotted path of the member to replace. + obj : Any + Replacement value. + + Returns + ------- + Bag + Fresh Bag with `obj` at `name` in place of the old member. + + Raises + ------ + KeyError + If `name` does not resolve in `bag`. + + Author: B.G (07/2026) + """ + parts = name.split(".") + + def _rebuild(b: "Bag", remaining: list[str]) -> "Bag": + head = remaining[0] + if head not in b: + raise KeyError(f"'{name}' not found in bag (no '{head}')") + result = Bag() + for iname, item in b.items(): + if iname != head: + result.add(iname, item) + continue + if len(remaining) == 1: + result.add(iname, obj) + else: + if not isinstance(item, Bag): + raise KeyError(f"replace: cannot descend into '{head}': not a Bag") + result.add(iname, _rebuild(item, remaining[1:])) + return result + + return _rebuild(bag, parts) diff --git a/pyfastflow/experimental/core/context/bk.py b/pyfastflow/experimental/core/context/bk.py new file mode 100644 index 0000000..aca4383 --- /dev/null +++ b/pyfastflow/experimental/core/context/bk.py @@ -0,0 +1,171 @@ +""" +`ctx.bk`: the reserved backend-intrinsics namespace for Taichi/Quadrants +templates that need real transcendental functions or a typed literal cast. + +Why this exists +---------------- +Most closure templates get by on plain arithmetic and the handful of python +builtins Taichi/Quadrants trace natively (`abs`, `min`, `max`, `int`, +`float`). A few need more: transcendental math (`sqrt`, `atan2`, `cos`, +`sin`, `floor`) and typed dtype casts, neither of which has a python +built-in equivalent that traces correctly - bare `math.sqrt`/`math.floor` +inside a `ti.func` raises `TaichiTypeError: must be real number, not Taichi +Expression`, since Taichi's AST transformer only special-cases its own short +builtin list, never the `math` module. + +`ctx.bk` exposes this surface as a reserved, always-present member of `ctx` +on the closure backends only: `ctx.bk.sqrt(x)`, `ctx.bk.atan2(y, x)`, +`ctx.bk.u32(0x846CA68B)`. One template text works against both Taichi and +Quadrants because `make_closure_bk(backend)` resolves the same attribute +surface against whichever module (`ti` or `qd`) `backend` is. + +A free global spliced into a template's namespace (bypassing `ctx` +entirely) was considered and rejected: every reference a template makes +should be visible in its derived Contract, not smuggled in through +`__globals__` - see ctx.py/contract.py for the grammar this keeps intact. + +Reserved, not a slot +--------------------- +`ctx.bk.*` is a builtin the grammar recognises structurally, not a +capability a template's Contract requires satisfied: contract.py's AST walk +drops any chain rooted at `RESERVED_BK_NAME` before it reaches +`Contract.chains`, so `ctx.bk.sqrt(...)` never shows up in `inspect()` or +`unmet()`, and `ingest()` never demands a wired "bk" slot for it. +Symmetrically, `_Builder._wire()`/`compose()` (builder.py) raise if a caller +tries to declare a slot or compose a sub-structure named "bk" - the name is +reserved and can never mean anything else. `bk` is exposed on every level of +the ctx tree compile_closure.py builds, not just the root, since a helper +several levels deep may need it just as much as its caller. + +cupy is unaffected: its templates are raw CUDA text, where the native C +spelling (`sqrtf`, `atan2f`, `floorf`, a plain `0x846CA68Bu` literal) is +already the natural way to write this - `ctx.bk` is never resolved against +a cupy compile. + +Surface +------- +`sqrt`, `atan2`, `cos`, `sin`, `floor` - transcendental math. +`u32`, `i32`, `i64`, `f32` - dtype tokens, exposed as the backend's own +dtype objects (see "Typed literal casts" below). +`bit_cast`, `select`, `cast` - the IEEE-754 bit-flip trick (reinterpreting a +float's bits as u32 and back) plus a branchless ternary-style select. +`atomic_min`, `atomic_max`, `atomic_add` - the three atomic ops Taichi/ +Quadrants expose, for genuinely concurrent writes into a DATA buffer. +`grouped`, `Vector` - `ti.grouped`/`ti.Vector` (and `qd.` equivalents) +passed through unchanged, for dimensionality-agnostic iteration and small +per-thread fixed-size local arrays. + +`abs`/`min`/`max`/`int`/`float` stay plain python builtins rather than +`ctx.bk` members - both backends special-case them for casting the same way +they special-case `abs`/`min`, and they work uniformly on a bare literal or +an already-traced expression, which the dtype objects below do not. + +Typed literal casts +-------------------- +`u32`/`i32`/`i64`/`f32` are exposed as the backend's own dtype objects +themselves (`ti.u32`, never a `lambda x: ti.cast(x, ti.u32)` wrapper) - +`ctx.bk.u32(0x846CA68B)` is called exactly as `ti.u32(0x846CA68B)` would be. +This is load-bearing: Taichi's frontend recognises a call whose callee +resolves, by identity, to one of its own dtype objects, and exempts that +call's literal argument from default-int-type inference even reached +through an attribute chain - `ti.u32(2221713035)` compiles, but wrapping the +same cast in an ordinary python callable breaks it (`Integer literal +2221713035 exceeded the range of default_ip: i32`), because the literal is +type-inferred as a bare argument before the wrapper body ever runs. So +`ctx.bk.u32(...)` must resolve to the real `ti.u32` object one attribute hop +away, not to a function that happens to produce the same value. Only `u32`/ +`i64` currently need the oversized-literal exemption in practice; `i32`/ +`f32` are exposed the same way for symmetry, since `ctx.bk.cast`'s second +argument is always one of these four dtype tokens regardless. + +Author: B.G (08/2026) +""" + +from typing import Any + +RESERVED_BK_NAME = "bk" +"""The reserved ctx member name for the backend-intrinsics namespace - see +the module docstring. Never wirable as a slot, never composable as a root.""" + + +class BkError(Exception): + """ + Raised by an unknown `ctx.bk.*` attribute - naming it and listing what is + actually available, rather than letting a typo fall through to a bare + AttributeError deep inside backend trace machinery. + + Author: B.G (08/2026) + """ + + +_BK_METHOD_NAMES = ( + "sqrt", "atan2", "cos", "sin", "floor", "u32", + "bit_cast", "select", "cast", "atomic_min", "atomic_max", "atomic_add", "i32", "i64", "f32", + "grouped", "Vector", +) +"""Every name `ctx.bk` resolves - see the module docstring's "Surface" section.""" + + +class ClosureBkNode: + """ + `ctx.bk` itself, for one closure backend module (`ti` or `qd`) - see the + module docstring. Every entry in `_BK_METHOD_NAMES` is resolved once, at + construction, straight to the backend's own object (`ti.sqrt`, `ti.u32`, + ...) - never wrapped, so a closure backend's own trace-time recognition + of e.g. `ti.sqrt` as a builtin op, or `ti.u32` as a dtype-cast callee that + exempts its own literal argument from default-int-type inference (see + the module docstring), is identity-based and applies exactly as it would + to a template calling `ti.sqrt`/`ti.u32` directly. + + Author: B.G (08/2026) + """ + + __slots__ = ("_backend", "_fns") + + def __init__(self, backend: Any): + self._backend = backend + self._fns = { + "sqrt": backend.sqrt, + "atan2": backend.atan2, + "cos": backend.cos, + "sin": backend.sin, + "floor": backend.floor, + "u32": backend.u32, + "bit_cast": backend.bit_cast, + "select": backend.select, + "cast": backend.cast, + "atomic_min": backend.atomic_min, + "atomic_max": backend.atomic_max, + "atomic_add": backend.atomic_add, + "i32": backend.i32, + "i64": backend.i64, + "f32": backend.f32, + "grouped": backend.grouped, + "Vector": backend.Vector, + } + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError(name) + try: + return self._fns[name] + except KeyError: + raise BkError( + f"ctx.bk.{name} is not a recognised backend intrinsic - available: " + f"{', '.join(_BK_METHOD_NAMES)}" + ) from None + + def __repr__(self) -> str: + return f"ClosureBkNode(backend={self._backend.__name__}, provides={_BK_METHOD_NAMES})" + + +def make_closure_bk(backend: Any) -> ClosureBkNode: + """ + `ctx.bk` for one closure compile - `backend` is the `taichi` or + `quadrants` module (`BoundKernel.compile()`'s own `backend` argument, + compile_closure.py). Stateless and cheap; built once per compile() and + shared by every node in that compile's own ctx tree. + + Author: B.G (08/2026) + """ + return ClosureBkNode(backend) diff --git a/pyfastflow/experimental/core/context/bound.py b/pyfastflow/experimental/core/context/bound.py new file mode 100644 index 0000000..a55d21d --- /dev/null +++ b/pyfastflow/experimental/core/context/bound.py @@ -0,0 +1,1022 @@ +""" +BoundKernel / BoundHelper: the bind phase - build() a frozen builder into one +of these, then bind()/wire() its slots freely, any number of times, in any +order. See parameter.py's module docstring for the overall build -> bind -> +compile scheme; builder.py/frozen.py are the build phase this continues from. + +build(frozen) - also reachable as `frozen.build()`, see frozen.py - walks +`frozen`'s whole composition tree and mints one independently-bindable slot +per full dotted path it finds: `frozen`'s own top-level PARAM/DATA slots get +single-segment addresses, and every composed root recurses with its own name +prefixed on, all the way down through nested composed FrozenHelpers. This is +where instancing happens - one FrozenHelper composed into eighty different +KernelBuilders is one frozen object (frozen.py), but eighty separate calls to +`.build()` each mint their own, independently-bindable copy of its address +tree, because each call allocates a fresh table. Every wired HELPER slot +still reachable at this point (see builder.py: ingest() does not require one +to be composed, only that a template's actual usage of it be locally known) +MUST be composed by now, or this raises naming the exact address - see the +module docstring of frozen.py for why the check waits until here rather than +running at ingest(): a HELPER slot's content is never filled by bind() (it is +fixed structurally at build time, per the design this module implements), so +build() - which is about to fix the address tree for good - is the last and +only point left where "nothing was ever composed here" can still be caught. + +Addressing is by qualified dotted path, always rooted at the explicit name a +slot or a compose() root was given - `flux.grad.z`, never a positional +`step0.*`. Every address a BoundKernel/BoundHelper accepts is represented +internally as a tuple of segments (`Address`), not a bare string, precisely +so a future pattern/glob layer over these paths (explicitly deferred, not +built here) can match per-segment without this module's own addressing +scheme standing in the way; `parse_address`/`format_address` are the only +places a dotted string and a segment tuple convert between each other. Only +PARAM/DATA leaves are ever minted an address - a prefix that names a +composed sub-structure rather than one of its leaves (`flux.grad.grid`, as +opposed to `flux.grad.grid.NX`) was never given a table entry at all, so +addressing it anywhere (bind, wire, inspect) raises the same "unknown +address" any other typo would. + +bind(addr, obj) fills a slot. Rebinding is normal, not an error - there is no +freeze-once rule at this layer, since immutability is a property of the +*compiled* artifact, not of a slot. What bind() checks depends on the +slot's kind: PARAM accepts any Parameter, of any mode (no mode constraint - +see slot.py's module docstring, genericity across modes is the entire point +of a PARAM slot); DATA checks the dtype declared at wire_data(..., dtype=...) +time, if one was declared, against whatever is bound. + +wire(addr_a, addr_b) makes two slots the same thing, as an equivalence +relation resolved *before* any value is looked at: binding either address +afterwards is visible through the other, symmetrically, and a chain of +wire() calls merges transitively (an ordinary union-find over addresses). +Wiring two slots of different kinds raises; nothing else is guarded - the +caller is trusted about which slots ought to mean the same thing. + +inspect() is the framework's primary debugging surface - the only place a +caller sees the whole binding contract as pasteable addresses, current state, +and wired equivalences at a glance. See its own docstring for the exact +output shape; cross-path leaf-name collisions are reported there as +informational text, never as an error - two unrelated slots happening to +share a last segment (two different `z`s at two different addresses) is +completely ordinary and not a caller's mistake to fix. + +Build-phase sharing (FrozenGroup.shared) +------------------------------------------- +A composed FrozenGroup (frozen.py) may declare, via GroupBuilder.share() +(builder.py), that several of its own composed children's PARAM slots read +the "same" quantity as one of the group's own top-level PARAM slots - grid's +`neighbour_raw.row.NX` and `is_on_edge.row.NX` (among many others) both mean +the grid's own `NX`, structurally, because a device template can only call +what is composed directly onto its own scope (builder.py's module docstring) +and so grid's own public helpers each end up re-composing the same private +`row`/`col` blocks under their own local names. Left alone, build() would +mint one independent address per occurrence - correct, but a caller then has +to bind (or wire-then-bind) every one of them by hand for what is +conceptually one value. + +`_walk_group`/`_walk_group_subtree` are `_walk`'s own group-aware variant: +walking into a composed FrozenGroup with any `.shared` entries switches to +these, which mint the group's own top-level PARAM/HELPER slots exactly as +`_walk` always has, but - for every relative path any canonical's `.shared` +set names - do NOT mint an independent table entry at all; instead they +record `full_address -> canonical_address` in a side `redirect` table +threaded alongside the usual one. The net effect: `build()` mints exactly +ONE address (the canonical) for the whole equivalence class, by default - not +`bind()`-time wire()-together-many, which still leaves every address +independently listed (see wire(), above) - this is coarser and happens +before a caller ever sees the address tree at all. + +Nested groups: `_ShareScope` and outermost-wins +-------------------------------------------------- +A composed FrozenGroup may itself compose another FrozenGroup that also +carries `.shared` entries (visu's hillshade group composing grid, itself +composed independently under each of two private gradient blocks - see +visu/__init__.py's own module docstring for the concrete case this was +designed against). Both layers' sharing must apply at once: grid's own +internal declarations still collapse its own private duplicates to its own +top-level canonical (`...GRID.dist_from_k.DX` -> `...GRID.DX`), and the +enclosing group's declarations may additionally redirect a path that reaches +INTO that nested group (`share("DX", "at.grad_x.GRID.DX", ...)`) to the +enclosing group's own canonical (`...GRID.DX` -> `hillshade.DX`). + +`_ShareScope` is one enclosing group's own sharing declarations, tagged with +the full address (`start`) its own paths are expressed relative to. Walking +into a shared FrozenGroup pushes a new scope onto an ordered list threaded +through the recursion - outermost first, most-recently-entered last - so an +enclosing group's own declarations stay active when a nested group's own +boundary is crossed, rather than being dropped by a single active scope. +`_resolve_shared` is the one place every PARAM leaf - a +group's own top-level name (`_walk_group`) or one reached descending its +composed subtree (`_walk_group_subtree`) - is checked against every active +scope at once. + +Overlapping scopes: checked OUTERMOST first, and the first match wins +outright - an inner group's own declaration for the same leaf is never even +consulted once an outer one has already claimed it. This is deliberate, not +an artifact of iteration order, and there is exactly one reason for it: +an outer scope's own canonical is always a genuinely, unconditionally minted +address (a group's own top-level PARAM loop never itself redirects, by +construction - see `_walk_group` below), so resolving outer-first can never +produce a redirect that points at another redirect. Resolving inner-first +could: `...GRID.dist_from_k.DX` might be collapsed by grid's own inner scope +to `...GRID.DX`, which is *itself* one of the outer scope's declared paths - +inner-first would leave a `redirect` entry pointing at `...GRID.DX`, which is +not itself in `table` (having been redirected further, to `hillshade.DX`), +and `value_at()`'s single-hop lookup does not chase a redirect chain. Outer- +first sidesteps this rather than requiring one: `...GRID.dist_from_k.DX` is +checked against the outer scope first, matches directly (its own full +address, translated relative to the outer scope's `start`, is one of the +outer scope's own declared paths too - grid/noise/visu's own `_find_param_ +paths` helpers do not stop at the shallowest match, so both the nested +canonical and its own private occurrences typically end up declared at the +outer scope as well), and redirects straight to `hillshade.DX` in one hop, +without ever consulting grid's own inner scope for that leaf at all. + +`split_paths` stays scope-local: a leaf exempted from one scope's own +collapse (that scope's own compose-site `split=`) simply falls through to +the next scope in line (or, if none claim it, mints as its own independent +address) - splitting a leaf out of the outer scope's collapse does not +affect whether some inner scope still collapses it, and vice versa. + +`redirect` is consulted only by `value_at()` - the read used internally by +compile_closure.py/compile_cupy.py/compile_shared.py's own structural walks, +which always compute the FULL address as they descend the frozen tree and +need SOME resolution at every PARAM leaf they reach, collapsed or not. It is +never consulted by `bind()`/`wire()`/`unmet()`/`addresses()`/`inspect()` - a +collapsed address is not independently bindable and does not appear in any +of those, which is the intended, visible consequence of collapsing it: `. +addresses()` after composing a D8 grid reports one `NX`, not seventeen. + +compose(name, frozen, split=[...]) (builder.py) opts specific relative paths +back OUT of a composed FrozenGroup's collapse, at the point that group is +composed into some other builder: `_walk`/`_walk_group_subtree` mint those +paths as ordinary, independent addresses instead of adding them to +`redirect`, exactly as if they had never been declared shared at that +compose site. This is a build-time decision, recorded on the composing +object's own `.split` (frozen.py) and read back only while `build()` walks +that one composed occurrence - never something bind() or a caller after the +fact can change; splitting the same shared path back out at a second, +different compose() site of the same group is independent and unaffected. + +Sharing across two separately-built composites (`kA.grid.NX` and +`kB.grid.NX`, two different KernelBuilders each composing their own +occurrence of the same FrozenGroup) is not this mechanism at all - those are +already two different address trees by construction (two separate `build()` +calls, per frozen.py's own instancing guarantee), and reconciling them, if +ever wanted, is ordinary bind-phase wire() between the two BoundKernels' +own addresses. + +Author: B.G (08/2026) +""" + +from typing import Any, NamedTuple + +import numpy as np + +from ..pool.base import new_uid +from .frozen import FrozenKernel, _Frozen +from .slot import SlotKind + +Address = tuple[str, ...] + + +class BindError(Exception): + """ + Raised by anything in the bind phase: an unknown address, a wire() + between mismatched slot kinds, a bind() of the wrong kind of object or + the wrong dtype, or build() finding a wired HELPER slot nothing was ever + composed into. Every case names the exact address involved. + + Author: B.G (08/2026) + """ + + +def parse_address(addr: str) -> Address: + """`"flux.grad.z"` -> `("flux", "grad", "z")`. Raises on an empty string.""" + if not addr: + raise BindError("address must not be empty") + return tuple(addr.split(".")) + + +def format_address(addr: Address) -> str: + """`("flux", "grad", "z")` -> `"flux.grad.z"`.""" + return ".".join(addr) + + +class _LeafInfo(NamedTuple): + """ + The fixed, never-rebound metadata build() mints for one address: which + kind of slot it is, and - DATA only - the dtype declared at wire_data() + time (None if left open). Distinct from the *bound value* itself, which + lives in `_Bound._values` and is free to change via bind()/rebind(). + + Author: B.G (08/2026) + """ + + kind: SlotKind + dtype: Any + + +class _ShareScope(NamedTuple): + """ + One enclosing group's own build-phase-sharing declarations, active while + walking anywhere inside that group's own composed subtree - see the + module docstring's "Nested groups" section. + + `start` is the full address at which this scope's own group root sits - + every path in `shared_paths`/`split_paths` is expressed relative to it, + so a leaf's own relative path for this scope is always `full_addr[len( + start):]`. `shared_paths` maps a relative PARAM path to this scope's own + canonical full address; `split_paths` is the (also relative) set this + scope's own compose-site `split=` opted back out of that collapse. + + Author: B.G (08/2026) + """ + + start: Address + shared_paths: "dict[Address, Address]" + split_paths: frozenset + + +def _resolve_shared(full_addr: Address, scopes: "list[_ShareScope]") -> "Address | None": + """ + The full address this PARAM leaf should redirect to, per every currently + active enclosing group's own sharing declarations, or None if none of + them claim it - see the module docstring's "Overlapping scopes" rule for + why `scopes` (ordered outermost-first by construction - see + `_walk_group`) is checked in that order, with the first match winning + outright. + + Author: B.G (08/2026) + """ + for scope in scopes: + rel = full_addr[len(scope.start) :] + canonical = scope.shared_paths.get(rel) + if canonical is not None and rel not in scope.split_paths: + return canonical + return None + + +def _walk( + prefix: Address, + frozen: _Frozen, + table: dict[Address, _LeafInfo], + redirect: "dict[Address, Address] | None" = None, +) -> None: + """ + Populate `table` with one entry per PARAM/DATA leaf reachable from + `frozen`, at its full dotted path under `prefix`, recursing into every + composed root. Raises BindError, naming the address, for a wired HELPER + slot with nothing composed into it - see the module docstring for why + this is where that gets caught. + + A composed child with any `.shared` entries of its own (a FrozenGroup, + typically, but a FrozenHelper composed as a child may also carry them - + see `_Builder.share()`, builder.py) is walked by `_walk_group` instead - + see the module docstring's "Build-phase sharing" section. `redirect`, optional, collects the + collapsed-address -> canonical-address table that mechanism needs; + every caller that does not care about it (routine.py/sequence.py/ + host_block.py's own direct `_walk()` calls, which only ever want a + reduced `table`) may simply omit it. + + Author: B.G (08/2026) + """ + if redirect is None: + redirect = {} + for name in frozen.slots.names(SlotKind.PARAM): + table[prefix + (name,)] = _LeafInfo(SlotKind.PARAM, None) + for name in frozen.slots.names(SlotKind.DATA): + table[prefix + (name,)] = _LeafInfo(SlotKind.DATA, frozen.slots[name].dtype) + + helper_roots = frozen.slots.names(SlotKind.HELPER) | set(frozen.composed) + for name in helper_roots: + addr = prefix + (name,) + if name not in frozen.composed: + raise BindError( + f"'{format_address(addr)}' is a wired HELPER slot with nothing composed into " + f"it - compose() a frozen helper under that name before build()" + ) + child = frozen.composed[name] + if child.shared: + _walk_group(addr, child, table, redirect, frozen.split.get(name, frozenset())) + else: + _walk(addr, child, table, redirect) + + +def _walk_group( + prefix: Address, + group: _Frozen, + table: dict[Address, _LeafInfo], + redirect: dict[Address, Address], + split_paths: frozenset, + scopes: "list[_ShareScope]" = (), +) -> None: + """ + `_walk`'s group-aware entry point: mints `group`'s own top-level PARAM/ + DATA/HELPER slots exactly as `_walk` always has - except each PARAM name + is first checked against `scopes`, the ENCLOSING scopes already active + when this group was reached (empty at the outermost group in a tree), + since an enclosing group's own sharing may claim this group's own + top-level name for further redirection (see the module docstring's + "Nested groups" section) - then descends into its composed subtree via + `_walk_group_subtree`, pushing this group's own scope (built from + `group.shared`/`split_paths`) onto `scopes` for that descent. + + `group` is not always a GroupBuilder's own FrozenGroup (which indeed + never carries DATA - GroupBuilder.wire_data always raises): a + KernelBuilder that calls `share()` on itself produces a FrozenKernel that + also reaches this function, as build()'s own top-level dispatch (see + build(), below) - and a FrozenKernel's own DATA slots are exactly as real + as a plain `_walk` would mint them, so this mints them here too rather + than silently dropping them. + + Author: B.G (08/2026) + """ + own_shared_paths: dict[Address, Address] = {} + for canonical, paths in group.shared.items(): + canonical_addr = prefix + (canonical,) + for p in paths: + own_shared_paths[p] = canonical_addr + own_scope = _ShareScope(prefix, own_shared_paths, frozenset(split_paths)) + nested_scopes = list(scopes) + [own_scope] + + for name in group.slots.names(SlotKind.PARAM): + full = prefix + (name,) + canonical = _resolve_shared(full, scopes) + if canonical is not None: + redirect[full] = canonical + else: + table[full] = _LeafInfo(SlotKind.PARAM, None) + for name in group.slots.names(SlotKind.DATA): + table[prefix + (name,)] = _LeafInfo(SlotKind.DATA, group.slots[name].dtype) + + helper_roots = group.slots.names(SlotKind.HELPER) | set(group.composed) + for name in helper_roots: + addr = prefix + (name,) + if name not in group.composed: + raise BindError( + f"'{format_address(addr)}' is a wired HELPER slot with nothing composed into " + f"it - compose() a frozen helper under that name before build()" + ) + _walk_group_subtree(addr, group.composed[name], table, redirect, nested_scopes) + + +def _walk_group_subtree( + prefix: Address, + frozen: _Frozen, + table: dict[Address, _LeafInfo], + redirect: dict[Address, Address], + scopes: "list[_ShareScope]", +) -> None: + """ + One level of `_walk_group`'s own descent into a group's composed + subtree. `scopes` carries every enclosing group's own sharing + declarations, outermost first (see the module docstring's "Nested + groups" section) - a PARAM leaf whose full address resolves against any + of them (`_resolve_shared`) is collapsed, a `redirect` entry rather than + an independent `table` entry; every other PARAM/DATA leaf mints + normally, exactly as plain `_walk` would. A nested composed FrozenGroup + (with its own `.shared`) pushes its own scope onto `scopes` for its own + descent (`_walk_group`) rather than starting a fresh, disconnected one - + sharing is declared per group, but an enclosing group's own declarations + stay active reaching through a group-within-a-group. + + Author: B.G (08/2026) + """ + for name in frozen.slots.names(SlotKind.PARAM): + full = prefix + (name,) + canonical = _resolve_shared(full, scopes) + if canonical is not None: + redirect[full] = canonical + else: + table[full] = _LeafInfo(SlotKind.PARAM, None) + for name in frozen.slots.names(SlotKind.DATA): + table[prefix + (name,)] = _LeafInfo(SlotKind.DATA, frozen.slots[name].dtype) + + helper_roots = frozen.slots.names(SlotKind.HELPER) | set(frozen.composed) + for name in helper_roots: + addr = prefix + (name,) + if name not in frozen.composed: + raise BindError( + f"'{format_address(addr)}' is a wired HELPER slot with nothing composed into " + f"it - compose() a frozen helper under that name before build()" + ) + child = frozen.composed[name] + child_split = frozen.split.get(name, frozenset()) + if child.shared: + _walk_group(addr, child, table, redirect, child_split, scopes=scopes) + else: + _walk_group_subtree(addr, child, table, redirect, scopes) + + +def build(frozen: _Frozen) -> "BoundKernel | BoundHelper": + """ + Walk `frozen`'s composition tree and return a fresh BoundKernel (if + `frozen` is a FrozenKernel) or BoundHelper (FrozenHelper), with one + independently-bindable slot minted per full address found - collapsed + per any build-phase sharing reachable in the tree (module docstring, + "Build-phase sharing"). Also reachable as `frozen.build()` (frozen.py). + + Parameters + ---------- + frozen : FrozenKernel or FrozenHelper + The frozen recipe to build. + + Returns + ------- + BoundKernel or BoundHelper + + Author: B.G (08/2026) + """ + table: dict[Address, _LeafInfo] = {} + redirect: dict[Address, Address] = {} + if frozen.shared: + # `frozen` is itself the object being build()-ed directly (e.g. for + # standalone inspection, or a KernelBuilder/HelperBuilder that + # declared its own share() - see _Builder.share(), builder.py) + # rather than reached as someone else's composed child - no + # enclosing object exists to have declared a `split`, so there is + # none. + _walk_group((), frozen, table, redirect, frozenset()) + else: + _walk((), frozen, table, redirect) + cls = BoundKernel if isinstance(frozen, FrozenKernel) else BoundHelper + return cls(frozen, table, redirect) + + +def _format_state(info: _LeafInfo, value: Any) -> str: + """ + The state column of one inspect() line - see _Bound.inspect. + + Author: B.G (08/2026) + """ + if value is None: + return "UNBOUND" + if info.kind is SlotKind.PARAM: + mode = getattr(value, "mode", None) + if mode == "const": + return f"bound(const {value.get()})" + if mode is not None: + return f"bound({mode})" + return "bound" + + +_DTYPE_SHORT = { + "float32": "f32", "float64": "f64", + "int32": "i32", "int64": "i64", + "uint8": "u8", "uint32": "u32", +} + + +def _short_dtype(dtype: Any) -> str: + """ + A dtype in the short spelling this package writes everywhere else + ("f32", "i64", ...) rather than python's own repr - `` is not a pasteable dtype, it is noise. Tries a numpy + coercion first (covers numpy dtypes/dtype classes and the cupy backend's + own dtype objects, which already are numpy dtypes); falls back to + `str(dtype)` for anything numpy cannot make sense of (a Taichi/Quadrants + dtype token, which already prints short - `ti.f32` reprs as `f32`). + + Author: B.G (08/2026) + """ + try: + name = np.dtype(dtype).name + except TypeError: + return str(dtype) + return _DTYPE_SHORT.get(name, name) + + +class _Bound: + """ + Shared machinery behind BoundKernel/BoundHelper. Not instantiated + directly - see build(). + + Author: B.G (08/2026) + """ + + def __init__( + self, + frozen: _Frozen, + table: dict[Address, _LeafInfo], + redirect: "dict[Address, Address] | None" = None, + ): + self._uid = new_uid() + self._frozen = frozen + self._table = table + # union-find over addresses: wire() merges groups, bind()/inspect() + # always resolve through _find() first, so a value lives once per + # group regardless of which member address it was bound through. + self._parent: dict[Address, Address] = {addr: addr for addr in table} + self._values: dict[Address, Any] = {} + # build-phase-collapsed addresses (module docstring, "Build-phase + # sharing") -> their canonical table address. Consulted by value_at() + # only - never by bind()/wire()/unmet()/addresses(), so a collapsed + # address stays genuinely absent from every caller-facing listing. + self._redirect: dict[Address, Address] = dict(redirect) if redirect else {} + + @property + def uid(self) -> int: + """Process-wide identity assigned at construction. See Parameter.uid (parameter.py).""" + return self._uid + + @property + def frozen(self) -> _Frozen: + """The FrozenKernel/FrozenHelper this object was build()-ed from.""" + return self._frozen + + def addresses(self) -> set[Address]: + """Every address this object has a slot for - the full, fixed address tree build() minted.""" + return set(self._table) + + def value_at(self, addr: "Address | str") -> Any: + """ + The object currently bound at `addr` (following wire()-d equivalence + to its group's representative), or None if that group is unbound. + Read-only counterpart to bind() - for the compile phase's use, and + for anything else that wants to read a binding without going + through inspect()'s formatted report. + + `addr` may be a build-phase-collapsed address (module docstring, + "Build-phase sharing") even though it is not one of `.addresses()`' + own members - compile_closure.py/compile_cupy.py/compile_shared.py's + structural walks compute the full address at every PARAM leaf they + reach regardless of whether build() minted it independently or + redirected it, and this is the one read path required to resolve + transparently either way. `bind()`/`wire()` do not get this + treatment - a collapsed address is not independently bindable. + + Parameters + ---------- + addr : Address or str + Dotted path or address tuple, possibly build-phase-collapsed. + + Returns + ------- + Any + The bound object, or None if unbound. + + Author: B.G (08/2026) + """ + a = self._addr_or_redirect(addr) + return self._values.get(self._find(a)) + + def _addr_or_redirect(self, addr: "Address | str") -> Address: + """ + `addr`, validated against `.addresses()` as `_addr()` always has, OR + - if `addr` is not itself one of this object's minted addresses - + its build-phase-collapsed canonical address, if one was recorded. + Raises the same "unknown address" `_addr()` always has if neither + applies. See value_at()'s own docstring for why only that method + uses this instead of `_addr()` directly. + + Author: B.G (08/2026) + """ + a = parse_address(addr) if isinstance(addr, str) else tuple(addr) + if a in self._table: + return a + redirected = self._redirect.get(a) + if redirected is not None: + return redirected + raise BindError( + f"unknown address {format_address(a)!r} - not one of this object's slots " + f"(see .addresses() for the full set)" + ) + + def slot_info(self, addr: "Address | str") -> _LeafInfo: + """This address's fixed kind/dtype, as minted by build() - never changes after that.""" + a = self._addr(addr) + return self._table[self._find(a)] + + def unmet(self) -> list[Address]: + """ + Every address whose equivalence group has no bound value yet, sorted. + Empty means every slot build() minted is filled - the precondition + compile() checks first (see compile_shared.py's check_unmet). + + Returns + ------- + list[Address] + Unbound addresses, sorted. + + Author: B.G (08/2026) + """ + return sorted(addr for addr in self._table if self._values.get(self._find(addr)) is None) + + def _addr(self, addr: "Address | str") -> Address: + a = parse_address(addr) if isinstance(addr, str) else tuple(addr) + if a not in self._table: + raise BindError( + f"unknown address {format_address(a)!r} - not one of this object's slots " + f"(see .addresses() for the full set)" + ) + return a + + def _find(self, addr: Address) -> Address: + parent = self._parent + root = addr + while parent[root] != root: + root = parent[root] + while parent[addr] != root: + parent[addr], addr = root, parent[addr] + return root + + def bind(self, addr: "Address | str", obj: Any) -> "_Bound": + """ + Fill the slot at `addr` (or the whole equivalence group it belongs + to, if wire()-d) with `obj`. Rebinding is normal - see the module + docstring - and simply overwrites what was there. + + A PARAM slot accepts any Parameter, of any mode. A DATA slot with a + declared dtype (wire_data(..., dtype=...)) checks `obj.dtype` + against it, when `obj` has one; an open (dtype=None) DATA slot + accepts anything. There is no HELPER case to handle here: build() + only ever mints table entries for PARAM/DATA leaves (see the module + docstring and `_walk`) - a dotted prefix that names a composed + sub-structure rather than one of its leaves (e.g. `flux.grad.grid`, + as opposed to `flux.grad.grid.NX`) was never minted an address at + all, so `_addr` above already raises "unknown address" for it before + this method's own kind dispatch ever runs. + + Parameters + ---------- + addr : Address or str + obj : Any + A Parameter for a PARAM slot; any dtype-matching object for a + DATA slot. + + Raises + ------ + BindError + `addr` is unknown, or `obj` is the wrong kind/dtype for its slot. + + Author: B.G (08/2026) + """ + a = self._addr(addr) + r = self._find(a) + info = self._table[r] + if info.kind is SlotKind.PARAM: + from .parameter import Parameter + + if not isinstance(obj, Parameter): + raise BindError( + f"{format_address(a)!r} is a PARAM slot; expected a Parameter, got " + f"{type(obj).__name__}" + ) + else: + assert info.kind is SlotKind.DATA + if info.dtype is not None: + obj_dtype = getattr(obj, "dtype", None) + if obj_dtype is not None and obj_dtype != info.dtype: + raise BindError( + f"{format_address(a)!r}: dtype mismatch, slot declares {info.dtype}, " + f"got {obj_dtype}" + ) + self._values[r] = obj + return self + + def bind_leaf(self, mapping: dict[str, Any], *, prefix: "Address | str" = (), strict: bool = False) -> "_Bound": + """ + Bind every address under `prefix` whose last segment is a key of + `mapping`, to that key's value - one bind() call per match, same + checks (PARAM/DATA kind, dtype) as an ordinary bind(). + + The bulk-bind counterpart to hand-writing one bind() per address: a + caller that does not know in advance exactly how many addresses under + `prefix` will need a given name bound - because that count depends on + backend/method/config, not on anything the caller controls (a + FrozenRoutine/FrozenSequence built differently per combination - see + e.g. _cupy_depressions.py's module docstring on real-launch + splitting) - matches against whatever `.addresses()` actually + minted, instead of a hand-typed address list that would need to + change with the combination. + + `prefix`, if given, restricts the match to addresses starting with + it - this is what resolves a leaf name recurring under two different + meanings at two different prefixes (make_accumulation's + pointer_jump_push ping-pong, where "rec_curr" means one buffer under + "step_a" and a different one under "step_b" - two bind_leaf() calls, + one per prefix, rather than one call that would bind both to + whichever value came last). + + `strict=True` raises if any `mapping` key matched no address under + `prefix` at all - off by default, since an existing caller may + deliberately pass one mapping wider than what a particular + combination's own address tree contains (leaf-name binding across + several method/reroute combinations relies on exactly this). A key + that genuinely never matches anything under any combination is + almost always a typo, though - pass `strict=True` wherever the + address set is known fixed. + + Parameters + ---------- + mapping : dict[str, Any] + Leaf name -> value to bind. + prefix : Address or str, optional + Restricts matches to addresses starting with this prefix. + strict : bool, optional + Raise if a mapping key matches no address. + + Returns + ------- + _Bound + self, for chaining. + + Raises + ------ + BindError + If `strict` is True and some key matched no address. + + Author: B.G (08/2026) + """ + p = parse_address(prefix) if isinstance(prefix, str) else tuple(prefix) + plen = len(p) + matched: set[str] = set() + for addr in self._table: + if addr[:plen] == p and addr[-1] in mapping: + self.bind(addr, mapping[addr[-1]]) + matched.add(addr[-1]) + if strict: + unused = sorted(set(mapping) - matched) + if unused: + raise BindError(f"bind_leaf(prefix={format_address(p)!r}): {unused} matched no address") + return self + + def bind_pattern(self, pattern: str, obj: Any) -> "_Bound": + """ + Bind every address matching `pattern`, a dotted string the same + length as the addresses it may match - a `*` segment matches any one + segment, any other segment must match literally. `"step_a.*.rec_curr"` + matches `("step_a", "get_src", "rec_curr")` but neither + `("step_a", "rec_curr")` (too short) nor + `("step_a", "a", "b", "rec_curr")` (too long) - there is deliberately + no multi-segment wildcard (see the module docstring's note on Address + being a tuple precisely so a pattern/glob layer over these paths + could be added without disturbing the addressing scheme itself). + + Raises if `pattern` matches zero addresses: unlike bind_leaf's + `strict` flag (off by default, since one mapping is often + deliberately wider than one combination's own address tree), a + single bind_pattern() call names one specific intended match, so a + pattern that resolves to nothing is almost always a typo, not a + legitimately-absent combination. + + Parameters + ---------- + pattern : str + Dotted pattern, `*` matching any one segment. + obj : Any + Value to bind at every matching address. + + Returns + ------- + _Bound + self, for chaining. + + Raises + ------ + BindError + If `pattern` matches zero addresses. + + Author: B.G (08/2026) + """ + segs = tuple(pattern.split(".")) + matched = False + for addr in self._table: + if len(addr) == len(segs) and all(s == "*" or s == a for s, a in zip(segs, addr)): + self.bind(addr, obj) + matched = True + if not matched: + raise BindError(f"bind_pattern({pattern!r}): matched no address") + return self + + def wire(self, addr_a: "Address | str", addr_b: "Address | str") -> "_Bound": + """ + Make `addr_a` and `addr_b` the same slot: binding either afterwards + fills both, and any address already wired to either joins the same + group (transitively - an ordinary union-find). Raises if the two + resolve to different slot kinds, or if both sides are already bound + to different objects (which wiring them together could not resolve + without silently discarding one). No other guard runs - see the + module docstring. + + Parameters + ---------- + addr_a, addr_b : Address or str + The two addresses to merge into one equivalence group. + + Returns + ------- + _Bound + self, for chaining. + + Raises + ------ + BindError + If the two addresses resolve to different slot kinds, or both + are already bound to different objects. + + Author: B.G (08/2026) + """ + a, b = self._addr(addr_a), self._addr(addr_b) + ra, rb = self._find(a), self._find(b) + if ra == rb: + return self + if self._table[ra].kind is not self._table[rb].kind: + raise BindError( + f"wire({format_address(a)!r}, {format_address(b)!r}): kind mismatch " + f"({self._table[ra].kind.value} vs {self._table[rb].kind.value})" + ) + va, vb = self._values.get(ra), self._values.get(rb) + if va is not None and vb is not None and va is not vb: + raise BindError( + f"wire({format_address(a)!r}, {format_address(b)!r}): both sides are already " + f"bound to different objects - rebind one to match before wiring" + ) + self._parent[ra] = rb + if rb not in self._values and va is not None: + self._values[rb] = va + return self + + def inspect(self) -> str: + """ + The full binding contract, one line per address, as exact pasteable + addresses, columns aligned to whatever the actual addresses/types on + this object need (never a fixed width - an address only ever gets + longer as a composition tree grows deeper): + + flux.grad.dx PARAM - bound(const 30.0) + flux.grad.z PARAM - UNBOUND + flux.acc DATA f32 UNBOUND + update.dt PARAM - bound(scalar) + + PARAM and DATA rows share one column layout - address, kind, type, + state - rather than DATA carrying an extra field: a PARAM slot's + type column reads "-" (the slot itself declares no dtype - see + slot.py), a DATA slot's reads its declared dtype in this package's + short spelling ("f32", not "" - see + _short_dtype) or "any" if wire_data() left it open. + + A wire()-d address carries a trailing `[wired: ...]` note listing + every other address in its equivalence group; the state column is + only padded when at least one row needs that trailing note, so a + report with no wire()-d addresses at all has no dangling whitespace. + After the per-address lines, any leaf name (an address's own last + segment) shared by two or more addresses that are *not* in the same + wire()-d group is listed once more, under an "Informational" + heading - never as an error; see the module docstring for why this + is deliberately not a conflict. + + Returns + ------- + str + The formatted report. + + Author: B.G (08/2026) + """ + groups: dict[Address, list[Address]] = {} + for addr in self._table: + groups.setdefault(self._find(addr), []).append(addr) + + rows: list[tuple[str, str, str, str, str]] = [] + for addr in sorted(self._table): + info = self._table[addr] + root = self._find(addr) + state = _format_state(info, self._values.get(root)) + peers = sorted(a for a in groups[root] if a != addr) + wired = f"[wired: {', '.join(format_address(p) for p in peers)}]" if peers else "" + if info.kind is SlotKind.DATA: + type_col = _short_dtype(info.dtype) if info.dtype is not None else "any" + else: + type_col = "-" + rows.append((format_address(addr), info.kind.value.upper(), type_col, state, wired)) + + by_leaf: dict[str, list[Address]] = {} + for addr in self._table: + by_leaf.setdefault(addr[-1], []).append(addr) + collisions = [] + for leaf, addrs in sorted(by_leaf.items()): + if len(addrs) < 2 or len({self._find(a) for a in addrs}) < 2: + continue + collisions.append(f" '{leaf}': {', '.join(format_address(a) for a in sorted(addrs))}") + + if not rows: + report = "(no slots)" + else: + w_addr = max(len(r[0]) for r in rows) + w_kind = max(len(r[1]) for r in rows) + w_type = max(len(r[2]) for r in rows) + w_state = max(len(r[3]) for r in rows) + pad_state = any(r[4] for r in rows) + lines = [] + for addr_s, kind_s, type_s, state_s, wired_s in rows: + parts = [addr_s.ljust(w_addr), kind_s.ljust(w_kind), type_s.ljust(w_type)] + parts.append(state_s.ljust(w_state) if pad_state else state_s) + line = " ".join(parts) + if wired_s: + line += f" {wired_s}" + lines.append(line) + report = "\n".join(lines) + + if collisions: + report += "\n\nInformational - same leaf name at multiple, unwired addresses (not an error):\n" + report += "\n".join(collisions) + return report + + def __repr__(self) -> str: + return f"{type(self).__name__}(uid={self._uid}, slots={len(self._table)})" + + +class BoundKernel(_Bound): + """ + The bound result of build()-ing a FrozenKernel. See the module + docstring. + + Author: B.G (08/2026) + """ + + def compile(self, backend: str, **kwargs) -> Any: + """ + The compile phase - produce a frozen, immutable callable from + this object's current bindings. A snapshot: this BoundKernel stays + live and rebindable afterwards, and a later compile() (with or + without edits in between) produces an independent callable - see + compile_shared.py's module docstring for the full contract + (CompiledKernel, swap(), the legal-PARAM-accessor and unmet-slot + checks every backend runs first). + + `backend` is `"taichi"`, `"quadrants"` or `"cupy"` - the same three + names `backends.py`'s `backend_classes()` uses elsewhere in this + package. `**kwargs` is backend-specific: cupy's `compile_kernel` + accepts `grid=`/`block=` launch-dimension defaults (see + compile_cupy.py); the closure backends take none. + + Imported locally, per backend, to avoid importing taichi/quadrants/ + cupy at module load time for a caller that only uses one of them - + the same reasoning `backends.py.backend_classes` follows. + + Parameters + ---------- + backend : str + "taichi", "quadrants" or "cupy". + **kwargs + Backend-specific: cupy accepts `grid=`/`block=`. + + Returns + ------- + Any + A compiled, callable kernel object (CompiledKernel). + + Raises + ------ + BindError + If `backend` is not one of the three known names. + + Author: B.G (08/2026) + """ + if backend == "taichi": + import taichi as ti + + from . import compile_closure + + return compile_closure.compile_kernel(self, ti, **kwargs) + if backend == "quadrants": + import quadrants as qd + + from . import compile_closure + + return compile_closure.compile_kernel(self, qd, **kwargs) + if backend == "cupy": + from . import compile_cupy + + return compile_cupy.compile_kernel(self, **kwargs) + raise BindError(f"compile: unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + + +class BoundHelper(_Bound): + """ + The bound result of build()-ing a FrozenHelper. See the module + docstring. + + Author: B.G (08/2026) + """ + + def compile(self, backend: str, **kwargs) -> Any: + """ + Always raises: a device helper has no standalone compiled form, on + any backend - it is compiled as part of the BoundKernel that + composes it (see compile_shared.py/compile_closure.py/ + compile_cupy.py). Mirrors HelperBuilder.compile() (builder.py) at + the build phase. + + Raises + ------ + TypeError + Always. + + Author: B.G (08/2026) + """ + raise TypeError( + "BoundHelper.compile() is not supported: a device helper is compiled as part of " + "the BoundKernel that composes it, not on its own. Compose this helper's " + "FrozenHelper into a KernelBuilder and call compile() on the resulting BoundKernel." + ) diff --git a/pyfastflow/experimental/core/context/builder.py b/pyfastflow/experimental/core/context/builder.py new file mode 100644 index 0000000..add6cd3 --- /dev/null +++ b/pyfastflow/experimental/core/context/builder.py @@ -0,0 +1,657 @@ +""" +KernelBuilder / HelperBuilder: the build phase - wire_param()/wire_helper()/ +wire_data()/compose(), then ingest() to close it out. See parameter.py's +module docstring for what the overall scheme (build -> bind -> compile) is +for; this module is the first of those three phases only. + +A builder is mutable exactly until ingest() runs. wire_param(name)/ +wire_helper(name)/wire_data(name) each declare one local Slot (slot.py); +compose(name, frozen) attaches an already-frozen sub-structure (frozen.py) +under an explicit slot name - never positionally - so a template can reach +`flux.grad.z` once `flux` names what was composed. Every name - wired or +composed - lives in one flat namespace per builder; wiring the same name +twice, or composing over a name already used by a PARAM/DATA slot, raises. +compose() may target a name already wire_helper()'d, though: a HELPER slot +declares that a name *will* be reachable as `ctx.{name}(...)`, and compose() +is how that promise gets kept - see compose()'s own docstring for why this +is allowed while every other double-use of a name is not. + +compose(name, frozen, split=[...]) is the other half of build-phase sharing +(GroupBuilder.share(), below): when `frozen` is a FrozenGroup carrying +`.shared` paths, every one of them is collapsed into its own canonical +address by default (bound.py's build() - see that module's docstring for +the full mechanism), and `split` opts specific dotted relative paths back +out into their own, independently-bindable addresses again at THIS compose +site. Each path in `split` must already be one of `frozen.shared`'s own +declared relative paths, checked here, eagerly - naming the exact path if +not. `split` on anything that is not a FrozenGroup with `.shared` entries +raises: there is nothing to split. + +ingest(template) is where the local contract is checked and the structural +contract is derived (contract.py) - a python def by static AST walk, CUDA +source text by scanning its own `$ctx....$` spans, dispatched on +`isinstance(template, str)`. Every chain the contract requires must resolve: +its root is either a wired PARAM/HELPER slot (further segments trusted, +unchecked - see the class docstrings below) or a composed root, in which case +the *next* segment must be among what the composed candidate `.provides` +(frozen.py) - otherwise this raises naming exactly what is missing. A chain +rooted at a wired DATA slot is a contract violation of a different kind: a +DATA slot is a plain call argument of the template's own signature, never +reached through `ctx`, so this raises with a hint toward that instead of the +generic "no declared slot" message. ingest() does not itself require a wired +HELPER slot to already be composed - a template need not reference every +slot it declares, and an unreferenced, uncomposed HELPER slot is harmless to +ingest(). It becomes a hard requirement one phase later, at build() (see +frozen.py, bound.py): the address tree build() walks has no way to +represent "reachable, but nothing composed here yet". + +`RESERVED_BK_NAME` ("bk", bk.py) may never be wired (wire_param/wire_helper/ +wire_data, via `_wire`) or composed (`compose`) - `ctx.bk` is reserved, +backend-recognised grammar (the backend-intrinsics namespace: `ctx.bk.sqrt`, +`ctx.bk.atan2`, ...), not a name any factory's own template surface may +repurpose. See bk.py's module docstring for the full mechanism and +contract.py for the matching rule on the derivation side (a `ctx.bk.*` chain +is dropped before it ever becomes a contract requirement, so ingest() never +asks for a "bk" slot to be wired in the first place). + +ingest() returns a frozen, immutable FrozenKernel/FrozenHelper and freezes +the builder itself in the same call - every wire_*/compose/ingest afterwards +raises FrozenBuilderError. A builder is therefore used once, start to finish; +build a new one for a different template rather than trying to reuse an +ingested one: a builder holding live, still-mutable slot state after having +already handed out one frozen, immutable result would be exactly the kind +of aliasing hazard the frozen/mutable split exists to rule out. + +Author: B.G (08/2026) +""" + +from typing import Any + +from ..pool.base import new_uid +from .bk import RESERVED_BK_NAME +from .contract import Contract, ContractError, extract_cupy_contract, extract_python_contract +from .frozen import FrozenBuilderError, FrozenGroup, FrozenHelper, FrozenKernel, _Frozen +from .slot import DataSlot, HelperSlot, ParamSlot, Slot, SlotGroup, SlotGroupError, SlotKind + + +class _Builder: + """ + Shared build-phase machinery behind KernelBuilder/HelperBuilder. Not + instantiated directly. + + Author: B.G (08/2026) + """ + + def __init__(self): + self._uid = new_uid() + self._slots = SlotGroup() + self._composed: dict[str, _Frozen] = {} + self._split: dict[str, frozenset] = {} + self._shared: dict[str, list[tuple]] = {} + self._frozen = False + + @property + def uid(self) -> int: + """Process-wide identity assigned at construction. See Parameter.uid (parameter.py).""" + return self._uid + + @property + def slots(self) -> SlotGroup: + """This builder's currently wired slots. Read-only - go through wire_*() to add more.""" + return self._slots + + @property + def composed(self) -> dict[str, _Frozen]: + """This builder's currently composed {name: frozen sub-structure}. Read-only - go through compose().""" + return dict(self._composed) + + def _check_mutable(self) -> None: + if self._frozen: + raise FrozenBuilderError( + f"{type(self).__name__}(uid={self._uid}) has already been ingest()-ed and is " + f"frozen - build a new {type(self).__name__} instead of reusing this one" + ) + + def _wire(self, slot: Slot) -> "_Builder": + self._check_mutable() + if slot.name == RESERVED_BK_NAME: + raise SlotGroupError( + f"'{RESERVED_BK_NAME}' is reserved - ctx.{RESERVED_BK_NAME} is the " + f"backend-intrinsics namespace (bk.py) and can never be wired as a slot" + ) + if slot.name in self._composed: + raise SlotGroupError(f"'{slot.name}' is already composed on this builder") + self._slots.add(slot) + return self + + def wire_param(self, name: str) -> "_Builder": + """ + Declare a PARAM slot named `name`: reached in device code as + `ctx.{name}.get(...)` / `ctx.{name}.set_node(...)`, uniformly across + whatever mode the Parameter eventually bound to it has (see slot.py's + module docstring). Deliberately generic - a slot declares a place to + plug in a Parameter (parameter.py) later, not a shape; nothing + here constrains mode or dtype, which is the entire point of a + Parameter being able to move between const/scalar/field without + touching a template. + + Parameters + ---------- + name : str + Slot name. + + Returns + ------- + _Builder + self, for chaining. + + Author: B.G (08/2026) + """ + return self._wire(ParamSlot(name)) + + def wire_helper(self, name: str) -> "_Builder": + """ + Declare a HELPER slot named `name`: called in device code as + `ctx.{name}(...)`. Filled at build time via compose() under this + same name - see the module docstring and compose()'s own docstring + for why compose() (and only compose()) may target an already-wired + HELPER slot. + + Parameters + ---------- + name : str + Slot name. + + Returns + ------- + _Builder + self, for chaining. + + Author: B.G (08/2026) + """ + return self._wire(HelperSlot(name)) + + def wire_data(self, name: str, *, dtype: Any = None) -> "_Builder": + """ + Declare a DATA slot named `name`: a trusted call argument of the + compiled kernel/helper's own signature, never reached through `ctx`. + See slot.py's module docstring for the PARAM/HELPER/DATA distinction. + Overridden on HelperBuilder to always raise - a helper is + device-only and takes data only as its caller's own trusted + argument, never as a declared slot of its own. + + Parameters + ---------- + name : str + dtype : optional + Declares this slot's data-argument contract, checked later at + bind/compile time against whatever value is actually bound or + passed. Left as None, the slot stays open to any dtype. + + Returns + ------- + _Builder + self, for chaining. + + Author: B.G (08/2026) + """ + return self._wire(DataSlot(name, dtype=dtype)) + + def compose(self, name: str, frozen: _Frozen, *, split: "list[str] | None" = None) -> "_Builder": + """ + Attach an already-frozen sub-structure (a FrozenKernel, FrozenHelper + or FrozenGroup - frozen.py) under slot `name`, giving a template + reaching `ctx.{name}` access to whatever `frozen` itself provides - + `{name}.{member}` for any PARAM/HELPER slot or composed root + `frozen` carries at its own top level. + + `frozen` is stored by identity, not copied: compose the same object + into any number of builders and every one of them shares it. + + `name` may be either fresh (nothing wired or composed under it yet) + or a HELPER slot already declared via wire_helper() on this same + builder - composing there is how that slot's "reachable, filled in + later" promise is kept, and is the one case where a name may be used + twice: once to wire the slot, once to compose its content. Composing + under a name already composed, or already wired as PARAM/DATA + (a kind compose() has no business filling), raises. `frozen` must be + a FrozenHelper or a FrozenGroup - a FrozenKernel raises: a kernel is + a host entry point, not something device code can call, and on a GPU + backend a kernel cannot call another kernel. + + `split`, optional, is a list of dotted relative paths (e.g. + `"neighbour_raw.row.NX"`) that opt back out of `frozen`'s own + build-phase sharing (GroupBuilder.share(), FrozenGroup.shared) at + THIS compose site, re-minting each as its own independently-bindable + address instead of collapsing into its shared canonical - see the + module docstring and bound.py's module docstring for the full + mechanism. + + Parameters + ---------- + name : str + Slot name to compose `frozen` under. + frozen : FrozenHelper or FrozenGroup + Already-frozen sub-structure to attach. + split : list[str], optional + Dotted relative paths to opt back out of `frozen`'s build-phase + sharing at this compose site. + + Returns + ------- + _Builder + self, for chaining. + + Raises + ------ + TypeError + If `frozen` is not a `_Frozen`, or is a FrozenKernel. + SlotGroupError + If `name` is reserved, already composed, already wired as a + non-HELPER slot, or `split` names a path not in + `frozen.shared`, or `split` is given for a `frozen` with no + shared paths at all. + + Author: B.G (08/2026) + """ + self._check_mutable() + if name == RESERVED_BK_NAME: + raise SlotGroupError( + f"'{RESERVED_BK_NAME}' is reserved - ctx.{RESERVED_BK_NAME} is the " + f"backend-intrinsics namespace (bk.py) and can never be composed as a root" + ) + if not isinstance(frozen, _Frozen): + raise TypeError(f"compose({name!r}, ...): expected a FrozenKernel/FrozenHelper, got {type(frozen).__name__}") + if isinstance(frozen, FrozenKernel): + raise TypeError( + f"compose({name!r}, ...): got a FrozenKernel, not a FrozenHelper - a kernel is a " + f"host entry point, not a device-callable helper, and cannot be composed into " + f"another builder (on a GPU backend a kernel cannot call another kernel). Build " + f"the shared logic as a HelperBuilder instead." + ) + if name in self._composed: + raise SlotGroupError(f"'{name}' is already composed on this builder") + if name in self._slots and self._slots[name].kind is not SlotKind.HELPER: + raise SlotGroupError( + f"'{name}' is already wired on this builder as {self._slots[name]!r}; compose() " + f"only fills a HELPER slot (or a fresh name), never a PARAM/DATA one" + ) + self._composed[name] = frozen + if split: + shared = getattr(frozen, "shared", None) + if not shared: + raise SlotGroupError( + f"compose({name!r}, ..., split={split!r}): {name!r}'s frozen object has no " + f"build-phase-shared PARAM paths to split - split only applies to a " + f"FrozenGroup composed with at least one share() declaration" + ) + all_shared = {p for paths in shared.values() for p in paths} + resolved = set() + for path in split: + segs = tuple(path.split(".")) + if segs not in all_shared: + raise SlotGroupError( + f"compose({name!r}, ..., split=...): {path!r} is not a shared path on " + f"the composed group (shared paths: " + f"{sorted('.'.join(p) for p in all_shared)})" + ) + resolved.add(segs) + self._split[name] = frozenset(resolved) + return self + + def share(self, canonical: str, *paths: str) -> "_Builder": + """ + Declare that `canonical` - a PARAM slot already wire_param()'d on + THIS builder - is the same value as each dotted `paths`, a relative + address reaching a PARAM slot somewhere in this builder's own + already-composed subtree (e.g. `"neighbour_raw.row.NX"`: the `row` + helper composed inside the `neighbour_raw` helper composed on this + builder, its own `NX` slot). bound.py's build() acts on this: by + default, every declared path collapses into `canonical`'s own + address - only `canonical` is independently minted, not every + private occurrence - which is the whole point (see bound.py's + module docstring for why this needed a build-phase mechanism rather + than being left to bind-phase wire() or bulk/pattern binding). + + This is explicit and local to one builder's own authoring - never + name-based matching across independently-authored composites (which + is exactly the kind of accidental collision this architecture's + addressing exists to prevent). A caller composing this builder's + frozen result elsewhere opts specific paths back OUT of the collapse + via compose()'s own `split=` (`_Builder.compose()`). + + Available on KernelBuilder and HelperBuilder as well as GroupBuilder: + a kernel or helper that both reads a composed sub-structure's PARAM + slot directly (via its own wire_param()) and also composes something + that re-composes the same sub-structure may collapse those + occurrences itself, exactly as a GroupBuilder does for its own + composed children - no group wrapper needed purely to reach share(). + For a KernelBuilder, `canonical` is only actually usable as a + collapse target - and this builder's own `.shared` only actually + takes effect - when this object is later reached as build()'s own + top-level frozen argument (a FrozenKernel is never itself composed + as someone else's child, per compose()'s own FrozenKernel guard + above); for a HelperBuilder composed as a child elsewhere, its + `.shared` is consulted exactly as a FrozenGroup's is (see bound.py's + module docstring). + + Parameters + ---------- + canonical : str + PARAM slot, already wire_param()'d on this builder, that the + given `paths` collapse into. + *paths : str + Dotted relative addresses into this builder's own composed + subtree, each naming a PARAM slot to share with `canonical`. + + Returns + ------- + _Builder + self, for chaining. + + Raises + ------ + SlotGroupError + If `canonical` is not a PARAM slot wired on this builder, if a + path does not resolve (through this builder's already-composed + children) to a real PARAM slot, or if a path is already declared + shared under a different (or the same) canonical - each relative + path may be shared at most once. + + Author: B.G (08/2026) + """ + self._check_mutable() + if canonical not in self._slots or self._slots[canonical].kind is not SlotKind.PARAM: + raise SlotGroupError( + f"share({canonical!r}, ...): {canonical!r} is not a PARAM slot wired on this " + f"builder - call wire_param({canonical!r}) before share()" + ) + if not paths: + raise SlotGroupError(f"share({canonical!r}): at least one path is required") + + already_shared = {p for ps in self._shared.values() for p in ps} + resolved: list[tuple] = [] + for path in paths: + segs = tuple(path.split(".")) + if len(segs) < 2: + raise SlotGroupError( + f"share({canonical!r}, {path!r}): a shared path must reach into a composed " + f"child (at least 'child.PARAM'), got {path!r}" + ) + root = segs[0] + if root not in self._composed: + raise SlotGroupError(f"share({canonical!r}, {path!r}): {root!r} is not composed on this builder") + node: _Frozen = self._composed[root] + walked = root + for seg in segs[1:-1]: + if seg not in node.composed: + raise SlotGroupError(f"share({canonical!r}, {path!r}): {seg!r} is not composed under {walked!r}") + node = node.composed[seg] + walked = f"{walked}.{seg}" + leaf = segs[-1] + if leaf not in node.slots.names(SlotKind.PARAM): + raise SlotGroupError(f"share({canonical!r}, {path!r}): {leaf!r} is not a PARAM slot under {walked!r}") + if segs in already_shared: + raise SlotGroupError(f"share({canonical!r}, {path!r}): {path!r} is already shared") + resolved.append(segs) + already_shared.add(segs) + + self._shared.setdefault(canonical, []) + self._shared[canonical].extend(resolved) + return self + + def _derive_and_check(self, template: Any) -> tuple[SlotGroup, dict[str, _Frozen], Contract]: + """ + Derive `template`'s Contract and check every chain it requires + against this builder's wired slots and composed sub-structures - + see the module docstring for exactly what each chain shape needs. + Returns the (slots, composed, contract) triple ingest() freezes + into a FrozenKernel/FrozenHelper; raises nothing itself, letting + ContractError/SlotGroupError from the checks below propagate. + + Author: B.G (08/2026) + """ + contract = extract_cupy_contract(template) if isinstance(template, str) else extract_python_contract(template) + + param_and_helper_roots = self._slots.names(SlotKind.PARAM) | self._slots.names(SlotKind.HELPER) + data_roots = self._slots.names(SlotKind.DATA) + + for chain in contract.chains: + root = chain[0] + if root in self._composed: + contract.check_root(root, self._composed[root].provides) + elif root in param_and_helper_roots: + continue + elif root in data_roots: + raise ContractError( + f"ctx.{root} is not reachable: '{root}' is a wire_data slot, and data is " + f"a trusted call argument of the template's own signature, never reached " + f"through ctx - pass it as a plain parameter instead of wiring it as a slot" + ) + else: + raise ContractError( + f"ctx.{root} has no declared slot - call wire_param({root!r}) or " + f"wire_helper({root!r}) before ingest(), or compose({root!r}, ...) an " + f"already-frozen sub-structure" + ) + + return self._slots.copy(), dict(self._composed), contract + + +class HelperBuilder(_Builder): + """ + Builds a device helper: PARAM/HELPER slots only, no data of its own. See + the module docstring's local-contract rules and frozen.py for what + ingest() returns. + + A helper takes data only as a trusted argument passed by whatever calls + it - never a declared slot of its own - so wire_data always raises here. + + Author: B.G (08/2026) + """ + + def wire_data(self, name: str, *, dtype: Any = None) -> "HelperBuilder": + """ + Always raises: a HelperBuilder is device-only and carries PARAM and + HELPER slots only. Data reaches a helper as a trusted call argument + supplied by whatever calls it, never as a slot declared on the + helper itself. Declare the data slot on the enclosing KernelBuilder + instead. + + Author: B.G (08/2026) + """ + raise TypeError( + "HelperBuilder.wire_data() is not allowed: a helper is device-only and takes data " + "only as a trusted call argument of its caller. Declare wire_data on the enclosing " + "KernelBuilder, and pass the value through as an ordinary template argument." + ) + + def ingest(self, template: Any) -> FrozenHelper: + """ + Close out the build phase: derive and check `template`'s contract + (see the module docstring), freeze this builder, and return the + resulting FrozenHelper. + + Parameters + ---------- + template : Any + A python def (closure backends) or CUDA source text (cupy). + + Returns + ------- + FrozenHelper + + Raises + ------ + ContractError + If a chain `template` requires has no matching slot/composed + root. + + Author: B.G (08/2026) + """ + self._check_mutable() + slots, composed, contract = self._derive_and_check(template) + self._frozen = True + return FrozenHelper(template, slots, composed, contract, split=self._split, shared=self._shared) + + +class KernelBuilder(_Builder): + """ + Builds a kernel: PARAM/HELPER/DATA slots all allowed. See the module + docstring's local-contract rules and frozen.py for what ingest() returns. + + Author: B.G (08/2026) + """ + + def ingest(self, template: Any) -> FrozenKernel: + """ + Close out the build phase: derive and check `template`'s contract + (see the module docstring), freeze this builder, and return the + resulting FrozenKernel. + + Parameters + ---------- + template : Any + A python def (closure backends) or CUDA source text (cupy). + + Returns + ------- + FrozenKernel + + Raises + ------ + ContractError + If a chain `template` requires has no matching slot/composed + root. + + Author: B.G (08/2026) + """ + self._check_mutable() + slots, composed, contract = self._derive_and_check(template) + self._frozen = True + return FrozenKernel(template, slots, composed, contract, split=self._split, shared=self._shared) + + +class GroupBuilder(_Builder): + """ + Builds a non-callable, navigable composite: PARAM/HELPER slots and + composed sub-structures only, no template of its own and never callable + in device code - see frozen.py's FrozenGroup for what this closes into + and why it exists (a caller needing both `ctx.grid.neighbour(i, k)`, a + composed HELPER call, and `ctx.grid.NX.get(0)`, a PARAM leaf reached + straight through the same composite, one level in). + + `wire_data` always raises, for the same reason it does on HelperBuilder: + a group is device-structure-only, never a call argument's own signature. + + `share()` (inherited from `_Builder` - see its own docstring for the full + mechanism) is build-phase sharing: a group PARAM slot the group's own + author declares once, that stands in for the same value re-read by + several of the group's own composed children - frozen.py's FrozenGroup + is what it freezes into. + + Author: B.G (08/2026) + """ + + def wire_data(self, name: str, *, dtype: Any = None) -> "GroupBuilder": + """ + Always raises: a GroupBuilder declares PARAM/HELPER slots only. See + HelperBuilder.wire_data() (same reasoning) and frozen.py's + FrozenGroup. + + Author: B.G (08/2026) + """ + raise TypeError( + "GroupBuilder.wire_data() is not allowed: a group is a passive, device-structure-" + "only composite - it is never the template a call argument belongs to. Declare " + "wire_data on whichever KernelBuilder eventually composes this group." + ) + + def close(self) -> FrozenGroup: + """ + Close out the build phase and return the resulting FrozenGroup. + Unlike KernelBuilder.ingest()/HelperBuilder.ingest(), there is no + template to derive a Contract from - a group is never itself the + target of a ctx.* chain resolution of its own body (see frozen.py), + so its Contract is always empty. Every wired HELPER slot must still + end up composed by build() time (frozen.py/bound.py), exactly as for + a HelperBuilder/KernelBuilder - unreferenced here since there is no + template to check it against at this phase, but still enforced one + phase later. + + Returns + ------- + FrozenGroup + + Author: B.G (08/2026) + """ + self._check_mutable() + self._frozen = True + return FrozenGroup( + None, self._slots.copy(), dict(self._composed), Contract(frozenset()), + split=self._split, shared=self._shared, + ) + + +def find_param_paths(frozen: "_Frozen", leaf_name: str, prefix: tuple = ()) -> list: + """ + Every relative dotted path, as a `"a.b.NAME"` string, under `frozen`'s own + composed subtree whose PARAM slot is literally named `leaf_name` - the + itemized list `share_leaf` hands to GroupBuilder.share(). Recurses through + `.composed` only (a HELPER slot with nothing composed raises earlier, at + that structure's own ingest()/build(), never reached here). Generic over + whether a composed node is itself a FrozenHelper or a nested FrozenGroup. + + Shared by grid/noise/visu's own factories - see grid/__init__.py's module + docstring ("Build-phase sharing collapses the duplicate addresses") for + why this exists. + + Parameters + ---------- + frozen : _Frozen + Sub-structure to search. + leaf_name : str + PARAM slot name to find. + prefix : tuple, optional + Path segments prepended to every result; used internally for + recursion. + + Returns + ------- + list[str] + Dotted relative paths to every occurrence of `leaf_name`. + + Author: B.G (08/2026) + """ + paths = [] + if leaf_name in frozen.slots.names(SlotKind.PARAM): + paths.append(".".join(prefix + (leaf_name,))) + for name, child in frozen.composed.items(): + paths.extend(find_param_paths(child, leaf_name, prefix + (name,))) + return paths + + +def share_leaf(group: "GroupBuilder", canonical: str) -> None: + """ + Declare every occurrence of a PARAM slot named `canonical` anywhere in + `group`'s already-composed subtree as build-phase-shared with `group`'s + own top-level `canonical` slot. A no-op if `canonical` occurs nowhere in + the subtree (e.g. OUTLET_MASK when no block happens to reference it under + the current config) - share() itself requires at least one path, so this + only calls it when there is something to share. + + Parameters + ---------- + group : GroupBuilder + Builder whose own `canonical` PARAM slot every found occurrence + collapses into. + canonical : str + PARAM slot name to search for and share. + + Author: B.G (08/2026) + """ + paths = [] + for name, child in group.composed.items(): + paths.extend(find_param_paths(child, canonical, (name,))) + if paths: + group.share(canonical, *paths) diff --git a/pyfastflow/experimental/core/context/compile_closure.py b/pyfastflow/experimental/core/context/compile_closure.py new file mode 100644 index 0000000..a4cd0d5 --- /dev/null +++ b/pyfastflow/experimental/core/context/compile_closure.py @@ -0,0 +1,285 @@ +""" +Taichi/Quadrants compile phase: turns a BoundKernel into a CompiledKernel +by emitting real `ti.func`/`ti.kernel` (or `qd.func`/`qd.kernel`) objects. +Both backends share every line here - only which module `backend` points to +differs - mirroring `_closure_backend.py`'s own Taichi/Quadrants split. + +The `ctx` problem and how this solves it +----------------------------------------- +`ctx` is a template's literal first parameter - `def tmpl(ctx, +i): return ctx.grad(ctx.z.get(i), i)` - not a name spliced into the +template's globals the way `_closure_backend.py`'s `specialize_closure` +works for a Parameter's own device view. That distinction +matters here: `ctx` being a real parameter means it is read via `LOAD_FAST` +bytecode, not `LOAD_GLOBAL` - splicing a value into `__globals__` under the +name `ctx`, the `specialize_closure` technique, would have **no effect at +all**, since a local parameter's bytecode never consults globals for that +name regardless of what sits there. + +The fix used here is a source-level transform, not a globals trick: +`_compile_dropping_ctx` takes the template's own AST (via `capture_template_ +meta`, compile_shared.py's cached inspect.getsource + ast.parse), deletes +`ctx` from the FunctionDef's own +parameter list, unparses the result, and `exec()`s it with `ctx` now bound in +that exec's globals. The body is untouched - every `ctx.foo` reference in it +still reads the name `ctx`, which now resolves via `LOAD_GLOBAL` since the +compiled code object no longer declares it as a local/parameter. Registering +the unparsed source in `linecache` under a synthetic filename before `exec()` +is what lets Taichi/Quadrants re-inspect it via their own `inspect.getsource` +during a later inlined trace (a ti.func is re-traced, not compiled once, each +place it is inlined) - the exact technique `_closure_backend.py`'s +`_fuse_group` already relies on for the same reason, proven working here the +same way. + +The ctx tree +------------ +What `ctx` resolves to, at every level, is a plain python object built once +per compile() by `_build_ctx_node`, walking the FrozenKernel/FrozenHelper's +own composition tree in lock-step with `BoundKernel`'s address tree: + + - one attribute per PARAM slot at that level, holding the bound + Parameter's `device_view()` - `ctx.z.get(i)` reaches it exactly as any + other Taichi/Quadrants-compiled template already does. + - one attribute per composed HELPER slot, holding the **raw compiled + `ti.func`/`qd.func` object itself** - not a wrapper - with that helper's + own children (its PARAM device views and its own composed HELPER + children, recursively) attached as extra attributes on that same + function object. A compiled closure-backend function is an ordinary + python object and accepts arbitrary attribute assignment freely, so + `ctx.grid` is simultaneously callable (`ctx.grid(...)`, invoking the + compiled func directly, never a wrapper) and further + attribute-navigable (`ctx.grid.neighbour(...)`, since `neighbour` was + attached onto the same object as `.grid`'s own attribute). This is built + bottom-up: a composed helper's own children are compiled and attached + before that helper's own template is compiled, since its body may + reference them. + +Every composed FrozenHelper reachable from a BoundKernel is compiled +unconditionally as part of that BoundKernel's compile(), whether or not the +immediate parent's own contract calls it bare - a `ctx.grid.neighbour(...)` +reference needs `neighbour` compiled regardless of whether `grid` itself is +ever called directly, and helper call signatures are fully trusted (no +attempt is made here to prune what nothing in this particular tree happens to +call). + +`ctx.bk`, the reserved backend-intrinsics namespace (bk.py) - `ctx.bk.sqrt`, +`ctx.bk.atan2`, `ctx.bk.cast_u32`, ... - is attached to every node this +module builds, at every level of the ctx tree, not just the root: `bk` is +built once per compile() (`make_closure_bk(backend)`) and threaded through +`_build_ctx_node`'s own recursion, so it is reachable from a deeply composed +private block exactly as it is from the kernel's own template. See bk.py's +module docstring for why this namespace exists and contract.py for why it +never appears as a slot requirement. + +Author: B.G (08/2026) +""" + +import ast +import copy +import linecache +from types import FunctionType +from typing import Any + +from ..pool.base import new_uid +from .bk import make_closure_bk +from .bound import Address, BoundKernel, format_address +from .compile_shared import ( + CompiledKernel, + CompileError, + capture_template_meta, + check_data_signature, + check_legal_accessors, + check_unmet, +) +from .ctx import CTX_PARAM_NAME +from .frozen import FrozenGroup, _Frozen +from .slot import SlotKind + + +def _drop_ctx_param(func_def: ast.FunctionDef, label: str) -> None: + """ + Remove `ctx` from `func_def`'s own parameter list in place - see the + module docstring for why this, not a globals splice, is what makes `ctx` + resolve as a global inside the body that is left untouched. + + Author: B.G (08/2026) + """ + if func_def.args.posonlyargs and func_def.args.posonlyargs[0].arg == CTX_PARAM_NAME: + func_def.args.posonlyargs = func_def.args.posonlyargs[1:] + elif func_def.args.args and func_def.args.args[0].arg == CTX_PARAM_NAME: + func_def.args.args = func_def.args.args[1:] + else: + raise CompileError(f"template {label!r}: first parameter must be {CTX_PARAM_NAME!r}") + + +def _compile_dropping_ctx(template, ctx_obj: Any, label: str) -> FunctionType: + """ + Rebuild `template` with `ctx` removed from its own signature and bound + instead as a global (`ctx_obj`) the body's untouched `ctx.*` references + now resolve against. Registers the rebuilt source in `linecache` under a + synthetic filename, `exec()`s it, and returns the resulting function - + not yet decorated with `backend.func`/`backend.kernel`, see the callers + below. + + `exec()` gives the rebuilt function only the globals dict it is handed - + `template`'s own closure cells (a value captured lexically from an + enclosing factory function: a baked constant, a composed helper + reference, ...) are not carried forward by `__globals__` alone and would + otherwise raise `NameError` the first time the rebuilt body reads that + name. Fixed here by seeding the exec globals with `template.__code__. + co_freevars` zipped against `template.__closure__`'s own cell contents, + laid on top of `template.__globals__` so a free variable wins over a + same-named module global - the lexical binding is what the template's + author actually wrote. A captured value becomes an ordinary global in the + rebuilt function; that is semantically fine here since tracing reads it + once and a device template never writes back into an enclosing scope, but + it does mean two templates built from the same closure with different + captured values are never the same rebuilt function object - each + compile() call mints its own. + + A data argument's own type annotation (a `ti.template()`/`qd.template()` + marker, typically) is a second, distinct case closure cells do not cover: + `def f(x: T): ...` evaluates the name `T` eagerly, in the enclosing + frame, the moment the original `def` statement runs - confirmed + empirically - so `T` is never a `LOAD_DEREF` inside `f`'s own code object + and never appears in `co_freevars`/`__closure__` even when `T` is a local + of an enclosing factory function, unlike a name the body itself reads. + What survives instead is `template.__annotations__` (arg name -> already- + evaluated value), captured by the ORIGINAL `def` statement before this + function ever ran. Re-executing the unparsed source re-evaluates each + annotation expression fresh, in the new exec namespace, so it needs the + same values resolvable under the same names again: for every remaining + (post ctx-drop) parameter whose annotation unparses to a bare name, + that name is seeded into the exec globals from `template.__annotations__ + [that parameter's name]` - the value the original template's own + annotation evaluated to, not a guess. + + Author: B.G (08/2026) + """ + _, tree = capture_template_meta(template) + if tree is None: + raise CompileError(f"template {label!r}: no recoverable source to compile") + body = [n for n in tree.body if isinstance(n, ast.FunctionDef)] + if not body: + raise CompileError(f"template {label!r}: source is not a function definition") + func_def = copy.deepcopy(body[0]) + _drop_ctx_param(func_def, label) + + module = ast.fix_missing_locations(ast.Module(body=[func_def], type_ignores=[])) + source = ast.unparse(module) + filename = f"" + linecache.cache[filename] = (len(source), None, source.splitlines(keepends=True), filename) + + exec_globals: dict[str, Any] = dict(getattr(template, "__globals__", {})) + code_obj = getattr(template, "__code__", None) + closure = getattr(template, "__closure__", None) + if code_obj is not None and closure: + exec_globals.update(zip(code_obj.co_freevars, (cell.cell_contents for cell in closure))) + + orig_annotations = getattr(template, "__annotations__", {}) + all_args = ( + list(func_def.args.posonlyargs) + list(func_def.args.args) + list(func_def.args.kwonlyargs) + ) + for arg in all_args: + if isinstance(arg.annotation, ast.Name) and arg.arg in orig_annotations: + exec_globals[arg.annotation.id] = orig_annotations[arg.arg] + + exec_globals["ctx"] = ctx_obj + code = compile(source, filename, "exec") + exec(code, exec_globals) + return exec_globals[func_def.name] + + +class _CtxNode: + """ + What `ctx` (or one of its composed-helper children) resolves to inside a + specialized template body - a plain attribute bag. See the module + docstring's "The ctx tree" section for what gets attached and why a + composed HELPER child is the raw compiled func itself rather than an + instance of this class. + + Author: B.G (08/2026) + """ + + +def _build_ctx_node(prefix: Address, frozen: _Frozen, bound: BoundKernel, backend: Any, bk: Any) -> _CtxNode: + """ + Recursively build the ctx tree rooted at `frozen` (found at `prefix` in + `bound`'s address tree), compiling every composed HELPER child - bottom + up, so a child's own compiled func exists before its parent's template + (which may call it) is compiled - and attaching each as both a callable + and a further-navigable node on the returned object. See the module + docstring. + + `bk` (bk.py's `make_closure_bk(backend)`, built once per compile() and + threaded through every recursive call) is attached to every node at + every level - the reserved `ctx.bk` namespace is reachable from the + kernel's own root template and from any composed helper's, however deep, + since a private block many levels down is exactly where noise's/visu's + own use of it lives. See bk.py's module docstring. + + Author: B.G (08/2026) + """ + node = _CtxNode() + node.bk = bk + for name in frozen.slots.names(SlotKind.PARAM): + addr = prefix + (name,) + param = bound.value_at(addr) + setattr(node, name, param.device_view()) + + for name in frozen.slots.names(SlotKind.HELPER) | set(frozen.composed): + child_addr = prefix + (name,) + child_frozen = frozen.composed[name] + child_node = _build_ctx_node(child_addr, child_frozen, bound, backend, bk) + if isinstance(child_frozen, FrozenGroup): + # A FrozenGroup has no template of its own to compile - it is a + # passive, non-callable composite (frozen.py). `ctx.` is + # attached exactly as built: navigable (`ctx..`), + # never callable. + setattr(node, name, child_node) + continue + label = format_address(child_addr) + raw = _compile_dropping_ctx(child_frozen.template, child_node, label) + compiled = backend.func(raw) + # child_node's own attributes (its PARAM device views, its own + # composed HELPER children) are copied onto the compiled func object + # itself, so `ctx.` is simultaneously callable (invokes this + # func) and navigable (`ctx..`) - see the module + # docstring. + for attr_name, attr_val in vars(child_node).items(): + setattr(compiled, attr_name, attr_val) + setattr(node, name, compiled) + + return node + + +def compile_kernel(bound: BoundKernel, backend: Any) -> CompiledKernel: + """ + Checks unmet slots and legal PARAM accessors first (compile_shared.py), + then builds the whole ctx tree and compiles the kernel's own template as + `backend.kernel(...)`. + + Parameters + ---------- + bound : BoundKernel + backend : module + `taichi` or `quadrants`. + + Returns + ------- + CompiledKernel + + Author: B.G (08/2026) + """ + check_unmet(bound) + check_legal_accessors(bound) + + frozen = bound.frozen + data_names = check_data_signature(frozen.template, frozen.slots.names(SlotKind.DATA)) + bk = make_closure_bk(backend) + root_node = _build_ctx_node((), frozen, bound, backend, bk) + raw = _compile_dropping_ctx(frozen.template, root_node, "root") + compiled = backend.kernel(raw) + + data_order = [(name,) for name in data_names] + return CompiledKernel(bound, compiled, data_order, needs_launch_dims=False) diff --git a/pyfastflow/experimental/core/context/compile_cupy.py b/pyfastflow/experimental/core/context/compile_cupy.py new file mode 100644 index 0000000..e9a8caa --- /dev/null +++ b/pyfastflow/experimental/core/context/compile_cupy.py @@ -0,0 +1,320 @@ +""" +cupy compile phase: turns a BoundKernel into a CompiledKernel by assembling +CUDA source text and building a `cp.RawModule` from it. + +Reuses cupy_backend.py's pure text/emission utilities - dtype/literal +formatting, the `pf_params` constant-block and `__restrict__`-local +machinery, `__global__`/`__device__` function-name extraction - which know +only about Parameter objects and plain text (see cupy_backend.py's module +docstring for the block's exact shape). The span *resolver* lives here: a +`$ctx.path$` span resolves against one BoundKernel's address tree (bound.py) +- `$ctx.z.get(i)$`, `$ctx.grid.neighbour(i, k)$`, the same grammar +contract.py derives. + +Composed helpers become `__device__` functions +------------------------------------------------ +Every composed FrozenHelper reachable from `bound` gets its own `__device__` +function, unconditionally (mirroring compile_closure.py's reasoning: a +`ctx.grid.neighbour(...)` span needs `neighbour` emitted regardless of +whether `grid` itself is ever called bare). Its C name is derived from its +own full address (`pf_flux_grad_grid_neighbour` for address `flux.grad.grid. +neighbour`), which is unique within one compile by construction (build() +never mints two different composed subtrees under the same address) - no +uid-based mangling needed, unlike `_cupy_blocks.py`'s per-make_grid-call +`new_uid()` tag, since there is exactly one BoundKernel's address tree per +compile here, not several independently-built grids sharing one module. +`_emit_device_func` renames the template's own declared function name to +that address-derived name in the emitted text (the template author's own +choice of name in source is never seen by the caller); a helper already +emitted once in this compile (reachable from two different addresses - +uncommon here since addresses are already unique, but the memo guards a +cycle regardless) is reused, not re-emitted. + +Author: B.G (08/2026) +""" + +import re +from typing import Any + +import cupy as cp + +from .bound import Address, BoundKernel, format_address +from .compile_shared import CompiledKernel, CompileError, check_legal_accessors, check_unmet +from .cupy_backend import ( + _DEVICE_NAME_RE, + _KERNEL_NAME_RE, + _KERNEL_SIG_RE, + _cuda_literal, + _ctype, + _extract_name, + _insert_locals, + _param_argname, + _param_block_source, + _split_args, + _upload_param_block, +) +from .ctx import CTX_PARAM_NAME +from .frozen import FrozenGroup, _Frozen +from .parameter import Parameter +from .slot import SlotKind + +_SPAN_RE = re.compile(r"\$(.*?)\$", re.S) +_CALL_RE = re.compile(r"([\w.]+)\s*(?:\((.*)\))?\s*$", re.S) +_CONSTANT_DECL_RE = re.compile(r"__constant__\s+[\w:\*&]+\s+(\w+)\s*(?:\[[^\]]*\])?\s*=") + + +class _EmitState: + """ + Everything one compile() accumulates across every `__device__`/ + `__global__` body it parses - the pointer registry and its + first-encounter local-index map (handed straight to cupy_backend.py's + `_param_block_source`/`_upload_param_block`/`_insert_locals`), and the + dependency-first, dedup-by-name map of every composed helper's own + `__device__` source. + + Author: B.G (08/2026) + """ + + def __init__(self): + self.registry: dict[int, dict] = {} + self.local_index: dict[int, int] = {} + self.device_srcs: dict[str, "str | None"] = {} + # Finalization order, distinct from `device_srcs`' own insertion + # order: a name is reserved (`= None`) in `device_srcs` *before* + # `_ensure_emitted` recurses into whatever it calls, so a name's + # position in `device_srcs` itself is caller-before-callee - the + # wrong direction for C, which needs a callee's definition (or at + # least a declaration) above its caller. `emit_order` instead + # records a name only once its own body is fully resolved, i.e. + # child-before-parent, which is what `compile_kernel` must emit in. + self.emit_order: list[str] = [] + + +def _register_ptr(state: _EmitState, param: Parameter, write: bool, local_ptrs: dict[int, dict]) -> str: + uid = param.uid + entry = state.registry.get(uid) + if entry is None: + entry = {"ctype": _ctype(param.dtype), "write": False, "array": param.get().data} + state.registry[uid] = entry + if write: + entry["write"] = True + if uid not in state.local_index: + state.local_index[uid] = len(state.local_index) + local = local_ptrs.get(uid) + if local is None: + local_ptrs[uid] = {"ctype": entry["ctype"], "write": write} + elif write: + local["write"] = True + return _param_argname(param, state.local_index) + + +def _expand_param(state: _EmitState, param: Parameter, method: str, call_args: list[str], local_ptrs: dict) -> str: + if method == "get": + if param.mode == "const": + return _cuda_literal(param.get()) + argname = _register_ptr(state, param, write=False, local_ptrs=local_ptrs) + if param.mode == "scalar": + return f"{argname}[0]" + idx = call_args[0] if call_args else "0" + return f"{argname}[{idx}]" + if param.mode == "const": + raise CompileError(f"{param.name}: const parameter is read-only") + if len(call_args) != 2: + raise CompileError(f"{param.name}: set_node(node, value) takes two arguments") + argname = _register_ptr(state, param, write=True, local_ptrs=local_ptrs) + node, val = call_args + return f"{argname}[0] = {val}" if param.mode == "scalar" else f"{argname}[{node}] = {val}" + + +def _c_name(addr: Address) -> str: + return "pf_" + "_".join(addr) + + +def _resolve_chain( + state: _EmitState, + segs: list[str], + call_args: list[str], + argstr: "str | None", + prefix: Address, + frozen: _Frozen, + bound: BoundKernel, + local_ptrs: dict, +) -> str: + """ + Resolve one span's `ctx.` path, rooted at `frozen`/`prefix` in + `bound`'s address tree - see the module docstring for the two shapes + (PARAM leaf, composed HELPER call/descent). + + Author: B.G (08/2026) + """ + if not segs: + raise CompileError("span '$ctx$' names nothing") + root = segs[0] + addr = prefix + (root,) + + if root in frozen.slots.names(SlotKind.PARAM): + if len(segs) != 2 or segs[1] not in ("get", "set_node"): + raise CompileError( + f"{format_address(addr)!r}: illegal PARAM accessor 'ctx.{'.'.join(segs)}' - " + f"legal accessors are .get(...) and .set_node(...)" + ) + param = bound.value_at(addr) + return _expand_param(state, param, segs[1], call_args, local_ptrs) + + if root in frozen.slots.names(SlotKind.HELPER) or root in frozen.composed: + child_frozen = frozen.composed[root] + if len(segs) == 1: + if isinstance(child_frozen, FrozenGroup): + raise CompileError( + f"{format_address(addr)!r}: a FrozenGroup composite is not callable - " + f"reference one of its members instead (ctx.{'.'.join(segs)}.)" + ) + fname = _ensure_emitted(state, addr, child_frozen, bound) + return f"{fname}({argstr if argstr is not None else ''})" + return _resolve_chain(state, segs[1:], call_args, argstr, addr, child_frozen, bound, local_ptrs) + + raise CompileError(f"{format_address(addr)!r}: no such PARAM/HELPER slot on 'ctx.{'.'.join(segs)}'") + + +def _make_repl(state: "_EmitState", prefix: Address, frozen: "_Frozen", bound: BoundKernel, local_ptrs: dict): + def _repl(match: re.Match) -> str: + cm = _CALL_RE.match(match.group(1).strip()) + if cm is None: + raise CompileError(f"malformed span: ${match.group(1)}$") + path = cm.group(1).split(".") + argstr = cm.group(2) + call_args = _split_args(argstr) if argstr is not None else [] + if path[0] != CTX_PARAM_NAME: + raise CompileError(f"span '${match.group(1)}$' is not ctx-rooted") + return _resolve_chain(state, path[1:], call_args, argstr, prefix, frozen, bound, local_ptrs) + + return _repl + + +def _mangle_constants(body: str, c_name: str) -> str: + """ + Rename every `__constant__` symbol `body` itself declares to a name + derived from `c_name` (this device block's own address-mangled function + name), consistently everywhere it appears in `body` - the declaration + and every use, both already in `body` since this runs on one block's own + text. Extends `_ensure_emitted`'s existing per-address renaming (until + this, applied only to the block's own `__device__` function name) to any + *other* top-level symbol a template happens to declare - a `__constant__` + lookup table backing a runtime-data if-ladder, in practice (grid's own + `delta` block, _cupy_blocks.py) - which needs exactly the same + per-address uniqueness the function name already gets: a FrozenHelper + composed at two different addresses in one compile is emitted twice + (`_ensure_emitted` memoizes by the mangled *function* name, which already + differs per address), and without this, both emissions would declare the + identical `__constant__` symbol name and collide at NVRTC compile time. + + Author: B.G (08/2026) + """ + for orig in dict.fromkeys(_CONSTANT_DECL_RE.findall(body)): + body = re.sub(rf"\b{re.escape(orig)}\b", f"{c_name}_{orig}", body) + return body + + +def _ensure_emitted(state: _EmitState, addr: Address, frozen: _Frozen, bound: BoundKernel) -> str: + """ + This composed helper's own `__device__` C function name, emitting its + source into `state.device_srcs` on first reach (memoized by name, so a + cycle - or the same address reached twice, which cannot currently happen + since addresses are already unique per compile - never re-emits). + + Author: B.G (08/2026) + """ + name = _c_name(addr) + if name in state.device_srcs: + return name + state.device_srcs[name] = None # reserve, guards a helper cycle + orig_match = _DEVICE_NAME_RE.search(frozen.template) + if orig_match is None: + raise CompileError(f"{format_address(addr)!r}: template has no recoverable __device__ function name") + renamed = frozen.template[: orig_match.start(1)] + name + frozen.template[orig_match.end(1) :] + renamed = _mangle_constants(renamed, name) + local_ptrs: dict[int, dict] = {} + body = _SPAN_RE.sub(_make_repl(state, addr, frozen, bound, local_ptrs), renamed) + body = _insert_locals(body, local_ptrs, state.local_index) + state.device_srcs[name] = body + state.emit_order.append(name) + return name + + +def _check_cupy_data_signature(template: str, declared_names: set[str]) -> list[str]: + """ + The `__global__` kernel's own C parameter names, in source order, + validated to be exactly `declared_names` as a set. cupy's text-source + counterpart to compile_shared.py's check_data_signature (there is no + python `inspect.signature` to read here). + + Author: B.G (08/2026) + """ + match = _KERNEL_SIG_RE.search(template) + if match is None: + raise CompileError("template has no recoverable __global__ signature") + argstr = match.group(2).strip() + parts = _split_args(argstr) if argstr else [] + names = [p.strip().rsplit(None, 1)[-1].lstrip("*") for p in parts] + if set(names) != declared_names: + raise CompileError( + f"__global__ signature declares data argument(s) {names}, wire_data() declared " + f"{sorted(declared_names)} - these must match exactly" + ) + return names + + +def compile_kernel(bound: BoundKernel, *, grid: Any = None, block: Any = None) -> CompiledKernel: + """ + Compile `bound` to a cupy `cp.RawModule`. Checks unmet slots and legal + PARAM accessors first (compile_shared.py), then emits the kernel's own + `__global__` body and every composed helper's `__device__` source it + reaches, in dependency order, assembles the `pf_params` constant block, + builds the module and uploads the block. + + Parameters + ---------- + bound : BoundKernel + grid, block : optional + Launch-dimension defaults for the returned CompiledKernel (see + CompiledKernel.__call__) - cupy has no auto-ranging equivalent to + Taichi/Quadrants, so a caller must supply them here or at call time. + + Returns + ------- + CompiledKernel + + Author: B.G (08/2026) + """ + check_unmet(bound) + check_legal_accessors(bound) + + frozen = bound.frozen + template = frozen.template + data_names = _check_cupy_data_signature(template, frozen.slots.names(SlotKind.DATA)) + + state = _EmitState() + kernel_name = _extract_name(_KERNEL_NAME_RE, template, "__global__") + local_ptrs: dict[int, dict] = {} + kbody = _SPAN_RE.sub(_make_repl(state, (), frozen, bound, local_ptrs), template) + kbody = _insert_locals(kbody, local_ptrs, state.local_index) + if 'extern "C"' not in kbody: + kbody = kbody.replace("__global__", 'extern "C" __global__', 1) + + source = "\n".join( + [_param_block_source(state.registry, state.local_index)] + + [state.device_srcs[name] for name in state.emit_order] + + [kbody] + ) + + module = cp.RawModule(code=source) + _upload_param_block(module, state.registry, state.local_index) + raw = module.get_function(kernel_name) + + def launch(*args, grid, block): + g = (grid,) if isinstance(grid, int) else tuple(grid) + b = (block,) if isinstance(block, int) else tuple(block) + return raw(g, b, tuple(args)) + + data_order = [(name,) for name in data_names] + return CompiledKernel(bound, launch, data_order, needs_launch_dims=True, grid=grid, block=block) diff --git a/pyfastflow/experimental/core/context/compile_shared.py b/pyfastflow/experimental/core/context/compile_shared.py new file mode 100644 index 0000000..7b43133 --- /dev/null +++ b/pyfastflow/experimental/core/context/compile_shared.py @@ -0,0 +1,316 @@ +""" +Backend-agnostic pieces of the compile phase: the two checks every +backend's `BoundKernel.compile()` runs before emitting anything, and +CompiledKernel - the callable every backend's compile() returns. + +Legal PARAM accessors +---------------------- +The build and bind phases deliberately leave "does this chain spell a real +accessor" unenforced - only the emission layer knows what a PARAM slot may +legally do in device code, per slot.py's own module docstring. Settled here, +identically on every +backend (Taichi, Quadrants, cupy all resolve a PARAM chain the same way, so +there is one answer, not three): a PARAM-rooted chain must be exactly +`(name, "get")` or `(name, "set_node")` - two segments, nothing else. A bare +`ctx.z` with no accessor, a chain with a third segment, or any other method +name is illegal. `set_node` is further illegal against a slot currently bound +to a const-mode Parameter (const is baked into generated code as a literal; +there is nothing to write). check_legal_accessors walks the whole composition +tree - the kernel's own contract plus every composed FrozenHelper's own, +recursively - and raises naming the exact address and the exact chain, before +any backend touches source generation. This is a compile()-time convenience +on top of what the backends already refuse structurally on their own - +Taichi/Quadrants' ClosureParamDeviceView simply has no `set_node` attribute +for a const Parameter (AttributeError at trace time), cupy's span expander +only implements `get`/`set_node` - check_legal_accessors exists so the error +arrives before any tracing/emission starts, naming the address plainly +instead of surfacing as a trace-time AttributeError or a malformed-span +ValueError deep in generated text. + +Unmet slots +----------- +check_unmet raises listing every address BoundKernel.unmet() reports, +formatted exactly as inspect() would show them - pasteable, not paraphrased. + +Data argument signature +------------------------ +check_data_signature/the cupy-specific text equivalent in compile_cupy.py +validate that a template's own declared data arguments (its python +parameters after `ctx`, or a cupy `__global__`'s own C parameter names) +match this kernel's wire_data() slots by name exactly - this is what lets +CompiledKernel resolve DATA addresses to launch-argument *positions* without +either side (template author, wire_data caller) tracking an implicit order +by hand. + +CompiledKernel +-------------- +What every backend's compile() returns: an immutable snapshot around a +resolved data-address order and a `launch` callable. Data is bound by +address, never passed positionally at call time - `swap(addr, buf)` re-points +one DATA address's current buffer with a plain dict write, no re-trace, no +recompile, exactly the ping-pong cost `z`/`z_prime` needs. `__call__` reads +whatever `swap()` currently holds for every address, in the fixed order +data_order was built with, and passes those positionally to `launch` - this +positional pass-through is what makes swap() free: the *compiled* kernel +itself is never touched, only a python dict entry. + +Author: B.G (08/2026) +""" + +import ast +import inspect +import textwrap +from functools import lru_cache +from typing import Any, Callable + +from .bound import Address, BindError, _Bound, format_address, parse_address +from .ctx import CTX_PARAM_NAME +from .slot import SlotKind + +_LEGAL_PARAM_ACCESSORS = ("get", "set_node") + + +@lru_cache(maxsize=256) +def capture_template_meta(template) -> tuple[str | None, ast.AST | None]: + """ + Return (source_text, ast) for a template. A python def is introspected; a + raw string (CUDA source) is kept verbatim and has no AST. The source + returned as `source_text` is `inspect.getsource`'s own indentation + (whatever the def's nesting produced); the source parsed into `tree` is + dedented first, so a nested def's indented body parses instead of + raising an IndentationError - a no-op for a module-level def, which is + already at column 0. A def with no recoverable source at all (a lambda, + an exec'd function) still comes back with `tree = None`. + + Cached because compile_closure.py's `_compile_dropping_ctx` asks once per + template to recover its AST, and a miss costs an inspect.getsource plus a + parse. The tree handed back is shared by every caller of this function + against that template: treat it as read-only. + + The cache key is the template object itself, so the bound size matters - + unbounded, it would pin every dynamically generated template and every CUDA + source string for the life of the process. An eviction only costs one + re-parse. + + Parameters + ---------- + template : callable or str + A python def, or raw CUDA source text. + + Returns + ------- + source_text : str or None + None only if a python def had no recoverable source. + tree : ast.AST or None + None for CUDA source text, or a python def with no recoverable/ + parseable source. + + Author: B.G (07/2026) + """ + if isinstance(template, str): + return template, None + try: + source = inspect.getsource(template) + except (OSError, TypeError): + return None, None + try: + tree = ast.parse(textwrap.dedent(source)) + except SyntaxError: + tree = None + return source, tree + + +class CompileError(Exception): + """ + Raised by the compile phase: unmet slots, an illegal PARAM accessor, + a data-argument signature mismatch between a template and its wire_data() + slots, or any other structural problem caught before/while emitting + device code. Every case names the exact address (and, for accessors, the + exact chain) involved. + + Author: B.G (08/2026) + """ + + +def check_unmet(bound: _Bound) -> None: + """ + Raise, listing every still-unbound address exactly as inspect() would + print it, if `bound` has any. Every concrete `compile()` calls this + first. + + Author: B.G (08/2026) + """ + missing = bound.unmet() + if missing: + listing = ", ".join(format_address(a) for a in missing) + raise CompileError(f"compile: unbound slot(s): {listing}") + + +def check_legal_accessors(bound: _Bound) -> None: + """ + Walk `bound`'s whole composition tree (the kernel's own frozen object, + then every composed FrozenHelper, recursively) and raise on the first + illegal PARAM accessor found - see the module docstring for the exact + legal set and why it is identical on every backend. + + Author: B.G (08/2026) + """ + _walk_accessors((), bound.frozen, bound) + + +def _walk_accessors(prefix: Address, frozen, bound: _Bound) -> None: + param_names = frozen.slots.names(SlotKind.PARAM) + for chain in frozen.contract.chains: + root = chain[0] + if root not in param_names: + continue + addr = prefix + (root,) + if len(chain) != 2 or chain[1] not in _LEGAL_PARAM_ACCESSORS: + raise CompileError( + f"{format_address(addr)!r}: illegal PARAM accessor 'ctx.{'.'.join(chain)}' - " + f"legal accessors are .get(...) and .set_node(...)" + ) + if chain[1] == "set_node": + value = bound.value_at(addr) + if value is not None and getattr(value, "mode", None) == "const": + raise CompileError( + f"{format_address(addr)!r}: 'ctx.{'.'.join(chain)}' - set_node against a " + f"const-mode PARAM slot is illegal (const is a baked-in literal, nothing " + f"to write)" + ) + + for name in frozen.slots.names(SlotKind.HELPER) | set(frozen.composed): + _walk_accessors(prefix + (name,), frozen.composed[name], bound) + + +def check_data_signature(template, declared_names: set[str]) -> list[str]: + """ + A python template's own data-argument names, in declaration order (its + parameters after `ctx`), validated to be exactly `declared_names` as a + set - not a subset, not a superset. The order returned is what + CompiledKernel resolves DATA addresses against for positional launch. + + Parameters + ---------- + template : callable + declared_names : set[str] + The kernel's own wire_data() slot names. + + Returns + ------- + list[str] + `template`'s data-argument names, in declaration order. + + Raises + ------ + CompileError + `template`'s first parameter is not `ctx`, or its data-argument + names do not match `declared_names` exactly. + + Author: B.G (08/2026) + """ + label = getattr(template, "__name__", "?") + params = list(inspect.signature(template).parameters) + if not params or params[0] != CTX_PARAM_NAME: + raise CompileError(f"template {label!r}: first parameter must be {CTX_PARAM_NAME!r}") + data_params = params[1:] + if set(data_params) != declared_names: + raise CompileError( + f"template {label!r} declares data argument(s) {data_params}, wire_data() " + f"declared {sorted(declared_names)} - these must match exactly" + ) + return data_params + + +class CompiledKernel: + """ + The immutable callable every backend's `BoundKernel.compile()` returns. + See the module docstring. + + Author: B.G (08/2026) + """ + + def __init__( + self, + bound: _Bound, + launch: Callable, + data_order: list[Address], + *, + needs_launch_dims: bool = False, + grid: Any = None, + block: Any = None, + ): + self._bound = bound + self._launch = launch + self._data_order = list(data_order) + self._data: dict[Address, Any] = {addr: bound.value_at(addr) for addr in self._data_order} + self._needs_launch_dims = needs_launch_dims + self._grid = grid + self._block = block + + @property + def data_order(self) -> list[Address]: + """This kernel's DATA addresses, in the fixed positional order `launch` is called with.""" + return list(self._data_order) + + def data_at(self, addr: "Address | str") -> Any: + """The buffer `swap()` currently has parked at `addr` (or the value it was compiled with).""" + a = parse_address(addr) if isinstance(addr, str) else tuple(addr) + if a not in self._data: + raise BindError(f"data_at: {format_address(a)!r} is not one of this compiled kernel's data addresses") + return self._data[a] + + def swap(self, addr: "Address | str", buf: Any) -> "CompiledKernel": + """ + Re-point DATA address `addr` at `buf` - a dict write, nothing else: + no re-trace, no recompile, the compiled kernel itself is untouched. + + Parameters + ---------- + addr : Address or str + buf : Any + Validated against the slot's declared dtype + (wire_data(..., dtype=...)), if one was declared - the same + check bind() runs. + + Author: B.G (08/2026) + """ + a = parse_address(addr) if isinstance(addr, str) else tuple(addr) + if a not in self._data: + raise BindError( + f"swap: {format_address(a)!r} is not one of this compiled kernel's data " + f"addresses ({', '.join(format_address(x) for x in self._data_order)})" + ) + info = self._bound.slot_info(a) + if info.dtype is not None: + obj_dtype = getattr(buf, "dtype", None) + if obj_dtype is not None and obj_dtype != info.dtype: + raise BindError( + f"swap({format_address(a)!r}, ...): dtype mismatch, slot declares " + f"{info.dtype}, got {obj_dtype}" + ) + self._data[a] = buf + return self + + def __call__(self, *, grid: Any = None, block: Any = None): + """ + Launch with whatever `swap()` currently holds for every DATA + address, in `data_order`. `grid`/`block` matter only on a backend + that needs explicit launch dimensions (cupy); ignored otherwise. + + Author: B.G (08/2026) + """ + args = [self._data[addr] for addr in self._data_order] + if not self._needs_launch_dims: + return self._launch(*args) + g = grid if grid is not None else self._grid + b = block if block is not None else self._block + if g is None or b is None: + raise CompileError( + "this compiled kernel needs explicit launch dimensions - pass grid=/block= " + "to compile() or to this call" + ) + return self._launch(*args, grid=g, block=b) + + def __repr__(self) -> str: + return f"CompiledKernel(data={[format_address(a) for a in self._data_order]})" diff --git a/pyfastflow/experimental/core/context/contract.py b/pyfastflow/experimental/core/context/contract.py new file mode 100644 index 0000000..db1af47 --- /dev/null +++ b/pyfastflow/experimental/core/context/contract.py @@ -0,0 +1,331 @@ +""" +A composite's structural contract - the set of ctx.* chains its template +actually touches - derived, never hand-authored, from the template's own +source. + +Two extractors read that source and produce the same kind of thing, a +Contract: a frozen set of chains, each chain the maximal tuple of dotted +segments following `ctx` - every reference through ctx counts, whether or not +it is ever called. `ctx.z.get(i)` and a hypothetical bare `ctx.grid.nx` used +as a value both contribute a chain; nothing about the grammar privileges a +call over any other use. What a chain's own trailing segment means for +device emission - is it a required `.get`, is a bare non-call reference even +legal - is a backend question, deliberately not decided here (see slot.py's +module docstring on PARAM access being uniform across modes, and the compile +phase's `check_legal_accessors`, compile_shared.py, for where "is this +spelling legal" actually gets enforced). + +extract_python_contract(template) a python def, read by inspect.getsource + + ast - STATIC ANALYSIS, the template is + never called. A Taichi/Quadrants template + cannot run outside kernel-trace context, + so this is the only sound way to find out + what it touches, and it is also why `ctx` + must appear only in the template body + itself - a lambda, an exec'd function, or + `ctx` passed into a nested python def the + walk cannot see through, all raise + ContractError rather than silently + under-reporting the contract. +extract_cupy_contract(source) CUDA text already carrying `$...$` spans. + The template is an f-string fully + materialised at build time, so this reads + the final string directly - no + inspect.getsource, and none of the + python surface's restrictions apply, + since there is no python callable to lose + sight of in the first place. + +`ctx.bk` (RESERVED_BK_NAME, bk.py) is grammar the python surface recognises +and drops rather than records: a chain rooted at `bk` is the reserved +backend-intrinsics namespace (`ctx.bk.sqrt(x)`, ...), never a slot +requirement, so extract_python_contract never adds one to the returned +Contract - see bk.py's module docstring for the full mechanism, including why +this namespace exists at all and why cupy's own extractor is deliberately +left untouched (`ctx.bk` is not part of the cupy template surface). + +Contract.check_root(root, provided) is the candidate-check compose() (builder. +py) runs once a template's contract is known: for every chain this contract +requires under `root` (e.g. `grid.neighbour`), the composed candidate must +provide the next segment (`neighbour`) among its own top-level names, or this +raises naming exactly what is missing and what the candidate offers instead. + +Author: B.G (08/2026) +""" + +import ast +import inspect +import re +import textwrap +from typing import Callable + +from .bk import RESERVED_BK_NAME +from .ctx import CTX_PARAM_NAME + +Chain = tuple[str, ...] + + +class ContractError(Exception): + """ + Raised when a template's source cannot be turned into a contract (no + recoverable source, `ctx` not the first parameter, a malformed span), or + when a derived contract is checked against a candidate that does not + satisfy it. + + Author: B.G (08/2026) + """ + + +# --------------------------------------------------------------------------- +# Contract +# --------------------------------------------------------------------------- + + +class Contract: + """ + A composite's derived structural contract: the frozen set of ctx.* chains + its template touches. See the module docstring. + + Author: B.G (08/2026) + """ + + def __init__(self, chains: frozenset[Chain]): + self._chains = frozenset(chains) + + @property + def chains(self) -> frozenset[Chain]: + """Every chain this contract requires, as a segment tuple each.""" + return self._chains + + @property + def roots(self) -> set[str]: + """The first segment of every chain - the ctx.* names this contract references directly.""" + return {chain[0] for chain in self._chains if chain} + + def check_root(self, root: str, provided: set[str]) -> None: + """ + Verify a composed candidate for slot `root` satisfies every chain + this contract requires under it. + + A chain `("grid", "neighbour")` requires `"neighbour"` to be among + `provided` - the candidate's own top-level names (see frozen.py, + `_Frozen.provides`). A chain of length 1 rooted at `root` (bare + `ctx.root`, no further member) needs nothing from `provided` - the + root itself being composed is enough. + + Parameters + ---------- + root : str + The composed slot name being checked. + provided : set[str] + The candidate's own top-level PARAM/HELPER/composed names. + + Raises + ------ + ContractError + Some chain's next segment is absent from `provided` - names the + first missing member (and how many more, if any) and what the + candidate provides instead. + + Author: B.G (08/2026) + """ + missing = sorted( + {chain[1] for chain in self._chains if len(chain) > 1 and chain[0] == root and chain[1] not in provided} + ) + if not missing: + return + extra = f" (+{len(missing) - 1} more: {', '.join(missing[1:])})" if len(missing) > 1 else "" + raise ContractError( + f"requires {root}.{missing[0]}{extra}, candidate provides {sorted(provided)}" + ) + + def __repr__(self) -> str: + if not self._chains: + return "Contract()" + body = ", ".join("ctx." + ".".join(c) for c in sorted(self._chains)) + return f"Contract({body})" + + def __eq__(self, other) -> bool: + return isinstance(other, Contract) and self._chains == other._chains + + def __hash__(self) -> int: + return hash(self._chains) + + +# --------------------------------------------------------------------------- +# python surface: static AST walk +# --------------------------------------------------------------------------- + + +def _ctx_chain(node: ast.AST) -> Chain | None: + """ + If `node` is an Attribute/Name chain rooted at `ctx`, its segments in + source order (`ctx.grid.neighbour` -> `("grid", "neighbour")`); else None. + + Author: B.G (08/2026) + """ + segments: list[str] = [] + cur = node + while isinstance(cur, ast.Attribute): + segments.append(cur.attr) + cur = cur.value + if isinstance(cur, ast.Name) and cur.id == CTX_PARAM_NAME: + segments.reverse() + return tuple(segments) if segments else None + return None + + +class _ChainVisitor(ast.NodeVisitor): + """ + Collects every maximal ctx.* chain in a template body, called or not. + + Only `visit_Attribute` is overridden: when a node's own chain resolves + all the way down to `ctx` (see `_ctx_chain`), that is by construction the + maximal chain at this point in the tree - a `ctx.grid.neighbour` node's + `.value` is `ctx.grid`, a strict prefix, never a separate reference worth + recording on its own - so this records the chain and does not descend + into `node.value`. A node that does not resolve to ctx falls through to + generic_visit, so a ctx chain nested anywhere within it (a call argument, + a binary operand, ...) is still found by ordinary recursion. Whether the + chain is then the func of a Call or used bare as a value makes no + difference here - both shapes are recorded identically (see the module + docstring). + + A chain rooted at RESERVED_BK_NAME (`ctx.bk.sqrt(x)`, ...) is dropped + instead of recorded - see the module docstring and bk.py. Nothing further + down such a chain needs a visit of its own (it resolves entirely to + Attribute/Name nodes already fully consumed by `_ctx_chain`), so this is + a plain early return, not a call into generic_visit. + + Author: B.G (08/2026) + """ + + def __init__(self): + self.chains: set[Chain] = set() + + def visit_Attribute(self, node: ast.Attribute) -> None: + chain = _ctx_chain(node) + if chain is not None: + if chain[0] == RESERVED_BK_NAME: + return + self.chains.add(chain) + return + self.generic_visit(node) + + +def _get_function_ast(template: Callable) -> ast.FunctionDef: + """ + The single FunctionDef node for `template`'s own source, dedented and + parsed. Raises ContractError - naming the template - if the source + cannot be recovered (a lambda, an exec'd function) or does not parse down + to one function definition. + + Author: B.G (08/2026) + """ + name = getattr(template, "__name__", repr(template)) + try: + source = inspect.getsource(template) + except (OSError, TypeError) as exc: + raise ContractError( + f"template {name!r}: no recoverable source (a lambda or exec'd function cannot " + f"be statically analysed - ctx's contract can only be derived from a def with " + f"real source)" + ) from exc + try: + tree = ast.parse(textwrap.dedent(source)) + except SyntaxError as exc: + raise ContractError(f"template {name!r}: source does not parse: {exc}") from exc + body = [n for n in tree.body if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))] + if len(body) != 1: + raise ContractError( + f"template {name!r}: expected source recoverable to exactly one function " + f"definition, got {len(body)}" + ) + return body[0] + + +def extract_python_contract(template: Callable) -> Contract: + """ + The Contract a python template requires, by static AST walk over its own + source. Never calls `template` - see the module docstring for why. + + Enforces that `ctx` is the template's first parameter (positional or + positional-or-keyword) - a template reaching a `ctx` handed to it under + another name, or not receiving one as its first argument at all, is + rejected here rather than silently producing an empty or wrong contract. + + Parameters + ---------- + template : callable + A plain python `def` with real, statically recoverable source. + + Raises + ------ + ContractError + Source cannot be recovered, or `ctx` is not the first parameter. + + Author: B.G (08/2026) + """ + fn = _get_function_ast(template) + name = getattr(template, "__name__", fn.name) + params = fn.args.posonlyargs + fn.args.args + if not params or params[0].arg != CTX_PARAM_NAME: + got = params[0].arg if params else "(no parameters)" + raise ContractError( + f"template {name!r}: first parameter must be named {CTX_PARAM_NAME!r}, got {got!r}" + ) + visitor = _ChainVisitor() + visitor.visit(fn) + return Contract(frozenset(visitor.chains)) + + +# --------------------------------------------------------------------------- +# cupy surface: span text scan +# --------------------------------------------------------------------------- + +_SPAN_RE = re.compile(r"\$(.*?)\$", re.S) +_PATH_RE = re.compile(r"^([\w.]+)") + + +def extract_cupy_contract(source: str) -> Contract: + """ + The Contract a cupy (CUDA source text) template requires, by scanning its + already-materialised `$...$` spans for ones prefixed `ctx.` - see + compile_cupy.py for the span resolver this reads, and the module + docstring for why this needs no AST and none of the python surface's + restrictions: the template is a plain string by the time this runs, + nothing to lose sight of. + + A span not prefixed `ctx.` (a bound plain value, a bare const name used + outside any span) contributes nothing - only ctx-rooted spans are part of + this template's structural contract. A span's own trailing `(...)` is not + part of the recorded chain and its presence or absence makes no + difference - `$ctx.grid.neighbour(i, k)$` and a hypothetical bare + `$ctx.grid.neighbour$` both record `("grid", "neighbour")` (see the + module docstring: every reference through ctx counts, called or not). + + Parameters + ---------- + source : str + Fully materialised CUDA source text. + + Raises + ------ + ContractError + A `$...$` span's contents do not start with a dotted path. + + Author: B.G (08/2026) + """ + chains: set[Chain] = set() + for match in _SPAN_RE.finditer(source): + inner = match.group(1).strip() + path_match = _PATH_RE.match(inner) + if path_match is None: + raise ContractError(f"malformed span: ${inner}$") + parts = path_match.group(1).split(".") + if parts[0] != CTX_PARAM_NAME: + continue + segments = tuple(parts[1:]) + if segments: + chains.add(segments) + return Contract(frozenset(chains)) diff --git a/pyfastflow/experimental/core/context/ctx.py b/pyfastflow/experimental/core/context/ctx.py new file mode 100644 index 0000000..153c5e6 --- /dev/null +++ b/pyfastflow/experimental/core/context/ctx.py @@ -0,0 +1,130 @@ +""" +CtxProbe: a standalone, inspectable stand-in for the `ctx` a template +receives, defining - structurally, not by execution - the chain grammar both +template surfaces target. + +A template is written as `def tmpl(ctx, i): return ctx.grad(ctx.z.get(i), i)` +(python) or as CUDA text with `$ctx....$` spans (cupy). Both spellings name +the same shape: a dotted attribute path rooted at `ctx`. `ctx.z.get(i)` reads +slot `z`, `ctx.grad(...)` calls slot `grad`, `ctx.grid.neighbour(i, k)` calls +member `neighbour` of composed slot `grid`. Every reference through ctx is +part of the contract - whether or not it is ever called; a bare, uncalled +`ctx.grid.nx` counts exactly as much as a called `ctx.z.get(i)` (see +contract.py's module docstring). Nothing after a call is part of the chain - +the grammar has no case for chaining off a call's return value, and neither +extractor looks for one. + +Contract derivation itself (contract.py) never runs a template to find this +out: a Taichi template cannot be called outside kernel-trace context, so the +python surface is read statically, via an AST walk over the template's own +source, and the cupy surface is read directly off its already-materialised +`$...$` span text. CtxProbe is not part of that pipeline - it exists so the +grammar the two extractors agree on has one place, testable by driving a probe +by hand with no template, no AST and no backend involved: `CtxProbe().grid. +neighbour(1, 2)` records the chain `("grid", "neighbour")` in `.touched`, +exactly the shape `extract_python_contract`/`extract_cupy_contract` would +derive from source spelling the same access - and, matching the extractors' +"maximal chain only" rule, a probe never keeps both a chain and a longer one +that extends it: reaching `ctx.z.get` after having only touched `ctx.z` drops +the shorter entry, since `ctx.z.get` is what was actually referenced by the +time attribute access stopped extending it. + +CTX_PARAM_NAME is the literal name a python template's first parameter must +carry - see contract.py's `extract_python_contract`, which enforces this. +There is no cupy equivalent to enforce: a `$ctx....$` span already says which +name it means in its own text. + +`ctx.bk` (RESERVED_BK_NAME, bk.py) is a second piece of reserved grammar, on +the closure (Taichi/Quadrants) python surface only: the backend-intrinsics +namespace (`ctx.bk.sqrt(x)`, ...), recognised structurally by contract.py and +never a slot a template's Contract requires satisfied - see bk.py's module +docstring for the full mechanism and why it exists. + +Author: B.G (08/2026) +""" + +CTX_PARAM_NAME = "ctx" +"""The reserved first-parameter name every python template must use - see +extract_python_contract in contract.py.""" + + +class _ChainNode: + """ + One in-progress dotted path rooted at a CtxProbe, produced by attribute + access on the probe or on another node. + + Purely structural: `.path` is the tuple of segments walked so far, + `.dotted` the same joined with ".". Every attribute access - not just a + terminating call - records into the owning probe's `.touched` (see the + module docstring); calling a node records the same, already-recorded + path again, which is a no-op. + + Author: B.G (08/2026) + """ + + def __init__(self, probe: "CtxProbe", path: tuple[str, ...]): + self._probe = probe + self._path = path + + @property + def path(self) -> tuple[str, ...]: + return self._path + + @property + def dotted(self) -> str: + return ".".join(self._path) + + def __getattr__(self, name: str) -> "_ChainNode": + if name.startswith("_"): + raise AttributeError(name) + child = _ChainNode(self._probe, self._path + (name,)) + self._probe._record(child._path) + return child + + def __call__(self, *args, **kwargs) -> "_ChainNode": + self._probe._record(self._path) + return self + + def __repr__(self) -> str: + return f"" + + +class CtxProbe: + """ + The `ctx` placeholder, as a real inspectable object rather than only a + static-analysis convention. See the module docstring. + + `.touched` accumulates every maximal chain reached so far - the same set + shape (a set of segment tuples) Contract (contract.py) holds, which is + the point: a probe can be driven by hand to sanity-check the grammar + independently of any AST walk or span scan. "Maximal" is enforced on + every record: a newly recorded chain evicts any already-touched chain it + strictly extends, and is itself dropped instead of recorded if some + already-touched chain already extends it (attribute access happening in + left-to-right order means that second case is rare in practice, but + cheap to guard regardless). + + Author: B.G (08/2026) + """ + + def __init__(self): + self.touched: set[tuple[str, ...]] = set() + + def _record(self, path: tuple[str, ...]) -> None: + if not path: + return + if any(len(existing) > len(path) and existing[: len(path)] == path for existing in self.touched): + return + for existing in [e for e in self.touched if len(e) < len(path) and path[: len(e)] == e]: + self.touched.discard(existing) + self.touched.add(path) + + def __getattr__(self, name: str) -> _ChainNode: + if name.startswith("_") or name == "touched": + raise AttributeError(name) + node = _ChainNode(self, (name,)) + self._record(node._path) + return node + + def __repr__(self) -> str: + return f"CtxProbe(touched={sorted('.'.join(c) for c in self.touched)})" diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py new file mode 100644 index 0000000..f0b035d --- /dev/null +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -0,0 +1,382 @@ +""" +cupy implementation of Parameter, plus the text/emission utilities +compile_cupy.py reuses to compile a BoundKernel to a `cp.RawModule`. + +A template is CUDA source text rather than a python function, since +cp.RawModule compiles source and there is no function whose globals could be +patched. Bound objects are written into that source as `$...$` spans holding a +dotted path, which keeps the in-kernel spelling the same as on the other +backends: + + $p.get(i)$ read parameter p at flat index i + $p.set_node(i,v)$ write parameter p at flat index i + $grid.nx.get(i)$ reach a bag member + $helper(a, b)$ call a bound device helper + +Compiling substitutes each span according to the parameter's mode - a CUDA +literal for const, or a read/write through a pointer for scalar and field. +That pointer never travels as a kernel argument: every scalar/field Parameter +a compilation unit reaches - the kernel's own bindings plus, recursively, its +helpers' - is collected once, deduplicated by uid, into a module-scope +constant block: + + struct pf_params_t { float* p_; const float* p_; ... }; + __constant__ pf_params_t pf_params; + +`` is a per-compilation-unit local index (0, 1, 2, ...) assigned the +first time this compile's traversal reaches a given Parameter, not its +process-global `uid` - `uid` still identifies the Parameter for dedup, cycle +detection and the ptr registry's keys, but never appears in emitted text, so +an unrelated allocation upstream that shifts every uid does not change this +source at all. See compile_cupy.py's `_register_ptr`. + +uploaded once per compile() via cp.RawModule.get_global. A member is `const` +when nothing in the unit writes that parameter, `T*` otherwise. Every +`__global__` and `__device__` function in the module sees the same block, so a +helper reaches a bound Parameter exactly the way its caller does - there is no +argument to thread through and no call site to rewrite. At the top of each +function body one local is declared per pointer that function's own spans +reference: + + const float* __restrict__ p_ = pf_params.p_; + +read (or written, dropping `const`) through for the rest of that body. This is +what keeps a function's own accesses provably non-aliasing to the compiler, +the same guarantee a `__restrict__` kernel argument used to carry - reading +`pf_params.p_` directly, span by span, would lose it. + +A const parameter can also be used bare, outside any span, in which case it +arrives as a `#define`. Only names the source actually mentions are defined, +which keeps macros for common identifiers - N, DIM, EPS, min - from silently +rewriting unrelated code in the translation unit. + +One `cp.RawModule` is built per compilation unit: the constant block, every +`__device__` helper the unit reaches (each emitted once, however many call +sites share it), and the unit's `__global__` kernel. See compile_cupy.py's +module docstring for how a unit's source is assembled and cached. + +Author: B.G (07/2026) +""" + +import re +from typing import Any + +import cupy as cp +import numpy as np + +from .parameter import MODES, Parameter + +_KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") +# the return type is one-or-more tokens, matched non-greedily so the LAST one +# before the parameter list is the function name - `__device__ unsigned int f(` +# names f, not int. +_DEVICE_NAME_RE = re.compile(r"__device__\s+(?:[\w:\*&]+\s+)+?(\w+)\s*\(") +_KERNEL_SIG_RE = re.compile(r"(__global__\s+void\s+\w+\s*\()(.*?)(\))", re.S) + +_CTYPE = { + np.dtype(np.float32): "float", + np.dtype(np.float64): "double", + np.dtype(np.int32): "int", + np.dtype(np.int64): "long long", + np.dtype(np.uint8): "unsigned char", + np.dtype(np.uint32): "unsigned int", +} + + +def _ctype(dtype) -> str: + """ + CUDA scalar type name for a (numpy) dtype. + + Author: B.G (07/2026) + """ + return _CTYPE[np.dtype(dtype)] + + +def _cuda_literal(value) -> str: + """ + Format a resolved const value as a CUDA literal. + + Author: B.G (07/2026) + """ + if isinstance(value, bool): + return "1" if value else "0" + if isinstance(value, (int, np.integer)): + return str(int(value)) + if isinstance(value, (float, np.floating)): + return f"{float(value)}f" + return str(value) + + +def _extract_name(pattern: re.Pattern, template: str, kind: str) -> str: + """ + The `__global__`/`__device__` function's own name, read out of the source + text - that is the entry point cp.RawModule.get_function is looked up by. + + Author: B.G (07/2026) + """ + match = pattern.search(template) + if not match: + raise ValueError(f"could not find a {kind} function name in template source") + return match.group(1) + + +def _split_args(argstr: str) -> list[str]: + """ + Split a call-argument string on top-level commas (respecting nesting). + + Author: B.G (07/2026) + """ + parts, depth, cur = [], 0, "" + for ch in argstr: + if ch in "([{": + depth += 1 + elif ch in ")]}": + depth -= 1 + if ch == "," and depth == 0: + parts.append(cur.strip()) + cur = "" + else: + cur += ch + if cur.strip(): + parts.append(cur.strip()) + return parts + + +def _param_argname(param: Parameter, local_index: dict[int, int]) -> str: + """ + The struct member / local variable name a Parameter's pointer is reached + through - stable for the object's whole lifetime *within this compile* + since it is derived from `local_index[param.uid]`, a per-compilation-unit + index assigned in first-encounter order (see compile_cupy.py's + `_register_ptr`), not from `uid` itself. `uid` still identifies the Parameter for dedup (two + spans reaching the same Parameter under two different handles look up the + same local index and therefore compute the same argname/struct member), + but the emitted name no longer carries the process-global uid, which is + what keeps generated source byte-stable across runs regardless of + allocation order upstream. + + Author: B.G (07/2026) + """ + return f"p_{local_index[param.uid]}" + + +def _insert_locals(body: str, local_ptrs: dict[int, dict], local_index: dict[int, int]) -> str: + """ + Prepend one `__restrict__` local per pointer `body` itself references, + reading through the module's `pf_params` constant block, right after the + function's opening brace. + + Declared `const` unless this body writes that parameter anywhere - kept + per function rather than read off the struct member (which is `const` + only when *no* function in the whole unit writes it), so a function that + only reads a parameter another function in the same unit writes still + gets the non-aliasing benefit of a const-qualified local. + + Ordered by local index ascending (first-encounter order for this compile, + see compile_cupy.py's `_register_ptr`) rather than by uid, so this declaration + block's text does not depend on the process-global uid values a run + happened to assign upstream. + + Author: B.G (07/2026) + """ + if not local_ptrs: + return body + idx = body.find("{") + if idx == -1: + raise ValueError("could not find a function body to insert parameter locals into") + decls = "".join( + f" {'' if e['write'] else 'const '}{e['ctype']}* __restrict__ {_argname_for(local_index[uid])} = pf_params.{_argname_for(local_index[uid])};\n" + for uid, e in sorted(local_ptrs.items(), key=lambda kv: local_index[kv[0]]) + ) + return f"{body[: idx + 1]}\n{decls}{body[idx + 1 :]}" + + +def _argname_for(local_idx: int) -> str: + """ + The struct member / local name for a pointer already assigned local index + `local_idx` in this compile - see _param_argname, which this must stay in + lockstep with. + + Author: B.G (07/2026) + """ + return f"p_{local_idx}" + + +def _param_block_source(registry: dict[int, dict], local_index: dict[int, int]) -> str: + """ + The `pf_params_t` struct and its `__constant__` instance for one + compilation unit's pointer registry - empty when the unit reaches no + scalar/field Parameter, so a unit with only consts and bare helpers emits + no block at all. + + Member order is by local index, ascending - i.e. first-encounter order + during this compile's traversal (see compile_cupy.py's `_register_ptr`), not by + uid. This is what keeps the struct's text (and therefore the whole + generated source) independent of the process-global uid values, so an + unrelated allocation upstream that shifts every uid does not change this + text. _upload_param_block writes pointers in the same order. + + Author: B.G (07/2026) + """ + if not registry: + return "" + members = "".join( + f" {'' if e['write'] else 'const '}{e['ctype']}* {_argname_for(local_index[uid])};\n" + for uid, e in sorted(registry.items(), key=lambda kv: local_index[kv[0]]) + ) + return f"struct pf_params_t {{\n{members}}};\n__constant__ pf_params_t pf_params;\n" + + +def _upload_param_block(module: "cp.RawModule", registry: dict[int, dict], local_index: dict[int, int]) -> None: + """ + Copy the current pointer for every registered Parameter into the module's + `pf_params` constant block, in the same local-index order the struct was + emitted in (see _param_block_source). + + Runs once per compile(), synchronously - safe as an ordinary host->device + copy anywhere a kernel launch would be. + + Author: B.G (07/2026) + """ + if not registry: + return + global_ptr = module.get_global("pf_params") + ptrs = np.array( + [e["array"].data.ptr for _, e in sorted(registry.items(), key=lambda kv: local_index[kv[0]])], + dtype=np.uint64, + ) + view = cp.ndarray(ptrs.shape, dtype=np.uint64, memptr=global_ptr) + view.set(ptrs) + + +class CupyParameter(Parameter): + """ + Parameter backed by a const python value or a pooled CupyDataHandle. + + dtypes are numpy dtypes throughout, so they need no translation. There is + no device_view() either: a parameter reaches device code when the span + parser substitutes it into the source. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None): + """ + Declare and initialize one parameter. "scalar"/"field" modes allocate + pooled storage immediately via `pool`; "const" stays a plain python + value, read bare in a template body as a #define. + + Parameters + ---------- + name : str + dtype : numpy dtype + mode : str + "const", "scalar" or "field". + value : Any + Initial value. + pool : DataPool + Backing store for "scalar"/"field" modes. + n_flat : int, optional + Required for "field" mode - the number of nodes. + + Author: B.G (07/2026) + """ + if mode not in MODES: + raise ValueError(f"{name}: mode must be one of {sorted(MODES)}, got {mode!r}") + + super().__init__() + self.name = name + self.dtype = dtype + self.mode = mode + self._pool = pool + self._const_value: Any = None + self._handle = None + + if mode == "scalar": + self._handle = pool.get_data(dtype, ()) + elif mode == "field": + if n_flat is None: + raise ValueError(f"{name}: field mode requires n_flat") + self._handle = pool.get_data(dtype, (n_flat,)) + + self._store(value) + + def get(self): + """ + The python value for const mode, the backing CupyDataHandle otherwise. + + Author: B.G (07/2026) + """ + return self._const_value if self.mode == "const" else self._handle + + def set(self, value) -> None: + """ + Overwrite the whole value: a device write for scalar, a full + host->device copy for field. const is immutable - see Parameter.set. + + Author: B.G (07/2026) + """ + if self.mode == "const": + raise ValueError( + f"{self.name}: const parameter is immutable; build a new Parameter and " + f"replace() it into the bag, then recompile" + ) + self._store(value) + + def _store(self, value) -> None: + """ + Write `value` according to the mode, with no immutability check - the + one path that may set a const, used by __init__ to place its initial + value. + + Author: B.G (07/2026) + """ + if self.mode == "const": + self._const_value = np.dtype(self.dtype).type(value).item() + elif self.mode == "scalar": + self._handle.data[...] = value + else: # field + arr = np.asarray(value, dtype=self.dtype).reshape(-1) + self._handle.from_numpy(arr) + + def set_node(self, node, value) -> None: + """ + Host-side single-cell write. scalar ignores node; const is read-only. + + Author: B.G (07/2026) + """ + if self.mode == "const": + raise ValueError(f"{self.name}: const parameter is read-only") + if self.mode == "scalar": + self._handle.data[...] = value + else: # field + self._handle.data[node] = value + + def read(self): + """ + Host-side scalar read - see Parameter.read for the contract. dtypes + are numpy dtypes already here, so no translation is needed. + + Author: B.G (07/2026) + """ + if self.mode == "const": + return self._const_value + if self.mode == "field": + raise ValueError( + f"{self.name}: read() is for scalar/const only; a field is not meant to be " + f"read back to the host as a whole" + ) + return np.dtype(self.dtype).type(self._handle.data.get()).item() + + def destroy(self) -> None: + """ + Return any pooled storage to the pool. const mode owns none, so this + is a no-op there. + + Author: B.G (07/2026) + """ + if self._handle is not None: + self._pool.release_data(self._handle) + self._handle = None + + diff --git a/pyfastflow/experimental/core/context/frozen.py b/pyfastflow/experimental/core/context/frozen.py new file mode 100644 index 0000000..7ac28b0 --- /dev/null +++ b/pyfastflow/experimental/core/context/frozen.py @@ -0,0 +1,227 @@ +""" +FrozenKernel / FrozenHelper: what ingest() (builder.py) hands back - the +immutable, value-like result of a KernelBuilder/HelperBuilder's build phase. + +Both are produced only by KernelBuilder.ingest() / HelperBuilder.ingest() +(builder.py), never constructed directly. Each holds: + + template the ingested template, unchanged (a python def or CUDA text). + slots a SlotGroup snapshot (slot.py) - this builder's own wired + PARAM/HELPER/DATA slots, frozen at the size they had when + ingest() ran. + composed {name: FrozenKernel|FrozenHelper} - the already-frozen + sub-structures compose()d in during build, by identity: the very + object handed to compose() is what sits here, never a copy. + contract the Contract (contract.py) derived from `template` at ingest + time. + split {composed_name: frozenset(relative Address)} - which of a + composed FrozenGroup's own shared() paths (see FrozenGroup, + below) this object's own compose(name, frozen, split=[...]) + call opted back out of that group's default collapse, keyed by + the composed slot name they were declared under. Empty for a + composed child that either is not a FrozenGroup or was composed + with no `split=`. See bound.py's `_walk_group`/`_walk_group_ + subtree` for where this is actually consulted - build() time, + the only point split is decided (see GroupBuilder.share()'s own + docstring, builder.py). + shared {canonical PARAM slot name (wired directly on this object): + frozenset(relative Address)} - this object's own build-phase + sharing declarations (`_Builder.share()`, builder.py). Every + frozen object carries this, not only a FrozenGroup: a + FrozenKernel's own `.shared` matters when it is reached + directly as build()'s top-level argument, a FrozenHelper's + when it is composed as someone else's child - see bound.py's + module docstring for the mechanism both are walked with. + +Nothing here is a recipe any more - a frozen object is done being built. +Mutability alternates through the scheme this module is one step of: builder +mutable -> frozen builder immutable (here) -> bound object mutable -> +compiled callable immutable. `__setattr__` raises FrozenBuilderError +unconditionally after construction, so any code path that tries to poke a new +value into a frozen object - rather than building a new one - fails loudly +and by name. + +A frozen object is shared, not copied: compose() the same FrozenHelper into +two different builders and both results hold that one object, checked by +identity anywhere sameness matters (uid, `is`). This is what lets one grid +neighbour helper, built once, back eighty different kernels without eighty +copies of its recipe. + +`.provides` is what a *further-out* compose() sees when this object is itself +composed one level up: the set of this object's own top-level PARAM/HELPER +slot names, plus its own composed root names. DATA slots are excluded - a +DATA slot is never reached through `ctx.*` (see slot.py), so it is not part +of what a chain like `outer.this.member` could ever ask this object to +provide. + +`.build()` is the entry point into the bind phase (bound.py): it walks +this object's whole composition tree - recursing into every composed +FrozenHelper/FrozenKernel in turn - and mints one independently-bindable +slot per full dotted path it finds, returning a BoundKernel/BoundHelper. This +is where "one FrozenHelper composed into eighty kernels is one frozen object +but eighty independently-bindable slot sets" actually happens: `build()` +never mutates `self` (nothing here could - see `__setattr__` above) and +allocates a fresh bind-time table on every call. See bound.py's module +docstring for the walk itself, the address grammar, and why every wired +HELPER slot must already be composed by the time `build()` runs (deferred +here rather than in `ingest()` - see bound.py for the reasoning). + +Author: B.G (08/2026) +""" + +from typing import Any + +from ..pool.base import new_uid +from .contract import Contract +from .slot import SlotGroup, SlotKind + + +class FrozenBuilderError(Exception): + """ + Raised on any attempt to mutate a frozen object - a FrozenKernel/ + FrozenHelper directly, or a KernelBuilder/HelperBuilder that has already + been ingest()-ed (see builder.py's `_check_mutable`). + + Author: B.G (08/2026) + """ + + +class _Frozen: + """ + Shared base of FrozenKernel/FrozenHelper. Not instantiated directly - see + the module docstring. + + Author: B.G (08/2026) + """ + + def __init__( + self, + template: Any, + slots: SlotGroup, + composed: dict[str, "_Frozen"], + contract: Contract, + split: "dict[str, frozenset] | None" = None, + shared: "dict[str, list] | None" = None, + ): + object.__setattr__(self, "template", template) + object.__setattr__(self, "slots", slots) + object.__setattr__(self, "composed", dict(composed)) + object.__setattr__(self, "contract", contract) + object.__setattr__(self, "split", {k: frozenset(v) for k, v in (split or {}).items()}) + object.__setattr__(self, "shared", {k: frozenset(v) for k, v in (shared or {}).items()}) + object.__setattr__(self, "_uid", new_uid()) + + @property + def uid(self) -> int: + """ + Process-wide identity assigned at construction, from the same + counter as Parameter/Bag (parameter.py, bag.py). Two references to + one FrozenKernel/FrozenHelper share a uid; composing "the same" + frozen object into two builders never changes it. + + Author: B.G (08/2026) + """ + return self._uid + + @property + def provides(self) -> set[str]: + """ + This object's own top-level PARAM/HELPER slot names, plus its own + composed root names - what a compose() one level further out checks + a chain's next segment against. See the module docstring. + + Author: B.G (08/2026) + """ + return self.slots.names(SlotKind.PARAM) | self.slots.names(SlotKind.HELPER) | set(self.composed) + + def build(self) -> "Any": + """ + Walk this object's whole composition tree and return a + BoundKernel/BoundHelper minting one independently-bindable slot per + full dotted path. See the module docstring and bound.py. + + Imported locally to avoid a module-level import cycle (bound.py + itself imports FrozenKernel/FrozenHelper from here, to tell which of + the two `build()` produces). + + Author: B.G (08/2026) + """ + from .bound import build as _build + + return _build(self) + + def __setattr__(self, name: str, value: Any) -> None: + raise FrozenBuilderError( + f"{type(self).__name__}(uid={self._uid}) is frozen and cannot be mutated - " + f"build a new {type(self).__name__} instead" + ) + + def __delattr__(self, name: str) -> None: + raise FrozenBuilderError( + f"{type(self).__name__}(uid={self._uid}) is frozen and cannot be mutated - " + f"build a new {type(self).__name__} instead" + ) + + def __repr__(self) -> str: + return f"{type(self).__name__}(uid={self._uid}, provides={sorted(self.provides)})" + + +class FrozenKernel(_Frozen): + """ + The frozen result of a KernelBuilder's ingest(). See the module + docstring. + + Author: B.G (08/2026) + """ + + +class FrozenHelper(_Frozen): + """ + The frozen result of a HelperBuilder's ingest(). See the module + docstring. + + Author: B.G (08/2026) + """ + + +class FrozenGroup(_Frozen): + """ + The frozen result of a GroupBuilder's close() (builder.py): a + non-callable, navigable composite - PARAM/HELPER slots and composed + sub-structures only, `template` always None, never itself the target of + a device call. `ctx.grid.NX.get(0)` (a PARAM leaf reached through it) and + `ctx.grid.neighbour(i, k)` (a composed HELPER child called through it) + both resolve by ordinary chain recursion through `.slots`/`.composed` + exactly as they would through a FrozenHelper one level in - a + FrozenGroup differs only in having no template of its own to compile, + so `ctx.grid(...)` (calling it bare) is illegal: compile_closure.py's + `_build_ctx_node` attaches its built ctx node directly, uncompiled and + non-callable, instead of wrapping it in `backend.func`; compile_cupy.py's + `_resolve_chain` raises CompileError if a chain ever tries to call it + with no further segment. + + `.contract` is always empty (a group's own build phase derives nothing - + see GroupBuilder.close()), which is exactly right for + compile_shared.check_legal_accessors' walk: it recurses into a + FrozenGroup's own composed children (where real contracts live) but + finds no PARAM chain of the group's own to check. + + `.shared` is build-phase sharing (`_Builder.share()`): {canonical PARAM + slot name (wired directly on this group): frozenset(relative Address)}, + each Address a dotted path into this group's own composed subtree that + reads the "same" quantity as `canonical` - the private per-axis blocks a + public helper composes for its own use (e.g. `neighbour_raw`'s own `row`) + read `NX` again independently of the group's own top-level `NX` slot, + otherwise. bound.py's build() (`_walk_group`/`_walk_group_subtree`) is + what actually acts on this: by default, every Address in `.shared`'s + values is never independently minted at all - only `canonical` is - so + `grid.NX` is the one PARAM address a caller sees and binds, not + `grid.NX` plus every private occurrence. A composer may opt specific + paths back out at compose() time (`split=` - builder.py's `_Builder. + compose()`, recorded as the composing object's own `.split`), re-minting + them as independent addresses again - see bound.py's module docstring + for the full mechanism and why it needs no separate machinery beyond a + build-time redirect table alongside the usual address table. + + Author: B.G (08/2026) + """ diff --git a/pyfastflow/experimental/core/context/host_block.py b/pyfastflow/experimental/core/context/host_block.py new file mode 100644 index 0000000..697bdb9 --- /dev/null +++ b/pyfastflow/experimental/core/context/host_block.py @@ -0,0 +1,224 @@ +""" +HostBlockBuilder / FrozenHostBlock / BoundHostBlock: a leaf builder, in the +same build -> freeze(ingest) -> bind -> compile family as KernelBuilder/ +HelperBuilder (builder.py/frozen.py/bound.py), for host-side python code that +needs to read/write Parameters between device launches - the layer +Sequence's (sequence.py) loop control (`loop`'s `max_times`/`until`) and +inter-block bookkeeping (a depression solver zeroing a counter Parameter +before each pass) run on. + +A host block is a leaf, not a composite: it declares PARAM slots only. +wire_helper()/wire_data() both raise here - a device helper cannot run on the +host (there is nothing to trace it against), and data is never what a host +block reads; it reads Parameters. compose() raises too, for the same +"leaf" reason - there is no sub-structure to attach. + +`ctx` resolves unwrapped +-------------------------- +Where a kernel/helper's `ctx.z` resolves to a Parameter's `device_view()` +(compile_closure.py/compile_cupy.py), a host block's `ctx.z` resolves to the +bound Parameter itself - parameter.py's host-facing surface, `.get()`/ +`.set(value)`/`.read()`, not `.get(node)`/`.set_node(node, value)`. A device +view genuinely cannot run outside kernel-trace context, so there is nothing +else `ctx.z` could mean here. `check_legal_host_accessors` enforces the +matching legal-chain set - `(name, "get"|"set"|"read")`, two segments, one of +those three - the same role compile_shared.py's `check_legal_accessors` plays +for device code, just against a different legal set. + +One class for every backend +----------------------------- +"Compiling" a host block is resolving names, not emitting device code: build +the ctx tree of raw Parameters and return `lambda: template(ctx)`. There is +no Taichi/Quadrants/cupy variant of that, so `BoundHostBlock.compile()` takes +no backend argument (it accepts and ignores `backend`, only to keep a call +site that already carries a `backend` variable from having to special-case +this block kind). + +Author: B.G (08/2026) +""" + +import inspect +from typing import Any + +from .bound import _Bound, _walk +from .builder import _Builder +from .compile_shared import CompileError, check_unmet +from .ctx import CTX_PARAM_NAME +from .frozen import _Frozen +from .slot import SlotKind + +_LEGAL_HOST_ACCESSORS = ("get", "set", "read") + + +class HostBlockBuilder(_Builder): + """ + Builds a host block: PARAM slots only, no HELPER, no DATA, no compose(). + See the module docstring. + + Author: B.G (08/2026) + """ + + def wire_helper(self, name: str) -> "HostBlockBuilder": + """ + Always raises: a host block runs on the host, and a device helper + cannot run there - there is nothing to trace it against. Declare the + Parameters this block needs with wire_param() instead. + + Author: B.G (08/2026) + """ + raise TypeError( + "HostBlockBuilder.wire_helper() is not allowed: a host block is a PARAM-only leaf " + "that runs on the host, and a device helper has no host-side form to call." + ) + + def wire_data(self, name: str, *, dtype: Any = None) -> "HostBlockBuilder": + """ + Always raises: a host block is a PARAM-only leaf. A host block reads + state through bound Parameters (ctx.z.get()/.read()), never through a + trusted call argument the way a kernel's DATA slot does. + + Author: B.G (08/2026) + """ + raise TypeError( + "HostBlockBuilder.wire_data() is not allowed: a host block is a PARAM-only leaf - " + "it reads state through bound Parameters (ctx.z.get()/.read()), never as a call " + "argument." + ) + + def compose(self, name: str, frozen: _Frozen) -> "HostBlockBuilder": + """ + Always raises: a host block is a leaf. There is no sub-structure to + attach - compose a HelperBuilder into a KernelBuilder instead if + device-side composition is what is actually wanted. + + Author: B.G (08/2026) + """ + raise TypeError( + "HostBlockBuilder.compose() is not allowed: a host block is a leaf (PARAM slots " + "only) - there is nothing here for a sub-structure to attach to." + ) + + def ingest(self, template: Any) -> "FrozenHostBlock": + """ + Close out the build phase exactly as KernelBuilder/HelperBuilder.ingest() + does (builder.py): derive and check `template`'s contract, freeze this + builder, return the resulting FrozenHostBlock. `template` must be a + plain python `def tmpl(ctx): ...` - no data arguments, since + wire_data() can never have wired one - checked properly at compile() + time (see BoundHostBlock.compile). + + Author: B.G (08/2026) + """ + self._check_mutable() + slots, composed, contract = self._derive_and_check(template) + self._frozen = True + return FrozenHostBlock(template, slots, composed, contract) + + +class FrozenHostBlock(_Frozen): + """ + The frozen result of a HostBlockBuilder's ingest(). See the module + docstring. `composed` is always empty (compose() raises during build), so + `.provides` reports only this block's own wired PARAM names. + + Author: B.G (08/2026) + """ + + def build(self) -> "BoundHostBlock": + """ + Mint one bindable address per wired PARAM slot and return a + BoundHostBlock. Overrides `_Frozen.build()` (frozen.py), whose own + dispatch (bound.py's `build()`) only knows FrozenKernel/FrozenHelper + and would hand back a BoundHelper here, the wrong type. Reuses + bound.py's `_walk` directly instead - it only needs `.slots`/ + `.composed`, both of which a FrozenHostBlock has, `.composed` always + empty. + + Author: B.G (08/2026) + """ + table: dict = {} + _walk((), self, table) + return BoundHostBlock(self, table) + + +def check_legal_host_accessors(bound: "_Bound") -> None: + """ + Raise on the first PARAM chain in `bound`'s frozen contract that is not + exactly `(name, "get"|"set"|"read")` - the host-facing accessor set + (parameter.py), as opposed to compile_shared.py's device-facing + `(name, "get"|"set_node")`. A host block never composes anything, so - + unlike compile_shared.check_legal_accessors - there is no composition + tree to walk, just this block's own contract. + + Author: B.G (08/2026) + """ + frozen = bound.frozen + param_names = frozen.slots.names(SlotKind.PARAM) + for chain in frozen.contract.chains: + root = chain[0] + if root not in param_names: + continue + if len(chain) != 2 or chain[1] not in _LEGAL_HOST_ACCESSORS: + raise CompileError( + f"{root!r}: illegal host PARAM accessor 'ctx.{'.'.join(chain)}' - legal " + f"accessors on a host block are .get(), .set(...) and .read()" + ) + + +class _HostCtxNode: + """ + What `ctx` resolves to inside a compiled host block's template body - a + plain attribute bag holding this block's Parameters unwrapped. See the + module docstring's "ctx resolves unwrapped" section. + + Author: B.G (08/2026) + """ + + +class BoundHostBlock(_Bound): + """ + The bound result of build()-ing a FrozenHostBlock. See the module + docstring. + + Author: B.G (08/2026) + """ + + def compile(self, backend: "str | None" = None, **kwargs) -> Any: + """ + Resolve this block's ctx (each wired PARAM slot's bound Parameter, + unwrapped) and return `lambda: template(ctx)`. + + Checks unmet slots and legal host accessors first, and that + `template`'s own signature declares exactly one parameter (`ctx` - + wire_data() can never have wired a data argument here, so a + template declaring one more is always a mistake). + + Parameters + ---------- + backend : str, optional + Accepted and ignored - see the module docstring's "One class + for every backend" section for why. + + Returns + ------- + callable + Zero-argument `lambda: template(ctx)`. + + Author: B.G (08/2026) + """ + check_unmet(self) + check_legal_host_accessors(self) + frozen = self._frozen + template = frozen.template + params = list(inspect.signature(template).parameters) + if params != [CTX_PARAM_NAME]: + label = getattr(template, "__name__", "?") + raise CompileError( + f"host block template {label!r} must declare exactly one parameter " + f"({CTX_PARAM_NAME!r}) - got {params}; a host block has no DATA slots, so " + f"there is nothing else for a second parameter to mean" + ) + ctx = _HostCtxNode() + for name in frozen.slots.names(SlotKind.PARAM): + setattr(ctx, name, self.value_at((name,))) + return lambda: template(ctx) diff --git a/pyfastflow/experimental/core/context/parameter.py b/pyfastflow/experimental/core/context/parameter.py new file mode 100644 index 0000000..0aa16ba --- /dev/null +++ b/pyfastflow/experimental/core/context/parameter.py @@ -0,0 +1,299 @@ +""" +Backend-agnostic building blocks for describing GPU work once and compiling it +against Taichi, Quadrants or cupy (or any future one). + +What this is for +---------------- +A physics model rarely needs a new numerical scheme just because one of its +parameters changed shape. Take heat diffusion: the update is identical whether +the diffusion coefficient K is a spatially variable field, a single value the +host retunes between steps, or a constant fixed for the whole run. That choice +matters enormously on a GPU - a compile-time constant costs no memory traffic +and can be folded into the generated code, a field costs a fetch per node - but +it does not change the maths. Boundary conditions and stencils behave the same +way: making a grid periodic alters the neighbour logic, not the scheme built on +top of it. + +Writing one kernel per combination is the obvious way to handle this, and it +becomes unmanageable fast. So instead a template reads a Parameter the same way +whatever its mode, and calls a neighbour helper without knowing which topology +implements it. Which mode, and which helper, is settled at compile time - where +it can still turn into a literal or a specialised routine - and the kernel code +never changes. + +Jargon +---------- +Parameter One named, typed value. Its `mode` says where the value lives: + "const" (baked into the generated code, fixed at + construction), "scalar" (a single device cell, writable) or + "field" (a device array, one value per node, writable). +KernelBuilder The recipe for a launchable kernel: PARAM/HELPER/DATA slots + declared (`.compose()`/`.wire_data()`), a template ingested + (`.ingest()`), then frozen (`.build()` -> FrozenKernel) and + bound (`.build()`/`.bind()` -> BoundKernel) before + `.compile(backend)` emits the real ti.kernel/qd.kernel/CUDA + __global__. +HelperBuilder Same recipe shape, for a device-side helper callable only + from other device code (ti.func, qd.func, CUDA __device__) - + composed into an enclosing kernel's tree rather than compiled + standalone. +FrozenGroup A pure-structure composite with no data of its own (e.g. a + grid's neighbour/distance helpers) - composed into a kernel's + tree the same way a FrozenHelper is, reached in-kernel by + dotted path: ctx.phys.dx.get(i), ctx.ops.neighbour(i). + +There is deliberately no stateful context class: a `make_*` factory returns a +FrozenGroup (pure structure) plus, separately, the concrete Parameters a +caller owns and binds itself - see grid/__init__.py's own module docstring +for that structure/data split. + +Compiling something +------------------- +A template is written once, generically, with `ctx` as its first parameter - +the tree it composes from, PARAM/HELPER slots reached as `ctx.name.get(i)`/ +`ctx.name(...)` - and turned into something callable in three phases: a +KernelBuilder declares its slots and composes children (`.wire_data()`, +`.compose()`, `.ingest()`), `.build()` freezes that recipe into an inert +FrozenKernel, and `.build()` again (this time on the frozen object, via +`.build()`'s own BoundKernel) plus `.bind()` fills every slot with a concrete +Parameter/helper/data buffer before `.compile(backend)` emits the real +ti.kernel/qd.kernel/CUDA source: + + kernel = KernelBuilder().compose("phys", phys_group).wire_data("h_new", "h_old").ingest(update_height) + bound = kernel.build() + bound.bind(("phys", "dx"), dx_p) + compiled = bound.compile("taichi") + compiled(h_new=h_new_field, h_old=h_old_field) # bulk data passed at call time + +See core/context/builder.py, frozen.py and bound.py for the three phases in +full. + +A helper composed anywhere in a kernel's tree - directly or nested under +another composed group - is specialized once as part of that kernel's own +`.compile()`, against that same compile's bindings; reaching the same +FrozenHelper from two addresses in one kernel still specializes it once +(build-phase-shared via `.share()`) if the two occurrences were unified into +one address to begin with, or twice if they were composed independently - +either way there is no standalone compiled Helper object to hold onto +outside of a particular kernel's compile. + +Data at call time, configuration at compile time +------------------------------------------------ +Bound objects are injected into the template body and never appear in the call +signature. A compiled Kernel takes exactly the arguments its template declares, +and that is where bulk data travels - the buffers read and written each step. +Everything that *describes* the problem rather than *being* it - grid spacing, +timestep, gravity, which helper implements the neighbour lookup - is bound. + +Reading a Parameter in device code is uniform across modes: p.get(node) to +read, p.set_node(node, value) to write. + +What a device helper may bind +----------------------------- +A helper binds whatever a kernel binds, in any mode, on every backend. + +On Taichi and Quadrants, bound objects reach device code as globals, and a +helper is traced as part of the kernel that calls it, so alpha.get(i) reads +the same inside a helper as it does in the kernel body. + +On cupy, every scalar/field Parameter a compilation unit reaches - the +kernel's own bindings plus, recursively, every helper's - is collected into +one module-scope `__constant__` block, uploaded once per compile(). Every +`__global__` and `__device__` function compiled into that module sees the +same block, so a helper reaches a bound Parameter exactly the way its caller +does, with no pointer argument to thread through and no call site to rewrite. +See cupy_backend.py's module docstring for the block's exact shape. + +Lifetime of a compiled object +----------------------------- +compile() freezes what it was given: const Parameters are baked in as literals, +scalar and field Parameters as the storage behind their DataHandle. What may +change afterwards follows from that, and splits cleanly along the mode: + + - Writing to a scalar or field Parameter - set(), set_node(), or a device + write from inside a kernel - *is* visible to every kernel that binds it, + including ones compiled beforehand, since they all hold that same storage. + This is the normal way to feed changing data, and it needs no recompile. + - A const Parameter is immutable: its value is fixed at construction and + set() raises. To change one, build a new Parameter, replace() it into the + bag, and recompile whatever bound the old one. + - destroy() returns storage to the pool, which may hand the same buffer out + again. Never destroy a Parameter that a live kernel still binds. This one + is not enforced at runtime. + +So a Parameter's build-time identity - name, dtype, mode, const value - is +fixed at construction, and only its device storage is writable. Which is also +the line to design along: if a quantity changes per step, it is scalar or +field and you simply write it; if changing it demands a recompile, const says +so rather than silently missing the kernels already built. + +Where things live +----------------- +This module defines Parameter and the modes it may take. The rest of the +scheme described above is split by concern: + + builder.py KernelBuilder/HelperBuilder/GroupBuilder - declaring slots and + composing children, the build phase. + frozen.py FrozenKernel/FrozenHelper/FrozenGroup - the inert, immutable + recipe a builder's `.build()` produces. + bound.py BoundKernel - filling a frozen recipe's slots with concrete + Parameters/helpers/data before compile(). + compile_closure.py / compile_cupy.py / compile_shared.py + backend-specific and shared compile-phase logic, turning a + BoundKernel into the real callable. + bag.py Bag and its operators (merge, extract, trim, replace, ...), a + lighter-weight named-collection convenience some factories + still return alongside the above, unrelated to compilation. + +Author: B.G (07/2026) +""" + +from abc import ABC, abstractmethod +from typing import Any + +from ..pool.base import new_uid + +MODES = ("const", "scalar", "field") +"""The storage kinds a Parameter's `mode` may take, common to every backend.""" + + +class Parameter(ABC): + """ + One named, typed value owned by a context. + + `mode` decides where the value lives - "const" in the generated code, + "scalar" in a single device cell, "field" in a device array - and every + backend offers all three (see MODES). + + Two surfaces. From the host: get(), set(value), set_node(node, value). + From device code: device_view(), which returns a backend object whose + .get(node) / .set_node(node, val) let a kernel read and write the + parameter identically whatever its mode. + + Author: B.G (07/2026) + """ + + name: str + dtype: Any + + def __init__(self): + """ + Assign this parameter's process-wide uid and open its `mode` slot. + Concrete backends call this first, then set `self.mode = ...` once as + part of their own __init__ - see the `mode` property below. + + Author: B.G (07/2026) + """ + self._uid = new_uid() + self._mode: str | None = None + + @property + def uid(self) -> int: + """ + Process-wide identity assigned at construction, from the same counter + as every other Parameter, Bag, Helper and pool data handle. Two + references to one Parameter share a uid; two different Parameters + never do, even if they hold equal values. Not stable across processes + and never meant to appear in generated code or a cache key - see the + module docstring, "uid vs handle". + + Author: B.G (07/2026) + """ + return self._uid + + @property + def mode(self) -> str: + """ + Where the value lives - "const", "scalar" or "field". Set once, by + the backend's __init__; reassigning it raises. To change a + parameter's mode, construct a new Parameter and swap it into the bag + in place of this one. + + Author: B.G (07/2026) + """ + return self._mode + + @mode.setter + def mode(self, value: str) -> None: + if self._mode is not None: + raise AttributeError( + f"{getattr(self, 'name', '?')}: Parameter.mode is immutable once set (already " + f"{self._mode!r}); construct a new Parameter and swap it into the bag instead" + ) + self._mode = value + + @abstractmethod + def get(self): + """ + Host-side value: a python scalar for const mode, a DataHandle for scalar/field. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def set(self, value) -> None: + """ + Update the whole parameter value in place, according to its mode: one + device cell for scalar, a full host->device copy for field. The write + lands in storage every kernel binding this parameter already reads, so + no recompile is needed. + + const mode raises: its value is fixed at construction. Build a new + Parameter, replace() it into the bag and recompile - see the module + docstring, "Lifetime of a compiled object". + + Author: B.G (07/2026) + """ + ... + + def set_node(self, node, value) -> None: + """ + Host-side single-cell write. scalar ignores node; const is read-only. + Overridden by concrete backends; device-side writes go through + device_view().set_node instead. + + Author: B.G (07/2026) + """ + raise NotImplementedError(f"{type(self).__name__} does not implement host set_node") + + def device_view(self): + """ + An object whose .get(node) / .set_node(node, val) work inside device + code. Taichi and Quadrants compile one out of ti/qd funcs. cupy leaves + this unimplemented, having no use for it: its parser substitutes + parameters into the source directly. + + Author: B.G (07/2026) + """ + raise NotImplementedError(f"{type(self).__name__} does not implement device_view") + + def read(self): + """ + Host-side scalar read, returned as a plain python value regardless of + mode - unlike get(), which hands back a DataHandle for scalar/field. + + const mode: the stored python value, no device traffic. + scalar mode: a device->host read that synchronizes. That sync is the + whole cost model of any host-driven loop built on top of this - call + it only where a step actually needs the value on the host. + field mode: raises. Reading a whole field back to the host is not + what this is for; use device_view()/get() from device code, or copy + the field explicitly if the host genuinely needs all of it. + + Author: B.G (07/2026) + """ + raise NotImplementedError(f"{type(self).__name__} does not implement read") + + @abstractmethod + def destroy(self) -> None: + """ + Release any backing storage owned by this parameter. Unsafe while a + compiled kernel still binds it - see the module docstring, "Lifetime + of a compiled object". + + Author: B.G (07/2026) + """ + ... + + diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py new file mode 100644 index 0000000..71c6352 --- /dev/null +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -0,0 +1,48 @@ +""" +Quadrants implementation of Parameter. + +A Parameter's device view is a python def specialized by splicing bound +objects into its globals before handing it to qd.func - the mechanism in +_closure_backend.py, shared with Taichi. + +The kernel/helper compile path (KernelBuilder/HelperBuilder -> FrozenKernel -> +BoundKernel.compile("quadrants")) is compile_closure.py, which reaches this +module only for `qd` itself, imported directly there. There, a kernel +template may type its data arguments qd.Tensor to accept either a field- or +ndarray-backed value at call time - Taichi has no equivalent. Field-mode +Parameters, on the other hand, must be field-backed, because they reach +device code as globals and Quadrants rejects an ndarray referenced as a +global inside a func. + +Caching +------- +Leave Quadrants' src_ll fast cache off - do not mark templates compiled here +with @qd.pure or qd.kernel(fastcache=True), even to skip the python-side AST +trace. That cache keys a kernel on its source text, re-read from disk by file +path and line range (_fast_caching/function_hasher.py), together with argument +and config hashes. It never inspects __globals__, and globals are exactly what +distinguishes one specialization from another here. Two compiles of one +template with different bound consts, or a different helper under the same +name, hash identically; a hit then skips AST transformation, so nothing +downstream can catch the mismatch. + +The IR-keyed caches are safe and stay on: Quadrants' own offline_cache, and +Taichi's, hash generated IR, which contains the baked literals and therefore +tells specializations apart. + +Author: B.G (07/2026) +""" + +import quadrants as qd + +from ._closure_backend import ClosureBackendParameter + + +class QuadrantsParameter(ClosureBackendParameter): + """ + Parameter backed by a Quadrants const value or a pooled QuadrantsDataHandle. + + Author: B.G (07/2026) + """ + + _backend = qd diff --git a/pyfastflow/experimental/core/context/routine.py b/pyfastflow/experimental/core/context/routine.py new file mode 100644 index 0000000..897d812 --- /dev/null +++ b/pyfastflow/experimental/core/context/routine.py @@ -0,0 +1,301 @@ +""" +An ordered, device-only sequence of already-built kernels that share one +address space and launch back to back as a single unit. + +`RoutineBuilder` / `FrozenRoutine` / `BoundRoutine` / `CompiledRoutine` follow +the same build -> freeze -> bind -> compile lifecycle as a single kernel +(see builder.py, frozen.py, bound.py), one level up. + +Composing steps +---------------- +`compose(name, frozen_kernel)` appends a step: `name` is its address prefix +and its position in launch order is its position in composition order - there +is no separate ordering call. `FrozenRoutine.order` reads back exactly the +sequence `compose()` was called in. + +Composing the same `FrozenKernel` under two step names runs it twice with +independently bound data at each occurrence: + + rb.compose("diffuse1", diffuse).compose("diffuse2", diffuse) + +gives two addresses (`diffuse1.*`, `diffuse2.*`), two independently bindable +slot sets, and two independent `CompiledKernel` launches after `compile()`. + +`compose()` rejects a `FrozenHelper` - a helper has no standalone launch. +Compose it into a `KernelBuilder` first, then compose that kernel here. + +Addressing +----------- +`build()` walks each step's own composition tree, prefixed with that step's +name: `flux.grad.z` names the `z` PARAM slot of the `grad` helper composed +inside the kernel composed under `flux`. + +Compiling +---------- +`BoundRoutine.compile(backend, **kwargs)` checks this routine's own unmet +slots, then per step: builds a fresh `BoundKernel` from that step's +`FrozenKernel`, copies over whatever is bound at that step's addresses +(`name.*` -> the step's own local addresses), and compiles it. The result is +a `CompiledRoutine` wrapping each step's `CompiledKernel`, in order. + +Per-step launch config +------------------------ +`compose(name, frozen_kernel, launch=None)` takes an optional dict of +compile()-kwargs (cupy's `grid=`/`block=`) applied to that step only, +overriding the routine-level `compile(backend, **kwargs)` defaults key by +key - e.g. `ops`' scan kernels, which need a different launch shape than +the rest of a routine they share with. + +`CompiledRoutine.swap(addr, buf)` routes `name.*` to that step's own +`CompiledKernel.swap()`. Calling a `CompiledRoutine` launches every step in +order, each with whatever its own `swap()` state currently holds. + +Author: B.G (08/2026) +""" + +from typing import Any + +from ..pool.base import new_uid +from .bound import Address, BindError, _Bound, _walk, _walk_group, format_address, parse_address +from .compile_shared import CompileError, check_unmet +from .frozen import FrozenBuilderError, FrozenHelper, FrozenKernel + + +class RoutineBuilderError(Exception): + """ + Raised by the RoutineBuilder build phase: a step name reused, an + attempt to compose a non-FrozenKernel, or a mutation after freeze(). + + Author: B.G (08/2026) + """ + + +class RoutineBuilder: + """ + Collects an ordered set of named kernel steps and freeze()s them into a + FrozenRoutine. See the module docstring. + + Author: B.G (08/2026) + """ + + def __init__(self): + self._uid = new_uid() + self._order: list[str] = [] + self._composed: dict[str, FrozenKernel] = {} + self._launch: dict[str, dict] = {} + self._frozen = False + + @property + def uid(self) -> int: + """Process-wide identity assigned at construction. See Parameter.uid (parameter.py).""" + return self._uid + + def _check_mutable(self) -> None: + if self._frozen: + raise FrozenBuilderError( + f"RoutineBuilder(uid={self._uid}) has already been freeze()-ed and is frozen - " + f"build a new RoutineBuilder instead of reusing this one" + ) + + def compose(self, name: str, frozen_kernel: FrozenKernel, *, launch: "dict | None" = None) -> "RoutineBuilder": + """ + Append a step named `name`, launching `frozen_kernel` at this + position in the routine's launch order. + + Parameters + ---------- + name : str + Address prefix for this step. Must be unique within the routine. + frozen_kernel : FrozenKernel + launch : dict, optional + compile()-kwargs (cupy's `grid=`/`block=`) applied to this step + only, overriding the routine-level default. See the module + docstring's "Per-step launch config" section. + + Author: B.G (08/2026) + """ + self._check_mutable() + if isinstance(frozen_kernel, FrozenHelper): + raise TypeError( + f"compose({name!r}, ...): got a FrozenHelper, not a FrozenKernel - a helper has " + f"no standalone launch and cannot be a routine step. Compose it into a " + f"KernelBuilder first, then compose that kernel's FrozenKernel here." + ) + if not isinstance(frozen_kernel, FrozenKernel): + raise TypeError(f"compose({name!r}, ...): expected a FrozenKernel, got {type(frozen_kernel).__name__}") + if name in self._composed: + raise RoutineBuilderError(f"'{name}' is already composed on this routine") + self._composed[name] = frozen_kernel + self._launch[name] = dict(launch) if launch else {} + self._order.append(name) + return self + + def freeze(self) -> "FrozenRoutine": + """ + Close out the build phase and return the resulting FrozenRoutine. + + Raises + ------ + RoutineBuilderError + No step was ever composed - an empty routine has nothing to + launch. + + Author: B.G (08/2026) + """ + self._check_mutable() + if not self._order: + raise RoutineBuilderError("freeze: routine has no steps - compose() at least one kernel first") + self._frozen = True + return FrozenRoutine(self._order, self._composed, self._launch) + + +class FrozenRoutine: + """ + The frozen result of a RoutineBuilder's freeze(): an ordered, immutable + {name: FrozenKernel}, plus each step's own launch-kwargs override. See + the module docstring. + + Author: B.G (08/2026) + """ + + def __init__(self, order: list, composed: dict, launch: "dict | None" = None): + self._uid = new_uid() + self._order = tuple(order) + self._composed = dict(composed) + self._launch = dict(launch) if launch else {} + + @property + def uid(self) -> int: + """Process-wide identity assigned at construction. See Parameter.uid (parameter.py).""" + return self._uid + + @property + def order(self) -> tuple: + """Step names in launch order (= composition order).""" + return self._order + + @property + def composed(self) -> dict: + """{step name: FrozenKernel}, read-only copy.""" + return dict(self._composed) + + @property + def launch(self) -> dict: + """{step name: launch-kwargs override dict}, read-only copy. See compose()'s `launch=`.""" + return dict(self._launch) + + def build(self) -> "BoundRoutine": + """ + Return a fresh BoundRoutine with every step's addresses walked and + prefixed by that step's name. See the module docstring's + "Addressing" section. + + A step's own `share()` declarations (builder.py) are honoured the + same way whether the step is built standalone or as part of a + routine. + + Author: B.G (08/2026) + """ + table: dict[Address, Any] = {} + for name in self._order: + step = self._composed[name] + if step.shared: + _walk_group((name,), step, table, {}, frozenset()) + else: + _walk((name,), step, table) + return BoundRoutine(self, table) + + def __repr__(self) -> str: + return f"FrozenRoutine(uid={self._uid}, steps={list(self._order)})" + + +class BoundRoutine(_Bound): + """ + The bound result of build()-ing a FrozenRoutine - bind()/wire()/ + inspect() work exactly as on a BoundKernel (_Bound, bound.py), over the + routine's whole `name.*` address space. See the module docstring's + "Compiling" section for compile(). + + Author: B.G (08/2026) + """ + + def compile(self, backend: str, **kwargs) -> "CompiledRoutine": + """ + Compile every step and return the resulting CompiledRoutine. See + the module docstring's "Compiling" section. + + Parameters + ---------- + backend : str + "taichi", "quadrants" or "cupy". + **kwargs + Routine-level compile()-kwargs, the default every step falls + back to unless overridden by its own `launch=`. + + Author: B.G (08/2026) + """ + check_unmet(self) + frozen: FrozenRoutine = self._frozen + steps: list[tuple[str, Any]] = [] + for name in frozen.order: + step_frozen = frozen.composed[name] + step_bound = step_frozen.build() + for local_addr in step_bound.addresses(): + val = self.value_at((name,) + local_addr) + if val is not None: + step_bound.bind(local_addr, val) + step_kwargs = {**kwargs, **frozen.launch.get(name, {})} + compiled = step_bound.compile(backend, **step_kwargs) + steps.append((name, compiled)) + return CompiledRoutine(steps) + + +class CompiledRoutine: + """ + An immutable, ordered sequence of already-compiled kernels, ready to + launch as one unit. See the module docstring. + + Author: B.G (08/2026) + """ + + def __init__(self, steps: list): + self._steps = list(steps) + self._by_name = dict(steps) + + @property + def step_names(self) -> list: + """Step names in launch order.""" + return [name for name, _ in self._steps] + + def swap(self, addr: "Address | str", buf: Any) -> "CompiledRoutine": + """ + Re-point one step's DATA address at `buf`. + + Parameters + ---------- + addr : Address or str + `name.*`, routed to step `name`'s own CompiledKernel.swap(). + buf : Any + Replacement buffer. + + Author: B.G (08/2026) + """ + a = parse_address(addr) if isinstance(addr, str) else tuple(addr) + if not a: + raise BindError("swap: address must not be empty") + name, local = a[0], a[1:] + if name not in self._by_name: + raise BindError( + f"swap: {format_address(a)!r} - no such routine step {name!r} " + f"(steps: {sorted(self._by_name)})" + ) + self._by_name[name].swap(local, buf) + return self + + def __call__(self) -> None: + """Launch every step in order, each with whatever its own swap() state currently holds.""" + for _, compiled in self._steps: + compiled() + + def __repr__(self) -> str: + return f"CompiledRoutine(steps={[n for n, _ in self._steps]})" diff --git a/pyfastflow/experimental/core/context/sequence.py b/pyfastflow/experimental/core/context/sequence.py new file mode 100644 index 0000000..f91424b --- /dev/null +++ b/pyfastflow/experimental/core/context/sequence.py @@ -0,0 +1,496 @@ +""" +The host-driven layer above Routine (routine.py): an ordered list of blocks +(kernel, whole routine, host block) plus loops whose trip count and break +are evaluated on the host. + +`SequenceBuilder` / `FrozenSequence` / `BoundSequence` / `CompiledSequence` +follow the same build -> freeze -> bind -> compile lifecycle as everything +else in this package. + +What this is for +------------------ +A Routine is device-only and linear - fixed steps, no python between them, +nothing about repeat count decided at run time. That is the wrong shape for +an outer pass whose trip count is not known until the device has been asked - +depression routing reads a pass count back from the device and either goes +round again or stops. A Sequence runs blocks in order, calls host code +between them, and loops with a host-evaluated predicate; `Parameter.read()` +(parameter.py) underpins every such predicate, and it synchronizes - the +layer's whole cost model, paid at block boundaries, never inside a block. + +Composition vs. order +----------------------- +Unlike RoutineBuilder, composing a block here (`compose(name, frozen)`) does +not by itself place it in execution order - a name may be composed once and +then referenced from `step(name)` and/or from inside `loop(...)`'s body more +than once (a callback run once before a loop and again every iteration is +exactly this shape). `step(name)` appends `name` to the top-level order; +`loop(body, max_times, until=None)` appends one loop entry whose `body` is a +sequence of already-composed names, run in order, `max_times` times, stopping +early when `until` returns True. + +`max_times`/`until` are each either a plain value (an int for `max_times`, +`None` for `until`, meaning "run to completion") or the *name* of an +already-composed host block (host_block.py) - a FrozenHostBlock, never a bare +python callable, since a name is the only handle this layer's addressing +scheme has to bind that block's own Parameters through. That host block's +compiled, zero-argument callable is invoked once per check; its return value +is coerced with `int()` for `max_times`, `bool()` for `until`. + +compose() accepts a FrozenKernel, a FrozenRoutine (routine.py) or a +FrozenHostBlock (host_block.py); a FrozenHelper raises, matching +RoutineBuilder.compose() - a device helper has no standalone host-callable +form on its own. + +Addressing +----------- +`build()` walks every composed block's own tree under that block's compose() +name: bound.py's `_walk` directly for a FrozenKernel or FrozenHostBlock (both +real `_Frozen` objects), and, for a FrozenRoutine, one more level of +recursion through its own `.composed` steps - so a routine composed under +`saddlesort` and internally stepping `label`/`sort` reaches +`saddlesort.label.*`/`saddlesort.sort.*`, exactly the address a standalone +Routine's own `build()` would have minted, with the sequence-level compose() +name prefixed on top. + +Compiling +---------- +`BoundSequence.compile(backend, **kwargs)` checks this sequence's own unmet +slots first, then compiles each composed name at most once (cached by name, +since one composed block may be referenced from several places in order): a +fresh bound object from that block's own `.build()`, filled from this +BoundSequence's current values at that block's addresses, then that object's +own `.compile()` - a FrozenHostBlock ignores `backend` (host_block.py), +everything else takes it. The result, CompiledSequence, is an ordered list +of zero-argument callables (blocks already resolved to their own compiled +form) plus loop entries carrying their own body/max_times/until, evaluated +on the host at call time. + +Per-block launch config +-------------------------- +`compose(name, frozen, launch=None)` accepts the same optional launch-kwargs +override `RoutineBuilder.compose()` does (routine.py) - a dict merged over +this sequence's own `compile(backend, **kwargs)` call, `{**kwargs, **launch}`, +for that one composed block only. Composing a whole FrozenRoutine under `name` +with a `launch` override hands that merged dict to the routine's own +`compile()` as *its* default - which the routine's own per-step `launch` +overrides then apply on top of, exactly as they would against any other +default. + +`CompiledSequence.swap(addr, buf)` routes `name.*` to the matching compiled +block's own `.swap()` (CompiledKernel.swap / CompiledRoutine.swap); raises if +that block has nothing to swap (a host block has no DATA of its own - see +host_block.py). + +Author: B.G (08/2026) +""" + +from typing import Any + +from ..pool.base import new_uid +from .bound import Address, BindError, _Bound, _walk, _walk_group, format_address, parse_address +from .compile_shared import CompileError, check_unmet +from .frozen import FrozenBuilderError, FrozenHelper, FrozenKernel +from .host_block import BoundHostBlock, FrozenHostBlock +from .routine import BoundRoutine, FrozenRoutine + + +class SequenceBuilderError(Exception): + """ + Raised by the SequenceBuilder build phase: a name reused or unknown, an + attempt to compose an unsupported frozen type, a malformed loop, or a + mutation after freeze(). + + Author: B.G (08/2026) + """ + + +def _walk_leaf(prefix: Address, frozen: Any, table: dict) -> None: + """ + `_walk` a single frozen object (FrozenKernel or FrozenHostBlock), + honouring its own top-level `.shared` (`_Builder.share()`, builder.py) + exactly as bound.py's own top-level `build()` does for a standalone + object - dispatching to `_walk_group` rather than plain `_walk` when it + has any - so a block composed with its own share() declarations + collapses identically whether built standalone or as part of a + sequence/routine. + + Author: B.G (08/2026) + """ + if frozen.shared: + _walk_group(prefix, frozen, table, {}, frozenset()) + else: + _walk(prefix, frozen, table) + + +def _walk_block(prefix: Address, frozen: Any, table: dict) -> None: + """ + Populate `table` with every PARAM/DATA leaf reachable from `frozen`, at + its full path under `prefix` - dispatching on which of the three + supported block kinds `frozen` is. See the module docstring's + "Addressing" section. + + Author: B.G (08/2026) + """ + if isinstance(frozen, FrozenRoutine): + for name, step_frozen in frozen.composed.items(): + _walk_leaf(prefix + (name,), step_frozen, table) + else: + _walk_leaf(prefix, frozen, table) + + +class SequenceBuilder: + """ + Collects a set of named blocks and an ordered list of steps/loops over + them, and freeze()s them into a FrozenSequence. See the module docstring. + + Author: B.G (08/2026) + """ + + def __init__(self): + self._uid = new_uid() + self._composed: dict[str, Any] = {} + self._launch: dict[str, dict] = {} + self._order: list[tuple] = [] + self._frozen = False + + @property + def uid(self) -> int: + """Process-wide identity assigned at construction. See Parameter.uid (parameter.py).""" + return self._uid + + def _check_mutable(self) -> None: + if self._frozen: + raise FrozenBuilderError( + f"SequenceBuilder(uid={self._uid}) has already been freeze()-ed and is frozen - " + f"build a new SequenceBuilder instead of reusing this one" + ) + + def _require_composed(self, name: str) -> Any: + if name not in self._composed: + raise SequenceBuilderError(f"{name!r} is not composed on this sequence - call compose({name!r}, ...) first") + return self._composed[name] + + def compose(self, name: str, frozen: Any, *, launch: "dict | None" = None) -> "SequenceBuilder": + """ + Register `frozen` under `name`, without placing it in execution + order - see step()/loop() for that, and the module docstring for why + the two are separate calls here (unlike RoutineBuilder.compose()). + + Parameters + ---------- + name : str + frozen : FrozenKernel, FrozenRoutine or FrozenHostBlock + launch : dict, optional + compile()-kwargs overriding this sequence's own compile()-level + default for this block only. Ignored for a FrozenHostBlock + (BoundHostBlock.compile() takes no backend-specific kwargs), + accepted here regardless so a caller need not special-case which + kind of block it is composing. + + Author: B.G (08/2026) + """ + self._check_mutable() + if isinstance(frozen, FrozenHelper): + raise TypeError( + f"compose({name!r}, ...): got a FrozenHelper, not a FrozenKernel/FrozenRoutine/" + f"FrozenHostBlock - a helper has no standalone host-callable form. Compose it " + f"into a KernelBuilder first." + ) + if not isinstance(frozen, (FrozenKernel, FrozenRoutine, FrozenHostBlock)): + raise TypeError( + f"compose({name!r}, ...): expected a FrozenKernel, FrozenRoutine or " + f"FrozenHostBlock, got {type(frozen).__name__}" + ) + if name in self._composed: + raise SequenceBuilderError(f"'{name}' is already composed on this sequence") + self._composed[name] = frozen + self._launch[name] = dict(launch) if launch else {} + return self + + def step(self, name: str) -> "SequenceBuilder": + """ + Append a top-level step launching the block composed under `name`. + + Author: B.G (08/2026) + """ + self._check_mutable() + self._require_composed(name) + self._order.append(("step", name)) + return self + + def loop(self, body, max_times, until: "str | None" = None) -> "SequenceBuilder": + """ + Append a loop running the composed blocks named in `body`, in order, + `max_times` times, stopping early once `until` reports True. + + Parameters + ---------- + body : Sequence[str] + Names of already-composed blocks, run in order each iteration. + max_times : int or str + Trip count, or the name of a composed host block whose + zero-argument callable is invoked once on entry and coerced + with `int()`. + until : str, optional + Name of a composed host block invoked after each iteration and + coerced with `bool()`; `None` runs to completion. + + Author: B.G (08/2026) + """ + self._check_mutable() + body = tuple(body) + if not body: + raise SequenceBuilderError("loop: body is empty") + for name in body: + self._require_composed(name) + if isinstance(max_times, str): + frozen = self._require_composed(max_times) + if not isinstance(frozen, FrozenHostBlock): + raise TypeError(f"loop: max_times={max_times!r} must name a host block, got {type(frozen).__name__}") + elif not isinstance(max_times, int): + raise TypeError("loop: max_times must be an int or the name of a composed host block") + if until is not None: + if not isinstance(until, str): + raise TypeError("loop: until must be None or the name of a composed host block") + frozen = self._require_composed(until) + if not isinstance(frozen, FrozenHostBlock): + raise TypeError(f"loop: until={until!r} must name a host block, got {type(frozen).__name__}") + self._order.append(("loop", body, max_times, until)) + return self + + def freeze(self) -> "FrozenSequence": + """ + Close out the build phase and return the resulting FrozenSequence. + + Raises + ------ + SequenceBuilderError + No step()/loop() was ever recorded. + + Author: B.G (08/2026) + """ + self._check_mutable() + if not self._order: + raise SequenceBuilderError("freeze: sequence has no steps - call step()/loop() at least once") + self._frozen = True + return FrozenSequence(self._composed, self._order, self._launch) + + +class FrozenSequence: + """ + The frozen result of a SequenceBuilder's freeze(): an immutable + {name: block} composition, each block's own launch-kwargs override, plus + the ordered step/loop list. See the module docstring. + + Author: B.G (08/2026) + """ + + def __init__(self, composed: dict, order: list, launch: "dict | None" = None): + self._uid = new_uid() + self._composed = dict(composed) + self._launch = dict(launch) if launch else {} + self._order = list(order) + + @property + def uid(self) -> int: + """Process-wide identity assigned at construction. See Parameter.uid (parameter.py).""" + return self._uid + + @property + def composed(self) -> dict: + """{name: FrozenKernel|FrozenRoutine|FrozenHostBlock}, read-only copy.""" + return dict(self._composed) + + @property + def launch(self) -> dict: + """{name: launch-kwargs override dict}, read-only copy. See compose()'s `launch=`.""" + return dict(self._launch) + + @property + def order(self) -> list: + """The ordered step/loop list, read-only copy.""" + return list(self._order) + + def build(self) -> "BoundSequence": + """ + Walk every composed block's own tree (_walk_block, prefixed with its + compose() name) and return a fresh BoundSequence. See the module + docstring's "Addressing" section. + + Author: B.G (08/2026) + """ + table: dict[Address, Any] = {} + for name, frozen in self._composed.items(): + _walk_block((name,), frozen, table) + return BoundSequence(self, table) + + def __repr__(self) -> str: + return f"FrozenSequence(uid={self._uid}, blocks={sorted(self._composed)})" + + +class BoundSequence(_Bound): + """ + The bound result of build()-ing a FrozenSequence - bind()/wire()/ + inspect() work exactly as on a BoundKernel (_Bound, bound.py), over the + sequence's whole `name.*` address space. See the module docstring's + "Compiling" section for compile(). + + Author: B.G (08/2026) + """ + + def compile(self, backend: str, **kwargs) -> "CompiledSequence": + """ + Compile every composed block and return the resulting + CompiledSequence. See the module docstring's "Compiling" section. + + Parameters + ---------- + backend : str + "taichi", "quadrants" or "cupy". + **kwargs + Sequence-level compile()-kwargs, the default every block falls + back to unless overridden by its own `launch=`. + + Author: B.G (08/2026) + """ + check_unmet(self) + frozen: FrozenSequence = self._frozen + compiled_blocks: dict[str, Any] = {} + + def _compile_name(name: str) -> Any: + if name in compiled_blocks: + return compiled_blocks[name] + child = frozen.composed[name] + child_bound = child.build() + for local_addr in child_bound.addresses(): + val = self.value_at((name,) + local_addr) + if val is not None: + child_bound.bind(local_addr, val) + block_kwargs = {**kwargs, **frozen.launch.get(name, {})} + compiled = child_bound.compile() if isinstance(child, FrozenHostBlock) else child_bound.compile(backend, **block_kwargs) + compiled_blocks[name] = compiled + return compiled + + entries: list[_SeqEntry] = [] + for item in frozen.order: + if item[0] == "step": + _, name = item + entries.append(_SeqEntry("run", run=_compile_name(name))) + else: + _, body, max_times, until = item + body_compiled = tuple(_compile_name(n) for n in body) + mt = max_times if isinstance(max_times, int) else _compile_name(max_times) + un = None if until is None else _compile_name(until) + entries.append(_SeqEntry("loop", body=body_compiled, max_times=mt, until=un)) + + return CompiledSequence(entries, compiled_blocks) + + +class _SeqEntry: + """ + One entry of a CompiledSequence: a "run" entry wraps a single already- + resolved zero-argument callable; a "loop" entry carries its own compiled + body, max_times and until (each itself a plain value or a zero-argument + callable). Not constructed directly outside BoundSequence.compile(). + + Author: B.G (08/2026) + """ + + __slots__ = ("kind", "run", "body", "max_times", "until") + + def __init__(self, kind: str, run: Any = None, body: tuple = (), max_times: Any = None, until: Any = None): + self.kind = kind + self.run = run + self.body = body + self.max_times = max_times + self.until = until + + +class CompiledSequence: + """ + An immutable, ordered list of resolved blocks and host-evaluated loops, + ready to run. See the module docstring. + + `last_trip_counts` reports how many body iterations each loop entry took + on the most recent call, in the order the loop entries appear. + + Author: B.G (08/2026) + """ + + def __init__(self, entries: list, compiled_blocks: dict): + self._entries = entries + self._compiled_blocks = compiled_blocks + self._last_trip_counts: tuple = () + + @property + def last_trip_counts(self) -> tuple: + """Body iterations taken by each loop entry on the most recent call, in entry order.""" + return self._last_trip_counts + + def swap(self, addr: "Address | str", buf: Any) -> "CompiledSequence": + """ + Re-point one composed block's DATA address at `buf`. + + Parameters + ---------- + addr : Address or str + `name.*`, routed to that block's own compiled `.swap()`. + buf : Any + Replacement buffer. + + Raises + ------ + BindError + `name` is not composed, or the named block has nothing to swap + (a compiled host block is a plain callable with no data + addresses of its own). + + Author: B.G (08/2026) + """ + a = parse_address(addr) if isinstance(addr, str) else tuple(addr) + if not a: + raise BindError("swap: address must not be empty") + name, local = a[0], a[1:] + if name not in self._compiled_blocks: + raise BindError( + f"swap: {format_address(a)!r} - no such composed block {name!r} " + f"(blocks: {sorted(self._compiled_blocks)})" + ) + target = self._compiled_blocks[name] + if not hasattr(target, "swap"): + raise BindError(f"swap: {format_address(a)!r} - block {name!r} has no data to swap (it is a host block)") + target.swap(local, buf) + return self + + def __call__(self) -> None: + """Run every entry in order - see the module docstring's cost-model paragraph.""" + trips: list[int] = [] + for entry in self._entries: + if entry.kind == "loop": + trips.append(self._run_loop(entry)) + else: + entry.run() + self._last_trip_counts = tuple(trips) + + def _run_loop(self, entry: _SeqEntry) -> int: + """ + Evaluate `max_times` once on entry, run the body that many times, + evaluating `until` after each iteration and stopping when it returns + True. Returns the number of iterations actually run. + + Author: B.G (08/2026) + """ + max_times = entry.max_times + times = int(max_times()) if callable(max_times) else int(max_times) + taken = 0 + for _ in range(max(0, times)): + for inner in entry.body: + inner() + taken += 1 + if entry.until is not None and bool(entry.until()): + break + return taken + + def __repr__(self) -> str: + return f"CompiledSequence(entries={len(self._entries)})" diff --git a/pyfastflow/experimental/core/context/slot.py b/pyfastflow/experimental/core/context/slot.py new file mode 100644 index 0000000..eb9026f --- /dev/null +++ b/pyfastflow/experimental/core/context/slot.py @@ -0,0 +1,214 @@ +""" +Slot: one named place declared on a builder during the build phase, plus +SlotGroup, the flat container that holds a builder's own set of them. + +This is the vocabulary wire_param()/wire_helper()/wire_data() (builder.py) +speak in. A Slot carries a name and a kind - PARAM, HELPER or DATA - plus, +for a DATA slot only, an optional dtype (see DataSlot). Nothing else: no +mode, no backend concern. The kind says how the slot is reached once the +template is written: + + PARAM read through ctx in device code (ctx.z.get(i)) - a Parameter is + bound to this slot later, in the bind phase, not here. Access + is uniform across a Parameter's modes: ctx.z.get(i) reads z + whether it is const, scalar or field, with `i` simply ignored for + const/scalar - the same "read all three modes identically" design + Parameter itself states (parameter.py's module docstring). This is + a stated convention, not something enforced here or in the bind + phase: which accessor spellings are legal on device code is a + backend-emission question, settled where the backend actually + knows what it can generate. + HELPER called through ctx in device code (ctx.grad(...)) - a HelperBuilder + is composed or bound to this slot, likewise later. + DATA a trusted device call argument of the compiled kernel/helper's own + signature, never reached through ctx at all - see the module + docstring of builder.py for why wire_data raises on a HelperBuilder. + An optional dtype declared at wire_data(name, dtype=...) time is + the one thing checked early: a wrong-dtype buffer is what corrupts + a launch or segfaults silently. Where that check actually runs + against a call-time value is bind/compile work; the dtype is only + declared and carried here. + +A Slot is local to the one builder that declared it via wire_*(); nothing +here is process-wide or shared the way a Parameter's own uid is +(parameter.py). Two builders wanting "the same" slot get there through +compose() (an already-frozen sub-structure, shared by identity - see +frozen.py) or, through bind()'s addressing, never by two Slots comparing +equal. + +SlotGroup is one level of depth on purpose - unlike Bag (bag.py), which +nests and is reached in-kernel by dotted path, a SlotGroup only ever +enumerates one builder's own local names. Depth comes from compose(), which +attaches an entire other frozen builder's own address tree under one slot +name; SlotGroup itself never nests. + +Naming note: "Handle" is reserved exclusively for pool.base.DataHandle +(core/pool/) - the device buffer handle behind a Parameter's scalar/field +storage. This module never uses that word for anything of its own; every +declared place at this layer is a Slot. + +Author: B.G (08/2026) +""" + +from enum import Enum +from typing import Iterator + + +class SlotKind(Enum): + """ + What a Slot's place will eventually hold. See the module docstring for + how each kind is reached from device code. + + Author: B.G (08/2026) + """ + + PARAM = "param" + HELPER = "helper" + DATA = "data" + + +class SlotGroupError(Exception): + """ + Raised when a builder's local slot namespace is misused - wiring a name + twice (as a slot or as a compose() root), or looking up a name that was + never wired. + + Author: B.G (08/2026) + """ + + +class Slot: + """ + One named place of a given SlotKind, local to the builder that declared + it. See the module docstring. + + Author: B.G (08/2026) + """ + + __slots__ = ("name", "kind") + + def __init__(self, name: str, kind: SlotKind): + self.name = name + self.kind = kind + + def __repr__(self) -> str: + return f"Slot({self.name!r}, kind={self.kind.value})" + + def __eq__(self, other) -> bool: + return isinstance(other, Slot) and self.name == other.name and self.kind is other.kind + + def __hash__(self) -> int: + return hash((self.name, self.kind)) + + +class ParamSlot(Slot): + """A PARAM slot - wire_param()'s own slot type. See the module docstring.""" + + __slots__ = () + + def __init__(self, name: str): + super().__init__(name, SlotKind.PARAM) + + +class HelperSlot(Slot): + """A HELPER slot - wire_helper()'s own slot type. See the module docstring.""" + + __slots__ = () + + def __init__(self, name: str): + super().__init__(name, SlotKind.HELPER) + + +class DataSlot(Slot): + """ + A DATA slot - wire_data()'s own slot type. See the module docstring. + + `dtype` is optional: None means the slot stays open (any dtype accepted + whenever it is eventually checked against a call-time value), anything + else declares a contract wire_data(name, dtype=...) callers can rely on + being validated downstream, at bind/compile time. + + Author: B.G (08/2026) + """ + + __slots__ = ("dtype",) + + def __init__(self, name: str, dtype=None): + super().__init__(name, SlotKind.DATA) + self.dtype = dtype + + def __repr__(self) -> str: + if self.dtype is None: + return f"Slot({self.name!r}, kind=data)" + return f"Slot({self.name!r}, kind=data, dtype={self.dtype})" + + +class SlotGroup: + """ + The flat {name: Slot} namespace behind one builder's wire_*() calls. + + Add-only during the build phase: a name may be wired once, checked here + rather than left to collide silently later. copy() gives the snapshot a + frozen builder (frozen.py) keeps for itself once ingest() has run, so a + later mutation of the *builder's* group (which cannot happen anyway, + since ingest() freezes the builder - see builder.py) could never reach + back into an already-frozen result even if that guard were ever loosened. + + Author: B.G (08/2026) + """ + + def __init__(self): + self._slots: dict[str, Slot] = {} + + def add(self, slot: Slot) -> None: + """ + Register `slot` under its own name. Raises if that name is already + wired on this group, as a Slot or otherwise. + + Author: B.G (08/2026) + """ + if slot.name in self._slots: + raise SlotGroupError( + f"'{slot.name}' is already wired on this builder " + f"(as {self._slots[slot.name]!r})" + ) + self._slots[slot.name] = slot + + def __contains__(self, name: str) -> bool: + return name in self._slots + + def __getitem__(self, name: str) -> Slot: + return self._slots[name] + + def __iter__(self) -> Iterator[Slot]: + return iter(self._slots.values()) + + def __len__(self) -> int: + return len(self._slots) + + def names(self, kind: SlotKind | None = None) -> set[str]: + """ + Every wired name, or just those of one `kind` if given. + + Author: B.G (08/2026) + """ + if kind is None: + return set(self._slots) + return {name for name, slot in self._slots.items() if slot.kind is kind} + + def copy(self) -> "SlotGroup": + """ + A fresh SlotGroup holding the same Slot objects (Slot is itself + immutable data, so nothing needs a deeper copy). + + Author: B.G (08/2026) + """ + new = SlotGroup() + new._slots = dict(self._slots) + return new + + def __repr__(self) -> str: + if not self._slots: + return "SlotGroup()" + body = ", ".join(repr(s) for s in self._slots.values()) + return f"SlotGroup({body})" diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py new file mode 100644 index 0000000..e7ffd81 --- /dev/null +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -0,0 +1,28 @@ +""" +Taichi implementation of Parameter. + +A Parameter's device view is a python def specialized by rebuilding it with +the bound objects spliced into its globals, then handed to ti.func; +_closure_backend.py holds that machinery, shared with Quadrants. Everything +below just names Taichi as the backend to use. + +The kernel/helper compile path (KernelBuilder/HelperBuilder -> FrozenKernel -> +BoundKernel.compile("taichi")) is compile_closure.py, which reaches this +module only for `ti` itself, imported directly there. + +Author: B.G (07/2026) +""" + +import taichi as ti + +from ._closure_backend import ClosureBackendParameter + + +class TaichiParameter(ClosureBackendParameter): + """ + Parameter backed by a Taichi const value or a pooled TaichiDataHandle. + + Author: B.G (07/2026) + """ + + _backend = ti diff --git a/pyfastflow/experimental/core/pool/__init__.py b/pyfastflow/experimental/core/pool/__init__.py new file mode 100644 index 0000000..f376c18 --- /dev/null +++ b/pyfastflow/experimental/core/pool/__init__.py @@ -0,0 +1,5 @@ +""" +New backend-agnostic pool architecture (DataHandle/Pool ABCs + backends). + +Author: B.G (07/2026) +""" diff --git a/pyfastflow/experimental/core/pool/_bucketed_pool.py b/pyfastflow/experimental/core/pool/_bucketed_pool.py new file mode 100644 index 0000000..e544d54 --- /dev/null +++ b/pyfastflow/experimental/core/pool/_bucketed_pool.py @@ -0,0 +1,56 @@ +""" +Shared bucketed Pool implementation, parameterized by a DataHandle subclass. + +Author: B.G (07/2026) +""" + +from typing import Any, ClassVar + +from .base import DataHandle, Pool + + +class BucketedPool(Pool): + """ + Pool manager bucketed by (dtype, shape). Subclasses only pin + `_handle_cls` to the DataHandle implementation they allocate. + + Author: B.G (07/2026) + """ + + _handle_cls: ClassVar[type] + + def __init__(self): + self._buckets: dict[tuple[Any, tuple[int, ...]], list[DataHandle]] = {} + + def get_data(self, dtype, shape) -> DataHandle: + key = (dtype, tuple(shape)) + bucket = self._buckets.setdefault(key, []) + for handle in bucket: + if not handle.in_use: + handle.acquire() + return handle + handle = self._handle_cls(dtype, key[1]) + handle.acquire() + bucket.append(handle) + return handle + + def release_data(self, handle: DataHandle) -> None: + handle.release() + + def clear_unused(self) -> None: + for bucket in self._buckets.values(): + for handle in bucket[:]: + if not handle.in_use: + handle.destroy() + bucket.remove(handle) + + def clear_all(self) -> None: + for bucket in self._buckets.values(): + for handle in bucket[:]: + handle.destroy() + bucket.remove(handle) + + def stats(self) -> dict: + total = sum(len(bucket) for bucket in self._buckets.values()) + in_use = sum(1 for bucket in self._buckets.values() for h in bucket if h.in_use) + return {"total": total, "in_use": in_use, "available": total - in_use} diff --git a/pyfastflow/experimental/core/pool/_fields_handle.py b/pyfastflow/experimental/core/pool/_fields_handle.py new file mode 100644 index 0000000..60384f6 --- /dev/null +++ b/pyfastflow/experimental/core/pool/_fields_handle.py @@ -0,0 +1,90 @@ +""" +Shared DataHandle implementation for FieldsBuilder-based backends. + +Taichi and Quadrants expose an identical field/FieldsBuilder API; subclasses +only pin `_backend` to their module (ti or qd). + +Author: B.G (07/2026) +""" + +from typing import Any, ClassVar + +from .base import DataHandle, new_uid + + +class FieldsBuilderDataHandle(DataHandle): + """ + DataHandle backed by one field allocated via FieldsBuilder. + + Composition, not inheritance: kernels take the raw field via `.data`, + not the handle itself - see pool/base.py design notes on why + subclassing a field type was rejected. + + Author: B.G (07/2026) + """ + + _backend: ClassVar[Any] + _next_id = 0 + + def __init__(self, dtype: Any, shape: tuple[int, ...]): + """ + Allocate a field of the given dtype/shape via FieldsBuilder. + + shape=() allocates a 0D scalar field, indexed as field[None]. + + Author: B.G (07/2026) + """ + cls = type(self) + cls._next_id += 1 + self.alloc_id = cls._next_id + self._uid = new_uid() + self.dtype = dtype + self.shape = tuple(shape) + self.in_use = False + + backend = self._backend + self._builder = backend.FieldsBuilder() + self._field = backend.field(dtype) + + if len(self.shape) == 0: + self._builder.place(self._field) + elif len(self.shape) == 1: + self._builder.dense(backend.i, self.shape).place(self._field) + elif len(self.shape) == 2: + self._builder.dense(backend.ij, self.shape).place(self._field) + else: + raise ValueError(f"Unsupported field dimensionality: {len(self.shape)}D. Only 0D, 1D, 2D supported.") + + self._snodetree = self._builder.finalize() + + @property + def data(self): + """ + Return the underlying field, for passing straight into kernels or + binding as a global. + + Author: B.G (07/2026) + """ + return self._field + + def acquire(self) -> None: + self.in_use = True + + def release(self) -> None: + self.in_use = False + + def destroy(self) -> None: + """ + Free the field's GPU memory. Unusable afterwards. + + Author: B.G (07/2026) + """ + if self._snodetree is not None: + self._snodetree.destroy() + self._snodetree = None + + def to_numpy(self): + return self._field.to_numpy() + + def from_numpy(self, arr) -> None: + self._field.from_numpy(arr) diff --git a/pyfastflow/experimental/core/pool/base.py b/pyfastflow/experimental/core/pool/base.py new file mode 100644 index 0000000..366760a --- /dev/null +++ b/pyfastflow/experimental/core/pool/base.py @@ -0,0 +1,180 @@ +""" +Backend-agnostic pool contracts. + +Defines the blueprint that every pool backend (Taichi fields, ndarrays, +quadrants, cupy, ...) must implement. No allocation logic here +- this is the interface only. + +Author: B.G (07/2026) +""" + +import itertools +from abc import ABC, abstractmethod +from typing import Any + +_uid_counter = itertools.count() + + +def new_uid() -> int: + """ + Return the next value from the process-wide identity counter. + + Every Parameter, Bag, Helper (device-function builder and its compiled + artifact) and pool DataHandle is assigned one of these at construction, + exposed as a read-only `uid` property. uids are plain integers drawn from + this single shared counter - not stable across processes, and + deliberately so: they identify an object within one running process and + must never appear in generated code or a cache key. + + Author: B.G (07/2026) + """ + return next(_uid_counter) + + +class DataHandle(ABC): + """ + Opaque handle to one pooled backend resource (a Taichi field, ndarray, ...). + + Owns the acquire/release lifecycle: `release()` returns the handle to its + pool for reuse without freeing memory; `destroy()` actually frees it. + + Attributes: + alloc_id: Per-backend allocation counter, assigned by the backend - used + for pool bookkeeping and not unique across backends. + uid: Process-wide identity from the shared counter (new_uid()) - unique + across every Parameter, Bag, Helper and DataHandle regardless of + backend. Concrete handles set self._uid in their own __init__. + dtype: Backend-native or common dtype tag for this resource. + shape: Resource dimensions. () for a scalar. + in_use: True between acquire() and release(). + + Two handles from different pools can share an alloc_id; only uid identifies + a handle on its own. + + Author: B.G (07/2026) + """ + + alloc_id: int + dtype: Any + shape: tuple[int, ...] + in_use: bool + + @property + def uid(self) -> int: + """ + Process-wide identity assigned at construction. See new_uid(). + + Author: B.G (07/2026) + """ + return self._uid + + @property + @abstractmethod + def data(self): + """ + Return the raw backend object (ti.field, np.ndarray, ...). + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def acquire(self) -> None: + """ + Mark this handle in_use. Called by the owning pool on checkout. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def release(self) -> None: + """ + Mark this handle available for reuse. Backend memory is kept. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def destroy(self) -> None: + """ + Free the underlying backend memory. Handle is unusable afterwards. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def to_numpy(self): + """ + Copy the resource out to a numpy array. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def from_numpy(self, arr) -> None: + """ + Copy a numpy array into the resource in place. + + Author: B.G (07/2026) + """ + ... + + +class Pool(ABC): + """ + Blueprint for a backend-specific pool manager. + + Implementations keep handles bucketed by (dtype, shape) and reuse + released handles before allocating new ones. + + Author: B.G (07/2026) + """ + + @abstractmethod + def get_data(self, dtype, shape) -> DataHandle: + """ + Return an available handle matching (dtype, shape), allocating one if needed. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def release_data(self, handle: DataHandle) -> None: + """ + Return a handle to the pool for reuse. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def clear_unused(self) -> None: + """ + Destroy and drop all handles currently not in_use. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def clear_all(self) -> None: + """ + Destroy and drop every handle, regardless of in_use state. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def stats(self) -> dict: + """ + Return {"total", "in_use", "available"} handle counts. + + Author: B.G (07/2026) + """ + ... diff --git a/pyfastflow/experimental/core/pool/cupy_handle.py b/pyfastflow/experimental/core/pool/cupy_handle.py new file mode 100644 index 0000000..e85ca9d --- /dev/null +++ b/pyfastflow/experimental/core/pool/cupy_handle.py @@ -0,0 +1,66 @@ +""" +Cupy backend implementation of DataHandle. + +Author: B.G (07/2026) +""" + +from typing import Any + +import cupy as cp + +from .base import DataHandle, new_uid + + +class CupyDataHandle(DataHandle): + """ + DataHandle backed by one cupy ndarray. + + Author: B.G (07/2026) + """ + + _next_id = 0 + + def __init__(self, dtype: Any, shape: tuple[int, ...]): + """ + Allocate a cupy ndarray of the given dtype/shape. + + Author: B.G (07/2026) + """ + CupyDataHandle._next_id += 1 + self.alloc_id = CupyDataHandle._next_id + self._uid = new_uid() + self.dtype = dtype + self.shape = tuple(shape) + self.in_use = False + self._array = cp.empty(self.shape, dtype=dtype) + + @property + def data(self): + """ + Return the underlying cupy ndarray, for passing straight into a + RawKernel launch. + + Author: B.G (07/2026) + """ + return self._array + + def acquire(self) -> None: + self.in_use = True + + def release(self) -> None: + self.in_use = False + + def destroy(self) -> None: + """ + Drop the reference; cupy's own memory pool reclaims the block for + reuse. Unusable afterwards. + + Author: B.G (07/2026) + """ + self._array = None + + def to_numpy(self): + return cp.asnumpy(self._array) + + def from_numpy(self, arr) -> None: + self._array[...] = cp.asarray(arr) diff --git a/pyfastflow/experimental/core/pool/cupy_pool.py b/pyfastflow/experimental/core/pool/cupy_pool.py new file mode 100644 index 0000000..2ac609f --- /dev/null +++ b/pyfastflow/experimental/core/pool/cupy_pool.py @@ -0,0 +1,18 @@ +""" +Cupy backend implementation of Pool. + +Author: B.G (07/2026) +""" + +from ._bucketed_pool import BucketedPool +from .cupy_handle import CupyDataHandle + + +class CupyPool(BucketedPool): + """ + Pool manager for CupyDataHandle, bucketed by (dtype, shape). + + Author: B.G (07/2026) + """ + + _handle_cls = CupyDataHandle diff --git a/pyfastflow/experimental/core/pool/quadrants_handle.py b/pyfastflow/experimental/core/pool/quadrants_handle.py new file mode 100644 index 0000000..e304f67 --- /dev/null +++ b/pyfastflow/experimental/core/pool/quadrants_handle.py @@ -0,0 +1,19 @@ +""" +Quadrants backend implementation of DataHandle. + +Author: B.G (07/2026) +""" + +import quadrants as qd + +from ._fields_handle import FieldsBuilderDataHandle + + +class QuadrantsDataHandle(FieldsBuilderDataHandle): + """ + DataHandle backed by one Quadrants field. + + Author: B.G (07/2026) + """ + + _backend = qd diff --git a/pyfastflow/experimental/core/pool/quadrants_pool.py b/pyfastflow/experimental/core/pool/quadrants_pool.py new file mode 100644 index 0000000..babae74 --- /dev/null +++ b/pyfastflow/experimental/core/pool/quadrants_pool.py @@ -0,0 +1,18 @@ +""" +Quadrants backend implementation of Pool. + +Author: B.G (07/2026) +""" + +from ._bucketed_pool import BucketedPool +from .quadrants_handle import QuadrantsDataHandle + + +class QuadrantsPool(BucketedPool): + """ + Pool manager for QuadrantsDataHandle, bucketed by (dtype, shape). + + Author: B.G (07/2026) + """ + + _handle_cls = QuadrantsDataHandle diff --git a/pyfastflow/experimental/core/pool/taichi_handle.py b/pyfastflow/experimental/core/pool/taichi_handle.py new file mode 100644 index 0000000..c4f73a4 --- /dev/null +++ b/pyfastflow/experimental/core/pool/taichi_handle.py @@ -0,0 +1,19 @@ +""" +Taichi backend implementation of DataHandle. + +Author: B.G (07/2026) +""" + +import taichi as ti + +from ._fields_handle import FieldsBuilderDataHandle + + +class TaichiDataHandle(FieldsBuilderDataHandle): + """ + DataHandle backed by one Taichi field. + + Author: B.G (07/2026) + """ + + _backend = ti diff --git a/pyfastflow/experimental/core/pool/taichi_pool.py b/pyfastflow/experimental/core/pool/taichi_pool.py new file mode 100644 index 0000000..442571f --- /dev/null +++ b/pyfastflow/experimental/core/pool/taichi_pool.py @@ -0,0 +1,18 @@ +""" +Taichi backend implementation of Pool. + +Author: B.G (07/2026) +""" + +from ._bucketed_pool import BucketedPool +from .taichi_handle import TaichiDataHandle + + +class TaichiPool(BucketedPool): + """ + Pool manager for TaichiDataHandle, bucketed by (dtype, shape). + + Author: B.G (07/2026) + """ + + _handle_cls = TaichiDataHandle diff --git a/pyfastflow/experimental/flow/__init__.py b/pyfastflow/experimental/flow/__init__.py new file mode 100644 index 0000000..88bd867 --- /dev/null +++ b/pyfastflow/experimental/flow/__init__.py @@ -0,0 +1,1256 @@ +""" +make_receivers: the SFD (single-flow-direction) receiver factory, built on +the new builder/frozen/bound stack (../core/context/builder.py, frozen.py, +bound.py) and on a grid FrozenGroup from ../grid's make_grid_group. + +Like grid/noise/ops there is no stateful context class - make_receivers +returns a dict of unbuilt structures (FrozenKernel/FrozenHelper): a +`receivers` FrozenKernel plus the distance/slope helpers it is made of, so a +caller can recombine them into its own kernel or routine rather than being +stuck with only the compiled receivers kernel. A caller `.build()`s the +member it wants, binds its PARAM/DATA addresses, `.compile()`s: + + grid = make_grid_group("taichi", topology="D8") + recv = make_receivers("taichi", grid, topology="D8", mode="steepest") + bound = recv["receivers"].build() + bound.bind_leaf(grid_params) # NX/NY/DX/N_NEIGHBOURS - see below + bound.bind("z", z_field) + bound.bind("rec", rec_field) + receivers_kernel = bound.compile("taichi") + receivers_kernel() + +`mode` ("steepest"|"stochastic") and `h_aware` (False: kernel takes (z, rec) +and slopes read h as 0; True: kernel takes (z, h, rec) and slopes use +(zi-zj)+(hi-hj)) each pick one of four kernel body variants at build time - +see _closure_receivers.py/_cupy_receivers.py's build_receivers. `topology` +("D4"|"D8") must match whatever `grid` was itself built with (see +../grid/__init__.py's own module docstring on the two-call structure/data +split) - it is not readable off `grid` itself, a bare FrozenGroup carrying no +Parameter values yet, and only matters here for +`diagonal_partition_correction` (below); the neighbour loop itself is +already parametrised by `grid.N_NEIGHBOURS`, read as ordinary device data. + +mode="stochastic" additionally needs a Parameter bound to the built +receivers kernel's `rand_unit.SEED` address (`rand_unit`'s own wired PARAM +slot, any mode - see rand_unit in the block modules; not build-phase-shared +with anything, so it stays at that nested address rather than being +promoted to the kernel's own top level the way `grid`'s own PARAM names +are) after `.build()`, exactly like any other PARAM slot - there is no Need +indirection in this stack - the host bumps the underlying Parameter between calls for a fresh draw. + +`diagonal_partition_correction` only changes anything when `topology == +"D8"`: it adds the sqrt(2) correction inside dist_from_k_corrected/ +dist_between_nodes_corrected (see _closure_receivers.py's +build_distance_slope_helpers for exactly which k values count as diagonal and +why). Off, or on a D4 grid, dist_from_k_corrected/dist_between_nodes_corrected +call straight through to `grid`'s own dist_from_k/dist_between_nodes, no +correction applied - either way these two helpers always independently +compose their own occurrence of `grid` (see _closure_receivers.py's module +docstring for why: a uniform, always-two-occurrences shape lets +`build_receivers` collapse them with `share()` unconditionally, never a +variable number of occurrences depending on the flag). + +Returned dict: `receivers` (FrozenKernel, data args (z, rec) or (z, h, rec)), +`dist_from_k_corrected`, `dist_between_nodes_corrected`, `slope_from_values_k`, +`slope_between_nodes` (FrozenHelper), plus `rand_unit` (FrozenHelper) only +when mode="stochastic". `receivers`'s own top-level PARAM slots are every +name `grid` itself wires (NX/NY/DX/N_NEIGHBOURS, plus NODATA_MASK/ +OUTLET_MASK if `grid` has them) - bind those bare names once on the built +receivers kernel, not once per occurrence (`grid`'s FrozenGroup is composed +twice inside `receivers`'s own tree - once directly, once nested under +`slope.dist_from_k_corrected` - and build-phase-shared, `_share_leaf`, into +one address each). + +A node with no downslope neighbour keeps `rec[i] = i` - a self-receiver, the +same convention a can_out (base level) node uses. This is the pit convention +depression handling later depends on. + +rand_unit's hash is keyed on (node, k, seed) rather than (node, seed): legacy +draws once per (node, k) candidate inside the neighbour loop +(`ti.random()`), so a node-keyed hash would scale every candidate at a node +by the same factor and weaken the randomisation rather than reproduce it. +Reproducibility itself diverges from legacy (hash-based vs ti.random()'s +counter-based PRNG) - only the selection distribution's shape is preserved. + +make_accumulation: the SFD downstream-accumulation factory, on the new +builder/frozen/bound stack (see its own docstring below for the per-method +details). `source` is bound directly, post-`.build()`, at the returned +kernel's own `SOURCE` PARAM slot - any mode (const, scalar or field all work +with no variant code, since every template reads `source.get(i)`) - there is +no Need indirection anywhere in this stack: + + source_p = TaichiParameter("SOURCE", dtype=ti.f32, mode="const", value=1.0, pool=pool) + accum = make_accumulation("taichi", grid, method="atomic", n_flat=n_flat) + bound = accum["accum"].build() + bound.bind("SOURCE", source_p) + bound.bind("rec", rec.data) + bound.bind("q", q.data) + accum_kernel = bound.compile("taichi") + accum_kernel() + +`method`: + - "atomic": on taichi/quadrants, one KernelBuilder ("accum", data args + (rec, q)) - two top-level for-loops (q[i] = source.get(i), then the + descent), which the closure backends already launch as two barrier- + separated GPU dispatches. On cupy, two KernelBuilders ("q_init", data + arg (q); "accum", data args (rec, q)) - a single CUDA __global__ has no + portable grid-wide barrier the way two consecutive Taichi/Quadrants + for-loops do, so "q_init" must be launched, and finish, before "accum". + Either way: every node walks its receiver chain to the root, atomic- + adding its own weight into each downstream node. Requires an acyclic + receiver graph (run after depression handling) - a cycle degrades the + result (the walk gives up once its guard counter reaches n_flat) rather + than hanging. + - "rake_compress" (ported to the new builder/frozen/bound/sequence stack, + ../core/context/builder.py/frozen.py/bound.py/sequence.py): a + SequenceBuilder (see _closure_accum.py's/_cupy_accum.py's + build_rake_compress) plus its constituent KernelBuilders, keyed + "zero_init", "reset_iteration", "decrement_iteration", "q_init", + "receivers_to_donors", "rake_compress_accum", "fuse_accum_buffers" (plus + "bump_iteration" on cupy only - see below). Composed sequence steps: + "zero_init" -> "reset_iteration" -> "q_init" -> "receivers_to_donors" -> + a loop over "rake_step" (the same rake_compress_accum kernel, `max_times + = ceil(log2(n_flat)) + 2`) -> "decrement_iteration" (undoes the loop's + last bump) -> "fuse_accum_buffers". On closure backends, the iteration + bump is rake_compress_accum's own second top-level `for` loop, folded in + rather than a separate single-thread kernel (two consecutive top-level + `for` loops inside one compiled Taichi/Quadrants kernel are already + separate offloaded tasks launched in order); on cupy, a single CUDA + `__global__` gives no such guarantee, so the loop body is + `["rake_step", "bump_iteration"]`, two real launches per round - see + _cupy_accum.py's build_rake_compress. No `source`/`iteration_p` argument + to this factory at all - `SOURCE`/`ITER` are bare wired PARAM slots (any + mode), bound by the caller on the built sequence after `.build()`, + exactly like make_receivers' `rand_unit.SEED`; see build_rake_compress's + own docstring (per backend) for the exact addresses (four independent + `ITER` addresses after rake_compress_accum's own `share()` collapses its + two composed ping-pong helpers' ITER occurrences into its own, one + `SOURCE` address). + - "pointer_jump_push" (ported the same way): a SequenceBuilder (see + build_pointer_jump_push) plus its constituent KernelBuilders, keyed + "q_init", "copy_rec_to_work", and (closure backends) + "accum_pointer_jump_push_step" or (cupy, split into two launches for a + real barrier between the copy and the push - see _cupy_accum.py) + "accum_pointer_jump_push_step_copy"/"accum_pointer_jump_push_step_core". + The ping-pong between rounds is two independently-bound occurrences of + the same step kernel ("step_a"/"step_b" - closure; "step_a_copy"/ + "step_a_core"/"step_b_copy"/"step_b_core" - cupy), alternated by a + sequence loop of `rounds // 2` iterations (`rounds`, computed here, + already rounded up to even) - no runtime swap() needed, unlike + routine.py's old add_swap. Same no-`source`-argument contract as + rake_compress; see build_pointer_jump_push's own docstring for its + addresses. + +Both factories return {"sequence": SequenceBuilder, **kernel_builders} - a +Bag, not a compiled object (these factories export builders, not compiled +kernels - see CLAUDE.md). The caller `.freeze()`s (or lets `.compile()` +freeze implicitly - SequenceBuilder has no separate freeze() call exposed +here beyond what `.build()`/`.compile()` already do internally), `.build()`s, +binds every PARAM/DATA address named above and in each build_* docstring, +then `.compile(backend, grid=..., block=...)` on cupy (grid/block size the +sequence's own default launch dims; single-thread steps override their own +via `launch=` at compose() time, already baked in by the factory) or +`.compile(backend)` on closure backends (no launch dims needed). The +compiled CompiledSequence takes no arguments; call it, then read +`last_trip_counts` if wanted (always exactly one loop entry per compiled +sequence here, so `last_trip_counts[0]` is the only entry, and is always the +full requested count - unlike depression routing's own use of the same loop +machinery, nothing here ever breaks out early via `until`). + +Why a SequenceBuilder loop rather than routine.py's unroll-N-times idiom +(../ops/_closure_blocks.py's build_scan_routine, for its own log-depth +passes): a scan pass's kernel body differs every round (`stride` baked in as +a build-time constant), so unrolling costs nothing beyond the kernels +themselves; here the SAME kernel body runs unchanged every round, so +unrolling would only multiply the number of addresses a caller has to bind +(once per round) for no benefit - a loop keeps that count fixed regardless of +round count (`ceil(log2(n_flat))+2` rounds at 1024x1024 is ~22). + +`n_flat` is REQUIRED for both, like `method="atomic"`'s own required `n_flat` +- `grid` is a bare `make_grid_group` FrozenGroup with no bound Parameter +values to read it off at build time (see ../grid/__init__.py's own module +docstring), so there is no way to read it from `grid` itself. `rake_compress` +also needs `n_neighbours` explicitly, for the same reason (sizing the fixed +per-node donor arrays) - `pointer_jump_push` needs neither `grid` nor +`n_neighbours` at all. + +make_depressions: the depression-handling factory, on the builder/frozen/ +bound stack. Two orthogonal build flags: + + ndep_p = TaichiParameter("NDEP", dtype=ti.i32, mode="scalar", value=0, pool=pool) + deps = make_depressions("taichi", grid, ndep_p, method="vanilla", reroute="carve", n_flat=n_flat) + +`method` ("vanilla"|"optimized") picks how basins are labelled and, for +reroute="carve", how the carve itself runs; `reroute` ("carve"|"jump") picks +how a resolved basin's pit is reconnected to its outlet. All four +combinations build. `depression_counter_p` is a bare caller-allocated scalar +i32 Parameter, bound directly at the returned kernel's own `NDEP` PARAM +slot, post-`.build()` - there is no Need indirection anywhere in this stack; +the underlying Parameter is +not built here, since this factory takes no pool: every scratch buffer below +is a caller-supplied data arg, never a bound field Parameter. + +Every buffer is n_flat-sized, since a per-basin array is indexed by basin id +and basin id = pit index + 1 (bid/basin_saddlenode/outlet range over the +same 0..n_flat-1 index space as every per-node buffer). Required data args, +by Bag member: + + "ndep": the `ndep_p` scalar Parameter itself (bag.ndep.read()) + "depression_counter": closure: (rec,) - accumulates into ndep_p, bound + directly as a raw field. cupy: (rec, ndep) - ndep_p + is only ever reached through $...$ get() spans + there, which registers it read-only in the + constant block (see compile_cupy.py's + _register_ptr), so the caller instead + passes `ndep_p.get().data` positionally, same as + `rec` - see build_depression_counter in + _cupy_depressions.py. Either way the caller must + ndep_p.set(0) before each launch (mirrors + ops.Reduce.run_sum). + "copy_field": (src, dst) i32 -> dst[:] = src[:] + "label_basins" (vanilla): Routine, data names ("rec", "bid", "rec_jump") + "label_basins" (optimized, closure): Kernel, data args (rec, rec_jump, bid) + "label_basins" (optimized, cupy): Routine, data names ("rec", "rec_jump", "bid") + - see _closure_depressions.py/_cupy_depressions.py's build_basin_labelling_* for + why cupy needs three real launches where a closure backend needs one. + "saddlesort": Routine (6 kernels, unchanged by `method`), data + names ("bid", "z", "z_prime", "is_border", + "basin_saddle", "basin_saddlenode", "outlet") - + the six constituent KernelBuilders are also + exposed under "saddlesort_". + "reroute" (carve, vanilla): Routine, data names ("rec", "rec_work", + "rec_jump", "tag", "tag_alt", "bid", + "basin_saddlenode", "outlet", "rerouted") + "reroute" (carve, optimized): Kernel, data args (rec, basin_saddlenode, outlet) + "reroute" (jump, closure): Kernel, data args (rec, outlet, rerouted) + "reroute" (jump, cupy): Routine, data names ("rec", "outlet", "rerouted") + - split for the same real-launch-barrier reason as label_basins above. + +Every "label_basins"/"reroute" constituent KernelBuilder is also exposed +under "label_basins_"/"reroute_" respectively, mirroring +make_accumulation's own routine+constituent-kernels convention. + +`reroute_jump`'s pit write is deliberately `rec[i - 1]`, not `rec[i]`: the +loop is over basin ids and basin id = pit index + 1 - ported exactly as +legacy has it, not "fixed" - see _closure_depressions.py's build_reroute_jump. + +`ops.bitpack` (pack/unpack_value/unpack_index) replaces legacy's +f32_i32_struct module for the lexicographic (elevation, target-basin) and +(elevation, node) argmins saddlesort's atomic_min passes need; on cupy, +i64 atomic_min is a CAS loop (CUDA has no native atomicMin over signed long +long) - see _cupy_depressions.py's build_atomic_min_ll. + +make_depressions builds the routines/kernels only. make_depression_solver +wraps that Bag into the outer host-driven loop the algorithm needs, as a +compiled Sequence: + + solver = make_depression_solver( + "taichi", deps, method="vanilla", reroute="carve", + rec=rec.data, z=z.data, bid=bid.data, ...) + solver() + solver.last_trip_counts # passes the loop actually took + + depression_counter -> ndep; ndep == 0 -> nothing to do + loop max_times = ceil(log2(max(2, ndep))) + 2: +