Repository navigation
Conversation
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 10, 2026 17:49
e035493 to
05b0448
Compare
maleadt
added this pull request to stack #1019
October 10, 2026 17:49
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## ka-0.10 #998 +/- ##
===========================================
- Coverage 82.95% 82.68% -0.28%
===========================================
Files 86 86
Lines 6173 6231 +58
===========================================
+ Hits 5121 5152 +31
- Misses 1052 1079 +27 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 10, 2026 18:46
05b0448 to
7cb4016
Compare
KernelAbstractions 0.10 builds on KernelInterface, which defines what a back end provides: memory and device management, compiling and launching kernels, and the device-side intrinsics. KernelAbstractions then launches `@kernel` kernels itself on any KernelInterface back end, which replaces Metal's copy of that launch path: partitioning the ndrange, building the kernel's context, and sizing the threadgroups. `MetalBackend` now implements KernelInterface, and Metal depends on it instead of on KernelAbstractions. What KernelAbstractions still needs from a back end moves to an extension: the stack allocation behind `@private`, and the Adapt rule for moving arrays to the CPU. KernelInterface's default `launch_configuration` uses the pipeline's `maxTotalThreadsPerThreadgroup`, as Metal's launch did. `KI.kernel_function` receives the callable unconverted, and keeps it as the kernel's source, which is converted again at every launch, as with `@metal`: the buffers a closure captures are declared to the encoder and kept alive. `KI.launch` passes the kernel arguments on as a tuple, to the same launch function `HostKernel` calls, so kernels with many arguments aren't splatted. Like `HostKernel`, it first waits for another task's queue that still uses one of the arguments' buffers. It passes Metal's `queue` and `submit` options on, and rejects `threads` and `groups`, which would override the launch geometry that KernelInterface validated. `KI.copyto!` also accepts contiguous views of host arrays, and `adapt(MetalBackend(), x)` moves any array to the GPU, as `adapt(MtlArray, x)` does, rather than only `Array`s. KernelAbstractions now computes indices in 32 bits where they fit, so its conversion tests pass and are no longer skipped. Its test of a kernel with 41 arguments is skipped instead, since a Metal kernel takes at most 31 buffers. Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Sub-groups are SIMD-groups: implement the sub-group queries, `sub_group_barrier` and `shfl_down`, and report their support to KernelInterface, with `shfl_down` for the types `simd_shuffle_down` supports. The sub-group width that `KI.sub_group_size` promises is 32, which `kernel_function` checks against the pipeline's `threadExecutionWidth`. KernelInterface leaves unspecified how work-items are grouped into sub-groups; Metal forms SIMD-groups from consecutive linear thread indices, so the last one of a threadgroup can be partial, which a test checks. Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Implement `KI.record_event` with an `MTLSharedEvent` that a command buffer signals after the task's queued work, and `KI.wait_event` by encoding a wait for it in the task's open batch of work, so that the GPU waits instead of the host. `KernelAbstractions.@spawn` uses them to order a new task's work after the work its parent had queued, without synchronizing the parent; the default is a full synchronization. KernelInterface's testsuite checks them once `record_event` returns an event. Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
`KI.versioninfo(MetalBackend())` prints `Metal.versioninfo()`.
… main branch Neither is registered yet. Julia 1.10 ignores [sources] and Julia 1.11 doesn't support workspaces, so Buildkite sets up the test environment by hand there, and runs the tests in it directly. Drop this commit once both are registered.
KernelAbstractions now backs @Private itself, with an overlay in GPUCompiler's shared method table that returns a PrivateArray over a stack alloca, so drop our Scratchpad override. This relies on Metal's compiler consulting that shared method table, i.e., declaring its tables through GPUCompiler.method_tables rather than overriding method_table_view.
Get KernelAbstractions and KernelInterface from KernelAbstractions.jl's `vc/ki-subgroup-ops` branch instead of `main`. Julia 1.10 ignores [sources], so Buildkite checks out and develops them there. Drop this commit once JuliaGPU/KernelAbstractions.jl#831 is merged. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 10, 2026 20:28
7cb4016 to
a5f25ee
Compare
Implement the sub-group communication of KernelInterface (JuliaGPU/KernelAbstractions.jl#831): - `shfl`, `shfl_down`, `shfl_up` and `shfl_xor` map to the SIMD-group shuffles, only for the primitive types Metal supports natively, so that other types reach KernelInterface's fallbacks (64-bit integers as `UInt32` words, structs field by field). Offsets past the SIMD-group width return the work-item's own value, also ones that don't fit in the intrinsics' `Int16`. - `sub_group_ballot` is `simd_ballot` masked to the lanes of the SIMD-group, since MSL leaves the upper bits undefined, and `sub_group_any`/`sub_group_all` are derived from it. - `get_max_sub_group_size` returns the constant `SIMD_WIDTH`, which `kernel_function` checks the pipeline's execution width against. - Sub-groups are reported as linear and independent: Metal doesn't specify the layout, but Apple GPUs form SIMD-groups from consecutive threads, which the testsuite checks. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 12, 2026 07:52
a170129 to
527def0
Compare
maleadt
changed the title
KernelInterface's functions with a `width` only need the caller's segment of the SIMD-group if the back end reports `supports_partitionable_subgroups`, which lets KernelAbstractions pack several narrow tiles into a SIMD-group. The shuffles and votes already only involve the active threads. MSL has no barrier for part of a SIMD-group, so the narrow `sub_group_barrier(width)` is a SIMD-group-scope fence and relies on Apple GPUs executing a SIMD-group with one program counter, so that the lanes of a segment in segment-uniform control flow execute in lock-step. That is the hardware's execution model (as documented by the Asahi reverse engineering) rather than an MSL guarantee; a stress test of diverging segments exchanging values through threadgroup and device memory found no errors in over 10^10 checks on M1 and M3, and the fence keeps the compiler from moving memory accesses across it.
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 12, 2026 09:08
527def0 to
906b562
Compare
An error occurred while trying to automatically change base from
ka-0.10
to
main
October 13, 2026 06:37
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements KernelInterface's sub-group communication (JuliaGPU/KernelAbstractions.jl#831)
on Metal:
shfl,shfl_down,shfl_up,shfl_xorfor the types Metal shuffles natively(
Float32,Float16, 8/16/32-bit integers; other types through KernelInterface'sfallbacks, as packed
UInt32words),sub_group_ballot/any/all, and a constantget_max_sub_group_size.supports_linear_subgroupsandsupports_independent_subgroupsaretrue.supports_partitionable_subgroupsistrue, so that KernelAbstractions can pack severaltiles narrower than 32 into a SIMD-group. The shuffles and votes with a
widthonlyinvolve the active threads already. MSL has no barrier for part of a SIMD-group, so
KI.sub_group_barrier(width)forwidth < 32is a SIMD-group-scopeatomic_thread_fenceon device and threadgroup memory, which relies on Apple GPUsexecuting a SIMD-group with one program counter (divergent paths are serialized, so the
lanes of a segment in segment-uniform control flow execute in lock-step). That is the
hardware's execution model, as described by the Asahi reverse engineering, not an MSL
guarantee. At width 32 it stays
simdgroup_barrier.The lock-step assumption is checked by a stress test of diverging segments that repeatedly
exchange values through threadgroup and device memory, with segment-dependent trip counts
and early exits: no errors in over 10^10 checks each on an M1 (G13) and an M3 Pro (G15),
while the controls (reading before writing, reading another segment's slot) do fail.
KernelInterface's testsuite has a smaller version of it. Apple's compiler emits no
instruction for the fence; it keeps LLVM and the AIR compiler from moving memory accesses
across it. M2 and M4 are untested, and Buildkite's Metal agents are all M1s.
Stacked on #994. The
[TEMP]commit goes once JuliaGPU/KernelAbstractions.jl#831 is merged.