[Lang] Add layout= to ndrange by hughperkins · Pull Request #710 · Genesis-Embodied-AI/quadrants

hughperkins · 2026-05-20T13:11:27Z

Issue: #

Brief Summary

copilot:summary

Walkthrough

copilot:walkthrough

qd.ndrange(*args, layout=None) now accepts a permutation of range(N) describing the iteration nesting order, outermost (slowest-varying) first. Mirrors layout= on qd.tensor exactly: canonical-preserving (loop variables stay bound to canonical axes; only visit order changes), identity / None is the no-op default. Implementation is pure-Python and confined to two files: * python/quadrants/lang/_ndrange.py: validate layout up front (QuadrantsSyntaxError on wrong length / non-permutation), normalize identity to None, and store bounds / dimensions / acc_dimensions in *physical* order (with _physical_to_canonical capturing the per-level remap). __iter__ keeps yielding canonical-ordered tuples in the new physical visit order. * python/quadrants/lang/ast/ast_transformer.py: build_ndrange_for and build_grouped_ndrange_for assign each decomposed flat-thread index to targets[physical_to_canonical[p]] instead of targets[p]. For the identity case physical_to_canonical is range(n) and the emitted IR matches the pre-layout codegen byte-for-byte (no perf regression on existing kernels). Primary motivation: pairing qd.ndrange(layout=L) with qd.tensor(..., layout=L) lines adjacent flat threads up with adjacent physical memory slots, restoring coalesced / cache-friendly access for layout-tagged tensors while keeping canonical indexing in the body. Design doc: perso_hugh/doc/ndrange_layout.md.

New subsection in parallelization.md after the qd.grouped section, with a worked example pairing qd.tensor(layout=(1, 0)) with qd.ndrange(layout=(1, 0)) to align iteration order with physical memory order. Notes that the layout permutation has the exact same meaning as on qd.tensor (canonical axis index at each successive nesting level, outermost first), and that loop variables remain bound to canonical axes regardless of layout. Cross-reference added from the "Controlling physical layout" section of tensor.md so users discovering qd.tensor(layout=...) find the matching iteration-order knob.

Covers: * layout=None / identity equivalence with the no-keyword form (in-kernel, field backend) * non-identity layouts (rank 2 transposed, rank 3 with (2, 0, 1)): canonical loop targets, full coverage of the index space, preserved (begin, end) offsets on tuple bounds * flat-index decomposition matches the physical iteration order * qd.grouped(...) returns canonical indices regardless of layout * qd.static(qd.grouped(...)) (unrolled path) also sees canonical indices in the layout-induced order * pairing with qd.tensor(layout=...) (the documented primary use case) * Python-side iteration outside @qd.kernel, including a helper that reconstructs the expected canonical-tuple sequence from (dims, layout) * introspection: layout=None and identity both normalize self.layout to None; non-identity preserves the user-supplied tuple * degenerate 1-D and 0-D ndranges * error cases: wrong length, not a permutation, out-of-range entries

``qd.grouped`` is decorated ``@quadrants_scope`` and raises when called outside a kernel, so test the underlying ``_Ndrange.grouped()`` method directly. Equivalent coverage, no qd.grouped-from-Python-scope assertion.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 24778270d0

ℹ️ 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 · 2026-05-20T13:15:26Z

+                    f"qd.ndrange(layout={layout_t!r}) has {len(layout_t)} entries "
+                    f"but ndrange was called with {n} dimension argument(s); they must match"
+                )
+            if sorted(layout_t) != list(range(n)):


Validate layout element types before permutation sort

sorted(layout_t) can raise a raw Python TypeError when layout contains non-integer or mixed-type entries (for example layout=(0, "1")), so users get an internal exception instead of the documented qd.QuadrantsSyntaxError for invalid layouts. Add an explicit integer-type check for each element before sorting/permutation validation so all invalid layouts are rejected consistently with Quadrants error types.

Useful? React with 👍 / 👎.

Reframe the parallelization.md subsection so that ``layout=`` is presented as an iteration-order control in its own right, not "the matching keyword for qd.Tensor". Add an explicit "independent of what's in the loop body" note and move the qd.tensor pairing into a "When is layout= useful?" subsection as the motivating but not gating use case. Also update the ``ndrange`` docstring to match: layout= works with field, ndarray, tensor, vector/matrix variants, or no tensor at all; the tensor pairing is the motivating use case, not the only one.

