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/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..be47f12 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py @@ -0,0 +1,351 @@ +""" +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 add_swap standing in for the python-level `T0, T1 = T1, T0`: + + 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) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .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..0b5ef93 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py @@ -0,0 +1,312 @@ +""" +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 +add_swap standing in for the python-level `T0, T1 = T1, T0`: + + 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) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .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..2f7e8b5 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py @@ -0,0 +1,312 @@ +""" +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 +add_swap standing in for the python-level `T0, T1 = T1, T0`: + + 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) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(apply_source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .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..2ccf0e8 --- /dev/null +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -0,0 +1,717 @@ +""" +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 Bag 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. + +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) +""" + +import ast +import copy +import hashlib +import inspect +import linecache +from types import FunctionType +from typing import Any, ClassVar + +import numpy as np + +from .compile import ( + HelperBuilder, + Kernel, + KernelBuilder, + _SpecializedHelper, + _SpecializeCtx, + capture_template_meta, + filter_bindings, + resolve_binding, +) +from .parameter import MODES, Parameter +from .routine import Routine, RoutineBuilder, _CompiledStep, _template_label + + +def specialize_closure(template, bindings: dict[str, Any], ctx: _SpecializeCtx) -> FunctionType: + """ + Rebuild `template` as a new function whose globals carry the resolved + bindings, 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. `ctx` is the compile this specialization + belongs to - it is what lets a bound HelperBuilder be specialized here, + against these same bindings, rather than standing for a stale one. + + Author: B.G (07/2026) + """ + resolved = {name: resolve_binding(value, ctx) for name, value in bindings.items()} + source = getattr(template, "__wrapped__", template) + func_globals = dict(source.__globals__) + func_globals.update(resolved) + + 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. + + 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. + + 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. + + 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 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] + + # No Parameter/HelperBuilder/Bag ever appears in these two bindings + # dicts, so the ctx each specialize_closure call needs is never + # actually consulted; a throwaway one per call is enough. + get_fn = backend.func( + specialize_closure( + get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle, "STATIC": backend.static}, _SpecializeCtx() + ) + ) + + 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}, _SpecializeCtx()) + ) + + return ClosureParamDeviceView(self.name, get_fn, set_fn) + + +class ClosureHelper(_SpecializedHelper): + """ + A device helper's specialization, compiled to a ti.func or qd.func. + Produced by a ClosureHelperBuilder as part of an enclosing kernel's + compile(); see HelperBuilder. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + super().__init__() + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + The raw ti.func/qd.func, for binding into another template's body. + + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, **kwargs): + """ + A ti.func/qd.func only runs inside kernel/func scope; callers use + `.compiled` there. + + Author: B.G (07/2026) + """ + raise RuntimeError(f"Helper '{self.name}' is only callable from kernel/func scope, not host Python") + + +class ClosureKernel(Kernel): + """ + A launchable kernel compiled to a ti.kernel or qd.kernel. + + Its call signature is the template's own, which declares data arguments + only - `def template(out: ti.template()): ...` - since bound objects reach + the body through globals instead. See compile.py. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + super().__init__() + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + The raw ti.kernel/qd.kernel behind this Kernel's __call__. + + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, **kwargs): + """ + Launches the compiled kernel. Args are data fields only. + + Author: B.G (07/2026) + """ + return self._compiled(*args, **kwargs) + + +class ClosureHelperBuilder(HelperBuilder): + """ + Recipe for a device helper compiled to a ti.func/qd.func. Subclasses pin + `_backend`. Specialized only as part of an enclosing kernel's compile() + - see HelperBuilder; compile() itself raises. + + Author: B.G (07/2026) + """ + + _backend: ClassVar[Any] + + def _specialize(self, ctx: _SpecializeCtx) -> ClosureHelper: + """ + Splice the referenced bindings into the template's globals against + `ctx`, and compile the result as a device func. + + Author: B.G (07/2026) + """ + specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) + return ClosureHelper(specialised.__name__, self._backend.func(specialised)) + + +class ClosureKernelBuilder(KernelBuilder): + """ + Compiles an ingested def into a launchable kernel. Subclasses pin + `_backend`. + + Author: B.G (07/2026) + """ + + _backend: ClassVar[Any] + + def compile(self) -> ClosureKernel: + """ + Splice the referenced bindings into the template's globals - each + compile() opening a fresh _SpecializeCtx, so every HelperBuilder + reachable from this kernel's bindings is specialized once, against + these bindings - and compile the result as a launchable kernel. + + Author: B.G (07/2026) + """ + ctx = _SpecializeCtx() + specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) + return ClosureKernel(specialised.__name__, self._backend.kernel(specialised)) + + +class _RenameNames(ast.NodeTransformer): + """ + Substitute every `ast.Name` whose id is a key of `mapping` with a fresh + Name carrying the mapped id, leaving everything else - including the + ctx (Load/Store/Del) of the node being replaced - untouched. + + Used to retarget a step's template argument names (`T_out`, `T_in`) onto + the routine's own data names (`T1`, `T0`) inside that step's body only; + bound names (`heat`, `N`, ...) are never touched since they are never + template arguments. + + Author: B.G (07/2026) + """ + + def __init__(self, mapping: dict[str, str]): + self._mapping = mapping + + def visit_Name(self, node: ast.Name) -> ast.Name: + new_id = self._mapping.get(node.id) + if new_id is None: + return node + return ast.copy_location(ast.Name(id=new_id, ctx=node.ctx), node) + + +def _top_level_assigned_names(stmt: ast.stmt) -> set[str]: + """ + Names a top-level Assign/AnnAssign/AugAssign statement binds. + + Only `ast.Name` targets count - `x = 1` binds `x`, including through + tuple/list unpacking (`a, b = ...`); `arr[i] = v` and `obj.attr = v` + write through a subscript or attribute and bind no new top-level python + name, so they are not collected. + + Author: B.G (07/2026) + """ + if isinstance(stmt, ast.Assign): + targets = stmt.targets + elif isinstance(stmt, (ast.AnnAssign, ast.AugAssign)): + targets = [stmt.target] + else: + return set() + + names: set[str] = set() + + def collect(target: ast.expr) -> None: + if isinstance(target, ast.Name): + names.add(target.id) + elif isinstance(target, (ast.Tuple, ast.List)): + for elt in target.elts: + collect(elt) + + for target in targets: + collect(target) + return names + + +def _check_fusable(label: str, func_def: ast.FunctionDef) -> None: + """ + Enforce the two structural fusable-template constraints on one step's + body, raising with `label` naming the offending step. + + No return anywhere: a return inside a spliced body would exit the whole + fused kernel, not just this step's contribution to it. Flat splice only: + once a top-level `for` loop has started, every remaining top-level + statement must also be a `for` loop - an optional preamble is allowed + before the first one, but nothing may follow the loops or sit between + them, since wrapping or interleaving would stop Taichi parallelizing + them as independent loops. + + Author: B.G (07/2026) + """ + for node in ast.walk(func_def): + if isinstance(node, ast.Return): + raise ValueError(f"fuse: step {label!r} contains a return statement, which would exit the whole fused kernel") + + seen_loop = False + for stmt in func_def.body: + if isinstance(stmt, ast.For): + seen_loop = True + elif seen_loop: + raise ValueError( + f"fuse: step {label!r} has a statement after a top-level for loop; a fusable template " + "must be an optional preamble followed only by top-level for loops, spliced flat" + ) + + +def _annotation_root_name(annotation: "ast.expr | None") -> "str | None": + """ + The bare name at the root of an annotation expression, e.g. `ti` out of + `ti.template()` or `qd.Tensor`. None if the annotation is missing or not + shaped as a dotted/called attribute access. + + Author: B.G (07/2026) + """ + node = annotation + if isinstance(node, ast.Call): + node = node.func + while isinstance(node, ast.Attribute): + node = node.value + return node.id if isinstance(node, ast.Name) else None + + +class ClosureRoutineBuilder(RoutineBuilder): + """ + Compiles a linear sequence of Taichi/Quadrants kernels sharing one bag. + + A step's data arity is read straight off its template's own python + signature - `def diffuse(T_out, T_in)` declares two - since bound objects + never appear there. Launching a compiled step is just calling it: Taichi + and Quadrants kernels derive their own launch range from the template, so + `grid`/`block` (cupy-only - see CupyRoutineBuilder) are accepted and + ignored. + + compile() defaults to fused=True: consecutive steps (up to a split() + boundary) are spliced into one generated kernel rather than launched as + separate ones - see compile() and _fuse_group(). + + Author: B.G (07/2026) + """ + + def _data_arity(self, kernel_builder: KernelBuilder) -> int: + template = kernel_builder.template + if template is None: + raise ValueError("add_kernel: kernel_builder has no ingested template") + return len(inspect.signature(template).parameters) + + def _make_caller(self, compiled_kernel, grid, block): + return compiled_kernel + + def compile(self, fused: bool = True, dump_source: str | None = None) -> "Routine": + """ + Validate (RoutineBuilder._validate) and compile every step. + + fused=False falls back to the base implementation: one kernel per + step, each an ordinary specialize_closure compile, exactly as + before fusion existed. This is the reference the fused path is + diffed against, and stays reachable as a runtime switch. + + fused=True (the default) compiles each split()-delimited group of + steps into one generated kernel: every step's top-level `for` loops + are concatenated flat into a single generated `def`, in the order + the steps were added, and that `def` is compiled as one + ti.kernel/qd.kernel. See _fuse_group for the mechanics and the + constraints this enforces. + + `dump_source`, fused mode only, is a file path; when given, every + generated group's source is appended to it (truncated first), + separated by a header naming the group. No file is written when + `dump_source` is None. + + Author: B.G (07/2026) + """ + if not fused: + return super().compile(fused=False) + + self._validate() + + compiled_steps: list[_CompiledStep] = [] + data_names: list[str] = [] + for group_index, group in enumerate(self._grouped_steps()): + kernel, group_data_names = self._fuse_group(group, group_index, dump_source) + caller = self._make_caller(kernel, None, None) + compiled_steps.append(_CompiledStep(caller, tuple(group_data_names))) + for name in group_data_names: + if name not in data_names: + data_names.append(name) + + defaults = {name: self._data[name] for name in data_names} + return Routine(compiled_steps, tuple(data_names), defaults) + + def _fuse_group(self, group: list, group_index: int, dump_source: "str | None"): + """ + Splice one split()-delimited group of steps into a single generated + kernel. + + One `_SpecializeCtx` covers the whole group, so a HelperBuilder + reachable from two of these steps is specialized once and both call + sites in the generated body share the specialized object - the same + guarantee a single ordinary kernel compile gives within itself. + + Per step, in order: the template's own AST is deep-copied out of the + capture_template_meta cache (never mutated in place - that cache is + shared with every other compile of the same template); the two + structural fusable-template constraints are checked + (_check_fusable); its top-level assignments are checked against + every earlier step in this group for a name collision + (_top_level_assigned_names); its template argument names are + substituted for this step's canonical_refs, positionally + (_RenameNames) - the same mapping data_handle_ref set up at + add_kernel time; and its (now renamed) body statements are appended + to the generated function's body, flat. + + The generated function's parameters are this group's data names, in + first-appearance order, each carrying the annotation its first + occurrence used (`ti.template()`, `qd.Tensor`, ...). The result is + unparsed to source, registered in linecache under a synthetic + filename - required for Taichi/Quadrants to re-parse it via + inspect.getsource when they run their own AST transform - exec'd, + and decorated as a kernel with the group's resolved bindings plus + every step's own template module globals injected. + + Author: B.G (07/2026) + """ + backend = group[0].kernel_builder._backend + ctx = _SpecializeCtx() + body_stmts: list[ast.stmt] = [] + data_names: list[str] = [] + annotations: dict[str, "ast.expr | None"] = {} + assigned_by: dict[str, str] = {} + module_globals: dict[str, Any] = {} + resolved_globals: dict[str, Any] = {} + + for step_index, step in enumerate(group): + kernel_builder = step.kernel_builder + template = kernel_builder.template + label = f"group{group_index}/step{step_index}:{_template_label(template)}" + + _, tree = capture_template_meta(template) + if tree is None or not tree.body or not isinstance(tree.body[0], ast.FunctionDef): + raise ValueError(f"fuse: step {label!r} has no recoverable source to fuse") + func_def = copy.deepcopy(tree.body[0]) + + _check_fusable(label, func_def) + + for stmt in func_def.body: + for name in _top_level_assigned_names(stmt): + prior = assigned_by.get(name) + if prior is not None: + raise ValueError( + f"fuse: top-level name '{name}' is assigned by both {prior!r} and {label!r}" + ) + assigned_by[name] = label + + params = [a.arg for a in func_def.args.args] + if len(params) != len(step.canonical_refs): + raise ValueError( + f"fuse: step {label!r} declares {len(params)} data argument(s), " + f"routine gives it {len(step.canonical_refs)}" + ) + rename_map = dict(zip(params, step.canonical_refs)) + + for arg_node, canon in zip(func_def.args.args, step.canonical_refs): + if canon not in data_names: + data_names.append(canon) + annotations[canon] = arg_node.annotation + + renamer = _RenameNames(rename_map) + renamed_body = [renamer.visit(stmt) for stmt in func_def.body] + body_stmts.extend(renamed_body) + + filtered = filter_bindings(template, kernel_builder.bindings) + module_globals.update(dict(getattr(template, "__globals__", {}))) + resolved_globals.update({name: resolve_binding(value, ctx) for name, value in filtered.items()}) + + if not body_stmts: + raise ValueError(f"fuse: group {group_index} produced an empty body") + + # module_globals is a base layer (ti/qd/np/builtins/plain helper functions + # each step's own module happens to define) applied once, in step order; + # resolved_globals - the actual bound objects this group's steps + # reference - is applied last so a name a later step's unrelated module + # globals happen to share (e.g. two templates in one file both seeing + # `heat` in their module namespace) never clobbers an earlier step's + # resolved binding. + exec_globals: dict[str, Any] = {} + exec_globals.update(module_globals) + exec_globals.update(resolved_globals) + + for annotation in annotations.values(): + alias = _annotation_root_name(annotation) + if alias is not None: + exec_globals.setdefault(alias, backend) + + # Name the fused function deterministically from its own content + # rather than a process-global uid: build it under a stable + # placeholder name first, unparse, hash that text, then rename to + # `_fused_group{group_index}_{hash}` and re-unparse so the source, + # linecache entry, and compiled filename all agree on the real name. + # Identical routine -> identical hash -> identical name -> the + # backend's own offline cache (Taichi/Quadrants re-parse the kernel + # via inspect.getsource off this source+filename) hits across process + # restarts; an unrelated upstream allocation shift no longer touches + # this name since it never depended on uid in the first place. + placeholder_name = f"_fused_group{group_index}_placeholder" + args_node = ast.arguments( + posonlyargs=[], + args=[ast.arg(arg=name, annotation=annotations.get(name)) for name in data_names], + vararg=None, + kwonlyargs=[], + kw_defaults=[], + kwarg=None, + defaults=[], + ) + fused_def = ast.FunctionDef( + name=placeholder_name, args=args_node, body=body_stmts, decorator_list=[], returns=None + ) + module = ast.fix_missing_locations(ast.Module(body=[fused_def], type_ignores=[])) + placeholder_source = ast.unparse(module) + digest = hashlib.sha256(placeholder_source.encode()).hexdigest()[:12] + func_name = f"_fused_group{group_index}_{digest}" + fused_def.name = func_name + source = ast.unparse(module) + + filename = f"" + linecache.cache[filename] = (len(source), None, source.splitlines(keepends=True), filename) + + if dump_source: + mode = "w" if group_index == 0 else "a" + with open(dump_source, mode) as fh: + fh.write(f"# --- group {group_index} ({filename}) ---\n{source}\n\n") + + code = compile(source, filename, "exec") + exec(code, exec_globals) + fused_fn = exec_globals[func_name] + + krn = ClosureKernel(func_name, backend.kernel(fused_fn)) + return krn, data_names diff --git a/pyfastflow/experimental/core/context/backends.py b/pyfastflow/experimental/core/context/backends.py new file mode 100644 index 0000000..ae932d2 --- /dev/null +++ b/pyfastflow/experimental/core/context/backends.py @@ -0,0 +1,82 @@ +""" +Per-backend wiring shared by every Bag factory (make_grid, make_noise, ...). + +A factory needs three things to build its Parameters and Helpers against a +chosen backend name: the backend module itself (for a block that needs e.g. +`backend_mod.abs`), the Parameter/HelperBuilder classes to construct with, +and the backend's own dtype objects keyed by the short names factories write +their Parameters with ("i32", "f32", "u8", "u32"). backend_classes() is the +one place that knows the mapping from a backend name to those three 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. + +make_helper() is the HelperBuilder assembly every block module's +build_helpers() performs once per block. + +Author: B.G (07/2026) +""" + +from typing import Any + +import numpy as np + +from .compile import HelperBuilder + + +def backend_classes(backend: str): + """ + (backend_module_or_None, ParameterCls, HelperBuilderCls, dtypes) for one + backend name. + + `backend_module_or_None` is `ti`/`qd` for the closure backends, `None` for + cupy - cupy blocks call plain C, never a bound backend module. `dtypes` + maps "i32"/"f32"/"u8"/"u32" to that backend's own dtype objects (ti.*/qd.* + for the closure backends, numpy dtypes for cupy) - every name either + make_grid or make_noise currently needs, plus the obvious siblings. + + 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 TaichiHelperBuilder, TaichiParameter + + return ti, TaichiParameter, TaichiHelperBuilder, { + "i32": ti.i32, "f32": ti.f32, "u8": ti.u8, "u32": ti.u32, + } + if backend == "quadrants": + import quadrants as qd + + from .quadrants_backend import QuadrantsHelperBuilder, QuadrantsParameter + + return qd, QuadrantsParameter, QuadrantsHelperBuilder, { + "i32": qd.i32, "f32": qd.f32, "u8": qd.u8, "u32": qd.u32, + } + if backend == "cupy": + from .cupy_backend import CupyHelperBuilder, CupyParameter + + return None, CupyParameter, CupyHelperBuilder, { + "i32": np.int32, "f32": np.float32, "u8": np.uint8, "u32": np.uint32, + } + raise ValueError(f"unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + + +def make_helper(HelperCls, template, **binds: Any) -> HelperBuilder: + """ + A HelperBuilder ingesting `template` with every entry of `binds` applied. + + The assembly every block module's build_helpers() needs for each of its + blocks, so a new block module does not carry its own copy. + + Author: B.G (07/2026) + """ + builder = HelperCls().ingest(template) + for name, obj in binds.items(): + builder.bind(name, obj) + return builder diff --git a/pyfastflow/experimental/core/context/bag.py b/pyfastflow/experimental/core/context/bag.py new file mode 100644 index 0000000..8bf705a --- /dev/null +++ b/pyfastflow/experimental/core/context/bag.py @@ -0,0 +1,364 @@ +""" +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 at compile time, by resolve_binding in compile.py. That is why this module +stands on its own: it depends on nothing else here beyond the shared uid +counter. + +merge/extract/trim/replace/from_builder 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. + + Two ways to use one: 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 bind_bag(bag) to merge every member in at top level under its own name. + + 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`. Raises if `name` is already taken. + + 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'. + + 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. + + 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. + + 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. + + Raises if any segment is missing or if 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. + + 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. Raises if any path does not resolve. + + 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. Raises 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. + + `name` may be a dotted path into nested Bags. + + 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. + + 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) + + +def from_builder(builder: "CompileBuilder") -> "Bag": + """ + A new Bag holding a builder's current bindings under their existing + names. + + A snapshot at call time: later bind() / rebind() calls on `builder` do + not retroactively change the returned Bag, and adding to the Bag does + not reach back into the builder. + + Author: B.G (07/2026) + """ + return Bag(dict(builder.bindings)) diff --git a/pyfastflow/experimental/core/context/compile.py b/pyfastflow/experimental/core/context/compile.py new file mode 100644 index 0000000..475b940 --- /dev/null +++ b/pyfastflow/experimental/core/context/compile.py @@ -0,0 +1,515 @@ +""" +Turning a template plus its bindings into something the backend can run. + +A builder collects bindings and a template; compile() specializes them into a +Specializable - a launchable Kernel, or, for a helper, the internal object an +enclosing kernel's compile() produces. resolve_binding is the hub of it: it +decides what a bound object becomes inside a template body, dispatching on +whether it is a Parameter, a HelperBuilder, an already-specialized object, or a +Bag. + +Everything here is one interlocking piece - _SpecializeCtx specializes a +HelperBuilder, which resolves its own bindings back through resolve_binding, +which may reach another HelperBuilder through a _LazyBagView. It reads +Parameter (parameter.py) and Bag (bag.py) but neither reads it back. + +Only the abstract builders live here; the concrete Taichi*, Quadrants* and +Cupy* ones sit alongside in their own modules. See parameter.py's module docstring +for what the whole scheme is for. + +Author: B.G (07/2026) +""" + +import ast +import inspect +import warnings +from abc import ABC, abstractmethod +from functools import lru_cache +from typing import Any + +from ..pool.base import new_uid +from .bag import Bag, from_builder +from .parameter import Parameter + + +class Specializable(ABC): + """ + Anything produced by specializing a template against bindings: a + launchable Kernel, or the internal object a HelperBuilder specializes + into as part of an enclosing kernel's compile(). + + The builder that produced this object stays authoritative for its recipe + - template and bindings - see CompileBuilder. + + Author: B.G (07/2026) + """ + + name: str + + def __init__(self): + """ + Assign this object's process-wide uid. Concrete backends call this + first in their own __init__. + + Author: B.G (07/2026) + """ + self._uid = new_uid() + + @property + def uid(self) -> int: + """ + Process-wide identity assigned at construction. See Parameter.uid. + + Author: B.G (07/2026) + """ + return self._uid + + @property + @abstractmethod + def compiled(self): + """ + Raw backend callable/source (ti.func / qd.func object, CUDA text), + for injection into another template's bindings. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def __call__(self, *args, **kwargs): ... + + +class _SpecializedHelper(Specializable): + """ + A device helper's specialized backend object (e.g. a ti.func + specialization), produced by a HelperBuilder as part of an enclosing + kernel's compile(). No public class holds this between compiles - it + lives only inside a _SpecializeCtx.compiled and the Kernel body being + built alongside it; see HelperBuilder and _SpecializeCtx. + + Not necessarily callable from host Python - backends where device + helpers can only run inside kernel/func scope raise on __call__. + + Author: B.G (07/2026) + """ + + +class Kernel(Specializable): + """ + Compiled entry point (e.g. a ti.kernel specialization). + + Unlike a helper's specialization, __call__ works from host Python - that + is how compute gets launched. Its arguments are the template's own + declared data arguments; see the module docstring on data at call time. + + Author: B.G (07/2026) + """ + + +class _SpecializeCtx: + """ + The state shared by every resolution happening inside one compile(). + + A HelperBuilder is a recipe, not something compiled ahead of time - it is + specialized here, on demand, the first time this compile reaches it. + `specialize` memoizes on the builder's uid so a helper reachable from two + places in one compile - bound flat and inside a Bag, or under two + different names - is specialized exactly once and both call sites share + the same object. The memo lives only as long as this ctx, i.e. one + compile(): a later compile against different bindings gets its own ctx + and specializes afresh, which is what lets a recompiled kernel pick up a + changed const in a helper it binds. + + `_active` catches a helper cycle - builder A binding builder B which + (directly or transitively) binds A back - by raising instead of + recursing forever. + + Author: B.G (07/2026) + """ + + def __init__(self): + self._memo: dict[int, Any] = {} + self._active: set[int] = set() + + def specialize(self, builder: "HelperBuilder") -> Any: + """ + This builder's specialized object for the compile this ctx belongs + to, specializing it on first request and returning the memoized + result on every later one. + + Author: B.G (07/2026) + """ + uid = builder.uid + cached = self._memo.get(uid) + if cached is not None: + return cached + if uid in self._active: + name = getattr(builder.template, "__name__", builder.template) + raise RecursionError(f"helper cycle detected while specializing '{name}' (uid {uid})") + self._active.add(uid) + try: + specialized = builder._specialize(self) + finally: + self._active.discard(uid) + self._memo[uid] = specialized + return specialized + + +class _LazyBagView: + """ + What a Bag looks like from inside a template body. + + `phys.g` resolves member `g` on first access and caches the result in the + instance dict, so later lookups skip __getattr__ altogether. Resolving a + member is not free - for a Parameter it compiles a device view, for a + HelperBuilder it specializes the helper - and a template usually touches + only a few members of the bag it binds, so resolution is deferred to the + members actually named. Members that are themselves Bags resolve to + another _LazyBagView, keeping nested bags lazy all the way down, and + carry the same ctx so a HelperBuilder reached through a nested Bag + specializes against the same compile as one bound flat. + + Author: B.G (07/2026) + """ + + def __init__(self, bag: "Bag", ctx: "_SpecializeCtx"): + object.__setattr__(self, "_bag", bag) + object.__setattr__(self, "_ctx", ctx) + + def __getattr__(self, name: str) -> Any: + # only called on a genuine miss (cached hits never reach here) + bag = object.__getattribute__(self, "_bag") + ctx = object.__getattribute__(self, "_ctx") + if name not in bag: + raise AttributeError(name) + resolved = resolve_binding(bag[name], ctx) + self.__dict__[name] = resolved + return resolved + + +def resolve_binding(value, ctx: "_SpecializeCtx"): + """ + Turn a bound object into what a template body should see in its place. + + Parameter device_view() - the carrier of .get / .set_node for + in-kernel access. + HelperBuilder specialized against `ctx` (memoized - see _SpecializeCtx) + and replaced with its .compiled backend callable or + source. + Specializable its .compiled backend callable or source. + Bag a _LazyBagView carrying the same `ctx`, so a dotted path + like grid.nx.get(i) traces as plain attribute lookups and + a helper reached that way shares the ctx's memo. + anything else passed through untouched. + + `ctx` is the compile this resolution belongs to - see _SpecializeCtx. It + threads through every nested Bag and every helper-calling-helper + resolution so the whole compile shares one memo. + + Author: B.G (07/2026) + """ + if isinstance(value, Parameter): + return value.device_view() + if isinstance(value, HelperBuilder): + return ctx.specialize(value).compiled + if isinstance(value, Specializable): + return value.compiled + if isinstance(value, Bag): + return _LazyBagView(value, ctx) + return value + + +@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. + + Cached because every compile() asks once to filter bindings, and a miss + costs an inspect.getsource plus a parse. The tree handed back is + therefore shared by every Specializable built from 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. + + 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(source) + except SyntaxError: + tree = None + return source, tree + + +def _used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: + """ + The subset of `bindings` whose name appears in the template body. + + Collecting every ast.Name id in the tree is enough: the root of an + attribute chain - `phys` in `phys.g.get(0)` - is itself an ast.Name. + + With no AST to consult (a CUDA source string, or a def whose source cannot + be recovered) this returns `bindings` unchanged, rather than dropping a + binding it cannot prove is unused. + + Author: B.G (07/2026) + """ + _, tree = capture_template_meta(template) + if tree is None: + return bindings + used = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)} + return {name: value for name, value in bindings.items() if name in used} + + +def filter_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: + """ + The bindings a template actually references, ready to inject. + + Anything bound but never referenced is reported in a single warning per + compile, which is what catches a misspelled bind() name in a context with + many parameters. Nothing is reported when there is no AST to check + against, since _used_bindings then treats every binding as used. + + Author: B.G (07/2026) + """ + filtered = _used_bindings(template, bindings) + unused = sorted(set(bindings) - set(filtered)) + if unused: + warnings.warn( + f"template '{getattr(template, '__name__', '?')}': bound but unused: {unused}", + UserWarning, + stacklevel=3, + ) + return filtered + + +class CompileBuilder(ABC): + """ + Collects dependencies and a template, and compiles them into one + Specializable. + + A builder is used as a chain: any number of bind() calls, one ingest(), + then compile(). Bound objects are not inspected as they arrive - what + each one is (Parameter, Helper, Bag, handle, plain value) is worked out + when the template is specialized, so bind() accepts anything. + + The builder stays authoritative for the recipe throughout its life: + `template` and `bindings` below are a read-only view onto the same state + compile() reads, so a later layer can inspect a builder without reaching + into its private attributes. compile() does not consume or mutate that + state - bind() again, ingest() a different template, or just call + compile() again, and every callable made earlier stays exactly as it was. + Recompiling a builder that has not changed since its last compile() + produces an equivalent callable; there is no reason to do it, though + nothing breaks if it happens. + + Everything here is backend-independent. A backend supplies compile(), and + may override ingest() if its templates need different handling. + + Author: B.G (07/2026) + """ + + def __init__(self): + self._uid = new_uid() + self._bindings: dict[str, Any] = {} + self._bag_names: set[str] = set() + self._template = None + + @property + def uid(self) -> int: + """ + Process-wide identity assigned at construction. See Parameter.uid. + + Author: B.G (07/2026) + """ + return self._uid + + @property + def template(self): + """ + The currently ingested template. Read-only - go through ingest() to + change it. + + Author: B.G (07/2026) + """ + return self._template + + @property + def bindings(self) -> dict[str, Any]: + """ + The current name -> object bindings, in bind() order. Read-only - go + through bind() / bind_bag() to change them. This is a live view onto + the builder's own dict, not a snapshot - the builder stays + authoritative for its recipe (see the class docstring). + + Author: B.G (07/2026) + """ + return self._bindings + + def bind(self, name: str, obj: Any) -> "CompileBuilder": + """ + Register `obj` under `name` for injection into the template body. + + Binding a name a second time replaces what it pointed to - handy for + editing a builder in place before recompiling. The one case this + refuses is rebinding a name that arrived through bind_bag(): which + bag member is meant is ambiguous once the bag itself may have + changed, so this raises instead of guessing. + + Author: B.G (07/2026) + """ + if name in self._bag_names: + raise KeyError(f"'{name}' was bound via bind_bag() and cannot be rebound directly") + self._bindings[name] = obj + return self + + def bind_bag(self, bag: "Bag") -> "CompileBuilder": + """ + Bind every member of `bag` at top level under its own name (flat), for + when a kernel refers to members directly rather than via a bag path. + + Author: B.G (07/2026) + """ + for name, item in bag.items(): + self.bind(name, item) + self._bag_names.add(name) + return self + + def ingest(self, template) -> "CompileBuilder": + """ + Take the generic template (a python def for closure backends, a CUDA + source string for cupy). Backends may override with their own handling. + + Author: B.G (07/2026) + """ + self._template = template + return self + + def as_bag(self) -> "Bag": + """ + This builder's current bindings, regrouped as a Bag under their + existing names. + + This is extraction for reuse, not a rewrite: the template body still + reads the same names, so the returned bag's members are exactly what + the builder already binds. Merge it into another bag and rebind() + against the result to move a builder's dependencies around without + touching the template. + + Author: B.G (07/2026) + """ + return from_builder(self) + + def rebind(self, bag: "Bag") -> "CompileBuilder": + """ + Re-resolve every name this builder currently binds against `bag`, + replacing each binding with what `bag` holds under that name. + + Every bound name must be present in `bag`; if any are missing this + raises once, listing all of them rather than stopping at the first. + A name whose current binding is itself a Bag (bound with bind() as a + nested group) requires `bag` to carry a Bag under that name too - a + template reaching it by dotted path needs the same shape on the + other end. + + This replaces bindings regardless of how they arrived, including + names bound via bind_bag() - superseding those is the point, so it + does not go through bind() and does not hit its bag-name raise. + Which names came from bind_bag() is unchanged by this call: rebind + swaps values, not the origin bookkeeping that governs future bind() + calls. + + Author: B.G (07/2026) + """ + missing = [name for name in self._bindings if name not in bag] + if missing: + raise KeyError(f"rebind: not found in bag: {sorted(missing)}") + for name, old in self._bindings.items(): + new = bag[name] + if isinstance(old, Bag) and not isinstance(new, Bag): + raise TypeError( + f"rebind: '{name}' is bound to a nested Bag; replacement must " + f"also be a Bag, got {type(new).__name__}" + ) + self._bindings[name] = new + return self + + @abstractmethod + def compile(self) -> Specializable: + """ + Produce a compiled Kernel from the builder's current template and + bindings. Does not consume or mutate the builder - the same builder + may be compiled again, with or without edits in between, and every + callable produced this way is independent of the others. + + On HelperBuilder this raises instead - see HelperBuilder.compile. + + Author: B.G (07/2026) + """ + ... + + +class HelperBuilder(CompileBuilder): + """ + Builds a device helper: template plus bindings, held purely as a recipe. + + A helper has no independent compiled form to keep between compiles - bind + this builder into a KernelBuilder, directly under a name or as a member + of a bound Bag, and the enclosing kernel's compile() specializes it + against that same compile's bindings. The same builder reached twice in + one compile is specialized once and shared; a later compile against + different bindings specializes it afresh. See the module docstring and + _SpecializeCtx. + + compile() therefore raises here: there is no standalone Helper for it to + return. _specialize(ctx) is the real entry point, called only by + _SpecializeCtx.specialize. + + Author: B.G (07/2026) + """ + + def compile(self) -> Specializable: + """ + Always raises - a HelperBuilder is never specialized on its own. + Bind it into a KernelBuilder (flat or inside a Bag) and call + compile() on that instead; the kernel's compile() specializes every + HelperBuilder it can reach as part of producing the kernel. + + Author: B.G (07/2026) + """ + raise TypeError( + "HelperBuilder.compile() is not supported: a device helper is specialized " + "by the kernel that binds it, not on its own. Bind this builder into a " + "KernelBuilder (directly or inside a Bag) and call compile() on that builder." + ) + + @abstractmethod + def _specialize(self, ctx: "_SpecializeCtx") -> Specializable: + """ + Produce this helper's specialized backend object for the compile + `ctx` belongs to. Called at most once per compile, by + _SpecializeCtx.specialize, which memoizes the result on this + builder's uid - not meant to be called directly. + + Author: B.G (07/2026) + """ + ... + + +class KernelBuilder(CompileBuilder): + """ + Builds a Kernel. compile() -> host-callable Kernel. + + Author: B.G (07/2026) + """ + + diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py new file mode 100644 index 0000000..4fcfa63 --- /dev/null +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -0,0 +1,985 @@ +""" +cupy implementations of Parameter, Kernel and their builders, plus +CupyHelperBuilder - the recipe for a device helper, specialized as part of +whichever kernel binds it (see compile.py, HelperBuilder). + +Here 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 _SpanParser._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 - see _SpecializeCtx), and the unit's `__global__` kernel. +`CupyKernel._raw_cache` keys the module on its final source text, same as the +single-kernel cache did before - the whole specialization still reduces to +that text, local-index-qualified struct members included, so it remains a +sound key, and since that index no longer depends on uid the key (and so the +cache hit rate) is stable across process restarts for an unchanged program. +A cache hit skips recompilation but the constant block is re-uploaded +regardless, since the pointers a compile's bindings currently resolve to are +not part of what the cache key captures. + +Author: B.G (07/2026) +""" + +import re +from typing import Any + +import cupy as cp +import numpy as np + +from .bag import Bag +from .compile import HelperBuilder, Kernel, KernelBuilder, _SpecializedHelper, _SpecializeCtx +from .parameter import MODES, Parameter +from .routine import Routine, RoutineBuilder, _CompiledStep + +_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) +_SPAN_RE = re.compile(r"\$(.*?)\$", re.S) +_CALL_RE = re.compile(r"([\w.]+)\s*(?:\((.*)\))?\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 _walk(path: list[str], bindings: dict[str, Any]): + """ + Resolve a dotted path against bindings, descending Bag members. + + Author: B.G (07/2026) + """ + obj = bindings[path[0]] + for seg in path[1:]: + obj = obj[seg] if isinstance(obj, Bag) else getattr(obj, seg) + return obj + + +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 _SpanParser._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]}" + + +class _SpanParser: + """ + Expands the `$...$` spans in a CUDA template body. + + A span reaching a scalar or field Parameter registers it into `ctx`'s + shared pointer registry (`ctx.cupy_ptr_registry`, one dict per compile() + - see CupyKernelBuilder.compile) rather than generating a call argument: + the registry is what becomes the module's `pf_params` constant block. + `local_ptrs` tracks, separately, only the pointers *this* body actually + used - CupyKernelBuilder.compile and CupyHelperBuilder._specialize use it + to prepend that function's own `__restrict__` locals, so a function + declares locals for what it reads or writes and nothing else. `ctx` is the + compile this parse belongs to - a span reaching a CupyHelperBuilder + specializes it against `ctx` (memoized there - see _SpecializeCtx), so a + helper reached twice in one compile is specialized once, and any pointer + it registers lands in the same shared registry as the kernel's own. + + A span reaching a CupyHelper does not collect its source here: that would + mean two parents sharing one leaf helper each carry a full copy of the + leaf's text in their own body, so the translation unit ends up with that + `__device__` function defined twice the moment both parents are bound + into one kernel. Every reachable helper's own (non-nested) source is + instead registered once, by name, into `ctx.cupy_device_srcs` - see + CupyHelperBuilder._specialize and CupyKernelBuilder.compile, which is + where the deduplicated, dependency-ordered set for the whole unit is + assembled. + + Author: B.G (07/2026) + """ + + def __init__(self, bindings: dict[str, Any], *, ctx: _SpecializeCtx): + self.bindings = bindings + self.ctx = ctx + self.local_ptrs: dict[int, dict] = {} # uid -> {ctype, write} + + def _register_ptr(self, param: Parameter, write: bool) -> str: + registry = getattr(self.ctx, "cupy_ptr_registry", None) + if registry is None: + registry = self.ctx.cupy_ptr_registry = {} + local_index = getattr(self.ctx, "cupy_local_index", None) + if local_index is None: + local_index = self.ctx.cupy_local_index = {} + uid = param.uid + entry = registry.get(uid) + if entry is None: + entry = {"ctype": _ctype(param.dtype), "write": False, "array": param.get().data} + registry[uid] = entry + if write: + entry["write"] = True + if uid not in local_index: + # first encounter of this Parameter in this compile - assign it + # the next local index, in traversal order (see _param_argname). + local_index[uid] = len(local_index) + local = self.local_ptrs.get(uid) + if local is None: + self.local_ptrs[uid] = {"ctype": entry["ctype"], "write": write} + elif write: + local["write"] = True + return _param_argname(param, local_index) + + def _expand_param(self, param: Parameter, method: str, call_args: list[str]) -> str: + if method == "get": + if param.mode == "const": + return _cuda_literal(param.get()) + argname = self._register_ptr(param, write=False) + if param.mode == "scalar": + return f"{argname}[0]" + idx = call_args[0] if call_args else "0" + return f"{argname}[{idx}]" + # set_node + if param.mode == "const": + raise ValueError(f"{param.name}: const parameter is read-only") + if len(call_args) != 2: + raise ValueError(f"{param.name}: set_node(node, value) takes two arguments") + argname = self._register_ptr(param, write=True) + node, val = call_args + return f"{argname}[0] = {val}" if param.mode == "scalar" else f"{argname}[{node}] = {val}" + + def _repl(self, match: re.Match) -> str: + cm = _CALL_RE.match(match.group(1).strip()) + if cm is None: + raise ValueError(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[-1] in ("get", "set_node"): + target = _walk(path[:-1], self.bindings) + if isinstance(target, Parameter): + return self._expand_param(target, path[-1], call_args) + + target = _walk(path, self.bindings) + if isinstance(target, CupyHelperBuilder): + target = self.ctx.specialize(target) + if isinstance(target, CupyHelper): + return f"{target.name}({argstr if argstr is not None else ''})" + if isinstance(target, Parameter): + return self._expand_param(target, "get", ["0"]) + return _cuda_literal(target) + + def parse(self, body: str) -> str: + """ + Expand every `$...$` span in `body`, accumulating local_ptrs (and, + transitively, `ctx.cupy_device_srcs` - see CupyHelperBuilder._specialize) + as a side effect, then prepend the `__restrict__` locals this body's + own spans implied - see _insert_locals. + + Author: B.G (07/2026) + """ + expanded = _SPAN_RE.sub(self._repl, body) + local_index = getattr(self.ctx, "cupy_local_index", None) or {} + return _insert_locals(expanded, self.local_ptrs, local_index) + + +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 _SpanParser._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 _const_defines(bindings: dict[str, Any], body: str) -> list[str]: + """ + A `#define NAME literal` for each top-level const Parameter whose name + appears in `body`, matched on word boundaries. + + Pass the span-expanded text, not the raw template: by then a span such as + `$phys.dx.get(0)$` is a numeric literal and no longer reads as a bare + mention of a const name. Names the source never mentions + are left undefined, keeping macros for identifiers like N, DIM, EPS or min + out of the translation unit, where they would rewrite unrelated code. + + Author: B.G (07/2026) + """ + return [ + f"#define {name} {_cuda_literal(obj.get())}" + for name, obj in bindings.items() + if isinstance(obj, Parameter) and obj.mode == "const" and re.search(rf"\b{re.escape(name)}\b", body) + ] + + +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 _SpanParser._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, but not inside CUDA graph capture + (see CupyRoutineBuilder.compile, which compiles every step - and so + performs every upload - before capture starts). + + 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. + + 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 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 + + +class CupyHelper(_SpecializedHelper): + """ + A device helper's specialization, held as CUDA `__device__` source text. + Produced by a CupyHelperBuilder as part of an enclosing kernel's + compile(); see HelperBuilder. + + Nothing is compiled at this stage: RawModule compiles whole translation + units, so `.compiled` hands back source, and the helper is compiled as part + of each kernel that splices it in. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, source: str): + super().__init__() + self.name = name + self._compiled_source = source + + @property + def compiled(self): + """ + The spliced `__device__` source text, for pasting into whatever + kernel or helper binds this one. + + Author: B.G (07/2026) + """ + return self._compiled_source + + def __call__(self, *args, **kwargs): + """ + CUDA source is not a python callable - a helper only runs inside a + compiled kernel's device code. + + Author: B.G (07/2026) + """ + raise RuntimeError( + f"Helper '{self.name}' is CUDA source, only callable from a compiled kernel's device code" + ) + + +class CupyKernel(Kernel): + """ + A launchable kernel, backed by a function pulled from a compiled + cp.RawModule (see CupyKernelBuilder.compile). + + Launch dimensions are explicit: a RawModule function has nothing like + Taichi's auto-ranging, so __call__ takes grid and block. The positional + arguments are exactly the template's own declared data arguments - every + bound scalar/field Parameter reaches this kernel through the module's + constant block instead of a launch argument, so there is nothing else to + append here the way an earlier version of this class needed to. + + Author: B.G (07/2026) + """ + + _raw_cache: dict[str, "cp.RawModule"] = {} + + def __init__(self, name: str, compiled, module: "cp.RawModule"): + super().__init__() + self.name = name + self._compiled = compiled + self._module = module + + @property + def compiled(self): + """ + The underlying RawModule function this Kernel's __call__ launches. + + Author: B.G (07/2026) + """ + return self._compiled + + @property + def module(self) -> "cp.RawModule": + """ + The cp.RawModule this kernel's function was pulled from - shared with + every other kernel compiled from the same final source text (see + CupyKernel._raw_cache). + + Author: B.G (07/2026) + """ + return self._module + + def __call__(self, *args, grid, block, **kwargs): + """ + Launches the compiled kernel. `grid`/`block` are int or tuple launch + dims, required since a RawModule function has no default range. + + Author: B.G (07/2026) + """ + grid = (grid,) if isinstance(grid, int) else tuple(grid) + block = (block,) if isinstance(block, int) else tuple(block) + return self._compiled(grid, block, tuple(args)) + + +class CupyHelperBuilder(HelperBuilder): + """ + Recipe for a device helper compiled from CUDA `__device__` source. A span + inside one may reach any Parameter mode and any other device helper, + exactly like a kernel's own spans - see the module docstring's constant + block. Specialized only as part of an enclosing kernel's compile() - see + HelperBuilder; compile() itself raises. + + Author: B.G (07/2026) + """ + + def _specialize(self, ctx: _SpecializeCtx) -> CupyHelper: + """ + Expand the template's spans against `ctx` and prepend the const + #defines it turned out to need, giving this helper's own `__device__` + text - its own body only, not any dependency's. Any scalar/field + Parameter a span here reaches is registered into `ctx.cupy_ptr_registry` + exactly as one reached from the enclosing kernel would be (see + _SpanParser) - the block that results is shared, so this helper and + its caller read the same pointer. + + This helper's own text is also registered, by name, into + `ctx.cupy_device_srcs` - the deduplicated, dependency-ordered set of + every `__device__` function the whole compile reaches (see + CupyKernelBuilder.compile). parser.parse() above resolves this + helper's own spans first, which is what recursively specializes (and + so registers) every helper *this* one calls before this line runs - + so by the time this helper registers itself, its own dependencies are + already in the registry, ahead of it. A helper already reached once in + this compile - by this call site or an earlier one - keeps its first + registration; `setdefault` leaves it untouched rather than duplicating + or reordering it. + + Author: B.G (07/2026) + """ + template = self._template + name = _extract_name(_DEVICE_NAME_RE, template, "__device__") + parser = _SpanParser(self._bindings, ctx=ctx) + body = parser.parse(template) + source = "\n".join(_const_defines(self._bindings, body) + [body]) + device_srcs = getattr(ctx, "cupy_device_srcs", None) + if device_srcs is None: + device_srcs = ctx.cupy_device_srcs = {} + device_srcs.setdefault(name, source) + return CupyHelper(name, source) + + +class CupyKernelBuilder(KernelBuilder): + """ + Builds a CupyKernel from CUDA `__global__` source. Spans expand and, for + scalar/field params, auto-generate the matching pointer args + launch + arrays. + + Author: B.G (07/2026) + """ + + def compile(self) -> CupyKernel: + """ + Expand the template's spans - the kernel's own and, recursively, + every CupyHelperBuilder they reach - prepend the whole unit's + `__device__` helper sources (deduplicated by name, dependency-first - + see CupyHelperBuilder._specialize and `ctx.cupy_device_srcs`), the + kernel's own const #defines, and the `pf_params` constant block the + whole unit's scalar/field Parameters collected into, then build (or + reuse, keyed by final source text) the cp.RawModule and pull this + kernel's function out of it. + + Opens a fresh _SpecializeCtx for this compile, so every + CupyHelperBuilder this kernel's spans reach is specialized once, + against these bindings, and every scalar/field Parameter any of them + binds lands in one shared pointer registry + (`ctx.cupy_ptr_registry`) - see _SpanParser and _param_block_source. + The same ctx is what lets a helper reachable from two of this + kernel's own bindings - or from two different helpers this kernel + reaches - contribute its `__device__` text to `ctx.cupy_device_srcs` + exactly once: emitting a shared leaf's definition once per + translation unit, rather than once per parent that calls it, is what + the module docstring's "reaches it exactly the way its caller does" + promise requires, and repeated emission is a compile error a CUDA + translation unit does not tolerate the way a repeated `#define` of + an identical macro does. + + The registry - and so the set of pointers to upload - is rebuilt by + parsing on every call, cache hit or not, since the final source text + (the cache key) does not capture what a Parameter's storage currently + points at. A cache hit therefore skips recompilation but never skips + the upload; see _upload_param_block. + + Author: B.G (07/2026) + """ + ctx = _SpecializeCtx() + ctx.cupy_ptr_registry = {} + ctx.cupy_local_index = {} + ctx.cupy_device_srcs = {} + template = self._template + name = _extract_name(_KERNEL_NAME_RE, template, "__global__") + parser = _SpanParser(self._bindings, ctx=ctx) + body = parser.parse(template) + # extern "C" linkage so cp.RawModule finds the entry by its plain name + # (C++ would name-mangle it); helper __device__ funcs stay mangled and + # link fine within the same translation unit. + if 'extern "C"' not in body: + body = body.replace("__global__", 'extern "C" __global__', 1) + registry = ctx.cupy_ptr_registry + local_index = ctx.cupy_local_index + source = "\n".join( + [_param_block_source(registry, local_index)] + + list(ctx.cupy_device_srcs.values()) + + _const_defines(self._bindings, body) + + [body] + ) + + module = CupyKernel._raw_cache.get(source) + if module is None: + module = cp.RawModule(code=source) + CupyKernel._raw_cache[source] = module + _upload_param_block(module, registry, local_index) + raw = module.get_function(name) + krn = CupyKernel(name, raw, module) + krn._final_source = source + return krn + + +class _CapturedRoutine(Routine): + """ + A Routine whose steps have been recorded into a CUDA graph; calling it + replays that graph instead of re-issuing each step's kernel launch. + + Built by CupyRoutineBuilder.compile(captured=True) (the default there). + Holds the same steps/data_names/defaults an uncaptured Routine would, so + introspection agrees between the two, plus the captured cp.cuda.Graph and + the private stream the capture was recorded on. + + Capture needs a stream of its own - the default stream cannot be put in + capture mode - but replay does not, and the graph is launched on the + caller's current stream. A routine that replayed on its private stream + would order against nothing the caller had queued, so reading a result + straight after a call would need a device-wide synchronize to be correct, + which no other launch in this package asks for. Launching on the current + stream keeps a captured Routine interchangeable with an uncaptured one and + with a plain Kernel: queue work, call, read. + + A captured graph bakes in the device pointers it was captured with - + every launch it holds already has its bulk-data argument list resolved, + the same arguments a step's `data_handle_ref` maps onto. A step's bound + scalar/field Parameters are not part of that argument list at all: they + reach the kernel through its module's `pf_params` constant block (see the + module docstring), read fresh by the device on every execution rather + than captured as part of the launch. Replay therefore sees whatever that + block currently holds - which is exactly what the block held at the last + compile(), since nothing but a compile()'s upload ever writes to it - so + the staleness rules below apply the same way whether a pointer reached a + step through its launch arguments or through its constant block. Call-time + data handle overrides (the `rout(A, B)` form) cannot be honoured against + a captured graph's already-resolved bulk-data arguments: doing so would + either use the wrong buffers silently or require a fresh capture on a + call site that looks like an ordinary launch, hiding a real performance + cliff behind what reads as a cheap replay. This raises instead - + compile(captured=False) for a routine that overrides are meant to work + against. + + See the module docstring's "Contract: no set()/destroy() mid-routine" for + the staleness rules a Routine has always had; capture adds one more way a + compiled Routine can go stale without anything checking for it: + - a write to a scalar or field Parameter reached by this routine's bag + goes through the same storage both the graph's launch arguments and its + steps' constant blocks already point at, so replay keeps seeing the new + value - this is the intended way to feed a captured routine changing + data, same as for an uncaptured one; + - set() on a const Parameter changes generated source, which the graph + never re-reads - the routine (and the graph baked into it) must be + recompiled; + - destroy() on any Parameter or data handle this routine reaches, or + anything else that returns a buffer to the pool, invalidates a pointer + the graph's launches were captured with or a step's constant block was + last uploaded with. Recompile. + None of this is enforced at runtime; it is exactly the discipline a + Kernel or an uncaptured Routine already asks for, just extended to a + graph's baked-in pointers - launch arguments and constant blocks alike - + as an added way "recompile after this" applies. + + Author: B.G (07/2026) + """ + + def __init__(self, steps: list, data_names: tuple, defaults: dict, graph, stream): + super().__init__(steps, data_names, defaults) + self._graph = graph + self._stream = stream + + def __call__(self, *args) -> None: + """ + Replay the captured graph on the caller's current stream, so the call + orders against surrounding work exactly as an uncaptured Routine's + launches would. Takes no arguments - see the class docstring for why + call-time data handle overrides are rejected rather than honoured or + silently re-captured. + + Author: B.G (07/2026) + """ + if args: + raise RuntimeError( + "Routine: this routine was compiled with captured=True; call-time data " + "handle overrides are not supported against a captured CUDA graph, since " + "the graph's launches already have their pointers baked in. Compile with " + "captured=False for a routine meant to be called with overrides, or build " + "a second routine over the override handles and capture that one." + ) + self._graph.launch() + + +class CupyRoutineBuilder(RoutineBuilder): + """ + Compiles an ordered sequence of compiled-kernel launches sharing one bag + into a Routine. + + A step's data arity is read off the `__global__` signature as written in + the ingested source - the template author's own declared data arguments, + the same ones data_handle_ref maps onto. Every bound scalar/field + Parameter a step's spans reach travels through its module's constant + block instead (see the module docstring), so nothing is appended to this + count at compile time the way an earlier version of this class needed to + account for. + + `grid`/`block` have no auto-ranging equivalent on cupy the way Taichi and + Quadrants derive one from the template, so they are resolved once per + step: whatever add_kernel(..., grid=..., block=...) gave that step, else + the default passed to this builder's own constructor. Neither given + raises at compile time, naming the step. + + Author: B.G (07/2026) + """ + + def __init__(self, *, grid=None, block=None): + super().__init__() + self._default_grid = grid + self._default_block = block + + def _data_arity(self, kernel_builder: KernelBuilder) -> int: + template = kernel_builder.template + if template is None: + raise ValueError("add_kernel: kernel_builder has no ingested template") + match = _KERNEL_SIG_RE.search(template) + if not match: + raise ValueError("add_kernel: could not find a __global__ signature in template source") + argstr = match.group(2).strip() + return len(_split_args(argstr)) if argstr else 0 + + def _make_caller(self, compiled_kernel, grid, block): + grid = grid if grid is not None else self._default_grid + block = block if block is not None else self._default_block + if grid is None or block is None: + raise ValueError( + f"CupyRoutineBuilder: no grid/block for step '{compiled_kernel.name}' - " + "pass grid=/block= to the builder or to this step's add_kernel" + ) + + def caller(*args): + return compiled_kernel(*args, grid=grid, block=block) + + return caller + + def compile(self, captured: bool = True, dump_source: str | None = None) -> Routine: + """ + Validate (RoutineBuilder._validate), compile every step's kernel, and + either return a Routine that launches them in order (captured=False) + or capture that same sequence of launches into a CUDA graph and + return a Routine that replays it (captured=True, the default here). + + `dump_source` is accepted for signature parity with the closure + backend's fused compile() and ignored - there is no generated source + on this backend either way. + + captured=False is exactly RoutineBuilder.compile's base behaviour: + one host-side launch per step, every call. It is the reference the + captured path is diffed against, and stays reachable as a runtime + switch - in particular it is the only way to get a Routine that + accepts call-time data handle overrides (see _CapturedRoutine). + + captured=True compiles every step exactly as captured=False does, + then: + 1. Warms up each step by launching it once, for real, on the + default stream. Each step's compile() (just above, in this same + method) already built its cp.RawModule and uploaded its constant + block before this point, so the warmup is not covering a lazy + module load the way it would have to for a JIT that only happened + on first launch; it remains cheap insurance against any other + first-launch cost CUDA's capture machinery would rather not see + mid-capture (a launch that fails or behaves unusually the first + time either fails the capture outright or bakes in a broken graph + node). + 2. Restores every data buffer this routine reaches (add_data's + handles) to the values it captured a copy of before warming up + - a warmup launch is a real launch and actually computes into + those buffers, and compile() otherwise would not be the + side-effect-free operation every other compile() in this + package is. + 3. Captures the same step sequence again, on a dedicated + non-blocking stream, via cp.cuda.Stream.begin_capture() / + end_capture(). Nothing on that stream executes during capture - + only the graph is built - so this pass leaves the (already + restored) buffers untouched. Each step's constant block (see the + module docstring) was already uploaded synchronously by its own + compile() call, above, well before this point - a synchronous + copy issued while a stream is being captured is illegal, so that + upload could not happen here even if it needed to. + 4. Checks the default cupy memory pool's used_bytes() is the same + before and after capture and raises if not. Every data handle a + routine launches with is already allocated by the time compile() + runs (add_data takes an existing handle), so nothing captured + here should need the pool; growth here means some step + allocated during capture, which CUDA graph capture does not + support - this is caught rather than silently producing an + unusable graph. + + The returned _CapturedRoutine keeps the dedicated stream and the + cp.cuda.Graph alive; __call__ replays the graph with no arguments - + overriding data handles at call time is rejected, see + _CapturedRoutine.__call__. + + Author: B.G (07/2026) + """ + if not captured: + return super().compile(fused=False) + + self._validate() + + compiled_steps: list[_CompiledStep] = [] + data_names: list[str] = [] + for step in self._steps: + compiled = step.kernel_builder.compile() + caller = self._make_caller(compiled, step.grid, step.block) + compiled_steps.append(_CompiledStep(caller, step.canonical_refs)) + for name in step.canonical_refs: + if name not in data_names: + data_names.append(name) + + defaults = {name: self._data[name] for name in data_names} + + def _launch_all(): + for step in compiled_steps: + step.caller(*(defaults[name] for name in step.canonical_refs)) + + # 1. warm up every kernel with one real launch, so any remaining + # first-launch cost happens before capture starts, not during it - + # the module itself is already built and its constant block already + # uploaded, by compile() just above. + snapshots = {name: buf.copy() for name, buf in defaults.items()} + _launch_all() + cp.cuda.Device().synchronize() + + # 2. undo the warmup launch's real effect - compile() must not leave + # the caller's buffers different from how it found them. + for name, buf in defaults.items(): + buf[...] = snapshots[name] + cp.cuda.Device().synchronize() + + # 3. capture the same sequence on a dedicated stream. + mempool = cp.get_default_memory_pool() + used_before = mempool.used_bytes() + stream = cp.cuda.Stream(non_blocking=True) + with stream: + stream.begin_capture() + _launch_all() + graph = stream.end_capture() + + # 4. no allocation should have happened while capturing. + used_after = mempool.used_bytes() + if used_after != used_before: + raise RuntimeError( + "CupyRoutineBuilder.compile(captured=True): the default memory pool's " + f"used_bytes() changed during capture ({used_before} -> {used_after}); a " + "step allocated instead of using an already-initialised data handle, which " + "CUDA graph capture does not support" + ) + + return _CapturedRoutine(compiled_steps, tuple(data_names), defaults, graph, stream) diff --git a/pyfastflow/experimental/core/context/parameter.py b/pyfastflow/experimental/core/context/parameter.py new file mode 100644 index 0000000..d1328c9 --- /dev/null +++ b/pyfastflow/experimental/core/context/parameter.py @@ -0,0 +1,276 @@ +""" +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). +HelperBuilder The recipe for a device-side helper: a small routine callable + only from other device code (ti.func, qd.func, CUDA + __device__). Bind it into a kernel - flat or inside a Bag - + and the kernel's own compile() specializes it; there is no + standalone compiled Helper object to hold onto. +Kernel A compiled entry point - what the host launches (ti.kernel, + qd.kernel, CUDA __global__). +Bag A named collection of any of the above, mixed freely, so a + group travels as one object and is reached in-kernel by dotted + path: phys.dx.get(i), ops.neighbour(i). + +A "context" is any concrete class - GridContext, FlowContext, ... - that groups +Parameters and registers Helpers. There is deliberately no base Context +class: a context needing another context's parameters binds them explicitly, +rather than reaching through a registry of stored connections. + +Compiling something +------------------- +Templates are written once, generically, and specialized by a builder: + + kernel = (TaichiKernelBuilder() + .bind("phys", phys) # a Bag of parameters + .bind("ops", ops) # a Bag of HelperBuilders + .ingest(update_height) # the template + .compile()) + kernel(h_new, h_old) # bulk data passed at call time + +bind(name, obj) makes `obj` visible inside the template body under `name`. +ingest() takes the template - a python def for Taichi/Quadrants, a CUDA source +string for cupy. compile() returns a Kernel. Only the abstract HelperBuilder / +KernelBuilder live here; the concrete Taichi*, Quadrants* and Cupy* builders +sit alongside this module. + +A HelperBuilder bound anywhere in a KernelBuilder's bindings - directly under +a name, or as a member of a bound Bag - is specialized as part of that +kernel's compile(), against that same compile's bindings. This is what lets a +helper reading a const Parameter pick up a different value after the const is +swapped and the *kernel* is recompiled, with the helper's own builder never +touched. Reaching the same HelperBuilder from two places in one kernel - bound +flat and inside a Bag, or under two different names - specializes it once; the +same specialized object is shared at both call sites. A HelperBuilder has no +compiled form of its own to keep between compiles: it is a recipe, always +specialized fresh as part of whatever kernel currently binds it. + +The builder is the recipe: its template and bindings can be inspected, and +compile() may be called again after a bind() edit, each call producing a new, +independent callable. Nothing about compile() consumes or mutates the +builder - recompiling a builder that has not changed since its last compile() +just repeats work for an equivalent result, which is pointless and best +avoided, though harmless if it happens. + +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: + + compile.py Specializable/Kernel, the abstract HelperBuilder and + KernelBuilder, and resolve_binding - everything involved in + turning a template plus bindings into a compiled object. + bag.py Bag and its operators (merge, extract, trim, replace, ...), + which know nothing of 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") + + @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..13f9eac --- /dev/null +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -0,0 +1,81 @@ +""" +Quadrants implementations of Parameter and the two builders. + +Templates are python defs, specialized by splicing bound objects into their +globals before handing them to qd.func or qd.kernel - the mechanism in +_closure_backend.py, shared with Taichi. + +Two things differ from Taichi. Kernel templates may type their data arguments +qd.Tensor, and one such template then accepts 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, + ClosureHelperBuilder, + ClosureKernelBuilder, + ClosureRoutineBuilder, +) + + +class QuadrantsParameter(ClosureBackendParameter): + """ + Parameter backed by a Quadrants const value or a pooled QuadrantsDataHandle. + + Author: B.G (07/2026) + """ + + _backend = qd + + +class QuadrantsHelperBuilder(ClosureHelperBuilder): + """ + Compiles a device helper with qd.func. Parameters bound into it must be + field-backed, since Quadrants rejects an ndarray referenced as a global. + + Author: B.G (07/2026) + """ + + _backend = qd + + +class QuadrantsKernelBuilder(ClosureKernelBuilder): + """ + Compiles a launchable kernel with qd.kernel. Type the template's data + arguments qd.Tensor to accept field- or ndarray-backed values alike. + + Author: B.G (07/2026) + """ + + _backend = qd + + +class QuadrantsRoutineBuilder(ClosureRoutineBuilder): + """ + Compiles an ordered sequence of Quadrants kernels sharing one bag into a + Routine. + + Author: B.G (07/2026) + """ diff --git a/pyfastflow/experimental/core/context/routine.py b/pyfastflow/experimental/core/context/routine.py new file mode 100644 index 0000000..3e86471 --- /dev/null +++ b/pyfastflow/experimental/core/context/routine.py @@ -0,0 +1,476 @@ +""" +A Routine is an ordered, device-only, linear sequence of kernels compiled and +launched as one unit, sharing a single bag of Parameters and Helpers across +every step. + +What this is for +----------------- +A single kernel is one pass over the grid. Most models need several passes +run back to back, over the same underlying parameters, every substep - heat +diffusion needs a diffuse pass then a source pass, repeated. Wiring that by +hand means keeping several already-compiled Kernels around plus a python loop +that launches them in order and juggles which buffer currently holds what. +A Routine collects that sequence once, as a builder, and compiles it into one +callable that launches every step in order. + +A Routine is built from KernelBuilders that are otherwise ordinary - built +exactly as they would be for a standalone compile, bind()ed against whatever +Parameters and Helpers the step needs. What is different is that a Routine's +steps do not each keep their own bindings forever: at compile() time, every +step is rebound (see compile.py, CompileBuilder.rebind) against the one bag the +whole Routine shares, so a Parameter reached under a given name means the same +object in every step that reaches it. check_handles (bag.py) is run across +every step's own bindings, as authored, before that rebind happens - so a +step built against one object under a name another step expects to mean +something else is caught at compile time, even though rebind would otherwise +silently make both agree with whatever the routine's bag holds. + +Bulk data - the buffers each step reads and writes - is handled separately +from the bag. add_data(name, handle) gives a buffer a routine-local name; +add_kernel's data_handle_ref maps positional template arguments onto those +names. add_swap(a, b) relabels which buffer a name currently means, for every +step added after it - it is bookkeeping only, resolved once at compile time, +and costs nothing at launch. Because a Routine is called and re-called without +returning to Python between steps, there is no place to do the python-level +`T0, T1 = T1, T0` a hand-written ping-pong loop uses; add_swap is what stands +in for it, and the compiled Routine keeps re-running correctly precisely +because the net effect of every swap in it is required to be the identity - +see RoutineBuilder.compile. + +Executing a routine +-------------------- + routine = (TaichiRoutineBuilder() + .add_data("T0", T0.data) + .add_data("T1", T1.data) + .bind_bag(shared_bag) + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .add_kernel(diffuse_builder, data_handle_ref=("T1", "T0")) + .add_kernel(source_builder, data_handle_ref=("T1",)) + .add_swap("T0", "T1") + .compile()) + routine() # uses the add_data() defaults + routine(T0.data, T1.data) # or override them positionally, per data_names + +compile() rebinds and compiles each step, in the order added, into a Routine: +an inert sequence of already-compiled kernels plus, per step, which of the +routine's data names it launches with. Calling a Routine simply launches its +steps in that order. This base implementation keeps every step a fully +separate kernel launch; both backend subclasses build on it rather than +replacing it outright. The closure-backend subclass (Taichi/Quadrants) +overrides compile() to splice consecutive steps' loop bodies into one +generated kernel by default; split() marks where that splicing should stop +and a new generated kernel should start. CupyRoutineBuilder, having no +source fusion available, instead defaults to capturing the compiled steps' +launches into a CUDA graph and replaying that graph on call - see +cupy_backend.py, CupyRoutineBuilder.compile and _CapturedRoutine. See +RoutineBuilder.compile, ClosureRoutineBuilder.compile and +CupyRoutineBuilder.compile for what a repeated call composes to and why the +swaps must balance either way. + +Contract: no set()/destroy() mid-routine +------------------------------------------ +A compiled Routine holds the same kind of frozen snapshot a Kernel does (see +parameter.py, "Lifetime of a compiled object"): scalar/field Parameters are baked +in by their storage, const Parameters by their literal value. Calling set() or +destroy() on any Parameter the routine's bag reaches, between two calls to the +routine (or between two of its steps, if that were possible), is undefined: +a scalar/field write is visible immediately since the routine holds the same +storage, a const write is not since the literal was already baked into every +step at compile time, and destroy() invalidates the storage a step still +points at. None of this is enforced at runtime - recompile the routine after +such a change, the same way a Kernel is recompiled. + +Contract: a captured graph bakes in pointers, not just storage +----------------------------------------------------------------- +A CupyRoutineBuilder-compiled Routine (captured=True, its default - see +cupy_backend.py) has everything above still true, plus one more thing a +CUDA graph adds on top of what a plain kernel launch already has: every +step's launch arguments, pointers included, are baked into the graph at +capture time, not re-read at replay time. + +- A write to a scalar or field Parameter's storage - the ordinary set() - + still lands where the graph's launches already point, so it is seen on the + very next replay with no recompile needed; this is the intended way to + feed a captured routine changing data, exactly as for an uncaptured one. +- set() on a const Parameter still only changes generated source the graph + never re-reads, so it goes stale exactly as an uncaptured Routine's kernels + would - recompile. +- destroy(), or anything else that returns a data handle's buffer to the + pool, invalidates a pointer the graph's launches were captured with. + Recompile - there is no cheaper fix, since the graph does not know which + of its baked-in pointers came from which handle. +None of this is enforced at runtime; see cupy_backend.py, _CapturedRoutine +for the exact rule set and why. + +Author: B.G (07/2026) +""" + +import re +from abc import ABC, abstractmethod +from typing import Any + +from .bag import Bag, check_handles +from .compile import CompileBuilder + +_C_FUNC_NAME_RE = re.compile(r"(?:__global__|__device__)\s+[\w:\*&]*\s*(\w+)\s*\(") + + +def _template_label(template) -> str: + """ + A short, human-readable name for a template, for error messages: a + python def's own __name__, or the __global__/__device__ entry point read + out of CUDA source text. Falls back to a truncated repr if neither + applies. + + Author: B.G (07/2026) + """ + name = getattr(template, "__name__", None) + if name is not None: + return name + if isinstance(template, str): + match = _C_FUNC_NAME_RE.search(template) + if match: + return match.group(1) + return repr(template)[:60] + + +def _flatten_bindings(bindings: dict[str, Any]) -> dict[str, Any]: + """ + A step's bound names, expanded with a dotted entry for every member of + any bound Bag - the same shape Bag.walk() produces for a single bag, + covering a step's whole binding set instead of one bag's contents. + + Author: B.G (07/2026) + """ + flat: dict[str, Any] = {} + for name, obj in bindings.items(): + flat[name] = obj + if isinstance(obj, Bag): + flat.update(dict(obj.walk(name))) + return flat + + +class _Step: + """ + One entry recorded by add_kernel: the kernel builder as given, and the + data names it launches with, resolved through the swap table as it stood + at the moment this step was added - see RoutineBuilder.add_kernel. + + Author: B.G (07/2026) + """ + + __slots__ = ("kernel_builder", "canonical_refs", "grid", "block") + + def __init__(self, kernel_builder, canonical_refs: tuple, grid, block): + self.kernel_builder = kernel_builder + self.canonical_refs = canonical_refs + self.grid = grid + self.block = block + + +class _CompiledStep: + """ + One step of a compiled Routine: a callable that launches the already + compiled kernel (backend launch convention baked in by + RoutineBuilder._make_caller), plus the data names it is called with. + + Author: B.G (07/2026) + """ + + __slots__ = ("caller", "canonical_refs") + + def __init__(self, caller, canonical_refs: tuple): + self.caller = caller + self.canonical_refs = canonical_refs + + +class Routine: + """ + A compiled, ordered sequence of kernels sharing one bag, ready to launch. + + Inert, like every other compiled object in this package (see compile.py, + Specializable): it retains the steps it was built from and nothing reads + back from it to drive a later compile. Go through the RoutineBuilder that + produced it to change anything and compile again. + + `data_names` lists the distinct names add_kernel's steps referenced, in + the order they first appear across the steps as they were added. Calling + the routine with no arguments launches every step using the handles given + to add_data(); calling it with `len(data_names)` positional arguments + overrides all of them for that call, matched up by position against + `data_names`. A repeated call, with or without override arguments, is + exactly what add_swap's net-identity requirement (see + RoutineBuilder.compile) makes safe: nothing about a Routine's own state + changes between calls, so calling it twice launches its steps against the + same starting arrangement of buffers twice. + + Author: B.G (07/2026) + """ + + def __init__(self, steps: list[_CompiledStep], data_names: tuple, defaults: dict[str, Any]): + self._steps = steps + self._data_names = data_names + self._defaults = defaults + + @property + def data_names(self) -> tuple: + """ + The distinct data names this routine's steps reference, in + first-appearance order across the steps as they were added. + + Author: B.G (07/2026) + """ + return self._data_names + + def __call__(self, *args) -> None: + """ + Launch every step in order. With no arguments, each data name + resolves to the handle given to add_data(); with + `len(data_names)` positional arguments, those override the defaults + for this call, matched by position against `data_names`. Any other + argument count raises. + + Author: B.G (07/2026) + """ + if args and len(args) != len(self._data_names): + raise ValueError( + f"Routine: expected 0 or {len(self._data_names)} argument(s) " + f"matching data_names={self._data_names}, got {len(args)}" + ) + table = dict(self._defaults) + if args: + table.update(zip(self._data_names, args)) + for step in self._steps: + step.caller(*(table[name] for name in step.canonical_refs)) + + +class RoutineBuilder(ABC): + """ + Collects data names, a shared bag, and an ordered list of kernel steps, + and compiles them into a Routine. + + add_data(name, handle) registers a routine-local name for a pooled data + handle. add_kernel(kernel_builder, data_handle_ref=(...)) appends a step: + `data_handle_ref` is a tuple of names, previously registered with + add_data, mapped positionally onto the template's own declared data + arguments. add_swap(a, b) is a build-time relabeling of which handle `a` + and `b` currently mean - it emits no code, costs nothing at launch, and + only affects steps added after it; a step added earlier keeps whatever + the table resolved to at the time it was added. split() records a fusion + boundary and otherwise does nothing; see its own docstring. bind_bag(bag) + sets the one bag every step is rebound against at compile time. + + Author: B.G (07/2026) + """ + + def __init__(self): + self._data: dict[str, Any] = {} + self._perm: dict[str, str] = {} + self._steps: list[_Step] = [] + self._splits: set[int] = set() + self._bag: "Bag | None" = None + + def add_data(self, name: str, handle: Any) -> "RoutineBuilder": + """ + Register `handle` under routine-local `name`, and the default it + resolves to unless a call to the compiled Routine overrides it. + + Author: B.G (07/2026) + """ + if name in self._data: + raise KeyError(f"add_data: '{name}' is already registered") + self._data[name] = handle + self._perm[name] = name + return self + + def add_kernel( + self, + kernel_builder: CompileBuilder, + data_handle_ref: tuple = (), + *, + grid=None, + block=None, + ) -> "RoutineBuilder": + """ + Append a step launching `kernel_builder`'s compiled kernel with the + data handles named in `data_handle_ref`, mapped positionally onto + the template's own declared data arguments. + + Each name in `data_handle_ref` must already be registered via + add_data(); which underlying handle it currently means is resolved + against the swap table as it stands right now; a later add_swap does + not reach back and change this step. `grid`/`block` are accepted on + every backend but only meaningful on cupy - see CupyRoutineBuilder. + + Author: B.G (07/2026) + """ + data_handle_ref = tuple(data_handle_ref) + arity = self._data_arity(kernel_builder) + if arity != len(data_handle_ref): + name = _template_label(kernel_builder.template) + raise ValueError( + f"add_kernel: template {name!r} declares {arity} data argument(s), " + f"data_handle_ref gives {len(data_handle_ref)}" + ) + unknown = [n for n in data_handle_ref if n not in self._perm] + if unknown: + raise KeyError(f"add_kernel: not registered via add_data: {sorted(unknown)}") + canonical = tuple(self._perm[n] for n in data_handle_ref) + self._steps.append(_Step(kernel_builder, canonical, grid, block)) + return self + + def add_swap(self, a: str, b: str) -> "RoutineBuilder": + """ + Relabel the table so `a` and `b` swap which handle they currently + mean. Emits no code and has no runtime cost - it only changes what a + later add_kernel resolves its data_handle_ref against; steps added + before this call are unaffected. See compile() for the requirement + that every swap in a routine nets out to the identity. + + Author: B.G (07/2026) + """ + if a not in self._perm or b not in self._perm: + missing = sorted(n for n in (a, b) if n not in self._perm) + raise KeyError(f"add_swap: not registered via add_data: {missing}") + self._perm[a], self._perm[b] = self._perm[b], self._perm[a] + return self + + def split(self) -> "RoutineBuilder": + """ + Mark the boundary between this step and the next as a fusion + boundary: a fused compile() launches everything before this call and + everything after it as separate generated kernels, back to back, + rather than splicing them into one. It has no effect on an unfused + compile() - every step there is already its own kernel, in the order + added, whether or not split() was called between them. + + Author: B.G (07/2026) + """ + self._splits.add(len(self._steps)) + return self + + def bind_bag(self, bag: "Bag") -> "RoutineBuilder": + """ + Set the one bag every step is rebound against at compile time. + + Author: B.G (07/2026) + """ + self._bag = bag + return self + + @abstractmethod + def _data_arity(self, kernel_builder: CompileBuilder) -> int: + """ + The number of data arguments `kernel_builder`'s ingested template + declares, for add_kernel's arity check. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def _make_caller(self, compiled_kernel, grid, block): + """ + A callable(*data_args) that launches `compiled_kernel` the way this + backend requires - straight through for Taichi/Quadrants, with + grid/block supplied for cupy. + + Author: B.G (07/2026) + """ + ... + + def _validate(self) -> None: + """ + Everything a compile needs checked before any step is actually + compiled, fused or not. + + In order: check_handles (bag.py) runs across every step's own + bindings, as authored - so two steps disagreeing about what one + handle means is caught here, before rebind would otherwise silently + make both agree with the routine's bag. Each step is then rebound + against the bag set by bind_bag(); a step whose current bindings the + bag cannot satisfy raises naming that step and everything missing. + The swap table's net permutation is then checked: every add_swap in + this routine must compose to the identity, since a Routine is called + and re-called without ever returning to Python to swap buffers + itself - an unbalanced set of swaps would compute into the wrong + buffer on the second call. + + Author: B.G (07/2026) + """ + if self._bag is None: + raise ValueError("compile: no bag bound - call bind_bag() first") + if not self._steps: + raise ValueError("compile: routine has no steps") + + units = {} + for i, step in enumerate(self._steps): + label = f"step{i}:{_template_label(step.kernel_builder.template)}" + units[label] = _flatten_bindings(step.kernel_builder.bindings) + check_handles(units) + + for i, step in enumerate(self._steps): + try: + step.kernel_builder.rebind(self._bag) + except KeyError as exc: + label = _template_label(step.kernel_builder.template) + raise KeyError(f"compile: step {i} ({label!r}) cannot be satisfied by the routine's bag: {exc}") from exc + + drift = {name: target for name, target in self._perm.items() if target != name} + if drift: + raise ValueError(f"compile: net swap permutation is not the identity, still swapped: {drift}") + + def _grouped_steps(self) -> list[list[_Step]]: + """ + `self._steps` partitioned at every split() boundary, in order. With + no split() calls this is a single group holding every step. + + Author: B.G (07/2026) + """ + groups: list[list[_Step]] = [] + current: list[_Step] = [] + for i, step in enumerate(self._steps): + if i in self._splits and current: + groups.append(current) + current = [] + current.append(step) + if current: + groups.append(current) + return groups + + def compile(self, fused: bool = False, dump_source: str | None = None) -> Routine: + """ + Validate (see _validate) and compile every step into a Routine. + + With fused=False (the default here - see the closure-backend + subclass for a backend where fusion is actually available and + defaults on), each step compiles to its own kernel and the Routine + launches them in order, exactly one host-side call per step. + `dump_source` is accepted for signature parity with a fusing + subclass and ignored, since there is no generated source to dump. + fused=True raises here: this base implementation backs cupy, which + has no source fusion. + + Author: B.G (07/2026) + """ + if fused: + raise NotImplementedError( + f"{type(self).__name__}: source fusion is not supported on this backend; " + "call compile(fused=False) (the default) instead" + ) + self._validate() + + compiled_steps: list[_CompiledStep] = [] + data_names: list[str] = [] + for step in self._steps: + compiled = step.kernel_builder.compile() + caller = self._make_caller(compiled, step.grid, step.block) + compiled_steps.append(_CompiledStep(caller, step.canonical_refs)) + for name in step.canonical_refs: + if name not in data_names: + data_names.append(name) + + defaults = {name: self._data[name] for name in data_names} + return Routine(compiled_steps, tuple(data_names), defaults) diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py new file mode 100644 index 0000000..ce2aa10 --- /dev/null +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -0,0 +1,61 @@ +""" +Taichi implementations of Parameter and the two builders. + +Templates are python defs. A builder specializes one by rebuilding it with the +bound objects spliced into its globals, then hands the result to ti.func or +ti.kernel; _closure_backend.py holds that machinery, shared with Quadrants. +Everything below just names Taichi as the backend to use. + +Kernel templates declare their data arguments the usual Taichi way, typically +ti.template(). Bound Parameters and helpers are not arguments - see parameter.py. + +Author: B.G (07/2026) +""" + +import taichi as ti + +from ._closure_backend import ( + ClosureBackendParameter, + ClosureHelperBuilder, + ClosureKernelBuilder, + ClosureRoutineBuilder, +) + + +class TaichiParameter(ClosureBackendParameter): + """ + Parameter backed by a Taichi const value or a pooled TaichiDataHandle. + + Author: B.G (07/2026) + """ + + _backend = ti + + +class TaichiHelperBuilder(ClosureHelperBuilder): + """ + Compiles a device helper with ti.func. + + Author: B.G (07/2026) + """ + + _backend = ti + + +class TaichiKernelBuilder(ClosureKernelBuilder): + """ + Compiles a launchable kernel with ti.kernel. + + Author: B.G (07/2026) + """ + + _backend = ti + + +class TaichiRoutineBuilder(ClosureRoutineBuilder): + """ + Compiles an ordered sequence of Taichi kernels sharing one bag into a + Routine. + + Author: B.G (07/2026) + """ 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/grid/__init__.py b/pyfastflow/experimental/grid/__init__.py new file mode 100644 index 0000000..fa3d08e --- /dev/null +++ b/pyfastflow/experimental/grid/__init__.py @@ -0,0 +1,167 @@ +""" +make_grid: the GridContext-equivalent Bag factory, built on the +backend-agnostic core (see ..core.context: parameter.py for Parameter, compile.py for HelperBuilder, bag.py for Bag). + +There is no stateful Context class here - by design, the core has none (see +core/context/parameter.py's module docstring). make_grid just builds a Bag once: a +uniform public surface (grid.nx, grid.neighbour(i, k), ...) whatever the +backend and whatever the grid's own topology/boundary/nodata/outlet config. + +Two kinds of knobs: + - value params (nx, ny, dx) - mode-overridable (const/scalar, dx also + field), always read in device code through `.get(...)`. Default mode is + "const" for all three. + - structural selectors (topology, boundary, nodata, outlet) - each one + picks which variant of a private block gets bound into the public + composite helpers; see _closure_blocks.py / _cupy_blocks.py. + +Masks are independent, optional bag members: nodata_mask (u8, 1 == inactive) +when nodata=True, outlet_mask (u8, 1 == outlet) when outlet=="mask". Neither +exists in the bag when its feature is off, so a caller that never asked for +nodata/mask-outlet never sees them. + +Author: B.G (07/2026) +""" + +import numpy as np + +from ..core.context.backends import backend_classes +from ..core.context.bag import Bag + +_TOPOLOGIES = {"D4": 4, "D8": 8} +_BOUNDARIES = frozenset({"normal", "periodic_EW", "periodic_NS"}) +_OUTLETS = frozenset({"edge", "mask"}) + + +def _blocks_for(backend: str): + """ + The private block module implementing make_grid's device code for one + backend name: the closure blocks (shared by Taichi and Quadrants) or the + cupy blocks. + + Author: B.G (07/2026) + """ + if backend in ("taichi", "quadrants"): + from . import _closure_blocks as blocks + elif backend == "cupy": + from . import _cupy_blocks as blocks + else: + raise ValueError(f"make_grid: unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + return blocks + + +def make_grid( + backend: str, + pool, + nx: int, + ny: int, + dx: float, + *, + topology: str = "D8", + boundary: str = "normal", + nodata: bool = False, + outlet: str = "edge", + nx_mode: str = "const", + ny_mode: str = "const", + dx_mode: str = "const", +) -> Bag: + """ + Build one grid's Bag: nx/ny/dx/n_neighbours params, the optional + nodata_mask/outlet_mask fields, and the neighbour/distance/edge helper + surface - all uniform by name regardless of backend or config. + + `topology` "D4"|"D8", `boundary` "normal"|"periodic_EW"|"periodic_NS", + `outlet` "edge"|"mask" pick block variants at build time (see + _closure_blocks.py / _cupy_blocks.py). `nodata` allocates and folds in + nodata_mask (u8, 1 == inactive) wherever a block needs it. + + `nx_mode`/`ny_mode` default "const", may be overridden to "scalar". + `dx_mode` defaults "const", may be overridden to "scalar" or "field" - a + field-mode dx is allocated (one cell per node, caller fills it) but the + public helpers that read dx (dist_from_k, dist_between_nodes) only ever + read index 0: neither's signature carries a node to key a per-node value + off, so a genuinely spatially-varying dx is not wired through those two + helpers as things stand - only reachable by reading grid.dx.get(i) + directly in a caller's own template. + + Author: B.G (07/2026) + """ + if topology not in _TOPOLOGIES: + raise ValueError(f"make_grid: topology must be one of {sorted(_TOPOLOGIES)}, got {topology!r}") + if boundary not in _BOUNDARIES: + raise ValueError(f"make_grid: boundary must be one of {sorted(_BOUNDARIES)}, got {boundary!r}") + if outlet not in _OUTLETS: + raise ValueError(f"make_grid: outlet must be one of {sorted(_OUTLETS)}, got {outlet!r}") + if nx_mode not in ("const", "scalar"): + raise ValueError(f"make_grid: nx_mode must be 'const' or 'scalar', got {nx_mode!r}") + if ny_mode not in ("const", "scalar"): + raise ValueError(f"make_grid: ny_mode must be 'const' or 'scalar', got {ny_mode!r}") + if dx_mode not in ("const", "scalar", "field"): + raise ValueError(f"make_grid: dx_mode must be 'const', 'scalar' or 'field', got {dx_mode!r}") + + backend_mod, ParamCls, HelperCls, dtypes = backend_classes(backend) + blocks = _blocks_for(backend) + n_flat = int(nx) * int(ny) + + nx_p = ParamCls("GRID_NX", dtype=dtypes["i32"], mode=nx_mode, value=int(nx), pool=pool) + ny_p = ParamCls("GRID_NY", dtype=dtypes["i32"], mode=ny_mode, value=int(ny), pool=pool) + + if dx_mode == "field": + dx_p = ParamCls( + "GRID_DX", + dtype=dtypes["f32"], + mode="field", + value=np.full(n_flat, dx, dtype=np.float32), + pool=pool, + n_flat=n_flat, + ) + else: + dx_p = ParamCls("GRID_DX", dtype=dtypes["f32"], mode=dx_mode, value=float(dx), pool=pool) + + n_neighbours_p = ParamCls( + "GRID_NNEIGHBOURS", dtype=dtypes["i32"], mode="const", value=_TOPOLOGIES[topology], pool=pool + ) + + nodata_mask_p = None + if nodata: + nodata_mask_p = ParamCls( + "GRID_NODATA_MASK", + dtype=dtypes["u8"], + mode="field", + value=np.zeros(n_flat, dtype=np.uint8), + pool=pool, + n_flat=n_flat, + ) + + outlet_mask_p = None + if outlet == "mask": + outlet_mask_p = ParamCls( + "GRID_OUTLET_MASK", + dtype=dtypes["u8"], + mode="field", + value=np.zeros(n_flat, dtype=np.uint8), + pool=pool, + n_flat=n_flat, + ) + + helpers = blocks.build_helpers( + HelperCls, + nx_p=nx_p, + ny_p=ny_p, + dx_p=dx_p, + nodata_mask_p=nodata_mask_p, + outlet_mask_p=outlet_mask_p, + topology=topology, + boundary=boundary, + nodata=nodata, + outlet=outlet, + backend_mod=backend_mod, + ) + + items = {"nx": nx_p, "ny": ny_p, "dx": dx_p, "n_neighbours": n_neighbours_p} + if nodata_mask_p is not None: + items["nodata_mask"] = nodata_mask_p + if outlet_mask_p is not None: + items["outlet_mask"] = outlet_mask_p + items.update(helpers) + return Bag(items) diff --git a/pyfastflow/experimental/grid/_closure_blocks.py b/pyfastflow/experimental/grid/_closure_blocks.py new file mode 100644 index 0000000..a4ba189 --- /dev/null +++ b/pyfastflow/experimental/grid/_closure_blocks.py @@ -0,0 +1,540 @@ +""" +Taichi/Quadrants (closure) block templates behind make_grid. + +Every private block below is one plain python def, PICKED - never branched on +inside a single function body - by build_helpers() according to the grid's +config: topology (D4/D8), boundary (normal/periodic_EW/periodic_NS), nodata +(on/off), outlet (edge/mask). Composability is per axis: a periodic boundary +swaps in the "periodic" variant of a row or column block, never both, and the +untouched axis keeps its "identity"/"bounded" variant - there is no +ti.static/#if choosing between them inside one function. The one runtime +if-ladder is _delta(k): k is genuine per-call device data, not a structural +choice, so it cannot be resolved by picking a python function ahead of time. +A dynamically-indexed local array for that ladder would spill to local memory +on GPU, hence the explicit if/elif chain instead. + +Every public helper below is a HelperBuilder that binds the private blocks it +needs BY NAME - helper binds helper - so a block reached from two composites +(e.g. _row reached from both neighbour_raw and dist_between_nodes) is +specialized once per compile and shared at both call sites (see compile.py, +_SpecializeCtx). + +nx/ny/dx are read exclusively through `.get(...)`, uniformly across whatever +mode they end up in (const, scalar, field) - see parameter.py, "Reading a +Parameter in device code is uniform across modes." This is what lets any of +them be overridden to a runtime-modifiable mode without touching a single +block template. + +Author: B.G (07/2026) +""" + +import functools + +from ..core.context.backends import make_helper + +# --------------------------------------------------------------------------- +# geometry +# --------------------------------------------------------------------------- + + +def _row_tmpl(i): + return i // NX.get(0) + + +def _col_tmpl(i): + return i % NX.get(0) + + +def _index_tmpl(row, col): + return row * NX.get(0) + col + + +def _in_bounds_tmpl(row, col): + ok = 0 + if row >= 0 and row < NY.get(0) and col >= 0 and col < NX.get(0): + ok = 1 + return ok + + +# --------------------------------------------------------------------------- +# topology: _delta(k) - runtime if-ladder, k is per-call device data +# --------------------------------------------------------------------------- + + +def _delta_d4_tmpl(k): + dr = 0 + dc = 0 + if k == 0: + dr, dc = -1, 0 + elif k == 1: + dr, dc = 0, -1 + elif k == 2: + dr, dc = 0, 1 + else: + dr, dc = 1, 0 + return dr, dc + + +def _delta_d8_tmpl(k): + dr = 0 + dc = 0 + if k == 0: + dr, dc = -1, -1 + elif k == 1: + dr, dc = -1, 0 + elif k == 2: + dr, dc = -1, 1 + elif k == 3: + dr, dc = 0, -1 + elif k == 4: + dr, dc = 0, 1 + elif k == 5: + dr, dc = 1, -1 + elif k == 6: + dr, dc = 1, 0 + else: + dr, dc = 1, 1 + return dr, dc + + +# --------------------------------------------------------------------------- +# per-axis wrap (boundary): identity or periodic, chosen per axis +# --------------------------------------------------------------------------- + + +def _row_wrap_identity_tmpl(row): + return row + + +def _row_wrap_periodic_tmpl(row): + r = row + if r < 0: + r += NY.get(0) + elif r >= NY.get(0): + r -= NY.get(0) + return r + + +def _col_wrap_identity_tmpl(col): + return col + + +def _col_wrap_periodic_tmpl(col): + c = col + if c < 0: + c += NX.get(0) + elif c >= NX.get(0): + c -= NX.get(0) + return c + + +def _wrap_tmpl(row, col): + return _ROWWRAP(row), _COLWRAP(col) + + +# --------------------------------------------------------------------------- +# per-axis edge gate (boundary): is a candidate coordinate still in range on +# that axis - a periodic axis never blocks, since it wraps instead. +# --------------------------------------------------------------------------- + + +def _row_edge_ok_bounded_tmpl(row): + ok = 0 + if row >= 0 and row < NY.get(0): + ok = 1 + return ok + + +def _row_edge_ok_periodic_tmpl(row): + return 1 + + +def _col_edge_ok_bounded_tmpl(col): + ok = 0 + if col >= 0 and col < NX.get(0): + ok = 1 + return ok + + +def _col_edge_ok_periodic_tmpl(col): + return 1 + + +# --------------------------------------------------------------------------- +# source-cell nodata gate: on/off +# --------------------------------------------------------------------------- + + +def _source_ok_always_tmpl(i): + return 1 + + +def _source_ok_nodata_tmpl(i): + ok = 1 + if NODATA_MASK.get(i) == 1: + ok = 0 + return ok + + +# --------------------------------------------------------------------------- +# _move_allowed(i, k): per-axis edge gate x source-cell nodata gate +# --------------------------------------------------------------------------- + + +def _move_allowed_tmpl(i, k): + row = _ROW(i) + col = _COL(i) + dr, dc = _DELTA(k) + return _ROWEDGEOK(row + dr) * _COLEDGEOK(col + dc) * _SOURCEOK(i) + + +# --------------------------------------------------------------------------- +# _valid(j): in range, and (nodata ? active : true) +# --------------------------------------------------------------------------- + + +def _valid_no_nodata_tmpl(j): + ok = 0 + if j >= 0 and j < NX.get(0) * NY.get(0): + ok = 1 + return ok + + +def _valid_nodata_tmpl(j): + ok = 0 + if j >= 0 and j < NX.get(0) * NY.get(0): + ok = 1 + if NODATA_MASK.get(j) == 1: + ok = 0 + return ok + + +# --------------------------------------------------------------------------- +# public: neighbour / neighbour_raw +# --------------------------------------------------------------------------- + + +def _neighbour_raw_tmpl(i, k): + row = _ROW(i) + col = _COL(i) + dr, dc = _DELTA(k) + wrow, wcol = _WRAP(row + dr, col + dc) + return _INDEX(wrow, wcol) + + +def _neighbour_tmpl(i, k): + j = -1 + if _MOVEALLOWED(i, k) == 1: + cand = _NEIGHBOURRAW(i, k) + if _VALID(cand) == 1: + j = cand + return j + + +# --------------------------------------------------------------------------- +# public: is_active / nodata +# --------------------------------------------------------------------------- + + +def _is_active_always_tmpl(i): + return 1 + + +def _is_active_mask_tmpl(i): + ok = 1 + if NODATA_MASK.get(i) == 1: + ok = 0 + return ok + + +def _nodata_tmpl(i): + return 1 - _ISACTIVE(i) + + +# --------------------------------------------------------------------------- +# public: is_on_edge / which_edge (per-axis, mirrors _wrap/_edge_ok) +# --------------------------------------------------------------------------- + + +def _row_is_edge_active_tmpl(row): + e = 0 + if row == 0 or row == NY.get(0) - 1: + e = 1 + return e + + +def _row_is_edge_periodic_tmpl(row): + return 0 + + +def _col_is_edge_active_tmpl(col): + e = 0 + if col == 0 or col == NX.get(0) - 1: + e = 1 + return e + + +def _col_is_edge_periodic_tmpl(col): + return 0 + + +def _is_on_edge_tmpl(i): + row = _ROW(i) + col = _COL(i) + e = 0 + if _ROWISEDGE(row) == 1 or _COLISEDGE(col) == 1: + e = 1 + return e + + +def _row_edge_code_active_tmpl(row): + code = -1 + if row == 0: + code = 0 + elif row == NY.get(0) - 1: + code = 3 + return code + + +def _row_edge_code_periodic_tmpl(row): + return -1 + + +def _col_edge_code_active_tmpl(col): + code = -1 + if col == 0: + code = 1 + elif col == NX.get(0) - 1: + code = 2 + return code + + +def _col_edge_code_periodic_tmpl(col): + return -1 + + +def _which_edge_tmpl(i): + row = _ROW(i) + col = _COL(i) + code = _ROWEDGECODE(row) + if code == -1: + code = _COLEDGECODE(col) + return code + + +# --------------------------------------------------------------------------- +# public: can_out +# --------------------------------------------------------------------------- + + +def _can_out_mask_tmpl(i): + out = 0 + if OUTLET_MASK.get(i) == 1: + out = 1 + return out + + +def _can_out_edge_tmpl(i): + return _ISONEDGE(i) + + +# --------------------------------------------------------------------------- +# public: dist_from_k / dist_between_nodes +# --------------------------------------------------------------------------- + + +def _dist_from_k_d4_tmpl(k): + return DX.get(0) + + +def _dist_from_k_d8_tmpl(k): + d = DX.get(0) + if k == 0 or k == 2 or k == 5 or k == 7: + d = DX.get(0) * SQRT2 + return d + + +def _row_dist_normal_tmpl(raw): + return _BK.abs(raw) + + +def _row_dist_periodic_tmpl(raw): + d = _BK.abs(raw) + return _BK.min(d, NY.get(0) - d) + + +def _col_dist_normal_tmpl(raw): + return _BK.abs(raw) + + +def _col_dist_periodic_tmpl(raw): + d = _BK.abs(raw) + return _BK.min(d, NX.get(0) - d) + + +def _dist_between_d4_tmpl(i, j): + out = -1.0 + if j >= 0: + dr = _ROWDIST(_ROW(j) - _ROW(i)) + dc = _COLDIST(_COL(j) - _COL(i)) + if dr == 0 and dc == 1: + out = DX.get(0) + elif dr == 1 and dc == 0: + out = DX.get(0) + return out + + +def _dist_between_d8_tmpl(i, j): + out = -1.0 + if j >= 0: + dr = _ROWDIST(_ROW(j) - _ROW(i)) + dc = _COLDIST(_COL(j) - _COL(i)) + if dr == 0 and dc == 1: + out = DX.get(0) + elif dr == 1 and dc == 0: + out = DX.get(0) + elif dr == 1 and dc == 1: + out = DX.get(0) * SQRT2 + return out + + +# --------------------------------------------------------------------------- +# public: neighbour_and_distance +# --------------------------------------------------------------------------- + + +def _neighbour_and_distance_tmpl(i, k): + j = _NEIGHBOUR(i, k) + d = -1.0 + if j != -1: + d = _DISTFROMK(k) + return j, d + + +def build_helpers( + HelperCls, + *, + nx_p, + ny_p, + dx_p, + nodata_mask_p, + outlet_mask_p, + topology, + boundary, + nodata, + outlet, + backend_mod, +): + """ + Wire one grid's private blocks and public composites for a closure + backend (Taichi or Quadrants), picking each block's variant from + `topology`/`boundary`/`nodata`/`outlet` and binding private blocks into + public ones by name. + + Returns {public_name: HelperBuilder}, meant to be merged straight into + the Bag make_grid() returns. + + Author: B.G (07/2026) + """ + import math + + sqrt2 = math.sqrt(2.0) + d8 = topology == "D8" + + mk = functools.partial(make_helper, HelperCls) + + row = mk(_row_tmpl, NX=nx_p) + col = mk(_col_tmpl, NX=nx_p) + index = mk(_index_tmpl, NX=nx_p) + + delta = mk(_delta_d8_tmpl if d8 else _delta_d4_tmpl) + + row_wrap = mk(_row_wrap_periodic_tmpl if boundary == "periodic_NS" else _row_wrap_identity_tmpl, NY=ny_p) + col_wrap = mk(_col_wrap_periodic_tmpl if boundary == "periodic_EW" else _col_wrap_identity_tmpl, NX=nx_p) + wrap = mk(_wrap_tmpl, _ROWWRAP=row_wrap, _COLWRAP=col_wrap) + + row_edge_ok = mk( + _row_edge_ok_periodic_tmpl if boundary == "periodic_NS" else _row_edge_ok_bounded_tmpl, NY=ny_p + ) + col_edge_ok = mk( + _col_edge_ok_periodic_tmpl if boundary == "periodic_EW" else _col_edge_ok_bounded_tmpl, NX=nx_p + ) + + source_ok = mk(_source_ok_nodata_tmpl, NODATA_MASK=nodata_mask_p) if nodata else mk(_source_ok_always_tmpl) + + move_allowed = mk( + _move_allowed_tmpl, + _ROW=row, + _COL=col, + _DELTA=delta, + _ROWEDGEOK=row_edge_ok, + _COLEDGEOK=col_edge_ok, + _SOURCEOK=source_ok, + ) + + valid_binds = {"NX": nx_p, "NY": ny_p} + if nodata: + valid_binds["NODATA_MASK"] = nodata_mask_p + valid = mk(_valid_nodata_tmpl if nodata else _valid_no_nodata_tmpl, **valid_binds) + + neighbour_raw = mk(_neighbour_raw_tmpl, _ROW=row, _COL=col, _DELTA=delta, _WRAP=wrap, _INDEX=index) + neighbour = mk(_neighbour_tmpl, _MOVEALLOWED=move_allowed, _NEIGHBOURRAW=neighbour_raw, _VALID=valid) + + is_active = mk(_is_active_mask_tmpl, NODATA_MASK=nodata_mask_p) if nodata else mk(_is_active_always_tmpl) + nodata_fn = mk(_nodata_tmpl, _ISACTIVE=is_active) + + row_is_edge = mk( + _row_is_edge_periodic_tmpl if boundary == "periodic_NS" else _row_is_edge_active_tmpl, NY=ny_p + ) + col_is_edge = mk( + _col_is_edge_periodic_tmpl if boundary == "periodic_EW" else _col_is_edge_active_tmpl, NX=nx_p + ) + is_on_edge = mk(_is_on_edge_tmpl, _ROW=row, _COL=col, _ROWISEDGE=row_is_edge, _COLISEDGE=col_is_edge) + + row_edge_code = mk( + _row_edge_code_periodic_tmpl if boundary == "periodic_NS" else _row_edge_code_active_tmpl, NY=ny_p + ) + col_edge_code = mk( + _col_edge_code_periodic_tmpl if boundary == "periodic_EW" else _col_edge_code_active_tmpl, NX=nx_p + ) + which_edge = mk( + _which_edge_tmpl, _ROW=row, _COL=col, _ROWEDGECODE=row_edge_code, _COLEDGECODE=col_edge_code + ) + + if outlet == "mask": + can_out = mk(_can_out_mask_tmpl, OUTLET_MASK=outlet_mask_p) + else: + can_out = mk(_can_out_edge_tmpl, _ISONEDGE=is_on_edge) + + dist_from_k = mk(_dist_from_k_d8_tmpl if d8 else _dist_from_k_d4_tmpl, DX=dx_p, SQRT2=sqrt2) + + row_dist = mk( + _row_dist_periodic_tmpl if boundary == "periodic_NS" else _row_dist_normal_tmpl, NY=ny_p, _BK=backend_mod + ) + col_dist = mk( + _col_dist_periodic_tmpl if boundary == "periodic_EW" else _col_dist_normal_tmpl, NX=nx_p, _BK=backend_mod + ) + dist_between = mk( + _dist_between_d8_tmpl if d8 else _dist_between_d4_tmpl, + _ROW=row, + _COL=col, + _ROWDIST=row_dist, + _COLDIST=col_dist, + DX=dx_p, + SQRT2=sqrt2, + ) + + neighbour_and_distance = mk(_neighbour_and_distance_tmpl, _NEIGHBOUR=neighbour, _DISTFROMK=dist_from_k) + + return { + "neighbour": neighbour, + "neighbour_raw": neighbour_raw, + "nodata": nodata_fn, + "is_active": is_active, + "can_out": can_out, + "dist_from_k": dist_from_k, + "dist_between_nodes": dist_between, + "is_on_edge": is_on_edge, + "which_edge": which_edge, + "neighbour_and_distance": neighbour_and_distance, + } diff --git a/pyfastflow/experimental/grid/_cupy_blocks.py b/pyfastflow/experimental/grid/_cupy_blocks.py new file mode 100644 index 0000000..d34cd71 --- /dev/null +++ b/pyfastflow/experimental/grid/_cupy_blocks.py @@ -0,0 +1,438 @@ +""" +cupy (CUDA source) block templates behind make_grid. + +Mirrors _closure_blocks.py block for block - same private/public split, same +per-axis composability (a periodic boundary swaps in the "periodic" __device__ +variant of a row or column block, the untouched axis keeps its +"identity"/"bounded" variant) - written as CUDA text instead of python defs, +since that is what CupyHelperBuilder/CupyKernelBuilder compile (see +cupy_backend.py's module docstring for the `$...$` span mechanism). + +_delta(k) is the one runtime if-ladder equivalent: k is per-call device data, +not a structural choice, so here it is a `__constant__` int table indexed by +k at runtime instead - a dynamically-indexed local array would live in local +memory on GPU, `__constant__` does not. Callers that loop over k statically +should add `#pragma unroll` at the call site; that is a caller concern, not +something a block enforces. + +Every device function name is prefixed with this grid's own tag (a fresh +new_uid()), so two make_grid() calls in one process never collide inside a +single compiled cupy module even if both are bound into the same kernel. + +Author: B.G (07/2026) +""" + +import functools +import math + +from ..core.context.backends import make_helper +from ..core.pool.base import new_uid + + +def build_helpers( + HelperCls, + *, + nx_p, + ny_p, + dx_p, + nodata_mask_p, + outlet_mask_p, + topology, + boundary, + nodata, + outlet, + backend_mod=None, +): + """ + Wire one grid's private blocks and public composites for the cupy + backend, picking each block's variant from + `topology`/`boundary`/`nodata`/`outlet` and binding private blocks into + public ones by name (a bound name resolves to the real emitted C symbol + at span-expansion time - see cupy_backend.py's _SpanParser). + + Returns {public_name: HelperBuilder}, meant to be merged straight into + the Bag make_grid() returns. `backend_mod` is accepted for signature + parity with the closure backend's build_helpers and unused here - cupy + templates call plain C (abs, ternaries) rather than a bound backend + module. + + Author: B.G (07/2026) + """ + t = f"pf{new_uid()}" + d8 = topology == "D8" + sqrt2 = math.sqrt(2.0) + + mk = functools.partial(make_helper, HelperCls) + + row = mk( + f"__device__ int {t}_row(int i) {{ return i / $NX.get(0)$; }}", + NX=nx_p, + ) + col = mk( + f"__device__ int {t}_col(int i) {{ return i % $NX.get(0)$; }}", + NX=nx_p, + ) + index = mk( + f"__device__ int {t}_index(int row, int col) {{ return row * $NX.get(0)$ + col; }}", + NX=nx_p, + ) + + if d8: + delta = mk( + f""" +__constant__ int {t}_DELTA_DR[8] = {{-1, -1, -1, 0, 0, 1, 1, 1}}; +__constant__ int {t}_DELTA_DC[8] = {{-1, 0, 1, -1, 1, -1, 0, 1}}; +__device__ void {t}_delta(int k, int* dr, int* dc) {{ + *dr = {t}_DELTA_DR[k]; + *dc = {t}_DELTA_DC[k]; +}} +""" + ) + else: + delta = mk( + f""" +__constant__ int {t}_DELTA_DR[4] = {{-1, 0, 0, 1}}; +__constant__ int {t}_DELTA_DC[4] = {{0, -1, 1, 0}}; +__device__ void {t}_delta(int k, int* dr, int* dc) {{ + *dr = {t}_DELTA_DR[k]; + *dc = {t}_DELTA_DC[k]; +}} +""" + ) + + row_wrap = ( + mk(f"__device__ int {t}_row_wrap(int row) {{ return row; }}") + if boundary != "periodic_NS" + else mk( + f""" +__device__ int {t}_row_wrap(int row) {{ + int r = row; + if (r < 0) r += $NY.get(0)$; + else if (r >= $NY.get(0)$) r -= $NY.get(0)$; + return r; +}} +""", + NY=ny_p, + ) + ) + col_wrap = ( + mk(f"__device__ int {t}_col_wrap(int col) {{ return col; }}") + if boundary != "periodic_EW" + else mk( + f""" +__device__ int {t}_col_wrap(int col) {{ + int c = col; + if (c < 0) c += $NX.get(0)$; + else if (c >= $NX.get(0)$) c -= $NX.get(0)$; + return c; +}} +""", + NX=nx_p, + ) + ) + wrap = mk( + f""" +__device__ void {t}_wrap(int row, int col, int* wrow, int* wcol) {{ + *wrow = $row_wrap(row)$; + *wcol = $col_wrap(col)$; +}} +""", + row_wrap=row_wrap, + col_wrap=col_wrap, + ) + + row_edge_ok = ( + mk(f"__device__ int {t}_row_edge_ok(int row) {{ return 1; }}") + if boundary == "periodic_NS" + else mk( + f"__device__ int {t}_row_edge_ok(int row) {{ return (row >= 0 && row < $NY.get(0)$) ? 1 : 0; }}", + NY=ny_p, + ) + ) + col_edge_ok = ( + mk(f"__device__ int {t}_col_edge_ok(int col) {{ return 1; }}") + if boundary == "periodic_EW" + else mk( + f"__device__ int {t}_col_edge_ok(int col) {{ return (col >= 0 && col < $NX.get(0)$) ? 1 : 0; }}", + NX=nx_p, + ) + ) + + source_ok = ( + mk( + f"__device__ int {t}_source_ok(int i) {{ return ($NODATA_MASK.get(i)$ == 1) ? 0 : 1; }}", + NODATA_MASK=nodata_mask_p, + ) + if nodata + else mk(f"__device__ int {t}_source_ok(int i) {{ return 1; }}") + ) + + move_allowed = mk( + f""" +__device__ int {t}_move_allowed(int i, int k) {{ + int row = $row(i)$; + int col = $col(i)$; + int dr, dc; + $delta(k, &dr, &dc)$; + int a = $row_edge_ok(row + dr)$; + int b = $col_edge_ok(col + dc)$; + int c = $source_ok(i)$; + return a * b * c; +}} +""", + row=row, + col=col, + delta=delta, + row_edge_ok=row_edge_ok, + col_edge_ok=col_edge_ok, + source_ok=source_ok, + ) + + valid_binds = {"NX": nx_p, "NY": ny_p} + if nodata: + valid_binds["NODATA_MASK"] = nodata_mask_p + valid = mk( + f""" +__device__ int {t}_valid(int j) {{ + if (j < 0 || j >= $NX.get(0)$ * $NY.get(0)$) return 0; + return ($NODATA_MASK.get(j)$ == 1) ? 0 : 1; +}} +""", + **valid_binds, + ) + else: + valid = mk( + f"__device__ int {t}_valid(int j) {{ return (j >= 0 && j < $NX.get(0)$ * $NY.get(0)$) ? 1 : 0; }}", + **valid_binds, + ) + + neighbour_raw = mk( + f""" +__device__ int {t}_neighbour_raw(int i, int k) {{ + int row = $row(i)$; + int col = $col(i)$; + int dr, dc; + $delta(k, &dr, &dc)$; + int wrow, wcol; + $wrap(row + dr, col + dc, &wrow, &wcol)$; + return $index(wrow, wcol)$; +}} +""", + row=row, + col=col, + delta=delta, + wrap=wrap, + index=index, + ) + + neighbour = mk( + f""" +__device__ int {t}_neighbour(int i, int k) {{ + int j = -1; + if ($move_allowed(i, k)$ == 1) {{ + int cand = $neighbour_raw(i, k)$; + if ($valid(cand)$ == 1) j = cand; + }} + return j; +}} +""", + move_allowed=move_allowed, + neighbour_raw=neighbour_raw, + valid=valid, + ) + + is_active = ( + mk( + f"__device__ int {t}_is_active(int i) {{ return ($NODATA_MASK.get(i)$ == 1) ? 0 : 1; }}", + NODATA_MASK=nodata_mask_p, + ) + if nodata + else mk(f"__device__ int {t}_is_active(int i) {{ return 1; }}") + ) + nodata_fn = mk( + f"__device__ int {t}_nodata(int i) {{ return 1 - $is_active(i)$; }}", + is_active=is_active, + ) + + row_is_edge = ( + mk(f"__device__ int {t}_row_is_edge(int row) {{ return 0; }}") + if boundary == "periodic_NS" + else mk( + f"__device__ int {t}_row_is_edge(int row) {{ return (row == 0 || row == $NY.get(0)$ - 1) ? 1 : 0; }}", + NY=ny_p, + ) + ) + col_is_edge = ( + mk(f"__device__ int {t}_col_is_edge(int col) {{ return 0; }}") + if boundary == "periodic_EW" + else mk( + f"__device__ int {t}_col_is_edge(int col) {{ return (col == 0 || col == $NX.get(0)$ - 1) ? 1 : 0; }}", + NX=nx_p, + ) + ) + is_on_edge = mk( + f""" +__device__ int {t}_is_on_edge(int i) {{ + int row = $row(i)$; + int col = $col(i)$; + return ($row_is_edge(row)$ == 1 || $col_is_edge(col)$ == 1) ? 1 : 0; +}} +""", + row=row, + col=col, + row_is_edge=row_is_edge, + col_is_edge=col_is_edge, + ) + + row_edge_code = ( + mk(f"__device__ int {t}_row_edge_code(int row) {{ return -1; }}") + if boundary == "periodic_NS" + else mk( + f""" +__device__ int {t}_row_edge_code(int row) {{ + if (row == 0) return 0; + if (row == $NY.get(0)$ - 1) return 3; + return -1; +}} +""", + NY=ny_p, + ) + ) + col_edge_code = ( + mk(f"__device__ int {t}_col_edge_code(int col) {{ return -1; }}") + if boundary == "periodic_EW" + else mk( + f""" +__device__ int {t}_col_edge_code(int col) {{ + if (col == 0) return 1; + if (col == $NX.get(0)$ - 1) return 2; + return -1; +}} +""", + NX=nx_p, + ) + ) + which_edge = mk( + f""" +__device__ int {t}_which_edge(int i) {{ + int row = $row(i)$; + int col = $col(i)$; + int code = $row_edge_code(row)$; + if (code == -1) code = $col_edge_code(col)$; + return code; +}} +""", + row=row, + col=col, + row_edge_code=row_edge_code, + col_edge_code=col_edge_code, + ) + + if outlet == "mask": + can_out = mk( + f"__device__ int {t}_can_out(int i) {{ return ($OUTLET_MASK.get(i)$ == 1) ? 1 : 0; }}", + OUTLET_MASK=outlet_mask_p, + ) + else: + can_out = mk( + f"__device__ int {t}_can_out(int i) {{ return $is_on_edge(i)$; }}", + is_on_edge=is_on_edge, + ) + + if d8: + dist_from_k = mk( + f""" +__device__ float {t}_dist_from_k(int k) {{ + float d = $DX.get(0)$; + if (k == 0 || k == 2 || k == 5 || k == 7) d = $DX.get(0)$ * {sqrt2}f; + return d; +}} +""", + DX=dx_p, + ) + else: + dist_from_k = mk( + f"__device__ float {t}_dist_from_k(int k) {{ return $DX.get(0)$; }}", + DX=dx_p, + ) + + row_dist = ( + mk(f"__device__ int {t}_row_dist(int raw) {{ return abs(raw); }}") + if boundary != "periodic_NS" + else mk( + f""" +__device__ int {t}_row_dist(int raw) {{ + int d = abs(raw); + int n = $NY.get(0)$; + return d < (n - d) ? d : (n - d); +}} +""", + NY=ny_p, + ) + ) + col_dist = ( + mk(f"__device__ int {t}_col_dist(int raw) {{ return abs(raw); }}") + if boundary != "periodic_EW" + else mk( + f""" +__device__ int {t}_col_dist(int raw) {{ + int d = abs(raw); + int n = $NX.get(0)$; + return d < (n - d) ? d : (n - d); +}} +""", + NX=nx_p, + ) + ) + + diag_line = f" else if (dr == 1 && dc == 1) out = $DX.get(0)$ * {sqrt2}f;\n" if d8 else "" + dist_between = mk( + f""" +__device__ float {t}_dist_between(int i, int j) {{ + float out = -1.0f; + if (j >= 0) {{ + int ri = $row(i)$; + int ci = $col(i)$; + int rj = $row(j)$; + int cj = $col(j)$; + int dr = $row_dist(rj - ri)$; + int dc = $col_dist(cj - ci)$; + if (dr == 0 && dc == 1) out = $DX.get(0)$; + else if (dr == 1 && dc == 0) out = $DX.get(0)$; +{diag_line} }} + return out; +}} +""", + row=row, + col=col, + row_dist=row_dist, + col_dist=col_dist, + DX=dx_p, + ) + + neighbour_and_distance = mk( + f""" +__device__ void {t}_neighbour_and_distance(int i, int k, int* j_out, float* d_out) {{ + int j = $neighbour(i, k)$; + float d = -1.0f; + if (j != -1) d = $dist_from_k(k)$; + *j_out = j; + *d_out = d; +}} +""", + neighbour=neighbour, + dist_from_k=dist_from_k, + ) + + return { + "neighbour": neighbour, + "neighbour_raw": neighbour_raw, + "nodata": nodata_fn, + "is_active": is_active, + "can_out": can_out, + "dist_from_k": dist_from_k, + "dist_between_nodes": dist_between, + "is_on_edge": is_on_edge, + "which_edge": which_edge, + "neighbour_and_distance": neighbour_and_distance, + } diff --git a/pyfastflow/experimental/noise/__init__.py b/pyfastflow/experimental/noise/__init__.py new file mode 100644 index 0000000..4b7ef51 --- /dev/null +++ b/pyfastflow/experimental/noise/__init__.py @@ -0,0 +1,203 @@ +""" +make_noise: the NoiseContext-equivalent Bag factory, built on the +backend-agnostic core (see ..core.context: parameter.py for Parameter, compile.py for HelperBuilder, bag.py for Bag) +and on a grid Bag from ..grid. + +Like make_grid there is no stateful context class - make_noise builds a Bag +once and hands it back. The bag is device helpers plus their parameters, +nothing else: no allocation, no fill kernel, no host-side generate_*. A caller +binds the bag into its own kernel and reads `noise.at(i)` inline, so noise +composes with whatever that kernel is already doing rather than forcing a +separate pass over a temporary field. + + noise = make_noise("taichi", pool, grid, kind="perlin", octaves=4) + + def init_template(z: ti.template()): + for i in z: + z[i] = noise.at(i) + + TaichiKernelBuilder().bind("noise", noise).ingest(init_template).compile() + +Two kinds of knobs, same split as make_grid: + - value params (amplitude, seed, frequency_x/y, octaves, persistence) - + mode-overridable, always read in device code through `.get(0)`. + - one structural selector, `kind` ("white" or "perlin") - it picks which + chain of private blocks the public `at` is wired to. The public surface + is `at(i)` whatever the kind, so swapping generators is a build-time + config change and the calling template never moves. + +Mode defaults are "const" - the same "flexible at build time, dense at +runtime" stance make_grid takes - with one exception: `seed` (white noise +only) defaults to "scalar", so reseeding is a host write rather than a +rebuild. That keeps it symmetric with Perlin, whose seed is not a device +parameter at all: it drives the permutation table, which is built on the host +and uploaded. Reseed Perlin by refilling that field: + + noise.perm.set(permutation_table(7)) + +Bag members are `at` plus whatever the kind actually uses: white carries +`amplitude`, `seed`, `white_unit`; perlin carries `amplitude`, `perm`, +`frequency_x`, `frequency_y`, `octaves`, `persistence`, `perlin_at`. A member +that a kind does not use is absent from the bag rather than present and inert. + +Values match pyfastflow/noise/ for the same seed and settings - this is a port +of that arithmetic, not a reimplementation. White noise lands in +[-amplitude, amplitude]; Perlin is the octave-averaged lattice noise scaled by +amplitude. + +Author: B.G (07/2026) +""" + +import numpy as np + +from ..core.context.backends import backend_classes +from ..core.context.bag import Bag + +_KINDS = frozenset({"white", "perlin"}) +_MODES = ("const", "scalar") + + +def permutation_table(seed: int) -> np.ndarray: + """ + The 512-entry Perlin permutation table for one seed: a Fisher-Yates + shuffle of 0..255 concatenated with itself, so a lattice lookup can index + past 255 without wrapping by hand. + + Same construction (and same numpy Generator) as + pyfastflow/noise/noisecontext.py, so a given seed yields the same table + and therefore the same noise. + + Author: B.G (07/2026) + """ + rng = np.random.default_rng(seed) + perm = np.arange(256, dtype=np.int32) + for i in range(255, 0, -1): + j = rng.integers(0, i + 1) + perm[i], perm[j] = perm[j], perm[i] + return np.concatenate([perm, perm]) + + +def _blocks_for(backend: str): + """ + The private block module implementing make_noise's device code for one + backend name: the closure blocks (shared by Taichi and Quadrants) or the + cupy blocks. + + Author: B.G (07/2026) + """ + if backend in ("taichi", "quadrants"): + from . import _closure_blocks as blocks + elif backend == "cupy": + from . import _cupy_blocks as blocks + else: + raise ValueError(f"make_noise: unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + return blocks + + +def make_noise( + backend: str, + pool, + grid: Bag, + *, + kind: str = "perlin", + amplitude: float = 1.0, + seed: int = 42, + frequency: float = 8.0, + frequency_x: float | None = None, + frequency_y: float | None = None, + octaves: int = 4, + persistence: float = 0.5, + amplitude_mode: str = "const", + seed_mode: str = "scalar", + frequency_mode: str = "const", + octaves_mode: str = "const", + persistence_mode: str = "const", +) -> Bag: + """ + Build one noise bag: the public `at(i)` device helper and the parameters + the chosen `kind` needs, all reading row/column off the `grid` bag rather + than carrying their own geometry. + + `kind` "white"|"perlin" picks the block chain (see _closure_blocks.py / + _cupy_blocks.py). `frequency` sets both axes; `frequency_x`/`frequency_y` + override one axis each. `octaves`/`persistence`/`frequency_*` are Perlin + only and are not allocated for white noise. + + The `*_mode` arguments are "const" or "scalar" and decide whether a value + is folded in at compile time or lives in a one-cell device field the host + can retune. Defaults are "const" except `seed_mode`, which is "scalar" - + see the module docstring. + + Author: B.G (07/2026) + """ + if kind not in _KINDS: + raise ValueError(f"make_noise: kind must be one of {sorted(_KINDS)}, got {kind!r}") + for label, mode in ( + ("amplitude_mode", amplitude_mode), + ("seed_mode", seed_mode), + ("frequency_mode", frequency_mode), + ("octaves_mode", octaves_mode), + ("persistence_mode", persistence_mode), + ): + if mode not in _MODES: + raise ValueError(f"make_noise: {label} must be 'const' or 'scalar', got {mode!r}") + + backend_mod, ParamCls, HelperCls, dtypes = backend_classes(backend) + blocks = _blocks_for(backend) + + amplitude_p = ParamCls( + "NOISE_AMPLITUDE", dtype=dtypes["f32"], mode=amplitude_mode, value=float(amplitude), pool=pool + ) + + seed_p = None + perm_p = None + frequency_x_p = None + frequency_y_p = None + octaves_p = None + persistence_p = None + + if kind == "white": + seed_p = ParamCls("NOISE_SEED", dtype=dtypes["u32"], mode=seed_mode, value=int(seed), pool=pool) + members = {"amplitude": amplitude_p, "seed": seed_p} + else: + perm_p = ParamCls( + "NOISE_PERM", + dtype=dtypes["i32"], + mode="field", + value=permutation_table(seed), + pool=pool, + n_flat=512, + ) + fx = float(frequency_x if frequency_x is not None else frequency) + fy = float(frequency_y if frequency_y is not None else frequency) + frequency_x_p = ParamCls("NOISE_FX", dtype=dtypes["f32"], mode=frequency_mode, value=fx, pool=pool) + frequency_y_p = ParamCls("NOISE_FY", dtype=dtypes["f32"], mode=frequency_mode, value=fy, pool=pool) + octaves_p = ParamCls("NOISE_OCTAVES", dtype=dtypes["i32"], mode=octaves_mode, value=int(octaves), pool=pool) + persistence_p = ParamCls( + "NOISE_PERSISTENCE", dtype=dtypes["f32"], mode=persistence_mode, value=float(persistence), pool=pool + ) + members = { + "amplitude": amplitude_p, + "perm": perm_p, + "frequency_x": frequency_x_p, + "frequency_y": frequency_y_p, + "octaves": octaves_p, + "persistence": persistence_p, + } + + helpers = blocks.build_helpers( + HelperCls, + grid=grid, + kind=kind, + amplitude_p=amplitude_p, + seed_p=seed_p, + perm_p=perm_p, + frequency_x_p=frequency_x_p, + frequency_y_p=frequency_y_p, + octaves_p=octaves_p, + persistence_p=persistence_p, + backend_mod=backend_mod, + ) + + members.update(helpers) + return Bag(members) diff --git a/pyfastflow/experimental/noise/_closure_blocks.py b/pyfastflow/experimental/noise/_closure_blocks.py new file mode 100644 index 0000000..f99d870 --- /dev/null +++ b/pyfastflow/experimental/noise/_closure_blocks.py @@ -0,0 +1,216 @@ +""" +Taichi/Quadrants (closure) block templates behind make_noise. + +Same shape as the grid's _closure_blocks.py: every block is a plain python +def, PICKED by build_helpers() from the noise config rather than branched on +inside one body. Here the only structural selector is `kind` - it decides +whether the public `at(i)` is wired to the white-noise chain or the Perlin +chain, and nothing downstream of `at` is shared between the two. + +The arithmetic is a port of pyfastflow/noise/white_noise.py and +perlin_noise.py, kept value-for-value identical: same integer hash constants, +same fade/lerp/grad, same octave accumulation, and the same argument order +(column first, row second) into the hash, so a bag built here reproduces what +NoiseContext produced for the same seed. + +Row/column come from the bound grid Bag (GRID.nx/GRID.ny read through +`.get(0)` like any other value param), so a noise bag inherits whatever mode +the grid's geometry is in and never carries its own copy of nx/ny. + +`_BK` is the backend module (taichi or quadrants) - both expose the same +cast/u32/i32/f32/floor surface, which is all these blocks need. + +Author: B.G (07/2026) +""" + +import functools + +from ..core.context.backends import make_helper + +# --------------------------------------------------------------------------- +# shared: flat index -> row / column, off the bound grid bag +# --------------------------------------------------------------------------- + + +def _row_tmpl(i): + return i // GRID.nx.get(0) + + +def _col_tmpl(i): + return i % GRID.nx.get(0) + + +# --------------------------------------------------------------------------- +# white noise +# --------------------------------------------------------------------------- + + +def _hash_u32_tmpl(x): + h = x + h ^= h >> _BK.u32(16) + h *= _BK.u32(0x7FEB352D) + h ^= h >> _BK.u32(15) + h *= _BK.u32(0x846CA68B) + h ^= h >> _BK.u32(16) + return h + + +def _white_unit_tmpl(i): + # column first, row second - the argument order white_noise.py hashes in + col = _COL(i) + row = _ROW(i) + key = _BK.u32(SEED.get(0)) + key ^= _BK.u32(col) * _BK.u32(374761393) + key ^= _BK.u32(row) * _BK.u32(668265263) + hashed = _HASH(key) + return _BK.cast(hashed, _BK.f32) / 4294967296.0 + + +def _at_white_tmpl(i): + return (_WHITEUNIT(i) - 0.5) * 2.0 * AMP.get(0) + + +# --------------------------------------------------------------------------- +# perlin noise +# --------------------------------------------------------------------------- + + +def _fade_tmpl(t): + return t * t * t * (t * (t * 6.0 - 15.0) + 10.0) + + +def _lerp_tmpl(t, a, b): + return a + t * (b - a) + + +def _grad_tmpl(hash_val, dx, dy): + idx = hash_val & 7 + gx = 0.0 + gy = 0.0 + + if idx == 0: + gx, gy = 1.0, 1.0 + elif idx == 1: + gx, gy = -1.0, 1.0 + elif idx == 2: + gx, gy = 1.0, -1.0 + elif idx == 3: + gx, gy = -1.0, -1.0 + elif idx == 4: + gx, gy = 1.0, 0.0 + elif idx == 5: + gx, gy = -1.0, 0.0 + elif idx == 6: + gx, gy = 0.0, 1.0 + else: + gx, gy = 0.0, -1.0 + + return gx * dx + gy * dy + + +def _perlin_at_tmpl(x, y): + x_floor = _BK.floor(x) + y_floor = _BK.floor(y) + + X = _BK.cast(x_floor, _BK.i32) & 255 + Y = _BK.cast(y_floor, _BK.i32) & 255 + + x_local = x - x_floor + y_local = y - y_floor + + u = _FADE(x_local) + v = _FADE(y_local) + + A = PERM.get(X) + Y + B = PERM.get((X + 1) & 255) + Y + AA = PERM.get(A & 255) + AB = PERM.get((A + 1) & 255) + BA = PERM.get(B & 255) + BB = PERM.get((B + 1) & 255) + + return _LERP( + v, + _LERP(u, _GRAD(AA, x_local, y_local), _GRAD(BA, x_local - 1.0, y_local)), + _LERP(u, _GRAD(AB, x_local, y_local - 1.0), _GRAD(BB, x_local - 1.0, y_local - 1.0)), + ) + + +def _at_perlin_tmpl(i): + nx_f = _BK.cast(GRID.nx.get(0), _BK.f32) + ny_f = _BK.cast(GRID.ny.get(0), _BK.f32) + + x = _BK.cast(_COL(i), _BK.f32) * FX.get(0) / nx_f + y = _BK.cast(_ROW(i), _BK.f32) * FY.get(0) / ny_f + + total = 0.0 + max_value = 0.0 + current_amplitude = 1.0 + current_frequency = 1.0 + + for _ in range(OCTAVES.get(0)): + total += _PERLINAT(x * current_frequency, y * current_frequency) * current_amplitude + max_value += current_amplitude + current_amplitude *= PERSISTENCE.get(0) + current_frequency *= 2.0 + + out = 0.0 + if max_value > 0.0: + out = (total / max_value) * AMP.get(0) + return out + + +def build_helpers( + HelperCls, + *, + grid, + kind, + amplitude_p, + seed_p, + perm_p, + frequency_x_p, + frequency_y_p, + octaves_p, + persistence_p, + backend_mod, +): + """ + Wire one noise bag's private blocks and its public `at` for a closure + backend (Taichi or Quadrants), picking the white or Perlin chain from + `kind` and binding private blocks into public ones by name. + + Returns {public_name: HelperBuilder}, meant to be merged straight into + the Bag make_noise() returns. The parameters not used by the chosen + `kind` arrive as None and are simply never bound. + + Author: B.G (07/2026) + """ + + mk = functools.partial(make_helper, HelperCls) + + row = mk(_row_tmpl, GRID=grid) + col = mk(_col_tmpl, GRID=grid) + + if kind == "white": + hash_u32 = mk(_hash_u32_tmpl, _BK=backend_mod) + white_unit = mk(_white_unit_tmpl, _ROW=row, _COL=col, _HASH=hash_u32, SEED=seed_p, _BK=backend_mod) + at = mk(_at_white_tmpl, _WHITEUNIT=white_unit, AMP=amplitude_p) + return {"at": at, "white_unit": white_unit} + + fade = mk(_fade_tmpl) + lerp = mk(_lerp_tmpl) + grad = mk(_grad_tmpl) + perlin_at = mk(_perlin_at_tmpl, _FADE=fade, _LERP=lerp, _GRAD=grad, PERM=perm_p, _BK=backend_mod) + at = mk( + _at_perlin_tmpl, + _ROW=row, + _COL=col, + _PERLINAT=perlin_at, + GRID=grid, + FX=frequency_x_p, + FY=frequency_y_p, + OCTAVES=octaves_p, + PERSISTENCE=persistence_p, + AMP=amplitude_p, + _BK=backend_mod, + ) + return {"at": at, "perlin_at": perlin_at} diff --git a/pyfastflow/experimental/noise/_cupy_blocks.py b/pyfastflow/experimental/noise/_cupy_blocks.py new file mode 100644 index 0000000..2160004 --- /dev/null +++ b/pyfastflow/experimental/noise/_cupy_blocks.py @@ -0,0 +1,210 @@ +""" +cupy (CUDA source) block templates behind make_noise. + +Mirrors _closure_blocks.py block for block - same private/public split, same +`kind` selector deciding which chain `at(i)` is wired to - written as CUDA +text instead of python defs, since that is what CupyHelperBuilder compiles +(see cupy_backend.py's module docstring for the `$...$` span mechanism). + +The arithmetic is the same port of pyfastflow/noise/white_noise.py and +perlin_noise.py the closure blocks carry, so the two backends agree bit for +bit on the white-noise hash and octave for octave on Perlin. + +Every device function name is prefixed with this noise bag's own tag (a fresh +new_uid()), so two make_noise() calls in one process never collide inside a +single compiled cupy module even if both are bound into the same kernel. + +Perlin's permutation lookups go through a named local (`int ia = A & 255;`) +before the span rather than inlining the expression into `$PERM.get(...)$` - +spans are resolved as text, and keeping their argument a bare identifier +keeps that resolution unambiguous. + +Author: B.G (07/2026) +""" + +import functools + +from ..core.context.backends import make_helper +from ..core.pool.base import new_uid + + +def build_helpers( + HelperCls, + *, + grid, + kind, + amplitude_p, + seed_p, + perm_p, + frequency_x_p, + frequency_y_p, + octaves_p, + persistence_p, + backend_mod=None, +): + """ + Wire one noise bag's private blocks and its public `at` for the cupy + backend, picking the white or Perlin chain from `kind` and binding + private blocks into public ones by name (a bound name resolves to the + real emitted C symbol at span-expansion time - see cupy_backend.py's + _SpanParser). + + Returns {public_name: HelperBuilder}, meant to be merged straight into + the Bag make_noise() returns. `backend_mod` is accepted for signature + parity with the closure backend's build_helpers and unused here - cupy + templates call plain C (floorf, casts) rather than a bound backend + module. + + Author: B.G (07/2026) + """ + t = f"pn{new_uid()}" + + mk = functools.partial(make_helper, HelperCls) + + row = mk(f"__device__ int {t}_row(int i) {{ return i / $GRID.nx.get(0)$; }}", GRID=grid) + col = mk(f"__device__ int {t}_col(int i) {{ return i % $GRID.nx.get(0)$; }}", GRID=grid) + + if kind == "white": + hash_u32 = mk( + f""" +__device__ unsigned int {t}_hash_u32(unsigned int x) {{ + unsigned int h = x; + h ^= h >> 16u; + h *= 0x7FEB352Du; + h ^= h >> 15u; + h *= 0x846CA68Bu; + h ^= h >> 16u; + return h; +}} +""" + ) + white_unit = mk( + f""" +__device__ float {t}_white_unit(int i) {{ + // column first, row second - the argument order white_noise.py hashes in + int c = $col(i)$; + int r = $row(i)$; + unsigned int key = (unsigned int)$SEED.get(0)$; + key ^= (unsigned int)c * 374761393u; + key ^= (unsigned int)r * 668265263u; + unsigned int hashed = $hash_u32(key)$; + return (float)hashed / 4294967296.0f; +}} +""", + row=row, + col=col, + hash_u32=hash_u32, + SEED=seed_p, + ) + at = mk( + f""" +__device__ float {t}_at(int i) {{ + return ($white_unit(i)$ - 0.5f) * 2.0f * $AMP.get(0)$; +}} +""", + white_unit=white_unit, + AMP=amplitude_p, + ) + return {"at": at, "white_unit": white_unit} + + fade = mk(f"__device__ float {t}_fade(float t) {{ return t * t * t * (t * (t * 6.0f - 15.0f) + 10.0f); }}") + lerp = mk(f"__device__ float {t}_lerp(float t, float a, float b) {{ return a + t * (b - a); }}") + grad = mk( + f""" +__device__ float {t}_grad(int hash_val, float dx, float dy) {{ + int idx = hash_val & 7; + float gx = 0.0f; + float gy = 0.0f; + if (idx == 0) {{ gx = 1.0f; gy = 1.0f; }} + else if (idx == 1) {{ gx = -1.0f; gy = 1.0f; }} + else if (idx == 2) {{ gx = 1.0f; gy = -1.0f; }} + else if (idx == 3) {{ gx = -1.0f; gy = -1.0f; }} + else if (idx == 4) {{ gx = 1.0f; gy = 0.0f; }} + else if (idx == 5) {{ gx = -1.0f; gy = 0.0f; }} + else if (idx == 6) {{ gx = 0.0f; gy = 1.0f; }} + else {{ gx = 0.0f; gy = -1.0f; }} + return gx * dx + gy * dy; +}} +""" + ) + perlin_at = mk( + f""" +__device__ float {t}_perlin_at(float x, float y) {{ + float x_floor = floorf(x); + float y_floor = floorf(y); + + int X = ((int)x_floor) & 255; + int Y = ((int)y_floor) & 255; + + float x_local = x - x_floor; + float y_local = y - y_floor; + + float u = $fade(x_local)$; + float v = $fade(y_local)$; + + int iX1 = (X + 1) & 255; + int A = $PERM.get(X)$ + Y; + int B = $PERM.get(iX1)$ + Y; + + int iAA = A & 255; + int iAB = (A + 1) & 255; + int iBA = B & 255; + int iBB = (B + 1) & 255; + + int AA = $PERM.get(iAA)$; + int AB = $PERM.get(iAB)$; + int BA = $PERM.get(iBA)$; + int BB = $PERM.get(iBB)$; + + float gaa = $grad(AA, x_local, y_local)$; + float gba = $grad(BA, x_local - 1.0f, y_local)$; + float gab = $grad(AB, x_local, y_local - 1.0f)$; + float gbb = $grad(BB, x_local - 1.0f, y_local - 1.0f)$; + + float lo = $lerp(u, gaa, gba)$; + float hi = $lerp(u, gab, gbb)$; + return $lerp(v, lo, hi)$; +}} +""", + fade=fade, + lerp=lerp, + grad=grad, + PERM=perm_p, + ) + at = mk( + f""" +__device__ float {t}_at(int i) {{ + float nx_f = (float)$GRID.nx.get(0)$; + float ny_f = (float)$GRID.ny.get(0)$; + + float x = (float)$col(i)$ * $FX.get(0)$ / nx_f; + float y = (float)$row(i)$ * $FY.get(0)$ / ny_f; + + float total = 0.0f; + float max_value = 0.0f; + float current_amplitude = 1.0f; + float current_frequency = 1.0f; + + int octaves = $OCTAVES.get(0)$; + for (int o = 0; o < octaves; ++o) {{ + total += $perlin_at(x * current_frequency, y * current_frequency)$ * current_amplitude; + max_value += current_amplitude; + current_amplitude *= $PERSISTENCE.get(0)$; + current_frequency *= 2.0f; + }} + + if (max_value > 0.0f) return (total / max_value) * $AMP.get(0)$; + return 0.0f; +}} +""", + row=row, + col=col, + perlin_at=perlin_at, + GRID=grid, + FX=frequency_x_p, + FY=frequency_y_p, + OCTAVES=octaves_p, + PERSISTENCE=persistence_p, + AMP=amplitude_p, + ) + return {"at": at, "perlin_at": perlin_at}