Conversation
* next step in kernel migration from parents to offsets * linter fix * add jax bincount * add cpp kernel to convert parents to offsets * fix typo * fix typo in an auto-generated test * almost there * nearly there * cleanup * fix remaining kernels * final cleanup * format json data * format * initialize maxnextparents = -1 (a sentinel meaning "no bin was touched") * update the kernels to work on offsets! * format * migrate cupy rawkernels from parents to offsets * migrate jax reducers * fix for platfroms where the int64 counts can't be safely cast * add bincount for cupy backend * compact loc * fix windows build and add optional OpenMP support * update cuda kernels * feat: add `missing_repeat` kernel implementation using cuda.compute (#3922) * feat: add missing_repeat implementation using cuda.compute * keep the same name as in cpu kernel * add tests for repetitions>1 and regularsize>1 * style: pre-commit fixes * add keywords * style fixes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * feat: add `index_rpad_and_clip` kernels implementations using cuda.compute (#3923) * feat: add `index_rpad_and_clip` kernels implementation using cuda.compute * style fix * add a test for `index_rpad_and_clip_axis0` that would have target>length * add keyword names * style * Apply suggestions from Ianna Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch> * style: pre-commit fixes --------- Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * fix: array attrs not being validated at creation and being of inconsistent type (#3996) * convert dict attrs to Attrs type * add test * fix: ensure jax backend uses arrays on cpu only (#3990) * ensure jax uses cpu by default * remove now redundant jax_platform_name setting * change error messages * better errors * remove DeviceArray mentions as that does not exist while we're at it too * almost there * cleanup * cleanup * remove old data file * update cupy kernels * remove exlicit test for depricated kernel * add complex kernels and port them to use offsets * fix complex reducers * add remaining kernels * add kernels for complex and bool sum * move to segmented_reduce * fix typo * promote type * fix complex bool reducer * try another algo * missed one * use type inference * make numba happy * remove reducer overloads * use sum op func for complex types * remove test of depricated code * avoid using == to compare floating-point products * fix boundary tests * cleanup * cleanup * remove dead code * revert lexsort to argmin/max * handrolled lexsort * remove dead code * remove bincount * import numba cuda for jit * fix tests --------- Co-authored-by: maxymnaumchyk <maxymnaumchyk@gmail.com> Co-authored-by: maxymnaumchyk <70752300+maxymnaumchyk@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Iason Krommydas <iason.krom@gmail.com>
|
@ariostas should definitely take a look too (when he's back to wor). I do not fully trust myself with c++ yet but can definitely review the python part changes. |
|
One question I have is how are the kernel test data generated? I'm not familiar with this part. |
yes, we can trust the data because these are the data from the tests. |
|
I've ran this against HiggsDNA and getting the same results for my analysis. Only small floating point differences that |
Thanks @ikrommyd ! Floating point comparison is indeed a tricky one. |
|
Yeah totally reasonable. Floating point ops are not commutative. Anyways I'll test this out a bit more with analyses examples and I also wanna if there's any noticeable performance hit. |
Thanks! Don’t worry about the performance yet — that’s the next on my list. First we must insure that the results are correct.👍 |
|
The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR4056 |
add benchmark script
* feat: add `awkward_IndexedArray_overlay_mask` kernel using cuda.compute * feat: add `awkward_IndexedArray_reduce_next_64` kernel using cuda.compute * feat: add `IndexedArray_reduce_next_nonlocal_nextshifts_64` kernel using cuda.compute * feat: add `ByteMaskedArray_getitem_nextcarry` kernel using cuda.compute * feat: add `awkward_ByteMaskedArray_numnull` kernel using cuda.compute * feat: add `awkward_RegularArray_getitem_jagged_expand` kernel using cuda.compute * add an upper bound * style: pre-commit fixes * feat: add `awkward_RegularArray_getitem_jagged_expand` kernel using cuda.compute * feat: add `awkward_UnionArray_simplify_one` kernel using cuda.compute * feat: add `awkward_ListArray_broadcast_tooffsets` kernel using cuda.compute * feat: add `awkward_ListArray_localindex` kernel using cuda.compute * fix the impl * feat: add `awkward_ListArray_compact_offsets` kernel using cuda.compute * feat: add `awkward_ListArray_combinations_length` kernel using cuda.compute * feat: add `awkward_ListArray_combinations` kernel using cuda.compute * feat: add `awkward_UnionArray_nestedfill_tags_index` kernel using cuda.compute * fix the tests for kernels that are deliberately raising errors * compare `starts` and `stops` separately * ignore `memptr` argument for pylint * style: pre-commit fixes * Add functions for indexing and repeating arrays * style: pre-commit fixes * return unary_transform call for segment_ids * style: pre-commit fixes * update the `awkward_IndexedArray_reduce_next_64` to work with offsets * style: pre-commit fixes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch>
* feat: add reimplementation of the `awkward_RegularArray` kernels using cuda.compute * reimplement more kernels using cuda.compute * reimplement more kernels using cuda.compute * style: pre-commit fixes * style fix * use the named functions instead of lambdas * add new kernels * add new kernels * fix a type mismatch * style: pre-commit fixes * add new kernels * fix the overflow error * new batch of kernels (draft) * add comments and cleanup * fix the tests messages * add new `BitMaskedArray` kernels * modify the `test_1381_check_errors.py` test * resolve the conflicts --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
|
@ianna what's the plan with this PR? Is this something that you would want to merge soon or is it to keep like a live comparison of the "Awkward3" changes?
If this is going in |
The plan is to discuss it at AS meeting.
ASIS performance is good enough. And we already know the bottlenecks -- where we can optimize it -- I want to have confirmation from the users that validation passes. |
|
What I'd like to see here is a benchmark showing that all or most of the high level ops ( It is important to know this before we can merge this change. The benchmarks can be done either with synthetic randomly generated data or a reasonably big root file (between 100K and 500K events which is the reasonable analysis chunk size atm). |
ianna
changed the title
parents to offsets (#3988)parents to offsets
|
@TaiSakuma -- could you please check the test failure? It's an |
|
@ianna @TaiSakuma this is a reproducer of the CI failure. Didn't have time to look if such an array is bad enough that should never be generated by hypothesis-awkward or whether we have a bug in awkward code. import awkward as ak
import numpy as np
layout = ak.contents.UnionArray(
ak.index.Index8(np.empty(0, dtype=np.int8)),
ak.index.Index64(np.empty(0, dtype=np.int64)),
[
ak.contents.BitMaskedArray(
ak.index.IndexU8(np.empty(0, dtype=np.uint8)),
ak.contents.EmptyArray(),
valid_when=True,
length=0,
lsb_order=False,
),
ak.contents.BitMaskedArray(
ak.index.IndexU8(np.empty(0, dtype=np.uint8)),
ak.contents.EmptyArray(),
valid_when=True,
length=0,
lsb_order=False,
),
],
)
ak.array_equal(layout, layout) |
|
🤖 AI text below 🤖 All four audits are complete. Here's the full review. Review: PR #4056 — feat: complete CPU kernel migration from
|
|
Pref analysis with Opus 4.8: 🤖 AI text below 🤖 CPU performance: better — substantially, for reductionsYes, it's better. I built Overall geomean: ~4.1× faster (after vs. before), across 3 data shapes × the reduction ops. The big wins — segmented reductions along
|
| Op | Typical speedup |
|---|---|
sum / prod |
2–19× |
max / min |
2–18× |
count_nonzero |
1.7–14× |
argmax / argmin (int64) |
1.8–3.7× |
sum (complex128) |
3–12× |
any / all / sum (bool) |
2–700× |
The improvement is biggest for "few long lists" and "balanced" shapes (the old parents-based code scaled poorly there); even the hardest case, "2M short lists," is still 1.7–4× faster. The extreme bool any/all numbers (e.g. 730×) are a genuine but somewhat special case — the new offsets-based path has an effectively free fast path for boolean reductions, so I'd treat those as outliers rather than the headline.
The one regression — sort / argsort
These are mixed-to-slightly-worse:
sort axis=1: 0.79–0.92× (up to ~20% slower) for balanced/long-list shapes, but 1.25–1.28× faster for many-short-lists.argsort axis=1: roughly neutral (0.93–1.05×), faster only for many-short-lists.- A couple of
float64argmax/argmincases in the long-list shape are ~8% slower.
So sort/argsort didn't benefit from the migration the way the reducers did, and lost a bit on long-list workloads.
Bottom line
For the operations this PR targets (segmented reductions), CPU performance is much better — typically several times faster, with no reduction regressions. The only thing to watch is sort/argsort, which is roughly neutral overall but ~10–20% slower on long-list inputs.
Caveats: single Apple-silicon machine, best-of-7 timings, Python 3.13. I couldn't test GPU (no CUDA here), and this measures only the kernels touched by the branch, not end-to-end workloads. The two built venvs are at /tmp/v_before and /tmp/v_after and the benchmark is /tmp/bench_cpu.py if you want to re-run or extend it.
|
@henryiii — fantastic! Can you ask for the memory profile comparison as well? |
|
@henryiii what did the agent benchmark? Cause the kernels themselves change but also the number of how many we need to run for each op might change too. We'd want to benchmark highlevel ops like |
|
Here's the benchmark script ( """CPU reduction/sort benchmark targeting the parents->offsets kernel migration."""
import gc
import sys
import time
import numpy as np
import awkward as ak
print("awkward", ak.__version__, "| awkward_cpp", getattr(__import__("awkward_cpp"), "__version__", "?"))
rng = np.random.default_rng(12345)
def make_ragged(n_lists, avg_len, dtype):
counts = rng.integers(0, 2 * avg_len + 1, size=n_lists)
total = int(counts.sum())
if dtype == bool:
data = rng.integers(0, 2, size=total).astype(bool)
elif dtype == np.complex128:
data = (rng.standard_normal(total) + 1j * rng.standard_normal(total)).astype(np.complex128)
elif np.issubdtype(dtype, np.integer):
data = rng.integers(-1000, 1000, size=total).astype(dtype)
else:
data = rng.standard_normal(total).astype(dtype)
offsets = np.zeros(n_lists + 1, dtype=np.int64)
offsets[1:] = np.cumsum(counts)
layout = ak.contents.ListOffsetArray(
ak.index.Index64(offsets), ak.contents.NumpyArray(data)
)
return ak.Array(layout)
def bench(fn, arr, repeat=7, inner=3):
# warmup
fn(arr)
best = float("inf")
for _ in range(repeat):
gc.disable()
t0 = time.perf_counter()
for _ in range(inner):
fn(arr)
t1 = time.perf_counter()
gc.enable()
best = min(best, (t1 - t0) / inner)
return best
# scales: (n_lists, avg_len)
SCALES = {
"few_long": (1_000, 10_000), # few segments, long lists
"many_short": (2_000_000, 5), # many segments, short lists
"balanced": (200_000, 50), # balanced
}
OPS = []
def reg(name, fn, dtypes):
OPS.append((name, fn, dtypes))
reg("sum_axis1", lambda a: ak.sum(a, axis=1), [np.int64, np.float64])
reg("sum_axisNone", lambda a: ak.sum(a, axis=None), [np.int64, np.float64])
reg("prod_axis1", lambda a: ak.prod(a, axis=1), [np.int64, np.float64])
reg("max_axis1", lambda a: ak.max(a, axis=1), [np.int64, np.float64])
reg("min_axis1", lambda a: ak.min(a, axis=1), [np.int64, np.float64])
reg("argmax_axis1", lambda a: ak.argmax(a, axis=1), [np.int64, np.float64])
reg("argmin_axis1", lambda a: ak.argmin(a, axis=1), [np.int64, np.float64])
reg("count_nonzero_axis1", lambda a: ak.count_nonzero(a, axis=1), [np.int64, np.float64])
reg("any_axis1", lambda a: ak.any(a, axis=1), [bool])
reg("all_axis1", lambda a: ak.all(a, axis=1), [bool])
reg("sum_bool_axis1", lambda a: ak.sum(a, axis=1), [bool])
reg("sum_complex_axis1", lambda a: ak.sum(a, axis=1), [np.complex128])
reg("sort_axis1", lambda a: ak.sort(a, axis=1), [np.int64, np.float64])
reg("argsort_axis1", lambda a: ak.argsort(a, axis=1), [np.int64, np.float64])
results = {}
for scale_name, (n_lists, avg_len) in SCALES.items():
cache = {}
for name, fn, dtypes in OPS:
for dt in dtypes:
key = id(dt) if not isinstance(dt, type) else dt
if dt not in cache:
cache[dt] = make_ragged(n_lists, avg_len, dt)
arr = cache[dt]
dtn = np.dtype(dt).name
t = bench(fn, arr)
results[(scale_name, name, dtn)] = t
print(f"{scale_name:12s} {name:22s} {dtn:12s} {t*1e3:10.3f} ms", flush=True)
del cache
gc.collect()
# dump machine-readable
import json
with open(sys.argv[1] if len(sys.argv) > 1 else "/tmp/bench_out.json", "w") as f:
json.dump({f"{k[0]}|{k[1]}|{k[2]}": v for k, v in results.items()}, f, indent=2)
print("WROTE", sys.argv[1] if len(sys.argv) > 1 else "/tmp/bench_out.json")A few notes on the methodology choices:
The comparison/analysis snippet wasn't saved to a file — it was run inline. Want me to save a cleaned-up version of both the benchmark and the comparison logic into the repo (e.g. under |
|
🤖 This comment was written by Claude Code. Changes to existing tests These are findings on the changes this PR makes to existing test files (modified files only; the newly added InventoryFour existing test files are modified, plus two internal test-infrastructure files:
Changes to expected values and assertions
Public API vs internal
Justification status of the public-API test changes
SummaryApart from the regenerated kernel-snapshot data, the only hand-written test expected value rewritten to match modified-kernel output is |
|
I'm a bit confused. I thought at first D is where 8 is sufracing to the public API. But D seems to be for integers and has nothing to do with 8. There is no test that fails for 8. Does that mean all tests for the public API that 8 surfaces already use |
|
🤖 This comment was written by Claude Code. On concern 8 — floating-point reassociation The offsets-pipeline reducer changes the order of summation, so float How it surfaces. Comparing released awkward against this PR, exact equality ( Why CI stays green. The reduction tests compare float Magnitude. The difference stays within tolerance: Disposition. This is an inherent consequence of the reassociation, accepted in review as not a bug. Because the shift stays within tolerance, nothing in the suite detects it — so a code change of this kind in an earlier release would have gone unnoticed as well, and exact float-reduction values have not been guaranteed stable across versions. |
|
🤖 This comment was written by Claude Code. Status of the eight merge-condition items Status of the eight items listed as conditions for merging — 1, 3, 6, 8, 19, 20, 23, 25 — as of head
Where this leaves the eight
|
|
Thanks @ianna for the comment and commits. So, 1, 3, 6, 8, 20, and 25 are closed. Two are open:
|
|
From my point of view,
|
@TaiSakuma - thank you so much for detailed review! IMHO about 19: for now we handle For D and E: I reverted the flat case to |
|
@TaiSakuma -- macos-latest, 3.10, full fails with: |
|
I re-ran |
|
🤖 This comment was written by Claude Code. Status: 19, D, E Status of concern 19 and of rows D and E from the test-changes comment, as of head ✅ 19 — CCCL ( ✅ D — ✅ E — |
TaiSakuma
left a comment
TaiSakuma
left a comment
There was a problem hiding this comment.
All concerns raised in this PR that risk breaking user code have been addressed.
I'll merge this tomorrow if it is still open and nobody objects.
Optionally, we can open an issue to track the remaining open concerns that have no possibility of breaking user code.
ikrommyd
left a comment
ikrommyd
left a comment
There was a problem hiding this comment.
Not actually requesting any changes here but I think @ariostas and me should do another pass here after all the recent commits and additionally I think this needs some physics validation too before it's merged. I can definitely run HiggsDNA and we can ask other framework authors to run their CIs against this branch.
Thank you Ianna and Tai for the recent progress and active development and feedback here.
I had time and have merged the PR. You had the opportunity to review the substantive changes prior to Friday, and the subsequent commits affected only the test suite. As this PR was blocking further merges in the workflow, integration at this stage was required. Please proceed with any additional validation on your side. |
|
What you actually said earlier above was to "Please post your thoughts here by the end of Friday" regarding whether the migration from parents to offsets should go into the main branch to which I voted yes. You did not say "you have until Friday to review and then I'll merge". In any case, physics validation should definitely have taken place here. It's definitely not nice to potentially silently break physics results if there was a bug here. Luckily, coffea, dask-awkward and HiggsDNA are fine so far. We'll see what the integration tests say soon. It's not the responsibility of analysis framework maintainers to deal with the fallout if there are any silent bugs here. It's awkward's responsibility to ensure that this doesn't happen as much as we can with testing/reviewing. The good part is that so far from testing for the past hour, I haven't seen anything yet. |
I think my earlier message could have been clearer. My intention was to ask for feedback on whether the migration should go into the main branch, not to promise that I would wait until Friday before merging. I can see how the wording could reasonably have been interpreted that way. That said, the migration itself had been open for review for several weeks, followed our normal review process, and required a review before it could be merged. It wasn't a last-minute or silent change. If your concern is about the more recent changes, those were confined to the test suite and were made in response to review and validation work. They didn't change the implementation itself. I completely agree that avoiding regressions in downstream physics analyses is important. That's why the implementation had unit tests, underwent human review, received additional scrutiny from AI-assisted code reviews, and is now being validated against downstream projects. Integration testing is an important part of building confidence, and I'm glad coffea, dask-awkward, and HiggsDNA appear to be unaffected so far. |
Just to add one point on wording and interpretation: phrases like “silent bugs” and suggestions that physics validation “should definitely have taken place” can imply that no validation was done on our side. That is not accurate. All changes went through the standard review process, including required approval before merge, and were supported by unit tests and additional validation work. Downstream integration testing is of course essential, but it is not correct to frame this as if there was no prior physics validation effort. Going forward, I think it would help to be more precise in distinguishing between (1) potential gaps in downstream coverage and (2) the existence of validation and review before merge. Conflating the two makes the discussion harder than it needs to be. |
|
@ikrommyd - I'm curious to know why "physics validation should definitely have taken place" on this PR. How does the validation work? How do you run it? Are the validation tools public that I can look at? |
|
The reason I say "silent bugs" is that the only thing that can actually be wrong in this PR when all the CIs are green is something that is untested and the result could have just changed without error without us knowing. I did not say anywhere that this hasn't been tested within the awkward world. It had tests and benchmarks but my understanding is that we didn't do any testing here of the broader scikit-hep ecosystem or analysis frameworks. @TaiSakuma because the only thing that could potentially be wrong here is just silently wrong numbers. I don't think that it happened but if it did, it would be pretty bad for physicists. And that's why the broader scikit-hep ecosystem and analysis frameworks should have been validated. As to how to do this, for packages you just run their test suite from awkward installed from this branch. For more complicated physics frameworks you cannot so easily so you need to ask the maintainers to try out awkward from this branch and report back if they get identical physics results. This is what I do for HiggsDNA for example, I validate that the numbers are identical down to floating point precision in its output. |
@ikrommyd — your comments came across as a personal attack. The claims about carelessness and “silent wrong numbers” are not supported by the facts. I need you to apologise for that framing. |
6 participants
(Back-port from the
awkward3branch.)What this changes
Reducer, sort, and unique machinery now passes contiguous
offsetsranges to kernels instead of per-elementparentslabels.The offsets contract. Every kernel that previously received
parentsnow receivesoffsetsof lengthoutlength + 1satisfying:offsets[0] == 0,offsets[outlength] ==total number of rows covered, and monotone non-decreasing (empty bins allowed). Binkis the contiguous range[offsets[k], offsets[k+1]).awkward-cppagainst the new contract;kernel-specification.ymlandkernel-test-data.jsonmigrated to match._reduce_next/_sort_next/_uniquesignatures (parentsparameter removed).ak.index.EmptyIndex/ak.index.ZeroIndexremoved, andak.index.LazyIndex(now unused) along with them (note: these were importable in released 2.x — low but nonzero third-party risk; see "Behavior changes").cuda.compute(CCCL Python); see "Decisions for maintainers to confirm" below.IndexedArray._uniquesimplified to carry-and-delegate, replacing ~120 lines of parents-era logic. Now covered by a regression test (tests/test_4056_indexedarray_unique.py) asserting the delegated result matches the materialized (project()) array across reordering, repeats, both axes, several dtypes, the empty fast-path, and a nested-list case — maintainer sign-off still welcome.Threading. This revision keeps the rewritten kernels serial; it does not include the OpenMP parallelization prototyped earlier. As a result, published
manylinuxwheels remain single-threaded and the auditwheel-bundled-libgomp / dask-oversubscription consideration does not apply at this head. It would return if OpenMP is reintroduced, and should be a deliberate decision at that point.Behavior changes
ak.sum/ak.prod(and quantities derived from them) onfloat32/float64/complex data can differ at the bit level from previous releases. Results stay correct to within normal rounding —ak.isclose/numpy.isclosecontinue to pass — but exact comparisons (==,ak.array_equalwithout a tolerance) may now report a difference.min,max,argmin,argmax,count, andcountnonzeroare unaffected (order-independent or exact). Code relying on bit-for-bit reproducibility against an older Awkward should compare with a tolerance.ak.index.EmptyIndex,ak.index.ZeroIndex, andak.index.LazyIndex. Third-party code importing them directly will raiseImportError. No supported, documented path used them.Decisions for maintainers to confirm
cuda_kernels_implindev/generate-kernel-signatures.py), so reducer dispatch on CUDA always usescuda.compute. The CUDA backend already relies on an undeclared, user-supplied prerequisite — CuPy is not declared inpyproject.tomleither — so this follows existing precedent rather than introducing a new pattern. What changes is that the required-but-undeclared set grows from{CuPy}to{CuPy, CCCL}, now covering reductions and ~70 structural kernels, with no CuPy-only fallback if CCCL is absent. CCCL is documented for CI inrequirements-test-gpu.txt(cuda-cccl>=1.0.2). Whether to declare it (e.g. as a GPU extra), re-enable the hand-written kernels, or delete them is left open here, pending GPU validation.awkward_sorting_ranges,awkward_reduce_local_nextparents_64, …) remain in the spec with no live CPU caller — kept for now as deprecated; can be dropped in a follow-up.Release management
awkward-cpp53 → 54 inawkward-cpp/pyproject.tomland the matching pin inpyproject.toml(both still read53). The C++ kernel signatures changed, so anawkward-cpp54 must be published before releasing thisawkward; otherwiseawkward_cpp==53pulls the old kernels against the new Python and breaks on install.Testing
tests-spec,tests-cpu-kernelspass; migratedkernel-test-data.jsonentries validated (offsets well-formed, empty bins → identity, argmin/argmax absolute-index conventions).tests/test_4056_record_reducer_shifts.py— positional record reducers with missing values inside lists, incl. typetracer.tests/test_4056_indexedarray_unique.py—IndexedArray._uniqueparity withproject()(regression cover for the carry-and-delegate rewrite).awkward_RegularArray_reduce_nonlocal_preparenext_64(7 cases covering uneven and empty bins) generates coverage on CPU and CUDA backends.awkward_sortkernel guards the empty-offsets case (offsetslength == 0) and clamps bin bounds againstlength, closing a potential out-of-bounds read on malformed offsets.tests-cuda,tests-cuda-kernels*) still need a run on real hardware before merge.tests-cuda/test_3149...test_block_boundary_prod_complex13fixes a pre-existing out-of-bounds read.AI disclosure
Parts of this PR were developed with AI assistance (per CONTRIBUTING.md).