proto[next]: ambient binding of the offset provider - #71
Conversation
Prototype for the 'ambient values' idea in the knowledge base: a program with no 'offset_provider=' takes its connectivities from whatever is bound to an ambient namespace at call time. Binding half only; ambient *fields* referenced by name inside an operator are not implemented.
A frozen element carries a content hash computed once at freeze time, and 'hash_offset_provider_items_by_id' prefers it over 'id(v)'. Two semantically equal meshes then share a compiled program instead of triggering a recompile. Unfrozen elements keep the previous id-based behaviour.
…tors
Adds a second spelling next to the 'gtx.bind' context manager: a 'bind={ns: value}'
kwarg where 'offset_provider=' goes today, scoped to that one call. Both program
and field operator entry points now resolve the offset provider from the ambient
context when the caller passes none, so direct field-operator calls work too.
'__call__' becomes a thin wrapper around '_invoke' so the existing bodies are
untouched; 'ProgramWithBoundArgs' overrides '_invoke' accordingly.
|
Added The first prototype only patched Two spellings, same mechanismwith gtx.bind(mesh, m): # scoped to a region
prog(a, out)
prog(a, out, bind={mesh: m}) # scoped to one call, where offset_provider= goes today
All six combinations plus embedded: Implementation note
Tests
Incidentally, writing those tests reproduced the very bug |
Two defects the backend matrix exposed, neither visible in a single-connectivity
demo or in embedded-only tests:
- 'offset_provider_of' built the mapping from 'dir()' (alphabetical) rather than
the bound object's own '__dict__' (insertion order). gt4py's offset provider is
order-sensitive, so a multi-connectivity mesh got wrong results on gtfn and a
segfault elsewhere.
- 'freeze' hashed via 'np.asarray', which refuses device arrays; every GPU
backend failed. Uses 'asnumpy()' now.
Adds a backend-matrix test covering {program, field_operator} x
{offset_provider=, context manager, bind=}.
A declaration carries the type, so an operator referring to an ambient value
types at decoration without the value being bound: 'from_value' dispatches on
'__gt_type__', so no type-system change is needed.
Execution is not wired up yet. The reference survives FOAST but dies at
'itir.Program' construction ('Symbols {SymbolRef(dx)} not found'), because eve
validates symbol refs before any remap can run. Both forms therefore need the
reference to become a real parameter first.
An operator can refer to 'dx = Static[float]' without it being a parameter, so it need not be threaded through nested operators. Three pieces: - embedded: a bound declaration behaves as the scalar it stands for. - compiled: '_SubstituteAmbientValues' replaces the reference by its value before lowering, since 'ClosureVarFolding' runs at decoration when nothing is bound yet. - caching: bound values enter the compiled-program key, and are mirrored onto the declaration so the lowering cache (which fingerprints closure variables) cannot serve one value's code for another. 'Extern[T]' currently folds like 'Static[T]'; making it a runtime argument needs a synthesised program parameter, as does the ambient-field case.
|
Added dx = gtx.Static[gtx.float64] # declared once, never a parameter
@gtx.field_operator
def delta_x(f: IJFloatField) -> IJFloatField:
"""Forward difference in x."""
return (1.0 / dx) * (f(IDim + 1) - f)
run_delta_x(f, out, bind={dx: 0.5}) # or: with gtx.bind(dx, 0.5):
How it works
Correction to the earlier planI said Honest gap
Tests
Verified on CPU backends only this round (47 passed ambient, 162 passed |
The lowering cache did not distinguish two bindings of the same declaration because its fingerprinter walks closure variables and sees only the declaration object. That was worked around by mirroring the bound value onto the declaration's '__dict__', which made the binding process-wide. Register a deconstructor for 'AmbientValue' with the frontend fingerprinter instead: it reads the current binding from the ContextVar, so the cache distinguishes values and bindings stay context-local.
Replaces the bespoke machinery with a parameter, added once when the program is
defined ('func_to_past'). From there ambient values travel the ordinary path:
type checking, lowering, 'static_params' and the compiled-program key all treat
them as arguments, and only the value is supplied per call.
The two forms now differ in one place only — whether the synthesised parameter
is listed as static:
- 'Extern[T]': ordinary runtime argument, one compiled program for all values.
- 'Static[T]': static argument, so the existing fold bakes it in and the
existing key specialises on it — one program per distinct value.
Measured on gtfn and dace, 2 values: Static 2 variants, Extern 1, both correct.
Deletes '_SubstituteAmbientValues', 'current_static_key' and the ambient
fingerprint deconstructor: the operator's IR no longer holds the value, so the
lowering cache needs no help. A free symbol in a lowered operator resolves
against the program's parameters, so no threading into operator signatures is
needed.
|
What changedA declaration becomes a program parameter once, when the program is defined ( The two forms then differ in exactly one place: whether that parameter is listed as a static one.
Two findings that made it small
Deleted
That last one is the satisfying part: the cache-staleness problem is gone rather than worked around, because the operator's IR no longer contains the value at all. Known gap: canonical identifiersSynthesised parameters currently take the closure variable's local name. Two operators referring to the same declaration under different local names, or two modules that both call theirs Parameters are already sorted by name so the signature is deterministic. Tests
Verified on CPU backends: 80 ambient tests, plus 328 in |
Connectivities and values were bound by two different rules: a 'Namespace'
harvested connectivities from a bound object by attribute *name*, while values
were keyed by declaration identity. Now there is one rule — the declaration is
the key:
prog(a, out, bind={V2E: connectivity, dx: 0.5})
A 'FieldOffset' already names the offset and fixes its source and target, so it
is the declaration; the offset provider is assembled from the bound ones. This
retires 'Namespace', 'offset_provider_of' and the name-collision check it
needed, and lets a container supply the very offset an operator refers to rather
than one that merely shares its name.
The offset provider itself is unchanged: it is still assembled and passed as
today, so only the binding surface moves.
| @gtx.field_operator | ||
| def delta_x(f: IJFloatField) -> IJFloatField: | ||
| """Forward difference in x.""" | ||
| return (1.0 / dx) * (f(IDim + 1) - f) |
There was a problem hiding this comment.
| return (1.0 / dx) * (f(IDim + 1) - f) | |
| return (1.0 / grid.dx) * (f(IDim + 1) - f) |
where the __getattr__ does a contextvar.get()
…ithmetic protocol
Declarations now live in a container and are read as 'grid.dx'. The declaration
is a descriptor: class access ('Grid.dx') yields the declaration, which is what
'bind=' takes as a key; instance access yields the bound value. Embedded
execution therefore sees a plain scalar, and the 12 arithmetic dunders that made
'AmbientValue' impersonate one are gone — they were also incomplete (no
comparisons, no '%', no numpy interop).
A container-qualified name ('Grid_dx') gives the synthesised parameter an
identifier that cannot collide across modules, which was the known gap.
Only declarations an operator actually *reads* become parameters: a container is
a place to declare things, and an operator that never reads 'grid.dx' must not
acquire it — for a 'Static[T]' that would specialise the compiled program on an
unused value.
The frontend needed two hooks: a container types itself as a namespace over its
own class (so 'grid.dx' resolves to the declared type with nothing bound), and
the lowering rewrites the attribute to a reference to the synthesised parameter.
Declarations are annotations in a container ('dx: Static[float]') and bind to a
'contextvars.ContextVar'. 'Static[T]'/'Extern[T]' are PEP 695 aliases over
'Annotated', so a type checker sees plain 'T' while the binding machinery reads
the marker — the bespoke 'AmbientValue' object is gone, and with it the
descriptor it needed.
'Grid.dx' (class access) is the variable, which 'bind=' takes as a key;
'grid.dx' (instance access) is its value. Binding is stdlib set/reset, so the
dict-in-a-ContextVar store disappears; a 'FieldOffset' carries its own variable
instead of a central registry.
The offset provider is now assembled from the offsets *this* program references
rather than from everything currently bound. That removes global state and fixes
a real defect: an unrelated bound mesh leaked into every program's offset
provider, where it also perturbed the compiled-program key and forced spurious
recompiles.
'Grid(dx=0.5, nu=1e-3)' carries values and binds every declaration it holds, so the caller provides the grid (or the mesh) as one thing and each program picks the parts it needs — rather than the caller tracking which program reads what. The two uses of a container instance do not collide: one constructed with values keeps them in its instance dict, so attribute access finds them directly; one constructed empty has nothing there, so access falls through to '__getattr__' and reads the bound variable. That is the instance an operator reads through. 'gtx.bind' now also takes containers, and an undeclared keyword is rejected at construction rather than silently ignored.
Two containers with the same class name in different modules produced the same synthesised parameter, and the result was silently wrong rather than an error: both declarations collapsed onto one parameter and one binding won. A container class name is not unique, so the parameter now carries a stable digest of the fully qualified name, and the declaration keeps the qualified name for diagnostics. A second collision sat behind it: '_get_closure_vars_recursively' merges by name, so two modules that both call their container 'grid' shadowed one another. Every caller now uses one per-operator walk that looks at each operator's own closure variables instead of a merged mapping.
Container identity ends up in the synthesised parameter name, which changes the stage fingerprint and therefore the build-cache key. 'id()' would be unique but would differ on every interpreter restart, so the cache would never hit; module and qualified name are stable but not always unique. Two containers built by the same factory share both and nothing stable distinguishes them, so they are now rejected at definition rather than silently sharing a parameter. 'class Grid(Container, name=...)' separates them when that is intended. The registry backing the check is weak, write-once at class definition, and never consulted on the execution path — unlike the offset registry removed earlier, which sat on the lookup path and leaked.
1 participant
Prototype, not for upstream. Backs the "ambient values" stretch goal in
gt4py_knowledge#29 — the
document should link this branch rather than vendoring a mock, since the point
is that
offset_provider=disappears from the real call path.What it does
A program called without
offset_providercollects one from whatever is boundto an ambient namespace: every
Connectivity/Dimensionreachable as a publicattribute of the bound object, keyed by attribute name. Colliding names across
namespaces raise rather than silently picking one mesh's table.
Binding is a
contextvars.ContextVar, so it nests, unwinds, and is per-context —which is what lets the same programs run against a second mesh in one process.
gt4py already carries
_offset_providerthis way inembedded/context.py; thisgeneralizes it to the compiled path and makes it user-facing.
Content hashing of frozen elements
Ambient values are static for the jitted programs that see them, so they can
identify by content rather than by
id.gtx.freeze(conn)computes a contenthash once and caches it on the element;
common.hash_offset_provider_items_by_idprefers it over
id(v)and falls back for anything unfrozen, so existingbehaviour is unchanged.
This fixes a real recompile: that function's own docstring warns it "could
generate different hashes for two offset providers that are semantically equal".
Measured with
tmp/ambient_demo.pyover 4 calls / 3 mesh objects, two of whichhold byte-identical tables:
id-keyed (before)The O(size) cost is paid once at freeze time, never per call.
Findings worth keeping
readonly=Truefreezing breaks the gtfn bindings. Marking the buffernon-writeable is what would make the cached hash trustworthy, but the generated
pybind11 signatures take mutable
ndarrayparameters and reject a read-onlyarray outright (
TypeError: run(): incompatible function arguments). So it isoff by default and the hash is currently only as stable as the caller's
discipline. Fixing the bindings to accept read-only arrays is the real
prerequisite for immutable ambient values.
KeyError: "Offset 'V2E' not found in offset provider."from deep in lowering, naming neither the namespace nor thebinding. Should become a proper
DSLError.arguments.py:148TODO (temporary pass still needs runtime offset-providerinformation) is why the key change is opt-in via freezing rather than a
wholesale switch to type-based keying.
Scope
Binding half only. Referring to ambient fields by name inside an operator —
mesh.edge_length— is not implemented; that needs the hidden-parametermechanism (
GTCallable.__gt_implicit_args__) from the read-only-field stretchgoal. Worth noting for that slice:
type_translation.pyalready hasNamespaceProxyoverFrozenNamespace | EnumMeta | ModuleType, socontainer-attribute, bare-name and module-prefix access are the same existing
mechanism — what's missing there is fields as values and late binding, not the
access syntax.
Tests
tests/next_tests/unit_tests/test_ambient.py— 8 backend-free tests (binding,nesting, collisions, freeze, cache-key collapse and separation, readonly).
Regression:
test_arg_call_interface.py297 passed;offset_provider=Nonestillresolves to
{}when nothing is bound, so the default path is unchanged.