Skip to content

Document the contract of custom NDRange mappings - #781

Open
giordano wants to merge 1 commit into
mainfrom
mg/document-mapping-contract
Open

giordano wants to merge 1 commit into
mainfrom
mg/document-mapping-contract

Conversation

@giordano

Copy link
Copy Markdown
Collaborator

No description provided.

Since #771 the `mapping` of an `NDRange` decides how blocked indices map to
`ndrange` indices, and every backend derives the index, validity and linear
index of a work item from `expand`, `in` and `linear_index`. That makes it
possible for a package to launch kernels over an iteration space of its
own, such as a list of indices, by specializing `partition` and `cartesian`
for its launch object and `expand` for its mapping, but nothing said so,
and a backend overriding `__validindex` for a generic context would
silently break such an extension.

The `NDRange` docstring now lists the functions a custom mapping defines,
`partition` and the new `cartesian` and `expand` docstrings point to it,
and the notes for backend implementations state that a backend must go
through `expand`, `in` and `linear_index` and must not assume the `ndrange`
of a context is a `CartesianIndices`.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHciC8x39gm97sABrSvBkt
@giordano giordano added the documentation Improvements or additions to documentation label Sep 14, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Benchmark Results

Show table
main 7f06431... main / 7f06431...
launch/3D static workgroup, dynamic ndrange 23.5 ± 17 μs 23.5 ± 20 μs 1 ± 1.1
launch/3D static workgroup, static ndrange 23.1 ± 16 μs 23.6 ± 20 μs 0.979 ± 1.1
launch/dynamic workgroup, dynamic ndrange 22.6 ± 8.4 μs 22.8 ± 4.9 μs 0.99 ± 0.43
launch/dynamic workgroup, dynamic ndrange, workgroupsize given 23.7 ± 17 μs 23.8 ± 20 μs 0.997 ± 1.1
launch/static workgroup, dynamic ndrange 23.7 ± 18 μs 23.4 ± 21 μs 1.02 ± 1.2
launch/static workgroup, static ndrange 22.9 ± 15 μs 23.4 ± 19 μs 0.983 ± 1
partition/dynamic workgroup, dynamic ndrange 0.0524 ± 0.0073 μs 0.0535 ± 0.0053 μs 0.979 ± 0.17
partition/static workgroup, dynamic ndrange 0.0612 ± 0.011 μs 0.0605 ± 0.011 μs 1.01 ± 0.25
partition/static workgroup, static ndrange 2.01 ± 0.005 ns 1.35 ± 0.009 ns 1.48 ± 0.011
saxpy/default/Float16/1024 27.2 ± 25 μs 25.2 ± 22 μs 1.08 ± 1.4
saxpy/default/Float16/1048576 0.365 ± 0.021 ms 0.365 ± 0.019 ms 0.999 ± 0.078
saxpy/default/Float16/16384 0.0503 ± 0.025 ms 0.0468 ± 0.025 ms 1.07 ± 0.79
saxpy/default/Float16/2048 0.0372 ± 0.023 ms 0.0362 ± 0.023 ms 1.03 ± 0.92
saxpy/default/Float16/256 23.5 ± 8.3 μs 23.5 ± 5 μs 1 ± 0.41
saxpy/default/Float16/262144 0.126 ± 0.023 ms 0.127 ± 0.023 ms 0.994 ± 0.26
saxpy/default/Float16/32768 0.0472 ± 0.023 ms 0.0505 ± 0.025 ms 0.933 ± 0.64
saxpy/default/Float16/4096 0.0504 ± 0.023 ms 0.0499 ± 0.023 ms 1.01 ± 0.66
saxpy/default/Float16/512 24.1 ± 14 μs 23.7 ± 8.2 μs 1.02 ± 0.68
saxpy/default/Float16/64 23.9 ± 16 μs 23.4 ± 3 μs 1.02 ± 0.69
saxpy/default/Float16/65536 0.0567 ± 0.022 ms 0.0575 ± 0.023 ms 0.986 ± 0.55
saxpy/default/Float32/1024 27.9 ± 27 μs 26.3 ± 26 μs 1.06 ± 1.5
saxpy/default/Float32/1048576 0.367 ± 0.046 ms 0.359 ± 0.045 ms 1.02 ± 0.18
saxpy/default/Float32/16384 0.049 ± 0.026 ms 0.0432 ± 0.024 ms 1.13 ± 0.87
saxpy/default/Float32/2048 0.0489 ± 0.023 ms 0.0477 ± 0.024 ms 1.02 ± 0.7
saxpy/default/Float32/256 23.4 ± 14 μs 23.7 ± 5.4 μs 0.988 ± 0.63
saxpy/default/Float32/262144 0.111 ± 0.019 ms 0.111 ± 0.018 ms 1 ± 0.24
saxpy/default/Float32/32768 0.0477 ± 0.025 ms 0.0476 ± 0.025 ms 1 ± 0.73
saxpy/default/Float32/4096 0.0509 ± 0.022 ms 0.0512 ± 0.022 ms 0.995 ± 0.61
saxpy/default/Float32/512 25.8 ± 27 μs 28 ± 26 μs 0.923 ± 1.3
saxpy/default/Float32/64 23.2 ± 6.1 μs 23.4 ± 3.3 μs 0.991 ± 0.29
saxpy/default/Float32/65536 0.0591 ± 0.026 ms 0.0568 ± 0.025 ms 1.04 ± 0.65
saxpy/default/Float64/1024 0.0347 ± 0.029 ms 27.8 ± 27 μs 1.24 ± 1.6
saxpy/default/Float64/1048576 0.684 ± 0.083 ms 0.663 ± 0.083 ms 1.03 ± 0.18
saxpy/default/Float64/16384 0.049 ± 0.026 ms 0.0487 ± 0.027 ms 1.01 ± 0.77
saxpy/default/Float64/2048 0.0401 ± 0.027 ms 0.0377 ± 0.029 ms 1.06 ± 1.1
saxpy/default/Float64/256 23.5 ± 12 μs 24.5 ± 13 μs 0.958 ± 0.69
saxpy/default/Float64/262144 0.194 ± 0.025 ms 0.189 ± 0.027 ms 1.03 ± 0.2
saxpy/default/Float64/32768 0.0563 ± 0.028 ms 0.0567 ± 0.028 ms 0.994 ± 0.69
saxpy/default/Float64/4096 0.045 ± 0.028 ms 0.0528 ± 0.023 ms 0.852 ± 0.64
saxpy/default/Float64/512 23.7 ± 13 μs 24.3 ± 8.8 μs 0.974 ± 0.64
saxpy/default/Float64/64 23.5 ± 7.4 μs 23.6 ± 3.2 μs 0.996 ± 0.34
saxpy/default/Float64/65536 0.071 ± 0.026 ms 0.0732 ± 0.025 ms 0.969 ± 0.49
saxpy/static workgroup=(1024,)/Float16/1024 29.6 ± 25 μs 27.4 ± 25 μs 1.08 ± 1.4
saxpy/static workgroup=(1024,)/Float16/1048576 0.363 ± 0.019 ms 0.363 ± 0.02 ms 1 ± 0.076
saxpy/static workgroup=(1024,)/Float16/16384 0.0404 ± 0.022 ms 0.0409 ± 0.022 ms 0.986 ± 0.76
saxpy/static workgroup=(1024,)/Float16/2048 0.0378 ± 0.023 ms 0.0357 ± 0.022 ms 1.06 ± 0.93
saxpy/static workgroup=(1024,)/Float16/256 27.4 ± 25 μs 26.2 ± 23 μs 1.04 ± 1.3
saxpy/static workgroup=(1024,)/Float16/262144 0.124 ± 0.023 ms 0.124 ± 0.023 ms 1 ± 0.26
saxpy/static workgroup=(1024,)/Float16/32768 0.0471 ± 0.023 ms 0.0483 ± 0.024 ms 0.975 ± 0.68
saxpy/static workgroup=(1024,)/Float16/4096 0.0491 ± 0.02 ms 0.0483 ± 0.021 ms 1.02 ± 0.61
saxpy/static workgroup=(1024,)/Float16/512 27.1 ± 25 μs 26.2 ± 22 μs 1.04 ± 1.3
saxpy/static workgroup=(1024,)/Float16/64 26.8 ± 25 μs 26.3 ± 22 μs 1.02 ± 1.3
saxpy/static workgroup=(1024,)/Float16/65536 0.0549 ± 0.023 ms 0.0549 ± 0.022 ms 1 ± 0.58
saxpy/static workgroup=(1024,)/Float32/1024 27 ± 27 μs 27.6 ± 26 μs 0.976 ± 1.3
saxpy/static workgroup=(1024,)/Float32/1048576 0.372 ± 0.046 ms 0.362 ± 0.042 ms 1.03 ± 0.17
saxpy/static workgroup=(1024,)/Float32/16384 0.0478 ± 0.025 ms 0.0447 ± 0.024 ms 1.07 ± 0.8
saxpy/static workgroup=(1024,)/Float32/2048 0.048 ± 0.021 ms 0.0467 ± 0.023 ms 1.03 ± 0.68
saxpy/static workgroup=(1024,)/Float32/256 25.9 ± 24 μs 23.1 ± 13 μs 1.12 ± 1.2
saxpy/static workgroup=(1024,)/Float32/262144 0.109 ± 0.018 ms 0.109 ± 0.018 ms 1 ± 0.23
saxpy/static workgroup=(1024,)/Float32/32768 0.0456 ± 0.024 ms 0.0459 ± 0.025 ms 0.993 ± 0.75
saxpy/static workgroup=(1024,)/Float32/4096 0.0487 ± 0.021 ms 0.0487 ± 0.021 ms 1 ± 0.61
saxpy/static workgroup=(1024,)/Float32/512 30.2 ± 27 μs 26.7 ± 26 μs 1.13 ± 1.5
saxpy/static workgroup=(1024,)/Float32/64 26 ± 24 μs 23.7 ± 17 μs 1.1 ± 1.3
saxpy/static workgroup=(1024,)/Float32/65536 0.0569 ± 0.025 ms 0.0573 ± 0.024 ms 0.993 ± 0.6
saxpy/static workgroup=(1024,)/Float64/1024 0.0341 ± 0.028 ms 27.5 ± 27 μs 1.24 ± 1.6
saxpy/static workgroup=(1024,)/Float64/1048576 0.685 ± 0.077 ms 0.683 ± 0.071 ms 1 ± 0.15
saxpy/static workgroup=(1024,)/Float64/16384 0.047 ± 0.024 ms 0.0435 ± 0.024 ms 1.08 ± 0.82
saxpy/static workgroup=(1024,)/Float64/2048 0.0413 ± 0.027 ms 0.0435 ± 0.026 ms 0.951 ± 0.85
saxpy/static workgroup=(1024,)/Float64/256 0.0318 ± 0.024 ms 27.7 ± 25 μs 1.15 ± 1.4
saxpy/static workgroup=(1024,)/Float64/262144 0.185 ± 0.026 ms 0.185 ± 0.025 ms 1 ± 0.2
saxpy/static workgroup=(1024,)/Float64/32768 0.052 ± 0.024 ms 0.0547 ± 0.026 ms 0.95 ± 0.63
saxpy/static workgroup=(1024,)/Float64/4096 0.0491 ± 0.022 ms 0.0505 ± 0.022 ms 0.971 ± 0.61
saxpy/static workgroup=(1024,)/Float64/512 27.9 ± 26 μs 24.1 ± 19 μs 1.16 ± 1.4
saxpy/static workgroup=(1024,)/Float64/64 25.9 ± 24 μs 24 ± 17 μs 1.08 ± 1.3
saxpy/static workgroup=(1024,)/Float64/65536 0.0684 ± 0.024 ms 0.0694 ± 0.024 ms 0.985 ± 0.49
time_to_load 0.774 ± 0.016 s 0.759 ± 0.016 s 1.02 ± 0.029
main 7f06431... main / 7f06431...
launch/3D static workgroup, dynamic ndrange 7 allocs: 0.297 kB 7 allocs: 0.297 kB 1
launch/3D static workgroup, static ndrange 8 allocs: 0.328 kB 8 allocs: 0.328 kB 1
launch/dynamic workgroup, dynamic ndrange 17 allocs: 0.469 kB 17 allocs: 0.469 kB 1
launch/dynamic workgroup, dynamic ndrange, workgroupsize given 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
launch/static workgroup, dynamic ndrange 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
launch/static workgroup, static ndrange 4 allocs: 0.203 kB 4 allocs: 0.203 kB 1
partition/dynamic workgroup, dynamic ndrange 2 allocs: 0.0625 kB 2 allocs: 0.0625 kB 1
partition/static workgroup, dynamic ndrange 2 allocs: 32 B 2 allocs: 32 B 1
partition/static workgroup, static ndrange 0 allocs: 0 B 0 allocs: 0 B
saxpy/default/Float16/1024 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/1048576 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/16384 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/2048 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/256 18 allocs: 0.484 kB 18 allocs: 0.484 kB 1
saxpy/default/Float16/262144 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/32768 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/4096 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/512 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float16/64 18 allocs: 0.484 kB 18 allocs: 0.484 kB 1
saxpy/default/Float16/65536 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/1024 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/1048576 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/16384 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/2048 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/256 18 allocs: 0.484 kB 18 allocs: 0.484 kB 1
saxpy/default/Float32/262144 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/32768 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/4096 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/512 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float32/64 18 allocs: 0.484 kB 18 allocs: 0.484 kB 1
saxpy/default/Float32/65536 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/1024 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/1048576 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/16384 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/2048 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/256 18 allocs: 0.484 kB 18 allocs: 0.484 kB 1
saxpy/default/Float64/262144 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/32768 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/4096 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/512 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/default/Float64/64 18 allocs: 0.484 kB 18 allocs: 0.484 kB 1
saxpy/default/Float64/65536 26 allocs: 0.609 kB 26 allocs: 0.609 kB 1
saxpy/static workgroup=(1024,)/Float16/1024 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/1048576 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/16384 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/2048 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/256 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
saxpy/static workgroup=(1024,)/Float16/262144 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/32768 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/4096 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/512 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float16/64 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
saxpy/static workgroup=(1024,)/Float16/65536 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/1024 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/1048576 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/16384 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/2048 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/256 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
saxpy/static workgroup=(1024,)/Float32/262144 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/32768 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/4096 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/512 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float32/64 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
saxpy/static workgroup=(1024,)/Float32/65536 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/1024 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/1048576 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/16384 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/2048 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/256 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
saxpy/static workgroup=(1024,)/Float64/262144 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/32768 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/4096 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/512 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
saxpy/static workgroup=(1024,)/Float64/64 3 allocs: 0.172 kB 3 allocs: 0.172 kB 1
saxpy/static workgroup=(1024,)/Float64/65536 6 allocs: 0.219 kB 6 allocs: 0.219 kB 1
time_to_load 0.201 k allocs: 11.8 kB 0.201 k allocs: 11.8 kB 1

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

@codecov

codecov Bot commented Sep 14, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 64.89%. Comparing base (e5f5f61) to head (7f06431).

Additional details and impacted files
@@           Coverage Diff           @@
##             main     #781   +/-   ##
=======================================
  Coverage   64.89%   64.89%           
=======================================
  Files          23       23           
  Lines        2011     2011           
=======================================
  Hits         1305     1305           
  Misses        706      706           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.


A backend must not assume that `__ndrange(ctx)` is a `CartesianIndices` or that
`expand` is an affine map: the `mapping` field of the `NDRange` lets a package
define its own iteration space, for example a list of indices to visit, by

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a new contract, I am not supper happy about. These are private implementation details and a packing extending them is doing fishy work, and potentially type piracy.

I would be much happier if we had a public interface that folks could extend

abstract type Mapping end

And then Oceananigans could do its subtype of that, and so ndrange=Mapping()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants