[Types] Allow dataclass subclasses as kernel params

[D0liphin]

Issue: #862

Brief Summary

Allow kernels that annotate with a certain dataclass T to also allow passing subclasses of T. Hopefully little to no performance degradation.

Walkthrough

The difference is easily explained by the following:

@dataclass
class QdDomain:
    world: qd.types.NDArray[qd.i32, 3]
    
@dataclass
class Domain(QdDomain):
    world_name: str

@qd.kernel
def do_some_cool_simulation(domain: QdDomain) -> None: ...

do_some_cool_simulation(Domain(...))  # previously: not allowed

Typing-wise, this is allowed by all type checkers, so wouldn't be surprising to a user. One argument for its addition is that it is the 'correct' thing to do.

What is still in question is whether or not this is a pattern that should be encouraged from a language design POV. IMO, this is better than any of the alternatives...

Alternatives

• Encourage the user to contain a subset of their type as _quadrants_safe: QdDomain. Cons: The user would have to access through this attribute, or create a @property for each attribute when dealing with the parent.
• Encourage the user to make an explicit converter .to_quadrants_safe(). Cons: This requires a new dataclass and method.
• Allow dataclasses with invalid member types, but throw a runtime error at kernel compile. This has the downside that we cannot catch this error at typecheck time in python.

Implementation Considerations

The main fix here should be clear (adding the subclass check), and I add a comment why in the codebase (please see). Also some housekeeping swapping the runtime error to a type error to follow the general pattern. But there are some tricky considerations with the caching that we should check. There are 4 caches that we care about:

1. (global) _frozen_dc_plans: for a given annotated type + kernel, precomputes launch argument layout
2. (instance) _qd_dc_unwrapped: map of field name to unwrapped value, so launch skips the getattr +
   _unwrap() on every field, every call.
3. (instance) _qd_all_field: a bool recording whether all active fields unwrap to Fields; if so we skip
   _recursive_set_args entirely.
4. (instance) arg._key — the template specialization key

We solve the frozen dc plans for free:

_frozen_dc_plans is fine, because the cache is keyed on the annotated type. Subclassing only changes which provided types are accepted, but doesn't change the annotation type. It also doesn't change what attributes are validly accessible on the provided argument, and therefore a cache for some class is also valid for any subclass of it! "But what if you change the type of an attribute?". This still works. The fields of the annotated type are used, and you would get a type failure at runtime. however, it is illegal anyway to do this from a typechecker POV.

We solve all the instance things by replacing the cache with a dict[annotated-type, cached-value]. The same instance can now be passed under different views. This obviously has the drawback that now we have to do an extra dict lookup (though I don't think that is OOM relevant here), but it's clearly correct and avoids thrashing.

Testing

Added whitebox cache tests to match the pattern of how the repository currently tests these caches. Updated existing cache tests in the most minor way I could.

Added some tests to ensure that subclassing works. Also added a test I thought to be missing that just checks if the dataclass type is correct at all.

[D0liphin]

Just marking this as a draft while I self-review and wait for CI to go through.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5b2c782450

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 37f7045bad

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Disable frozen caches for mutable subclasses

When the annotated dataclass is frozen but the accepted value is a plain subclass that restores mutation (for example, by overriding __setattr__ to call object.__setattr__), this branch accepts it while the frozen fast paths still derive immutability from the annotated base type. Both _extract_arg and _get_frozen_dc_unwrapped then permanently cache the original field binding, so after sub.x is rebound, subsequent kernel launches silently continue operating on the old array. Since non-dataclass subclasses are explicitly supported, either verify the runtime class is safely frozen or bypass the frozen caches for subclass instances.

Useful? React with 👍 / 👎.

[D0liphin]

Thanks robot, will check