PERF: visit the inner dimension of transposed copies in cache-sized chunks - #32672
beatakouchnir wants to merge 6 commits into
Conversation
|
What are the units for the numbers in the table? Can you |
|
@mattip, it's GB/s; apologies for leaving it out. I've updated the PR body with it as well as with the ASV measurements. |
|
Hi @crusaderky, in the issue you said you had vibe coded something. Is this similar to what you had done? And also is there any change you could potentially run the benchmark that gave you the plots in the issue using this PR branch? |
| * After the raw iterator sorts and coalesces the axes, a transposed copy | ||
| * (F->C, C->F, or any operand whose inner stride jumps between cache lines | ||
| * while its next outer stride stays within one) touches one cache line per | ||
| * inner element and then touches the very same lines again on every outer | ||
| * iteration. Once that line set outgrows the private cache, every line is | ||
| * fetched once per element it holds, and the copy runs at memory bandwidth | ||
| * divided by the number of elements per line. Splitting the inner dimension | ||
| * into chunks whose line set fits in cache keeps the lines resident across | ||
| * the outer iterations. The 1-D transfer function is unchanged; only the | ||
| * order in which the raw iteration visits the array changes. |
There was a problem hiding this comment.
This sort of paragraph comment is typical of LLM-generated code. Please rewrite all comments in this PR in your own words and try to limit yourself to comments that would be useful for a future reader. That means comments shouldn't narrate the implementation (the code should speak for itself IMO) and shouldn't discuss how the codebase used to work. Comments should be reserved for documenting non-obvious non-local facts that aid understanding.
LLM-generated text is very difficult to digest for humans, try to make sure you understand it well enough to explain it in your own words.
There was a problem hiding this comment.
Thank you for your feedback, @ngoldbaum. I have updated the code comments and release note to only carry non-obvious, non-local information.
|
Can you merge with or rebase on current |
There was a problem hiding this comment.
The performance improvement is real but IMO this code change is in the wrong spot.
It would be better to add a tiled traversal mode to the iterator machinery in NumPy. Code paths could then internally opt in to the tiled mode.
I also think that this needs much more careful testing on a wide variety of hardware before merging.
I also want to caution you about opening vibe-coded PRs in projects you don't regularly contribute to. Especially for big, possibly risky changes like this it's much better to have a discussion about the pros and cons of various designs before just going and writing code. LLMs make it way too easy to write code that solves a problem but not necessarily the problem.
| #define NPY_COPY_CACHE_LINE 128 /* covers both 64- and 128-byte lines */ | ||
| #define NPY_COPY_CHUNK_BUDGET (256 * 1024) /* 256 KiB chosen by measurement */ |
There was a problem hiding this comment.
I wouldn't be at all surprised to learn that these choices are hardware-dependent and need tuning. This whole PR needs validation on a broad range of hardware IMO.
There was a problem hiding this comment.
Good point; I'm only able to test on my own M5 Max, but happy to change these if other contributors can provide their measurements.
| if (ndim >= 2) { | ||
| chunk = transposed_copy_chunk(n_inner, src_strides_it, dst_strides_it); | ||
| } |
There was a problem hiding this comment.
Doing this for all dtypes leads to observable behavior changes:
import numpy as np
for dtype in (object, np.dtypes.StringDType()):
src = np.full((4097, 16), 1, dtype=dtype).T
src[0, 2048] = "bad"
dst = np.zeros(src.shape, dtype=np.int64)
try:
np.copyto(dst, src, casting="unsafe")
except ValueError:
print(np.count_nonzero(dst))
# Before: 2048, 2048
# This PR: 32768, 32768
you should probably limit this optimization to operations that are known statically to be safe (e.g. only numeric built-in dtypes perhaps). Future work could enable the optimization for more dtypes after auditing for issues like this.
There was a problem hiding this comment.
Fixed; chunking is now limited to built-in numeric and bool dtypes.
…hunks A transposed copy (F->C, C->F, or any operand whose inner stride crosses cache lines while its next outer stride stays within one) touches one line per inner element and touches the same lines again on every outer iteration. Once that line set outgrows the private cache, each line is fetched once per element it holds and the copy runs at memory bandwidth divided by the elements per line: on an M5 Max, out[:] = a.T on a (rows, 128) float64 array falls from 12.7 GB/s at 256 KiB to 1.2 GB/s at 256 MiB while a C->C copy holds 60-80 GB/s. raw_array_assign_array now visits the inner dimension in chunks of 2048 elements when either operand shows that stride pattern and the line set would exceed 256 KiB, so the lines stay resident across the outer iterations. The 1-D transfer function and the raw iteration over the outer dimensions are unchanged; a copy that does not match the pattern takes the previous single-chunk path. Measured on the same machine: 6.6 GB/s at 256 MiB for both F->C and C->F (5.6x), square transposes 5.0 -> 7.7 GB/s, sizes below 4 MiB unchanged. A budget sweep from 64 KiB to 1 MiB put 256 KiB within 10% of the best for the reporter's shape and best for square shapes. Tests cover both directions at inner lengths around the chunk boundary for seven dtypes including object and void, negative strides and offsets, a three-dimensional non-coalescable case, and object reference counts. Closes numpygh-32453. Co-Authored-By: Claude Fable 5.1 <[email protected]>
The chunk loop wrote shape_it[0] and the data pointers per chunk even when there was one chunk, which cost a few nanoseconds per inner call on small and in-cache copies (2-D tiny copy +6%, 100x100 transpose +10% in an A/B against main). The iteration is now a static inline helper; the unchunked case calls it once with the original arguments, and only real chunking pays for the bookkeeping. transposed_copy_chunk also tests the cheap size condition before the stride pattern. Co-Authored-By: Claude Fable 5.1 <[email protected]>
…iplying, which can overflow a 32-bit npy_intp
I don't want to take the discussion some place it shouldn't be, so about tiled traversal that should be a separate issue but I just wanted to say that adding tilted traversal would probably be hard to implement and review so if this is the ONLY consumer ever, I'd bet that it's probably not worth it. At the same time tiled traversal could enable other copy entry points to use it and can also be used on elementwise ufuncs of mixed layouts like |
A cast that can fail part-way (object, StringDType, same-value casting) now takes the original loop, so the elements written before the error are the same as before numpygh-32672. The release note says numeric arrays. Co-Authored-By: Claude Fable 5.1 <[email protected]>
46897bf to
33f0a1c
Compare
I personally don't think we should spend effort replacing use of numpy's interation infrastructure with one-off hacks. Instead we should improve the infrastructure. It's not just this once place that has this problem: any operation that needs an iterator over a transposed source has this issue. |
|
Rebased on current Regarding the approach, I followed the copy-specific fast path suggestion made by @seberg in #32453, but if the maintainers choose to go in a different direction, I'm happy to close this PR and contribute my measurements to the issue. |
|
I think given Sebastian said it might be an interesting vibe experiment I'd also like to see someone try the more general approach. Maybe a minimal version of that could allow someone to tackle this more generally later? |
|
I would be very curious if you can get this into That said, this code path doesn't actually hit That said, I believe this code needs to do a bit deeper; making this nice isn't just a quick vibe-code but needs thought. My main two things to consider:
So, I think that probably means a helper/option for create-sorted-strides that is necessary to get the best performance here by just chunking up the last dimension (in many cases!). |
|
I can take a stab at the helper if I can get consensus that this is the preferred approach. I just locally tested a sorted-strides helper for 3-D - 6-D copies, currently unimproved by this PR, and saw speedups of 2.1-3.8x:
(F->C copies of float64 arrays, 70-130 MiB each, M5 Max; C->F within 10% of each row) The C->C copy speed is unchanged from
One limitation I have discovered on my machine is that source strides that are a multiple of the page size (16 KiB), a 4096 x 4096 float64 transpose being the common case, don't see a speedup without padding:
This is unchanged from |
This branch has not been deployed
5 participants
PR summary
Fixes #32453
Copying a transposed C-order (rows, 128) float64 array into a contiguous destination makes the inner loop stride through the source at 1024 bytes, touching one cache line per element while using only 8 bytes of it; the next outer pass revisits the same lines 8 bytes over. While all
rowslines fit in the private cache, each is fetched once and reused across 16 (8 on x86) outer passes; once they don't, each line is re-fetched per element it holds, so the copy reads 16x (8x) more from memory than it writes. C→F has the same problem on the destination side via write-allocate.transposed_copy_chunkdetects this pattern when, after sorting and coalescing, either operand's inner stride is at least a cache line (128 bytes) and its next-outer stride is smaller; if the inner length times line size also exceeds 256 KiB, the inner dimension is processed in 2048-element chunks, each running the full outer iteration withshape_it[0]and the data pointers adjusted, so a chunk's lines stay cache-resident. The transfer function, outer iteration order, overlap, and threading handling are unchanged, and non-matching, 1-D, or short copies take the old path as a single chunk.raw_array_wheremasked_assign_arrayandNpyIter-based copies are untouched.Copy throughput in GB/s (bytes per nanosecond), main → this branch, best of 3 runs with a 256 MiB cache flush before each, single core on an M5 Max:
Sizes of 4 MiB and below are unchanged within noise.
ASV, interleaved, three rounds, main vs the pushed head:
First time contributor introduction
I'm an applied ML scientist and have been using numpy for over a decade.
AI Disclosure
Claude was used to write the code and take the measurements; I have reviewed the code.