Compound types#
Overview#
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— lightweight container of tensors and primitives; can contain ndarrays@qd.data_oriented— enables object-oriented programming with Quadrants:@qd.kernelcan decorate instance methods, with tensors stored onselfinstead of being passed in as kernel arguments or held as global fields@qd.dataclass— for structures that are embedded into the kernel, and don’t contain ndarrays
property |
|
|
|
|---|---|---|---|
Exists as a struct in the kernel |
no (members flattened to args) |
no (members flattened to args) |
yes (fixed memory layout) |
Can be used as tensor element type |
no |
no |
yes |
Members can be tensors (field, ndarray) |
yes |
yes |
no |
|
no |
yes |
no |
|
no |
yes |
yes |
Member declaration |
type-annotated class fields |
live attributes (no annotations) |
type-annotated class fields |
Kernel-arg annotation |
|
|
|
Primitive members are |
runtime args, or compile-time constants when annotated |
compile-time constants by default (opt out) |
fields of the in-kernel struct |
⚠️ Deprecation:
@dataclasses.dataclassinstance passed viaqd.Template. Passing a@dataclasses.dataclassinstance into aqd.Template-annotated kernel parameter is not supported and emits aDeprecationWarningat compile time. In a future release it will become an error.
See Nesting compatibility below for a per-container × per-member-type breakdown, including the constraints on the outer kernel-arg annotation and ndarray reassignment.
How to choose a compound type?#
It’s of course very subjective, but some guidelines you could consider:
if you are trying to write a python class that runs on the GPU => use a
@qd.data_orientedif you are trying to write typed dataclasses, for passing data around between the
@data_orientedclasses, and between methods of the same@data_orientedclass => use@dataclasses.dataclassesif you are writing a static configuration object - a bag of flags and sizes fixed at setup time, used to specialize kernels => use a frozen
@dataclasses.dataclasswithFinal[T]fields@qd.dataclassis used to create structured element types for field tensors. We also use it to create the Cholesky tiles.
dataclasses.dataclass#
dataclasses.dataclass allows you to create structs containing:
ndarrays
fields
primitive types
These structs:
can be passed into kernels (
@qd.kernel) and sub-functions (@qd.func)can be combined with other parameters in the function signature
do not affect runtime performance compared to passing elements directly as parameters
can be nested (a dataclass can contain other dataclasses)
The members are read-only. However, ndarrays and fields are stored as references (pointers), so the contents of the ndarrays and fields can be freely mutated by the kernels and @qd.funcs.
Example#
import quadrants as qd
from dataclasses import dataclass
qd.init(arch=qd.gpu)
@dataclass
class MyStruct:
a: qd.types.NDArray[qd.i32, 1]
b: qd.types.NDArray[qd.i32, 1]
a = qd.ndarray(qd.i32, shape=(55,))
b = qd.ndarray(qd.i32, shape=(57,))
@qd.kernel
def k1(my_struct: MyStruct) -> None:
my_struct.a[35] += 3
my_struct.b[37] += 5
my_struct = MyStruct(a=a, b=b)
k1(my_struct)
print("my_struct.a[35]", my_struct.a[35])
print("my_struct.b[37]", my_struct.b[37])
Output:
my_struct.a[35] 3
my_struct.b[37] 5
Nesting#
Dataclasses can contain other dataclasses:
@dataclass
class Inner:
x: qd.types.NDArray[qd.f32, 1]
@dataclass
class Outer:
inner: Inner
y: qd.types.NDArray[qd.f32, 1]
@qd.kernel
def k2(s: Outer) -> None:
s.inner.x[0] = 1.0
s.y[0] = 2.0
Passing nested sub-structs to a qd.func#
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:
@dataclass
class Inner:
x: qd.types.NDArray[qd.f32, 1]
@dataclass
class Outer:
inner: Inner
y: qd.types.NDArray[qd.f32, 1]
@qd.func
def touch_inner(inner: Inner) -> None:
inner.x[0] += 1.0
@qd.func
def touch_outer(s: Outer) -> None:
s.y[0] += 10.0
touch_inner(s.inner) # call site inside a qd.func body
@qd.kernel
def k(s: Outer) -> None:
touch_outer(s) # whole-struct call
touch_inner(s.inner) # sub-struct call
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)andf(inner=s.inner))call sites both directly inside
@qd.kernelbodies and inside other@qd.funcbodiespruning 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.
Frozen vs non-frozen#
A dataclasses.dataclass may be either non-frozen (the default) or frozen (@dataclass(frozen=True)). Both work as kernel arguments, but kernel launch is faster with frozen=True (because it enables some optimizations that would otherwise not be possible). Recommend frozen=True unless you specifically need to rebind members after construction. Note that rebinding members after construction contradicts certain best practices; for example, it is typically incompatible with type linters such as pyright and mypy.
Subclasses#
You are able to pass any subclass of the annotated dataclass type. This is especially useful for cases where only a subset of your class is valid in quadrants, as below:
@dataclass(frozen=True) # the base need not be frozen for subclassing to work
class Position:
x: qd.types.NDArray[qd.math.vec3, 1]
# A dataclass subclass of a frozen base must itself be frozen (Python forbids inheriting a non-frozen dataclass from a
# frozen one). A plain, non-dataclass subclass also works and needs no decorator.
@dataclass(frozen=True)
class Particle(Position):
name: list[str]
@qd.kernel
def update_positions(pos: Position) -> None:
for i in range(pos.x.shape[0]):
update_position(pos.x, i)
Note: when the annotated type is frozen, the flattened launch arguments are cached on the instance on first use. Do not rebind a field after the first launch; the kernel would keep using the stale cached value. Frozen dataclasses forbid rebinding by design, so this only affects code that deliberately bypasses that (e.g. object.__setattr__, or an unsafe_hash=True non-frozen dataclass).
Compile-time constant fields: typing.Final#
By default a primitive dataclass field becomes a runtime kernel argument: you can change its value between launches without recompiling, and the kernel reads the fresh value each launch. That also means the value is not known at compile time, so it cannot be used inside qd.static(...), as a static loop bound, or to eliminate a branch at compile time.
Annotate the field with typing.Final[T] to make it a compile-time constant instead. The value is baked into the compiled kernel:
import dataclasses
from typing import Final
@dataclasses.dataclass(frozen=True)
class SimConfig:
enable_gravity: Final[bool] # compile-time constant
n_substeps: Final[int] # compile-time constant
gain: int # ordinary runtime kernel argument
@qd.kernel
def integrate(config: SimConfig, positions: qd.types.NDArray[qd.i32, 1]):
# Legal because n_substeps is Final: qd.static requires a compile-time constant.
steps = qd.static(config.n_substeps)
for i in positions:
# This branch is resolved at compile time; the untaken side is not compiled at all.
if qd.static(config.enable_gravity):
positions[i] -= 1
positions[i] += config.gain * steps # gain read at runtime
integrate(SimConfig(enable_gravity=True, n_substeps=4, gain=3), positions)
This is the recommended pattern for static configuration objects - bags of flags and sizes that are fixed once at setup time and used to specialize kernels. It replaces the older approach of declaring such a config as @qd.data_oriented and passing it via qd.Template purely to obtain compile-time member reads.
Semantics of a Final[T] field:
Baked into the kernel.
config.fieldinside a kernel body (or inside a@qd.funccalled from one) resolves at compile time to the field’s actual Python value.Each distinct value compiles a separate kernel. Quadrants decides which compiled kernel to reuse by looking at the field’s value.
Not a kernel argument. A
Finalfield occupies no kernel argument slot and costs nothing at launch.Mixing is fine. Final and ordinary fields coexist in the same dataclass, at any nesting depth.
Restrictions:
Every class on the path to the field must be frozen (
frozen=True, orunsafe_hash=Trueif you must keep one mutable). A baked value must not be reassignable, so a non-frozen@dataclassis rejected both when it declares aFinalfield itself and when it merely holds another dataclass that has one - rebinding that inner object would change a baked value without recompiling.unsafe_hash=Truedoes not snapshot the instance: aFinal-bearing instance is re-read on every launch (Quadrants deliberately skips its per-instance caches so theFinal-value checks re-run), so a later reassignment is not ignored - changing aFinalfield selects a different compiled kernel, and changing an ordinary field supplies its new runtime value. Because each distinctFinalvalue compiles a separate kernel, treatFinalfields as fixed and reserveunsafe_hashfor mutating the ordinary ones.Tmust be a value Quadrants can bake as a literal:bool,int,float,str, or a member of anenum.Enumsubclass. Arrays,qd.dataclassstructs,qd.Tensor, nested dataclasses and arbitrary objects are rejected. To make a nested dataclass’s leaves compile-time, mark those leaf fieldsFinalrather than the nested field itself.An
enumbaked this way must be plain: state-free, behavior-free and scalar-valued. Quadrants keys a baked member on its class plus its name/value, so it rejects (at first launch) any enum that carries something a kernel could additionally observe but the key cannot capture: per-member state (attributes set in__init__, or a populated__slots__slot - e.g.config.mode.unit), user-defined methods, properties, class variables or overridden operators on the class (e.g.config.mode.label, a custom==), or a membervaluethat is not itself a bakeable scalar. PlainEnum/IntEnum/IntFlagmembers (including composites) are fine. A locally-defined or otherwise dynamically-recreated enum class is still accepted, but it is distinguished only within a single process, so it will not reuse another process’s on-disk cached kernel; define enums at module scope if you want that cross-process reuse.Final[float]requires theraise_on_templated_floatsqd.initoption to be off (its default). That option makes Quadrants rejectfloatvalues that drive kernel specialization, and aFinal[float]field does exactly that, so it is rejected when the option is enabled.String annotations are rejected.
from __future__ import annotationsleaves the annotation unresolved, so Quadrants cannot see theFinaland would silently treat the field as an ordinary runtime argument. This raises an error rather than quietly doing the wrong thing.
If you want the opposite trade-off for a @qd.data_oriented class (primitive members that are runtime rather than baked), see Runtime primitives: template_primitives=False.
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. Inside a kernel, tensor members are read-write through indexing (s.x[i] = ...), but the member binding itself (s.x = other_tensor) cannot be reassigned from inside a kernel.
Reserved field names (_qd_ prefix). Field names beginning with _qd_ are reserved for Quadrants’ internal state; defining a dataclass field with such a name is unsupported and gives undefined behavior.
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.
@qd.data_oriented
class Simulation:
def __init__(self, n):
self.x = qd.field(qd.f32, shape=n)
@qd.kernel
def step(self):
for i in self.x:
self.x[i] += 1.0
sim = Simulation(100)
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).
stable_members#
Recommended for any @qd.data_oriented class whose ndarray members are allocated once (typically in __init__) and not subsequently rebound - the common case. Decorate with stable_members=True:
@qd.data_oriented(stable_members=True)
class Simulation:
def __init__(self, n):
self.x = qd.ndarray(qd.f32, shape=(n,))
self.v = qd.ndarray(qd.f32, shape=(n,))
# ... more ndarray / field / primitive members
This skips a per-call walk that Quadrants otherwise runs to detect ndarray member rebinding between kernel launches. The walk is O(number of ndarray members) per kernel call, so the savings scale with the container’s size.
Microbenchmark on an RTX PRO 6000 Blackwell with a container holding 30 qd.ndarray members across two nesting levels, calling a trivial kernel that takes the container as a qd.template() arg:
Per-launch Python overhead |
|
|---|---|
|
18.5 us/call |
|
13.5 us/call |
-5 us/call (-28%) |
Trade-off: with stable_members=True, reassigning an ndarray member on an instance is undefined behavior - the previously compiled kernel will be reused even if the new ndarray has a different dtype, ndim, or layout, silently bit-reinterpreting the new array’s storage. Set it only on classes whose ndarray members are allocated once (typically in __init__) and never rebound. See Reassigning ndarray members below for the supported alternative.
stable_members concerns only ndarray members and is independent of template_primitives, which concerns primitive members; the two can be combined, e.g. @qd.data_oriented(stable_members=True, template_primitives=False).
Setting the class attribute _qd_stable_members = True has the same effect as the decorator argument, for the case where you need to opt a class in without controlling its decoration. It is read from the class itself, not from its bases, so each subclass that wants the hint has to set it (or be decorated) too.
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 compiled kernel.
@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
Runtime primitives: template_primitives=False#
If you want to change a primitive member’s value from Python between launches without triggering a recompile, decorate the class with @qd.data_oriented(template_primitives=False). Every primitive member the kernel actually accesses (int, float, bool, including those reached through nested dataclasses.dataclass / @qd.data_oriented members) is then lifted into a runtime scalar kernel argument and read fresh on every launch, instead of being compiled into the kernel as a constant. The member stays read-only inside the kernel (it is a kernel argument, not a writable variable): you mutate it in Python, and the kernel sees the new value on the next launch.
@qd.data_oriented(template_primitives=False)
class Simulation:
def __init__(self):
self.n = 100
self.x = qd.ndarray(qd.f32, shape=(256,))
@qd.kernel
def step(self):
for i in range(self.n):
self.x[i] += 1.0
sim = Simulation()
sim.step() # compiles once; n is a runtime kernel argument
sim.n = 200 # no recompilation
sim.step() # the new value of n takes effect immediately
Notes and restrictions:
dtype follows the runtime defaults - an
int/boolmember becomes the default integer type (qd.i32, unless you override the runtime default integer type inqd.init()) and afloatmember becomes the default float type (qd.f32, unless you override the runtime default float type inqd.init()). A member whose value falls outside the default integer range will overflow where a baked literal would not; if you need exact wide-integer constants, keep the default (baked) behaviour.dtype is fixed at first compile: the kernel-argument dtype is chosen from the member’s Python type the first time the kernel compiles, and is not re-specialised if you later reassign the member to a different type. The live value is coerced to that dtype on every launch, so binding a
floatto a member that was first seen as aninttruncates it (exactly as passing afloatto anint-typed kernel argument would). Keep a lifted member’s type stable across launches; if the type itself must vary, use the default (baked) behaviour, which re-specialises the kernel per type.Pruning: only the primitives the kernel actually reads are turned into kernel arguments, so a class with many primitive members does not blow up the kernel argument count.
qd.staticis an error: a lifted primitive cannot be used insideqd.static(...), because that context requires a compile-time constant. Doing so raisesQuadrantsSyntaxError. Use the defaulttemplate_primitives=Truefor values that must be baked (e.g. unrolled loop bounds).No re-specialisation on value change: because the value is a runtime argument, mutating it never triggers a recompile. Distinct instances of the class still compile separately (the kernel is keyed per instance), exactly as with the default.
This is opt-in: the default
@qd.data_orientedcontinues to bake primitive members as shown above.
Tensor members#
@qd.data_oriented classes may hold tensor members of any backend: qd.field, qd.ndarray, or qd.Tensor.
@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 Appendix — 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 by default, so each distinct primitive value compiles a new specialized kernel - unless the class is decorated @qd.data_oriented(template_primitives=False), in which case the accessed primitives become runtime kernel arguments (see Runtime primitives above).
qd.dataclass / qd.types.struct#
Unlike @qd.data_oriented and @dataclasses.dataclass, @qd.dataclass creates a struct type that is available inside the kernels themselves. The other two compound types only exist on the Python side, before compilation, and don’t appear in compiled kernel code at all.
@qd.dataclass members can only be:
primitives (
qd.f32,qd.i32,qd.bool, etc.)fixed-size vectors (
qd.types.vector(N, dtype))fixed-size matrices (
qd.types.matrix(M, N, dtype))
A qd.dataclass is analogous to a C struct. All members, including the fixed-size vectors and fixed-size matrices, are laid out within the struct itself. They are not pointers to tensors allocated elsewhere. Changing the size of a vector or matrix changes thus the size of the qd.dataclass.
A consequence of a qd.dataclass containing all members within itself, rather than being pointers, is that a qd.dataclass cannot contain qd.field or qd.ndarray members.
A @qd.dataclass can be turned into a tensor of structs (e.g. MyStruct.field(shape=(N,))) with two possible memory layouts:
Struct-of-arrays (SoA) (
qd.Layout.SOA): extrudes each member of the struct into its own tensor of lengthN.Array-of-structs (AoS) (
qd.Layout.AOS): the storage is an array ofNstruct cells laid out contiguously in memory. AoS is only available withqd.fieldbacking.
Note that although a @qd.dataclass’s members can’t themselves be tensors, allocating one in SoA layout (MyStruct.field(shape=(N,), layout=qd.Layout.SOA)) extrudes each member into its own length-N tensor — so the resulting collection effectively behaves like a struct of parallel tensors, even though the @qd.dataclass type itself doesn’t have tensor-typed members.
@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 `@qd.dataclass`. `@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).
@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 a @qd.dataclass without a class body.
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 summarizes which member types are allowed inside which container type. “yes” means the member is handled correctly when the container is passed to a kernel; “no” means the member is ignored or the combination raises an error.
Container ↓ / Member → |
|
|
primitive |
|
|
|
|---|---|---|---|---|---|---|
|
yes |
yes |
yes |
yes |
deprecated |
yes |
|
yes |
yes |
yes |
yes |
yes |
yes |
|
no |
yes |
yes |
no |
no |
yes |
Outer kernel-arg annotation#
The outermost annotation you put on the kernel parameter should match the parameter type as follows:
Kernel parameter compound type |
Annotation |
|---|---|
|
|
|
|
Reassigning ndarray members#
For @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 specialized kernel is compiled and cached for the new shape; subsequent launches with the original shape continue to use the original cached kernel. (For @dataclasses.dataclass containers — passed via the dataclass-type annotation — the member binding follows the standard dataclass mutability rules: frozen dataclasses can’t rebind, non-frozen ones can, and a rebind triggers a fresh kernel arg setup on the next launch.)
This support is only available on @qd.data_oriented classes without the stable_members=True opt-in. Setting stable_members=True is a promise that ndarray members on instances of the class are never reassigned; if you break that promise the previously compiled kernel is silently reused against the new ndarray.
Restrictions#
@qd.dataclasscannot containqd.ndarrayorqd.fieldmembers. See the@qd.dataclasssection above for the full list of allowed member types. (The function-form factoryqd.types.struct(...)has the same restrictions.)A typed-dataclass kernel-arg annotation cannot have a
@qd.data_orientedmember type — errors clearly at compile timeDeclare all ndarray members on a
@qd.data_orientedclass in__init__.Deleting an ndarray attribute that was present on an
@qd.data_orientedinstance’s first launch raisesAttributeErroron the next launch on that instance.Adding a new ndarray attribute after first launch on a given
@qd.data_orientedinstance will cause incorrect undefined behavior.