[DataOriented] Fix ndarrays on data oriented

docs/source/user_guide/compound_types.md

@@ -5,21 +5,19 @@
 It can be useful to combine multiple ndarrays or fields together into a single struct-like object that can be passed into kernels, and into @qd.func's.
 
 The following compound types are available:
-- `dataclasses.dataclass` — **recommended**
-- `@qd.data_oriented` — for classes that define `@qd.kernel` methods, cannot contain ndarrays
-- `@qd.struct` / `@qd.dataclass` — legacy, field-only
+- `@dataclasses.dataclass` — lightweight container of tensors and primitives; can contain ndarrays
+- `@qd.data_oriented` — for creating objects with `self` that define `@qd.kernel` methods
+- `@qd.dataclass` — for structures that are embedded into the kernel, and don't contain ndarrays
 
-| type                               | can be passed to qd.kernel? | can be passed to qd.func? | can contain ndarray? | can contain field? | can be nested? | supports differentiation? |
-|------------------------------------|:---------------------------:|:-------------------------:|:--------------------:|:------------------:|:--------------:|:-------------------------:|
-| `dataclasses.dataclass`            | yes                         | yes                       | yes                  | yes                | yes            | no [*1]                   |
-| `@qd.data_oriented`               | yes                         | yes                       | no                   | yes                | yes            | yes                       |
-| `@qd.struct`, `@qd.dataclass`     | yes                         | yes                       | no                   | yes                | yes            | yes                       |
+| property                            | `@dataclasses.dataclass`     | `@qd.data_oriented`              | `@qd.dataclass`                     |
+|-------------------------------------|:----------------------------:|:--------------------------------:|:-----------------------------------:|
+| Kernel-side representation          | none (flattened away)        | none (flattened away)            | real type with fixed memory layout  |
+| Can be tensor element type          | no                           | no                               | yes                                 |
+| Can hold ndarrays                   | yes                          | yes                              | no                                  |
+| `@qd.kernel` methods on `self`      | no                           | yes                              | no                                  |
+| Member declaration                  | type-annotated class fields  | live attributes (no annotations) | type-annotated class fields         |
 
