feat[next]!: connectivities declared as classes, class-keyed offset providers, table_types; remove FieldOffset

docs/development/ADRs/next/0019-Connectivities.md

@@ -9,6 +9,10 @@ tags: []
 - **Created**: 2024-11-08
 - **Updated**: 2026-05-27
 
+> The `FieldOffset` part of this record is superseded by
+> [ADR 0030](0030-Connectivities_As_Types.md): connectivities are declared as
+> `NeighborConnectivity` classes, and offset providers are keyed by them.
+
 The representation of Connectivities (neighbor tables, `NeighborTableOffsetProvider`) and their identifier (offset tag, `FieldOffset`, etc.) was extended and modified based on the needs of different parts of the toolchain. Here we outline the ideas for consolidating the different closely-related concepts.
 
 ## History

docs/development/ADRs/next/0029-Dimensions_As_Nominal_Types.md

@@ -53,7 +53,8 @@ string equality and are never checked against each other at declaration time: th
 name, and the `offset_provider` key. Whichever one reaches
 `common.get_offset` depends on the execution path and the operation. Making a
 dimension's identity its Python type is the prerequisite for collapsing those
-names into one declaration (a follow-up ADR covers the connectivity half).
+names into one declaration ([ADR 0030](0030-Connectivities_As_Types.md) covers the
+connectivity half).
 
 ## Decision
 
@@ -82,10 +83,11 @@ disappears.
 
 1. **Reconstruction from the IR is an import.** `common.resolve(tag)` imports the
    module and walks the qualname; nested declarations resolve naturally. The IR
-   references a Python type exactly the way `pickle` references a class. It is
-   memoized, because type inference calls it once per `AxisLiteral`. An
-   `AxisLiteral` stores only the tag: its `kind` is the resolved dimension's, so the
-   two cannot disagree.
+   references a Python type exactly the way `pickle` references a class. Where the
+   module path ends is memoized, because type inference calls it once per
+   `AxisLiteral`; the attribute walk is repeated, so a redefined declaration is
+   found. An `AxisLiteral` stores only the tag: its `kind` is the resolved
+   dimension's, so the two cannot disagree.
 
    A purely dotted tag does not record where the module path ends and the
    qualname begins, so `resolve` tries the *longest importable prefix* and walks
@@ -216,7 +218,8 @@ The last row needs `DimensionMeta.__add__` / `__sub__` declared with the self-ty
 site and both reject it at the definition site, with different diagnostics (mypy
 `[misc]`, pyright `reportGeneralTypeIssues`), so it costs two separately spelled
 suppressions. The runtime check covers unannotated code; hand-written iterator IR,
-which names dimensions by tag, is not checked. Comparisons are deliberately *not* restricted: `D == n`
+which names dimensions by tag, is not checked. `as_offset(dim, field)` needs index
+arithmetic too and takes an `AnyCartesianAxisIndex`. Comparisons are deliberately *not* restricted: `D == n`
 and `D < n` build a `Domain` on every dimension, as `concat_where` over a mesh
 location requires.
 

docs/development/ADRs/next/0030-Connectivities_As_Types.md

@@ -75,7 +75,8 @@ and skip values were never checked against the `FieldOffset` declaration.
   `(Domain, V2E.Local)`, the codomain is `Codomain`, the dtype is integral, and
   the neighbor counts and skip values agree. Skip values are checked on the
   table's type: a table with a `skip_value` counts as having skip values whether
-  or not an entry uses it.
+  or not an entry uses it. Programs run the check on the tables they are given,
+  see below.
 
 `Domain` and `Codomain` name the two index spaces the declaration maps between.
 A bound table is a field over `(Domain, Local)` with values in `Codomain`: the
@@ -93,8 +94,7 @@ of a table bound to a declaration is a `common.NeighborTableType`:
 `domain` and `codomain` are derived from the declaration,
 `(connectivity.domain, local_dimension_of(connectivity))` and
 `connectivity.codomain`, so they cannot disagree with it. The mapping from
-offset-provider keys to these records is `common.TableTypes`, and it can be
-given instead of the tables for ahead-of-time compilation.
+offset-provider keys to these records is `common.TableTypes`.
 
 A table cannot tell which declaration it is bound to: the table of a sharer
 (`C2CE`) has the same domain as its owner's (`C2E`), with another codomain. So a
