[Perf] Skip _recursive_set_args for all-Field frozen dataclass structs

[hughperkins]

When a frozen dataclass struct has only qd.Field entries (no Ndarrays, no templates that consume launch-context slots), the entire struct contributes zero arguments to the kernel launch context.

Two changes eliminate the per-launch overhead for these structs:

1. _func_base.py: After computing the field plan, identify which entries are "non-field" (consume slots). Cache as _qd_non_field_plan on the struct instance. If empty, return (0, True) without recursion.

2. kernel.py: Before calling _recursive_set_args, check for a cached empty _qd_non_field_plan. If present, skip the call entirely, avoiding even function-call overhead.

On the anymal_uniform_kinematic benchmark (two robots, random control each step, field backend), this eliminates a -11% regression caused by struct-class annotations replacing qd.template() in kernel signatures.

Issue: #

Brief Summary

copilot:summary

Walkthrough

copilot:walkthrough

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4b2c4ad4ec

ℹ️ 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]

Scope cached non-field plans to kernel specialization

Do not reuse _qd_non_field_plan across all launches without a key: this cache is currently stored per struct instance and then consumed unconditionally, but the active field plan depends on used_py_dataclass_parameters and py_dataclass_basename. If the same frozen dataclass instance is first launched in a context where only qd.Field members are active (empty plan), a later launch that needs ndarray/non-field members will incorrectly skip argument emission (including the kernel.py fast-skip path), producing a wrong launch context. Please key this cache by specialization context (at least used-params/basename) or recompute per call.

Useful? React with 👍 / 👎.

[hughperkins]

@claude review

[claude]

Additional findings (outside current diff — PR may have been updated during review):

• 🔴 python/quadrants/lang/kernel.py:494-496 — The new kernel.py launch-loop fast-skip at kernel.py:494 reads the wrong attribute name: _qd_non_field_plan (singular), but _func_base.py only ever sets _qd_non_field_plans (plural — a per-instance dict keyed by id(plan)). getattr(val, "_qd_non_field_plan", None) always returns None, the if _nfp is not None and not _nfp: branch is unreachable, and the launch-loop fast-skip is dead code — silently defeating one of the two perf optimizations advertised in the PR description ("avoiding even function-call overhead" for all-Field structs). Note that simply renaming to _qd_non_field_plans is insufficient since that attribute is a dict (truthy whenever any plan has been populated); the fix needs either a per-plan_id lookup or a separate boolean fast-path attribute (e.g. _qd_non_field_empty) stored alongside the dict.

  Extended reasoning...

  What the bug is

  _func_base.py (lines 637/641 in the merged version) writes/reads the per-instance cache under the plural attribute _qd_non_field_plans — a dict keyed by id(plan) (the per-specialization fix from commit 4a8c552). kernel.py:494, however, reads the singular _qd_non_field_plan:

  _nfp = getattr(val, "_qd_non_field_plan", None)
  if _nfp is not None and not _nfp:
      continue

  A grep across the entire repo finds exactly three references to _qd_non_field*: the two plural reads/writes in _func_base.py, and this one singular read in kernel.py. The singular form is never written anywhere, so getattr always returns None, the predicate _nfp is not None and not _nfp is unreachable, and the fast-skip never fires.

  How this happened

  Commit 4b2c4ad ("Skip _recursive_set_args for all-Field frozen dataclass structs") added the kernel.py probe with the singular attribute name. Commit 4a8c552 ("Key non-field plan cache by specialization to avoid cross-kernel staleness") then renamed the _func_base.py cache to plural to fix the cross-kernel staleness flagged by reviewers — but did not update the matching kernel.py probe. The two ends of the optimization no longer agree on the attribute name.

  Step-by-step proof

  Take an all-Field frozen dataclass struct (the exact PR target — anymal_uniform_kinematic, two robots, field backend):

  1. Launch kernel for the first time with struct s. launch_kernel reaches the new probe at kernel.py:494: _nfp = getattr(s, "_qd_non_field_plan", None) → returns None (the attribute doesn't exist, never has).
  2. if _nfp is not None and not _nfp: → False (because _nfp is None). Fast-skip does not trigger.
  3. Falls through to self._recursive_set_args(...) — paying full call overhead.
  4. Inside _recursive_set_args, hits the frozen-dc branch (_func_base.py:625+), computes plan, allocates/looks up the _qd_non_field_plans plural dict on s, computes _nf_plan = (), stores it, then hits if not _nf_plan: return 0, True.
  5. Second launch repeats the same dance — the kernel.py-level fast-skip still sees no _qd_non_field_plan (singular), so it still does not fire. Every launch pays the cost of the call into _recursive_set_args, the dataclass-fields branch detection, the _get_frozen_dc_plan + _get_frozen_dc_unwrapped lookups, and the _qd_non_field_plans dict lookup before bottoming out at the inner early-return.

  Impact

  Correctness is preserved — the inner if not _nf_plan: return 0, True early-return inside _recursive_set_args still produces the right slot count (0). But the kernel.py change contributes nothing in the merged form; the PR description explicitly advertises this hop as "avoiding even function-call overhead" before _recursive_set_args, and that benefit is silently lost. The PR's own benchmark claim (eliminating the -11% regression on anymal_uniform_kinematic) was presumably measured against a version where the two attribute names matched.

  Why a simple rename is not enough

  Renaming the kernel.py probe to read _qd_non_field_plans (plural) would not work: that attribute is a dict keyed by id(plan), so it becomes truthy as soon as any plan entry is populated, and not _nfp would not select the all-Field empty-plan case. Two correct fixes:

  • Store an additional per-instance boolean fast-path attribute (e.g. _qd_non_field_empty: bool) alongside the dict, so the launch loop can check if getattr(val, "_qd_non_field_empty", False): continue in O(1) without per-launch dict logic. But this still needs to be keyed by plan_id (a different specialization may not be empty), so really this would have to be a set of plan_ids or a dict mapping plan_id -> bool.
  • Look up the dict by the active plan_id here in launch_kernel — but plan_id is computed inside _recursive_set_args from a plan that itself depends on used_py_dataclass_parameters (keyed by self.used_py_dataclass_parameters_by_key_enforcing[key]) and py_dataclass_basename. Doing the lookup at the kernel.py level requires hoisting the _get_frozen_dc_plan call (or its memoized result) out, which partially defeats the point of skipping the function call.

  The cleanest path is to attach the empty-flag to the entry stored in the module-level _frozen_dc_plans dict itself (which is already correctly keyed by specialization) and have kernel.py look that up — or simply remove the kernel.py fast-skip and accept the slightly higher function-call cost, since the inner early-return already preserves correctness.

[github-actions]

Coverage Report (2941eefce)

File	Coverage	Missing
🟢 python/quadrants/lang/_func_base.py	100%
🔴 python/quadrants/lang/kernel.py	50%	495

Diff coverage: 88% · Overall: 74% · 8 lines, 1 missing

Full annotated report

[hughperkins]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 80fd717c11

ℹ️ 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]

