Skip to content

Forward to a lower-dimensional kernel for trailing singleton dimensions - #98

Draft
lkdvos wants to merge 3 commits into
mainfrom
ld-kernel-trim
Draft

lkdvos wants to merge 3 commits into
mainfrom
ld-kernel-trim

Conversation

@lkdvos

@lkdvos lkdvos commented Oct 29, 2026

Copy link
Copy Markdown
Member

Builds on #97.

order_and_fuse_dims keeps fused dimensions as trailing dimensions of size 1, and the kernel still runs a loop for each of them. This PR makes the kernel forward to the kernel for its non-singleton dimensions instead.

Compiling a lower-dimensional kernel for every dimensionality would add to the compile time of every first call. To limit that, the views are rewrapped with the lower dimensionality via StridedViews.unsafe_stridedview, so the forwarded kernel is the same method instance that arrays of that dimensionality compile. The kernel only indexes the parents with the strides it is given, so the view strides need no normalization. This requires StridedViews 0.5.5.

Trade-off

The runtime gain is modest, and the compile-time cost is real, so this is a draft to discuss whether it is worth it.

Runtime, C = α * permutedims(A, p) via _mapreducedim!, ns, against #97:

shape elements #97 Float64 PR Float64 #97 ComplexF64 PR ComplexF64
(2, 2, 1, 1, 1, 1) 4 182 164 208 181
(2, 2, 2, 2, 1, 1) 16 208 215 233 247
(2, 2, 2, 2, 2, 2) 64 301 308 344 350

The difference beyond 4 elements is within the noise. The gain is larger for direct kernel calls that leave many trailing singletons, such as the permutedims! in the follow-up PR.

First calls in a fresh session:

main #97 PR
one permutation per rank 6, 5, 4, 3, 2 (same function and element type) 2.25 s 2.10 s 2.17 s
permutation, broadcast, reduction, map! (different functions and types) 2.35 s 2.27 s 3.77 s

The rewrapping makes the lower-dimensional kernels shared when the function and element types repeat, but each distinct call still compiles about 1.5× more. The test suite also takes noticeably longer to compile.

The tests pass, single- and multi-threaded (4 threads), including the JLArray and Aqua tests.

🤖 Generated with Claude Code

lkdvos and others added 3 commits October 29, 2026 09:53
Constructing the step range `1:blocks[i]:dims[i]` computes its last element with an integer
division, which is paid for every block loop in the nest, also when the arrays fit in a
single block. A while loop only needs comparisons and additions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nels

Whether the innermost dimension has unit stride in all arrays is loop invariant, but it was
checked every time the innermost loop started. For arrays with short innermost dimensions
this is a large part of the work. Non-reducing kernels now generate the loop nest for both
cases and select one up front.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Fused dimensions are kept as trailing dimensions of size 1, which still cost a loop each in
the kernel. The kernel now forwards to the kernel for its non-singleton dimensions. The views
are rewrapped with the lower dimensionality, so that this kernel is the same one that arrays
of that dimensionality compile, rather than one extra kernel per dimensionality and type.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@codecov

codecov Bot commented Oct 30, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.30769% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/mapreduce.jl 92.30% 2 Missing ⚠️
Files with missing lines Coverage Δ
src/mapreduce.jl 93.63% <92.30%> (-0.12%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Base automatically changed from ld-small-kernel-overhead to main November 3, 2026 12:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant