From 900f01de9aa238dcea53a3bef665ce4de270b920 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 20 Jul 2026 11:01:28 +0200 Subject: [PATCH 01/31] Adds an abstract pool and data class for pre-multi-backend refactor --- pyfastflow/pool/base.py | 145 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 pyfastflow/pool/base.py diff --git a/pyfastflow/pool/base.py b/pyfastflow/pool/base.py new file mode 100644 index 0000000..1f1596b --- /dev/null +++ b/pyfastflow/pool/base.py @@ -0,0 +1,145 @@ +""" +Backend-agnostic pool contracts. + +Defines the blueprint that every pool backend (Taichi fields, numpy arrays, +future tree/sparse structures, ...) must implement. No allocation logic here +- this is the interface only. + +Author: B.G (07/2026) +""" + +from abc import ABC, abstractmethod +from typing import Any + + +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: + id: Unique handle identifier, assigned by the backend. + dtype: Backend-native or common dtype tag for this resource. + shape: Resource dimensions. () for a scalar. + in_use: True between acquire() and release(). + + Author: B.G (07/2026) + """ + + id: int + dtype: Any + shape: tuple[int, ...] + in_use: bool + + @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) + """ + ... From 815c33a8d98013ac70112753b8fc4ad0461e4ba6 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 20 Jul 2026 13:11:45 +0200 Subject: [PATCH 02/31] Add the taichi field data handle --- pyfastflow/pool/taichi_handle.py | 86 ++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 pyfastflow/pool/taichi_handle.py diff --git a/pyfastflow/pool/taichi_handle.py b/pyfastflow/pool/taichi_handle.py new file mode 100644 index 0000000..ce0a820 --- /dev/null +++ b/pyfastflow/pool/taichi_handle.py @@ -0,0 +1,86 @@ +""" +Taichi backend implementation of DataHandle. + +Wraps one Taichi field allocated via FieldsBuilder for explicit lifetime +control (snodetree.destroy() actually frees GPU memory; release() just +returns it to the pool). Composition, not inheritance: kernels take the +raw field via `.data`, not the handle itself - see pool/base.py design +notes on why subclassing ti.Field was rejected. + +Author: B.G (07/2026) +""" + +from typing import Any + +import taichi as ti + +from .base import DataHandle + + +class TaichiDataHandle(DataHandle): + """ + DataHandle backed by one Taichi field. + + Author: B.G (07/2026) + """ + + _next_id = 0 + + def __init__(self, dtype: Any, shape: tuple[int, ...]): + """ + Allocate a Taichi field of the given dtype/shape via FieldsBuilder. + + shape=() allocates a 0D scalar field, indexed as field[None]. + + Author: B.G (07/2026) + """ + TaichiDataHandle._next_id += 1 + self.id = TaichiDataHandle._next_id + self.dtype = dtype + self.shape = tuple(shape) + self.in_use = False + + self._builder = ti.FieldsBuilder() + self._field = ti.field(dtype) + + if len(self.shape) == 0: + self._builder.place(self._field) + elif len(self.shape) == 1: + self._builder.dense(ti.i, self.shape).place(self._field) + elif len(self.shape) == 2: + self._builder.dense(ti.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 ti.field, for passing straight into kernels. + + 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) From 64718658e57a903d956f2888db81a967ddde0aa8 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 20 Jul 2026 13:12:22 +0200 Subject: [PATCH 03/31] Lay out the abstract classes for compiling the functions, helper/kernel/params --- pyfastflow/context/base.py | 190 +++++++++++++++++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 pyfastflow/context/base.py diff --git a/pyfastflow/context/base.py b/pyfastflow/context/base.py new file mode 100644 index 0000000..48f2d8e --- /dev/null +++ b/pyfastflow/context/base.py @@ -0,0 +1,190 @@ +""" +Backend-agnostic context building blocks. + +No "Context" class here on purpose: a context is just whatever concrete class +(GridContext, FlowContext, ...) groups a set of Parameters and registers +DeviceFunctions. Cross-context references are plain explicit bindings passed +at compile() time, not a stored connection registry. + +Author: B.G (07/2026) +""" + +from abc import ABC, abstractmethod +from typing import Any, ClassVar + +from ..pool.base import DataHandle + + +class Parameter(ABC): + """ + One named, typed value owned by a context. + + REQUIRED_MODES is the baseline every backend must support; a backend + widens SUPPORTED_MODES to add more storage kinds. Enforced at subclass + definition time via __init_subclass__, not at instantiation. + + Author: B.G (07/2026) + """ + + REQUIRED_MODES: ClassVar[frozenset[str]] = frozenset({"const", "scalar", "field"}) + SUPPORTED_MODES: ClassVar[frozenset[str]] = REQUIRED_MODES + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + missing = Parameter.REQUIRED_MODES - cls.SUPPORTED_MODES + if missing: + raise TypeError(f"{cls.__name__} must support modes {sorted(missing)}") + + name: str + dtype: Any + mode: str + + @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 parameter's value in place, according to its mode. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def destroy(self) -> None: + """ + Release any backing storage owned by this parameter. + + Author: B.G (07/2026) + """ + ... + + +class Specializable(ABC): + """ + Shared compile/call contract for DeviceFunction and (later) Kernel. + + compile() specializes a template with explicit bindings injected as + globals - the mechanism that replaces a stored connection registry. + + Author: B.G (07/2026) + """ + + name: str + + @classmethod + @abstractmethod + def compile(cls, template, *, bindings: dict[str, Any]) -> "Specializable": + """ + Specialize `template` with `bindings` injected as globals. + + Author: B.G (07/2026) + """ + ... + + @property + @abstractmethod + def compiled(self): + """ + Raw backend callable (e.g. the ti.func/ti.kernel object), for + injection as a global into another template's bindings. + + Author: B.G (07/2026) + """ + ... + + @abstractmethod + def __call__(self, *args, **kwargs): ... + + +class DeviceFunction(Specializable): + """ + Compiled device-side helper (e.g. a ti.func specialization). + + Not necessarily callable from host Python - backends where device + functions can only run inside kernel/func scope may raise on __call__. + + Author: B.G (07/2026) + """ + + +def resolve_binding(value): + """ + Unwrap a Parameter/Specializable to its backend-workable object. + + Parameter -> get() -> .data if that's a DataHandle, else the raw value. + Specializable (DeviceFunction/Kernel) -> .compiled. + Anything else passes through unchanged. + + Author: B.G (07/2026) + """ + if isinstance(value, Parameter): + resolved = value.get() + return resolved.data if isinstance(resolved, DataHandle) else resolved + if isinstance(value, Specializable): + return value.compiled + return value + + +class Bag: + """ + Simple named collection, mergeable into a Specializable.compile() bindings dict via as_bindings(). + + Author: B.G (07/2026) + """ + + def __init__(self, items: dict[str, Any] | None = None): + self._items: dict[str, Any] = dict(items or {}) + + 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 __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 as_bindings(self) -> dict[str, Any]: + """ + Return {name: item} for merging into a compile() bindings dict. + + Author: B.G (07/2026) + """ + return dict(self._items) + + +class ParamBag(Bag): + """ + Named collection of Parameter objects. + + Author: B.G (07/2026) + """ + + +class HelperBag(Bag): + """ + Named collection of DeviceFunction objects. + + Author: B.G (07/2026) + """ From 58143a92870f397443ef3f31cfd5dfa3e5d5c253 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 20 Jul 2026 16:11:13 +0200 Subject: [PATCH 04/31] Implement taichi pool --- pyfastflow/pool/base.py | 4 +-- pyfastflow/pool/taichi_pool.py | 54 ++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) create mode 100644 pyfastflow/pool/taichi_pool.py diff --git a/pyfastflow/pool/base.py b/pyfastflow/pool/base.py index 1f1596b..4539378 100644 --- a/pyfastflow/pool/base.py +++ b/pyfastflow/pool/base.py @@ -1,8 +1,8 @@ """ Backend-agnostic pool contracts. -Defines the blueprint that every pool backend (Taichi fields, numpy arrays, -future tree/sparse structures, ...) must implement. No allocation logic here +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) diff --git a/pyfastflow/pool/taichi_pool.py b/pyfastflow/pool/taichi_pool.py new file mode 100644 index 0000000..c86aea9 --- /dev/null +++ b/pyfastflow/pool/taichi_pool.py @@ -0,0 +1,54 @@ +""" +Taichi backend implementation of Pool. + +Author: B.G (07/2026) +""" + +from typing import Any + +from .base import DataHandle, Pool +from .taichi_handle import TaichiDataHandle + + +class TaichiPool(Pool): + """ + Pool manager for TaichiDataHandle, bucketed by (dtype, shape). + + Author: B.G (07/2026) + """ + + def __init__(self): + self._buckets: dict[tuple[Any, tuple[int, ...]], list[TaichiDataHandle]] = {} + + 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 = TaichiDataHandle(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} From 105f9cad2f657cabd6c7ba7183c777e818358b9b Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 20 Jul 2026 16:16:22 +0200 Subject: [PATCH 05/31] Implementing the taichi backend for the new context factory --- pyfastflow/context/taichi_backend.py | 147 +++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 pyfastflow/context/taichi_backend.py diff --git a/pyfastflow/context/taichi_backend.py b/pyfastflow/context/taichi_backend.py new file mode 100644 index 0000000..a957a2a --- /dev/null +++ b/pyfastflow/context/taichi_backend.py @@ -0,0 +1,147 @@ +""" +Taichi backend implementation of Parameter and DeviceFunction. + +Author: B.G (07/2026) +""" + +from types import FunctionType +from typing import Any + +import numpy as np +import taichi as ti + +from ..pool.base import Pool +from .base import DeviceFunction, Parameter, resolve_binding + + +def _numpy_dtype(dtype): + """ + Map a Taichi dtype to the numpy dtype used for host-side (de)serialization. + + Author: B.G (07/2026) + """ + if dtype == ti.u8: + return np.uint8 + if dtype == ti.i32: + return np.int32 + if dtype == ti.i64: + return np.int64 + return np.float32 + + +class TaichiParameter(Parameter): + """ + Parameter backed by a Taichi scalar/const value or a pooled TaichiDataHandle. + + Author: B.G (07/2026) + """ + + SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) + + def __init__(self, name: str, *, dtype, mode: str, value, pool: 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. + + Author: B.G (07/2026) + """ + if mode not in self.SUPPORTED_MODES: + raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") + + 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.set(value) + + def get(self): + """ + Author: B.G (07/2026) + """ + return self._const_value if self.mode == "const" else self._handle + + def set(self, value) -> None: + """ + Author: B.G (07/2026) + """ + if self.mode == "const": + self._const_value = _numpy_dtype(self.dtype)(value).item() + elif self.mode == "scalar": + self._handle.data[None] = value + else: # field + arr = np.asarray(value, dtype=_numpy_dtype(self.dtype)).reshape(-1) + self._handle.data.from_numpy(arr) + + def destroy(self) -> None: + """ + Author: B.G (07/2026) + """ + if self._handle is not None: + self._pool.release_data(self._handle) + self._handle = None + + +class TaichiDeviceFunction(DeviceFunction): + """ + DeviceFunction backed by a compiled ti.func. + + compile() clones the template's code object with a globals dict where + sentinels are replaced by resolved bindings, then decorates with ti.func - + same specialization mechanism as the legacy CallableFactory.compile. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @classmethod + def compile(cls, template, *, bindings: dict[str, Any]) -> "TaichiDeviceFunction": + """ + Author: B.G (07/2026) + """ + resolved = {name: resolve_binding(value) 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 cls(source.__name__, ti.func(specialised)) + + @property + def compiled(self): + """ + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, **kwargs): + """ + A ti.func cannot be invoked from host Python - it only runs inside + kernel/func scope, where callers use `.compiled` directly. + + Author: B.G (07/2026) + """ + raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") From 9b90f7c47ee7871894144d5b1acdb7cba98ee10d Mon Sep 17 00:00:00 2001 From: bgailleton Date: Tue, 21 Jul 2026 10:42:47 +0200 Subject: [PATCH 06/31] prepare the new file structure for the module. Separate core form the rest of the package --- src/core/context/__init__.py | 9 + {pyfastflow => src/core}/context/base.py | 13 ++ src/core/context/quadrants_backend.py | 208 ++++++++++++++++++ .../core}/context/taichi_backend.py | 91 ++++++-- src/core/pool/__init__.py | 9 + {pyfastflow => src/core}/pool/base.py | 0 src/core/pool/quadrants_handle.py | 86 ++++++++ src/core/pool/quadrants_pool.py | 54 +++++ .../core}/pool/taichi_handle.py | 0 {pyfastflow => src/core}/pool/taichi_pool.py | 0 10 files changed, 450 insertions(+), 20 deletions(-) create mode 100644 src/core/context/__init__.py rename {pyfastflow => src/core}/context/base.py (92%) create mode 100644 src/core/context/quadrants_backend.py rename {pyfastflow => src/core}/context/taichi_backend.py (60%) create mode 100644 src/core/pool/__init__.py rename {pyfastflow => src/core}/pool/base.py (100%) create mode 100644 src/core/pool/quadrants_handle.py create mode 100644 src/core/pool/quadrants_pool.py rename {pyfastflow => src/core}/pool/taichi_handle.py (100%) rename {pyfastflow => src/core}/pool/taichi_pool.py (100%) diff --git a/src/core/context/__init__.py b/src/core/context/__init__.py new file mode 100644 index 0000000..366edf6 --- /dev/null +++ b/src/core/context/__init__.py @@ -0,0 +1,9 @@ +""" +New backend-agnostic context architecture (Parameter/Specializable ABCs + backends). + +Physically lives under ./src/core/context, exposed at import time as +pyfastflow.experimental.core.context via the path shim in +pyfastflow/experimental/core/__init__.py. + +Author: B.G (07/2026) +""" diff --git a/pyfastflow/context/base.py b/src/core/context/base.py similarity index 92% rename from pyfastflow/context/base.py rename to src/core/context/base.py index 48f2d8e..6674c9c 100644 --- a/pyfastflow/context/base.py +++ b/src/core/context/base.py @@ -115,6 +115,19 @@ class DeviceFunction(Specializable): """ +class Kernel(Specializable): + """ + Compiled entry point (e.g. a ti.kernel specialization). + + Unlike DeviceFunction, __call__ is expected to work from host Python - + that's how the compute actually gets launched. Its template's own + parameters are data fields only; params/helpers arrive via bindings + and are resolved into the kernel body, not passed at call time. + + Author: B.G (07/2026) + """ + + def resolve_binding(value): """ Unwrap a Parameter/Specializable to its backend-workable object. diff --git a/src/core/context/quadrants_backend.py b/src/core/context/quadrants_backend.py new file mode 100644 index 0000000..f00ca23 --- /dev/null +++ b/src/core/context/quadrants_backend.py @@ -0,0 +1,208 @@ +""" +Quadrants backend implementation of Parameter, DeviceFunction and Kernel. + +Mirrors taichi_backend.py; the specialization mechanism (clone template code +object, patch globals, decorate) is unchanged between the two - verified +empirically before writing this file. One addition specific to this backend: +Kernel templates may type their own data-field arguments as `qd.Tensor`, +which accepts either a field- or ndarray-backed value at call time with no +change to the compiled template - Taichi has no equivalent for this. + +Author: B.G (07/2026) +""" + +from types import FunctionType +from typing import Any + +import numpy as np +import quadrants as qd + +from ..pool.base import Pool +from .base import DeviceFunction, Kernel, Parameter, resolve_binding + + +def _numpy_dtype(dtype): + """ + Map a Quadrants dtype to the numpy dtype used for host-side (de)serialization. + + Author: B.G (07/2026) + """ + if dtype == qd.u8: + return np.uint8 + if dtype == qd.i32: + return np.int32 + if dtype == qd.i64: + return np.int64 + return np.float32 + + +class QuadrantsParameter(Parameter): + """ + Parameter backed by a Quadrants scalar/const value or a pooled QuadrantsDataHandle. + + Author: B.G (07/2026) + """ + + SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) + + def __init__(self, name: str, *, dtype, mode: str, value, pool: 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. + + Author: B.G (07/2026) + """ + if mode not in self.SUPPORTED_MODES: + raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") + + 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.set(value) + + def get(self): + """ + Author: B.G (07/2026) + """ + return self._const_value if self.mode == "const" else self._handle + + def set(self, value) -> None: + """ + Author: B.G (07/2026) + """ + if self.mode == "const": + self._const_value = _numpy_dtype(self.dtype)(value).item() + elif self.mode == "scalar": + self._handle.data[None] = value + else: # field + arr = np.asarray(value, dtype=_numpy_dtype(self.dtype)).reshape(-1) + self._handle.data.from_numpy(arr) + + def destroy(self) -> None: + """ + Author: B.G (07/2026) + """ + if self._handle is not None: + self._pool.release_data(self._handle) + self._handle = None + + +def _specialize(template, bindings: dict[str, Any]): + """ + Clone `template`'s code object with a globals dict where sentinels are + replaced by resolved bindings. Shared by DeviceFunction and Kernel - + they only differ in the qd.func/qd.kernel decorator applied on top. + + Author: B.G (07/2026) + """ + resolved = {name: resolve_binding(value) 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 QuadrantsDeviceFunction(DeviceFunction): + """ + DeviceFunction backed by a compiled qd.func. + + Only field-backed Parameters/DeviceFunctions can be resolved into a + template this way - Quadrants rejects ndarrays referenced as globals + (verified: raises "Ndarray used in kernel scope but not registered as + a kernel parameter" even for the unwrapped object). + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @classmethod + def compile(cls, template, *, bindings: dict[str, Any]) -> "QuadrantsDeviceFunction": + """ + Author: B.G (07/2026) + """ + specialised = _specialize(template, bindings) + return cls(specialised.__name__, qd.func(specialised)) + + @property + def compiled(self): + """ + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, **kwargs): + """ + A qd.func cannot be invoked from host Python - it only runs inside + kernel/func scope, where callers use `.compiled` directly. + + Author: B.G (07/2026) + """ + raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") + + +class QuadrantsKernel(Kernel): + """ + Kernel backed by a compiled qd.kernel. + + The template's own data-field arguments should be typed `qd.Tensor` - + unlike Taichi's ti.template()/ti.types.ndarray() split, a single + qd.Tensor-typed template accepts either a field- or ndarray-backed + value at call time (verified). Params/helpers still arrive via + bindings, resolved into the kernel body - callers only ever pass data + fields at call time. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @classmethod + def compile(cls, template, *, bindings: dict[str, Any]) -> "QuadrantsKernel": + """ + Author: B.G (07/2026) + """ + specialised = _specialize(template, bindings) + return cls(specialised.__name__, qd.kernel(specialised)) + + @property + def compiled(self): + """ + 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) diff --git a/pyfastflow/context/taichi_backend.py b/src/core/context/taichi_backend.py similarity index 60% rename from pyfastflow/context/taichi_backend.py rename to src/core/context/taichi_backend.py index a957a2a..3f3e6da 100644 --- a/pyfastflow/context/taichi_backend.py +++ b/src/core/context/taichi_backend.py @@ -1,5 +1,5 @@ """ -Taichi backend implementation of Parameter and DeviceFunction. +Taichi backend implementation of Parameter, DeviceFunction and Kernel. Author: B.G (07/2026) """ @@ -11,7 +11,34 @@ import taichi as ti from ..pool.base import Pool -from .base import DeviceFunction, Parameter, resolve_binding +from .base import DeviceFunction, Kernel, Parameter, resolve_binding + + +def _specialize(template, bindings: dict[str, Any]): + """ + Clone `template`'s code object with a globals dict where sentinels are + replaced by resolved bindings. Shared by DeviceFunction and Kernel - + they only differ in the ti.func/ti.kernel decorator applied on top. + + Author: B.G (07/2026) + """ + resolved = {name: resolve_binding(value) 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 def _numpy_dtype(dtype): @@ -111,24 +138,8 @@ def compile(cls, template, *, bindings: dict[str, Any]) -> "TaichiDeviceFunction """ Author: B.G (07/2026) """ - resolved = {name: resolve_binding(value) 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 cls(source.__name__, ti.func(specialised)) + specialised = _specialize(template, bindings) + return cls(specialised.__name__, ti.func(specialised)) @property def compiled(self): @@ -145,3 +156,43 @@ def __call__(self, *args, **kwargs): Author: B.G (07/2026) """ raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") + + +class TaichiKernel(Kernel): + """ + Kernel backed by a compiled ti.kernel. + + The template's own signature should only declare data-field arguments + (e.g. `def template(out: ti.template()): ...`) - any params/helpers it + needs are resolved into the kernel body via `bindings`, so callers only + ever pass data fields at call time, never params/helpers. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @classmethod + def compile(cls, template, *, bindings: dict[str, Any]) -> "TaichiKernel": + """ + Author: B.G (07/2026) + """ + specialised = _specialize(template, bindings) + return cls(specialised.__name__, ti.kernel(specialised)) + + @property + def compiled(self): + """ + 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) diff --git a/src/core/pool/__init__.py b/src/core/pool/__init__.py new file mode 100644 index 0000000..2772901 --- /dev/null +++ b/src/core/pool/__init__.py @@ -0,0 +1,9 @@ +""" +New backend-agnostic pool architecture (DataHandle/Pool ABCs + backends). + +Physically lives under ./src/core/pool, exposed at import time as +pyfastflow.experimental.core.pool via the path shim in +pyfastflow/experimental/core/__init__.py. + +Author: B.G (07/2026) +""" diff --git a/pyfastflow/pool/base.py b/src/core/pool/base.py similarity index 100% rename from pyfastflow/pool/base.py rename to src/core/pool/base.py diff --git a/src/core/pool/quadrants_handle.py b/src/core/pool/quadrants_handle.py new file mode 100644 index 0000000..aaaced8 --- /dev/null +++ b/src/core/pool/quadrants_handle.py @@ -0,0 +1,86 @@ +""" +Quadrants backend implementation of DataHandle. + +Mirrors taichi_handle.py - Quadrants' FieldsBuilder/SNodeTree lifecycle is +identical to Taichi's. Composition, not inheritance, same rationale as the +Taichi backend: kernels take the raw field via `.data`. + +Author: B.G (07/2026) +""" + +from typing import Any + +import quadrants as qd + +from .base import DataHandle + + +class QuadrantsDataHandle(DataHandle): + """ + DataHandle backed by one Quadrants field. + + Author: B.G (07/2026) + """ + + _next_id = 0 + + def __init__(self, dtype: Any, shape: tuple[int, ...]): + """ + Allocate a Quadrants field of the given dtype/shape via FieldsBuilder. + + shape=() allocates a 0D scalar field, indexed as field[None]. + + Author: B.G (07/2026) + """ + QuadrantsDataHandle._next_id += 1 + self.id = QuadrantsDataHandle._next_id + self.dtype = dtype + self.shape = tuple(shape) + self.in_use = False + + self._builder = qd.FieldsBuilder() + self._field = qd.field(dtype) + + if len(self.shape) == 0: + self._builder.place(self._field) + elif len(self.shape) == 1: + self._builder.dense(qd.i, self.shape).place(self._field) + elif len(self.shape) == 2: + self._builder.dense(qd.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 qd.field, for passing straight into kernels + or binding as a global (ndarrays cannot be bound as globals - + Parameter/DeviceFunction wiring stays field-backed for that reason). + + 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/src/core/pool/quadrants_pool.py b/src/core/pool/quadrants_pool.py new file mode 100644 index 0000000..ea0e6a8 --- /dev/null +++ b/src/core/pool/quadrants_pool.py @@ -0,0 +1,54 @@ +""" +Quadrants backend implementation of Pool. + +Author: B.G (07/2026) +""" + +from typing import Any + +from .base import DataHandle, Pool +from .quadrants_handle import QuadrantsDataHandle + + +class QuadrantsPool(Pool): + """ + Pool manager for QuadrantsDataHandle, bucketed by (dtype, shape). + + Author: B.G (07/2026) + """ + + def __init__(self): + self._buckets: dict[tuple[Any, tuple[int, ...]], list[QuadrantsDataHandle]] = {} + + 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 = QuadrantsDataHandle(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/pool/taichi_handle.py b/src/core/pool/taichi_handle.py similarity index 100% rename from pyfastflow/pool/taichi_handle.py rename to src/core/pool/taichi_handle.py diff --git a/pyfastflow/pool/taichi_pool.py b/src/core/pool/taichi_pool.py similarity index 100% rename from pyfastflow/pool/taichi_pool.py rename to src/core/pool/taichi_pool.py From ed7a8f040b7a77481dbde70c083ff7a7d432e095 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Wed, 22 Jul 2026 16:38:27 +0200 Subject: [PATCH 07/31] Massive refactore of the experimental core. Getting toward stable api --- pyfastflow/experimental/__init__.py | 5 + pyfastflow/experimental/core/__init__.py | 5 + .../experimental/core/context/__init__.py | 5 + .../core/context/_closure_backend.py | 150 ++++++ pyfastflow/experimental/core/context/base.py | 349 ++++++++++++++ .../experimental/core/context/cupy_backend.py | 437 ++++++++++++++++++ .../core/context/quadrants_backend.py | 180 ++++++++ .../core/context/taichi_backend.py | 171 +++++++ pyfastflow/experimental/core/pool/__init__.py | 5 + .../experimental/core/pool/_bucketed_pool.py | 16 +- .../experimental/core/pool/_fields_handle.py | 41 +- .../experimental}/core/pool/base.py | 0 .../experimental/core/pool/cupy_handle.py | 65 +++ .../experimental/core/pool/cupy_pool.py | 18 + .../core/pool/quadrants_handle.py | 19 + .../experimental/core/pool/quadrants_pool.py | 18 + .../experimental/core/pool/taichi_handle.py | 19 + .../experimental/core/pool/taichi_pool.py | 18 + src/core/context/__init__.py | 9 - src/core/context/base.py | 203 -------- src/core/context/quadrants_backend.py | 208 --------- src/core/context/taichi_backend.py | 198 -------- src/core/pool/__init__.py | 9 - src/core/pool/taichi_handle.py | 86 ---- src/core/pool/taichi_pool.py | 54 --- 25 files changed, 1495 insertions(+), 793 deletions(-) create mode 100644 pyfastflow/experimental/__init__.py create mode 100644 pyfastflow/experimental/core/__init__.py create mode 100644 pyfastflow/experimental/core/context/__init__.py create mode 100644 pyfastflow/experimental/core/context/_closure_backend.py create mode 100644 pyfastflow/experimental/core/context/base.py create mode 100644 pyfastflow/experimental/core/context/cupy_backend.py create mode 100644 pyfastflow/experimental/core/context/quadrants_backend.py create mode 100644 pyfastflow/experimental/core/context/taichi_backend.py create mode 100644 pyfastflow/experimental/core/pool/__init__.py rename src/core/pool/quadrants_pool.py => pyfastflow/experimental/core/pool/_bucketed_pool.py (74%) rename src/core/pool/quadrants_handle.py => pyfastflow/experimental/core/pool/_fields_handle.py (56%) rename {src => pyfastflow/experimental}/core/pool/base.py (100%) create mode 100644 pyfastflow/experimental/core/pool/cupy_handle.py create mode 100644 pyfastflow/experimental/core/pool/cupy_pool.py create mode 100644 pyfastflow/experimental/core/pool/quadrants_handle.py create mode 100644 pyfastflow/experimental/core/pool/quadrants_pool.py create mode 100644 pyfastflow/experimental/core/pool/taichi_handle.py create mode 100644 pyfastflow/experimental/core/pool/taichi_pool.py delete mode 100644 src/core/context/__init__.py delete mode 100644 src/core/context/base.py delete mode 100644 src/core/context/quadrants_backend.py delete mode 100644 src/core/context/taichi_backend.py delete mode 100644 src/core/pool/__init__.py delete mode 100644 src/core/pool/taichi_handle.py delete mode 100644 src/core/pool/taichi_pool.py 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..af91431 --- /dev/null +++ b/pyfastflow/experimental/core/__init__.py @@ -0,0 +1,5 @@ +""" +New backend-agnostic core (Parameter/DeviceFunction/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..fcbb697 --- /dev/null +++ b/pyfastflow/experimental/core/context/__init__.py @@ -0,0 +1,5 @@ +""" +New backend-agnostic context architecture (Parameter/Specializable ABCs + backends). + +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..20de1ad --- /dev/null +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -0,0 +1,150 @@ +""" +Shared machinery for backends that compile python function templates by +patching globals on a cloned code object (Taichi, Quadrants). Not used by +backends without that mechanism (e.g. the cupy/RawKernel backend). + +Author: B.G (07/2026) +""" + +from types import FunctionType +from typing import Any + +import numpy as np + +from .base import Parameter, resolve_binding + + +def specialize_closure(template, bindings: dict[str, Any]) -> FunctionType: + """ + Clone `template`'s code object with a globals dict where sentinels are + replaced by resolved bindings. + + Author: B.G (07/2026) + """ + resolved = {name: resolve_binding(value) 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: + """ + Device-facing view of a Parameter for closure backends: `.get` and (unless + const) `.set_node` are the raw compiled device funcs (ti.func / qd.func), + so a template body traces `p.get(i)` / `p.set_node(i, v)` as plain + attribute lookups + func calls, uniform across modes. Const params expose + no `.set_node` (read-only), so touching it in a kernel raises at trace time. + + 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 value or a pooled DataHandle, for backends + sharing the closure-specialization mechanism. Subclasses pin `_numpy_dtype` + and implement `device_view` with their own func decorator. + + Author: B.G (07/2026) + """ + + SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) + + def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): + """ + Declare and initialize one parameter. "scalar"/"field" modes allocate + pooled storage immediately via `pool`; "const" stays a plain python + value. solo=True (const only) lets the parameter be read bare in a + template body (no .get()) - it resolves to a compile-time literal. + + Author: B.G (07/2026) + """ + if mode not in self.SUPPORTED_MODES: + raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") + if solo and mode != "const": + raise ValueError(f"{name}: solo access is const-only, got mode {mode!r}") + + self.name = name + self.dtype = dtype + self.mode = mode + self.solo = solo + 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.set(value) + + @staticmethod + def _numpy_dtype(dtype): + """ + Map a backend dtype to the numpy dtype used for host-side (de)serialization. + Overridden per backend. + + Author: B.G (07/2026) + """ + raise NotImplementedError + + def get(self): + """ + Author: B.G (07/2026) + """ + return self._const_value if self.mode == "const" else self._handle + + def set(self, value) -> None: + """ + 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: + """ + Author: B.G (07/2026) + """ + if self._handle is not None: + self._pool.release_data(self._handle) + self._handle = None diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py new file mode 100644 index 0000000..732c509 --- /dev/null +++ b/pyfastflow/experimental/core/context/base.py @@ -0,0 +1,349 @@ +""" +Backend-agnostic context building blocks. + +No "Context" class here on purpose: a context is just whatever concrete class +(GridContext, FlowContext, ...) groups a set of Parameters and registers +DeviceFunctions. Cross-context references are plain explicit bindings passed +at bind() time, not a stored connection registry. + +The compile surface is a two-layer builder: an abstract DeviceFunctionBuilder / +KernelBuilder here, backend variants (Taichi*/Quadrants*/Cupy*) elsewhere. A +builder collects bind()ed dependencies + one ingest()ed template and produces a +compiled DeviceFunction/Kernel. Bound Parameters/helpers/bags are injected into +the template body (never passed at call time); only a template's own explicit +data-field arguments are passed to the front callable. + +Author: B.G (07/2026) +""" + +import ast +import inspect +from abc import ABC, abstractmethod +from types import SimpleNamespace +from typing import Any, ClassVar + +from ..pool.base import DataHandle + + +class Parameter(ABC): + """ + One named, typed value owned by a context. + + REQUIRED_MODES is the baseline every backend must support; a backend + widens SUPPORTED_MODES to add more storage kinds. Enforced at subclass + definition time via __init_subclass__, not at instantiation. + + Host surface: get() / set(value) / set_node(node, value). Device surface: + device_view(), returning a backend object whose .get(node) / .set_node( + node, val) are usable from inside device code - the uniform accessor that + lets one kernel read/write a Parameter identically no matter its mode. + + Author: B.G (07/2026) + """ + + REQUIRED_MODES: ClassVar[frozenset[str]] = frozenset({"const", "scalar", "field"}) + SUPPORTED_MODES: ClassVar[frozenset[str]] = REQUIRED_MODES + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + missing = Parameter.REQUIRED_MODES - cls.SUPPORTED_MODES + if missing: + raise TypeError(f"{cls.__name__} must support modes {sorted(missing)}") + + name: str + dtype: Any + mode: str + solo: bool = False + + @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. + + 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): + """ + Return a backend device-view object (get/set_node usable in device + code). Implemented per backend; closure backends compile ti/qd funcs, + cupy returns parser metadata. + + 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. + + Author: B.G (07/2026) + """ + ... + + +class Specializable(ABC): + """ + Shared call contract for DeviceFunction and Kernel, plus the hidden raw + material (source text, AST, dependency manifest) a later fusion/graph + compiler can reuse. Built by a *Builder, not by a classmethod. + + Author: B.G (07/2026) + """ + + name: str + _source: str | None = None + _ast: ast.AST | None = None + _dependencies: dict[str, str] | None = None + + @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 DeviceFunction(Specializable): + """ + Compiled device-side helper (e.g. a ti.func specialization). + + Not necessarily callable from host Python - backends where device + functions 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 DeviceFunction, __call__ works from host Python - that's how the + compute gets launched. Its template's own parameters are data fields only; + params/helpers arrive via bindings and are resolved into the kernel body. + + Author: B.G (07/2026) + """ + + +def resolve_binding(value): + """ + Unwrap a bound object to what a closure-backend template body should see. + + Parameter -> literal if solo (const), else device_view() (a .get/.set_node + carrier for uniform in-kernel access). + Specializable (DeviceFunction/Kernel) -> .compiled. + Bag -> a namespace whose attributes are each member resolved the same way, + so dotted paths (grid.nx.get(i)) trace as plain attribute lookups. + Anything else passes through unchanged. + + Author: B.G (07/2026) + """ + if isinstance(value, Parameter): + if getattr(value, "solo", False): + resolved = value.get() + return resolved.data if isinstance(resolved, DataHandle) else resolved + return value.device_view() + if isinstance(value, Specializable): + return value.compiled + if isinstance(value, Bag): + return SimpleNamespace(**{name: resolve_binding(item) for name, item in value.items()}) + return value + + +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 with no AST (not python). + + 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 attach_meta(obj: Specializable, template, bindings: dict[str, Any]) -> None: + """ + Stash hidden raw material on a freshly built Specializable for later reuse + by a higher-level compiler (kernel fusion / graph building). + + Author: B.G (07/2026) + """ + obj._source, obj._ast = capture_template_meta(template) + obj._dependencies = {name: type(value).__name__ for name, value in bindings.items()} + + +class CompileBuilder(ABC): + """ + Two-layer compile surface: collect bind()ed dependencies + one ingest()ed + template, then compile() to a backend Specializable. bind() detects the + kind of each object at resolution time (Parameter / DeviceFunction / Bag / + handle / plain value); backends implement compile() (and may override + ingest()) their own way. + + Author: B.G (07/2026) + """ + + def __init__(self): + self._bindings: dict[str, Any] = {} + self._template = None + + def bind(self, name: str, obj: Any) -> "CompileBuilder": + """ + Register `obj` under `name` for injection into the template body. + Raises if `name` is already bound. + + Author: B.G (07/2026) + """ + if name in self._bindings: + raise KeyError(f"'{name}' is already bound") + 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) + 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 + + @abstractmethod + def compile(self) -> Specializable: + """ + Produce the compiled DeviceFunction/Kernel. + + Author: B.G (07/2026) + """ + ... + + +class DeviceFunctionBuilder(CompileBuilder): + """ + Builds a DeviceFunction. compile() -> DeviceFunction. + + Author: B.G (07/2026) + """ + + +class KernelBuilder(CompileBuilder): + """ + Builds a Kernel. compile() -> host-callable Kernel. + + Author: B.G (07/2026) + """ + + +class Bag: + """ + Simple named collection, mergeable into a builder via bind_bag() or bound + whole (bind('grid', bag)) for dotted-path access in the template body. + + Author: B.G (07/2026) + """ + + def __init__(self, items: dict[str, Any] | None = None): + self._items: dict[str, Any] = dict(items or {}) + + def add(self, name: str, item: Any) -> None: + """ + Register `item` under `name`. Raises if `name` is already registered. + + Author: B.G (07/2026) + """ + if name in self._items: + raise KeyError(f"'{name}' is already registered in this bag") + self._items[name] = item + + def __getattr__(self, name: str) -> Any: + try: + return self._items[name] + except KeyError: + raise AttributeError(name) + + def __getitem__(self, name: str) -> Any: + return self._items[name] + + def __contains__(self, name: str) -> bool: + return name in self._items + + def __iter__(self): + return iter(self._items) + + def items(self): + return self._items.items() + + def as_bindings(self) -> dict[str, Any]: + """ + Return {name: item} for merging into a builder's bindings. + + Author: B.G (07/2026) + """ + return dict(self._items) + + +class ParamBag(Bag): + """ + Named collection of Parameter objects. + + Author: B.G (07/2026) + """ + + +class HelperBag(Bag): + """ + Named collection of DeviceFunction objects. + + 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..71b654e --- /dev/null +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -0,0 +1,437 @@ +""" +Cupy backend implementation of Parameter, DeviceFunction, Kernel and their +builders. + +Templates are raw CUDA source text, not python callables - there is no +closure mechanism for cp.RawKernel. Params/helpers are referenced inside the +source through `$...$` spans holding a dotted path, uniform with the closure +backends' in-kernel API: + + $p.get(i)$ -> read param p at flat index i + $p.set_node(i,v)$ -> write param p at flat index i (device-side) + $grid.nx.get(i)$ -> bag member access (dotted head) + $helper(a, b)$ -> call a bound CupyDeviceFunction + +The parser expands each span AND, for scalar/field params, auto-generates the +matching pointer argument into the __global__ signature plus the launch-time +array - the source never hand-declares those. Expansion by mode: + - const -> a CUDA literal (and a `#define NAME literal` for bare use of a + top-level const param outside any span) + - scalar -> NAME[0] (a 0-d device pointer) + - field -> NAME[i] +set_node on a const is a compile error; a param read only -> `const T*`, a +param written anywhere in the kernel -> non-const `T*`. + +Device functions may only reference const params / other device functions +inside spans - a spliced __device__ function has no way to receive an extra +pointer argument. + +Author: B.G (07/2026) +""" + +import re +from typing import Any + +import cupy as cp +import numpy as np + +from .base import DeviceFunction, DeviceFunctionBuilder, Kernel, KernelBuilder, Parameter, attach_meta +from .base import Bag + +_KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") +_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", +} + + +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: + """ + 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 + + +class _SpanParser: + """ + Expands `$...$` spans in a CUDA template body, accumulating the pointer + params and helper sources they imply. allow_arrays=False (device + functions) rejects scalar/field params and set_node. + + Author: B.G (07/2026) + """ + + def __init__(self, bindings: dict[str, Any], *, allow_arrays: bool): + self.bindings = bindings + self.allow_arrays = allow_arrays + self.ptr_params: dict[str, dict] = {} # argname -> {ctype, write, array} + self.helper_srcs: dict[str, str] = {} # name -> source + + def _register_ptr(self, argname: str, param: Parameter, write: bool) -> None: + entry = self.ptr_params.get(argname) + if entry is None: + self.ptr_params[argname] = {"ctype": _ctype(param.dtype), "write": write, "array": param.get().data} + elif write: + entry["write"] = True + + def _expand_param(self, param: Parameter, method: str, argname: str, call_args: list[str]) -> str: + if method == "get": + if param.mode == "const": + return _cuda_literal(param.get()) + if not self.allow_arrays: + raise ValueError(f"{param.name}: scalar/field param cannot be used in a device function (no pointer arg)") + self._register_ptr(argname, 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 not self.allow_arrays: + raise ValueError(f"{param.name}: set_node cannot be used in a device function (no pointer arg)") + if len(call_args) != 2: + raise ValueError(f"{param.name}: set_node(node, value) takes two arguments") + self._register_ptr(argname, 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], "_".join(path[:-1]), call_args) + + target = _walk(path, self.bindings) + if isinstance(target, CupyDeviceFunction): + self.helper_srcs[target.name] = target.compiled + return f"{target.name}({argstr if argstr is not None else ''})" + if isinstance(target, Parameter): + return self._expand_param(target, "get", "_".join(path), ["0"]) + return _cuda_literal(target) + + def parse(self, body: str) -> str: + """ + Author: B.G (07/2026) + """ + return _SPAN_RE.sub(self._repl, body) + + +def _const_defines(bindings: dict[str, Any]) -> list[str]: + """ + `#define NAME literal` for each top-level const Parameter, so bare uses of + a const param (outside any span) compile. + + 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" + ] + + +def _inject_signature(kernel_src: str, ptr_params: dict[str, dict]) -> str: + """ + Append the generated pointer params to the __global__ signature. + + Author: B.G (07/2026) + """ + if not ptr_params: + return kernel_src + decls = [ + f"{'' if e['write'] else 'const '}{e['ctype']}* {name}" + for name, e in ptr_params.items() + ] + joined = ", ".join(decls) + + def _sub(m: re.Match) -> str: + existing = m.group(2).strip() + sep = ", " if existing else "" + return f"{m.group(1)}{m.group(2)}{sep}{joined}{m.group(3)}" + + return _KERNEL_SIG_RE.sub(_sub, kernel_src, count=1) + + +class CupyParameter(Parameter): + """ + Parameter backed by a const python value or a pooled CupyDataHandle. + + Cupy dtypes are numpy dtypes already, so no dtype-mapping hook is needed. + Resolved into templates by the `$...$` parser (not device_view). + + Author: B.G (07/2026) + """ + + SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) + + def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): + """ + Author: B.G (07/2026) + """ + if mode not in self.SUPPORTED_MODES: + raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") + if solo and mode != "const": + raise ValueError(f"{name}: solo access is const-only, got mode {mode!r}") + + self.name = name + self.dtype = dtype + self.mode = mode + self.solo = solo + 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.set(value) + + def get(self): + """ + Author: B.G (07/2026) + """ + return self._const_value if self.mode == "const" else self._handle + + def set(self, value) -> None: + """ + 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: + """ + Author: B.G (07/2026) + """ + if self._handle is not None: + self._pool.release_data(self._handle) + self._handle = None + + +class CupyDeviceFunction(DeviceFunction): + """ + DeviceFunction backed by a CUDA `__device__` function's source text. + + There's no separately-compiled device function in the RawKernel model - + `.compiled` returns source, spliced into whatever kernel/device function + binds it. Never callable from host Python. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, source: str): + self.name = name + # note: distinct from Specializable._source (the raw template, set by + # attach_meta) - this is the spliced __device__ source `.compiled` serves. + self._compiled_source = source + + @property + def compiled(self): + """ + Author: B.G (07/2026) + """ + return self._compiled_source + + def __call__(self, *args, **kwargs): + """ + Author: B.G (07/2026) + """ + raise RuntimeError( + f"DeviceFunction '{self.name}' is CUDA source, only callable from a compiled kernel's device code" + ) + + +class CupyKernel(Kernel): + """ + Kernel backed by a compiled cp.RawKernel. + + __call__ requires explicit grid/block launch dims - cp.RawKernel has no + auto-ranging. Call-time positional args are the kernel's own data-field + arguments; the parser-generated pointer arrays are appended after them. + + compile() caches the RawKernel by final spliced source text. + + Author: B.G (07/2026) + """ + + _raw_cache: dict[str, "cp.RawKernel"] = {} + + def __init__(self, name: str, compiled, bound_arrays: list): + self.name = name + self._compiled = compiled + self._bound_arrays = bound_arrays + + @property + def compiled(self): + """ + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, grid, block, **kwargs): + """ + Launches the compiled kernel. `grid`/`block` are int or tuple launch + dims, required since RawKernel 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) + tuple(self._bound_arrays)) + + +class CupyDeviceFunctionBuilder(DeviceFunctionBuilder): + """ + Builds a CupyDeviceFunction from CUDA `__device__` source. Spans may only + reference const params / other device functions (no pointer args). + + Author: B.G (07/2026) + """ + + def compile(self) -> CupyDeviceFunction: + """ + Author: B.G (07/2026) + """ + template = self._template + name = _extract_name(_DEVICE_NAME_RE, template, "__device__") + parser = _SpanParser(self._bindings, allow_arrays=False) + body = parser.parse(template) + source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings) + [body]) + fn = CupyDeviceFunction(name, source) + attach_meta(fn, template, self._bindings) + return fn + + +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: + """ + Author: B.G (07/2026) + """ + template = self._template + name = _extract_name(_KERNEL_NAME_RE, template, "__global__") + parser = _SpanParser(self._bindings, allow_arrays=True) + body = parser.parse(template) + body = _inject_signature(body, parser.ptr_params) + # extern "C" linkage so cp.RawKernel 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) + source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings) + [body]) + bound_arrays = [e["array"] for e in parser.ptr_params.values()] + + raw = CupyKernel._raw_cache.get(source) + if raw is None: + raw = cp.RawKernel(source, name) + CupyKernel._raw_cache[source] = raw + krn = CupyKernel(name, raw, bound_arrays) + krn._final_source = source + attach_meta(krn, template, self._bindings) + return krn diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py new file mode 100644 index 0000000..d6b4698 --- /dev/null +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -0,0 +1,180 @@ +""" +Quadrants backend implementation of Parameter, DeviceFunction, Kernel and +their builders. + +Mirrors taichi_backend.py; both share the closure-specialization mechanism +via _closure_backend.py. One difference specific to this backend: Kernel +templates may type their own data-field arguments as `qd.Tensor`, which +accepts either a field- or ndarray-backed value at call time with no change +to the compiled template - Taichi has no equivalent for this. Note that +field-mode Parameters must be field-backed (they close over the field as a +global): Quadrants rejects ndarrays referenced as globals inside a func. + +Author: B.G (07/2026) +""" + +import numpy as np +import quadrants as qd + +from ._closure_backend import ClosureBackendParameter, ClosureParamDeviceView, specialize_closure +from .base import DeviceFunction, DeviceFunctionBuilder, Kernel, KernelBuilder, attach_meta + + +class QuadrantsParameter(ClosureBackendParameter): + """ + Parameter backed by a Quadrants const value or a pooled QuadrantsDataHandle. + + Author: B.G (07/2026) + """ + + @staticmethod + def _numpy_dtype(dtype): + """ + Map a Quadrants dtype to the numpy dtype used for host-side (de)serialization. + + Author: B.G (07/2026) + """ + if dtype == qd.u8: + return np.uint8 + if dtype == qd.i32: + return np.int32 + if dtype == qd.i64: + return np.int64 + return np.float32 + + def device_view(self) -> ClosureParamDeviceView: + """ + Compile this parameter's uniform device accessors as qd.funcs. + + get(node) dispatches on mode at qd.func trace time via qd.static, so + only the taken arm compiles: const returns a baked literal, scalar + reads HANDLE[None], field reads HANDLE[node]. set_node(node, val) is + built only for scalar/field (const is read-only, exposes no setter). + MODE/VALUE/HANDLE are plain python values, not Quadrants values. + + Author: B.G (07/2026) + """ + mode = self.mode + value = self._const_value + handle = self._handle.data if self._handle is not None else None + + def get_template(node): + if qd.static(MODE == "const"): + return VALUE + elif qd.static(MODE == "scalar"): + return HANDLE[None] + else: + return HANDLE[node] + + get_fn = qd.func(specialize_closure(get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle})) + + set_fn = None + if mode != "const": + + def set_node_template(node, val): + if qd.static(MODE == "scalar"): + HANDLE[None] = val + else: + HANDLE[node] = val + + set_fn = qd.func(specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle})) + + return ClosureParamDeviceView(self.name, get_fn, set_fn) + + +class QuadrantsDeviceFunction(DeviceFunction): + """ + DeviceFunction backed by a compiled qd.func. Built by QuadrantsDeviceFunctionBuilder. + + Only field-backed Parameters/DeviceFunctions can be resolved into a + template this way - Quadrants rejects ndarrays referenced as globals. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, **kwargs): + """ + A qd.func only runs inside kernel/func scope; callers use `.compiled` there. + + Author: B.G (07/2026) + """ + raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") + + +class QuadrantsKernel(Kernel): + """ + Kernel backed by a compiled qd.kernel. Built by QuadrantsKernelBuilder. + + The template's own data-field arguments should be typed `qd.Tensor` - a + single qd.Tensor-typed template accepts either a field- or ndarray-backed + value at call time. Params/helpers arrive via bind() and are resolved into + the kernel body; callers pass data fields only. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + 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 QuadrantsDeviceFunctionBuilder(DeviceFunctionBuilder): + """ + Builds a QuadrantsDeviceFunction: specialize the ingested def with bound + globals, decorate with qd.func. + + Author: B.G (07/2026) + """ + + def compile(self) -> QuadrantsDeviceFunction: + """ + Author: B.G (07/2026) + """ + specialised = specialize_closure(self._template, self._bindings) + fn = QuadrantsDeviceFunction(specialised.__name__, qd.func(specialised)) + attach_meta(fn, self._template, self._bindings) + return fn + + +class QuadrantsKernelBuilder(KernelBuilder): + """ + Builds a QuadrantsKernel: specialize the ingested def with bound globals, + decorate with qd.kernel. + + Author: B.G (07/2026) + """ + + def compile(self) -> QuadrantsKernel: + """ + Author: B.G (07/2026) + """ + specialised = specialize_closure(self._template, self._bindings) + krn = QuadrantsKernel(specialised.__name__, qd.kernel(specialised)) + attach_meta(krn, self._template, self._bindings) + return krn diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py new file mode 100644 index 0000000..796df0e --- /dev/null +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -0,0 +1,171 @@ +""" +Taichi backend implementation of Parameter, DeviceFunction, Kernel and their +builders. + +Author: B.G (07/2026) +""" + +import numpy as np +import taichi as ti + +from ._closure_backend import ClosureBackendParameter, ClosureParamDeviceView, specialize_closure +from .base import DeviceFunction, DeviceFunctionBuilder, Kernel, KernelBuilder, attach_meta + + +class TaichiParameter(ClosureBackendParameter): + """ + Parameter backed by a Taichi const value or a pooled TaichiDataHandle. + + Author: B.G (07/2026) + """ + + @staticmethod + def _numpy_dtype(dtype): + """ + Map a Taichi dtype to the numpy dtype used for host-side (de)serialization. + + Author: B.G (07/2026) + """ + if dtype == ti.u8: + return np.uint8 + if dtype == ti.i32: + return np.int32 + if dtype == ti.i64: + return np.int64 + return np.float32 + + def device_view(self) -> ClosureParamDeviceView: + """ + Compile this parameter's uniform device accessors as ti.funcs. + + get(node) dispatches on mode at ti.func trace time via ti.static, so + only the taken arm compiles: const returns a baked literal, scalar + reads HANDLE[None], field reads HANDLE[node]. set_node(node, val) is + built only for scalar/field (const is read-only, exposes no setter). + MODE/VALUE/HANDLE are plain python values, not Taichi values. + + Const getters bake VALUE as a compile-time literal: a later .set() + needs the view (and anything binding it) rebuilt to take effect. + + Author: B.G (07/2026) + """ + mode = self.mode + value = self._const_value + handle = self._handle.data if self._handle is not None else None + + def get_template(node): + if ti.static(MODE == "const"): + return VALUE + elif ti.static(MODE == "scalar"): + return HANDLE[None] + else: + return HANDLE[node] + + get_fn = ti.func(specialize_closure(get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle})) + + set_fn = None + if mode != "const": + + def set_node_template(node, val): + if ti.static(MODE == "scalar"): + HANDLE[None] = val + else: + HANDLE[node] = val + + set_fn = ti.func(specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle})) + + return ClosureParamDeviceView(self.name, get_fn, set_fn) + + +class TaichiDeviceFunction(DeviceFunction): + """ + DeviceFunction backed by a compiled ti.func. Built by TaichiDeviceFunctionBuilder. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + Author: B.G (07/2026) + """ + return self._compiled + + def __call__(self, *args, **kwargs): + """ + A ti.func only runs inside kernel/func scope; callers use `.compiled` there. + + Author: B.G (07/2026) + """ + raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") + + +class TaichiKernel(Kernel): + """ + Kernel backed by a compiled ti.kernel. Built by TaichiKernelBuilder. + + The template's own signature declares data-field arguments only (e.g. + `def template(out: ti.template()): ...`); params/helpers arrive via bind() + and are resolved into the kernel body, so callers pass data fields only. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + 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 TaichiDeviceFunctionBuilder(DeviceFunctionBuilder): + """ + Builds a TaichiDeviceFunction: specialize the ingested def with bound + globals, decorate with ti.func. + + Author: B.G (07/2026) + """ + + def compile(self) -> TaichiDeviceFunction: + """ + Author: B.G (07/2026) + """ + specialised = specialize_closure(self._template, self._bindings) + fn = TaichiDeviceFunction(specialised.__name__, ti.func(specialised)) + attach_meta(fn, self._template, self._bindings) + return fn + + +class TaichiKernelBuilder(KernelBuilder): + """ + Builds a TaichiKernel: specialize the ingested def with bound globals, + decorate with ti.kernel. + + Author: B.G (07/2026) + """ + + def compile(self) -> TaichiKernel: + """ + Author: B.G (07/2026) + """ + specialised = specialize_closure(self._template, self._bindings) + krn = TaichiKernel(specialised.__name__, ti.kernel(specialised)) + attach_meta(krn, self._template, self._bindings) + return krn 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/src/core/pool/quadrants_pool.py b/pyfastflow/experimental/core/pool/_bucketed_pool.py similarity index 74% rename from src/core/pool/quadrants_pool.py rename to pyfastflow/experimental/core/pool/_bucketed_pool.py index ea0e6a8..e544d54 100644 --- a/src/core/pool/quadrants_pool.py +++ b/pyfastflow/experimental/core/pool/_bucketed_pool.py @@ -1,24 +1,26 @@ """ -Quadrants backend implementation of Pool. +Shared bucketed Pool implementation, parameterized by a DataHandle subclass. Author: B.G (07/2026) """ -from typing import Any +from typing import Any, ClassVar from .base import DataHandle, Pool -from .quadrants_handle import QuadrantsDataHandle -class QuadrantsPool(Pool): +class BucketedPool(Pool): """ - Pool manager for QuadrantsDataHandle, bucketed by (dtype, shape). + 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[QuadrantsDataHandle]] = {} + self._buckets: dict[tuple[Any, tuple[int, ...]], list[DataHandle]] = {} def get_data(self, dtype, shape) -> DataHandle: key = (dtype, tuple(shape)) @@ -27,7 +29,7 @@ def get_data(self, dtype, shape) -> DataHandle: if not handle.in_use: handle.acquire() return handle - handle = QuadrantsDataHandle(dtype, key[1]) + handle = self._handle_cls(dtype, key[1]) handle.acquire() bucket.append(handle) return handle diff --git a/src/core/pool/quadrants_handle.py b/pyfastflow/experimental/core/pool/_fields_handle.py similarity index 56% rename from src/core/pool/quadrants_handle.py rename to pyfastflow/experimental/core/pool/_fields_handle.py index aaaced8..c9d8a6d 100644 --- a/src/core/pool/quadrants_handle.py +++ b/pyfastflow/experimental/core/pool/_fields_handle.py @@ -1,52 +1,56 @@ """ -Quadrants backend implementation of DataHandle. +Shared DataHandle implementation for FieldsBuilder-based backends. -Mirrors taichi_handle.py - Quadrants' FieldsBuilder/SNodeTree lifecycle is -identical to Taichi's. Composition, not inheritance, same rationale as the -Taichi backend: kernels take the raw field via `.data`. +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 - -import quadrants as qd +from typing import Any, ClassVar from .base import DataHandle -class QuadrantsDataHandle(DataHandle): +class FieldsBuilderDataHandle(DataHandle): """ - DataHandle backed by one Quadrants field. + 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 Quadrants field of the given dtype/shape via FieldsBuilder. + 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) """ - QuadrantsDataHandle._next_id += 1 - self.id = QuadrantsDataHandle._next_id + cls = type(self) + cls._next_id += 1 + self.id = cls._next_id self.dtype = dtype self.shape = tuple(shape) self.in_use = False - self._builder = qd.FieldsBuilder() - self._field = qd.field(dtype) + 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(qd.i, self.shape).place(self._field) + self._builder.dense(backend.i, self.shape).place(self._field) elif len(self.shape) == 2: - self._builder.dense(qd.ij, self.shape).place(self._field) + 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.") @@ -55,9 +59,8 @@ def __init__(self, dtype: Any, shape: tuple[int, ...]): @property def data(self): """ - Return the underlying qd.field, for passing straight into kernels - or binding as a global (ndarrays cannot be bound as globals - - Parameter/DeviceFunction wiring stays field-backed for that reason). + Return the underlying field, for passing straight into kernels or + binding as a global. Author: B.G (07/2026) """ diff --git a/src/core/pool/base.py b/pyfastflow/experimental/core/pool/base.py similarity index 100% rename from src/core/pool/base.py rename to pyfastflow/experimental/core/pool/base.py diff --git a/pyfastflow/experimental/core/pool/cupy_handle.py b/pyfastflow/experimental/core/pool/cupy_handle.py new file mode 100644 index 0000000..0d99d41 --- /dev/null +++ b/pyfastflow/experimental/core/pool/cupy_handle.py @@ -0,0 +1,65 @@ +""" +Cupy backend implementation of DataHandle. + +Author: B.G (07/2026) +""" + +from typing import Any + +import cupy as cp + +from .base import DataHandle + + +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.id = CupyDataHandle._next_id + 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/src/core/context/__init__.py b/src/core/context/__init__.py deleted file mode 100644 index 366edf6..0000000 --- a/src/core/context/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -""" -New backend-agnostic context architecture (Parameter/Specializable ABCs + backends). - -Physically lives under ./src/core/context, exposed at import time as -pyfastflow.experimental.core.context via the path shim in -pyfastflow/experimental/core/__init__.py. - -Author: B.G (07/2026) -""" diff --git a/src/core/context/base.py b/src/core/context/base.py deleted file mode 100644 index 6674c9c..0000000 --- a/src/core/context/base.py +++ /dev/null @@ -1,203 +0,0 @@ -""" -Backend-agnostic context building blocks. - -No "Context" class here on purpose: a context is just whatever concrete class -(GridContext, FlowContext, ...) groups a set of Parameters and registers -DeviceFunctions. Cross-context references are plain explicit bindings passed -at compile() time, not a stored connection registry. - -Author: B.G (07/2026) -""" - -from abc import ABC, abstractmethod -from typing import Any, ClassVar - -from ..pool.base import DataHandle - - -class Parameter(ABC): - """ - One named, typed value owned by a context. - - REQUIRED_MODES is the baseline every backend must support; a backend - widens SUPPORTED_MODES to add more storage kinds. Enforced at subclass - definition time via __init_subclass__, not at instantiation. - - Author: B.G (07/2026) - """ - - REQUIRED_MODES: ClassVar[frozenset[str]] = frozenset({"const", "scalar", "field"}) - SUPPORTED_MODES: ClassVar[frozenset[str]] = REQUIRED_MODES - - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - missing = Parameter.REQUIRED_MODES - cls.SUPPORTED_MODES - if missing: - raise TypeError(f"{cls.__name__} must support modes {sorted(missing)}") - - name: str - dtype: Any - mode: str - - @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 parameter's value in place, according to its mode. - - Author: B.G (07/2026) - """ - ... - - @abstractmethod - def destroy(self) -> None: - """ - Release any backing storage owned by this parameter. - - Author: B.G (07/2026) - """ - ... - - -class Specializable(ABC): - """ - Shared compile/call contract for DeviceFunction and (later) Kernel. - - compile() specializes a template with explicit bindings injected as - globals - the mechanism that replaces a stored connection registry. - - Author: B.G (07/2026) - """ - - name: str - - @classmethod - @abstractmethod - def compile(cls, template, *, bindings: dict[str, Any]) -> "Specializable": - """ - Specialize `template` with `bindings` injected as globals. - - Author: B.G (07/2026) - """ - ... - - @property - @abstractmethod - def compiled(self): - """ - Raw backend callable (e.g. the ti.func/ti.kernel object), for - injection as a global into another template's bindings. - - Author: B.G (07/2026) - """ - ... - - @abstractmethod - def __call__(self, *args, **kwargs): ... - - -class DeviceFunction(Specializable): - """ - Compiled device-side helper (e.g. a ti.func specialization). - - Not necessarily callable from host Python - backends where device - functions can only run inside kernel/func scope may raise on __call__. - - Author: B.G (07/2026) - """ - - -class Kernel(Specializable): - """ - Compiled entry point (e.g. a ti.kernel specialization). - - Unlike DeviceFunction, __call__ is expected to work from host Python - - that's how the compute actually gets launched. Its template's own - parameters are data fields only; params/helpers arrive via bindings - and are resolved into the kernel body, not passed at call time. - - Author: B.G (07/2026) - """ - - -def resolve_binding(value): - """ - Unwrap a Parameter/Specializable to its backend-workable object. - - Parameter -> get() -> .data if that's a DataHandle, else the raw value. - Specializable (DeviceFunction/Kernel) -> .compiled. - Anything else passes through unchanged. - - Author: B.G (07/2026) - """ - if isinstance(value, Parameter): - resolved = value.get() - return resolved.data if isinstance(resolved, DataHandle) else resolved - if isinstance(value, Specializable): - return value.compiled - return value - - -class Bag: - """ - Simple named collection, mergeable into a Specializable.compile() bindings dict via as_bindings(). - - Author: B.G (07/2026) - """ - - def __init__(self, items: dict[str, Any] | None = None): - self._items: dict[str, Any] = dict(items or {}) - - 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 __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 as_bindings(self) -> dict[str, Any]: - """ - Return {name: item} for merging into a compile() bindings dict. - - Author: B.G (07/2026) - """ - return dict(self._items) - - -class ParamBag(Bag): - """ - Named collection of Parameter objects. - - Author: B.G (07/2026) - """ - - -class HelperBag(Bag): - """ - Named collection of DeviceFunction objects. - - Author: B.G (07/2026) - """ diff --git a/src/core/context/quadrants_backend.py b/src/core/context/quadrants_backend.py deleted file mode 100644 index f00ca23..0000000 --- a/src/core/context/quadrants_backend.py +++ /dev/null @@ -1,208 +0,0 @@ -""" -Quadrants backend implementation of Parameter, DeviceFunction and Kernel. - -Mirrors taichi_backend.py; the specialization mechanism (clone template code -object, patch globals, decorate) is unchanged between the two - verified -empirically before writing this file. One addition specific to this backend: -Kernel templates may type their own data-field arguments as `qd.Tensor`, -which accepts either a field- or ndarray-backed value at call time with no -change to the compiled template - Taichi has no equivalent for this. - -Author: B.G (07/2026) -""" - -from types import FunctionType -from typing import Any - -import numpy as np -import quadrants as qd - -from ..pool.base import Pool -from .base import DeviceFunction, Kernel, Parameter, resolve_binding - - -def _numpy_dtype(dtype): - """ - Map a Quadrants dtype to the numpy dtype used for host-side (de)serialization. - - Author: B.G (07/2026) - """ - if dtype == qd.u8: - return np.uint8 - if dtype == qd.i32: - return np.int32 - if dtype == qd.i64: - return np.int64 - return np.float32 - - -class QuadrantsParameter(Parameter): - """ - Parameter backed by a Quadrants scalar/const value or a pooled QuadrantsDataHandle. - - Author: B.G (07/2026) - """ - - SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) - - def __init__(self, name: str, *, dtype, mode: str, value, pool: 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. - - Author: B.G (07/2026) - """ - if mode not in self.SUPPORTED_MODES: - raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") - - 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.set(value) - - def get(self): - """ - Author: B.G (07/2026) - """ - return self._const_value if self.mode == "const" else self._handle - - def set(self, value) -> None: - """ - Author: B.G (07/2026) - """ - if self.mode == "const": - self._const_value = _numpy_dtype(self.dtype)(value).item() - elif self.mode == "scalar": - self._handle.data[None] = value - else: # field - arr = np.asarray(value, dtype=_numpy_dtype(self.dtype)).reshape(-1) - self._handle.data.from_numpy(arr) - - def destroy(self) -> None: - """ - Author: B.G (07/2026) - """ - if self._handle is not None: - self._pool.release_data(self._handle) - self._handle = None - - -def _specialize(template, bindings: dict[str, Any]): - """ - Clone `template`'s code object with a globals dict where sentinels are - replaced by resolved bindings. Shared by DeviceFunction and Kernel - - they only differ in the qd.func/qd.kernel decorator applied on top. - - Author: B.G (07/2026) - """ - resolved = {name: resolve_binding(value) 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 QuadrantsDeviceFunction(DeviceFunction): - """ - DeviceFunction backed by a compiled qd.func. - - Only field-backed Parameters/DeviceFunctions can be resolved into a - template this way - Quadrants rejects ndarrays referenced as globals - (verified: raises "Ndarray used in kernel scope but not registered as - a kernel parameter" even for the unwrapped object). - - Author: B.G (07/2026) - """ - - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @classmethod - def compile(cls, template, *, bindings: dict[str, Any]) -> "QuadrantsDeviceFunction": - """ - Author: B.G (07/2026) - """ - specialised = _specialize(template, bindings) - return cls(specialised.__name__, qd.func(specialised)) - - @property - def compiled(self): - """ - Author: B.G (07/2026) - """ - return self._compiled - - def __call__(self, *args, **kwargs): - """ - A qd.func cannot be invoked from host Python - it only runs inside - kernel/func scope, where callers use `.compiled` directly. - - Author: B.G (07/2026) - """ - raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") - - -class QuadrantsKernel(Kernel): - """ - Kernel backed by a compiled qd.kernel. - - The template's own data-field arguments should be typed `qd.Tensor` - - unlike Taichi's ti.template()/ti.types.ndarray() split, a single - qd.Tensor-typed template accepts either a field- or ndarray-backed - value at call time (verified). Params/helpers still arrive via - bindings, resolved into the kernel body - callers only ever pass data - fields at call time. - - Author: B.G (07/2026) - """ - - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @classmethod - def compile(cls, template, *, bindings: dict[str, Any]) -> "QuadrantsKernel": - """ - Author: B.G (07/2026) - """ - specialised = _specialize(template, bindings) - return cls(specialised.__name__, qd.kernel(specialised)) - - @property - def compiled(self): - """ - 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) diff --git a/src/core/context/taichi_backend.py b/src/core/context/taichi_backend.py deleted file mode 100644 index 3f3e6da..0000000 --- a/src/core/context/taichi_backend.py +++ /dev/null @@ -1,198 +0,0 @@ -""" -Taichi backend implementation of Parameter, DeviceFunction and Kernel. - -Author: B.G (07/2026) -""" - -from types import FunctionType -from typing import Any - -import numpy as np -import taichi as ti - -from ..pool.base import Pool -from .base import DeviceFunction, Kernel, Parameter, resolve_binding - - -def _specialize(template, bindings: dict[str, Any]): - """ - Clone `template`'s code object with a globals dict where sentinels are - replaced by resolved bindings. Shared by DeviceFunction and Kernel - - they only differ in the ti.func/ti.kernel decorator applied on top. - - Author: B.G (07/2026) - """ - resolved = {name: resolve_binding(value) 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 - - -def _numpy_dtype(dtype): - """ - Map a Taichi dtype to the numpy dtype used for host-side (de)serialization. - - Author: B.G (07/2026) - """ - if dtype == ti.u8: - return np.uint8 - if dtype == ti.i32: - return np.int32 - if dtype == ti.i64: - return np.int64 - return np.float32 - - -class TaichiParameter(Parameter): - """ - Parameter backed by a Taichi scalar/const value or a pooled TaichiDataHandle. - - Author: B.G (07/2026) - """ - - SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) - - def __init__(self, name: str, *, dtype, mode: str, value, pool: 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. - - Author: B.G (07/2026) - """ - if mode not in self.SUPPORTED_MODES: - raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") - - 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.set(value) - - def get(self): - """ - Author: B.G (07/2026) - """ - return self._const_value if self.mode == "const" else self._handle - - def set(self, value) -> None: - """ - Author: B.G (07/2026) - """ - if self.mode == "const": - self._const_value = _numpy_dtype(self.dtype)(value).item() - elif self.mode == "scalar": - self._handle.data[None] = value - else: # field - arr = np.asarray(value, dtype=_numpy_dtype(self.dtype)).reshape(-1) - self._handle.data.from_numpy(arr) - - def destroy(self) -> None: - """ - Author: B.G (07/2026) - """ - if self._handle is not None: - self._pool.release_data(self._handle) - self._handle = None - - -class TaichiDeviceFunction(DeviceFunction): - """ - DeviceFunction backed by a compiled ti.func. - - compile() clones the template's code object with a globals dict where - sentinels are replaced by resolved bindings, then decorates with ti.func - - same specialization mechanism as the legacy CallableFactory.compile. - - Author: B.G (07/2026) - """ - - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @classmethod - def compile(cls, template, *, bindings: dict[str, Any]) -> "TaichiDeviceFunction": - """ - Author: B.G (07/2026) - """ - specialised = _specialize(template, bindings) - return cls(specialised.__name__, ti.func(specialised)) - - @property - def compiled(self): - """ - Author: B.G (07/2026) - """ - return self._compiled - - def __call__(self, *args, **kwargs): - """ - A ti.func cannot be invoked from host Python - it only runs inside - kernel/func scope, where callers use `.compiled` directly. - - Author: B.G (07/2026) - """ - raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") - - -class TaichiKernel(Kernel): - """ - Kernel backed by a compiled ti.kernel. - - The template's own signature should only declare data-field arguments - (e.g. `def template(out: ti.template()): ...`) - any params/helpers it - needs are resolved into the kernel body via `bindings`, so callers only - ever pass data fields at call time, never params/helpers. - - Author: B.G (07/2026) - """ - - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @classmethod - def compile(cls, template, *, bindings: dict[str, Any]) -> "TaichiKernel": - """ - Author: B.G (07/2026) - """ - specialised = _specialize(template, bindings) - return cls(specialised.__name__, ti.kernel(specialised)) - - @property - def compiled(self): - """ - 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) diff --git a/src/core/pool/__init__.py b/src/core/pool/__init__.py deleted file mode 100644 index 2772901..0000000 --- a/src/core/pool/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -""" -New backend-agnostic pool architecture (DataHandle/Pool ABCs + backends). - -Physically lives under ./src/core/pool, exposed at import time as -pyfastflow.experimental.core.pool via the path shim in -pyfastflow/experimental/core/__init__.py. - -Author: B.G (07/2026) -""" diff --git a/src/core/pool/taichi_handle.py b/src/core/pool/taichi_handle.py deleted file mode 100644 index ce0a820..0000000 --- a/src/core/pool/taichi_handle.py +++ /dev/null @@ -1,86 +0,0 @@ -""" -Taichi backend implementation of DataHandle. - -Wraps one Taichi field allocated via FieldsBuilder for explicit lifetime -control (snodetree.destroy() actually frees GPU memory; release() just -returns it to the pool). Composition, not inheritance: kernels take the -raw field via `.data`, not the handle itself - see pool/base.py design -notes on why subclassing ti.Field was rejected. - -Author: B.G (07/2026) -""" - -from typing import Any - -import taichi as ti - -from .base import DataHandle - - -class TaichiDataHandle(DataHandle): - """ - DataHandle backed by one Taichi field. - - Author: B.G (07/2026) - """ - - _next_id = 0 - - def __init__(self, dtype: Any, shape: tuple[int, ...]): - """ - Allocate a Taichi field of the given dtype/shape via FieldsBuilder. - - shape=() allocates a 0D scalar field, indexed as field[None]. - - Author: B.G (07/2026) - """ - TaichiDataHandle._next_id += 1 - self.id = TaichiDataHandle._next_id - self.dtype = dtype - self.shape = tuple(shape) - self.in_use = False - - self._builder = ti.FieldsBuilder() - self._field = ti.field(dtype) - - if len(self.shape) == 0: - self._builder.place(self._field) - elif len(self.shape) == 1: - self._builder.dense(ti.i, self.shape).place(self._field) - elif len(self.shape) == 2: - self._builder.dense(ti.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 ti.field, for passing straight into kernels. - - 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/src/core/pool/taichi_pool.py b/src/core/pool/taichi_pool.py deleted file mode 100644 index c86aea9..0000000 --- a/src/core/pool/taichi_pool.py +++ /dev/null @@ -1,54 +0,0 @@ -""" -Taichi backend implementation of Pool. - -Author: B.G (07/2026) -""" - -from typing import Any - -from .base import DataHandle, Pool -from .taichi_handle import TaichiDataHandle - - -class TaichiPool(Pool): - """ - Pool manager for TaichiDataHandle, bucketed by (dtype, shape). - - Author: B.G (07/2026) - """ - - def __init__(self): - self._buckets: dict[tuple[Any, tuple[int, ...]], list[TaichiDataHandle]] = {} - - 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 = TaichiDataHandle(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} From d7785f43213ad2f888bcea21c0d5bb9698ff93bd Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 11:15:02 +0200 Subject: [PATCH 08/31] Simplify taichi vs quadrants backend as they are very similar in syntax --- .../core/context/_closure_backend.py | 232 +++++++++++++++++- pyfastflow/experimental/core/context/base.py | 115 ++++++++- .../experimental/core/context/cupy_backend.py | 23 +- .../core/context/quadrants_backend.py | 136 ++-------- .../core/context/taichi_backend.py | 139 ++--------- 5 files changed, 382 insertions(+), 263 deletions(-) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 20de1ad..d4293e5 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -7,11 +7,21 @@ """ from types import FunctionType -from typing import Any +from typing import Any, ClassVar import numpy as np -from .base import Parameter, resolve_binding +from .base import ( + Bag, + DeviceFunction, + DeviceFunctionBuilder, + Kernel, + KernelBuilder, + Parameter, + attach_meta, + filter_bindings, + resolve_binding, +) def specialize_closure(template, bindings: dict[str, Any]) -> FunctionType: @@ -61,13 +71,17 @@ def __init__(self, name: str, get_fn, set_fn=None): class ClosureBackendParameter(Parameter): """ Parameter backed by a const value or a pooled DataHandle, for backends - sharing the closure-specialization mechanism. Subclasses pin `_numpy_dtype` - and implement `device_view` with their own func decorator. + sharing the closure-specialization mechanism (Taichi, Quadrants). + Subclasses pin `_backend` (the `ti`/`qd` module) - `_numpy_dtype` and + `_build_device_view` are shared here since `ti.u8/i32/i64`, + `ti.func`/`ti.static` and their `qd.` equivalents are identical in name + and behaviour across both modules. Author: B.G (07/2026) """ SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) + _backend: ClassVar[Any] def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): """ @@ -90,6 +104,7 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No 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, ()) @@ -100,15 +115,22 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No self.set(value) - @staticmethod - def _numpy_dtype(dtype): + @classmethod + def _numpy_dtype(cls, dtype): """ - Map a backend dtype to the numpy dtype used for host-side (de)serialization. - Overridden per backend. + Map a backend dtype (`ti.*`/`qd.*`) to the numpy dtype used for + host-side (de)serialization. Author: B.G (07/2026) """ - raise NotImplementedError + 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): """ @@ -122,6 +144,7 @@ def set(self, value) -> None: """ if self.mode == "const": self._const_value = self._numpy_dtype(self.dtype)(value).item() + self._device_view = None # a cached view would bake the stale literal elif self.mode == "scalar": self._handle.data[None] = value else: # field @@ -148,3 +171,194 @@ def destroy(self) -> None: 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: + """ + Cached uniform device accessor for this parameter. Built once via + `_build_device_view` (backend-specific) and memoized on the instance: + rebuilding it per compile is pointless python work at O(params x + kernels) since the ti.func/qd.func objects are identical every time. + Invalidated (not refreshed) by set() on const mode and by destroy(), + the two places the baked literal/handle a cached view depends on can + change; scalar/field set() writes through the handle a cached view + already points at, so no invalidation is needed there. + + 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 uniform device accessors as backend funcs + (ti.func / qd.func). + + get(node) dispatches on mode at trace time via `_backend.static`, so + only the taken arm compiles: const returns a baked literal, scalar + reads HANDLE[None], field reads HANDLE[node]. set_node(node, val) is + built only for scalar/field (const is read-only, exposes no setter). + MODE/VALUE/HANDLE are plain python values, not backend values. + + Const getters bake VALUE as a compile-time literal: a later .set() + needs the view (and anything binding it) rebuilt to take effect. + + Called at most once per instance, memoized by device_view(). + + Author: B.G (07/2026) + """ + backend = self._backend + mode = self.mode + value = self._const_value + handle = self._handle.data if self._handle is not None else None + + def get_template(node): + if STATIC(MODE == "const"): + return VALUE + elif STATIC(MODE == "scalar"): + return HANDLE[None] + else: + return HANDLE[node] + + get_fn = backend.func( + specialize_closure(get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle, "STATIC": backend.static}) + ) + + set_fn = None + if mode != "const": + + def set_node_template(node, val): + if STATIC(MODE == "scalar"): + HANDLE[None] = val + else: + HANDLE[node] = val + + set_fn = backend.func( + specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle, "STATIC": backend.static}) + ) + + return ClosureParamDeviceView(self.name, get_fn, set_fn) + + +class ClosureDeviceFunction(DeviceFunction): + """ + DeviceFunction backed by a compiled closure-backend func (ti.func / + qd.func). Built by ClosureDeviceFunctionBuilder subclasses. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + 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"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") + + +class ClosureKernel(Kernel): + """ + Kernel backed by a compiled closure-backend kernel (ti.kernel / + qd.kernel). Built by ClosureKernelBuilder subclasses. + + The template's own signature declares data-field arguments only (e.g. + `def template(out: ti.template()): ...`); params/helpers arrive via bind() + and are resolved into the kernel body, so callers pass data fields only. + + Author: B.G (07/2026) + """ + + def __init__(self, name: str, compiled): + self.name = name + self._compiled = compiled + + @property + def compiled(self): + """ + 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) + + +def _check_const_only(bindings: dict[str, Any]) -> None: + """ + Enforce the framework rule (base.py module docstring): a device helper + only ever binds const-mode Parameters - any data it needs is passed to it + as an explicit argument by the calling kernel instead. Recurses into + bound Bags so a scalar/field param nested inside one is also caught. + + Author: B.G (07/2026) + """ + for name, value in bindings.items(): + if isinstance(value, Parameter): + if value.mode != "const": + raise ValueError( + f"device helper: bound parameter '{name}' has mode {value.mode!r}, but a device " + "helper may only bind const parameters - pass the data as an explicit argument to " + "the helper instead" + ) + elif isinstance(value, Bag): + _check_const_only(dict(value.items())) + + +class ClosureDeviceFunctionBuilder(DeviceFunctionBuilder): + """ + Builds a ClosureDeviceFunction: specialize the ingested def with bound + globals, decorate with `_backend.func`. Subclasses pin `_backend`. + + Author: B.G (07/2026) + """ + + _backend: ClassVar[Any] + + def compile(self) -> ClosureDeviceFunction: + """ + Author: B.G (07/2026) + """ + _check_const_only(self._bindings) + specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings)) + fn = ClosureDeviceFunction(specialised.__name__, self._backend.func(specialised)) + attach_meta(fn, self._template, self._bindings) + return fn + + +class ClosureKernelBuilder(KernelBuilder): + """ + Builds a ClosureKernel: specialize the ingested def with bound globals, + decorate with `_backend.kernel`. Subclasses pin `_backend`. + + Author: B.G (07/2026) + """ + + _backend: ClassVar[Any] + + def compile(self) -> ClosureKernel: + """ + Author: B.G (07/2026) + """ + specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings)) + krn = ClosureKernel(specialised.__name__, self._backend.kernel(specialised)) + attach_meta(krn, self._template, self._bindings) + return krn diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 732c509..8b251ab 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -6,6 +6,17 @@ DeviceFunctions. Cross-context references are plain explicit bindings passed at bind() time, not a stored connection registry. +Load-bearing rule: data flows at call time (a template's own explicit +arguments); configuration flows at compile time (bind()ed Parameters/helpers/ +bags, resolved into the template body). + +A device helper (DeviceFunction) may only bind const-mode Parameters. Any +data it needs (scalar/field buffers) is passed to it as an explicit argument +by the calling kernel, not bound in - a device helper has no way to receive +an extra pointer argument once spliced into a caller, so this holds across +all three backends alike. Kernels are unrestricted: they may bind params of +any mode. + The compile surface is a two-layer builder: an abstract DeviceFunctionBuilder / KernelBuilder here, backend variants (Taichi*/Quadrants*/Cupy*) elsewhere. A builder collects bind()ed dependencies + one ingest()ed template and produces a @@ -18,8 +29,9 @@ import ast import inspect +import warnings from abc import ABC, abstractmethod -from types import SimpleNamespace +from functools import lru_cache from typing import Any, ClassVar from ..pool.base import DataHandle @@ -155,6 +167,32 @@ class Kernel(Specializable): """ +class _LazyBagView: + """ + Lazy stand-in for a Bag in a template body: `phys.g` resolves `g` only on + first attribute access (via resolve_binding), then caches it in the + instance dict so later lookups bypass __getattr__ entirely. Avoids + filter_bindings' top-level-only reach: binding one large bag no longer + eagerly resolves (and device_view()-compiles) every member, only the ones + a template actually touches. Nested Bag members resolve to another + _LazyBagView, so nesting stays lazy all the way down. + + Author: B.G (07/2026) + """ + + def __init__(self, bag: "Bag"): + object.__setattr__(self, "_bag", bag) + + def __getattr__(self, name: str) -> Any: + # only called on a genuine miss (cached hits never reach here) + bag = object.__getattribute__(self, "_bag") + if name not in bag: + raise AttributeError(name) + resolved = resolve_binding(bag[name]) + self.__dict__[name] = resolved + return resolved + + def resolve_binding(value): """ Unwrap a bound object to what a closure-backend template body should see. @@ -162,8 +200,9 @@ def resolve_binding(value): Parameter -> literal if solo (const), else device_view() (a .get/.set_node carrier for uniform in-kernel access). Specializable (DeviceFunction/Kernel) -> .compiled. - Bag -> a namespace whose attributes are each member resolved the same way, - so dotted paths (grid.nx.get(i)) trace as plain attribute lookups. + Bag -> a _LazyBagView resolving each member on first attribute access, so + dotted paths (grid.nx.get(i)) trace as plain attribute lookups without + forcing every other member in the bag to resolve too. Anything else passes through unchanged. Author: B.G (07/2026) @@ -176,15 +215,21 @@ def resolve_binding(value): if isinstance(value, Specializable): return value.compiled if isinstance(value, Bag): - return SimpleNamespace(**{name: resolve_binding(item) for name, item in value.items()}) + return _LazyBagView(value) return value +@lru_cache(maxsize=None) 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 with no AST (not python). + Cached per template: every compile() asks twice (once to filter bindings, + once for attach_meta), and each miss costs an inspect.getsource + a parse. + The returned tree is therefore SHARED by every Specializable built from + that template - treat it as read-only. + Author: B.G (07/2026) """ if isinstance(template, str): @@ -200,6 +245,45 @@ def capture_template_meta(template) -> tuple[str | None, ast.AST | None]: return source, tree +def used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: + """ + Subset of `bindings` whose key is actually referenced by `template`'s + body (by name - an attribute chain's root, e.g. `phys` in + `phys.g.get(0)`, is itself an `ast.Name`, so collecting every `ast.Name` + id in the tree is sufficient). Falls back to returning `bindings` + unchanged whenever the AST isn't available (non-python template, source + unavailable) - never silently drop a binding we 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]: + """ + used_bindings(), plus one warning per compile listing everything bound but + never referenced - catches typos in a large context. Backend-agnostic: the + closure builders feed the result to their specializer. No warning when the + AST isn't available (used_bindings then returns `bindings` unchanged, so + nothing looks unused). + + 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 + + def attach_meta(obj: Specializable, template, bindings: dict[str, Any]) -> None: """ Stash hidden raw material on a freshly built Specializable for later reuse @@ -290,20 +374,35 @@ class Bag: Simple named collection, mergeable into a builder via bind_bag() or bound whole (bind('grid', bag)) for dotted-path access in the template body. + `_member_type` (None on the base class) restricts what a subclass's + members may be; a nested Bag is always accepted regardless, since bags + nest (a ParamBag may legitimately contain another bag of params). + Author: B.G (07/2026) """ + _member_type: ClassVar[type | None] = None + def __init__(self, items: dict[str, Any] | None = None): - self._items: dict[str, Any] = dict(items or {}) + self._items: dict[str, Any] = {} + for name, item in (items or {}).items(): + self.add(name, item) def add(self, name: str, item: Any) -> None: """ - Register `item` under `name`. Raises if `name` is already registered. + Register `item` under `name`. Raises if `name` is already registered + or if `item` doesn't match this bag's `_member_type` (a nested Bag is + always accepted). Author: B.G (07/2026) """ if name in self._items: raise KeyError(f"'{name}' is already registered in this bag") + if self._member_type is not None and not isinstance(item, (self._member_type, Bag)): + raise TypeError( + f"{type(self).__name__}: member '{name}' must be a {self._member_type.__name__} " + f"(or a Bag), got {type(item).__name__}" + ) self._items[name] = item def __getattr__(self, name: str) -> Any: @@ -340,6 +439,8 @@ class ParamBag(Bag): Author: B.G (07/2026) """ + _member_type = Parameter + class HelperBag(Bag): """ @@ -347,3 +448,5 @@ class HelperBag(Bag): Author: B.G (07/2026) """ + + _member_type = DeviceFunction diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 71b654e..a6b9b78 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -24,7 +24,11 @@ Device functions may only reference const params / other device functions inside spans - a spliced __device__ function has no way to receive an extra -pointer argument. +pointer argument. This isn't a cupy quirk: it's the backend-agnostic rule +(see base.py's module docstring) that a device helper only ever binds const +params, with any data it needs passed in as an explicit argument by the +calling kernel - the closure backends (Taichi, Quadrants) enforce the same +rule in ClosureDeviceFunctionBuilder.compile(). Author: B.G (07/2026) """ @@ -193,17 +197,22 @@ def parse(self, body: str) -> str: return _SPAN_RE.sub(self._repl, body) -def _const_defines(bindings: dict[str, Any]) -> list[str]: +def _const_defines(bindings: dict[str, Any], body: str) -> list[str]: """ - `#define NAME literal` for each top-level const Parameter, so bare uses of - a const param (outside any span) compile. + `#define NAME literal` for each top-level const Parameter whose identifier + actually appears (word-boundary match) in `body` - `body` must be the + already span-expanded text, so a span like `$phys.dx.get(0)$` has already + become a numeric literal and correctly does not count as a bare use of a + top-level const name. Skipping unused names avoids pasting unhygienic + macros (a const named N, I, DIM, EPS, min, ...) into the translation unit + where they'd silently rewrite unrelated identifiers. 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" + if isinstance(obj, Parameter) and obj.mode == "const" and re.search(rf"\b{re.escape(name)}\b", body) ] @@ -395,7 +404,7 @@ def compile(self) -> CupyDeviceFunction: name = _extract_name(_DEVICE_NAME_RE, template, "__device__") parser = _SpanParser(self._bindings, allow_arrays=False) body = parser.parse(template) - source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings) + [body]) + source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body]) fn = CupyDeviceFunction(name, source) attach_meta(fn, template, self._bindings) return fn @@ -424,7 +433,7 @@ def compile(self) -> CupyKernel: # link fine within the same translation unit. if 'extern "C"' not in body: body = body.replace("__global__", 'extern "C" __global__', 1) - source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings) + [body]) + source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body]) bound_arrays = [e["array"] for e in parser.ptr_params.values()] raw = CupyKernel._raw_cache.get(source) diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py index d6b4698..0fa1e57 100644 --- a/pyfastflow/experimental/core/context/quadrants_backend.py +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -2,8 +2,8 @@ Quadrants backend implementation of Parameter, DeviceFunction, Kernel and their builders. -Mirrors taichi_backend.py; both share the closure-specialization mechanism -via _closure_backend.py. One difference specific to this backend: Kernel +Shares the closure-specialization mechanism with taichi_backend.py via +_closure_backend.py. One difference specific to this backend: Kernel templates may type their own data-field arguments as `qd.Tensor`, which accepts either a field- or ndarray-backed value at call time with no change to the compiled template - Taichi has no equivalent for this. Note that @@ -13,11 +13,15 @@ Author: B.G (07/2026) """ -import numpy as np import quadrants as qd -from ._closure_backend import ClosureBackendParameter, ClosureParamDeviceView, specialize_closure -from .base import DeviceFunction, DeviceFunctionBuilder, Kernel, KernelBuilder, attach_meta +from ._closure_backend import ( + ClosureBackendParameter, + ClosureDeviceFunction, + ClosureDeviceFunctionBuilder, + ClosureKernel, + ClosureKernelBuilder, +) class QuadrantsParameter(ClosureBackendParameter): @@ -27,62 +31,10 @@ class QuadrantsParameter(ClosureBackendParameter): Author: B.G (07/2026) """ - @staticmethod - def _numpy_dtype(dtype): - """ - Map a Quadrants dtype to the numpy dtype used for host-side (de)serialization. - - Author: B.G (07/2026) - """ - if dtype == qd.u8: - return np.uint8 - if dtype == qd.i32: - return np.int32 - if dtype == qd.i64: - return np.int64 - return np.float32 - - def device_view(self) -> ClosureParamDeviceView: - """ - Compile this parameter's uniform device accessors as qd.funcs. - - get(node) dispatches on mode at qd.func trace time via qd.static, so - only the taken arm compiles: const returns a baked literal, scalar - reads HANDLE[None], field reads HANDLE[node]. set_node(node, val) is - built only for scalar/field (const is read-only, exposes no setter). - MODE/VALUE/HANDLE are plain python values, not Quadrants values. - - Author: B.G (07/2026) - """ - mode = self.mode - value = self._const_value - handle = self._handle.data if self._handle is not None else None - - def get_template(node): - if qd.static(MODE == "const"): - return VALUE - elif qd.static(MODE == "scalar"): - return HANDLE[None] - else: - return HANDLE[node] - - get_fn = qd.func(specialize_closure(get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle})) - - set_fn = None - if mode != "const": - - def set_node_template(node, val): - if qd.static(MODE == "scalar"): - HANDLE[None] = val - else: - HANDLE[node] = val - - set_fn = qd.func(specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle})) - - return ClosureParamDeviceView(self.name, get_fn, set_fn) - - -class QuadrantsDeviceFunction(DeviceFunction): + _backend = qd + + +class QuadrantsDeviceFunction(ClosureDeviceFunction): """ DeviceFunction backed by a compiled qd.func. Built by QuadrantsDeviceFunctionBuilder. @@ -92,27 +44,8 @@ class QuadrantsDeviceFunction(DeviceFunction): Author: B.G (07/2026) """ - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @property - def compiled(self): - """ - Author: B.G (07/2026) - """ - return self._compiled - - def __call__(self, *args, **kwargs): - """ - A qd.func only runs inside kernel/func scope; callers use `.compiled` there. - - Author: B.G (07/2026) - """ - raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") - -class QuadrantsKernel(Kernel): +class QuadrantsKernel(ClosureKernel): """ Kernel backed by a compiled qd.kernel. Built by QuadrantsKernelBuilder. @@ -124,27 +57,8 @@ class QuadrantsKernel(Kernel): Author: B.G (07/2026) """ - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @property - def compiled(self): - """ - 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 QuadrantsDeviceFunctionBuilder(DeviceFunctionBuilder): +class QuadrantsDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): """ Builds a QuadrantsDeviceFunction: specialize the ingested def with bound globals, decorate with qd.func. @@ -152,17 +66,10 @@ class QuadrantsDeviceFunctionBuilder(DeviceFunctionBuilder): Author: B.G (07/2026) """ - def compile(self) -> QuadrantsDeviceFunction: - """ - Author: B.G (07/2026) - """ - specialised = specialize_closure(self._template, self._bindings) - fn = QuadrantsDeviceFunction(specialised.__name__, qd.func(specialised)) - attach_meta(fn, self._template, self._bindings) - return fn + _backend = qd -class QuadrantsKernelBuilder(KernelBuilder): +class QuadrantsKernelBuilder(ClosureKernelBuilder): """ Builds a QuadrantsKernel: specialize the ingested def with bound globals, decorate with qd.kernel. @@ -170,11 +77,4 @@ class QuadrantsKernelBuilder(KernelBuilder): Author: B.G (07/2026) """ - def compile(self) -> QuadrantsKernel: - """ - Author: B.G (07/2026) - """ - specialised = specialize_closure(self._template, self._bindings) - krn = QuadrantsKernel(specialised.__name__, qd.kernel(specialised)) - attach_meta(krn, self._template, self._bindings) - return krn + _backend = qd diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py index 796df0e..e96c2be 100644 --- a/pyfastflow/experimental/core/context/taichi_backend.py +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -5,11 +5,15 @@ Author: B.G (07/2026) """ -import numpy as np import taichi as ti -from ._closure_backend import ClosureBackendParameter, ClosureParamDeviceView, specialize_closure -from .base import DeviceFunction, DeviceFunctionBuilder, Kernel, KernelBuilder, attach_meta +from ._closure_backend import ( + ClosureBackendParameter, + ClosureDeviceFunction, + ClosureDeviceFunctionBuilder, + ClosureKernel, + ClosureKernelBuilder, +) class TaichiParameter(ClosureBackendParameter): @@ -19,123 +23,26 @@ class TaichiParameter(ClosureBackendParameter): Author: B.G (07/2026) """ - @staticmethod - def _numpy_dtype(dtype): - """ - Map a Taichi dtype to the numpy dtype used for host-side (de)serialization. - - Author: B.G (07/2026) - """ - if dtype == ti.u8: - return np.uint8 - if dtype == ti.i32: - return np.int32 - if dtype == ti.i64: - return np.int64 - return np.float32 - - def device_view(self) -> ClosureParamDeviceView: - """ - Compile this parameter's uniform device accessors as ti.funcs. - - get(node) dispatches on mode at ti.func trace time via ti.static, so - only the taken arm compiles: const returns a baked literal, scalar - reads HANDLE[None], field reads HANDLE[node]. set_node(node, val) is - built only for scalar/field (const is read-only, exposes no setter). - MODE/VALUE/HANDLE are plain python values, not Taichi values. - - Const getters bake VALUE as a compile-time literal: a later .set() - needs the view (and anything binding it) rebuilt to take effect. - - Author: B.G (07/2026) - """ - mode = self.mode - value = self._const_value - handle = self._handle.data if self._handle is not None else None - - def get_template(node): - if ti.static(MODE == "const"): - return VALUE - elif ti.static(MODE == "scalar"): - return HANDLE[None] - else: - return HANDLE[node] - - get_fn = ti.func(specialize_closure(get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle})) - - set_fn = None - if mode != "const": - - def set_node_template(node, val): - if ti.static(MODE == "scalar"): - HANDLE[None] = val - else: - HANDLE[node] = val - - set_fn = ti.func(specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle})) - - return ClosureParamDeviceView(self.name, get_fn, set_fn) - - -class TaichiDeviceFunction(DeviceFunction): + _backend = ti + + +class TaichiDeviceFunction(ClosureDeviceFunction): """ DeviceFunction backed by a compiled ti.func. Built by TaichiDeviceFunctionBuilder. Author: B.G (07/2026) """ - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @property - def compiled(self): - """ - Author: B.G (07/2026) - """ - return self._compiled - - def __call__(self, *args, **kwargs): - """ - A ti.func only runs inside kernel/func scope; callers use `.compiled` there. - - Author: B.G (07/2026) - """ - raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") - -class TaichiKernel(Kernel): +class TaichiKernel(ClosureKernel): """ Kernel backed by a compiled ti.kernel. Built by TaichiKernelBuilder. - The template's own signature declares data-field arguments only (e.g. - `def template(out: ti.template()): ...`); params/helpers arrive via bind() - and are resolved into the kernel body, so callers pass data fields only. - Author: B.G (07/2026) """ - def __init__(self, name: str, compiled): - self.name = name - self._compiled = compiled - - @property - def compiled(self): - """ - 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 TaichiDeviceFunctionBuilder(DeviceFunctionBuilder): +class TaichiDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): """ Builds a TaichiDeviceFunction: specialize the ingested def with bound globals, decorate with ti.func. @@ -143,17 +50,10 @@ class TaichiDeviceFunctionBuilder(DeviceFunctionBuilder): Author: B.G (07/2026) """ - def compile(self) -> TaichiDeviceFunction: - """ - Author: B.G (07/2026) - """ - specialised = specialize_closure(self._template, self._bindings) - fn = TaichiDeviceFunction(specialised.__name__, ti.func(specialised)) - attach_meta(fn, self._template, self._bindings) - return fn + _backend = ti -class TaichiKernelBuilder(KernelBuilder): +class TaichiKernelBuilder(ClosureKernelBuilder): """ Builds a TaichiKernel: specialize the ingested def with bound globals, decorate with ti.kernel. @@ -161,11 +61,4 @@ class TaichiKernelBuilder(KernelBuilder): Author: B.G (07/2026) """ - def compile(self) -> TaichiKernel: - """ - Author: B.G (07/2026) - """ - specialised = specialize_closure(self._template, self._bindings) - krn = TaichiKernel(specialised.__name__, ti.kernel(specialised)) - attach_meta(krn, self._template, self._bindings) - return krn + _backend = ti From 4cd7d253e7570bee8f48bcbed8231304206753f7 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 11:16:30 +0200 Subject: [PATCH 09/31] clean now unused classes --- .../core/context/quadrants_backend.py | 38 ++++--------------- .../core/context/taichi_backend.py | 22 +---------- 2 files changed, 10 insertions(+), 50 deletions(-) diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py index 0fa1e57..095b7e8 100644 --- a/pyfastflow/experimental/core/context/quadrants_backend.py +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -17,9 +17,7 @@ from ._closure_backend import ( ClosureBackendParameter, - ClosureDeviceFunction, ClosureDeviceFunctionBuilder, - ClosureKernel, ClosureKernelBuilder, ) @@ -34,34 +32,12 @@ class QuadrantsParameter(ClosureBackendParameter): _backend = qd -class QuadrantsDeviceFunction(ClosureDeviceFunction): - """ - DeviceFunction backed by a compiled qd.func. Built by QuadrantsDeviceFunctionBuilder. - - Only field-backed Parameters/DeviceFunctions can be resolved into a - template this way - Quadrants rejects ndarrays referenced as globals. - - Author: B.G (07/2026) - """ - - -class QuadrantsKernel(ClosureKernel): - """ - Kernel backed by a compiled qd.kernel. Built by QuadrantsKernelBuilder. - - The template's own data-field arguments should be typed `qd.Tensor` - a - single qd.Tensor-typed template accepts either a field- or ndarray-backed - value at call time. Params/helpers arrive via bind() and are resolved into - the kernel body; callers pass data fields only. - - Author: B.G (07/2026) - """ - - class QuadrantsDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): """ - Builds a QuadrantsDeviceFunction: specialize the ingested def with bound - globals, decorate with qd.func. + Builds a ClosureDeviceFunction: specialize the ingested def with bound + globals, decorate with qd.func. Only field-backed Parameters/DeviceFunctions + can be resolved into a template this way - Quadrants rejects ndarrays + referenced as globals. Author: B.G (07/2026) """ @@ -71,8 +47,10 @@ class QuadrantsDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): class QuadrantsKernelBuilder(ClosureKernelBuilder): """ - Builds a QuadrantsKernel: specialize the ingested def with bound globals, - decorate with qd.kernel. + Builds a ClosureKernel: specialize the ingested def with bound globals, + decorate with qd.kernel. The template's own data-field arguments should be + typed `qd.Tensor` - a single qd.Tensor-typed template accepts either a + field- or ndarray-backed value at call time. Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py index e96c2be..f0da77d 100644 --- a/pyfastflow/experimental/core/context/taichi_backend.py +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -9,9 +9,7 @@ from ._closure_backend import ( ClosureBackendParameter, - ClosureDeviceFunction, ClosureDeviceFunctionBuilder, - ClosureKernel, ClosureKernelBuilder, ) @@ -26,25 +24,9 @@ class TaichiParameter(ClosureBackendParameter): _backend = ti -class TaichiDeviceFunction(ClosureDeviceFunction): - """ - DeviceFunction backed by a compiled ti.func. Built by TaichiDeviceFunctionBuilder. - - Author: B.G (07/2026) - """ - - -class TaichiKernel(ClosureKernel): - """ - Kernel backed by a compiled ti.kernel. Built by TaichiKernelBuilder. - - Author: B.G (07/2026) - """ - - class TaichiDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): """ - Builds a TaichiDeviceFunction: specialize the ingested def with bound + Builds a ClosureDeviceFunction: specialize the ingested def with bound globals, decorate with ti.func. Author: B.G (07/2026) @@ -55,7 +37,7 @@ class TaichiDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): class TaichiKernelBuilder(ClosureKernelBuilder): """ - Builds a TaichiKernel: specialize the ingested def with bound globals, + Builds a ClosureKernel: specialize the ingested def with bound globals, decorate with ti.kernel. Author: B.G (07/2026) From 602c88fe1b033c966397aa5bc80d61e2c03a76c0 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 11:29:42 +0200 Subject: [PATCH 10/31] Further clean and comments --- .../core/context/_closure_backend.py | 23 +++++++++++-- pyfastflow/experimental/core/context/base.py | 26 +++++--------- .../experimental/core/context/cupy_backend.py | 34 +++++++++++++++++++ 3 files changed, 63 insertions(+), 20 deletions(-) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index d4293e5..da6c078 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -134,12 +134,17 @@ def _numpy_dtype(cls, dtype): 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 cast python scalar for const, a device + write for scalar, a full host->device copy for field. + Author: B.G (07/2026) """ if self.mode == "const": @@ -166,6 +171,9 @@ def set_node(self, node, value) -> None: 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: @@ -256,6 +264,8 @@ def __init__(self, name: str, 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 @@ -275,9 +285,8 @@ class ClosureKernel(Kernel): Kernel backed by a compiled closure-backend kernel (ti.kernel / qd.kernel). Built by ClosureKernelBuilder subclasses. - The template's own signature declares data-field arguments only (e.g. - `def template(out: ti.template()): ...`); params/helpers arrive via bind() - and are resolved into the kernel body, so callers pass data fields only. + The template's own signature declares data-field arguments only, e.g. + `def template(out: ti.template()): ...` - see base.py's module docstring. Author: B.G (07/2026) """ @@ -289,6 +298,8 @@ def __init__(self, name: str, compiled): @property def compiled(self): """ + The raw ti.kernel/qd.kernel behind this Kernel's __call__. + Author: B.G (07/2026) """ return self._compiled @@ -335,6 +346,9 @@ class ClosureDeviceFunctionBuilder(DeviceFunctionBuilder): def compile(self) -> ClosureDeviceFunction: """ + Check the const-only rule, inject the referenced bindings into the + template's globals, decorate the result as a device func. + Author: B.G (07/2026) """ _check_const_only(self._bindings) @@ -356,6 +370,9 @@ class ClosureKernelBuilder(KernelBuilder): def compile(self) -> ClosureKernel: """ + Inject the referenced bindings into the template's globals and + decorate the result as a launchable kernel. + Author: B.G (07/2026) """ specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings)) diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 8b251ab..3bf837f 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -160,8 +160,8 @@ class Kernel(Specializable): Compiled entry point (e.g. a ti.kernel specialization). Unlike DeviceFunction, __call__ works from host Python - that's how the - compute gets launched. Its template's own parameters are data fields only; - params/helpers arrive via bindings and are resolved into the kernel body. + compute gets launched. Call-time arguments are data fields only, per the + module docstring's call-time/compile-time rule. Author: B.G (07/2026) """ @@ -245,7 +245,7 @@ def capture_template_meta(template) -> tuple[str | None, ast.AST | None]: return source, tree -def used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: +def _used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: """ Subset of `bindings` whose key is actually referenced by `template`'s body (by name - an attribute chain's root, e.g. `phys` in @@ -265,15 +265,15 @@ def used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: def filter_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: """ - used_bindings(), plus one warning per compile listing everything bound but - never referenced - catches typos in a large context. Backend-agnostic: the - closure builders feed the result to their specializer. No warning when the - AST isn't available (used_bindings then returns `bindings` unchanged, so - nothing looks unused). + Bindings the template actually references (see _used_bindings), plus one + warning per compile listing everything bound but never referenced - catches + typos in a large context. Backend-agnostic: the closure builders feed the + result to their specializer. No warning when the AST isn't available + (_used_bindings then returns `bindings` unchanged, so nothing looks unused). Author: B.G (07/2026) """ - filtered = used_bindings(template, bindings) + filtered = _used_bindings(template, bindings) unused = sorted(set(bindings) - set(filtered)) if unused: warnings.warn( @@ -423,14 +423,6 @@ def __iter__(self): def items(self): return self._items.items() - def as_bindings(self) -> dict[str, Any]: - """ - Return {name: item} for merging into a builder's bindings. - - Author: B.G (07/2026) - """ - return dict(self._items) - class ParamBag(Bag): """ diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index a6b9b78..30aeb5f 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -83,6 +83,9 @@ def _cuda_literal(value) -> str: 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.RawKernel is looked up by. + Author: B.G (07/2026) """ match = pattern.search(template) @@ -192,6 +195,9 @@ def _repl(self, match: re.Match) -> str: def parse(self, body: str) -> str: """ + Expand every `$...$` span in `body`, accumulating ptr_params and + helper_srcs as a side effect. + Author: B.G (07/2026) """ return _SPAN_RE.sub(self._repl, body) @@ -252,6 +258,11 @@ class CupyParameter(Parameter): def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): """ + Declare and initialize one parameter. "scalar"/"field" modes allocate + pooled storage immediately via `pool`; "const" stays a plain python + value. solo=True (const only) lets the parameter be read bare in a + template body - it becomes a #define rather than a span expansion. + Author: B.G (07/2026) """ if mode not in self.SUPPORTED_MODES: @@ -278,12 +289,17 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No 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 cast python scalar for const, a device + write for scalar, a full host->device copy for field. + Author: B.G (07/2026) """ if self.mode == "const": @@ -309,6 +325,9 @@ def set_node(self, node, value) -> None: 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: @@ -336,12 +355,18 @@ def __init__(self, name: str, source: str): @property def compiled(self): """ + The spliced `__device__` source text, for pasting into whatever + kernel or device function binds this helper. + 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( @@ -372,6 +397,8 @@ def __init__(self, name: str, compiled, bound_arrays: list): @property def compiled(self): """ + The underlying cp.RawKernel this Kernel's __call__ launches. + Author: B.G (07/2026) """ return self._compiled @@ -398,6 +425,9 @@ class CupyDeviceFunctionBuilder(DeviceFunctionBuilder): def compile(self) -> CupyDeviceFunction: """ + Expand the template's spans and prepend the helper sources and const + #defines it turned out to need, giving the final `__device__` text. + Author: B.G (07/2026) """ template = self._template @@ -421,6 +451,10 @@ class CupyKernelBuilder(KernelBuilder): def compile(self) -> CupyKernel: """ + Expand the template's spans, inject the pointer args they implied into + the __global__ signature, prepend helpers + const #defines, then build + (or reuse, keyed by final source text) the cp.RawKernel. + Author: B.G (07/2026) """ template = self._template From e85216d42a25e52143d15215f1280b9e1c801d19 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 14:54:27 +0200 Subject: [PATCH 11/31] Improve cache and a bit of cleaning --- pyfastflow/experimental/core/context/base.py | 32 ++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 3bf837f..6d00a04 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -17,6 +17,24 @@ all three backends alike. Kernels are unrestricted: they may bind params of any mode. +Staleness contract: compile() freezes its bindings. A const-mode Parameter is +baked into the compiled body (as a literal when solo, as a device view holding +a literal otherwise), and a scalar/field Parameter is baked as its DataHandle's +backing storage. Nothing propagates a later change back into an already +compiled Kernel/DeviceFunction: + + - set() on a const param, and destroy() on any param, drop that param's + cached device_view so the *next* compile() sees the change. Kernels + compiled before it keep the old value - recompile them to pick it up. + - destroy() also returns the handle to the pool, which may hand the same + buffer to another parameter, while kernels compiled earlier still write + through it. Do not destroy() a parameter that any live kernel binds. + - set() on scalar/field writes through the handle a compiled kernel already + points at, so those stay live and need no recompile - that is the normal + way to feed changing data without recompiling. + +None of this is checked at runtime: it is a convention, not an enforcement. + The compile surface is a two-layer builder: an abstract DeviceFunctionBuilder / KernelBuilder here, backend variants (Taichi*/Quadrants*/Cupy*) elsewhere. A builder collects bind()ed dependencies + one ingest()ed template and produces a @@ -80,6 +98,8 @@ def get(self): def set(self, value) -> None: """ Update the whole parameter value in place, according to its mode. + On const mode this does not reach already-compiled kernels - see the + module docstring's staleness contract. Author: B.G (07/2026) """ @@ -108,7 +128,9 @@ def device_view(self): @abstractmethod def destroy(self) -> None: """ - Release any backing storage owned by this parameter. + Release any backing storage owned by this parameter. Unsafe while any + compiled kernel still binds it - see the module docstring's staleness + contract. Author: B.G (07/2026) """ @@ -219,7 +241,7 @@ def resolve_binding(value): return value -@lru_cache(maxsize=None) +@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; @@ -230,6 +252,12 @@ def capture_template_meta(template) -> tuple[str | None, ast.AST | None]: The returned tree is therefore SHARED by every Specializable built from that template - treat it as read-only. + Bounded on purpose: the key is the template object itself, so an unbounded + cache would pin every dynamically generated template (and every CUDA source + string) for the process lifetime. 256 is far above any realistic count of + distinct templates alive at once, and this is a pure derived-value cache - + an eviction only costs one re-parse. + Author: B.G (07/2026) """ if isinstance(template, str): From 479624e7fb16bb28317a8a350152eecc799c9a15 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 14:57:52 +0200 Subject: [PATCH 12/31] Add a docstring warning about offline fastcache for quadrants backend - not possible because my param injections are ot visible by the latter --- .../core/context/quadrants_backend.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py index 095b7e8..cde7015 100644 --- a/pyfastflow/experimental/core/context/quadrants_backend.py +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -10,6 +10,22 @@ field-mode Parameters must be field-backed (they close over the field as a global): Quadrants rejects ndarrays referenced as globals inside a func. +DO NOT enable Quadrants' src_ll fast cache (`@qd.pure`, or +`qd.kernel(fastcache=True)`) on templates compiled through this framework - +not even as a way to skip the python-side AST trace. That cache keys on the +kernel's *source text*, re-read from disk by file path and line range +(_fast_caching/function_hasher.py), plus arg and config hashes. It never looks +at __globals__ - and injecting globals is precisely how specialize_closure +distinguishes one specialization from another. Two compiles of the same +template with different bound consts (or a different helper under the same +name) therefore produce the same key, and a fast-cache hit skips AST +transformation entirely, so nothing downstream catches the mismatch. + +Currently inert: that path is gated on is_pure, and the builders here call +plain qd.kernel / qd.func. The two IR-keyed caches (Quadrants' offline_cache, +and Taichi's) are safe by contrast - baked literals are part of the IR they +hash, so they discriminate specializations correctly. + Author: B.G (07/2026) """ From 565bc7789caa9678a0952453411538b84f69587b Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 15:25:42 +0200 Subject: [PATCH 13/31] Rework the comments (some were polluted by a LLM) --- .../core/context/_closure_backend.py | 130 +++++---- pyfastflow/experimental/core/context/base.py | 275 +++++++++++------- .../experimental/core/context/cupy_backend.py | 110 +++---- .../core/context/quadrants_backend.py | 65 ++--- .../core/context/taichi_backend.py | 17 +- 5 files changed, 335 insertions(+), 262 deletions(-) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index da6c078..cbcdcf8 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -1,7 +1,19 @@ """ -Shared machinery for backends that compile python function templates by -patching globals on a cloned code object (Taichi, Quadrants). Not used by -backends without that mechanism (e.g. the cupy/RawKernel backend). +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 ParamBag 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) """ @@ -26,8 +38,13 @@ def specialize_closure(template, bindings: dict[str, Any]) -> FunctionType: """ - Clone `template`'s code object with a globals dict where sentinels are - replaced by resolved bindings. + 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. Author: B.G (07/2026) """ @@ -52,11 +69,12 @@ def specialize_closure(template, bindings: dict[str, Any]) -> FunctionType: class ClosureParamDeviceView: """ - Device-facing view of a Parameter for closure backends: `.get` and (unless - const) `.set_node` are the raw compiled device funcs (ti.func / qd.func), - so a template body traces `p.get(i)` / `p.set_node(i, v)` as plain - attribute lookups + func calls, uniform across modes. Const params expose - no `.set_node` (read-only), so touching it in a kernel raises at trace time. + 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) """ @@ -70,12 +88,11 @@ def __init__(self, name: str, get_fn, set_fn=None): class ClosureBackendParameter(Parameter): """ - Parameter backed by a const value or a pooled DataHandle, for backends - sharing the closure-specialization mechanism (Taichi, Quadrants). - Subclasses pin `_backend` (the `ti`/`qd` module) - `_numpy_dtype` and - `_build_device_view` are shared here since `ti.u8/i32/i64`, - `ti.func`/`ti.static` and their `qd.` equivalents are identical in name - and behaviour across both modules. + 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) """ @@ -85,10 +102,12 @@ class ClosureBackendParameter(Parameter): def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): """ - Declare and initialize one parameter. "scalar"/"field" modes allocate - pooled storage immediately via `pool`; "const" stays a plain python - value. solo=True (const only) lets the parameter be read bare in a - template body (no .get()) - it resolves to a compile-time literal. + Declare one parameter and give it its initial value. + + scalar and field take pooled storage straight away; const stays a plain + python value. solo=True, available on const only, lets the parameter be + read bare in a template body - written `p` rather than `p.get(i)` - + because it resolves to a compile-time literal. Author: B.G (07/2026) """ @@ -183,14 +202,15 @@ def destroy(self) -> None: def device_view(self) -> ClosureParamDeviceView: """ - Cached uniform device accessor for this parameter. Built once via - `_build_device_view` (backend-specific) and memoized on the instance: - rebuilding it per compile is pointless python work at O(params x - kernels) since the ti.func/qd.func objects are identical every time. - Invalidated (not refreshed) by set() on const mode and by destroy(), - the two places the baked literal/handle a cached view depends on can - change; scalar/field set() writes through the handle a cached view - already points at, so no invalidation is needed there. + 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. Two things can invalidate it - + set() on a const mode, which changes the literal baked into the getter, + and destroy(), which releases the storage it reads. Both drop the view + so the next caller rebuilds; neither reaches kernels compiled earlier + (see base.py, "Lifetime of a compiled object"). A scalar or field set() + needs no invalidation, writing through the very storage the view reads. Author: B.G (07/2026) """ @@ -200,19 +220,13 @@ def device_view(self) -> ClosureParamDeviceView: def _build_device_view(self) -> ClosureParamDeviceView: """ - Compile this parameter's uniform device accessors as backend funcs - (ti.func / qd.func). - - get(node) dispatches on mode at trace time via `_backend.static`, so - only the taken arm compiles: const returns a baked literal, scalar - reads HANDLE[None], field reads HANDLE[node]. set_node(node, val) is - built only for scalar/field (const is read-only, exposes no setter). - MODE/VALUE/HANDLE are plain python values, not backend values. + Compile this parameter's device accessors as backend funcs. - Const getters bake VALUE as a compile-time literal: a later .set() - needs the view (and anything binding it) rebuilt to take effect. - - Called at most once per instance, memoized by device_view(). + 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) """ @@ -251,8 +265,7 @@ def set_node_template(node, val): class ClosureDeviceFunction(DeviceFunction): """ - DeviceFunction backed by a compiled closure-backend func (ti.func / - qd.func). Built by ClosureDeviceFunctionBuilder subclasses. + A device helper compiled to a ti.func or qd.func. Author: B.G (07/2026) """ @@ -282,11 +295,11 @@ def __call__(self, *args, **kwargs): class ClosureKernel(Kernel): """ - Kernel backed by a compiled closure-backend kernel (ti.kernel / - qd.kernel). Built by ClosureKernelBuilder subclasses. + A launchable kernel compiled to a ti.kernel or qd.kernel. - The template's own signature declares data-field arguments only, e.g. - `def template(out: ti.template()): ...` - see base.py's module docstring. + 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 base.py. Author: B.G (07/2026) """ @@ -315,10 +328,12 @@ def __call__(self, *args, **kwargs): def _check_const_only(bindings: dict[str, Any]) -> None: """ - Enforce the framework rule (base.py module docstring): a device helper - only ever binds const-mode Parameters - any data it needs is passed to it - as an explicit argument by the calling kernel instead. Recurses into - bound Bags so a scalar/field param nested inside one is also caught. + Raise unless every Parameter reachable in `bindings` is const mode. + + This enforces the rule in base.py that a device helper binds const + parameters only, passing any data through explicit arguments instead. + Bound Bags are searched too, so a scalar or field parameter tucked inside + one does not slip past. Author: B.G (07/2026) """ @@ -336,8 +351,7 @@ def _check_const_only(bindings: dict[str, Any]) -> None: class ClosureDeviceFunctionBuilder(DeviceFunctionBuilder): """ - Builds a ClosureDeviceFunction: specialize the ingested def with bound - globals, decorate with `_backend.func`. Subclasses pin `_backend`. + Compiles an ingested def into a device helper. Subclasses pin `_backend`. Author: B.G (07/2026) """ @@ -346,8 +360,8 @@ class ClosureDeviceFunctionBuilder(DeviceFunctionBuilder): def compile(self) -> ClosureDeviceFunction: """ - Check the const-only rule, inject the referenced bindings into the - template's globals, decorate the result as a device func. + Check the const-only rule, splice the referenced bindings into the + template's globals, and compile the result as a device func. Author: B.G (07/2026) """ @@ -360,8 +374,8 @@ def compile(self) -> ClosureDeviceFunction: class ClosureKernelBuilder(KernelBuilder): """ - Builds a ClosureKernel: specialize the ingested def with bound globals, - decorate with `_backend.kernel`. Subclasses pin `_backend`. + Compiles an ingested def into a launchable kernel. Subclasses pin + `_backend`. Author: B.G (07/2026) """ @@ -370,8 +384,8 @@ class ClosureKernelBuilder(KernelBuilder): def compile(self) -> ClosureKernel: """ - Inject the referenced bindings into the template's globals and - decorate the result as a launchable kernel. + Splice the referenced bindings into the template's globals and compile + the result as a launchable kernel. Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 6d00a04..191ff37 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -1,46 +1,78 @@ """ -Backend-agnostic context building blocks. - -No "Context" class here on purpose: a context is just whatever concrete class -(GridContext, FlowContext, ...) groups a set of Parameters and registers -DeviceFunctions. Cross-context references are plain explicit bindings passed -at bind() time, not a stored connection registry. - -Load-bearing rule: data flows at call time (a template's own explicit -arguments); configuration flows at compile time (bind()ed Parameters/helpers/ -bags, resolved into the template body). - -A device helper (DeviceFunction) may only bind const-mode Parameters. Any -data it needs (scalar/field buffers) is passed to it as an explicit argument -by the calling kernel, not bound in - a device helper has no way to receive -an extra pointer argument once spliced into a caller, so this holds across -all three backends alike. Kernels are unrestricted: they may bind params of -any mode. - -Staleness contract: compile() freezes its bindings. A const-mode Parameter is -baked into the compiled body (as a literal when solo, as a device view holding -a literal otherwise), and a scalar/field Parameter is baked as its DataHandle's -backing storage. Nothing propagates a later change back into an already -compiled Kernel/DeviceFunction: - - - set() on a const param, and destroy() on any param, drop that param's - cached device_view so the *next* compile() sees the change. Kernels - compiled before it keep the old value - recompile them to pick it up. - - destroy() also returns the handle to the pool, which may hand the same - buffer to another parameter, while kernels compiled earlier still write - through it. Do not destroy() a parameter that any live kernel binds. - - set() on scalar/field writes through the handle a compiled kernel already - points at, so those stay live and need no recompile - that is the normal - way to feed changing data without recompiling. - -None of this is checked at runtime: it is a convention, not an enforcement. - -The compile surface is a two-layer builder: an abstract DeviceFunctionBuilder / -KernelBuilder here, backend variants (Taichi*/Quadrants*/Cupy*) elsewhere. A -builder collects bind()ed dependencies + one ingest()ed template and produces a -compiled DeviceFunction/Kernel. Bound Parameters/helpers/bags are injected into -the template body (never passed at call time); only a template's own explicit -data-field arguments are passed to the front callable. +Backend-agnostic building blocks for describing GPU work once and compiling it +against Taichi, Quadrants or cupy (or any future one). + +Jargon +---------- +Parameter One named, typed value. Its `mode` says where the value lives: + "const" (compile-time constant, non modifiable mid-run), + "scalar" (single value modifiable) or "field" (a device array). +DeviceFunction A compiled device-side helper: a small routine callable only + from other device code (ti.func, qd.func, CUDA __device__). +Kernel A compiled entry point - what the host launches (ti.kernel, + qd.kernel, CUDA __global__). +Bag A named collection of Parameters (ParamBag) or DeviceFunctions + (HelperBag), so a group travels as one object and is read + in-kernel by dotted path: phys.dx.get(i). + +A "context" is any concrete class - GridContext, FlowContext, ... - that groups +Parameters and registers DeviceFunctions. 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 ParamBag + .bind("ops", ops) # a HelperBag + .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 or DeviceFunction. Only the +abstract DeviceFunctionBuilder / KernelBuilder live here; the concrete +Taichi*, Quadrants* and Cupy* builders sit alongside this module. + +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. A const Parameter declared solo=True +is the exception: it resolves to a bare compile-time literal, read as `p` with +no call. + +Device helpers bind const parameters only +----------------------------------------- +A DeviceFunction may only bind const-mode Parameters; any data it needs is +passed to it as an explicit argument by the calling kernel. A helper is +spliced into its caller and has no way to acquire a pointer argument of its +own, so this holds on all three backends alike. It is checked at compile time. +Kernels carry no such restriction and may bind any mode. + +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. So: + + - Writing to a scalar or field Parameter *is* visible to already-compiled + kernels, which hold that same storage. This is the normal way to feed + changing data. + - set() on a const Parameter is not. It drops the parameter's cached device + view so the next compile() picks the new value up, but kernels compiled + before it keep the old literal. Recompile them. + - destroy() returns storage to the pool, which may hand the same buffer out + again. Never destroy a Parameter that a live kernel still binds. + +None of this is enforced at runtime. Author: B.G (07/2026) """ @@ -59,14 +91,16 @@ class Parameter(ABC): """ One named, typed value owned by a context. - REQUIRED_MODES is the baseline every backend must support; a backend - widens SUPPORTED_MODES to add more storage kinds. Enforced at subclass - definition time via __init_subclass__, not at instantiation. + `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 must offer all three (REQUIRED_MODES). A backend may widen + SUPPORTED_MODES with further storage kinds; the check runs when the + subclass is defined, so an incomplete backend fails at import. - Host surface: get() / set(value) / set_node(node, value). Device surface: - device_view(), returning a backend object whose .get(node) / .set_node( - node, val) are usable from inside device code - the uniform accessor that - lets one kernel read/write a Parameter identically no matter its mode. + 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) """ @@ -97,9 +131,9 @@ def get(self): @abstractmethod def set(self, value) -> None: """ - Update the whole parameter value in place, according to its mode. - On const mode this does not reach already-compiled kernels - see the - module docstring's staleness contract. + Update the whole parameter value in place, according to its mode. On + const mode this does not reach already-compiled kernels - see the + module docstring, "Lifetime of a compiled object". Author: B.G (07/2026) """ @@ -117,9 +151,10 @@ def set_node(self, node, value) -> None: def device_view(self): """ - Return a backend device-view object (get/set_node usable in device - code). Implemented per backend; closure backends compile ti/qd funcs, - cupy returns parser metadata. + 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) """ @@ -128,9 +163,9 @@ def device_view(self): @abstractmethod def destroy(self) -> None: """ - Release any backing storage owned by this parameter. Unsafe while any - compiled kernel still binds it - see the module docstring's staleness - contract. + 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) """ @@ -139,9 +174,12 @@ def destroy(self) -> None: class Specializable(ABC): """ - Shared call contract for DeviceFunction and Kernel, plus the hidden raw - material (source text, AST, dependency manifest) a later fusion/graph - compiler can reuse. Built by a *Builder, not by a classmethod. + Anything a builder can compile: DeviceFunction or Kernel. + + Beyond the call contract, each instance keeps the raw material it was made + from - source text, AST, and a manifest of what it bound. Nothing in this + module reads those back; they are here so a higher layer (kernel fusion, a + graph compiler) can work from the originals instead of re-deriving them. Author: B.G (07/2026) """ @@ -181,9 +219,9 @@ class Kernel(Specializable): """ Compiled entry point (e.g. a ti.kernel specialization). - Unlike DeviceFunction, __call__ works from host Python - that's how the - compute gets launched. Call-time arguments are data fields only, per the - module docstring's call-time/compile-time rule. + Unlike DeviceFunction, __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) """ @@ -191,13 +229,15 @@ class Kernel(Specializable): class _LazyBagView: """ - Lazy stand-in for a Bag in a template body: `phys.g` resolves `g` only on - first attribute access (via resolve_binding), then caches it in the - instance dict so later lookups bypass __getattr__ entirely. Avoids - filter_bindings' top-level-only reach: binding one large bag no longer - eagerly resolves (and device_view()-compiles) every member, only the ones - a template actually touches. Nested Bag members resolve to another - _LazyBagView, so nesting stays lazy all the way down. + 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 - 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. Author: B.G (07/2026) """ @@ -217,15 +257,14 @@ def __getattr__(self, name: str) -> Any: def resolve_binding(value): """ - Unwrap a bound object to what a closure-backend template body should see. + Turn a bound object into what a template body should see in its place. - Parameter -> literal if solo (const), else device_view() (a .get/.set_node - carrier for uniform in-kernel access). - Specializable (DeviceFunction/Kernel) -> .compiled. - Bag -> a _LazyBagView resolving each member on first attribute access, so - dotted paths (grid.nx.get(i)) trace as plain attribute lookups without - forcing every other member in the bag to resolve too. - Anything else passes through unchanged. + Parameter a bare literal when solo, otherwise device_view() - the + carrier of .get / .set_node for in-kernel access. + Specializable its .compiled backend callable or source. + Bag a _LazyBagView, so a dotted path like grid.nx.get(i) traces + as plain attribute lookups. + anything else passed through untouched. Author: B.G (07/2026) """ @@ -244,19 +283,18 @@ def resolve_binding(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 with no AST (not python). + 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 per template: every compile() asks twice (once to filter bindings, - once for attach_meta), and each miss costs an inspect.getsource + a parse. - The returned tree is therefore SHARED by every Specializable built from - that template - treat it as read-only. + Cached because every compile() asks twice - once to filter bindings, once + for attach_meta - 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. - Bounded on purpose: the key is the template object itself, so an unbounded - cache would pin every dynamically generated template (and every CUDA source - string) for the process lifetime. 256 is far above any realistic count of - distinct templates alive at once, and this is a pure derived-value cache - - an eviction only costs one re-parse. + 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) """ @@ -275,12 +313,14 @@ def capture_template_meta(template) -> tuple[str | None, ast.AST | None]: def _used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: """ - Subset of `bindings` whose key is actually referenced by `template`'s - body (by name - an attribute chain's root, e.g. `phys` in - `phys.g.get(0)`, is itself an `ast.Name`, so collecting every `ast.Name` - id in the tree is sufficient). Falls back to returning `bindings` - unchanged whenever the AST isn't available (non-python template, source - unavailable) - never silently drop a binding we cannot prove is unused. + 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) """ @@ -293,11 +333,12 @@ def _used_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: def filter_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: """ - Bindings the template actually references (see _used_bindings), plus one - warning per compile listing everything bound but never referenced - catches - typos in a large context. Backend-agnostic: the closure builders feed the - result to their specializer. No warning when the AST isn't available - (_used_bindings then returns `bindings` unchanged, so nothing looks unused). + 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) """ @@ -314,8 +355,8 @@ def filter_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: def attach_meta(obj: Specializable, template, bindings: dict[str, Any]) -> None: """ - Stash hidden raw material on a freshly built Specializable for later reuse - by a higher-level compiler (kernel fusion / graph building). + Record on a freshly built Specializable what it was made from: its source, + its AST, and the type name of each thing it bound. See Specializable. Author: B.G (07/2026) """ @@ -325,11 +366,16 @@ def attach_meta(obj: Specializable, template, bindings: dict[str, Any]) -> None: class CompileBuilder(ABC): """ - Two-layer compile surface: collect bind()ed dependencies + one ingest()ed - template, then compile() to a backend Specializable. bind() detects the - kind of each object at resolution time (Parameter / DeviceFunction / Bag / - handle / plain value); backends implement compile() (and may override - ingest()) their own way. + Collects dependencies and a template, and compiles them into one + Specializable. + + A builder is used once, 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, DeviceFunction, Bag, handle, plain value) is + worked out when the template is specialized, so bind() accepts anything. + + Everything here is backend-independent. A backend supplies compile(), and + may override ingest() if its templates need different handling. Author: B.G (07/2026) """ @@ -399,12 +445,15 @@ class KernelBuilder(CompileBuilder): class Bag: """ - Simple named collection, mergeable into a builder via bind_bag() or bound - whole (bind('grid', bag)) for dotted-path access in the template body. + A named collection that can be handed to a builder in one go. + + Two ways to use one: bind it whole - bind("grid", bag) - and reach its + members by dotted path in the template body, or bind_bag(bag) to merge + every member in at top level under its own name. - `_member_type` (None on the base class) restricts what a subclass's - members may be; a nested Bag is always accepted regardless, since bags - nest (a ParamBag may legitimately contain another bag of params). + A subclass sets `_member_type` to restrict what it will hold. Bags are + always accepted whatever that restriction says, since bags nest: a ParamBag + may legitimately group its parameters into sub-bags. Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 30aeb5f..eba6fe4 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -1,34 +1,34 @@ """ -Cupy backend implementation of Parameter, DeviceFunction, Kernel and their -builders. - -Templates are raw CUDA source text, not python callables - there is no -closure mechanism for cp.RawKernel. Params/helpers are referenced inside the -source through `$...$` spans holding a dotted path, uniform with the closure -backends' in-kernel API: - - $p.get(i)$ -> read param p at flat index i - $p.set_node(i,v)$ -> write param p at flat index i (device-side) - $grid.nx.get(i)$ -> bag member access (dotted head) - $helper(a, b)$ -> call a bound CupyDeviceFunction - -The parser expands each span AND, for scalar/field params, auto-generates the -matching pointer argument into the __global__ signature plus the launch-time -array - the source never hand-declares those. Expansion by mode: - - const -> a CUDA literal (and a `#define NAME literal` for bare use of a - top-level const param outside any span) - - scalar -> NAME[0] (a 0-d device pointer) - - field -> NAME[i] -set_node on a const is a compile error; a param read only -> `const T*`, a -param written anywhere in the kernel -> non-const `T*`. - -Device functions may only reference const params / other device functions -inside spans - a spliced __device__ function has no way to receive an extra -pointer argument. This isn't a cupy quirk: it's the backend-agnostic rule -(see base.py's module docstring) that a device helper only ever binds const -params, with any data it needs passed in as an explicit argument by the -calling kernel - the closure backends (Taichi, Quadrants) enforce the same -rule in ClosureDeviceFunctionBuilder.compile(). +cupy implementations of Parameter, DeviceFunction, Kernel and their builders. + +Here a template is CUDA source text rather than a python function, since +cp.RawKernel 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, NAME[0] for scalar, NAME[i] for field - and, for scalar and +field, also generates the pointer argument that expansion implies. Those +arguments are appended to the __global__ signature and their arrays supplied at +launch, so a template never declares them by hand. A parameter only read is +declared `const T*`; one written anywhere in the kernel is declared `T*`. +Writing to a const parameter is an error. + +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. + +Spans inside a device function may only reach const parameters and other +device helpers. That is the framework rule from base.py rather than anything +specific to cupy: a helper is spliced into its caller and cannot take on a +pointer argument of its own. Author: B.G (07/2026) """ @@ -130,9 +130,13 @@ def _walk(path: list[str], bindings: dict[str, Any]): class _SpanParser: """ - Expands `$...$` spans in a CUDA template body, accumulating the pointer - params and helper sources they imply. allow_arrays=False (device - functions) rejects scalar/field params and set_node. + Expands the `$...$` spans in a CUDA template body. + + Expansion has side effects the builder needs afterwards: `ptr_params` + collects the pointer arguments the spans implied, and `helper_srcs` the + source of every device helper they called. Set allow_arrays=False when + parsing a device function, which may not reach scalar or field parameters + and so has no use for either. Author: B.G (07/2026) """ @@ -205,13 +209,14 @@ def parse(self, body: str) -> str: def _const_defines(bindings: dict[str, Any], body: str) -> list[str]: """ - `#define NAME literal` for each top-level const Parameter whose identifier - actually appears (word-boundary match) in `body` - `body` must be the - already span-expanded text, so a span like `$phys.dx.get(0)$` has already - become a numeric literal and correctly does not count as a bare use of a - top-level const name. Skipping unused names avoids pasting unhygienic - macros (a const named N, I, DIM, EPS, min, ...) into the translation unit - where they'd silently rewrite unrelated identifiers. + 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) """ @@ -248,8 +253,9 @@ class CupyParameter(Parameter): """ Parameter backed by a const python value or a pooled CupyDataHandle. - Cupy dtypes are numpy dtypes already, so no dtype-mapping hook is needed. - Resolved into templates by the `$...$` parser (not device_view). + 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) """ @@ -337,11 +343,11 @@ def destroy(self) -> None: class CupyDeviceFunction(DeviceFunction): """ - DeviceFunction backed by a CUDA `__device__` function's source text. + A device helper, held as CUDA `__device__` source text. - There's no separately-compiled device function in the RawKernel model - - `.compiled` returns source, spliced into whatever kernel/device function - binds it. Never callable from host Python. + Nothing is compiled at this stage: RawKernel 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) """ @@ -376,13 +382,15 @@ def __call__(self, *args, **kwargs): class CupyKernel(Kernel): """ - Kernel backed by a compiled cp.RawKernel. + A launchable kernel, backed by a compiled cp.RawKernel. - __call__ requires explicit grid/block launch dims - cp.RawKernel has no - auto-ranging. Call-time positional args are the kernel's own data-field - arguments; the parser-generated pointer arrays are appended after them. + Launch dimensions are explicit: RawKernel has nothing like Taichi's + auto-ranging, so __call__ takes grid and block. The positional arguments + are the template's own data arguments, and the arrays behind the generated + pointer parameters follow them. - compile() caches the RawKernel by final spliced source text. + Compiled kernels are cached on the class by final source text, which is + what the whole specialization reduces to and therefore a sound key. Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py index cde7015..104dc40 100644 --- a/pyfastflow/experimental/core/context/quadrants_backend.py +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -1,30 +1,31 @@ """ -Quadrants backend implementation of Parameter, DeviceFunction, Kernel and -their builders. - -Shares the closure-specialization mechanism with taichi_backend.py via -_closure_backend.py. One difference specific to this backend: Kernel -templates may type their own data-field arguments as `qd.Tensor`, which -accepts either a field- or ndarray-backed value at call time with no change -to the compiled template - Taichi has no equivalent for this. Note that -field-mode Parameters must be field-backed (they close over the field as a -global): Quadrants rejects ndarrays referenced as globals inside a func. - -DO NOT enable Quadrants' src_ll fast cache (`@qd.pure`, or -`qd.kernel(fastcache=True)`) on templates compiled through this framework - -not even as a way to skip the python-side AST trace. That cache keys on the -kernel's *source text*, re-read from disk by file path and line range -(_fast_caching/function_hasher.py), plus arg and config hashes. It never looks -at __globals__ - and injecting globals is precisely how specialize_closure -distinguishes one specialization from another. Two compiles of the same -template with different bound consts (or a different helper under the same -name) therefore produce the same key, and a fast-cache hit skips AST -transformation entirely, so nothing downstream catches the mismatch. - -Currently inert: that path is gated on is_pure, and the builders here call -plain qd.kernel / qd.func. The two IR-keyed caches (Quadrants' offline_cache, -and Taichi's) are safe by contrast - baked literals are part of the IR they -hash, so they discriminate specializations correctly. +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) """ @@ -50,10 +51,8 @@ class QuadrantsParameter(ClosureBackendParameter): class QuadrantsDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): """ - Builds a ClosureDeviceFunction: specialize the ingested def with bound - globals, decorate with qd.func. Only field-backed Parameters/DeviceFunctions - can be resolved into a template this way - Quadrants rejects ndarrays - referenced as globals. + 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) """ @@ -63,10 +62,8 @@ class QuadrantsDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): class QuadrantsKernelBuilder(ClosureKernelBuilder): """ - Builds a ClosureKernel: specialize the ingested def with bound globals, - decorate with qd.kernel. The template's own data-field arguments should be - typed `qd.Tensor` - a single qd.Tensor-typed template accepts either a - field- or ndarray-backed value at call time. + 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) """ diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py index f0da77d..d3e36b6 100644 --- a/pyfastflow/experimental/core/context/taichi_backend.py +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -1,6 +1,13 @@ """ -Taichi backend implementation of Parameter, DeviceFunction, Kernel and their -builders. +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 base.py. Author: B.G (07/2026) """ @@ -26,8 +33,7 @@ class TaichiParameter(ClosureBackendParameter): class TaichiDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): """ - Builds a ClosureDeviceFunction: specialize the ingested def with bound - globals, decorate with ti.func. + Compiles a device helper with ti.func. Author: B.G (07/2026) """ @@ -37,8 +43,7 @@ class TaichiDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): class TaichiKernelBuilder(ClosureKernelBuilder): """ - Builds a ClosureKernel: specialize the ingested def with bound globals, - decorate with ti.kernel. + Compiles a launchable kernel with ti.kernel. Author: B.G (07/2026) """ From 734eb9424a5a5bb56a6f2ff907de92ac41b07bf8 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 15:36:51 +0200 Subject: [PATCH 14/31] improve base.py docstring --- pyfastflow/experimental/core/context/base.py | 24 ++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 191ff37..1147b0d 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -2,11 +2,31 @@ 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" (compile-time constant, non modifiable mid-run), - "scalar" (single value modifiable) or "field" (a device array). + "const" (a compile-time constant, not modifiable mid-run), + "scalar" (a single value, modifiable) or "field" (a device + array, one value per node). DeviceFunction A compiled device-side helper: a small routine callable only from other device code (ti.func, qd.func, CUDA __device__). Kernel A compiled entry point - what the host launches (ti.kernel, From 8201eca79c099dfeca038926fcc5478fc617a7f5 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Thu, 23 Jul 2026 16:10:35 +0200 Subject: [PATCH 15/31] Merge the different bags in a single entity --- .../core/context/_closure_backend.py | 2 +- pyfastflow/experimental/core/context/base.py | 57 ++++++------------- 2 files changed, 18 insertions(+), 41 deletions(-) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index cbcdcf8..4273d85 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -4,7 +4,7 @@ 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 ParamBag when the backend traces it. The rebuilt function +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, diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 1147b0d..7255bc4 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -31,9 +31,9 @@ from other device code (ti.func, qd.func, CUDA __device__). Kernel A compiled entry point - what the host launches (ti.kernel, qd.kernel, CUDA __global__). -Bag A named collection of Parameters (ParamBag) or DeviceFunctions - (HelperBag), so a group travels as one object and is read - in-kernel by dotted path: phys.dx.get(i). +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 DeviceFunctions. There is deliberately no base Context @@ -45,8 +45,8 @@ Templates are written once, generically, and specialized by a builder: kernel = (TaichiKernelBuilder() - .bind("phys", phys) # a ParamBag - .bind("ops", ops) # a HelperBag + .bind("phys", phys) # a Bag of parameters + .bind("ops", ops) # a Bag of device helpers .ingest(update_height) # the template .compile()) kernel(h_new, h_old) # bulk data passed at call time @@ -467,19 +467,23 @@ 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, DeviceFunctions, 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, or bind_bag(bag) to merge - every member in at top level under its own name. + 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. - A subclass sets `_member_type` to restrict what it will hold. Bags are - always accepted whatever that restriction says, since bags nest: a ParamBag - may legitimately group its parameters into sub-bags. + Build it, grow it, bind it. There is no removal or reassignment: to change + the contents, build another bag. Author: B.G (07/2026) """ - _member_type: ClassVar[type | None] = None - def __init__(self, items: dict[str, Any] | None = None): self._items: dict[str, Any] = {} for name, item in (items or {}).items(): @@ -487,19 +491,12 @@ def __init__(self, items: dict[str, Any] | None = None): def add(self, name: str, item: Any) -> None: """ - Register `item` under `name`. Raises if `name` is already registered - or if `item` doesn't match this bag's `_member_type` (a nested Bag is - always accepted). + 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") - if self._member_type is not None and not isinstance(item, (self._member_type, Bag)): - raise TypeError( - f"{type(self).__name__}: member '{name}' must be a {self._member_type.__name__} " - f"(or a Bag), got {type(item).__name__}" - ) self._items[name] = item def __getattr__(self, name: str) -> Any: @@ -519,23 +516,3 @@ def __iter__(self): def items(self): return self._items.items() - - -class ParamBag(Bag): - """ - Named collection of Parameter objects. - - Author: B.G (07/2026) - """ - - _member_type = Parameter - - -class HelperBag(Bag): - """ - Named collection of DeviceFunction objects. - - Author: B.G (07/2026) - """ - - _member_type = DeviceFunction From 686bdfb2d2390e31f19727efe1054895f708cb8d Mon Sep 17 00:00:00 2001 From: bgailleton Date: Fri, 24 Jul 2026 12:35:54 +0200 Subject: [PATCH 16/31] Refactor to builder/composer state #1 --- .../core/context/_closure_backend.py | 2 + pyfastflow/experimental/core/context/base.py | 158 +++++++++++++++++- .../experimental/core/context/cupy_backend.py | 2 + .../experimental/core/pool/_fields_handle.py | 3 +- pyfastflow/experimental/core/pool/base.py | 34 +++- .../experimental/core/pool/cupy_handle.py | 3 +- 6 files changed, 197 insertions(+), 5 deletions(-) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 4273d85..8a14af5 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -116,6 +116,7 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No if solo and mode != "const": raise ValueError(f"{name}: solo access is const-only, got mode {mode!r}") + super().__init__() self.name = name self.dtype = dtype self.mode = mode @@ -271,6 +272,7 @@ class ClosureDeviceFunction(DeviceFunction): """ def __init__(self, name: str, compiled): + super().__init__() self.name = name self._compiled = compiled diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 7255bc4..301eb01 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -104,7 +104,7 @@ from functools import lru_cache from typing import Any, ClassVar -from ..pool.base import DataHandle +from ..pool.base import DataHandle, new_uid class Parameter(ABC): @@ -136,9 +136,54 @@ def __init_subclass__(cls, **kwargs): name: str dtype: Any - mode: str solo: bool = False + 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): """ @@ -234,6 +279,24 @@ class DeviceFunction(Specializable): Author: B.G (07/2026) """ + def __init__(self): + """ + Assign this helper'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 + class Kernel(Specializable): """ @@ -454,6 +517,19 @@ class DeviceFunctionBuilder(CompileBuilder): Author: B.G (07/2026) """ + def __init__(self): + super().__init__() + 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 + class KernelBuilder(CompileBuilder): """ @@ -485,10 +561,21 @@ class Bag: """ 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. @@ -516,3 +603,70 @@ def __iter__(self): def items(self): return self._items.items() + + 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}'" + ) diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index eba6fe4..e60a8d0 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -276,6 +276,7 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No if solo and mode != "const": raise ValueError(f"{name}: solo access is const-only, got mode {mode!r}") + super().__init__() self.name = name self.dtype = dtype self.mode = mode @@ -353,6 +354,7 @@ class CupyDeviceFunction(DeviceFunction): """ def __init__(self, name: str, source: str): + super().__init__() self.name = name # note: distinct from Specializable._source (the raw template, set by # attach_meta) - this is the spliced __device__ source `.compiled` serves. diff --git a/pyfastflow/experimental/core/pool/_fields_handle.py b/pyfastflow/experimental/core/pool/_fields_handle.py index c9d8a6d..9dbe18f 100644 --- a/pyfastflow/experimental/core/pool/_fields_handle.py +++ b/pyfastflow/experimental/core/pool/_fields_handle.py @@ -9,7 +9,7 @@ from typing import Any, ClassVar -from .base import DataHandle +from .base import DataHandle, new_uid class FieldsBuilderDataHandle(DataHandle): @@ -37,6 +37,7 @@ def __init__(self, dtype: Any, shape: tuple[int, ...]): cls = type(self) cls._next_id += 1 self.id = cls._next_id + self._uid = new_uid() self.dtype = dtype self.shape = tuple(shape) self.in_use = False diff --git a/pyfastflow/experimental/core/pool/base.py b/pyfastflow/experimental/core/pool/base.py index 4539378..c9536fc 100644 --- a/pyfastflow/experimental/core/pool/base.py +++ b/pyfastflow/experimental/core/pool/base.py @@ -8,9 +8,28 @@ 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): """ @@ -20,7 +39,11 @@ class DataHandle(ABC): pool for reuse without freeing memory; `destroy()` actually frees it. Attributes: - id: Unique handle identifier, assigned by the backend. + 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(). @@ -33,6 +56,15 @@ class DataHandle(ABC): 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): diff --git a/pyfastflow/experimental/core/pool/cupy_handle.py b/pyfastflow/experimental/core/pool/cupy_handle.py index 0d99d41..d8c4d96 100644 --- a/pyfastflow/experimental/core/pool/cupy_handle.py +++ b/pyfastflow/experimental/core/pool/cupy_handle.py @@ -8,7 +8,7 @@ import cupy as cp -from .base import DataHandle +from .base import DataHandle, new_uid class CupyDataHandle(DataHandle): @@ -28,6 +28,7 @@ def __init__(self, dtype: Any, shape: tuple[int, ...]): """ CupyDataHandle._next_id += 1 self.id = CupyDataHandle._next_id + self._uid = new_uid() self.dtype = dtype self.shape = tuple(shape) self.in_use = False From fb3c32e29e7215e7c8b23084a2cafe004a6ab86a Mon Sep 17 00:00:00 2001 From: bgailleton Date: Fri, 24 Jul 2026 13:50:05 +0200 Subject: [PATCH 17/31] Refactor to builder/composer state #2 --- pyfastflow/experimental/core/__init__.py | 2 +- .../core/context/_closure_backend.py | 55 +- pyfastflow/experimental/core/context/base.py | 593 +++++++++++++++--- .../experimental/core/context/cupy_backend.py | 63 +- .../core/context/quadrants_backend.py | 4 +- .../core/context/taichi_backend.py | 4 +- .../experimental/core/pool/_fields_handle.py | 2 +- pyfastflow/experimental/core/pool/base.py | 9 +- .../experimental/core/pool/cupy_handle.py | 2 +- 9 files changed, 611 insertions(+), 123 deletions(-) diff --git a/pyfastflow/experimental/core/__init__.py b/pyfastflow/experimental/core/__init__.py index af91431..0a2cdcd 100644 --- a/pyfastflow/experimental/core/__init__.py +++ b/pyfastflow/experimental/core/__init__.py @@ -1,5 +1,5 @@ """ -New backend-agnostic core (Parameter/DeviceFunction/Kernel/Pool ABCs + backends). +New backend-agnostic core (Parameter/Helper/Kernel/Pool ABCs + backends). Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 8a14af5..95fa8ae 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -25,18 +25,19 @@ from .base import ( Bag, - DeviceFunction, - DeviceFunctionBuilder, + HelperBuilder, Kernel, KernelBuilder, Parameter, + _SpecializedHelper, + _SpecializeCtx, attach_meta, filter_bindings, resolve_binding, ) -def specialize_closure(template, bindings: dict[str, Any]) -> FunctionType: +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. @@ -44,11 +45,13 @@ def specialize_closure(template, bindings: dict[str, Any]) -> FunctionType: 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. + 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) for name, value in bindings.items()} + 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) @@ -244,8 +247,13 @@ def get_template(node): 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}) + specialize_closure( + get_template, {"MODE": mode, "VALUE": value, "HANDLE": handle, "STATIC": backend.static}, _SpecializeCtx() + ) ) set_fn = None @@ -258,15 +266,17 @@ def set_node_template(node, val): HANDLE[node] = val set_fn = backend.func( - specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle, "STATIC": backend.static}) + specialize_closure(set_node_template, {"MODE": mode, "HANDLE": handle, "STATIC": backend.static}, _SpecializeCtx()) ) return ClosureParamDeviceView(self.name, get_fn, set_fn) -class ClosureDeviceFunction(DeviceFunction): +class ClosureHelper(_SpecializedHelper): """ - A device helper compiled to a ti.func or qd.func. + 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) """ @@ -292,7 +302,7 @@ def __call__(self, *args, **kwargs): Author: B.G (07/2026) """ - raise RuntimeError(f"DeviceFunction '{self.name}' is only callable from kernel/func scope, not host Python") + raise RuntimeError(f"Helper '{self.name}' is only callable from kernel/func scope, not host Python") class ClosureKernel(Kernel): @@ -307,6 +317,7 @@ class ClosureKernel(Kernel): """ def __init__(self, name: str, compiled): + super().__init__() self.name = name self._compiled = compiled @@ -351,25 +362,28 @@ def _check_const_only(bindings: dict[str, Any]) -> None: _check_const_only(dict(value.items())) -class ClosureDeviceFunctionBuilder(DeviceFunctionBuilder): +class ClosureHelperBuilder(HelperBuilder): """ - Compiles an ingested def into a device helper. Subclasses pin `_backend`. + 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 compile(self) -> ClosureDeviceFunction: + def _specialize(self, ctx: _SpecializeCtx) -> ClosureHelper: """ Check the const-only rule, splice the referenced bindings into the - template's globals, and compile the result as a device func. + template's globals against `ctx`, and compile the result as a device + func. Author: B.G (07/2026) """ _check_const_only(self._bindings) - specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings)) - fn = ClosureDeviceFunction(specialised.__name__, self._backend.func(specialised)) + specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) + fn = ClosureHelper(specialised.__name__, self._backend.func(specialised)) attach_meta(fn, self._template, self._bindings) return fn @@ -386,12 +400,15 @@ class ClosureKernelBuilder(KernelBuilder): def compile(self) -> ClosureKernel: """ - Splice the referenced bindings into the template's globals and compile - the result as a launchable kernel. + 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) """ - specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings)) + ctx = _SpecializeCtx() + specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) krn = ClosureKernel(specialised.__name__, self._backend.kernel(specialised)) attach_meta(krn, self._template, self._bindings) return krn diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 301eb01..dce759a 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -27,8 +27,11 @@ "const" (a compile-time constant, not modifiable mid-run), "scalar" (a single value, modifiable) or "field" (a device array, one value per node). -DeviceFunction A compiled device-side helper: a small routine callable only - from other device code (ti.func, qd.func, CUDA __device__). +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 @@ -36,7 +39,7 @@ path: phys.dx.get(i), ops.neighbour(i). A "context" is any concrete class - GridContext, FlowContext, ... - that groups -Parameters and registers DeviceFunctions. There is deliberately no base Context +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. @@ -46,16 +49,34 @@ kernel = (TaichiKernelBuilder() .bind("phys", phys) # a Bag of parameters - .bind("ops", ops) # a Bag of device helpers + .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 or DeviceFunction. Only the -abstract DeviceFunctionBuilder / KernelBuilder live here; the concrete -Taichi*, Quadrants* and Cupy* builders sit alongside this module. +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 ------------------------------------------------ @@ -72,11 +93,12 @@ Device helpers bind const parameters only ----------------------------------------- -A DeviceFunction may only bind const-mode Parameters; any data it needs is +A HelperBuilder may only bind const-mode Parameters; any data it needs is passed to it as an explicit argument by the calling kernel. A helper is spliced into its caller and has no way to acquire a pointer argument of its -own, so this holds on all three backends alike. It is checked at compile time. -Kernels carry no such restriction and may bind any mode. +own, so this holds on all three backends alike. It is checked when the helper +is specialized, i.e. as part of the enclosing kernel's compile(). Kernels +carry no such restriction and may bind any mode. Lifetime of a compiled object ----------------------------- @@ -239,20 +261,86 @@ def destroy(self) -> None: class Specializable(ABC): """ - Anything a builder can compile: DeviceFunction or Kernel. - - Beyond the call contract, each instance keeps the raw material it was made - from - source text, AST, and a manifest of what it bound. Nothing in this - module reads those back; they are here so a higher layer (kernel fusion, a - graph compiler) can work from the originals instead of re-deriving them. + 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(). + + Beyond the call contract, each instance keeps a read-only snapshot of the + raw material it was made from - template, source text, AST, and the real + bindings dict as it stood at this object's own compile time. Nothing in + this module reads those back to drive a later compile; they are here so a + higher layer (kernel fusion, a graph compiler) can work from the originals + instead of re-deriving them. The builder that produced this object stays + authoritative for anything that needs to change - see CompileBuilder. Author: B.G (07/2026) """ name: str + _template: Any = None _source: str | None = None _ast: ast.AST | None = None - _dependencies: dict[str, str] | None = None + _dependencies: dict[str, Any] | None = None + + 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 + def template(self): + """ + The template object this was compiled from - a python def for + Taichi/Quadrants, a CUDA source string for cupy. Read-only: this is a + snapshot for introspection, not a handle to recompile from. Change and + recompile through the builder instead. + + Author: B.G (07/2026) + """ + return self._template + + @property + def source(self) -> str | None: + """ + The template's source text, captured at compile time. See `template`. + + Author: B.G (07/2026) + """ + return self._source + + @property + def ast(self) -> "ast.AST | None": + """ + The template's parsed AST, captured at compile time. None for a + template with no recoverable source (e.g. cupy's CUDA text). See + `template`. + + Author: B.G (07/2026) + """ + return self._ast + + @property + def bindings(self) -> dict[str, Any]: + """ + The real bound objects - not type names - as they stood at this + object's own compile time. Read-only snapshot; see `template`. + + Author: B.G (07/2026) + """ + return self._dependencies @property @abstractmethod @@ -269,46 +357,81 @@ def compiled(self): def __call__(self, *args, **kwargs): ... -class DeviceFunction(Specializable): +class _SpecializedHelper(Specializable): """ - Compiled device-side helper (e.g. a ti.func specialization). + 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 - functions can only run inside kernel/func scope raise on __call__. + helpers can only run inside kernel/func scope raise on __call__. Author: B.G (07/2026) """ - def __init__(self): - """ - Assign this helper's process-wide uid. Concrete backends call this - first in their own __init__. - Author: B.G (07/2026) - """ - self._uid = new_uid() +class Kernel(Specializable): + """ + Compiled entry point (e.g. a ti.kernel specialization). - @property - def uid(self) -> int: - """ - Process-wide identity assigned at construction. See Parameter.uid. + 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) - """ - return self._uid + Author: B.G (07/2026) + """ -class Kernel(Specializable): +class _SpecializeCtx: """ - Compiled entry point (e.g. a ti.kernel specialization). - - Unlike DeviceFunction, __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. + 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: """ @@ -316,39 +439,51 @@ class _LazyBagView: `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 - 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. + 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"): + 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]) + resolved = resolve_binding(bag[name], ctx) self.__dict__[name] = resolved return resolved -def resolve_binding(value): +def resolve_binding(value, ctx: "_SpecializeCtx"): """ Turn a bound object into what a template body should see in its place. Parameter a bare literal when solo, otherwise 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, so a dotted path like grid.nx.get(i) traces - as plain attribute lookups. + 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): @@ -356,10 +491,12 @@ def resolve_binding(value): resolved = value.get() return resolved.data if isinstance(resolved, DataHandle) else resolved 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) + return _LazyBagView(value, ctx) return value @@ -438,13 +575,16 @@ def filter_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: def attach_meta(obj: Specializable, template, bindings: dict[str, Any]) -> None: """ - Record on a freshly built Specializable what it was made from: its source, - its AST, and the type name of each thing it bound. See Specializable. + Record on a freshly built Specializable what it was made from: the + template itself, its source, its AST, and the real bound objects (not + type names, so a caller can inspect and reuse what was actually bound). + See Specializable. Author: B.G (07/2026) """ + obj._template = template obj._source, obj._ast = capture_template_meta(template) - obj._dependencies = {name: type(value).__name__ for name, value in bindings.items()} + obj._dependencies = dict(bindings) class CompileBuilder(ABC): @@ -452,10 +592,20 @@ class CompileBuilder(ABC): Collects dependencies and a template, and compiles them into one Specializable. - A builder is used once, 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, DeviceFunction, Bag, handle, plain value) is - worked out when the template is specialized, so bind() accepts anything. + 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. @@ -464,18 +614,57 @@ class CompileBuilder(ABC): """ 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 copy of a frozen snapshot; a compiled + object's own `.bindings` is the frozen copy, taken at its compile + time. + + 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. - Raises if `name` is already bound. + + 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._bindings: - raise KeyError(f"'{name}' is already bound") + 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 @@ -488,6 +677,7 @@ def bind_bag(self, bag: "Bag") -> "CompileBuilder": """ for name, item in bag.items(): self.bind(name, item) + self._bag_names.add(name) return self def ingest(self, template) -> "CompileBuilder": @@ -500,35 +690,115 @@ def ingest(self, template) -> "CompileBuilder": 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 the compiled DeviceFunction/Kernel. + 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 DeviceFunctionBuilder(CompileBuilder): +class HelperBuilder(CompileBuilder): """ - Builds a DeviceFunction. compile() -> DeviceFunction. + 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 __init__(self): - super().__init__() - self._uid = new_uid() + 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. - @property - def uid(self) -> int: + Author: B.G (07/2026) """ - Process-wide identity assigned at construction. See Parameter.uid. + 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) """ - return self._uid + ... class KernelBuilder(CompileBuilder): @@ -544,7 +814,7 @@ 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, DeviceFunctions, further Bags, plain python values - mixed + 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 @@ -670,3 +940,182 @@ def check_handles(units: dict[str, dict[str, Any]]) -> None: 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 a Parameter's + mode is changed: mode is fixed at construction (see Parameter.mode), so + changing it means building a new Parameter and using replace() to swap it + into the bag in place of 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/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index e60a8d0..5a23645 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -1,5 +1,7 @@ """ -cupy implementations of Parameter, DeviceFunction, Kernel and their builders. +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 base.py, HelperBuilder). Here a template is CUDA source text rather than a python function, since cp.RawKernel compiles source and there is no function whose globals could be @@ -25,7 +27,7 @@ which keeps macros for common identifiers - N, DIM, EPS, min - from silently rewriting unrelated code in the translation unit. -Spans inside a device function may only reach const parameters and other +Spans inside a device helper may only reach const parameters and other device helpers. That is the framework rule from base.py rather than anything specific to cupy: a helper is spliced into its caller and cannot take on a pointer argument of its own. @@ -39,7 +41,7 @@ import cupy as cp import numpy as np -from .base import DeviceFunction, DeviceFunctionBuilder, Kernel, KernelBuilder, Parameter, attach_meta +from .base import HelperBuilder, Kernel, KernelBuilder, Parameter, _SpecializedHelper, _SpecializeCtx, attach_meta from .base import Bag _KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") @@ -135,15 +137,19 @@ class _SpanParser: Expansion has side effects the builder needs afterwards: `ptr_params` collects the pointer arguments the spans implied, and `helper_srcs` the source of every device helper they called. Set allow_arrays=False when - parsing a device function, which may not reach scalar or field parameters - and so has no use for either. + parsing a device helper, which may not reach scalar or field parameters + and so has no use for either. `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. Author: B.G (07/2026) """ - def __init__(self, bindings: dict[str, Any], *, allow_arrays: bool): + def __init__(self, bindings: dict[str, Any], *, allow_arrays: bool, ctx: _SpecializeCtx): self.bindings = bindings self.allow_arrays = allow_arrays + self.ctx = ctx self.ptr_params: dict[str, dict] = {} # argname -> {ctype, write, array} self.helper_srcs: dict[str, str] = {} # name -> source @@ -159,7 +165,7 @@ def _expand_param(self, param: Parameter, method: str, argname: str, call_args: if param.mode == "const": return _cuda_literal(param.get()) if not self.allow_arrays: - raise ValueError(f"{param.name}: scalar/field param cannot be used in a device function (no pointer arg)") + raise ValueError(f"{param.name}: scalar/field param cannot be used in a device helper (no pointer arg)") self._register_ptr(argname, param, write=False) if param.mode == "scalar": return f"{argname}[0]" @@ -169,7 +175,7 @@ def _expand_param(self, param: Parameter, method: str, argname: str, call_args: if param.mode == "const": raise ValueError(f"{param.name}: const parameter is read-only") if not self.allow_arrays: - raise ValueError(f"{param.name}: set_node cannot be used in a device function (no pointer arg)") + raise ValueError(f"{param.name}: set_node cannot be used in a device helper (no pointer arg)") if len(call_args) != 2: raise ValueError(f"{param.name}: set_node(node, value) takes two arguments") self._register_ptr(argname, param, write=True) @@ -190,7 +196,9 @@ def _repl(self, match: re.Match) -> str: return self._expand_param(target, path[-1], "_".join(path[:-1]), call_args) target = _walk(path, self.bindings) - if isinstance(target, CupyDeviceFunction): + if isinstance(target, CupyHelperBuilder): + target = self.ctx.specialize(target) + if isinstance(target, CupyHelper): self.helper_srcs[target.name] = target.compiled return f"{target.name}({argstr if argstr is not None else ''})" if isinstance(target, Parameter): @@ -342,9 +350,11 @@ def destroy(self) -> None: self._handle = None -class CupyDeviceFunction(DeviceFunction): +class CupyHelper(_SpecializedHelper): """ - A device helper, held as CUDA `__device__` source text. + 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: RawKernel compiles whole translation units, so `.compiled` hands back source, and the helper is compiled as part @@ -364,7 +374,7 @@ def __init__(self, name: str, source: str): def compiled(self): """ The spliced `__device__` source text, for pasting into whatever - kernel or device function binds this helper. + kernel or helper binds this one. Author: B.G (07/2026) """ @@ -378,7 +388,7 @@ def __call__(self, *args, **kwargs): Author: B.G (07/2026) """ raise RuntimeError( - f"DeviceFunction '{self.name}' is CUDA source, only callable from a compiled kernel's device code" + f"Helper '{self.name}' is CUDA source, only callable from a compiled kernel's device code" ) @@ -400,6 +410,7 @@ class CupyKernel(Kernel): _raw_cache: dict[str, "cp.RawKernel"] = {} def __init__(self, name: str, compiled, bound_arrays: list): + super().__init__() self.name = name self._compiled = compiled self._bound_arrays = bound_arrays @@ -425,27 +436,30 @@ def __call__(self, *args, grid, block, **kwargs): return self._compiled(grid, block, tuple(args) + tuple(self._bound_arrays)) -class CupyDeviceFunctionBuilder(DeviceFunctionBuilder): +class CupyHelperBuilder(HelperBuilder): """ - Builds a CupyDeviceFunction from CUDA `__device__` source. Spans may only - reference const params / other device functions (no pointer args). + Recipe for a device helper compiled from CUDA `__device__` source. Spans + may only reference const params / other device helpers (no pointer + args). Specialized only as part of an enclosing kernel's compile() - see + HelperBuilder; compile() itself raises. Author: B.G (07/2026) """ - def compile(self) -> CupyDeviceFunction: + def _specialize(self, ctx: _SpecializeCtx) -> CupyHelper: """ - Expand the template's spans and prepend the helper sources and const - #defines it turned out to need, giving the final `__device__` text. + Expand the template's spans against `ctx` and prepend the helper + sources and const #defines it turned out to need, giving the final + `__device__` text. Author: B.G (07/2026) """ template = self._template name = _extract_name(_DEVICE_NAME_RE, template, "__device__") - parser = _SpanParser(self._bindings, allow_arrays=False) + parser = _SpanParser(self._bindings, allow_arrays=False, ctx=ctx) body = parser.parse(template) source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body]) - fn = CupyDeviceFunction(name, source) + fn = CupyHelper(name, source) attach_meta(fn, template, self._bindings) return fn @@ -465,11 +479,16 @@ def compile(self) -> CupyKernel: the __global__ signature, prepend helpers + const #defines, then build (or reuse, keyed by final source text) the cp.RawKernel. + Opens a fresh _SpecializeCtx for this compile, so every + CupyHelperBuilder this kernel's spans reach is specialized once, + against these bindings. + Author: B.G (07/2026) """ + ctx = _SpecializeCtx() template = self._template name = _extract_name(_KERNEL_NAME_RE, template, "__global__") - parser = _SpanParser(self._bindings, allow_arrays=True) + parser = _SpanParser(self._bindings, allow_arrays=True, ctx=ctx) body = parser.parse(template) body = _inject_signature(body, parser.ptr_params) # extern "C" linkage so cp.RawKernel finds the entry by its plain name diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py index 104dc40..34dce11 100644 --- a/pyfastflow/experimental/core/context/quadrants_backend.py +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -34,7 +34,7 @@ from ._closure_backend import ( ClosureBackendParameter, - ClosureDeviceFunctionBuilder, + ClosureHelperBuilder, ClosureKernelBuilder, ) @@ -49,7 +49,7 @@ class QuadrantsParameter(ClosureBackendParameter): _backend = qd -class QuadrantsDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): +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. diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py index d3e36b6..ff9537a 100644 --- a/pyfastflow/experimental/core/context/taichi_backend.py +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -16,7 +16,7 @@ from ._closure_backend import ( ClosureBackendParameter, - ClosureDeviceFunctionBuilder, + ClosureHelperBuilder, ClosureKernelBuilder, ) @@ -31,7 +31,7 @@ class TaichiParameter(ClosureBackendParameter): _backend = ti -class TaichiDeviceFunctionBuilder(ClosureDeviceFunctionBuilder): +class TaichiHelperBuilder(ClosureHelperBuilder): """ Compiles a device helper with ti.func. diff --git a/pyfastflow/experimental/core/pool/_fields_handle.py b/pyfastflow/experimental/core/pool/_fields_handle.py index 9dbe18f..60384f6 100644 --- a/pyfastflow/experimental/core/pool/_fields_handle.py +++ b/pyfastflow/experimental/core/pool/_fields_handle.py @@ -36,7 +36,7 @@ def __init__(self, dtype: Any, shape: tuple[int, ...]): """ cls = type(self) cls._next_id += 1 - self.id = cls._next_id + self.alloc_id = cls._next_id self._uid = new_uid() self.dtype = dtype self.shape = tuple(shape) diff --git a/pyfastflow/experimental/core/pool/base.py b/pyfastflow/experimental/core/pool/base.py index c9536fc..366760a 100644 --- a/pyfastflow/experimental/core/pool/base.py +++ b/pyfastflow/experimental/core/pool/base.py @@ -39,8 +39,8 @@ class DataHandle(ABC): pool for reuse without freeing memory; `destroy()` actually frees it. Attributes: - id: Per-backend allocation counter, assigned by the backend - used for - pool bookkeeping and not unique across backends. + 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__. @@ -48,10 +48,13 @@ class DataHandle(ABC): 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) """ - id: int + alloc_id: int dtype: Any shape: tuple[int, ...] in_use: bool diff --git a/pyfastflow/experimental/core/pool/cupy_handle.py b/pyfastflow/experimental/core/pool/cupy_handle.py index d8c4d96..e85ca9d 100644 --- a/pyfastflow/experimental/core/pool/cupy_handle.py +++ b/pyfastflow/experimental/core/pool/cupy_handle.py @@ -27,7 +27,7 @@ def __init__(self, dtype: Any, shape: tuple[int, ...]): Author: B.G (07/2026) """ CupyDataHandle._next_id += 1 - self.id = CupyDataHandle._next_id + self.alloc_id = CupyDataHandle._next_id self._uid = new_uid() self.dtype = dtype self.shape = tuple(shape) From df9b78fc4e42c7e0ccf4e6ae1039e6008010ea89 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Fri, 24 Jul 2026 14:10:18 +0200 Subject: [PATCH 18/31] Refactor to builder/composer state #3 --- .../experimental/core/context/__init__.py | 4 +- .../core/context/_closure_backend.py | 26 ++++++++++ .../experimental/core/context/cupy_backend.py | 50 +++++++++++++++++++ .../core/context/quadrants_backend.py | 10 ++++ .../core/context/taichi_backend.py | 10 ++++ 5 files changed, 99 insertions(+), 1 deletion(-) diff --git a/pyfastflow/experimental/core/context/__init__.py b/pyfastflow/experimental/core/context/__init__.py index fcbb697..0eca8ba 100644 --- a/pyfastflow/experimental/core/context/__init__.py +++ b/pyfastflow/experimental/core/context/__init__.py @@ -1,5 +1,7 @@ """ -New backend-agnostic context architecture (Parameter/Specializable ABCs + backends). +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 index 95fa8ae..98bd82e 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -18,6 +18,7 @@ Author: B.G (07/2026) """ +import inspect from types import FunctionType from typing import Any, ClassVar @@ -35,6 +36,7 @@ filter_bindings, resolve_binding, ) +from .routine import RoutineBuilder def specialize_closure(template, bindings: dict[str, Any], ctx: _SpecializeCtx) -> FunctionType: @@ -412,3 +414,27 @@ def compile(self) -> ClosureKernel: krn = ClosureKernel(specialised.__name__, self._backend.kernel(specialised)) attach_meta(krn, self._template, self._bindings) return krn + + +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. + + 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 diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 5a23645..9cedb4c 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -43,6 +43,7 @@ from .base import HelperBuilder, Kernel, KernelBuilder, Parameter, _SpecializedHelper, _SpecializeCtx, attach_meta from .base import Bag +from .routine import RoutineBuilder _KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") _DEVICE_NAME_RE = re.compile(r"__device__\s+[\w:\*&]+\s+(\w+)\s*\(") @@ -507,3 +508,52 @@ def compile(self) -> CupyKernel: krn._final_source = source attach_meta(krn, template, self._bindings) return krn + + +class CupyRoutineBuilder(RoutineBuilder): + """ + Compiles an ordered sequence of cp.RawKernel launches sharing one bag + into a Routine. + + A step's data arity is read off the `__global__` signature as written in + the ingested source, before compile() appends the pointer arguments + spans imply - so it counts exactly the arguments the template author + declared, the same ones data_handle_ref maps onto. + + `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 diff --git a/pyfastflow/experimental/core/context/quadrants_backend.py b/pyfastflow/experimental/core/context/quadrants_backend.py index 34dce11..13f9eac 100644 --- a/pyfastflow/experimental/core/context/quadrants_backend.py +++ b/pyfastflow/experimental/core/context/quadrants_backend.py @@ -36,6 +36,7 @@ ClosureBackendParameter, ClosureHelperBuilder, ClosureKernelBuilder, + ClosureRoutineBuilder, ) @@ -69,3 +70,12 @@ class QuadrantsKernelBuilder(ClosureKernelBuilder): """ _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/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py index ff9537a..341fadd 100644 --- a/pyfastflow/experimental/core/context/taichi_backend.py +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -18,6 +18,7 @@ ClosureBackendParameter, ClosureHelperBuilder, ClosureKernelBuilder, + ClosureRoutineBuilder, ) @@ -49,3 +50,12 @@ class TaichiKernelBuilder(ClosureKernelBuilder): """ _backend = ti + + +class TaichiRoutineBuilder(ClosureRoutineBuilder): + """ + Compiles an ordered sequence of Taichi kernels sharing one bag into a + Routine. + + Author: B.G (07/2026) + """ From fa5b500eae651db2b2a55c9b648e761666c680bb Mon Sep 17 00:00:00 2001 From: bgailleton Date: Fri, 24 Jul 2026 15:22:31 +0200 Subject: [PATCH 19/31] Refactor to builder/composer state #4 --- .../core/context/_closure_backend.py | 289 ++++++++++- .../experimental/core/context/cupy_backend.py | 182 ++++++- .../experimental/core/context/routine.py | 475 ++++++++++++++++++ 3 files changed, 944 insertions(+), 2 deletions(-) create mode 100644 pyfastflow/experimental/core/context/routine.py diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 98bd82e..ab744e2 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -18,12 +18,16 @@ Author: B.G (07/2026) """ +import ast +import copy import inspect +import linecache from types import FunctionType from typing import Any, ClassVar import numpy as np +from ..pool.base import new_uid from .base import ( Bag, HelperBuilder, @@ -33,10 +37,11 @@ _SpecializedHelper, _SpecializeCtx, attach_meta, + capture_template_meta, filter_bindings, resolve_binding, ) -from .routine import RoutineBuilder +from .routine import Routine, RoutineBuilder, _CompiledStep, _template_label def specialize_closure(template, bindings: dict[str, Any], ctx: _SpecializeCtx) -> FunctionType: @@ -416,6 +421,108 @@ def compile(self) -> ClosureKernel: return krn +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. @@ -427,6 +534,10 @@ class ClosureRoutineBuilder(RoutineBuilder): `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) """ @@ -438,3 +549,179 @@ def _data_arity(self, kernel_builder: KernelBuilder) -> int: 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] = {} + raw_bindings: 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) + raw_bindings.update(filtered) + 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) + + func_name = f"_fused_group{group_index}_{new_uid()}" + 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=func_name, args=args_node, body=body_stmts, decorator_list=[], returns=None) + module = ast.fix_missing_locations(ast.Module(body=[fused_def], type_ignores=[])) + 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)) + attach_meta(krn, fused_fn, raw_bindings) + return krn, data_names diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 9cedb4c..1e1e2c1 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -43,7 +43,7 @@ from .base import HelperBuilder, Kernel, KernelBuilder, Parameter, _SpecializedHelper, _SpecializeCtx, attach_meta from .base import Bag -from .routine import RoutineBuilder +from .routine import Routine, RoutineBuilder, _CompiledStep _KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") _DEVICE_NAME_RE = re.compile(r"__device__\s+[\w:\*&]+\s+(\w+)\s*\(") @@ -510,6 +510,80 @@ def compile(self) -> CupyKernel: 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 argument list resolved. Call-time + data handle overrides (the `rout(A, B)` form) cannot be honoured against + that: 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 the graph's launches 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 the + pointers the graph's launches were captured 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 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 cp.RawKernel launches sharing one bag @@ -557,3 +631,109 @@ 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. cp.RawKernel compiles its module lazily, on + first launch; that JIT must not happen while capturing (a + RawKernel launched for the first time mid-capture either fails + the capture outright or bakes in a broken graph node - CUDA's + capture machinery assumes the module is already resident). + 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. + 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 first-launch + # module JIT happens before capture starts, not during it. + 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/routine.py b/pyfastflow/experimental/core/context/routine.py new file mode 100644 index 0000000..6307e9c --- /dev/null +++ b/pyfastflow/experimental/core/context/routine.py @@ -0,0 +1,475 @@ +""" +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 base.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 (base.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 +base.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 .base import Bag, CompileBuilder, check_handles + +_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 base.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 (base.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) From 65cbd098f87db5888fe8ed3668f243a24ae0ce43 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Fri, 24 Jul 2026 17:14:46 +0200 Subject: [PATCH 20/31] Refactor to builder/composer state #5 --- .../core/context/_closure_backend.py | 29 +- pyfastflow/experimental/core/context/base.py | 23 +- .../experimental/core/context/cupy_backend.py | 386 +++++++++++++----- 3 files changed, 292 insertions(+), 146 deletions(-) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index ab744e2..8c0c88c 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -346,29 +346,6 @@ def __call__(self, *args, **kwargs): return self._compiled(*args, **kwargs) -def _check_const_only(bindings: dict[str, Any]) -> None: - """ - Raise unless every Parameter reachable in `bindings` is const mode. - - This enforces the rule in base.py that a device helper binds const - parameters only, passing any data through explicit arguments instead. - Bound Bags are searched too, so a scalar or field parameter tucked inside - one does not slip past. - - Author: B.G (07/2026) - """ - for name, value in bindings.items(): - if isinstance(value, Parameter): - if value.mode != "const": - raise ValueError( - f"device helper: bound parameter '{name}' has mode {value.mode!r}, but a device " - "helper may only bind const parameters - pass the data as an explicit argument to " - "the helper instead" - ) - elif isinstance(value, Bag): - _check_const_only(dict(value.items())) - - class ClosureHelperBuilder(HelperBuilder): """ Recipe for a device helper compiled to a ti.func/qd.func. Subclasses pin @@ -382,13 +359,11 @@ class ClosureHelperBuilder(HelperBuilder): def _specialize(self, ctx: _SpecializeCtx) -> ClosureHelper: """ - Check the const-only rule, splice the referenced bindings into the - template's globals against `ctx`, and compile the result as a device - func. + Splice the referenced bindings into the template's globals against + `ctx`, and compile the result as a device func. Author: B.G (07/2026) """ - _check_const_only(self._bindings) specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) fn = ClosureHelper(specialised.__name__, self._backend.func(specialised)) attach_meta(fn, self._template, self._bindings) diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index dce759a..6b1a5f6 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -91,14 +91,21 @@ is the exception: it resolves to a bare compile-time literal, read as `p` with no call. -Device helpers bind const parameters only ------------------------------------------ -A HelperBuilder may only bind const-mode Parameters; any data it needs is -passed to it as an explicit argument by the calling kernel. A helper is -spliced into its caller and has no way to acquire a pointer argument of its -own, so this holds on all three backends alike. It is checked when the helper -is specialized, i.e. as part of the enclosing kernel's compile(). Kernels -carry no such restriction and may bind any mode. +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 ----------------------------- diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 1e1e2c1..10bfc07 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -4,7 +4,7 @@ whichever kernel binds it (see base.py, HelperBuilder). Here a template is CUDA source text rather than a python function, since -cp.RawKernel compiles source and there is no function whose globals could be +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: @@ -15,22 +15,44 @@ $helper(a, b)$ call a bound device helper Compiling substitutes each span according to the parameter's mode - a CUDA -literal for const, NAME[0] for scalar, NAME[i] for field - and, for scalar and -field, also generates the pointer argument that expansion implies. Those -arguments are appended to the __global__ signature and their arrays supplied at -launch, so a template never declares them by hand. A parameter only read is -declared `const T*`; one written anywhere in the kernel is declared `T*`. -Writing to a const parameter is an error. +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; + +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. -Spans inside a device helper may only reach const parameters and other -device helpers. That is the framework rule from base.py rather than anything -specific to cupy: a helper is spliced into its caller and cannot take on a -pointer argument of its own. +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, uid-qualified struct members included, so it remains a sound key. +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) """ @@ -87,7 +109,7 @@ def _cuda_literal(value) -> str: 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.RawKernel is looked up by. + text - that is the entry point cp.RawModule.get_function is looked up by. Author: B.G (07/2026) """ @@ -131,43 +153,69 @@ def _walk(path: list[str], bindings: dict[str, Any]): return obj +def _param_argname(param: Parameter) -> str: + """ + The struct member / local variable name a Parameter's pointer is reached + through - stable for the object's whole lifetime since it is derived from + `uid`, not from whatever name(s) it happens to be bound under. This is + what makes dedup by uid automatic: two spans reaching the same Parameter + under two different handles compute the same argname and therefore the + same struct member. + + Author: B.G (07/2026) + """ + return f"p_{param.uid}" + + class _SpanParser: """ Expands the `$...$` spans in a CUDA template body. - Expansion has side effects the builder needs afterwards: `ptr_params` - collects the pointer arguments the spans implied, and `helper_srcs` the - source of every device helper they called. Set allow_arrays=False when - parsing a device helper, which may not reach scalar or field parameters - and so has no use for either. `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. + 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. Author: B.G (07/2026) """ - def __init__(self, bindings: dict[str, Any], *, allow_arrays: bool, ctx: _SpecializeCtx): + def __init__(self, bindings: dict[str, Any], *, ctx: _SpecializeCtx): self.bindings = bindings - self.allow_arrays = allow_arrays self.ctx = ctx - self.ptr_params: dict[str, dict] = {} # argname -> {ctype, write, array} + self.local_ptrs: dict[int, dict] = {} # uid -> {ctype, write} self.helper_srcs: dict[str, str] = {} # name -> source - def _register_ptr(self, argname: str, param: Parameter, write: bool) -> None: - entry = self.ptr_params.get(argname) + 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 = {} + uid = param.uid + entry = registry.get(uid) if entry is None: - self.ptr_params[argname] = {"ctype": _ctype(param.dtype), "write": write, "array": param.get().data} - elif write: + entry = {"ctype": _ctype(param.dtype), "write": False, "array": param.get().data} + registry[uid] = entry + if write: entry["write"] = True + 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) - def _expand_param(self, param: Parameter, method: str, argname: str, call_args: list[str]) -> str: + 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()) - if not self.allow_arrays: - raise ValueError(f"{param.name}: scalar/field param cannot be used in a device helper (no pointer arg)") - self._register_ptr(argname, param, write=False) + argname = self._register_ptr(param, write=False) if param.mode == "scalar": return f"{argname}[0]" idx = call_args[0] if call_args else "0" @@ -175,11 +223,9 @@ def _expand_param(self, param: Parameter, method: str, argname: str, call_args: # set_node if param.mode == "const": raise ValueError(f"{param.name}: const parameter is read-only") - if not self.allow_arrays: - raise ValueError(f"{param.name}: set_node cannot be used in a device helper (no pointer arg)") if len(call_args) != 2: raise ValueError(f"{param.name}: set_node(node, value) takes two arguments") - self._register_ptr(argname, param, write=True) + 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}" @@ -194,7 +240,7 @@ def _repl(self, match: re.Match) -> str: if path[-1] in ("get", "set_node"): target = _walk(path[:-1], self.bindings) if isinstance(target, Parameter): - return self._expand_param(target, path[-1], "_".join(path[:-1]), call_args) + return self._expand_param(target, path[-1], call_args) target = _walk(path, self.bindings) if isinstance(target, CupyHelperBuilder): @@ -203,17 +249,55 @@ def _repl(self, match: re.Match) -> str: self.helper_srcs[target.name] = target.compiled return f"{target.name}({argstr if argstr is not None else ''})" if isinstance(target, Parameter): - return self._expand_param(target, "get", "_".join(path), ["0"]) + return self._expand_param(target, "get", ["0"]) return _cuda_literal(target) def parse(self, body: str) -> str: """ - Expand every `$...$` span in `body`, accumulating ptr_params and - helper_srcs as a side effect. + Expand every `$...$` span in `body`, accumulating local_ptrs and + helper_srcs as a side effect, then prepend the `__restrict__` locals + this body's own spans implied - see _insert_locals. Author: B.G (07/2026) """ - return _SPAN_RE.sub(self._repl, body) + expanded = _SPAN_RE.sub(self._repl, body) + return _insert_locals(expanded, self.local_ptrs) + + +def _insert_locals(body: str, local_ptrs: dict[int, dict]) -> 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. + + 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(uid)} = pf_params.{_argname_for(uid)};\n" + for uid, e in sorted(local_ptrs.items()) + ) + return f"{body[: idx + 1]}\n{decls}{body[idx + 1 :]}" + + +def _argname_for(uid: int) -> str: + """ + The struct member / local name for a pointer already registered under + `uid` - see _param_argname, which this must stay in lockstep with. + + Author: B.G (07/2026) + """ + return f"p_{uid}" def _const_defines(bindings: dict[str, Any], body: str) -> list[str]: @@ -236,26 +320,46 @@ def _const_defines(bindings: dict[str, Any], body: str) -> list[str]: ] -def _inject_signature(kernel_src: str, ptr_params: dict[str, dict]) -> str: +def _param_block_source(registry: dict[int, dict]) -> str: """ - Append the generated pointer params to the __global__ signature. + 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 uid, ascending - arbitrary but fixed, so it agrees with + the upload order _upload_param_block writes in. Author: B.G (07/2026) """ - if not ptr_params: - return kernel_src - decls = [ - f"{'' if e['write'] else 'const '}{e['ctype']}* {name}" - for name, e in ptr_params.items() - ] - joined = ", ".join(decls) + if not registry: + return "" + members = "".join( + f" {'' if e['write'] else 'const '}{e['ctype']}* {_argname_for(uid)};\n" + for uid, e in sorted(registry.items()) + ) + return f"struct pf_params_t {{\n{members}}};\n__constant__ pf_params_t pf_params;\n" - def _sub(m: re.Match) -> str: - existing = m.group(2).strip() - sep = ", " if existing else "" - return f"{m.group(1)}{m.group(2)}{sep}{joined}{m.group(3)}" - return _KERNEL_SIG_RE.sub(_sub, kernel_src, count=1) +def _upload_param_block(module: "cp.RawModule", registry: dict[int, dict]) -> None: + """ + Copy the current pointer for every registered Parameter into the module's + `pf_params` constant block, in the same uid order the struct was emitted + in. + + 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())], dtype=np.uint64) + view = cp.ndarray(ptrs.shape, dtype=np.uint64, memptr=global_ptr) + view.set(ptrs) class CupyParameter(Parameter): @@ -357,7 +461,7 @@ class CupyHelper(_SpecializedHelper): Produced by a CupyHelperBuilder as part of an enclosing kernel's compile(); see HelperBuilder. - Nothing is compiled at this stage: RawKernel compiles whole translation + 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. @@ -395,53 +499,65 @@ def __call__(self, *args, **kwargs): class CupyKernel(Kernel): """ - A launchable kernel, backed by a compiled cp.RawKernel. - - Launch dimensions are explicit: RawKernel has nothing like Taichi's - auto-ranging, so __call__ takes grid and block. The positional arguments - are the template's own data arguments, and the arrays behind the generated - pointer parameters follow them. + A launchable kernel, backed by a function pulled from a compiled + cp.RawModule (see CupyKernelBuilder.compile). - Compiled kernels are cached on the class by final source text, which is - what the whole specialization reduces to and therefore a sound key. + 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.RawKernel"] = {} + _raw_cache: dict[str, "cp.RawModule"] = {} - def __init__(self, name: str, compiled, bound_arrays: list): + def __init__(self, name: str, compiled, module: "cp.RawModule"): super().__init__() self.name = name self._compiled = compiled - self._bound_arrays = bound_arrays + self._module = module @property def compiled(self): """ - The underlying cp.RawKernel this Kernel's __call__ launches. + 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 RawKernel has no default range. + 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) + tuple(self._bound_arrays)) + return self._compiled(grid, block, tuple(args)) class CupyHelperBuilder(HelperBuilder): """ - Recipe for a device helper compiled from CUDA `__device__` source. Spans - may only reference const params / other device helpers (no pointer - args). Specialized only as part of an enclosing kernel's compile() - see + 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) @@ -451,13 +567,17 @@ def _specialize(self, ctx: _SpecializeCtx) -> CupyHelper: """ Expand the template's spans against `ctx` and prepend the helper sources and const #defines it turned out to need, giving the final - `__device__` text. + `__device__` text. 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. Author: B.G (07/2026) """ template = self._template name = _extract_name(_DEVICE_NAME_RE, template, "__device__") - parser = _SpanParser(self._bindings, allow_arrays=False, ctx=ctx) + parser = _SpanParser(self._bindings, ctx=ctx) body = parser.parse(template) source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body]) fn = CupyHelper(name, source) @@ -476,35 +596,53 @@ class CupyKernelBuilder(KernelBuilder): def compile(self) -> CupyKernel: """ - Expand the template's spans, inject the pointer args they implied into - the __global__ signature, prepend helpers + const #defines, then build - (or reuse, keyed by final source text) the cp.RawKernel. + Expand the template's spans - the kernel's own and, recursively, + every CupyHelperBuilder they reach - prepend the helpers' source, + 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. + 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 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 = {} template = self._template name = _extract_name(_KERNEL_NAME_RE, template, "__global__") - parser = _SpanParser(self._bindings, allow_arrays=True, ctx=ctx) + parser = _SpanParser(self._bindings, ctx=ctx) body = parser.parse(template) - body = _inject_signature(body, parser.ptr_params) - # extern "C" linkage so cp.RawKernel finds the entry by its plain name + # 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) - source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body]) - bound_arrays = [e["array"] for e in parser.ptr_params.values()] + registry = ctx.cupy_ptr_registry + source = "\n".join( + [_param_block_source(registry)] + + list(parser.helper_srcs.values()) + + _const_defines(self._bindings, body) + + [body] + ) - raw = CupyKernel._raw_cache.get(source) - if raw is None: - raw = cp.RawKernel(source, name) - CupyKernel._raw_cache[source] = raw - krn = CupyKernel(name, raw, bound_arrays) + 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) + raw = module.get_function(name) + krn = CupyKernel(name, raw, module) krn._final_source = source attach_meta(krn, template, self._bindings) return krn @@ -530,30 +668,43 @@ class _CapturedRoutine(Routine): 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 argument list resolved. Call-time + 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 - that: 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. + 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 the graph's launches 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; + 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 the - pointers the graph's launches were captured with. Recompile. + 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 as an added way "recompile after this" applies. + graph's baked-in pointers - launch arguments and constant blocks alike - + as an added way "recompile after this" applies. Author: B.G (07/2026) """ @@ -586,13 +737,16 @@ def __call__(self, *args) -> None: class CupyRoutineBuilder(RoutineBuilder): """ - Compiles an ordered sequence of cp.RawKernel launches sharing one bag + 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, before compile() appends the pointer arguments - spans imply - so it counts exactly the arguments the template author - declared, the same ones data_handle_ref maps onto. + 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 @@ -652,11 +806,15 @@ def compile(self, captured: bool = True, dump_source: str | None = None) -> Rout 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. cp.RawKernel compiles its module lazily, on - first launch; that JIT must not happen while capturing (a - RawKernel launched for the first time mid-capture either fails - the capture outright or bakes in a broken graph node - CUDA's - capture machinery assumes the module is already resident). + 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 @@ -667,7 +825,11 @@ def compile(self, captured: bool = True, dump_source: str | None = None) -> Rout 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. + 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() @@ -705,8 +867,10 @@ 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 first-launch - # module JIT happens before capture starts, not during it. + # 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() From 5b085e70eaa16a55f4da41a5ca463c9d2dad276b Mon Sep 17 00:00:00 2001 From: bgailleton Date: Fri, 24 Jul 2026 17:28:44 +0200 Subject: [PATCH 21/31] adding core example to demonstrate kernel param and helepr builders --- .../heat_diffusion/heat_diffusion_cupy.py | 341 +++++++++++++++++ .../heat_diffusion_pure_taichi.py | 194 ++++++++++ .../heat_diffusion_quadrants.py | 336 +++++++++++++++++ .../heat_diffusion_routine_cupy.py | 351 ++++++++++++++++++ .../heat_diffusion_routine_quadrants.py | 312 ++++++++++++++++ .../heat_diffusion_routine_taichi.py | 312 ++++++++++++++++ .../heat_diffusion/heat_diffusion_taichi.py | 336 +++++++++++++++++ .../core/shallow_water/shallow_water_cupy.py | 308 +++++++++++++++ .../shallow_water/shallow_water_quadrants.py | 279 ++++++++++++++ .../shallow_water/shallow_water_taichi.py | 278 ++++++++++++++ 10 files changed, 3047 insertions(+) create mode 100644 examples/core/heat_diffusion/heat_diffusion_cupy.py create mode 100644 examples/core/heat_diffusion/heat_diffusion_pure_taichi.py create mode 100644 examples/core/heat_diffusion/heat_diffusion_quadrants.py create mode 100644 examples/core/heat_diffusion/heat_diffusion_routine_cupy.py create mode 100644 examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py create mode 100644 examples/core/heat_diffusion/heat_diffusion_routine_taichi.py create mode 100644 examples/core/heat_diffusion/heat_diffusion_taichi.py create mode 100644 examples/core/shallow_water/shallow_water_cupy.py create mode 100644 examples/core/shallow_water/shallow_water_quadrants.py create mode 100644 examples/core/shallow_water/shallow_water_taichi.py 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..b5e0cd9 --- /dev/null +++ b/examples/core/heat_diffusion/heat_diffusion_cupy.py @@ -0,0 +1,341 @@ +""" +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.base 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, solo=True) +room_p = CupyParameter("ROOM", dtype=np.int32, mode="const", value=GRID_N // 4, pool=pool, solo=True) +wall_thick_p = CupyParameter("WALL_THICK", dtype=np.int32, mode="const", value=8, pool=pool, solo=True) +door_p = CupyParameter("DOOR", dtype=np.int32, mode="const", value=6, pool=pool, solo=True) +seed_p = CupyParameter("SEED", dtype=np.float32, mode="const", value=17.0, pool=pool, solo=True) + +dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool, solo=True) +dx2_p = CupyParameter("DX2", dtype=np.float32, mode="const", value=DX_M**2, pool=pool, solo=True) + +alpha_air_seed_p = CupyParameter("ALPHA_AIR_SEED", dtype=np.float32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) +alpha_wall_seed_p = CupyParameter("ALPHA_WALL_SEED", dtype=np.float32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) +t_bg_p = CupyParameter("T_BG", dtype=np.float32, mode="const", value=15.0, pool=pool, solo=True) + +src_i_p = CupyParameter("SRC_I", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_j_p = CupyParameter("SRC_J", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_r_p = CupyParameter("SRC_R", dtype=np.int32, mode="const", value=10, pool=pool, solo=True) + +# 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 - the solo consts 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$; + int dy = idx % N - $stove.at.j$; + if (dx * dx + dy * dy <= $stove.r$ * $stove.r$) { + 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 solo consts - 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$; + T_out[idx] = T_in[idx] + $heat.dt$ * 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 base.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..824a5cc --- /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: `wall`, `alpha` and `stove` are 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. Structural constants +(N, ROOM, ...) are declared solo=True and used bare as compile-time literals. + +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.base 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: solo=True -> read bare in kernel bodies as compile-time +# literals (no .get()), since they are never going to vary at runtime. +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) +room_p = QuadrantsParameter("ROOM", dtype=qd.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) +wall_thick_p = QuadrantsParameter("WALL_THICK", dtype=qd.i32, mode="const", value=8, pool=pool, solo=True) +door_p = QuadrantsParameter("DOOR", dtype=qd.i32, mode="const", value=6, pool=pool, solo=True) +seed_p = QuadrantsParameter("SEED", dtype=qd.f32, mode="const", value=17.0, pool=pool, solo=True) + +dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool, solo=True) # seconds +dx2_p = QuadrantsParameter("DX2", dtype=qd.f32, mode="const", value=DX_M**2, pool=pool, solo=True) # meters^2 + +# Seed values for the alpha field - used bare inside set_alpha. +alpha_air_seed_p = QuadrantsParameter("ALPHA_AIR_SEED", dtype=qd.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) +alpha_wall_seed_p = QuadrantsParameter("ALPHA_WALL_SEED", dtype=qd.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) +t_bg_p = QuadrantsParameter("T_BG", dtype=qd.f32, mode="const", value=15.0, pool=pool, solo=True) + +src_i_p = QuadrantsParameter("SRC_I", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_j_p = QuadrantsParameter("SRC_J", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_r_p = QuadrantsParameter("SRC_R", dtype=qd.i32, mode="const", value=10, pool=pool, solo=True) # 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 - 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 + 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, 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 = qd.cast(whash(vline, seg) * ROOM, qd.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 = qd.cast(whash(hline + 7919, seg) * ROOM, qd.i32) + gap = (i % ROOM) >= door and (i % ROOM) < door + DOOR + if not gap: + is_wall = 1 + wall.set_node(i * N + 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, N): + idx = i * N + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED) + else: + alpha.set_node(idx, ALPHA_AIR_SEED) + + +# 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 + + +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, and members of either level resolve on their +# own type - the solo consts become compile-time literals (read bare), the +# scalar Parameter keeps its .get(0). 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 + dy = j - stove.at.j + if dx * dx + dy * dy <= stove.r * stove.r: + 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 solo consts - 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, `heat.dx2` 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 + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2 + T_out[i, j] = T_in[i, j] + heat.dt * 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 base.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..ae8d0a4 --- /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.base 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, solo=True) +room_p = CupyParameter("ROOM", dtype=np.int32, mode="const", value=GRID_N // 4, pool=pool, solo=True) +wall_thick_p = CupyParameter("WALL_THICK", dtype=np.int32, mode="const", value=8, pool=pool, solo=True) +door_p = CupyParameter("DOOR", dtype=np.int32, mode="const", value=6, pool=pool, solo=True) +seed_p = CupyParameter("SEED", dtype=np.float32, mode="const", value=17.0, pool=pool, solo=True) + +dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool, solo=True) +dx2_p = CupyParameter("DX2", dtype=np.float32, mode="const", value=DX_M**2, pool=pool, solo=True) + +alpha_air_seed_p = CupyParameter("ALPHA_AIR_SEED", dtype=np.float32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) +alpha_wall_seed_p = CupyParameter("ALPHA_WALL_SEED", dtype=np.float32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) +t_bg_p = CupyParameter("T_BG", dtype=np.float32, mode="const", value=15.0, pool=pool, solo=True) + +src_i_p = CupyParameter("SRC_I", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_j_p = CupyParameter("SRC_J", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_r_p = CupyParameter("SRC_R", dtype=np.int32, mode="const", value=10, pool=pool, solo=True) + +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 base.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$; + int dy = idx % N - $stove.at.j$; + if (dx * dx + dy * dy <= $stove.r$ * $stove.r$) { + 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$; + T_out[idx] = T_in[idx] + $heat.dt$ * 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..e87cb0d --- /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.base 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, solo=True) +room_p = QuadrantsParameter("ROOM", dtype=qd.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) +wall_thick_p = QuadrantsParameter("WALL_THICK", dtype=qd.i32, mode="const", value=8, pool=pool, solo=True) +door_p = QuadrantsParameter("DOOR", dtype=qd.i32, mode="const", value=6, pool=pool, solo=True) +seed_p = QuadrantsParameter("SEED", dtype=qd.f32, mode="const", value=17.0, pool=pool, solo=True) + +dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool, solo=True) +dx2_p = QuadrantsParameter("DX2", dtype=qd.f32, mode="const", value=DX_M**2, pool=pool, solo=True) + +alpha_air_seed_p = QuadrantsParameter("ALPHA_AIR_SEED", dtype=qd.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) +alpha_wall_seed_p = QuadrantsParameter("ALPHA_WALL_SEED", dtype=qd.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) +t_bg_p = QuadrantsParameter("T_BG", dtype=qd.f32, mode="const", value=15.0, pool=pool, solo=True) + +src_i_p = QuadrantsParameter("SRC_I", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_j_p = QuadrantsParameter("SRC_J", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_r_p = QuadrantsParameter("SRC_R", dtype=qd.i32, mode="const", value=10, pool=pool, solo=True) + +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 - 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 + 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, 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 = qd.cast(whash(vline, seg) * ROOM, qd.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 = qd.cast(whash(hline + 7919, seg) * ROOM, qd.i32) + gap = (i % ROOM) >= door and (i % ROOM) < door + DOOR + if not gap: + is_wall = 1 + wall.set_node(i * N + 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, N): + idx = i * N + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED) + else: + alpha.set_node(idx, ALPHA_AIR_SEED) + + +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 + + +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 + dy = j - stove.at.j + if dx * dx + dy * dy <= stove.r * stove.r: + 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 base.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 + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2 + T_out[i, j] = T_in[i, j] + heat.dt * 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..7a6ebf4 --- /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.base 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, solo=True) +room_p = TaichiParameter("ROOM", dtype=ti.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) +wall_thick_p = TaichiParameter("WALL_THICK", dtype=ti.i32, mode="const", value=8, pool=pool, solo=True) +door_p = TaichiParameter("DOOR", dtype=ti.i32, mode="const", value=6, pool=pool, solo=True) +seed_p = TaichiParameter("SEED", dtype=ti.f32, mode="const", value=17.0, pool=pool, solo=True) + +dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool, solo=True) +dx2_p = TaichiParameter("DX2", dtype=ti.f32, mode="const", value=DX_M**2, pool=pool, solo=True) + +alpha_air_seed_p = TaichiParameter("ALPHA_AIR_SEED", dtype=ti.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) +alpha_wall_seed_p = TaichiParameter("ALPHA_WALL_SEED", dtype=ti.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) +t_bg_p = TaichiParameter("T_BG", dtype=ti.f32, mode="const", value=15.0, pool=pool, solo=True) + +src_i_p = TaichiParameter("SRC_I", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_j_p = TaichiParameter("SRC_J", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_r_p = TaichiParameter("SRC_R", dtype=ti.i32, mode="const", value=10, pool=pool, solo=True) + +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 - 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 + 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, 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.set_node(i * N + 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, N): + idx = i * N + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED) + else: + alpha.set_node(idx, ALPHA_AIR_SEED) + + +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 + + +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 + dy = j - stove.at.j + if dx * dx + dy * dy <= stove.r * stove.r: + 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 base.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 + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2 + T_out[i, j] = T_in[i, j] + heat.dt * 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..54e1a63 --- /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: `wall`, `alpha` and `stove` are 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. Structural constants +(N, ROOM, ...) are declared solo=True and used bare as compile-time literals. + +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.base 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: solo=True -> read bare in kernel bodies as compile-time +# literals (no .get()), since they are never going to vary at runtime. +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) +room_p = TaichiParameter("ROOM", dtype=ti.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) +wall_thick_p = TaichiParameter("WALL_THICK", dtype=ti.i32, mode="const", value=8, pool=pool, solo=True) +door_p = TaichiParameter("DOOR", dtype=ti.i32, mode="const", value=6, pool=pool, solo=True) +seed_p = TaichiParameter("SEED", dtype=ti.f32, mode="const", value=17.0, pool=pool, solo=True) + +dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool, solo=True) # seconds +dx2_p = TaichiParameter("DX2", dtype=ti.f32, mode="const", value=DX_M**2, pool=pool, solo=True) # meters^2 + +# Seed values for the alpha field - used bare inside set_alpha. +alpha_air_seed_p = TaichiParameter("ALPHA_AIR_SEED", dtype=ti.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) +alpha_wall_seed_p = TaichiParameter("ALPHA_WALL_SEED", dtype=ti.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) +t_bg_p = TaichiParameter("T_BG", dtype=ti.f32, mode="const", value=15.0, pool=pool, solo=True) + +src_i_p = TaichiParameter("SRC_I", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_j_p = TaichiParameter("SRC_J", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) +src_r_p = TaichiParameter("SRC_R", dtype=ti.i32, mode="const", value=10, pool=pool, solo=True) # 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 - 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 + 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, 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.set_node(i * N + 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, N): + idx = i * N + j + if wall.get(idx) == 1: + alpha.set_node(idx, ALPHA_WALL_SEED) + else: + alpha.set_node(idx, ALPHA_AIR_SEED) + + +# 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 + + +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, and members of either level resolve on their +# own type - the solo consts become compile-time literals (read bare), the +# scalar Parameter keeps its .get(0). 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 + dy = j - stove.at.j + if dx * dx + dy * dy <= stove.r * stove.r: + 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 solo consts - 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, `heat.dx2` 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 + j + a = heat.alpha.get(idx) + lap = heat.lap(T_in, i, j) / heat.dx2 + T_out[i, j] = T_in[i, j] + heat.dt * 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 base.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/shallow_water/shallow_water_cupy.py b/examples/core/shallow_water/shallow_water_cupy.py new file mode 100644 index 0000000..71ff871 --- /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.base 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, solo=True) +rest_depth_p = CupyParameter("REST_DEPTH", dtype=np.float32, mode="const", value=REST_DEPTH_VAL, pool=pool, solo=True) +drop_r_p = CupyParameter("DROP_R", dtype=np.int32, mode="const", value=DROP_R_VAL, pool=pool, solo=True) + +# 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 base.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..49a7214 --- /dev/null +++ b/examples/core/shallow_water/shallow_water_quadrants.py @@ -0,0 +1,279 @@ +""" +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 solo=True: compile-time +literals used bare, no .get(). + +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.base 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: solo=True -> read bare as compile-time literals. +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) +rest_depth_p = QuadrantsParameter("REST_DEPTH", dtype=qd.f32, mode="const", value=REST_DEPTH_VAL, pool=pool, solo=True) +drop_r_p = QuadrantsParameter("DROP_R", dtype=qd.i32, mode="const", value=DROP_R_VAL, pool=pool, solo=True) + +# 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 - 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 + + +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 * DROP_R: + 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 - 1: + ue = u[i + 1, j] + vs = v[i, j] + vn = 0.0 + if j < N - 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 base.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..16a1197 --- /dev/null +++ b/examples/core/shallow_water/shallow_water_taichi.py @@ -0,0 +1,278 @@ +""" +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 solo=True: compile-time +literals used bare, no .get(). + +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.base 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: solo=True -> read bare as compile-time literals. +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) +rest_depth_p = TaichiParameter("REST_DEPTH", dtype=ti.f32, mode="const", value=REST_DEPTH_VAL, pool=pool, solo=True) +drop_r_p = TaichiParameter("DROP_R", dtype=ti.i32, mode="const", value=DROP_R_VAL, pool=pool, solo=True) + +# 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 - 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 + + +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 * DROP_R: + 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 - 1: + ue = u[i + 1, j] + vs = v[i, j] + vn = 0.0 + if j < N - 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 base.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") From adb318de2e66c4332617065a95ebba35bdc9aa99 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 27 Jul 2026 10:08:48 +0200 Subject: [PATCH 22/31] Add example with Rouine --- examples/core/lem/lem_routine_cupy.py | 299 +++++++++++++++++++++ examples/core/lem/lem_routine_quadrants.py | 258 ++++++++++++++++++ examples/core/lem/lem_routine_taichi.py | 258 ++++++++++++++++++ 3 files changed, 815 insertions(+) create mode 100644 examples/core/lem/lem_routine_cupy.py create mode 100644 examples/core/lem/lem_routine_quadrants.py create mode 100644 examples/core/lem/lem_routine_taichi.py diff --git a/examples/core/lem/lem_routine_cupy.py b/examples/core/lem/lem_routine_cupy.py new file mode 100644 index 0000000..f59c491 --- /dev/null +++ b/examples/core/lem/lem_routine_cupy.py @@ -0,0 +1,299 @@ +""" +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 solo consts, arriving as #defines and read + bare. DX is a non-solo const, read as $grid.dx.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.base 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 +# --------------------------------------------------------------------------- +# solo consts: emitted as #defines and read bare in the source +n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool, solo=True) +dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool, solo=True) +seed_a_p = CupyParameter("SEED_A", dtype=np.float32, mode="const", value=12.9898, pool=pool, solo=True) +seed_b_p = CupyParameter("SEED_B", dtype=np.float32, mode="const", value=78.233, pool=pool, solo=True) + +# non-solo const: still fixed at compile time, but read through a span like +# any other mode, so a template can be written without knowing it is const +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: grid.n is a solo const read bare, grid.dx a non-solo const reached +# through a span - 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 non-solo const 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 solo const under one name. +# hill.d is a span read, hill.lap a spliced call, hill.dt a bare literal. +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$ * $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 solo const reached through the bag, so it splices in as a + // bare literal here just as it would read bare if bound flat + int n = $grid.n$; + 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..77993bc --- /dev/null +++ b/examples/core/lem/lem_routine_quadrants.py @@ -0,0 +1,258 @@ +""" +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 and DT are solo consts, read bare as compile-time + literals. DX is a non-solo const, read as grid.dx.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. + 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.base 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 +# --------------------------------------------------------------------------- +# solo consts: folded into the generated code as bare literals, read as `N` +n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) +dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool, solo=True) +seed_a_p = QuadrantsParameter("SEED_A", dtype=qd.f32, mode="const", value=12.9898, pool=pool, solo=True) +seed_b_p = QuadrantsParameter("SEED_B", dtype=qd.f32, mode="const", value=78.233, pool=pool, solo=True) + +# non-solo const: still fixed at compile time, but read through .get(0) like +# any other mode, so a template can be written without knowing it is const +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: grid.n is a solo const read bare, grid.dx a non-solo const 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 - 1) + + +clampi_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clampi) + + +def laplacian(field_, i, j): + # calls another helper, and reads a non-solo const 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 + 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 + qd.cast(j, qd.f32) * SEED_B + 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 solo const under one name. +# hill.d is a device accessor, hill.lap a specialized func, hill.dt a literal. +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 * 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 + # base level: pin the east and west edges, so the ridge drains outward + if j == 0 or j == grid.n - 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..eac8178 --- /dev/null +++ b/examples/core/lem/lem_routine_taichi.py @@ -0,0 +1,258 @@ +""" +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 and DT are solo consts, read bare as compile-time + literals. DX is a non-solo const, read as grid.dx.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. + 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.base 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 +# --------------------------------------------------------------------------- +# solo consts: folded into the generated code as bare literals, read as `N` +n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) +dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool, solo=True) +seed_a_p = TaichiParameter("SEED_A", dtype=ti.f32, mode="const", value=12.9898, pool=pool, solo=True) +seed_b_p = TaichiParameter("SEED_B", dtype=ti.f32, mode="const", value=78.233, pool=pool, solo=True) + +# non-solo const: still fixed at compile time, but read through .get(0) like +# any other mode, so a template can be written without knowing it is const +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: grid.n is a solo const read bare, grid.dx a non-solo const 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 - 1) + + +clampi_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clampi) + + +def laplacian(field_, i, j): + # calls another helper, and reads a non-solo const 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 + 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 + ti.cast(j, ti.f32) * SEED_B + 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 solo const under one name. +# hill.d is a device accessor, hill.lap a specialized func, hill.dt a literal. +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 * 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 + # base level: pin the east and west edges, so the ridge drains outward + if j == 0 or j == grid.n - 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) From 179ea537c216e502b793a1177b1a859821f134cd Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 27 Jul 2026 11:37:47 +0200 Subject: [PATCH 23/31] Work on the cache. Taichi has some limitations --- README.md | 7 +- .../core/context/_closure_backend.py | 22 ++++- .../experimental/core/context/cupy_backend.py | 98 +++++++++++++------ 3 files changed, 91 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index b312953..807102b 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,12 @@ # PyFastFlow -**First full-GPU geomorphological and hydrodynamic toolbox powered by Taichi-lang** +**GPU-centred framework for geomorphological and hydrodynamic toolbox** + +- Developping, ditributing and reusing multi-backend GPU routines (Taichi lang, quadrants and cupy) +- Flow and flood routines for flexible grids and stencils (e.g. periodic conditions, no data areas, outlet location) +- [![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 diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 8c0c88c..0ca81fa 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -20,6 +20,7 @@ import ast import copy +import hashlib import inspect import linecache from types import FunctionType @@ -27,7 +28,6 @@ import numpy as np -from ..pool.base import new_uid from .base import ( Bag, HelperBuilder, @@ -671,7 +671,17 @@ def _fuse_group(self, group: list, group_index: int, dump_source: "str | None"): if alias is not None: exec_globals.setdefault(alias, backend) - func_name = f"_fused_group{group_index}_{new_uid()}" + # 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], @@ -681,8 +691,14 @@ def _fuse_group(self, group: list, group_index: int, dump_source: "str | None"): kwarg=None, defaults=[], ) - fused_def = ast.FunctionDef(name=func_name, args=args_node, body=body_stmts, decorator_list=[], returns=None) + 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"" diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 10bfc07..ad7c9f6 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -21,9 +21,16 @@ helpers' - is collected once, deduplicated by uid, into a module-scope constant block: - struct pf_params_t { float* p_; const float* p_; ... }; + 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 @@ -32,12 +39,12 @@ function body one local is declared per pointer that function's own spans reference: - const float* __restrict__ p_ = pf_params.p_; + 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. +`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, @@ -49,7 +56,9 @@ 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, uid-qualified struct members included, so it remains a sound key. +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. @@ -153,18 +162,22 @@ def _walk(path: list[str], bindings: dict[str, Any]): return obj -def _param_argname(param: Parameter) -> str: +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 since it is derived from - `uid`, not from whatever name(s) it happens to be bound under. This is - what makes dedup by uid automatic: two spans reaching the same Parameter - under two different handles compute the same argname and therefore the - same struct member. + 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_{param.uid}" + return f"p_{local_index[param.uid]}" class _SpanParser: @@ -197,6 +210,9 @@ 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: @@ -204,12 +220,16 @@ def _register_ptr(self, param: Parameter, write: bool) -> str: 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) + return _param_argname(param, local_index) def _expand_param(self, param: Parameter, method: str, call_args: list[str]) -> str: if method == "get": @@ -261,10 +281,11 @@ def parse(self, body: str) -> str: Author: B.G (07/2026) """ expanded = _SPAN_RE.sub(self._repl, body) - return _insert_locals(expanded, self.local_ptrs) + 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]) -> str: +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 @@ -276,6 +297,11 @@ def _insert_locals(body: str, local_ptrs: dict[int, dict]) -> str: 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: @@ -284,20 +310,21 @@ def _insert_locals(body: str, local_ptrs: dict[int, dict]) -> str: 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(uid)} = pf_params.{_argname_for(uid)};\n" - for uid, e in sorted(local_ptrs.items()) + 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(uid: int) -> str: +def _argname_for(local_idx: int) -> str: """ - The struct member / local name for a pointer already registered under - `uid` - see _param_argname, which this must stay in lockstep with. + 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_{uid}" + return f"p_{local_idx}" def _const_defines(bindings: dict[str, Any], body: str) -> list[str]: @@ -320,32 +347,36 @@ def _const_defines(bindings: dict[str, Any], body: str) -> list[str]: ] -def _param_block_source(registry: dict[int, dict]) -> str: +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 uid, ascending - arbitrary but fixed, so it agrees with - the upload order _upload_param_block writes in. + 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(uid)};\n" - for uid, e in sorted(registry.items()) + 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]) -> None: +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 uid order the struct was emitted - in. + `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 @@ -357,7 +388,10 @@ def _upload_param_block(module: "cp.RawModule", registry: dict[int, dict]) -> No if not registry: return global_ptr = module.get_global("pf_params") - ptrs = np.array([e["array"].data.ptr for _, e in sorted(registry.items())], dtype=np.uint64) + 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) @@ -619,6 +653,7 @@ def compile(self) -> CupyKernel: """ ctx = _SpecializeCtx() ctx.cupy_ptr_registry = {} + ctx.cupy_local_index = {} template = self._template name = _extract_name(_KERNEL_NAME_RE, template, "__global__") parser = _SpanParser(self._bindings, ctx=ctx) @@ -629,8 +664,9 @@ def compile(self) -> CupyKernel: 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)] + [_param_block_source(registry, local_index)] + list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body] @@ -640,7 +676,7 @@ def compile(self) -> CupyKernel: if module is None: module = cp.RawModule(code=source) CupyKernel._raw_cache[source] = module - _upload_param_block(module, registry) + _upload_param_block(module, registry, local_index) raw = module.get_function(name) krn = CupyKernel(name, raw, module) krn._final_source = source From 1725e9df23289b1a237af64504d41f6e5bdf771b Mon Sep 17 00:00:00 2001 From: bgailleton Date: Mon, 27 Jul 2026 15:00:35 +0200 Subject: [PATCH 24/31] Update README --- README.md | 145 +++++++++++++++--------------------------------------- 1 file changed, 40 insertions(+), 105 deletions(-) diff --git a/README.md b/README.md index 807102b..8fcd8a4 100644 --- a/README.md +++ b/README.md @@ -1,128 +1,63 @@ # PyFastFlow -**GPU-centred framework for geomorphological and hydrodynamic toolbox** +> ⚠️ **Experimental.** This describes the in-progress v1 core. The API is +> settling and will change. -- Developping, ditributing and reusing multi-backend GPU routines (Taichi lang, quadrants and cupy) -- Flow and flood routines for flexible grids and stencils (e.g. periodic conditions, no data areas, outlet location) -- +**GPU routines for Earth-surface processes — portable across backends.** -[![Python](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/) -[![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** - -**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). From 20130e7197d364741676e27fa0cd57937ba5761e Mon Sep 17 00:00:00 2001 From: bgailleton Date: Tue, 28 Jul 2026 10:07:31 +0200 Subject: [PATCH 25/31] Fix a cupy key error and other minor bug --- README.md | 2 +- .../experimental/core/context/cupy_backend.py | 75 ++++++++++++++----- 2 files changed, 57 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 8fcd8a4..7642f11 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ Two things in one: ## Why -Writing flexible GPU kernels for physics simulation can be time-consuming. Performance on GPU are highly affected by memory layout. Hence simple things as constant _vs_ 2D varying parameters or stencil operations (periodic boundaries, no data, ...) +Writing flexible GPU kernels for physics simulation can be time-consuming. Performance on GPU is highly affected by memory layout. Hence simple things such as a constant _vs_ a 2D varying parameter, or stencil operations (periodic boundaries, no-data, outlets, ...), each push you toward a different hand-tuned kernel — multiplied by every backend you want to run on. PyFastFlow lets you write a routine once and *bind* how each parameter and stencil is realised, then specialises it per backend and layout. ## Example diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index ad7c9f6..383ad88 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -197,6 +197,16 @@ class _SpanParser: 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) """ @@ -204,7 +214,6 @@ def __init__(self, bindings: dict[str, Any], *, ctx: _SpecializeCtx): self.bindings = bindings self.ctx = ctx self.local_ptrs: dict[int, dict] = {} # uid -> {ctype, write} - self.helper_srcs: dict[str, str] = {} # name -> source def _register_ptr(self, param: Parameter, write: bool) -> str: registry = getattr(self.ctx, "cupy_ptr_registry", None) @@ -266,7 +275,6 @@ def _repl(self, match: re.Match) -> str: if isinstance(target, CupyHelperBuilder): target = self.ctx.specialize(target) if isinstance(target, CupyHelper): - self.helper_srcs[target.name] = target.compiled return f"{target.name}({argstr if argstr is not None else ''})" if isinstance(target, Parameter): return self._expand_param(target, "get", ["0"]) @@ -274,9 +282,10 @@ def _repl(self, match: re.Match) -> str: def parse(self, body: str) -> str: """ - Expand every `$...$` span in `body`, accumulating local_ptrs and - helper_srcs as a side effect, then prepend the `__restrict__` locals - this body's own spans implied - see _insert_locals. + 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) """ @@ -599,13 +608,25 @@ class CupyHelperBuilder(HelperBuilder): def _specialize(self, ctx: _SpecializeCtx) -> CupyHelper: """ - Expand the template's spans against `ctx` and prepend the helper - sources and const #defines it turned out to need, giving the final - `__device__` text. 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. + 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) """ @@ -613,7 +634,11 @@ def _specialize(self, ctx: _SpecializeCtx) -> CupyHelper: name = _extract_name(_DEVICE_NAME_RE, template, "__device__") parser = _SpanParser(self._bindings, ctx=ctx) body = parser.parse(template) - source = "\n".join(list(parser.helper_srcs.values()) + _const_defines(self._bindings, body) + [body]) + 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) fn = CupyHelper(name, source) attach_meta(fn, template, self._bindings) return fn @@ -631,17 +656,28 @@ class CupyKernelBuilder(KernelBuilder): def compile(self) -> CupyKernel: """ Expand the template's spans - the kernel's own and, recursively, - every CupyHelperBuilder they reach - prepend the helpers' source, - 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. + 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 @@ -654,6 +690,7 @@ def compile(self) -> CupyKernel: 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) @@ -667,7 +704,7 @@ def compile(self) -> CupyKernel: local_index = ctx.cupy_local_index source = "\n".join( [_param_block_source(registry, local_index)] - + list(parser.helper_srcs.values()) + + list(ctx.cupy_device_srcs.values()) + _const_defines(self._bindings, body) + [body] ) From d8fbac8451de6cbbf28b5e548c26424a1ce01438 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Tue, 28 Jul 2026 11:24:07 +0200 Subject: [PATCH 26/31] Add the grid bag helpers --- pyfastflow/experimental/grid/__init__.py | 196 +++++++ .../experimental/grid/_closure_blocks.py | 540 ++++++++++++++++++ pyfastflow/experimental/grid/_cupy_blocks.py | 440 ++++++++++++++ 3 files changed, 1176 insertions(+) create mode 100644 pyfastflow/experimental/grid/__init__.py create mode 100644 pyfastflow/experimental/grid/_closure_blocks.py create mode 100644 pyfastflow/experimental/grid/_cupy_blocks.py diff --git a/pyfastflow/experimental/grid/__init__.py b/pyfastflow/experimental/grid/__init__.py new file mode 100644 index 0000000..fb863e1 --- /dev/null +++ b/pyfastflow/experimental/grid/__init__.py @@ -0,0 +1,196 @@ +""" +make_grid: the GridContext-equivalent Bag factory, built on the +backend-agnostic core (see ..core.context.base for Parameter/HelperBuilder/Bag). + +There is no stateful Context class here - by design, the core has none (see +core/context/base.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", non-solo, for all three - see the note below. + - 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. + +Note on nx/ny solo=True: earlier drafts of this design set nx/ny to a +default of const+solo=True (read bare, as a compile-time literal with no +`.get()`), matching how n_neighbours is read. That is incompatible with the +"read every value param mode-agnostically via `.get()`" requirement the +private blocks are built around - a solo Parameter resolves to a bare python +literal in device code (see base.py's resolve_binding), which has no `.get` +method, so `NX.get(0)` inside a shared geometry block would fail to trace +the moment nx defaulted to solo. To keep one geometry block working +whatever mode nx/ny/dx are in - the actual point of the mode-overridable +design - nx/ny/dx are built here as solo=False by default (same as dx's +own, unambiguous spec: "non-solo, read via .get(0)"). n_neighbours, which +has no override and is genuinely only ever a compile-time literal, keeps +solo=True. + +Author: B.G (07/2026) +""" + +import numpy as np + +from ..core.context.base import Bag + +_TOPOLOGIES = {"D4": 4, "D8": 8} +_BOUNDARIES = frozenset({"normal", "periodic_EW", "periodic_NS"}) +_OUTLETS = frozenset({"edge", "mask"}) + + +def _backend_classes(backend: str): + """ + (backend_module_or_None, ParameterCls, HelperBuilderCls, blocks_module, + dtypes) for one backend name. cupy's backend_module is None - its blocks + call plain C, never a bound backend module - see _cupy_blocks.py. + + Author: B.G (07/2026) + """ + if backend == "taichi": + import taichi as ti + + from ..core.context.taichi_backend import TaichiHelperBuilder, TaichiParameter + from . import _closure_blocks as blocks + + return ti, TaichiParameter, TaichiHelperBuilder, blocks, {"i32": ti.i32, "f32": ti.f32, "u8": ti.u8} + if backend == "quadrants": + import quadrants as qd + + from ..core.context.quadrants_backend import QuadrantsHelperBuilder, QuadrantsParameter + from . import _closure_blocks as blocks + + return qd, QuadrantsParameter, QuadrantsHelperBuilder, blocks, {"i32": qd.i32, "f32": qd.f32, "u8": qd.u8} + if backend == "cupy": + from ..core.context.cupy_backend import CupyHelperBuilder, CupyParameter + from . import _cupy_blocks as blocks + + return None, CupyParameter, CupyHelperBuilder, blocks, { + "i32": np.int32, + "f32": np.float32, + "u8": np.uint8, + } + raise ValueError(f"make_grid: unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + + +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, blocks, dtypes = _backend_classes(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, solo=True + ) + + 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..783c139 --- /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 base.py, +_SpecializeCtx). + +nx/ny/dx are read exclusively through `.get(...)`, uniformly across whatever +mode they end up in (const, scalar, field) - see base.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) +""" + +# --------------------------------------------------------------------------- +# 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" + + def mk(template, **binds): + b = HelperCls().ingest(template) + for name, obj in binds.items(): + b.bind(name, obj) + return b + + 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..2d6646a --- /dev/null +++ b/pyfastflow/experimental/grid/_cupy_blocks.py @@ -0,0 +1,440 @@ +""" +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 math + +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) + + def mk(template, **binds): + b = HelperCls().ingest(template) + for name, obj in binds.items(): + b.bind(name, obj) + return b + + 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, + } From 6bef1c17259b08e277762bf980042b036a48cd90 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Tue, 28 Jul 2026 16:17:48 +0200 Subject: [PATCH 27/31] gjj0 --- pyfastflow/experimental/core/context/base.py | 4 +--- pyfastflow/experimental/core/context/cupy_backend.py | 6 +++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/base.py index 6b1a5f6..9abfcd4 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -87,9 +87,7 @@ 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. A const Parameter declared solo=True -is the exception: it resolves to a bare compile-time literal, read as `p` with -no call. +read, p.set_node(node, value) to write. What a device helper may bind ----------------------------- diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 383ad88..ff535fe 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -77,7 +77,10 @@ from .routine import Routine, RoutineBuilder, _CompiledStep _KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") -_DEVICE_NAME_RE = re.compile(r"__device__\s+[\w:\*&]+\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) @@ -88,6 +91,7 @@ np.dtype(np.int32): "int", np.dtype(np.int64): "long long", np.dtype(np.uint8): "unsigned char", + np.dtype(np.uint32): "unsigned int", } From 95aaad65ccfb396c4156546d0b622a92e6fddaa2 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Tue, 28 Jul 2026 17:44:12 +0200 Subject: [PATCH 28/31] Clean some routine inherited from first designs and remove unused bit. Add the noise helpers for perlin and all --- .../heat_diffusion/heat_diffusion_cupy.py | 43 ++-- .../heat_diffusion_quadrants.py | 108 ++++----- .../heat_diffusion_routine_cupy.py | 36 +-- .../heat_diffusion_routine_quadrants.py | 78 +++---- .../heat_diffusion_routine_taichi.py | 78 +++---- .../heat_diffusion/heat_diffusion_taichi.py | 108 ++++----- examples/core/lem/lem_routine_cupy.py | 48 ++-- examples/core/lem/lem_routine_quadrants.py | 48 ++-- examples/core/lem/lem_routine_taichi.py | 48 ++-- .../core/shallow_water/shallow_water_cupy.py | 6 +- .../shallow_water/shallow_water_quadrants.py | 24 +- .../shallow_water/shallow_water_taichi.py | 24 +- .../core/context/_closure_backend.py | 29 +-- .../experimental/core/context/backends.py | 82 +++++++ pyfastflow/experimental/core/context/base.py | 110 ++------- .../experimental/core/context/cupy_backend.py | 23 +- pyfastflow/experimental/grid/__init__.py | 57 ++--- .../experimental/grid/_closure_blocks.py | 10 +- pyfastflow/experimental/grid/_cupy_blocks.py | 8 +- pyfastflow/experimental/noise/__init__.py | 203 ++++++++++++++++ .../experimental/noise/_closure_blocks.py | 216 ++++++++++++++++++ pyfastflow/experimental/noise/_cupy_blocks.py | 210 +++++++++++++++++ 22 files changed, 1095 insertions(+), 502 deletions(-) create mode 100644 pyfastflow/experimental/core/context/backends.py create mode 100644 pyfastflow/experimental/noise/__init__.py create mode 100644 pyfastflow/experimental/noise/_closure_blocks.py create mode 100644 pyfastflow/experimental/noise/_cupy_blocks.py diff --git a/examples/core/heat_diffusion/heat_diffusion_cupy.py b/examples/core/heat_diffusion/heat_diffusion_cupy.py index b5e0cd9..f3a542f 100644 --- a/examples/core/heat_diffusion/heat_diffusion_cupy.py +++ b/examples/core/heat_diffusion/heat_diffusion_cupy.py @@ -66,22 +66,22 @@ 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, solo=True) -room_p = CupyParameter("ROOM", dtype=np.int32, mode="const", value=GRID_N // 4, pool=pool, solo=True) -wall_thick_p = CupyParameter("WALL_THICK", dtype=np.int32, mode="const", value=8, pool=pool, solo=True) -door_p = CupyParameter("DOOR", dtype=np.int32, mode="const", value=6, pool=pool, solo=True) -seed_p = CupyParameter("SEED", dtype=np.float32, mode="const", value=17.0, pool=pool, solo=True) +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, solo=True) -dx2_p = CupyParameter("DX2", dtype=np.float32, mode="const", value=DX_M**2, pool=pool, solo=True) +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, solo=True) -alpha_wall_seed_p = CupyParameter("ALPHA_WALL_SEED", dtype=np.float32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) -t_bg_p = CupyParameter("T_BG", dtype=np.float32, mode="const", value=15.0, pool=pool, solo=True) +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, solo=True) -src_j_p = CupyParameter("SRC_J", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) -src_r_p = CupyParameter("SRC_R", dtype=np.int32, mode="const", value=10, pool=pool, solo=True) +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 @@ -215,7 +215,7 @@ # 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 - the solo consts expand to CUDA literals, the scalar +# 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( @@ -235,9 +235,9 @@ __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$; - int dy = idx % N - $stove.at.j$; - if (dx * dx + dy * dy <= $stove.r$ * $stove.r$) { + 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)$; } } @@ -247,8 +247,9 @@ ) # A MIXED Bag: everything the diffusion step needs, whatever kind it is - a -# field Parameter, a device helper, two solo consts - under one name. A bag has -# no member type; each member is resolved on its own when the spans expand. +# 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}) @@ -265,8 +266,8 @@ int i = idx / N; int j = idx % N; float a = $heat.alpha.get(idx)$; - float lap = $heat.lap(T_in, i, j)$ / $heat.dx2$; - T_out[idx] = T_in[idx] + $heat.dt$ * a * lap; + float lap = $heat.lap(T_in, i, j)$ / $heat.dx2.get(0)$; + T_out[idx] = T_in[idx] + $heat.dt.get(0)$ * a * lap; } """ ) diff --git a/examples/core/heat_diffusion/heat_diffusion_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_quadrants.py index 824a5cc..96bce93 100644 --- a/examples/core/heat_diffusion/heat_diffusion_quadrants.py +++ b/examples/core/heat_diffusion/heat_diffusion_quadrants.py @@ -17,11 +17,10 @@ - diffuse: explicit FTCS heat equation dT/dt = alpha(i,j) * lap(T), with a clamped (Neumann / no-flux) boundary Laplacian. -Uniform device surface: `wall`, `alpha` and `stove` are 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. Structural constants -(N, ROOM, ...) are declared solo=True and used bare as compile-time literals. +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); @@ -87,25 +86,25 @@ pool = QuadrantsPool() -# Structural constants: solo=True -> read bare in kernel bodies as compile-time -# literals (no .get()), since they are never going to vary at runtime. -n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) -room_p = QuadrantsParameter("ROOM", dtype=qd.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) -wall_thick_p = QuadrantsParameter("WALL_THICK", dtype=qd.i32, mode="const", value=8, pool=pool, solo=True) -door_p = QuadrantsParameter("DOOR", dtype=qd.i32, mode="const", value=6, pool=pool, solo=True) -seed_p = QuadrantsParameter("SEED", dtype=qd.f32, mode="const", value=17.0, pool=pool, solo=True) +# 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, solo=True) # seconds -dx2_p = QuadrantsParameter("DX2", dtype=qd.f32, mode="const", value=DX_M**2, pool=pool, solo=True) # meters^2 +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 - used bare inside set_alpha. -alpha_air_seed_p = QuadrantsParameter("ALPHA_AIR_SEED", dtype=qd.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) -alpha_wall_seed_p = QuadrantsParameter("ALPHA_WALL_SEED", dtype=qd.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) -t_bg_p = QuadrantsParameter("T_BG", dtype=qd.f32, mode="const", value=15.0, pool=pool, solo=True) +# 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, solo=True) -src_j_p = QuadrantsParameter("SRC_J", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) -src_r_p = QuadrantsParameter("SRC_R", dtype=qd.i32, mode="const", value=10, pool=pool, solo=True) # stove radius, cells +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). @@ -126,7 +125,7 @@ def clamp(i): - return min(max(i, 0), N - 1) + return min(max(i, 0), N.get(0) - 1) clamp_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clamp) @@ -145,7 +144,7 @@ def laplacian(field_, i, j): 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 + 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) @@ -158,25 +157,25 @@ def whash(a, b): def generate_walls_template(): - for i, j in qd.ndrange(N, N): + for i, j in qd.ndrange(N.get(0), N.get(0)): is_wall = 0 - if i < WALL_THICK or i >= N - WALL_THICK or j < WALL_THICK or j >= N - WALL_THICK: + 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 < WALL_THICK: - vline = i // ROOM - seg = j // ROOM - door = qd.cast(whash(vline, seg) * ROOM, qd.i32) - gap = (j % ROOM) >= door and (j % ROOM) < door + DOOR + 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 < WALL_THICK: - hline = j // ROOM - seg = i // ROOM - door = qd.cast(whash(hline + 7919, seg) * ROOM, qd.i32) - gap = (i % ROOM) >= door and (i % ROOM) < door + DOOR + 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 + j, is_wall) + wall.set_node(i * N.get(0) + j, is_wall) generate_walls_kernel = ( @@ -193,12 +192,12 @@ def generate_walls_template(): def set_alpha_template(): - for i, j in qd.ndrange(N, N): - idx = i * N + j + 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) + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) else: - alpha.set_node(idx, ALPHA_AIR_SEED) + alpha.set_node(idx, ALPHA_AIR_SEED.get(0)) # The two seed values are grouped on the host for tidiness, then merged in with @@ -221,17 +220,17 @@ def set_alpha_template(): def init_temperature_template(T: qd.Tensor): for i, j in T: - T[i, j] = T_BG + 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, and members of either level resolve on their -# own type - the solo consts become compile-time literals (read bare), the -# scalar Parameter keeps its .get(0). Everything else here still binds flat, so -# the two styles sit side by side in one file. +# 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}), @@ -243,9 +242,9 @@ def init_temperature_template(T: qd.Tensor): def apply_source_template(T: qd.Tensor): for i, j in T: - dx = i - stove.at.i - dy = j - stove.at.j - if dx * dx + dy * dy <= stove.r * stove.r: + 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) @@ -253,18 +252,19 @@ def apply_source_template(T: qd.Tensor): # A MIXED Bag: everything the diffusion step needs, whatever kind it is - a -# field Parameter, a device helper, two solo consts - 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, `heat.dx2` a literal. +# 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 + j + idx = i * N.get(0) + j a = heat.alpha.get(idx) - lap = heat.lap(T_in, i, j) / heat.dx2 - T_out[i, j] = T_in[i, j] + heat.dt * a * lap + 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() diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py index ae8d0a4..93cb55a 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py @@ -75,22 +75,22 @@ pool = CupyPool() -n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool, solo=True) -room_p = CupyParameter("ROOM", dtype=np.int32, mode="const", value=GRID_N // 4, pool=pool, solo=True) -wall_thick_p = CupyParameter("WALL_THICK", dtype=np.int32, mode="const", value=8, pool=pool, solo=True) -door_p = CupyParameter("DOOR", dtype=np.int32, mode="const", value=6, pool=pool, solo=True) -seed_p = CupyParameter("SEED", dtype=np.float32, mode="const", value=17.0, pool=pool, solo=True) +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, solo=True) -dx2_p = CupyParameter("DX2", dtype=np.float32, mode="const", value=DX_M**2, pool=pool, solo=True) +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, solo=True) -alpha_wall_seed_p = CupyParameter("ALPHA_WALL_SEED", dtype=np.float32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) -t_bg_p = CupyParameter("T_BG", dtype=np.float32, mode="const", value=15.0, pool=pool, solo=True) +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, solo=True) -src_j_p = CupyParameter("SRC_J", dtype=np.int32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) -src_r_p = CupyParameter("SRC_R", dtype=np.int32, mode="const", value=10, pool=pool, solo=True) +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) @@ -236,9 +236,9 @@ __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$; - int dy = idx % N - $stove.at.j$; - if (dx * dx + dy * dy <= $stove.r$ * $stove.r$) { + 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)$; } } @@ -261,8 +261,8 @@ int i = idx / N; int j = idx % N; float a = $heat.alpha.get(idx)$; - float lap = $heat.lap(T_in, i, j)$ / $heat.dx2$; - T_out[idx] = T_in[idx] + $heat.dt$ * a * lap; + float lap = $heat.lap(T_in, i, j)$ / $heat.dx2.get(0)$; + T_out[idx] = T_in[idx] + $heat.dt.get(0)$ * a * lap; } """ ) diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py index e87cb0d..b314fa4 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py @@ -69,22 +69,22 @@ pool = QuadrantsPool() -n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) -room_p = QuadrantsParameter("ROOM", dtype=qd.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) -wall_thick_p = QuadrantsParameter("WALL_THICK", dtype=qd.i32, mode="const", value=8, pool=pool, solo=True) -door_p = QuadrantsParameter("DOOR", dtype=qd.i32, mode="const", value=6, pool=pool, solo=True) -seed_p = QuadrantsParameter("SEED", dtype=qd.f32, mode="const", value=17.0, pool=pool, solo=True) +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, solo=True) -dx2_p = QuadrantsParameter("DX2", dtype=qd.f32, mode="const", value=DX_M**2, pool=pool, solo=True) +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, solo=True) -alpha_wall_seed_p = QuadrantsParameter("ALPHA_WALL_SEED", dtype=qd.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) -t_bg_p = QuadrantsParameter("T_BG", dtype=qd.f32, mode="const", value=15.0, pool=pool, solo=True) +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, solo=True) -src_j_p = QuadrantsParameter("SRC_J", dtype=qd.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) -src_r_p = QuadrantsParameter("SRC_R", dtype=qd.i32, mode="const", value=10, pool=pool, solo=True) +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) @@ -98,7 +98,7 @@ def clamp(i): - return min(max(i, 0), N - 1) + return min(max(i, 0), N.get(0) - 1) clamp_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clamp) @@ -117,7 +117,7 @@ def laplacian(field_, i, j): 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 + 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) @@ -131,25 +131,25 @@ def whash(a, b): def generate_walls_template(): - for i, j in qd.ndrange(N, N): + for i, j in qd.ndrange(N.get(0), N.get(0)): is_wall = 0 - if i < WALL_THICK or i >= N - WALL_THICK or j < WALL_THICK or j >= N - WALL_THICK: + 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 < WALL_THICK: - vline = i // ROOM - seg = j // ROOM - door = qd.cast(whash(vline, seg) * ROOM, qd.i32) - gap = (j % ROOM) >= door and (j % ROOM) < door + DOOR + 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 < WALL_THICK: - hline = j // ROOM - seg = i // ROOM - door = qd.cast(whash(hline + 7919, seg) * ROOM, qd.i32) - gap = (i % ROOM) >= door and (i % ROOM) < door + DOOR + 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 + j, is_wall) + wall.set_node(i * N.get(0) + j, is_wall) generate_walls_kernel = ( @@ -166,12 +166,12 @@ def generate_walls_template(): def set_alpha_template(): - for i, j in qd.ndrange(N, N): - idx = i * N + j + 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) + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) else: - alpha.set_node(idx, ALPHA_AIR_SEED) + 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}) @@ -189,7 +189,7 @@ def set_alpha_template(): def init_temperature_template(T: qd.Tensor): for i, j in T: - T[i, j] = T_BG + T[i, j] = T_BG.get(0) init_temperature_kernel = QuadrantsKernelBuilder().bind("T_BG", t_bg_p).ingest(init_temperature_template).compile() @@ -205,9 +205,9 @@ def init_temperature_template(T: qd.Tensor): def apply_source_template(T: qd.Tensor): for i, j in T: - dx = i - stove.at.i - dy = j - stove.at.j - if dx * dx + dy * dy <= stove.r * stove.r: + 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) @@ -222,10 +222,10 @@ def apply_source_template(T: qd.Tensor): def diffuse_template(T_out: qd.Tensor, T_in: qd.Tensor): for i, j in T_in: - idx = i * N + j + idx = i * N.get(0) + j a = heat.alpha.get(idx) - lap = heat.lap(T_in, i, j) / heat.dx2 - T_out[i, j] = T_in[i, j] + heat.dt * a * lap + 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) diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py index 7a6ebf4..8268770 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py @@ -69,22 +69,22 @@ pool = TaichiPool() -n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) -room_p = TaichiParameter("ROOM", dtype=ti.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) -wall_thick_p = TaichiParameter("WALL_THICK", dtype=ti.i32, mode="const", value=8, pool=pool, solo=True) -door_p = TaichiParameter("DOOR", dtype=ti.i32, mode="const", value=6, pool=pool, solo=True) -seed_p = TaichiParameter("SEED", dtype=ti.f32, mode="const", value=17.0, pool=pool, solo=True) +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, solo=True) -dx2_p = TaichiParameter("DX2", dtype=ti.f32, mode="const", value=DX_M**2, pool=pool, solo=True) +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, solo=True) -alpha_wall_seed_p = TaichiParameter("ALPHA_WALL_SEED", dtype=ti.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) -t_bg_p = TaichiParameter("T_BG", dtype=ti.f32, mode="const", value=15.0, pool=pool, solo=True) +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, solo=True) -src_j_p = TaichiParameter("SRC_J", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) -src_r_p = TaichiParameter("SRC_R", dtype=ti.i32, mode="const", value=10, pool=pool, solo=True) +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) @@ -98,7 +98,7 @@ def clamp(i): - return min(max(i, 0), N - 1) + return min(max(i, 0), N.get(0) - 1) clamp_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clamp) @@ -117,7 +117,7 @@ def laplacian(field_, i, j): 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 + 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) @@ -131,25 +131,25 @@ def whash(a, b): def generate_walls_template(): - for i, j in ti.ndrange(N, N): + for i, j in ti.ndrange(N.get(0), N.get(0)): is_wall = 0 - if i < WALL_THICK or i >= N - WALL_THICK or j < WALL_THICK or j >= N - WALL_THICK: + 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 < 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 + 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 < 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 + 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 + j, is_wall) + wall.set_node(i * N.get(0) + j, is_wall) generate_walls_kernel = ( @@ -166,12 +166,12 @@ def generate_walls_template(): def set_alpha_template(): - for i, j in ti.ndrange(N, N): - idx = i * N + j + 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) + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) else: - alpha.set_node(idx, ALPHA_AIR_SEED) + 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}) @@ -189,7 +189,7 @@ def set_alpha_template(): def init_temperature_template(T: ti.template()): for i, j in T: - T[i, j] = T_BG + T[i, j] = T_BG.get(0) init_temperature_kernel = TaichiKernelBuilder().bind("T_BG", t_bg_p).ingest(init_temperature_template).compile() @@ -205,9 +205,9 @@ def init_temperature_template(T: ti.template()): def apply_source_template(T: ti.template()): for i, j in T: - dx = i - stove.at.i - dy = j - stove.at.j - if dx * dx + dy * dy <= stove.r * stove.r: + 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) @@ -222,10 +222,10 @@ def apply_source_template(T: ti.template()): def diffuse_template(T_out: ti.template(), T_in: ti.template()): for i, j in T_in: - idx = i * N + j + idx = i * N.get(0) + j a = heat.alpha.get(idx) - lap = heat.lap(T_in, i, j) / heat.dx2 - T_out[i, j] = T_in[i, j] + heat.dt * a * lap + 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) diff --git a/examples/core/heat_diffusion/heat_diffusion_taichi.py b/examples/core/heat_diffusion/heat_diffusion_taichi.py index 54e1a63..e811cd3 100644 --- a/examples/core/heat_diffusion/heat_diffusion_taichi.py +++ b/examples/core/heat_diffusion/heat_diffusion_taichi.py @@ -17,11 +17,10 @@ - diffuse: explicit FTCS heat equation dT/dt = alpha(i,j) * lap(T), with a clamped (Neumann / no-flux) boundary Laplacian. -Uniform device surface: `wall`, `alpha` and `stove` are 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. Structural constants -(N, ROOM, ...) are declared solo=True and used bare as compile-time literals. +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); @@ -87,25 +86,25 @@ pool = TaichiPool() -# Structural constants: solo=True -> read bare in kernel bodies as compile-time -# literals (no .get()), since they are never going to vary at runtime. -n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) -room_p = TaichiParameter("ROOM", dtype=ti.i32, mode="const", value=GRID_N // 4, pool=pool, solo=True) -wall_thick_p = TaichiParameter("WALL_THICK", dtype=ti.i32, mode="const", value=8, pool=pool, solo=True) -door_p = TaichiParameter("DOOR", dtype=ti.i32, mode="const", value=6, pool=pool, solo=True) -seed_p = TaichiParameter("SEED", dtype=ti.f32, mode="const", value=17.0, pool=pool, solo=True) +# 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, solo=True) # seconds -dx2_p = TaichiParameter("DX2", dtype=ti.f32, mode="const", value=DX_M**2, pool=pool, solo=True) # meters^2 +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 - used bare inside set_alpha. -alpha_air_seed_p = TaichiParameter("ALPHA_AIR_SEED", dtype=ti.f32, mode="const", value=ALPHA_AIR_VAL, pool=pool, solo=True) -alpha_wall_seed_p = TaichiParameter("ALPHA_WALL_SEED", dtype=ti.f32, mode="const", value=ALPHA_WALL_VAL, pool=pool, solo=True) -t_bg_p = TaichiParameter("T_BG", dtype=ti.f32, mode="const", value=15.0, pool=pool, solo=True) +# 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, solo=True) -src_j_p = TaichiParameter("SRC_J", dtype=ti.i32, mode="const", value=GRID_N // 4 + GRID_N // 8, pool=pool, solo=True) -src_r_p = TaichiParameter("SRC_R", dtype=ti.i32, mode="const", value=10, pool=pool, solo=True) # stove radius, cells +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). @@ -126,7 +125,7 @@ def clamp(i): - return min(max(i, 0), N - 1) + return min(max(i, 0), N.get(0) - 1) clamp_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clamp) @@ -145,7 +144,7 @@ def laplacian(field_, i, j): 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 + 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) @@ -158,25 +157,25 @@ def whash(a, b): def generate_walls_template(): - for i, j in ti.ndrange(N, N): + for i, j in ti.ndrange(N.get(0), N.get(0)): is_wall = 0 - if i < WALL_THICK or i >= N - WALL_THICK or j < WALL_THICK or j >= N - WALL_THICK: + 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 < 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 + 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 < 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 + 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 + j, is_wall) + wall.set_node(i * N.get(0) + j, is_wall) generate_walls_kernel = ( @@ -193,12 +192,12 @@ def generate_walls_template(): def set_alpha_template(): - for i, j in ti.ndrange(N, N): - idx = i * N + j + 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) + alpha.set_node(idx, ALPHA_WALL_SEED.get(0)) else: - alpha.set_node(idx, ALPHA_AIR_SEED) + alpha.set_node(idx, ALPHA_AIR_SEED.get(0)) # The two seed values are grouped on the host for tidiness, then merged in with @@ -221,17 +220,17 @@ def set_alpha_template(): def init_temperature_template(T: ti.template()): for i, j in T: - T[i, j] = T_BG + 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, and members of either level resolve on their -# own type - the solo consts become compile-time literals (read bare), the -# scalar Parameter keeps its .get(0). Everything else here still binds flat, so -# the two styles sit side by side in one file. +# 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}), @@ -243,9 +242,9 @@ def init_temperature_template(T: ti.template()): def apply_source_template(T: ti.template()): for i, j in T: - dx = i - stove.at.i - dy = j - stove.at.j - if dx * dx + dy * dy <= stove.r * stove.r: + 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) @@ -253,18 +252,19 @@ def apply_source_template(T: ti.template()): # A MIXED Bag: everything the diffusion step needs, whatever kind it is - a -# field Parameter, a device helper, two solo consts - 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, `heat.dx2` a literal. +# 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 + j + idx = i * N.get(0) + j a = heat.alpha.get(idx) - lap = heat.lap(T_in, i, j) / heat.dx2 - T_out[i, j] = T_in[i, j] + heat.dt * a * lap + 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() diff --git a/examples/core/lem/lem_routine_cupy.py b/examples/core/lem/lem_routine_cupy.py index f59c491..21f0b95 100644 --- a/examples/core/lem/lem_routine_cupy.py +++ b/examples/core/lem/lem_routine_cupy.py @@ -11,10 +11,13 @@ needs all of them at once. The other examples each isolate one thing; this one carries the lot: - Parameter modes N and DT are solo consts, arriving as #defines and read - bare. DX is a non-solo const, read as $grid.dx.get(0)$. - D and SEA_LEVEL are scalars the host retunes between - frames. UPLIFT is a field, one rate per node. + 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 @@ -82,14 +85,15 @@ # --------------------------------------------------------------------------- # parameters - one of every mode # --------------------------------------------------------------------------- -# solo consts: emitted as #defines and read bare in the source -n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool, solo=True) -dt_p = CupyParameter("DT", dtype=np.float32, mode="const", value=DT_VAL, pool=pool, solo=True) -seed_a_p = CupyParameter("SEED_A", dtype=np.float32, mode="const", value=12.9898, pool=pool, solo=True) -seed_b_p = CupyParameter("SEED_B", dtype=np.float32, mode="const", value=78.233, pool=pool, solo=True) - -# non-solo const: still fixed at compile time, but read through a span like -# any other mode, so a template can be written without knowing it is const +# 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 @@ -107,8 +111,9 @@ # --------------------------------------------------------------------------- # bags # --------------------------------------------------------------------------- -# nested: grid.n is a solo const read bare, grid.dx a non-solo const reached -# through a span - members resolve on their own type, not the bag's +# 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 @@ -131,7 +136,7 @@ .ingest( r""" __device__ float laplacian(const float* f, int i, int j) { - // calls another helper, and reads a non-solo const out of a bound bag + // 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)$; @@ -185,8 +190,8 @@ # routine kernels # --------------------------------------------------------------------------- -# mixed bag: a scalar Parameter, a helper and a solo const under one name. -# hill.d is a span read, hill.lap a spliced call, hill.dt a bare literal. +# 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 = ( @@ -199,7 +204,7 @@ 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$ * $hill.d.get(0)$ * $hill.lap(z_in, i, j)$; + z_out[idx] = z_in[idx] + $hill.dt.get(0)$ * $hill.d.get(0)$ * $hill.lap(z_in, i, j)$; } """ ) @@ -214,9 +219,10 @@ .ingest( r""" extern "C" __global__ void uplift_bc(float* z) { - // grid.n is a solo const reached through the bag, so it splices in as a - // bare literal here just as it would read bare if bound flat - int n = $grid.n$; + // 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; diff --git a/examples/core/lem/lem_routine_quadrants.py b/examples/core/lem/lem_routine_quadrants.py index 77993bc..86d3371 100644 --- a/examples/core/lem/lem_routine_quadrants.py +++ b/examples/core/lem/lem_routine_quadrants.py @@ -11,10 +11,11 @@ needs all of them at once. The other examples each isolate one thing; this one carries the lot: - Parameter modes N and DT are solo consts, read bare as compile-time - literals. DX is a non-solo const, read as grid.dx.get(0). - D and SEA_LEVEL are scalars the host retunes between - frames. UPLIFT is a field, one rate per node. + 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 @@ -72,14 +73,14 @@ # --------------------------------------------------------------------------- # parameters - one of every mode # --------------------------------------------------------------------------- -# solo consts: folded into the generated code as bare literals, read as `N` -n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) -dt_p = QuadrantsParameter("DT", dtype=qd.f32, mode="const", value=DT_VAL, pool=pool, solo=True) -seed_a_p = QuadrantsParameter("SEED_A", dtype=qd.f32, mode="const", value=12.9898, pool=pool, solo=True) -seed_b_p = QuadrantsParameter("SEED_B", dtype=qd.f32, mode="const", value=78.233, pool=pool, solo=True) - -# non-solo const: still fixed at compile time, but read through .get(0) like -# any other mode, so a template can be written without knowing it is const +# 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 @@ -99,8 +100,8 @@ # --------------------------------------------------------------------------- # bags # --------------------------------------------------------------------------- -# nested: grid.n is a solo const read bare, grid.dx a non-solo const read -# through .get(0) - members resolve on their own type, not the bag's +# 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 @@ -112,14 +113,14 @@ def clampi(i): - return min(max(i, 0), N - 1) + 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 non-solo const out of a bound bag + # 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) @@ -134,7 +135,7 @@ def laplacian(field_, i, j): 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 + j) + return UPLIFT.get(i * N.get(0) + j) uplift_at_fn = QuadrantsHelperBuilder().bind("UPLIFT", uplift_p).bind("N", n_p).ingest(uplift_at) @@ -147,7 +148,7 @@ def uplift_at(i, j): 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 + qd.cast(j, qd.f32) * SEED_B + 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 @@ -158,14 +159,15 @@ def init_topo_template(z: qd.Tensor): # routine kernels # --------------------------------------------------------------------------- -# mixed bag: a scalar Parameter, a helper and a solo const under one name. -# hill.d is a device accessor, hill.lap a specialized func, hill.dt a literal. +# 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 * hill.d.get(0) * hill.lap(z_in, i, j) + 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) @@ -173,9 +175,9 @@ def diffuse_template(z_out: qd.Tensor, z_in: qd.Tensor): def uplift_template(z: qd.Tensor): for i, j in z: - z[i, j] += up(i, j) * DT + 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 - 1: + if j == 0 or j == grid.n.get(0) - 1: z[i, j] = SEA.get(0) diff --git a/examples/core/lem/lem_routine_taichi.py b/examples/core/lem/lem_routine_taichi.py index eac8178..640ca08 100644 --- a/examples/core/lem/lem_routine_taichi.py +++ b/examples/core/lem/lem_routine_taichi.py @@ -11,10 +11,11 @@ needs all of them at once. The other examples each isolate one thing; this one carries the lot: - Parameter modes N and DT are solo consts, read bare as compile-time - literals. DX is a non-solo const, read as grid.dx.get(0). - D and SEA_LEVEL are scalars the host retunes between - frames. UPLIFT is a field, one rate per node. + 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 @@ -72,14 +73,14 @@ # --------------------------------------------------------------------------- # parameters - one of every mode # --------------------------------------------------------------------------- -# solo consts: folded into the generated code as bare literals, read as `N` -n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) -dt_p = TaichiParameter("DT", dtype=ti.f32, mode="const", value=DT_VAL, pool=pool, solo=True) -seed_a_p = TaichiParameter("SEED_A", dtype=ti.f32, mode="const", value=12.9898, pool=pool, solo=True) -seed_b_p = TaichiParameter("SEED_B", dtype=ti.f32, mode="const", value=78.233, pool=pool, solo=True) - -# non-solo const: still fixed at compile time, but read through .get(0) like -# any other mode, so a template can be written without knowing it is const +# 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 @@ -99,8 +100,8 @@ # --------------------------------------------------------------------------- # bags # --------------------------------------------------------------------------- -# nested: grid.n is a solo const read bare, grid.dx a non-solo const read -# through .get(0) - members resolve on their own type, not the bag's +# 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 @@ -112,14 +113,14 @@ def clampi(i): - return min(max(i, 0), N - 1) + 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 non-solo const out of a bound bag + # 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) @@ -134,7 +135,7 @@ def laplacian(field_, i, j): 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 + j) + return UPLIFT.get(i * N.get(0) + j) uplift_at_fn = TaichiHelperBuilder().bind("UPLIFT", uplift_p).bind("N", n_p).ingest(uplift_at) @@ -147,7 +148,7 @@ def uplift_at(i, j): 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 + ti.cast(j, ti.f32) * SEED_B + 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 @@ -158,14 +159,15 @@ def init_topo_template(z: ti.template()): # routine kernels # --------------------------------------------------------------------------- -# mixed bag: a scalar Parameter, a helper and a solo const under one name. -# hill.d is a device accessor, hill.lap a specialized func, hill.dt a literal. +# 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 * hill.d.get(0) * hill.lap(z_in, i, j) + 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) @@ -173,9 +175,9 @@ def diffuse_template(z_out: ti.template(), z_in: ti.template()): def uplift_template(z: ti.template()): for i, j in z: - z[i, j] += up(i, j) * DT + 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 - 1: + if j == 0 or j == grid.n.get(0) - 1: z[i, j] = SEA.get(0) diff --git a/examples/core/shallow_water/shallow_water_cupy.py b/examples/core/shallow_water/shallow_water_cupy.py index 71ff871..0d32adf 100644 --- a/examples/core/shallow_water/shallow_water_cupy.py +++ b/examples/core/shallow_water/shallow_water_cupy.py @@ -71,9 +71,9 @@ # parameters # --------------------------------------------------------------------------- # Structural constants -> #define, used bare in the CUDA source. -n_p = CupyParameter("N", dtype=np.int32, mode="const", value=GRID_N, pool=pool, solo=True) -rest_depth_p = CupyParameter("REST_DEPTH", dtype=np.float32, mode="const", value=REST_DEPTH_VAL, pool=pool, solo=True) -drop_r_p = CupyParameter("DROP_R", dtype=np.int32, mode="const", value=DROP_R_VAL, pool=pool, solo=True) +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)$. diff --git a/examples/core/shallow_water/shallow_water_quadrants.py b/examples/core/shallow_water/shallow_water_quadrants.py index 49a7214..077c541 100644 --- a/examples/core/shallow_water/shallow_water_quadrants.py +++ b/examples/core/shallow_water/shallow_water_quadrants.py @@ -24,8 +24,9 @@ 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 solo=True: compile-time -literals used bare, no .get(). +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) """ @@ -77,10 +78,11 @@ # --------------------------------------------------------------------------- # parameters # --------------------------------------------------------------------------- -# Structural constants: solo=True -> read bare as compile-time literals. -n_p = QuadrantsParameter("N", dtype=qd.i32, mode="const", value=GRID_N, pool=pool, solo=True) -rest_depth_p = QuadrantsParameter("REST_DEPTH", dtype=qd.f32, mode="const", value=REST_DEPTH_VAL, pool=pool, solo=True) -drop_r_p = QuadrantsParameter("DROP_R", dtype=qd.i32, mode="const", value=DROP_R_VAL, pool=pool, solo=True) +# 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. @@ -102,7 +104,7 @@ def clamp(i): - return min(max(i, 0), N - 1) + return min(max(i, 0), N.get(0) - 1) clamp_fn = QuadrantsHelperBuilder().bind("N", n_p).ingest(clamp) @@ -127,7 +129,7 @@ def face_depth(up, down, vel): def init_height_template(h: qd.Tensor): for i, j in h: - h[i, j] = REST_DEPTH + h[i, j] = REST_DEPTH.get(0) init_height_kernel = QuadrantsKernelBuilder().bind("REST_DEPTH", rest_depth_p).ingest(init_height_template).compile() @@ -137,7 +139,7 @@ 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 * DROP_R: + if dxr * dxr + dyr * dyr <= DROP_R.get(0) * DROP_R.get(0): h[i, j] += drop.amp.get(0) @@ -174,11 +176,11 @@ def update_height_template(h_out: qd.Tensor, h_in: qd.Tensor, u: qd.Tensor, v: q # 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 - 1: + if i < N.get(0) - 1: ue = u[i + 1, j] vs = v[i, j] vn = 0.0 - if j < N - 1: + if j < N.get(0) - 1: vn = v[i, j + 1] ip = ops.clamp(i + 1) diff --git a/examples/core/shallow_water/shallow_water_taichi.py b/examples/core/shallow_water/shallow_water_taichi.py index 16a1197..36877f4 100644 --- a/examples/core/shallow_water/shallow_water_taichi.py +++ b/examples/core/shallow_water/shallow_water_taichi.py @@ -23,8 +23,9 @@ 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 solo=True: compile-time -literals used bare, no .get(). +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) """ @@ -76,10 +77,11 @@ # --------------------------------------------------------------------------- # parameters # --------------------------------------------------------------------------- -# Structural constants: solo=True -> read bare as compile-time literals. -n_p = TaichiParameter("N", dtype=ti.i32, mode="const", value=GRID_N, pool=pool, solo=True) -rest_depth_p = TaichiParameter("REST_DEPTH", dtype=ti.f32, mode="const", value=REST_DEPTH_VAL, pool=pool, solo=True) -drop_r_p = TaichiParameter("DROP_R", dtype=ti.i32, mode="const", value=DROP_R_VAL, pool=pool, solo=True) +# 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. @@ -101,7 +103,7 @@ def clamp(i): - return min(max(i, 0), N - 1) + return min(max(i, 0), N.get(0) - 1) clamp_fn = TaichiHelperBuilder().bind("N", n_p).ingest(clamp) @@ -126,7 +128,7 @@ def face_depth(up, down, vel): def init_height_template(h: ti.template()): for i, j in h: - h[i, j] = REST_DEPTH + h[i, j] = REST_DEPTH.get(0) init_height_kernel = TaichiKernelBuilder().bind("REST_DEPTH", rest_depth_p).ingest(init_height_template).compile() @@ -136,7 +138,7 @@ 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 * DROP_R: + if dxr * dxr + dyr * dyr <= DROP_R.get(0) * DROP_R.get(0): h[i, j] += drop.amp.get(0) @@ -173,11 +175,11 @@ def update_height_template(h_out: ti.template(), h_in: ti.template(), u: ti.temp # 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 - 1: + if i < N.get(0) - 1: ue = u[i + 1, j] vs = v[i, j] vn = 0.0 - if j < N - 1: + if j < N.get(0) - 1: vn = v[i, j + 1] ip = ops.clamp(i + 1) diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 0ca81fa..9e3b4b2 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -29,6 +29,7 @@ import numpy as np from .base import ( + MODES, Bag, HelperBuilder, Kernel, @@ -36,7 +37,6 @@ Parameter, _SpecializedHelper, _SpecializeCtx, - attach_meta, capture_template_meta, filter_bindings, resolve_binding, @@ -107,30 +107,24 @@ class ClosureBackendParameter(Parameter): Author: B.G (07/2026) """ - SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) _backend: ClassVar[Any] - def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): + 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. solo=True, available on const only, lets the parameter be - read bare in a template body - written `p` rather than `p.get(i)` - - because it resolves to a compile-time literal. + scalar and field take pooled storage straight away; const stays a + plain python value. Author: B.G (07/2026) """ - if mode not in self.SUPPORTED_MODES: - raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") - if solo and mode != "const": - raise ValueError(f"{name}: solo access is const-only, got mode {mode!r}") + 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.solo = solo self._pool = pool self._const_value: Any = None self._handle = None @@ -365,9 +359,7 @@ def _specialize(self, ctx: _SpecializeCtx) -> ClosureHelper: Author: B.G (07/2026) """ specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) - fn = ClosureHelper(specialised.__name__, self._backend.func(specialised)) - attach_meta(fn, self._template, self._bindings) - return fn + return ClosureHelper(specialised.__name__, self._backend.func(specialised)) class ClosureKernelBuilder(KernelBuilder): @@ -391,9 +383,7 @@ def compile(self) -> ClosureKernel: """ ctx = _SpecializeCtx() specialised = specialize_closure(self._template, filter_bindings(self._template, self._bindings), ctx) - krn = ClosureKernel(specialised.__name__, self._backend.kernel(specialised)) - attach_meta(krn, self._template, self._bindings) - return krn + return ClosureKernel(specialised.__name__, self._backend.kernel(specialised)) class _RenameNames(ast.NodeTransformer): @@ -607,7 +597,6 @@ def _fuse_group(self, group: list, group_index: int, dump_source: "str | None"): assigned_by: dict[str, str] = {} module_globals: dict[str, Any] = {} resolved_globals: dict[str, Any] = {} - raw_bindings: dict[str, Any] = {} for step_index, step in enumerate(group): kernel_builder = step.kernel_builder @@ -648,7 +637,6 @@ def _fuse_group(self, group: list, group_index: int, dump_source: "str | None"): body_stmts.extend(renamed_body) filtered = filter_bindings(template, kernel_builder.bindings) - raw_bindings.update(filtered) module_globals.update(dict(getattr(template, "__globals__", {}))) resolved_globals.update({name: resolve_binding(value, ctx) for name, value in filtered.items()}) @@ -714,5 +702,4 @@ def _fuse_group(self, group: list, group_index: int, dump_source: "str | None"): fused_fn = exec_globals[func_name] krn = ClosureKernel(func_name, backend.kernel(fused_fn)) - attach_meta(krn, fused_fn, raw_bindings) 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..e0c4174 --- /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 .base 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/base.py b/pyfastflow/experimental/core/context/base.py index 9abfcd4..93fff62 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -129,9 +129,12 @@ import warnings from abc import ABC, abstractmethod from functools import lru_cache -from typing import Any, ClassVar +from typing import Any -from ..pool.base import DataHandle, new_uid +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): @@ -140,9 +143,7 @@ class Parameter(ABC): `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 must offer all three (REQUIRED_MODES). A backend may widen - SUPPORTED_MODES with further storage kinds; the check runs when the - subclass is defined, so an incomplete backend fails at import. + 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 @@ -152,18 +153,8 @@ class Parameter(ABC): Author: B.G (07/2026) """ - REQUIRED_MODES: ClassVar[frozenset[str]] = frozenset({"const", "scalar", "field"}) - SUPPORTED_MODES: ClassVar[frozenset[str]] = REQUIRED_MODES - - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - missing = Parameter.REQUIRED_MODES - cls.SUPPORTED_MODES - if missing: - raise TypeError(f"{cls.__name__} must support modes {sorted(missing)}") - name: str dtype: Any - solo: bool = False def __init__(self): """ @@ -270,22 +261,13 @@ class Specializable(ABC): launchable Kernel, or the internal object a HelperBuilder specializes into as part of an enclosing kernel's compile(). - Beyond the call contract, each instance keeps a read-only snapshot of the - raw material it was made from - template, source text, AST, and the real - bindings dict as it stood at this object's own compile time. Nothing in - this module reads those back to drive a later compile; they are here so a - higher layer (kernel fusion, a graph compiler) can work from the originals - instead of re-deriving them. The builder that produced this object stays - authoritative for anything that needs to change - see CompileBuilder. + The builder that produced this object stays authoritative for its recipe + - template and bindings - see CompileBuilder. Author: B.G (07/2026) """ name: str - _template: Any = None - _source: str | None = None - _ast: ast.AST | None = None - _dependencies: dict[str, Any] | None = None def __init__(self): """ @@ -305,48 +287,6 @@ def uid(self) -> int: """ return self._uid - @property - def template(self): - """ - The template object this was compiled from - a python def for - Taichi/Quadrants, a CUDA source string for cupy. Read-only: this is a - snapshot for introspection, not a handle to recompile from. Change and - recompile through the builder instead. - - Author: B.G (07/2026) - """ - return self._template - - @property - def source(self) -> str | None: - """ - The template's source text, captured at compile time. See `template`. - - Author: B.G (07/2026) - """ - return self._source - - @property - def ast(self) -> "ast.AST | None": - """ - The template's parsed AST, captured at compile time. None for a - template with no recoverable source (e.g. cupy's CUDA text). See - `template`. - - Author: B.G (07/2026) - """ - return self._ast - - @property - def bindings(self) -> dict[str, Any]: - """ - The real bound objects - not type names - as they stood at this - object's own compile time. Read-only snapshot; see `template`. - - Author: B.G (07/2026) - """ - return self._dependencies - @property @abstractmethod def compiled(self): @@ -474,8 +414,8 @@ def resolve_binding(value, ctx: "_SpecializeCtx"): """ Turn a bound object into what a template body should see in its place. - Parameter a bare literal when solo, otherwise device_view() - the - carrier of .get / .set_node for in-kernel access. + 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. @@ -492,9 +432,6 @@ def resolve_binding(value, ctx: "_SpecializeCtx"): Author: B.G (07/2026) """ if isinstance(value, Parameter): - if getattr(value, "solo", False): - resolved = value.get() - return resolved.data if isinstance(resolved, DataHandle) else resolved return value.device_view() if isinstance(value, HelperBuilder): return ctx.specialize(value).compiled @@ -511,10 +448,10 @@ 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 twice - once to filter bindings, once - for attach_meta - 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. + 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 @@ -578,20 +515,6 @@ def filter_bindings(template, bindings: dict[str, Any]) -> dict[str, Any]: return filtered -def attach_meta(obj: Specializable, template, bindings: dict[str, Any]) -> None: - """ - Record on a freshly built Specializable what it was made from: the - template itself, its source, its AST, and the real bound objects (not - type names, so a caller can inspect and reuse what was actually bound). - See Specializable. - - Author: B.G (07/2026) - """ - obj._template = template - obj._source, obj._ast = capture_template_meta(template) - obj._dependencies = dict(bindings) - - class CompileBuilder(ABC): """ Collects dependencies and a template, and compiles them into one @@ -648,9 +571,8 @@ 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 copy of a frozen snapshot; a compiled - object's own `.bindings` is the frozen copy, taken at its compile - time. + the builder's own dict, not a snapshot - the builder stays + authoritative for its recipe (see the class docstring). Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index ff535fe..c877be4 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -72,7 +72,7 @@ import cupy as cp import numpy as np -from .base import HelperBuilder, Kernel, KernelBuilder, Parameter, _SpecializedHelper, _SpecializeCtx, attach_meta +from .base import MODES, HelperBuilder, Kernel, KernelBuilder, Parameter, _SpecializedHelper, _SpecializeCtx from .base import Bag from .routine import Routine, RoutineBuilder, _CompiledStep @@ -420,27 +420,21 @@ class CupyParameter(Parameter): Author: B.G (07/2026) """ - SUPPORTED_MODES = frozenset({"const", "scalar", "field"}) - - def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | None = None, solo: bool = False): + 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. solo=True (const only) lets the parameter be read bare in a - template body - it becomes a #define rather than a span expansion. + value, read bare in a template body as a #define. Author: B.G (07/2026) """ - if mode not in self.SUPPORTED_MODES: - raise ValueError(f"{name}: mode must be one of {sorted(self.SUPPORTED_MODES)}, got {mode!r}") - if solo and mode != "const": - raise ValueError(f"{name}: solo access is const-only, got mode {mode!r}") + 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.solo = solo self._pool = pool self._const_value: Any = None self._handle = None @@ -518,8 +512,6 @@ class CupyHelper(_SpecializedHelper): def __init__(self, name: str, source: str): super().__init__() self.name = name - # note: distinct from Specializable._source (the raw template, set by - # attach_meta) - this is the spliced __device__ source `.compiled` serves. self._compiled_source = source @property @@ -643,9 +635,7 @@ def _specialize(self, ctx: _SpecializeCtx) -> CupyHelper: if device_srcs is None: device_srcs = ctx.cupy_device_srcs = {} device_srcs.setdefault(name, source) - fn = CupyHelper(name, source) - attach_meta(fn, template, self._bindings) - return fn + return CupyHelper(name, source) class CupyKernelBuilder(KernelBuilder): @@ -721,7 +711,6 @@ def compile(self) -> CupyKernel: raw = module.get_function(name) krn = CupyKernel(name, raw, module) krn._final_source = source - attach_meta(krn, template, self._bindings) return krn diff --git a/pyfastflow/experimental/grid/__init__.py b/pyfastflow/experimental/grid/__init__.py index fb863e1..af77982 100644 --- a/pyfastflow/experimental/grid/__init__.py +++ b/pyfastflow/experimental/grid/__init__.py @@ -10,7 +10,7 @@ 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", non-solo, for all three - see the note below. + "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. @@ -20,25 +20,12 @@ exists in the bag when its feature is off, so a caller that never asked for nodata/mask-outlet never sees them. -Note on nx/ny solo=True: earlier drafts of this design set nx/ny to a -default of const+solo=True (read bare, as a compile-time literal with no -`.get()`), matching how n_neighbours is read. That is incompatible with the -"read every value param mode-agnostically via `.get()`" requirement the -private blocks are built around - a solo Parameter resolves to a bare python -literal in device code (see base.py's resolve_binding), which has no `.get` -method, so `NX.get(0)` inside a shared geometry block would fail to trace -the moment nx defaulted to solo. To keep one geometry block working -whatever mode nx/ny/dx are in - the actual point of the mode-overridable -design - nx/ny/dx are built here as solo=False by default (same as dx's -own, unambiguous spec: "non-solo, read via .get(0)"). n_neighbours, which -has no override and is genuinely only ever a compile-time literal, keeps -solo=True. - Author: B.G (07/2026) """ import numpy as np +from ..core.context.backends import backend_classes from ..core.context.base import Bag _TOPOLOGIES = {"D4": 4, "D8": 8} @@ -46,38 +33,21 @@ _OUTLETS = frozenset({"edge", "mask"}) -def _backend_classes(backend: str): +def _blocks_for(backend: str): """ - (backend_module_or_None, ParameterCls, HelperBuilderCls, blocks_module, - dtypes) for one backend name. cupy's backend_module is None - its blocks - call plain C, never a bound backend module - see _cupy_blocks.py. + 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 == "taichi": - import taichi as ti - - from ..core.context.taichi_backend import TaichiHelperBuilder, TaichiParameter - from . import _closure_blocks as blocks - - return ti, TaichiParameter, TaichiHelperBuilder, blocks, {"i32": ti.i32, "f32": ti.f32, "u8": ti.u8} - if backend == "quadrants": - import quadrants as qd - - from ..core.context.quadrants_backend import QuadrantsHelperBuilder, QuadrantsParameter + if backend in ("taichi", "quadrants"): from . import _closure_blocks as blocks - - return qd, QuadrantsParameter, QuadrantsHelperBuilder, blocks, {"i32": qd.i32, "f32": qd.f32, "u8": qd.u8} - if backend == "cupy": - from ..core.context.cupy_backend import CupyHelperBuilder, CupyParameter + elif backend == "cupy": from . import _cupy_blocks as blocks - - return None, CupyParameter, CupyHelperBuilder, blocks, { - "i32": np.int32, - "f32": np.float32, - "u8": np.uint8, - } - raise ValueError(f"make_grid: unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + else: + raise ValueError(f"make_grid: unknown backend {backend!r}, expected 'taichi', 'quadrants' or 'cupy'") + return blocks def make_grid( @@ -129,7 +99,8 @@ def make_grid( 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, blocks, dtypes = _backend_classes(backend) + 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) @@ -148,7 +119,7 @@ def make_grid( 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, solo=True + "GRID_NNEIGHBOURS", dtype=dtypes["i32"], mode="const", value=_TOPOLOGIES[topology], pool=pool ) nodata_mask_p = None diff --git a/pyfastflow/experimental/grid/_closure_blocks.py b/pyfastflow/experimental/grid/_closure_blocks.py index 783c139..854ea01 100644 --- a/pyfastflow/experimental/grid/_closure_blocks.py +++ b/pyfastflow/experimental/grid/_closure_blocks.py @@ -28,6 +28,10 @@ Author: B.G (07/2026) """ +import functools + +from ..core.context.backends import make_helper + # --------------------------------------------------------------------------- # geometry # --------------------------------------------------------------------------- @@ -437,11 +441,7 @@ def build_helpers( sqrt2 = math.sqrt(2.0) d8 = topology == "D8" - def mk(template, **binds): - b = HelperCls().ingest(template) - for name, obj in binds.items(): - b.bind(name, obj) - return b + mk = functools.partial(make_helper, HelperCls) row = mk(_row_tmpl, NX=nx_p) col = mk(_col_tmpl, NX=nx_p) diff --git a/pyfastflow/experimental/grid/_cupy_blocks.py b/pyfastflow/experimental/grid/_cupy_blocks.py index 2d6646a..d34cd71 100644 --- a/pyfastflow/experimental/grid/_cupy_blocks.py +++ b/pyfastflow/experimental/grid/_cupy_blocks.py @@ -22,8 +22,10 @@ Author: B.G (07/2026) """ +import functools import math +from ..core.context.backends import make_helper from ..core.pool.base import new_uid @@ -60,11 +62,7 @@ def build_helpers( d8 = topology == "D8" sqrt2 = math.sqrt(2.0) - def mk(template, **binds): - b = HelperCls().ingest(template) - for name, obj in binds.items(): - b.bind(name, obj) - return b + mk = functools.partial(make_helper, HelperCls) row = mk( f"__device__ int {t}_row(int i) {{ return i / $NX.get(0)$; }}", diff --git a/pyfastflow/experimental/noise/__init__.py b/pyfastflow/experimental/noise/__init__.py new file mode 100644 index 0000000..23d3268 --- /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.base for Parameter/HelperBuilder/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.base 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} From 10eb0bb135fd4beb5b0870d56d86702b5fce53b2 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Wed, 29 Jul 2026 10:48:37 +0200 Subject: [PATCH 29/31] Clean and restructure the core context base.py (1000+ LoC for multiple concepts) --- .../heat_diffusion/heat_diffusion_cupy.py | 2 +- .../heat_diffusion_quadrants.py | 2 +- .../heat_diffusion_routine_cupy.py | 2 +- .../heat_diffusion_routine_quadrants.py | 2 +- .../heat_diffusion_routine_taichi.py | 2 +- .../heat_diffusion/heat_diffusion_taichi.py | 2 +- examples/core/lem/lem_routine_cupy.py | 2 +- examples/core/lem/lem_routine_quadrants.py | 2 +- examples/core/lem/lem_routine_taichi.py | 2 +- .../core/shallow_water/shallow_water_cupy.py | 2 +- .../shallow_water/shallow_water_quadrants.py | 2 +- .../shallow_water/shallow_water_taichi.py | 2 +- .../core/context/_closure_backend.py | 8 +- .../experimental/core/context/backends.py | 2 +- pyfastflow/experimental/core/context/bag.py | 334 ++++++++ pyfastflow/experimental/core/context/base.py | 806 +----------------- .../experimental/core/context/compile.py | 515 +++++++++++ .../experimental/core/context/cupy_backend.py | 7 +- .../experimental/core/context/routine.py | 9 +- pyfastflow/experimental/grid/__init__.py | 4 +- .../experimental/grid/_closure_blocks.py | 2 +- pyfastflow/experimental/noise/__init__.py | 4 +- 22 files changed, 890 insertions(+), 825 deletions(-) create mode 100644 pyfastflow/experimental/core/context/bag.py create mode 100644 pyfastflow/experimental/core/context/compile.py diff --git a/examples/core/heat_diffusion/heat_diffusion_cupy.py b/examples/core/heat_diffusion/heat_diffusion_cupy.py index f3a542f..02fdf25 100644 --- a/examples/core/heat_diffusion/heat_diffusion_cupy.py +++ b/examples/core/heat_diffusion/heat_diffusion_cupy.py @@ -37,7 +37,7 @@ import matplotlib.pyplot as plt import numpy as np -from pyfastflow.experimental.core.context.base import Bag +from pyfastflow.experimental.core.context.bag import Bag from pyfastflow.experimental.core.context.cupy_backend import ( CupyHelperBuilder, CupyKernelBuilder, diff --git a/examples/core/heat_diffusion/heat_diffusion_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_quadrants.py index 96bce93..e9364b6 100644 --- a/examples/core/heat_diffusion/heat_diffusion_quadrants.py +++ b/examples/core/heat_diffusion/heat_diffusion_quadrants.py @@ -50,7 +50,7 @@ import numpy as np import quadrants as qd -from pyfastflow.experimental.core.context.base import Bag +from pyfastflow.experimental.core.context.bag import Bag from pyfastflow.experimental.core.context.quadrants_backend import ( QuadrantsHelperBuilder, QuadrantsKernelBuilder, diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py index 93cb55a..6dd4425 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py @@ -47,7 +47,7 @@ import matplotlib.pyplot as plt import numpy as np -from pyfastflow.experimental.core.context.base import Bag, merge +from pyfastflow.experimental.core.context.bag import Bag, merge from pyfastflow.experimental.core.context.cupy_backend import ( CupyHelperBuilder, CupyKernelBuilder, diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py index b314fa4..4d83720 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py @@ -42,7 +42,7 @@ import numpy as np import quadrants as qd -from pyfastflow.experimental.core.context.base import Bag, merge +from pyfastflow.experimental.core.context.bag import Bag, merge from pyfastflow.experimental.core.context.quadrants_backend import ( QuadrantsHelperBuilder, QuadrantsKernelBuilder, diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py index 8268770..1b64a1c 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py @@ -42,7 +42,7 @@ import numpy as np import taichi as ti -from pyfastflow.experimental.core.context.base import Bag, merge +from pyfastflow.experimental.core.context.bag import Bag, merge from pyfastflow.experimental.core.context.taichi_backend import ( TaichiHelperBuilder, TaichiKernelBuilder, diff --git a/examples/core/heat_diffusion/heat_diffusion_taichi.py b/examples/core/heat_diffusion/heat_diffusion_taichi.py index e811cd3..718332d 100644 --- a/examples/core/heat_diffusion/heat_diffusion_taichi.py +++ b/examples/core/heat_diffusion/heat_diffusion_taichi.py @@ -50,7 +50,7 @@ import numpy as np import taichi as ti -from pyfastflow.experimental.core.context.base import Bag +from pyfastflow.experimental.core.context.bag import Bag from pyfastflow.experimental.core.context.taichi_backend import ( TaichiHelperBuilder, TaichiKernelBuilder, diff --git a/examples/core/lem/lem_routine_cupy.py b/examples/core/lem/lem_routine_cupy.py index 21f0b95..cf5b38c 100644 --- a/examples/core/lem/lem_routine_cupy.py +++ b/examples/core/lem/lem_routine_cupy.py @@ -55,7 +55,7 @@ import matplotlib.pyplot as plt import numpy as np -from pyfastflow.experimental.core.context.base import Bag, merge +from pyfastflow.experimental.core.context.bag import Bag, merge from pyfastflow.experimental.core.context.cupy_backend import ( CupyHelperBuilder, CupyKernelBuilder, diff --git a/examples/core/lem/lem_routine_quadrants.py b/examples/core/lem/lem_routine_quadrants.py index 86d3371..906df3b 100644 --- a/examples/core/lem/lem_routine_quadrants.py +++ b/examples/core/lem/lem_routine_quadrants.py @@ -45,7 +45,7 @@ import numpy as np import quadrants as qd -from pyfastflow.experimental.core.context.base import Bag, merge +from pyfastflow.experimental.core.context.bag import Bag, merge from pyfastflow.experimental.core.context.quadrants_backend import ( QuadrantsHelperBuilder, QuadrantsKernelBuilder, diff --git a/examples/core/lem/lem_routine_taichi.py b/examples/core/lem/lem_routine_taichi.py index 640ca08..223e3df 100644 --- a/examples/core/lem/lem_routine_taichi.py +++ b/examples/core/lem/lem_routine_taichi.py @@ -45,7 +45,7 @@ import numpy as np import taichi as ti -from pyfastflow.experimental.core.context.base import Bag, merge +from pyfastflow.experimental.core.context.bag import Bag, merge from pyfastflow.experimental.core.context.taichi_backend import ( TaichiHelperBuilder, TaichiKernelBuilder, diff --git a/examples/core/shallow_water/shallow_water_cupy.py b/examples/core/shallow_water/shallow_water_cupy.py index 0d32adf..59e3986 100644 --- a/examples/core/shallow_water/shallow_water_cupy.py +++ b/examples/core/shallow_water/shallow_water_cupy.py @@ -31,7 +31,7 @@ import matplotlib.pyplot as plt import numpy as np -from pyfastflow.experimental.core.context.base import Bag +from pyfastflow.experimental.core.context.bag import Bag from pyfastflow.experimental.core.context.cupy_backend import ( CupyHelperBuilder, CupyKernelBuilder, diff --git a/examples/core/shallow_water/shallow_water_quadrants.py b/examples/core/shallow_water/shallow_water_quadrants.py index 077c541..142b7ab 100644 --- a/examples/core/shallow_water/shallow_water_quadrants.py +++ b/examples/core/shallow_water/shallow_water_quadrants.py @@ -38,7 +38,7 @@ import numpy as np import quadrants as qd -from pyfastflow.experimental.core.context.base import Bag +from pyfastflow.experimental.core.context.bag import Bag from pyfastflow.experimental.core.context.quadrants_backend import ( QuadrantsHelperBuilder, QuadrantsKernelBuilder, diff --git a/examples/core/shallow_water/shallow_water_taichi.py b/examples/core/shallow_water/shallow_water_taichi.py index 36877f4..43eb686 100644 --- a/examples/core/shallow_water/shallow_water_taichi.py +++ b/examples/core/shallow_water/shallow_water_taichi.py @@ -37,7 +37,7 @@ import numpy as np import taichi as ti -from pyfastflow.experimental.core.context.base import Bag +from pyfastflow.experimental.core.context.bag import Bag from pyfastflow.experimental.core.context.taichi_backend import ( TaichiHelperBuilder, TaichiKernelBuilder, diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 9e3b4b2..2f24198 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -28,13 +28,11 @@ import numpy as np -from .base import ( - MODES, - Bag, +from .base import MODES, Parameter +from .compile import ( HelperBuilder, Kernel, KernelBuilder, - Parameter, _SpecializedHelper, _SpecializeCtx, capture_template_meta, @@ -312,7 +310,7 @@ class ClosureKernel(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 base.py. + the body through globals instead. See compile.py. Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/core/context/backends.py b/pyfastflow/experimental/core/context/backends.py index e0c4174..ae932d2 100644 --- a/pyfastflow/experimental/core/context/backends.py +++ b/pyfastflow/experimental/core/context/backends.py @@ -23,7 +23,7 @@ import numpy as np -from .base import HelperBuilder +from .compile import HelperBuilder def backend_classes(backend: str): diff --git a/pyfastflow/experimental/core/context/bag.py b/pyfastflow/experimental/core/context/bag.py new file mode 100644 index 0000000..b1c0288 --- /dev/null +++ b/pyfastflow/experimental/core/context/bag.py @@ -0,0 +1,334 @@ +""" +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 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 a Parameter's + mode is changed: mode is fixed at construction (see Parameter.mode), so + changing it means building a new Parameter and using replace() to swap it + into the bag in place of 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/base.py b/pyfastflow/experimental/core/context/base.py index 93fff62..cdf8e46 100644 --- a/pyfastflow/experimental/core/context/base.py +++ b/pyfastflow/experimental/core/context/base.py @@ -121,14 +121,21 @@ None of this is enforced at runtime. +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) """ -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 @@ -255,794 +262,3 @@ def destroy(self) -> None: ... -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) - """ - - -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 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 a Parameter's - mode is changed: mode is fixed at construction (see Parameter.mode), so - changing it means building a new Parameter and using replace() to swap it - into the bag in place of 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..092e0e1 --- /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 (base.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 base.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 .bag import Bag +from .base import Parameter +from ..pool.base import new_uid + + +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 index c877be4..1f6afbf 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -1,7 +1,7 @@ """ 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 base.py, HelperBuilder). +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 @@ -72,8 +72,9 @@ import cupy as cp import numpy as np -from .base import MODES, HelperBuilder, Kernel, KernelBuilder, Parameter, _SpecializedHelper, _SpecializeCtx -from .base import Bag +from .base import MODES, Parameter +from .compile import HelperBuilder, Kernel, KernelBuilder, _SpecializedHelper, _SpecializeCtx +from .bag import Bag from .routine import Routine, RoutineBuilder, _CompiledStep _KERNEL_NAME_RE = re.compile(r"__global__\s+void\s+(\w+)\s*\(") diff --git a/pyfastflow/experimental/core/context/routine.py b/pyfastflow/experimental/core/context/routine.py index 6307e9c..19bc84d 100644 --- a/pyfastflow/experimental/core/context/routine.py +++ b/pyfastflow/experimental/core/context/routine.py @@ -17,9 +17,9 @@ 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 base.py, CompileBuilder.rebind) against the one bag the +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 (base.py) is run across +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 @@ -111,7 +111,8 @@ from abc import ABC, abstractmethod from typing import Any -from .base import Bag, CompileBuilder, check_handles +from .bag import Bag, check_handles +from .compile import CompileBuilder _C_FUNC_NAME_RE = re.compile(r"(?:__global__|__device__)\s+[\w:\*&]*\s*(\w+)\s*\(") @@ -385,7 +386,7 @@ def _validate(self) -> None: Everything a compile needs checked before any step is actually compiled, fused or not. - In order: check_handles (base.py) runs across every step's own + 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 diff --git a/pyfastflow/experimental/grid/__init__.py b/pyfastflow/experimental/grid/__init__.py index af77982..b497a05 100644 --- a/pyfastflow/experimental/grid/__init__.py +++ b/pyfastflow/experimental/grid/__init__.py @@ -1,6 +1,6 @@ """ make_grid: the GridContext-equivalent Bag factory, built on the -backend-agnostic core (see ..core.context.base for Parameter/HelperBuilder/Bag). +backend-agnostic core (see ..core.context: base.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/base.py's module docstring). make_grid just builds a Bag once: a @@ -26,7 +26,7 @@ import numpy as np from ..core.context.backends import backend_classes -from ..core.context.base import Bag +from ..core.context.bag import Bag _TOPOLOGIES = {"D4": 4, "D8": 8} _BOUNDARIES = frozenset({"normal", "periodic_EW", "periodic_NS"}) diff --git a/pyfastflow/experimental/grid/_closure_blocks.py b/pyfastflow/experimental/grid/_closure_blocks.py index 854ea01..e7e2624 100644 --- a/pyfastflow/experimental/grid/_closure_blocks.py +++ b/pyfastflow/experimental/grid/_closure_blocks.py @@ -16,7 +16,7 @@ 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 base.py, +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 diff --git a/pyfastflow/experimental/noise/__init__.py b/pyfastflow/experimental/noise/__init__.py index 23d3268..6632908 100644 --- a/pyfastflow/experimental/noise/__init__.py +++ b/pyfastflow/experimental/noise/__init__.py @@ -1,6 +1,6 @@ """ make_noise: the NoiseContext-equivalent Bag factory, built on the -backend-agnostic core (see ..core.context.base for Parameter/HelperBuilder/Bag) +backend-agnostic core (see ..core.context: base.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 @@ -51,7 +51,7 @@ def init_template(z: ti.template()): import numpy as np from ..core.context.backends import backend_classes -from ..core.context.base import Bag +from ..core.context.bag import Bag _KINDS = frozenset({"white", "perlin"}) _MODES = ("const", "scalar") From bfb279c71a4a0de88372b99921f913a5ccd0cac3 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Wed, 29 Jul 2026 11:39:02 +0200 Subject: [PATCH 30/31] Clean imports --- examples/core/heat_diffusion/heat_diffusion_cupy.py | 2 +- examples/core/heat_diffusion/heat_diffusion_quadrants.py | 2 +- .../core/heat_diffusion/heat_diffusion_routine_cupy.py | 2 +- .../heat_diffusion/heat_diffusion_routine_quadrants.py | 2 +- .../core/heat_diffusion/heat_diffusion_routine_taichi.py | 2 +- examples/core/heat_diffusion/heat_diffusion_taichi.py | 2 +- examples/core/shallow_water/shallow_water_cupy.py | 2 +- examples/core/shallow_water/shallow_water_quadrants.py | 2 +- examples/core/shallow_water/shallow_water_taichi.py | 2 +- pyfastflow/experimental/core/context/_closure_backend.py | 4 ++-- pyfastflow/experimental/core/context/compile.py | 8 ++++---- pyfastflow/experimental/core/context/cupy_backend.py | 4 ++-- .../experimental/core/context/{base.py => parameter.py} | 0 pyfastflow/experimental/core/context/routine.py | 4 ++-- pyfastflow/experimental/core/context/taichi_backend.py | 2 +- pyfastflow/experimental/grid/__init__.py | 4 ++-- pyfastflow/experimental/grid/_closure_blocks.py | 2 +- pyfastflow/experimental/noise/__init__.py | 2 +- 18 files changed, 24 insertions(+), 24 deletions(-) rename pyfastflow/experimental/core/context/{base.py => parameter.py} (100%) diff --git a/examples/core/heat_diffusion/heat_diffusion_cupy.py b/examples/core/heat_diffusion/heat_diffusion_cupy.py index 02fdf25..fc8d06e 100644 --- a/examples/core/heat_diffusion/heat_diffusion_cupy.py +++ b/examples/core/heat_diffusion/heat_diffusion_cupy.py @@ -334,7 +334,7 @@ # 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 base.py, "Lifetime of a compiled object"). +# 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) diff --git a/examples/core/heat_diffusion/heat_diffusion_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_quadrants.py index e9364b6..0c97d39 100644 --- a/examples/core/heat_diffusion/heat_diffusion_quadrants.py +++ b/examples/core/heat_diffusion/heat_diffusion_quadrants.py @@ -328,7 +328,7 @@ def diffuse_template(T_out: qd.Tensor, T_in: qd.Tensor): # 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 base.py, "Lifetime of a compiled object"). +# 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) diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py index 6dd4425..be47f12 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_cupy.py @@ -226,7 +226,7 @@ # 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 base.py). +# routine's add_kernel() - compile() does not consume it (see compile.py). apply_source_builder = ( CupyKernelBuilder() .bind("N", n_p) diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py index 4d83720..0b5ef93 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_quadrants.py @@ -213,7 +213,7 @@ def apply_source_template(T: qd.Tensor): # 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 base.py). +# 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() diff --git a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py index 1b64a1c..2f7e8b5 100644 --- a/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py +++ b/examples/core/heat_diffusion/heat_diffusion_routine_taichi.py @@ -213,7 +213,7 @@ def apply_source_template(T: ti.template()): # 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 base.py). +# 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() diff --git a/examples/core/heat_diffusion/heat_diffusion_taichi.py b/examples/core/heat_diffusion/heat_diffusion_taichi.py index 718332d..4f24567 100644 --- a/examples/core/heat_diffusion/heat_diffusion_taichi.py +++ b/examples/core/heat_diffusion/heat_diffusion_taichi.py @@ -328,7 +328,7 @@ def diffuse_template(T_out: ti.template(), T_in: ti.template()): # 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 base.py, "Lifetime of a compiled object"). +# 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) diff --git a/examples/core/shallow_water/shallow_water_cupy.py b/examples/core/shallow_water/shallow_water_cupy.py index 59e3986..48e7f4c 100644 --- a/examples/core/shallow_water/shallow_water_cupy.py +++ b/examples/core/shallow_water/shallow_water_cupy.py @@ -300,7 +300,7 @@ # 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 base.py, "Lifetime of a compiled object"). +# 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): diff --git a/examples/core/shallow_water/shallow_water_quadrants.py b/examples/core/shallow_water/shallow_water_quadrants.py index 142b7ab..33af678 100644 --- a/examples/core/shallow_water/shallow_water_quadrants.py +++ b/examples/core/shallow_water/shallow_water_quadrants.py @@ -273,7 +273,7 @@ def update_height_template(h_out: qd.Tensor, h_in: qd.Tensor, u: qd.Tensor, v: q # 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 base.py, "Lifetime of a compiled object"). +# 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): diff --git a/examples/core/shallow_water/shallow_water_taichi.py b/examples/core/shallow_water/shallow_water_taichi.py index 43eb686..cda70a2 100644 --- a/examples/core/shallow_water/shallow_water_taichi.py +++ b/examples/core/shallow_water/shallow_water_taichi.py @@ -272,7 +272,7 @@ def update_height_template(h_out: ti.template(), h_in: ti.template(), u: ti.temp # 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 base.py, "Lifetime of a compiled object"). +# 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): diff --git a/pyfastflow/experimental/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 2f24198..2c6adc8 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -28,7 +28,6 @@ import numpy as np -from .base import MODES, Parameter from .compile import ( HelperBuilder, Kernel, @@ -39,6 +38,7 @@ filter_bindings, resolve_binding, ) +from .parameter import MODES, Parameter from .routine import Routine, RoutineBuilder, _CompiledStep, _template_label @@ -212,7 +212,7 @@ def device_view(self) -> ClosureParamDeviceView: set() on a const mode, which changes the literal baked into the getter, and destroy(), which releases the storage it reads. Both drop the view so the next caller rebuilds; neither reaches kernels compiled earlier - (see base.py, "Lifetime of a compiled object"). A scalar or field set() + (see parameter.py, "Lifetime of a compiled object"). A scalar or field set() needs no invalidation, writing through the very storage the view reads. Author: B.G (07/2026) diff --git a/pyfastflow/experimental/core/context/compile.py b/pyfastflow/experimental/core/context/compile.py index 092e0e1..475b940 100644 --- a/pyfastflow/experimental/core/context/compile.py +++ b/pyfastflow/experimental/core/context/compile.py @@ -11,10 +11,10 @@ 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 (base.py) and Bag (bag.py) but neither reads it back. +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 base.py's module docstring +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) @@ -27,9 +27,9 @@ from functools import lru_cache from typing import Any -from .bag import Bag -from .base import Parameter from ..pool.base import new_uid +from .bag import Bag, from_builder +from .parameter import Parameter class Specializable(ABC): diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index 1f6afbf..ea9b1cb 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -72,9 +72,9 @@ import cupy as cp import numpy as np -from .base import MODES, Parameter -from .compile import HelperBuilder, Kernel, KernelBuilder, _SpecializedHelper, _SpecializeCtx 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*\(") diff --git a/pyfastflow/experimental/core/context/base.py b/pyfastflow/experimental/core/context/parameter.py similarity index 100% rename from pyfastflow/experimental/core/context/base.py rename to pyfastflow/experimental/core/context/parameter.py diff --git a/pyfastflow/experimental/core/context/routine.py b/pyfastflow/experimental/core/context/routine.py index 19bc84d..3e86471 100644 --- a/pyfastflow/experimental/core/context/routine.py +++ b/pyfastflow/experimental/core/context/routine.py @@ -72,7 +72,7 @@ Contract: no set()/destroy() mid-routine ------------------------------------------ A compiled Routine holds the same kind of frozen snapshot a Kernel does (see -base.py, "Lifetime of a compiled object"): scalar/field Parameters are baked +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: @@ -190,7 +190,7 @@ class Routine: """ A compiled, ordered sequence of kernels sharing one bag, ready to launch. - Inert, like every other compiled object in this package (see base.py, + 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. diff --git a/pyfastflow/experimental/core/context/taichi_backend.py b/pyfastflow/experimental/core/context/taichi_backend.py index 341fadd..ce2aa10 100644 --- a/pyfastflow/experimental/core/context/taichi_backend.py +++ b/pyfastflow/experimental/core/context/taichi_backend.py @@ -7,7 +7,7 @@ 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 base.py. +ti.template(). Bound Parameters and helpers are not arguments - see parameter.py. Author: B.G (07/2026) """ diff --git a/pyfastflow/experimental/grid/__init__.py b/pyfastflow/experimental/grid/__init__.py index b497a05..fa3d08e 100644 --- a/pyfastflow/experimental/grid/__init__.py +++ b/pyfastflow/experimental/grid/__init__.py @@ -1,9 +1,9 @@ """ make_grid: the GridContext-equivalent Bag factory, built on the -backend-agnostic core (see ..core.context: base.py for Parameter, compile.py for HelperBuilder, bag.py for Bag). +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/base.py's module docstring). make_grid just builds a Bag once: a +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. diff --git a/pyfastflow/experimental/grid/_closure_blocks.py b/pyfastflow/experimental/grid/_closure_blocks.py index e7e2624..a4ba189 100644 --- a/pyfastflow/experimental/grid/_closure_blocks.py +++ b/pyfastflow/experimental/grid/_closure_blocks.py @@ -20,7 +20,7 @@ _SpecializeCtx). nx/ny/dx are read exclusively through `.get(...)`, uniformly across whatever -mode they end up in (const, scalar, field) - see base.py, "Reading a +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. diff --git a/pyfastflow/experimental/noise/__init__.py b/pyfastflow/experimental/noise/__init__.py index 6632908..4b7ef51 100644 --- a/pyfastflow/experimental/noise/__init__.py +++ b/pyfastflow/experimental/noise/__init__.py @@ -1,6 +1,6 @@ """ make_noise: the NoiseContext-equivalent Bag factory, built on the -backend-agnostic core (see ..core.context: base.py for Parameter, compile.py for HelperBuilder, bag.py for Bag) +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 From 1f08dc1491d68231b353e3613e858047be106b50 Mon Sep 17 00:00:00 2001 From: bgailleton Date: Wed, 29 Jul 2026 13:48:40 +0200 Subject: [PATCH 31/31] More cleaning and add the grid manipulation example --- examples/grid/ball_walk_cupy.py | 250 ++++++++++++++++++ examples/grid/ball_walk_quadrants.py | 247 +++++++++++++++++ examples/grid/ball_walk_taichi.py | 243 +++++++++++++++++ .../core/context/_closure_backend.py | 34 ++- pyfastflow/experimental/core/context/bag.py | 38 ++- .../experimental/core/context/cupy_backend.py | 21 +- .../experimental/core/context/parameter.py | 44 +-- 7 files changed, 844 insertions(+), 33 deletions(-) create mode 100644 examples/grid/ball_walk_cupy.py create mode 100644 examples/grid/ball_walk_quadrants.py create mode 100644 examples/grid/ball_walk_taichi.py 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/core/context/_closure_backend.py b/pyfastflow/experimental/core/context/_closure_backend.py index 2c6adc8..2ccf0e8 100644 --- a/pyfastflow/experimental/core/context/_closure_backend.py +++ b/pyfastflow/experimental/core/context/_closure_backend.py @@ -135,7 +135,7 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No raise ValueError(f"{name}: field mode requires n_flat") self._handle = pool.get_data(dtype, (n_flat,)) - self.set(value) + self._store(value) @classmethod def _numpy_dtype(cls, dtype): @@ -164,14 +164,28 @@ def get(self): def set(self, value) -> None: """ - Overwrite the whole value: a cast python scalar for const, a device - write for scalar, a full host->device copy for field. + 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() - self._device_view = None # a cached view would bake the stale literal elif self.mode == "scalar": self._handle.data[None] = value else: # field @@ -208,12 +222,12 @@ 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. Two things can invalidate it - - set() on a const mode, which changes the literal baked into the getter, - and destroy(), which releases the storage it reads. Both drop the view - so the next caller rebuilds; neither reaches kernels compiled earlier - (see parameter.py, "Lifetime of a compiled object"). A scalar or field set() - needs no invalidation, writing through the very storage the view reads. + 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) """ diff --git a/pyfastflow/experimental/core/context/bag.py b/pyfastflow/experimental/core/context/bag.py index b1c0288..8bf705a 100644 --- a/pyfastflow/experimental/core/context/bag.py +++ b/pyfastflow/experimental/core/context/bag.py @@ -87,6 +87,34 @@ def __iter__(self): 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 @@ -291,10 +319,12 @@ 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 a Parameter's - mode is changed: mode is fixed at construction (see Parameter.mode), so - changing it means building a new Parameter and using replace() to swap it - into the bag in place of the old one. + `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) """ diff --git a/pyfastflow/experimental/core/context/cupy_backend.py b/pyfastflow/experimental/core/context/cupy_backend.py index ea9b1cb..4fcfa63 100644 --- a/pyfastflow/experimental/core/context/cupy_backend.py +++ b/pyfastflow/experimental/core/context/cupy_backend.py @@ -447,7 +447,7 @@ def __init__(self, name: str, *, dtype, mode: str, value, pool, n_flat: int | No raise ValueError(f"{name}: field mode requires n_flat") self._handle = pool.get_data(dtype, (n_flat,)) - self.set(value) + self._store(value) def get(self): """ @@ -459,8 +459,23 @@ def get(self): def set(self, value) -> None: """ - Overwrite the whole value: a cast python scalar for const, a device - write for scalar, a full host->device copy for field. + 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) """ diff --git a/pyfastflow/experimental/core/context/parameter.py b/pyfastflow/experimental/core/context/parameter.py index cdf8e46..d1328c9 100644 --- a/pyfastflow/experimental/core/context/parameter.py +++ b/pyfastflow/experimental/core/context/parameter.py @@ -24,9 +24,9 @@ Jargon ---------- Parameter One named, typed value. Its `mode` says where the value lives: - "const" (a compile-time constant, not modifiable mid-run), - "scalar" (a single value, modifiable) or "field" (a device - array, one value per node). + "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 - @@ -108,18 +108,25 @@ 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. So: - - - Writing to a scalar or field Parameter *is* visible to already-compiled - kernels, which hold that same storage. This is the normal way to feed - changing data. - - set() on a const Parameter is not. It drops the parameter's cached device - view so the next compile() picks the new value up, but kernels compiled - before it keep the old literal. Recompile them. +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. + again. Never destroy a Parameter that a live kernel still binds. This one + is not enforced at runtime. -None of this is 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 ----------------- @@ -221,9 +228,14 @@ def get(self): @abstractmethod def set(self, value) -> None: """ - Update the whole parameter value in place, according to its mode. On - const mode this does not reach already-compiled kernels - see the - module docstring, "Lifetime of a compiled object". + 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) """