Preserve argument validation before all-field fast-path

This early continue runs before _recursive_set_args, so any object carrying _qd_all_field=True (which is now cached on valid all-Field structs) bypasses all runtime type checks and argument emission for the current parameter. If such an object is accidentally passed to a non-struct parameter (for example, due to argument ordering mistakes), previous code raised QuadrantsRuntimeTypeError, but this change silently leaves the launch-context slot unset and can produce incorrect kernel execution instead of a clear error.

Useful? React with 👍 / 👎.

[github-actions]

Coverage Report (5522b8c66)

File	Coverage	Missing
🟢 python/quadrants/lang/_func_base.py	100%
🔴 python/quadrants/lang/kernel.py	67%	7
🟢 tests/python/test_all_field_shortcut.py	99%	345,376-377

Diff coverage: 98% · Overall: 74% · 226 lines, 4 missing

Full annotated report

[hughperkins]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5d4ac7c768

ℹ️ 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]

Preserve dataclass type validation before skipping args

This shortcut skips _recursive_set_args() solely when val._qd_all_field is true and the annotation is a dataclass, but _recursive_set_args() is where the strict provided_arg_type is needed_arg_type check lives. As a result, once an instance has _qd_all_field=True, passing it to a kernel parameter annotated with a different all-Field dataclass type can be silently accepted instead of raising the previous runtime type error, which weakens API/type safety and can route the wrong struct type through launch.

Useful? React with 👍 / 👎.

[github-actions]

Coverage Report (5d4ac7c76)

File	Coverage	Missing
🟢 python/quadrants/lang/_func_base.py	100%
🔴 python/quadrants/lang/kernel.py	67%	7
🟢 tests/python/test_all_field_shortcut.py	99%	345,376-377

Diff coverage: 98% · Overall: 74% · 226 lines, 4 missing

Full annotated report

[github-actions]

Coverage Report (af1b51ab9)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	90%	3362

Diff coverage: 90% · Overall: 74% · 10 lines, 1 missing

Full annotated report