-## Recommendation
-
-**Use `dataclasses.dataclass` for new code.** It supports both fields and ndarrays, can be nested, and uses standard Python — no Quadrants-specific decorator needed.
-
-The other compound types exist for historical reasons.
+See [Nesting compatibility](#nesting-compatibility) below for a per-container × per-member-type breakdown, including the constraints on the outer kernel-arg annotation and ndarray reassignment.
 
 ## dataclasses.dataclass
 
@@ -91,7 +89,7 @@ def k2(s: Outer) -> None:
 
 ### Passing nested sub-structs to a `qd.func`
 
-You can pass either a whole nested-dataclass argument or one of its sub-struct fields to a `qd.func`. The callee declares the sub-struct's type as the parameter annotation; the caller writes the attribute access at the call site:
+You can pass either a whole nested-dataclass argument or one of its sub-struct members to a `qd.func`. The callee declares the sub-struct's type as the parameter annotation; the caller writes the attribute access at the call site:
 
 ```python
 @dataclass
@@ -123,10 +121,14 @@ Sub-struct passing supports:
 - arbitrary nesting depth (`f(s.a.b.c)` where each level is a dataclass)
 - positional and keyword call sites (`f(s.inner)` and `f(inner=s.inner)`)
 - call sites both directly inside `@qd.kernel` bodies and inside other `@qd.func` bodies
-- pruning of the sub-struct's leaf fields that the callee never reads
+- pruning of the sub-struct's leaf members that the callee never reads
 
 Note: assigning a sub-struct to a local variable and then passing it (`t = s.inner; touch_inner(t)`) is **not** supported. Pass the attribute access directly at the call site.
 
+### Under the hood
+
+A `dataclasses.dataclass` is a Python-only container. The compiler reads it at compile time and flattens its members into individual kernel parameters — the container itself has no memory layout and doesn't exist on the kernel side. That's why members are read-only: the values are captured at compile time and re-assigning them afterwards has no effect on running kernels.
+
 ## qd.data_oriented
 
 `@qd.data_oriented` is designed for classes that define `@qd.kernel` methods as class members. It wraps these methods to correctly bind `self` during kernel compilation.
@@ -148,14 +150,146 @@ sim.step()
 
 `@qd.data_oriented` objects can also be passed as `qd.Template` parameters to kernels defined outside the class, and they support nesting (one `@qd.data_oriented` struct containing another).
 
-## qd.struct / qd.dataclass
+### Primitive members
+
+Primitive members on `self` (e.g. `int`, `float`, `bool`, `enum.Enum`) are supported, but they are treated as **template values**: each distinct primitive value across instances triggers a new kernel compilation, with the value baked into the kernel IR.
+
+```python
+@qd.data_oriented
+class Simulation:
+    def __init__(self, n):
+        self.n = n
+        self.x = qd.ndarray(qd.f32, shape=(n,))
+
+    @qd.kernel
+    def step(self):
+        for i in range(self.n):
+            self.x[i] += 1.0
+
+Simulation(100).step()   # compiles kernel #1 with n=100 baked in
+Simulation(200).step()   # compiles kernel #2 with n=200 baked in
+```
+
+### Tensor members
 
-`@qd.struct` (and its alias `@qd.dataclass`) is a Quadrants-native struct type. It can only contain fields and primitive types, not ndarrays.
+`@qd.data_oriented` classes may hold tensor members of any backend: `qd.field`, `qd.ndarray`, or `qd.Tensor`.
+
+```python
+@qd.data_oriented
+class State:
+    def __init__(self, n):
+        self.n = n
+        self.a = qd.field(qd.f32, shape=n)
+        self.b = qd.ndarray(qd.f32, shape=(n,))
+        self.c = qd.tensor(qd.f32, shape=(n,))
+
+    @qd.kernel
+    def step(self):
+        for i in range(self.n):
+            self.a[i] += 1.0
+            self.b[i] += 1.0
+            self.c[i] += 1.0
+
+state = State(100)
+state.step()
+```
+
+### Fastcache
+
+`@qd.kernel(fastcache=True)` is supported on methods of `@qd.data_oriented` classes, but is disabled for fields; see [Advanced — compound-type cache keying](fastcache.md#compound-type-cache-keying) for more information.
+
+### Under the hood
+
+Like `dataclasses.dataclass`, a `@qd.data_oriented` object is Python-only — the compiler flattens it into individual kernel parameters and the object itself has no kernel-side representation. Unlike `dataclasses.dataclass` it needs no member annotations: the compiler reads the live instance's attributes directly. Primitive members are baked into the kernel as constants, so each distinct primitive value compiles a new specialised kernel.
+
+## qd.dataclass / qd.types.struct
+
+Unlike `@qd.data_oriented` and `@dataclasses.dataclass`, `@qd.dataclass` creates a struct that is available within the kernels themselves. The former types are only used for structure on the python side, before compilation. `@qd.dataclass` can be used as the element type of fields. One key downside of `@qd.dataclass` is that they can only be used with fields and primitives, not with ndarray. This is because tensors are embedded in the struct by value, not as a reference pointer.
 
 ```python
 @qd.dataclass
 class Particle:
     pos: qd.types.vector(3, qd.f32)
     vel: qd.types.vector(3, qd.f32)
     mass: qd.f32
+
+# AOS layout: each element of `particles` is a (pos, vel, mass) cell contiguous in memory.
+# Only possible because Particle is a StructType — `@qd.data_oriented` and
+# `dataclasses.dataclass` containers can't be the element type of a tensor.
+particles = Particle.field(shape=(N,), layout=qd.Layout.AOS)
 ```
+
+Methods can be added to a `@qd.dataclass` and may be decorated with `@qd.func` so they can be called from kernels via `instance.method(...)` syntax (the call is inlined at compile time, like any other `@qd.func`).
+
+```python
+@qd.dataclass
+class Particle:
+    pos: qd.types.vector(3, qd.f32)
+    vel: qd.types.vector(3, qd.f32)
+    mass: qd.f32
+
+    @qd.func
+    def kinetic_energy(self):
+        return 0.5 * self.mass * self.vel.dot(self.vel)
+
+particles = Particle.field(shape=(N,))
+
+@qd.kernel
+def total_ke() -> qd.f32:
+    total = 0.0
+    for i in range(N):
+        total += particles[i].kinetic_energy()
+    return total
+```
+
+`qd.types.struct(name1=type1, ...)` is the function-form equivalent of `@qd.dataclass`: it builds the same `StructType` without a class body.
+
+```python
+vec3 = qd.types.vector(3, qd.f32)
+Particle = qd.types.struct(pos=vec3, vel=vec3, mass=qd.f32)
+particles = Particle.field(shape=(N,))
+```
+
+### Under the hood
+
+Unlike the other two compound types, `@qd.dataclass` is a real kernel-side type with a fixed memory layout. Each instance is laid out contiguously in memory, members are stored by value, and a tensor of the struct can be allocated (`Particle.field(...)`). Storing by value is also why ndarrays can't be members — ndarrays are heap-allocated buffers with dynamic shape and don't fit into a fixed-size cell.
+
+## Nesting compatibility
+
+This table summarises which member types are allowed inside which container type. "yes" means the member is walked correctly when the container is passed to a kernel; "no" means the member is ignored or the combination raises an error.
+
+| Container ↓     /     Member → | `qd.ndarray` | `qd.field` | primitive | `dataclasses.dataclass` | `@qd.data_oriented` | `@qd.dataclass` |
+|---|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dataclasses.dataclass`         | yes | yes | yes | yes | yes [\*1] | yes |
+| `@qd.data_oriented`             | yes | yes | yes | yes | yes      | yes |
+| `@qd.dataclass`                 | no  | yes | yes | no  | no       | yes |
+
+[\*1] A `dataclasses.dataclass` may *hold* a `@qd.data_oriented` member, but the **outer kernel-arg annotation** must be `qd.template()`, not the dataclass type itself. Passing a typed-dataclass kernel arg (`def k(s: Outer)`) whose member type is a `@qd.data_oriented` class raises a clear `QuadrantsSyntaxError` at compile time pointing you to `qd.template()`. The reason: typed-dataclass kernel args are flattened from annotations, but `@qd.data_oriented` carries no per-member annotations — its members are walked from the live instance, which only happens on the template path.
+
+### Outer kernel-arg annotation
+
+The outermost annotation you put on the kernel parameter determines how the container is walked:
+
+| Annotation | Kernel-arg walker | Notes |
+|---|---|---|
+| `qd.types.NDArray[...]`           | ndarray slot                                       | leaf-level only |
+| `MyDataclass` (dataclass type)    | per-member flatten using annotations               | needs every member to have a quadrants-typed annotation |
+| `qd.template()`                   | value-driven walk of `vars(self)` / dataclass members | supports the full nesting matrix above |
+
+Two practical consequences:
+
+- **Containers with `@qd.data_oriented` anywhere in the tree** must be passed via `qd.template()` (or be the `self` of a `@qd.kernel` method on a `@qd.data_oriented` class). Using a typed-dataclass annotation on the outermost arg errors.
+- **A non-frozen `dataclasses.dataclass`** can be passed via the typed-dataclass annotation, but cannot be the outer `qd.template()` arg — `qd.template()` uses the instance as a dict key inside the template-mapper and a non-frozen dataclass has `__hash__ = None`. Add `frozen=True` if you need to pass it as `qd.template()` (for example, when it holds `@qd.data_oriented` children).
+
+### Reassigning ndarray members
+
+For both `dataclasses.dataclass` and `@qd.data_oriented` containers passed via `qd.template()`, reassigning an ndarray member between kernel launches is supported, including changes to `dtype`, `ndim`, or layout. A new specialised kernel is compiled and cached for the new shape; subsequent launches with the original shape continue to use the original cached kernel.
+
+### Restrictions
+
+- **`@qd.dataclass` (the Quadrants `StructType` decorator) cannot contain ndarrays.** This is a legacy field-only type. Use `dataclasses.dataclass` or `@qd.data_oriented` instead. (The function-form factory `qd.types.struct(...)` produces the same `StructType` and has the same restrictions.)
+- **A typed-dataclass kernel-arg annotation cannot have a `@qd.data_oriented` member type** (see [\*1] above) — errors clearly at compile time.
+- **An outer `qd.template()` arg of dataclass type must be `frozen=True`** — non-frozen dataclasses are unhashable and the template-mapper cannot use them as cache keys.
+- **Declare all ndarray members on a `@qd.data_oriented` class in `__init__`.** The template-mapper caches the set of ndarray-attribute paths reachable from the first instance walked, per class. Adding *new* ndarray attributes on later instances of the same class is safe — the per-instance weakref in the spec key disambiguates them, and the compile-time walker registers all reachable ndarrays. But:
+  - **Deleting an ndarray attribute** that was present on the first launch raises `AttributeError` on the next launch (the cached path still tries to `getattr` the missing attribute).
+  - **Reassigning a post-first-walk ndarray attribute** (one not present on the first instance walked, then added later and re-assigned) to one with a different `dtype` / `ndim` is *not* detected by the in-memory invalidation tracker. The stale compiled kernel is silently reused, leading to bit-reinterpretation of the new array's storage.

docs/source/user_guide/fastcache.md

@@ -50,43 +50,13 @@ qd.init(arch=qd.gpu)
 # qd.init(arch=qd.gpu, print_non_pure=True)
 ```
 
-## Dataclass fields with cached values
-
-By default, for `dataclasses.dataclass` parameters, fastcache only includes the *types* of each field in the cache key, not their values. This is fine for fields like ndarrays, where the compiled kernel doesn't depend on the actual data, only the dtype and dimensionality.
-
-However, some dataclass fields hold configuration values that get baked into the compiled kernel — typically values used with `qd.static()`, such as loop bounds or feature flags:
-
-```python
-for i in qd.static(range(config.num_layers)):
-    ...
-```
-
-Here the value of `num_layers` is compiled into the kernel. Concretely the loop will be unrolled, at compile time. If `num_layers` changes, a different kernel must be compiled.
-
-Mark such fields with `add_value_to_cache_key` so their values are included in the cache key:
-
-```python
-import dataclasses
-from quadrants.lang._fast_caching import FIELD_METADATA_CACHE_VALUE
-
-@dataclasses.dataclass
-class SimConfig:
-    num_envs: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
-    dt: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
-    use_gravity: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
-```
-
-With this annotation, changing `num_envs` from 100 to 200 produces a different cache key so the correct compiled kernel is looked up (or compiled if not yet cached). Without it, the wrong kernel could be loaded.
-
-Note: `@qd.data_oriented` objects and `qd.Template` parameters already include primitive values in the cache key automatically — this annotation is only needed for `dataclasses.dataclass` fields.
-
 ## Constraints
 
 A kernel is eligible for fastcache only if all of the following hold:
 
 ### 1. All data flows through parameters
 
-The kernel must receive every piece of data it operates on as an explicit parameter. It must **not** capture variables from the enclosing Python scope (closures over fields, ndarrays, or mutable globals). This is the core "purity" constraint — the compiled kernel's behavior must be fully determined by its arguments.
+The kernel must receive every piece of data it operates on as an explicit parameter. It must **not** capture variables from the enclosing Python scope (closures over ndarrays, mutable globals, or any other external state). This is the core "purity" constraint — the compiled kernel's behavior must be fully determined by its arguments.
 
 ```python
 a = qd.ndarray(qd.f32, (10,))
@@ -125,8 +95,8 @@ Fastcache supports the following parameter types:
 | `qd.types.NDArray` (scalar, vector, matrix) | Yes | dtype, ndim, layout |
 | `torch.Tensor` | Yes | dtype, ndim |
 | `numpy.ndarray` | Yes | dtype, ndim |
-| `dataclasses.dataclass` | Yes | field types recursively; field values if annotated with `add_value_to_cache_key` (see [above](#dataclass-fields-with-cached-values)) |
-| `@qd.data_oriented` objects | Yes | member types and primitive member values recursively |
+| `dataclasses.dataclass` | Yes | member types recursively; member values if annotated with `FIELD_METADATA_CACHE_VALUE` (see [Advanced — compound-type cache keying](#compound-type-cache-keying)) |
+| `@qd.data_oriented` objects | Yes | member types recursively; primitive member types and values baked into kernel (see [Advanced — compound-type cache keying](#compound-type-cache-keying)) |
 | `qd.Template` primitives (int, float, bool) | Yes | type and value (baked into kernel) |
 | Non-template primitives (int, float, bool) | Yes | type only |
 | `enum.Enum` | Yes | name and value |
@@ -172,3 +142,31 @@ print(obs.cache_stored)         # True if the compiled kernel was stored to cach
 ```
 
 On the first run you'll see `cache_stored=True` but `cache_loaded=False`. On the second run (after `qd.init`), `cache_loaded=True`.
+
+### Compound-type cache keying
+
+The args hasher walks compound-type kernel parameters recursively. For each leaf member it decides what (if anything) contributes to the cache key. The headline rules:
+
+**`@qd.data_oriented`:** the walker descends into `vars(obj)`. For each child:
+
+- `qd.ndarray` member — `(dtype, ndim, layout)` is included in the cache key. Element values are not.
+- Primitive (`int` / `float` / `bool` / `enum.Enum`) member — value is baked into the kernel (same semantics as a `qd.Template` primitive). Two instances of the same class with different primitive member values get different cache entries.
+- Nested `@qd.data_oriented` member — recurses.
+- Nested `dataclasses.dataclass` member — recurses (with the dataclass rules below).
+- `qd.field` member — fastcache is disabled for the entire kernel call. The kernel still runs via normal compilation; a warn-level log line is emitted.
+
+**`dataclasses.dataclass`:** the walker descends into the declared members. For each member, only the *type* is included in the cache key by default — **not** the value. To include a member's value, annotate it:
+
+```python
+import dataclasses
+from quadrants.lang._fast_caching import FIELD_METADATA_CACHE_VALUE
+
+@dataclasses.dataclass
+class SimConfig:
+    num_layers: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
+    dt: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
+```
+
+This is necessary whenever the compiled kernel depends on the member's *value* rather than just its type (for example, when the value is used as a loop bound that the compiler bakes into the generated code). Without the annotation, two `SimConfig` instances with different `num_layers` values would share a fastcache key, and the second instance would silently load a kernel compiled for the wrong value.
+
+Note the asymmetry: `@qd.data_oriented` primitive members are baked into the kernel automatically (same semantics as `qd.Template`); `dataclasses.dataclass` members contribute only their *type* to the cache key unless you opt in per-member.