diff --git a/content/index.md b/content/index.md index 13d9474..552e520 100644 --- a/content/index.md +++ b/content/index.md @@ -29,7 +29,7 @@ Work-in-progress proposals, organized by contributor. - [[personal/havogt/dtype-generic-fields|Dtype-generic fields in gt4py.next]] — keywords: type-system, generics, dtype, frontend, foast, past, monomorphization, field-operators, type-checking - [[personal/havogt/dimension-generic-fields|Generic dimensions and statically typed staggering]] — keywords: type-system, generics, dimensions, staggering, type-checking, mypy, frontend, foast, monomorphization, unstructured - [[personal/havogt/dependent-local-dimensions|Dependent local dimensions and connectivity chains]] — keywords: type-system, dimensions, unstructured, connectivities, local-dimensions, reduction, neighbor-sum, type-checking -- [[personal/havogt/closure-variable-resolution|Closure variable resolution in gt4py.next]] — keywords: frontend, foast, past, closure-variables, name-resolution, constants, builtins, aliasing, gtir, lowering +- [[personal/havogt/closure-variable-resolution|Closure variable resolution in gt4py.next]] — keywords: frontend, foast, past, closure-variables, name-resolution, constants, builtins, aliasing, gtir, lowering, contextvar, ambient, fields, mesh, static-args - [[personal/havogt/scan-redesign|Redesign of the vertical scan in gt4py.next]] — keywords: scan, vertical, reduction, boundary-conditions, windows, k-caches, fusion, embedded, jax, frontend - [[personal/havogt/mesh-and-first-class-halos|A mesh concept with first-class halos]] — keywords: mesh, halos, unstructured, connectivities, offset-provider, domain-inference, distributed, halo-exchange, prior-art - [[personal/havogt/field-data-protocol|A FieldData protocol for gt4py.next embedded fields]] — keywords: fields, domain, data, protocol, embedded, function-fields, boundary-conditions, materialization, lazy, concat_where, origin, prior-art diff --git a/content/personal/havogt/closure-variable-resolution.md b/content/personal/havogt/closure-variable-resolution.md index 8debe15..0b7245a 100644 --- a/content/personal/havogt/closure-variable-resolution.md +++ b/content/personal/havogt/closure-variable-resolution.md @@ -1,7 +1,7 @@ --- title: Closure variable resolution in gt4py.next author: havogt -tags: [frontend, foast, past, closure-variables, name-resolution, constants, builtins, aliasing, gtir, lowering] +tags: [frontend, foast, past, closure-variables, name-resolution, constants, builtins, aliasing, gtir, lowering, contextvar, ambient, fields, mesh, static-args] created: 2026-06-12 status: draft --- @@ -194,6 +194,162 @@ could bake the buffer in as a true constant, layered on `StaticArg`-style descriptors (ADR 0021 — Argument Descriptors). This fits A (the canonicalization pass introduces the hidden parameter) and becomes more natural under B. +### Stretch goal: ambient values bound at execution time + +The stretch goal above binds at *decoration* time. The variant worth recording +binds later: a user declares a value once, globally, and any field operator +reaches it without it appearing in a signature — a `ContextVar` filled at +program-execution time (JIT time for compiled backends). + +> **Prototyped**: gt4py branch +> [`ambient-offset-provider`](https://github.com/havogt/gt4py/pull/71) (fork PR). +> Everything below marked *measured* comes from there; everything marked *open* +> is still open. + +The motivating case is mesh properties. They are scalars on a Cartesian grid +(`dx`, `dy`) and fields on an unstructured one (connectivities plus their +weights); both must today be carried explicitly — as extra operator parameters, +or as the `offset_provider` threaded down the call chain. Every operator between +the caller and the one that actually needs a weight has to name it. + +#### Surface + +One rule for everything ambient: **the declaration is the key**. + +```python +V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) +dx = gtx.Static[float] # a value, folded into the generated code +nu = gtx.Extern[float] # a value, passed at runtime + + +@gtx.field_operator +def delta_x(f: IJField) -> IJField: + return (1.0 / dx) * (f(I + 1) - f) # never a parameter + + +prog(f, out, bind={V2E: connectivity, dx: 0.5}) # or: with gtx.bind(dx, 0.5): ... +``` + +A `FieldOffset` is *already* a declaration — it names the offset and fixes its +source and target — so it binds exactly like a value, and a program called +without an `offset_provider` assembles one from the bound offsets. A container +may declare what it supplies, with the class attribute holding the declaration +and the instance attribute the value: + +```python +class Mesh: + V2E = V2E # this mesh supplies the V2E connectivity + dx = physics.dx + + +prog(f, out, bind=Mesh(...)) +``` + +Binding by declaration *identity* rather than by attribute name is what lets a +container supply the very offset an operator refers to, rather than one that +merely shares its name — and it means a container attribute need not be named +after the declaration at all (*measured*: an attribute called `spacing` resolves +`physics.dx` correctly, across modules). + +`bind=` is sugar over the `ContextVar`, scoped to one call, so the two spellings +compose rather than compete. + +#### How it works + +- **A declaration carries its type.** `Static[float]` implements `__gt_type__`, + and `type_translation.from_value` already dispatches on that — so an operator + referring to an ambient value type-checks when it is *defined*, with nothing + bound. **No type-system change was needed.** An untyped placeholder fails with + `DSLTypeError: Unexpected object ...`, and the failure is temporal, not + spatial: it happens in the same file, so putting the declaration in another + module changes nothing. +- **The reference becomes a synthesised program parameter**, added once in + `func_to_past`. From there it travels the ordinary path: type checking, + lowering, `static_params` and the compiled-program key all treat it as an + argument, and only the value is supplied per call. +- **The two forms differ in one place only** — whether that parameter is listed + as static. `Extern[T]` stays a runtime argument; `Static[T]` is a static + argument, so the *existing* fold in `past_to_itir` bakes it in and the + *existing* `StaticArg` key specialises on it. *Measured* on gtfn and dace with + two values: `Static` compiles 2 variants, `Extern` 1, both correct. +- **Connectivities identify by content, not by `id`.** `gtx.freeze(conn)` caches + a content hash once; `hash_offset_provider_items_by_id` prefers it. *Measured*: + 3 compiled variants drop to 2 when two structurally identical meshes stop + being keyed apart by object identity — that function's own docstring warns it + "could generate different hashes for two offset providers that are + semantically equal". + +Three things it turned out **not** to need, each of which was assumed at some +point in the design: a new FOAST specialization step (the existing static-argument +fold suffices); threading the parameter into operator signatures and call sites +(a free symbol in a lowered operator resolves against the *program's* parameters, +and both gtfn and dace codegen fine that way); and a per-operator inspection pass +(`transform_utils._get_closure_vars_recursively` already collects declarations +transitively through nested operators). + +#### Transition: retiring the second binding rule + +The prototype first grew *two* binding mechanisms with different identity rules: +connectivities were harvested from a bound object by attribute **name** (a +`Namespace`), values were keyed by declaration **identity**. Same surface, two +rules — and the name-based half needed a collision check, because two namespaces +could each offer an offset called `V2E`. + +Unifying on declarations removes `Namespace`, the attribute harvesting and that +collision check outright, and it is worth doing in two steps: + +1. **Move only the binding surface.** Offsets bind by declaration; the + `offset_provider` is still assembled from the bound ones and passed exactly as + today, so the frontend, the IR and the backends are untouched. *This step is + implemented.* +2. **Then consider making connectivities ordinary ambient values**, i.e. an + `Extern[Connectivity]` that becomes a synthesised parameter like any other, + at which point `offset_provider` stops being a separate concept. This is not + obviously right and should not be assumed: `offset_provider` carries things + the parameter path does not model today (its *type* drives domain inference + and connectivity typing, and `arguments.py` still notes the temporary pass + needs the runtime object). Worth a separate look rather than a follow-through. + +The staging matters because step 1 is a pure surface change with no risk to the +toolchain, while step 2 touches how connectivities are typed and inferred. + +#### Binding time + +"Static for the lifetime of the application" is the wrong unit; the useful one is +**static for a jitted program**. That makes the two forms a declaration-site +choice of what identifies a value in the jit key: `Static[T]` puts the *value* +there, `Extern[T]` only its type. Which also dissolves an apparent conflict with +[[personal/havogt/jax-connectivities/jax-connectivities|the JAX connectivities proposal]] — a +connectivity wants descriptor-identity (one program per mesh *class*), a scalar +wants value-identity, and nothing forces one policy on both. + +#### Open + +- **Canonical identifiers.** Synthesised parameters currently take the closure + variable's *local* name, so two operators naming the same declaration + differently, or two modules that both use `dx`, collide. Declarations should be + named explicitly (`Static[float]("dx")`) with a FOAST rename pass; a generated + counter will not do, since the name lands in the compiled signature and would + shift with import order. +- **Ambient fields** (`mesh.edge_length`) are not implemented, but now look like + an `Extern[Field[...]]` — a synthesised parameter with no folding, i.e. the + read-only-field stretch goal above with a later binding time. +- **Immutability.** `freeze(readonly=True)` would make the cached content hash + trustworthy, but the gtfn bindings are generated with mutable `ndarray` + parameters and reject a read-only array outright, so it is off by default. +- **Debuggability**, unchanged: errors when nothing is bound must be good, and + there should be a way to see what an operator depends on ambiently. + +The ceiling this aims at is a mesh concept built on top, where connectivities +and weights are properties of an ambient mesh rather than arguments — see +[[personal/havogt/mesh-and-first-class-halos|A mesh concept with first-class halos]]. +It would also be a plausible substrate for +[[personal/egparedes/discretization-independent-fd-syntax|A discretization-independent +surface syntax]], whose mesh-invariant surface needs weights and connectivity to +come from a declared mesh property rather than from operator arguments (that +proposal is a higher-level surface layered on gt4py's core concepts, not a +replacement for them). + ### Source material - Proposed ADR and **partial prototype of strategy A** on gt4py branch diff --git a/content/personal/havogt/dimension-generic-fields/__pycache__/test_typed_dimensions.cpython-312-pytest-9.1.1.pyc b/content/personal/havogt/dimension-generic-fields/__pycache__/test_typed_dimensions.cpython-312-pytest-9.1.1.pyc new file mode 100644 index 0000000..bdf5d5a Binary files /dev/null and b/content/personal/havogt/dimension-generic-fields/__pycache__/test_typed_dimensions.cpython-312-pytest-9.1.1.pyc differ diff --git a/content/personal/havogt/dimension-generic-fields/__pycache__/typed_dimensions.cpython-312.pyc b/content/personal/havogt/dimension-generic-fields/__pycache__/typed_dimensions.cpython-312.pyc new file mode 100644 index 0000000..2eeaaf9 Binary files /dev/null and b/content/personal/havogt/dimension-generic-fields/__pycache__/typed_dimensions.cpython-312.pyc differ