The motivating use case framing is just "align iteration with physical memory layout", not specifically with non-default layouts; the example that follows speaks for itself.

The subsection's example (qd.tensor + matching ndrange layout) belongs in the tensor user guide, not the parallelization one, where it muddies the "layout= is independent of the body" message.

Quadrants is the public repo; design notes live in a private one and shouldn't be cited from here.

Addresses codex review on PR #710: ``sorted(layout_t)`` raises a raw Python ``TypeError`` on mixed-type / non-integer entries (e.g. ``layout=(0, "1")``), exposing an internal exception instead of the documented Quadrants error type. Add an explicit per-element check that raises ``QuadrantsTypeError`` (matching the convention used for bounds validation above). ``bool`` is rejected explicitly even though it is an ``int`` subclass — accepting ``True`` / ``False`` as axis indices would be a foot-gun. Zero perf impact: ``_Ndrange.__init__`` runs at AST-build time, once per kernel compile (cached thereafter). The new ``isinstance`` loop only runs when ``layout is not None``; default ``layout=None`` short- circuits before any new check. Adds three tests: string entry, float entry, bool entry.

Reflow the "not a permutation" QuadrantsSyntaxError onto one line per black's default config in this repo (length still under the project limit). No semantic change.

duburcqa · 2026-05-20T13:29:39Z


 `I` is a `qd.Vector` with one element per dimension.

+### Controlling iteration order with `layout=`


You prefer 'layout' over 'axis' ?

I think it should be consistent with qd.Tensor.

(honestly though, in the absence of numpy using axis, I do actually prefer layout).

axis in numpy is a single integer, not a tuple. I think the meaning is slightly different:

Layout is weird. There is no layout involved here.

Like, it has nothing to do with qd.Tensor layout, why the name should be the same?

Numpy uses axes: https://numpy.org/doc/2.2/reference/generated/numpy.permute_dims.html

its the exact same concept

if you check the commit history AI was trying to exactly tie the motivation of using this functionality on ndarray to using the similar functionality on qd.Tensor

in practice, they will tend to be used together

without using the similar functionality on qd.Tensor, there is no obvious reason why we would want to be able to switch the dimensions on ndarray at runtime I feel?

I feel that our options are:

use layout for both Tensor and ndarray, or

use axes for both Tensor and ndarray

(I'm tempted to compromise with using layout for Tensor and axes for ndarray, but it's inconsistent, and inconsistency within Quadrants is worse than inconsistency between Quadrants and some other library I feel).

I suppose you could argue that one is permuting the physical layout, and one is permuting iteration order, and those are different things. But it's still about permutating dimensions I feel.

github-actions · 2026-05-20T14:01:06Z

Total: 3 file(s) changed, +304 -32 code lines.

github-actions · 2026-05-20T15:27:16Z

Diff coverage: 97% · 294 lines, 10 missing

hughperkins added 4 commits May 20, 2026 05:33

test: fix Python-scope grouped iteration test

2477827

``qd.grouped`` is decorated ``@quadrants_scope`` and raises when called outside a kernel, so test the underlying ``_Ndrange.grouped()`` method directly. Equivalent coverage, no qd.grouped-from-Python-scope assertion.

chatgpt-codex-connector Bot reviewed May 20, 2026

View reviewed changes

hughperkins added 6 commits May 20, 2026 06:16

docs: drop "non-default physical memory layout" framing

f890f34

The motivating use case framing is just "align iteration with physical memory layout", not specifically with non-default layouts; the example that follows speaks for itself.

docs: drop "When is layout= useful?" subsection from parallelization.md

bfd49e7

The subsection's example (qd.tensor + matching ndrange layout) belongs in the tensor user guide, not the parallelization one, where it muddies the "layout= is independent of the body" message.

test: drop reference to private design-notes location

370a86b

Quadrants is the public repo; design notes live in a private one and shouldn't be cited from here.

ndrange: black reformat

3fa8d29

Reflow the "not a permutation" QuadrantsSyntaxError onto one line per black's default config in this repo (length still under the project limit). No semantic change.

duburcqa reviewed May 20, 2026

View reviewed changes


		`I` is a `qd.Vector` with one element per dimension.

		### Controlling iteration order with `layout=`

Conversation

hughperkins commented May 20, 2026

Brief Summary

Walkthrough

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot May 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hughperkins May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented May 20, 2026

Uh oh!

github-actions Bot commented May 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hughperkins May 20, 2026 •

edited

Loading