Skip to content

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

Draft
lkdvos wants to merge 1 commit into
ld-small-kernel-overheadfrom
ld-kernel-trim
Draft

lkdvos wants to merge 1 commit into
ld-small-kernel-overheadfrom
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

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>
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