@@ -173,11 +173,11 @@ treating it as a type.
 
 ### Frontend integration
 
-A declaration is typed like the `FieldOffset` it replaces: `V2E.__gt_type__()`
-is a `ts.ShiftType`, which takes a field over the codomain to one over the
-domain, `Shift[<tag>: Edge -> (Vertex, V2E.Local)]`. `V2E[i]` has the domain
-`(Vertex,)`, and so does a Cartesian shift `KDim + 1`, over `KDim` and without a
-tag. The tag is the connectivity's `offset_tag`:
+A declaration is typed as a shift: `V2E.__gt_type__()` is a `ts.ShiftType`,
+which takes a field over the codomain to one over the domain,
+`Shift[<tag>: Edge -> (Vertex, V2E.Local)]`. `V2E[i]` has the domain `(Vertex,)`,
+and so do the Cartesian shifts `KDim + 1` and `as_offset(KDim, offsets)`, over
+`KDim` and without a tag. The tag is the connectivity's `offset_tag`:
 
 - **the local dimension's tag**, `V2E.Local.tag`, for the connectivity that
   declares it. This is the single string that shifts, neighbor reductions and
@@ -190,15 +190,64 @@ tag. The tag is the connectivity's `offset_tag`:
   over it (`common.connectivity_key_over`): the owner's if bound, else the
   sharer with the smallest tag. Connectivities sharing a local dimension must
   therefore have the same neighbor *structure* — the same count, and a skip value
-  at the same positions — which is what sharing a neighbor axis means.
+  at the same positions — which is what sharing a neighbor axis means;
+  `check_offset_provider` enforces it for the tables it is given.
 
-`V2E.Local` inside DSL code types as that local dimension, and
-`FieldOffset.Local` names the same thing on a legacy offset, so the spelling
-works for both. The other frontend touch points treat the class like the
-`FieldOffset` it derives: grid-type deduction (`transform_utils`, `past_to_itir`)
-counts it as unstructured, and embedded `premap` accepts it. `V2E[i]` subscripts the
+`V2E.Local` inside DSL code types as that local dimension. The other frontend
+touch points treat a declaration as an unstructured shift: grid-type deduction
+(`transform_utils`, `past_to_itir`) counts it as unstructured, and embedded
+`premap` accepts it. `V2E[i]` subscripts the
 metaclass, which forwards type-parameter subscription (`NeighborConnectivity[V, E]`) to `__class_getitem__`, since a metaclass `__getitem__` shadows it.
 
+### Offset providers are keyed by the declaration
+
+Users bind tables to declarations:
+
+```python
+program(..., offset_provider={V2E: v2e_table, C2E: c2e_table})
+```
+
+Every entry point of a program normalizes such a provider to the form the IR
+uses: each declaration is replaced by its `offset_tag`. Everything below the
+entry points — lowering, the backends, compiled-program caching — therefore keeps
+seeing a provider keyed by strings, which is also what hand-written IR uses.
+
+The frontend entry points (`Program.__call__`, `FieldOperator.__call__`,
+`compile`, `CompilationOptions.connectivities`) are *strict*: a string key must
+be a tag, i.e. a qualified name, and a bare name such as `"V2E"` is the removed
+`FieldOffset` spelling, rejected with a message pointing here. The IR-level hooks
+(`embedded.context.update`, the iterator `fendef`, DaCe's `get_sdfg_conn_args`)
+accept any string, because a hand-written program names its offsets itself.
+
+The types of the tables, a `common.TableTypes` under the same keys, are called
+`table_types` throughout: `CompileTimeArgs.table_types`, the IR passes and the code
+generators. Ahead-of-time compilation can take them in place of the tables,
+e.g. `{V2E: NeighborTableType(connectivity=V2E, dtype=int32, skip_value=None, max_neighbors=6)}` (`compile(offset_provider=...)` accepts either), and they are
+keyed and normalized exactly like an offset provider: by declarations, strictly,
+at the frontend; by tags at the IR level.
+
+Tables are checked against their declarations (`check_offset_provider`) at every
+entry point, but the result is remembered per set of bound tables, so repeated
+calls cost one hash. Reading the tables — comparing the skip-value positions of
+two connectivities that share a local dimension — is done only where a program is
+compiled, not on the call path. A tag that
+names no declared connectivity, as in hand-written IR, is not checked.
+
+### `FieldOffset` is removed
+
+`FieldOffset` and its export are gone. An unstructured connectivity is a
+`NeighborConnectivity`; a Cartesian shift is `Dim + i`, which the DSL already
+had; and `as_offset` takes the Cartesian axis to shift along,
+`as_offset(KDim, k_offsets)`, instead of a Cartesian `FieldOffset`.
+`scripts/python/migrate_connectivities.py` rewrites declarations and Cartesian
+offset uses, and reports the provider keys and other sites it cannot rewrite from
+the source alone. It decides per declaration whether a dimension is a Cartesian
+axis (`kind=VERTICAL`, a Cartesian offset, `as_offset`, index arithmetic or a
+staggered counterpart) or a mesh location (a domain or codomain of a neighbor
+offset), and reports a dimension with no evidence, or with evidence for both,
+instead of guessing. It drops aliases of the removed `DimensionKind.LOCAL` and
+reports its other uses.
+
 ## Consequences
 
 - An unstructured connectivity is spelled once. The provider key, the offset tag
@@ -211,8 +260,8 @@ metaclass, which forwards type-parameter subscription (`NeighborConnectivity[V,
 - A declaration is fingerprinted by its name *and* its declared dimensions and
   counts, so redefining it under the same name (e.g. re-running a notebook
   cell) does not reuse artifacts compiled for the old declaration.
-- `FieldOffset` remains during migration; a `FieldOffset` and a
-  `NeighborConnectivity` sharing a local dimension are interchangeable.
+- `FieldOffset` is removed, and offset providers are keyed by declarations: a
+  breaking change for every unstructured program, eased by the migration script.
 
 ## Alternatives considered
 

docs/user/next/QuickstartGuide.md

@@ -218,28 +218,28 @@ edge_values = gtx.as_field([EdgeDim], np.zeros((12,)))
 
 +++
 
-You can transform fields (or tuples of fields) over one domain to another domain by using the call operator of the source field with a _field offset_ as argument. This transform uses the connectivity between the source and target domains to find the values of adjacent mesh elements.
+You can transform fields (or tuples of fields) over one domain to another domain by using the call operator of the source field with a _connectivity_ as argument. This transform uses the connectivity between the source and target domains to find the values of adjacent mesh elements.
 
 To understand this transform, you can look at the edge-to-cell connectivity table `edge_to_cell_table` listed above. This table has the same shape as the output of the transform, that is, one dimension over the edges and another _local_ dimension. The table stores indices into a field over cells, the transform essentially gives you another field where the indices have been replaced with the values in the cell field at the corresponding indices.
 
 Another way to look at it is that transform uses the edge-to-cell connectivity to look up all the cell neighbors of edges, and associates the values of those neighbor cells with each edge.
 
-You can use the field offset `E2C` below to transform a field over cells to a field over edges using the edge-to-cell connectivities:
+You can use the connectivity `E2C` declared below to transform a field over cells to a field over edges using the edge-to-cell connectivities. It is declared as a class: for each edge (`EdgeDim`), a list of neighbor cells (`CellDim`), indexed by its nested local dimension `E2C.Local`:
 
 ```{code-cell} ipython3
-class E2CDim(gtx.LocalDimensionIndex): ...
-E2C = gtx.FieldOffset(E2CDim.tag, source=CellDim, target=(EdgeDim, E2CDim))
+class E2C(gtx.NeighborConnectivity[EdgeDim, CellDim]):
+    class Local(gtx.LocalDimensionIndex): ...
 ```
 
-The field offset is named by its local dimension's `tag`, and the offset provider below is keyed by the same `tag`, so all three refer to one connectivity. Note that the field offset does not contain the actual connectivity table, that's provided through an _offset provider_:
+Note that the declaration does not contain the actual connectivity table, that's provided through an _offset provider_, a dictionary from connectivity declarations to tables:
 
 ```{code-cell} ipython3
-E2C_offset_provider = gtx.as_connectivity([EdgeDim, E2CDim], codomain=CellDim, data=edge_to_cell_table, skip_value=-1)
+E2C_offset_provider = gtx.as_connectivity([EdgeDim, E2C.Local], codomain=CellDim, data=edge_to_cell_table, skip_value=-1)
 ```
 
 The field operator `nearest_cell_to_edge` below shows an example of applying this transform. There is a little twist though: the subscript in `E2C[0]` means that only the value of the first connected cell is taken, the second (if exists) is ignored.
 
-Pay attention to the syntax where the field offset `E2C` can be freely accessed in the field operator, but the offset provider `E2C_offset_provider` is passed in a dictionary to the program.
+Pay attention to the syntax where the connectivity `E2C` can be freely accessed in the field operator, but the offset provider `E2C_offset_provider` is passed in a dictionary to the program.
 
 ```{code-cell} ipython3
 @gtx.field_operator
@@ -250,7 +250,7 @@ def nearest_cell_to_edge(cell_values: gtx.Field[Dims[CellDim], float64]) -> gtx.
 def run_nearest_cell_to_edge(cell_values: gtx.Field[Dims[CellDim], float64], out : gtx.Field[Dims[EdgeDim], float64]):
     nearest_cell_to_edge(cell_values, out=out)
 
-run_nearest_cell_to_edge(cell_values, edge_values, offset_provider={E2CDim.tag: E2C_offset_provider})
+run_nearest_cell_to_edge(cell_values, edge_values, offset_provider={E2C: E2C_offset_provider})
 
 print("0th adjacent cell's value: {}".format(edge_values.asnumpy()))
 ```
@@ -265,19 +265,19 @@ Running the above snippet results in the following edge field:
 
 #### Using reductions on connected mesh elements
 
-Similarly to the previous example, the output is once again a field on edges. The difference is that this field operator does not take the first column of the transformed field, but sums the columns. In other words, the result is the sum of all the cells adjacent to an edge. You can achieve this by first transforming the cell field to a field over the cell neighbors of edges (i.e. a field of dimensions Edge × E2CDim) using `cells(E2C)`, then calling the `neighbor_sum` builtin function to sum along the `E2CDim` dimension.
+Similarly to the previous example, the output is once again a field on edges. The difference is that this field operator does not take the first column of the transformed field, but sums the columns. In other words, the result is the sum of all the cells adjacent to an edge. You can achieve this by first transforming the cell field to a field over the cell neighbors of edges (i.e. a field of dimensions Edge × E2C.Local) using `cells(E2C)`, then calling the `neighbor_sum` builtin function to sum along the `E2C.Local` dimension.
 
 ```{code-cell} ipython3
 @gtx.field_operator
 def sum_adjacent_cells(cells : gtx.Field[Dims[CellDim], float64]) -> gtx.Field[Dims[EdgeDim], float64]:
-    # type of cells(E2C) is gtx.Field[Dims[EdgeDim, E2CDim], float64]
-    return neighbor_sum(cells(E2C), axis=E2CDim)
+    # type of cells(E2C) is gtx.Field[Dims[EdgeDim, E2C.Local], float64]
+    return neighbor_sum(cells(E2C), axis=E2C.Local)
 
 @gtx.program
 def run_sum_adjacent_cells(cells : gtx.Field[Dims[CellDim], float64], out : gtx.Field[Dims[EdgeDim], float64]):
     sum_adjacent_cells(cells, out=out)
 
-run_sum_adjacent_cells(cell_values, edge_values, offset_provider={E2CDim.tag: E2C_offset_provider})
+run_sum_adjacent_cells(cell_values, edge_values, offset_provider={E2C: E2C_offset_provider})
 
 print("sum of adjacent cells: {}".format(edge_values.asnumpy()))
 ```
@@ -376,13 +376,13 @@ print("where nested tuple return: {}".format(((result_1.asnumpy(), result_2.asnu
 
 #### Implementing the pseudo-laplacian
 
-As explained in the section outline, the pseudo-laplacian needs the cell-to-edge connectivities as well in addition to the edge-to-cell connectivities. Though the connectivity table has been filled in above, you still need to define the local dimension, the field offset, and the offset provider that describe how to use the connectivity table. The procedure is identical to the edge-to-cell connectivity from before:
+As explained in the section outline, the pseudo-laplacian needs the cell-to-edge connectivities as well in addition to the edge-to-cell connectivities. Though the connectivity table has been filled in above, you still need to declare the connectivity and the offset provider that describe how to use the connectivity table. The procedure is identical to the edge-to-cell connectivity from before:
 
 ```{code-cell} ipython3
-class C2EDim(gtx.LocalDimensionIndex): ...
-C2E = gtx.FieldOffset(C2EDim.tag, source=EdgeDim, target=(CellDim, C2EDim))
+class C2E(gtx.NeighborConnectivity[CellDim, EdgeDim]):
+    class Local(gtx.LocalDimensionIndex): ...
 
-C2E_offset_provider = gtx.as_connectivity([CellDim, C2EDim], codomain=EdgeDim, data=cell_to_edge_table, skip_value=-1)
+C2E_offset_provider = gtx.as_connectivity([CellDim, C2E.Local], codomain=EdgeDim, data=cell_to_edge_table, skip_value=-1)
 ```
 
 **Weights of edge differences:**
@@ -410,7 +410,7 @@ edge_weights = np.array([
     [0, -1, -1], # cell 5
 ], dtype=np.float64)
 
-edge_weight_field = gtx.as_field([CellDim, C2EDim], edge_weights)
+edge_weight_field = gtx.as_field([CellDim, C2E.Local], edge_weights)
 ```
 
 Now you have everything to implement the pseudo-laplacian. Its field operator requires the cell field and the edge weights as inputs, and outputs a cell field of the same shape as the input.
@@ -422,17 +422,17 @@ The second lines first creates a temporary field using `edge_differences(C2E)`,
 ```{code-cell} ipython3
 @gtx.field_operator
 def pseudo_lap(cells : gtx.Field[Dims[CellDim], float64],
-               edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64]) -> gtx.Field[Dims[CellDim], float64]:
+               edge_weights : gtx.Field[Dims[CellDim, C2E.Local], float64]) -> gtx.Field[Dims[CellDim], float64]:
     edges = cells(E2C[0]) # type: gtx.Field[Dims[EdgeDim], float64]
-    return neighbor_sum(edges(C2E) * edge_weights, axis=C2EDim)
+    return neighbor_sum(edges(C2E) * edge_weights, axis=C2E.Local)
 ```
 
 The program itself is just a shallow wrapper over the `pseudo_lap` field operator. The significant part is how offset providers for both the edge-to-cell and cell-to-edge connectivities are supplied when the program is called:
 
 ```{code-cell} ipython3
 @gtx.program
 def run_pseudo_laplacian(cells : gtx.Field[Dims[CellDim], float64],
-                         edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64],
+                         edge_weights : gtx.Field[Dims[CellDim, C2E.Local], float64],
                          out : gtx.Field[Dims[CellDim], float64]):
     pseudo_lap(cells, edge_weights, out=out)
 
@@ -441,7 +441,7 @@ result_pseudo_lap = gtx.as_field([CellDim], np.zeros(shape=(6,)))
 run_pseudo_laplacian(cell_values,
                      edge_weight_field,
                      result_pseudo_lap,
-                     offset_provider={E2CDim.tag: E2C_offset_provider, C2EDim.tag: C2E_offset_provider})
+                     offset_provider={E2C: E2C_offset_provider, C2E: C2E_offset_provider})
 
 print("pseudo-laplacian: {}".format(result_pseudo_lap.asnumpy()))
 ```
@@ -451,7 +451,7 @@ As a closure, here is an example of chaining field operators, which is very simp
 ```{code-cell} ipython3
 @gtx.field_operator
 def pseudo_laplap(cells : gtx.Field[Dims[CellDim], float64],
-                  edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64]) -> gtx.Field[Dims[CellDim], float64]:
+                  edge_weights : gtx.Field[Dims[CellDim, C2E.Local], float64]) -> gtx.Field[Dims[CellDim], float64]:
     return pseudo_lap(pseudo_lap(cells, edge_weights), edge_weights)
 ```