Repository navigation
Conversation
Benchmark ResultsShow table
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
vchuravy
left a comment
There was a problem hiding this comment.
Very neat. One comment on the "magic number"
@kernel tile=N@kernel tile and @kernel tile=N
|
I tried tiles on DGSEM volume integral kernels as in Trixi.jl (tensor-product elements with Nq = 4 nodes per direction, so a 2D element has 16 nodes and a 3D slab 16, and a 2D element maps onto one tile). The benchmarks used c5b02cf with CUDA.jl's 1. Several sub-group-sized-or-smaller tiles per workgroup (N < W).
With these, 2. N-D ndrange / tile shape. Tiles need a 1-D ndrange. DG kernels are naturally written with
When Nq² does not divide the tile width (Nq = 3, 5), some lanes are in range but map to no node. This works because the shuffles take the lane modulo N. Documenting that idiom (clamp the loads, guard the stores, keep all lanes in the shuffles) would help. 3.
4. Struct shuffles on PoCL cost one work-group barrier per field. On PoCL, an
(Investigated with the assistance of Claude Code.) |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## vc/groupreduce #864 +/- ##
==================================================
- Coverage 82.49% 81.91% -0.58%
==================================================
Files 28 29 +1
Lines 2451 2594 +143
==================================================
+ Hits 2022 2125 +103
- Misses 429 469 +40 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
I let Claude rework this PR, #831, and the back-end PRs around those four points: 1. Several narrow tiles per sub-group. Tiles narrower than the sub-group now share it, though not quite as you sketched. Tiles may diverge: a tile-uniform
On an RTX 5080, a kernel like yours (ndrange 2. N-D ndranges. I took your first option:
3. Inlining. Right, with a narrower cause: 4. Struct shuffles on PoCL. Structs are now packed into 32-bit words, with small fields sharing a word. The words go through a hook that KA's PoCL backend implements with |
With a KernelAbstractions that has tiles (#864), add DGSEM kernels in which the N² work-items of an element (2D) or of a slab (3D) form a tile and exchange values with `tile_shfl` instead of `@localmem` and `@synchronize`: the 2D weak form, the 2D half-sweep flux differencing (each two-point flux computed once, cyclically, for uniform control flow) and the 3D slab kernel. "weak/node localmem tile_barrier" keeps the local memory and replaces `@synchronize` with `tile_barrier`. The kernels live in benchmark/dgsem_tiles.jl, which is only included if KernelAbstractions defines `tile_shfl`, so the suite still runs on versions without tiles. Assisted-by: Claude Code (Opus 5.5)
With a KernelAbstractions that has tiles (#864), add DGSEM kernels in which the N² work-items of an element (2D) or of a slab (3D) form a tile and exchange values with `tile_shfl` instead of `@localmem` and `@synchronize`: the 2D weak form, the 2D half-sweep flux differencing (each two-point flux computed once, cyclically, for uniform control flow) and the 3D slab kernel. "weak/node localmem tile_barrier" keeps the local memory and replaces `@synchronize` with `tile_barrier`. The kernels live in benchmark/dgsem_tiles.jl, which is only included if KernelAbstractions defines `tile_shfl`, so the suite still runs on versions without tiles. Assisted-by: Claude Code (Opus 5.5)
With a KernelAbstractions that has tiles (#864), add DGSEM kernels in which the N² work-items of an element (2D) or of a slab (3D) form a tile and exchange values with `tile_shfl` instead of `@localmem` and `@synchronize`: the 2D weak form, the 2D half-sweep flux differencing (each two-point flux computed once, cyclically, for uniform control flow) and the 3D slab kernel. "weak/node localmem tile_barrier" keeps the local memory and replaces `@synchronize` with `tile_barrier`. The kernels live in benchmark/dgsem_tiles.jl, which is only included if KernelAbstractions defines `tile_shfl`, so the suite still runs on versions without tiles. Assisted-by: Claude Code (Opus 5.5)
With a KernelAbstractions that has tiles (#864), add DGSEM kernels in which the N² work-items of an element (2D) or of a slab (3D) form a tile and exchange values with `tile_shfl` instead of `@localmem` and `@synchronize`: the 2D weak form, the 2D half-sweep flux differencing (each two-point flux computed once, cyclically, for uniform control flow) and the 3D slab kernel. "weak/node localmem tile_barrier" keeps the local memory and replaces `@synchronize` with `tile_barrier`. The kernels live in benchmark/dgsem_tiles.jl, which is only included if KernelAbstractions defines `tile_shfl`, so the suite still runs on versions without tiles. Assisted-by: Claude Code (Opus 5.5)
With a KernelAbstractions that has tiles (#864), add DGSEM kernels in which the N² work-items of an element (2D) or of a slab (3D) form a tile and exchange values with `tile_shfl` instead of `@localmem` and `@synchronize`: the 2D weak form, the 2D half-sweep flux differencing (each two-point flux computed once, cyclically, for uniform control flow) and the 3D slab kernel. "weak/node localmem tile_barrier" keeps the local memory and replaces `@synchronize` with `tile_barrier`. The kernels live in benchmark/dgsem_tiles.jl, which is only included if KernelAbstractions defines `tile_shfl`, so the suite still runs on versions without tiles. Assisted-by: Claude Code (Opus 5.5)
Add `shfl`, `shfl_up` and `shfl_xor` next to `shfl_down`, their variants within segments of a `width`, and the votes `sub_group_any`, `sub_group_all` and `sub_group_ballot`. Back ends implement the shuffles for their native primitive types; other primitive types are shuffled as `UInt32` or `UInt64` words and `isbits` structs field by field, with `supports_shuffle` derived accordingly. Which work-items form a sub-group stays unspecified. Back ends that form them from consecutive work-items report `supports_linear_subgroups`, and those whose sub-groups can communicate in control flow that differs between the sub-groups of a work-group report `supports_independent_subgroups`. The testsuite computes the expected results from the sub-groups that formed, so that it runs without either guarantee. The POCL back end implements the shuffles and votes with SPIRVIntrinsics, for widths of up to 64 so that the ballot covers every lane. PoCL forms sub-groups from consecutive work-items, but executes sub-group operations as work-group barriers, so it reports linear but not independent sub-groups. Co-authored-by: Tim Besard <tim@juliahub.com>
Add `sub_group_reduce`, `sub_group_scan` and `sub_group_exclusive_scan`, in the order of the lanes, for associative operators that needn't be commutative. The fallbacks are built on the shuffles. Back ends may implement them natively per operator and type, as long as the result is the documented one, except that floating-point `+` and `*` may be reassociated. The POCL back end uses PoCL's native reductions and scans where they have Julia's semantics: `+` on integers and floats, and `min`/`max` on integers. Co-authored-by: Tim Besard <tim@juliahub.com>
@kernel tile and @kernel tile=N@kernel tile
Structs and tuples were shuffled field by field, a shuffle per primitive
field, so that e.g. an `NTuple{2, Float16}` took two. Pack their fields
into `UInt32` words instead, with fields of less than 4 bytes sharing a
word, and shuffle the words through `shfl_word_tuple`. A backend can
implement that to exchange all words at once: on PoCL, where every
shuffle is a work-group barrier, that would help, once its CPU device
accepts vector shuffles from SPIR-V.
The generated fallbacks also weren't inlined: `@inline` on a generated
function doesn't apply to a body that the generator returns with an
explicit `return`, so shuffling an `SVector` in a larger kernel left a
call, with the value passed through memory. They now mark their bodies
for inlining.
Tiles narrower than a sub-group have to share it, and to be useful they have to be able to diverge from each other, like CUDA's `tiled_partition<N>`. Add the votes, collectives and barrier within segments of `width` lanes (the shuffles already had a `width`), and `supports_partitionable_subgroups`: whether a sub-group can be partitioned into segments that diverge independently, so that only the work-items of a segment have to execute these together. It extends `supports_independent_subgroups` to segments, and implies it. The fallbacks build on the functions without a `width`, so that they are correct wherever the whole sub-group executes them. Backends that report partitionable sub-groups have to make them work for a segment alone: CUDA needs a mask of the segment's lanes, other backends a ballot that only reports active lanes and a barrier that a segment can execute. OpenCL and Metal have no barrier for part of a sub-group, so there it takes hardware that executes a sub-group in lock-step, like Intel's and Apple's GPUs, and a fence. PoCL never reports it, as its sub-group functions are work-group barriers. The shuffles with a `width` of other types now shuffle `UInt32` words with a `width`, so that a backend's shuffles with a `width` handle them.
Every shuffle on PoCL is a work-group barrier, so a struct packed into words cost a barrier per word: tile kernels shuffling an `SVector` were several times slower than ones using local memory. Exchange the words with the vector shuffles of `SPV_INTEL_subgroups`, a barrier per up to 16 words, which PoCL's CPU device accepts with `cl_intel_subgroups` since pocl/pocl#2394. Four shuffles of an `SVector{4, Float32}` over 4M work-items take 1.7 ms instead of 4.6 ms. KernelInterface gets `shuffle_source`, the lane a shuffle reads from, for backends that exchange the words with a shuffle from a lane.
Work-group reductions and scans for `@kernel`, as collectives like `@synchronize`: they are statements of their own that every work-item of the workgroup reaches, and the split of the kernel body makes the work-items that pad a partial workgroup take part, contributing the neutral element without evaluating the value. The value is converted to the type of `neutral` before the call, so that a union-typed value doesn't split the collective, and `neutral` is evaluated once. The upper bound of a dynamic workgroup size is the keyword `groupsize`. `@groupreduce` needs a commutative operator, `@groupscan` only an associative one, in the order of `@index(Local, Linear)`. Both use local memory. Co-authored-by: Tim Besard <tim@juliahub.com>
On backends with sub-groups, `@groupreduce` reduces every sub-group with `KernelInterface.sub_group_reduce`, which backends can implement natively, and then the results of the sub-groups. That only relies on the sub-group ids, so it needs neither linear nor independent sub-groups. KernelAbstractions picks the algorithm instead of a user-facing option: kernels that use collectives get the backend's sub-group capabilities in their `CompilerMetadata`, and use the sub-groups for types the backend can shuffle. Floating-point types need the backend's support for that type, since a native shuffle may be selected by type alone; other primitive types and structs are shuffled as packed words. Kernels without collectives don't query the capabilities. Co-authored-by: Tim Besard <tim@juliahub.com>
A kernel constructed with a tile width, `f(backend; tile = N)`, divides its workgroups into tiles of `N` consecutive work-items (by `@index(Local, Linear)`), like CUDA's `tiled_partition<N>`. The work-items of a tile communicate with operations that take the tile (`tile_shfl` and its variants, `tile_any`, `tile_all`, `tile_ballot`, `tile_reduce`, `tile_scan`, `tile_exclusive_scan` and `tile_barrier`), in control flow that only has to be uniform over the tile. `@tile()` returns the work-item's tile and lane, and `@index(Tile)` the global index of its tile. Tiles aren't declared in the kernel: what sets a kernel with tiles apart is its width, which the host needs for the launch geometry, and which is passed when the kernel is constructed. A kernel that uses `@tile()` or `@index(Tile)` has to get one; the `@kernel` macro notes that use, so that construction rejects a missing width. There is no default: the sub-group width differs between backends (32 or 64 on AMD GPUs), so kernels name the width they're written for. The width is a type parameter of `Kernel` (`tile_width(kernel)`) and reaches the device through the kernel's `CompilerMetadata`, where `tile_width(t)` is a constant. Tiles are a capability: a power-of-two `N` of at most the sub-group width `W`, any number of tiles per workgroup, and tiles that diverge independently, also where several share a sub-group (KernelInterface's segments). `supports_tiles(backend, N)` reports it, and constructing a kernel with tiles a backend doesn't support is an error. All GPU backends support widths up to 32. The CPU backend has no tiles, as its sub-group functions are work-group barriers, which can't diverge. Falling back to a tile per workgroup there would leak into how kernels choose their workgroup size, and be slower than a kernel using local memory, which is what such backends need instead. The ndrange can have several dimensions, e.g. one per index of the nodes of an element, if the workgroups divide it into whole tiles. Such kernels are launched on a 1-D grid, so that the tiles stay on the backend's sub-groups. Without an explicit workgroup size, the launch picks one that does; construction and launch check the geometry before compiling.
Tiles give
@kernela grouping below the workgroup: a kernel constructed with a tile widthdivides each workgroup into tiles of
Nconsecutive work-items (by@index(Local, Linear)), like CUDA'stiled_partition<N>, whose work-items exchange values without going through local memory.No declaration, a width at construction:
f(backend; tile = N), also with a staticworkgroup size and
ndrange. What sets a kernel with tiles apart is the width, which the hostneeds for the launch geometry, so that's the only switch. A kernel that uses
@tile()or@index(Tile)has to get one:@kernelnotes the use, and construction rejects a missingwidth. There is no default, since the sub-group width
Wdiffers between back ends (32 or64 on AMD); width-agnostic kernels pass
tile = KernelInterface.sub_group_size(backend).A width that follows a run-time parameter (DGSEM's polynomial degree) works the same way.
tile_width(kernel)returns the width on the host,tile_width(t)in the kernel, where itis a constant; it is a type parameter of
Kerneland reaches the device throughCompilerMetadata.Tiles are a capability. Tiles of
Nwork-items need a power-of-twoN ≤ W; aworkgroup can hold any number of them, and they diverge independently, also where several
share a sub-group (
KernelInterface.supports_partitionable_subgroups, KernelInterface: sub-group shuffles, votes, reductions and scans #831).supports_tiles(backend, N)reports whether a backend provides that, and constructing akernel it doesn't support throws, with the reason. There is no fallback.
N ≤ W(Wis 32 or 64)N ≤ 32; the narrowtile_barrieris a SIMD-group-scope fence that relies on Apple GPUs executing SIMD-groups in lock-step (KernelInterface: sub-group shuffles, votes and segments Metal.jl#998)An earlier version of this PR gave narrow tiles a workgroup each where segments couldn't
diverge, with a
tiles_per_workgroupquery. That leaked into user code (the DGSEMbenchmarks in DGSEM benchmarks: variants with tiles (#864) #874 had to query it to choose the number of elements per workgroup), and
didn't help: unpacked narrow tiles were 1.7–2x slower than packed ones on CUDA, behind the
local-memory kernels, and on PoCL every tile operation is a work-group barrier, 2–7x slower
than local memory. Kernels that should also run on the CPU need a variant without tiles,
as with oneDPL, AdaptiveCpp or Kokkos.
Launch: the workgroup has a multiple of
Nwork-items, and every tile is entirelyinside or outside of the
ndrange. N-D ndranges work, e.g.(4, 4, nelements)in workgroupsof
(4, 4, k), and are launched on a 1-D grid so that tiles stay on sub-groups. Automaticsizing picks a valid shape. Construction and launch check the geometry before compiling.
Operations, taking the tile so that helper functions can use them:
tile_shfl,tile_shfl_{up,down,xor},tile_any/all/ballot, orderedtile_reduce,tile_scan,tile_exclusive_scan,tile_barrier. They need control flow that is uniform over the tileonly: tiles can diverge, loop differently or return. They lower to KernelInterface's
sub-group operations with a width of
N.@index(Tile)numbers tiles workgroup byworkgroup.
@kernelpasses what it finds in a kernel's body (use of tiles or collectives) as methodsof the kernel's own function, so that this also works for kernels in a local scope and in
precompiled packages with overloaded kernels.
Documented in the manual under Tiles and Choosing the tile width, plus a new section, Choosing how work-items communicate, that compares local memory, the workgroup collectives, tiles and KernelInterface kernels. The 0.10 release notes have an entry.
Tested with the KernelAbstractions testsuite on PoCL (tiles throw), CUDA (RTX 5080), AMDGPU
(gfx1036) and Metal (M1 and M3). Stacked on #830.