Implement groupreduce API #559
Draft
Conversation
This file contains 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
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Benchmark Results
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
|
UPD: #218 (comment) |
| end | ||
|
|
||
| function groupreduce_testsuite(backend, AT) | ||
| @testset "@groupreduce" begin |
There was a problem hiding this comment.
Suggested change
| @testset "@groupreduce" begin | |
| return @testset "@groupreduce" begin |
| end | ||
|
|
||
| function groupreduce_testsuite(backend, AT) | ||
| @testset "@groupreduce" begin |
There was a problem hiding this comment.
Suggested change
| @testset "@groupreduce" begin | |
| return @testset "@groupreduce" begin |
| groupsizes = "$backend" == "oneAPIBackend" ? | ||
| (256,) : | ||
| (256, 512, 1024) | ||
| @testset "@groupreduce" begin |
There was a problem hiding this comment.
Suggested change
| @testset "@groupreduce" begin | |
| return @testset "@groupreduce" begin |
vchuravy
reviewed
Feb 3, 2025
src/reduce.jl
Outdated
Comment on lines
15
to
22
| - `algo` specifies which reduction algorithm to use: | ||
| - `Reduction.thread`: | ||
| Perform thread group reduction (requires `groupsize * sizeof(T)` bytes of shared memory). | ||
| Available accross all backends. | ||
| - `Reduction.warp`: | ||
| Perform warp group reduction (requires `32 * sizeof(T)` bytes of shared memory). | ||
| Potentially faster, since requires fewer writes to shared memory. | ||
| To query if backend supports warp reduction, use `supports_warp_reduction(backend)`. |
There was a problem hiding this comment.
Why is that needed? Shouldn't the backend go and use warp reductions if it can?
There was a problem hiding this comment.
I'm now doing an auto-selection of the algorithm based on device function __supports_warp_reduction().
vchuravy
reviewed
Feb 3, 2025
src/reduce.jl
Outdated
Comment on lines
70
to
77
| while s > 0x00 | ||
| if (local_idx - 0x01) < s | ||
| other_idx = local_idx + s | ||
| if other_idx ≤ groupsize | ||
| @inbounds storage[local_idx] = op(storage[local_idx], storage[other_idx]) | ||
| end | ||
| end | ||
| @synchronize() |
There was a problem hiding this comment.
(I assume this code is GPU only anyways)
vchuravy
reviewed
Feb 3, 2025
src/reduce.jl
Outdated
Comment on lines
89
to
93
| macro shfl_down(val, offset) | ||
| return quote | ||
| $__shfl_down($(esc(val)), $(esc(offset))) | ||
| end | ||
| end |
There was a problem hiding this comment.
If it isn't user-facing or needs special CPU handling you don't need to introduce a new macro
vchuravy
reviewed
Feb 3, 2025
test/groupreduce.jl
Outdated
Comment on lines
1
to
13
| @kernel function groupreduce_1!(y, x, op, neutral, algo) | ||
| i = @index(Global) | ||
| val = i > length(x) ? neutral : x[i] | ||
| res = @groupreduce(op, val, neutral, algo) | ||
| i == 1 && (y[1] = res) | ||
| end | ||
|
|
||
| @kernel function groupreduce_2!(y, x, op, neutral, algo, ::Val{groupsize}) where {groupsize} | ||
| i = @index(Global) | ||
| val = i > length(x) ? neutral : x[i] | ||
| res = @groupreduce(op, val, neutral, algo, groupsize) | ||
| i == 1 && (y[1] = res) | ||
| end |
There was a problem hiding this comment.
These need to be cpu=false since you are using non-top-level @synchronize
Comment on lines
+1
to
+2
| @kernel cpu=false function groupreduce_1!(y, x, op, neutral) | ||
| i = @index(Global) |
There was a problem hiding this comment.
Suggested change
| @kernel cpu=false function groupreduce_1!(y, x, op, neutral) | |
| i = @index(Global) | |
| @kernel cpu = false function groupreduce_1!(y, x, op, neutral) | |
| @kernel cpu = false function groupreduce_2!(y, x, op, neutral, ::Val{groupsize}) where {groupsize} |
Comment on lines
+1
to
+2
| @kernel cpu=false function groupreduce_1!(y, x, op, neutral) | ||
| i = @index(Global) |
There was a problem hiding this comment.
Suggested change
| @kernel cpu=false function groupreduce_1!(y, x, op, neutral) | |
| i = @index(Global) | |
| @kernel cpu = false function groupreduce_1!(y, x, op, neutral) | |
| @kernel cpu = false function groupreduce_2!(y, x, op, neutral, ::Val{groupsize}) where {groupsize} |
| groupsizes = "$backend" == "oneAPIBackend" ? | ||
| (256,) : | ||
| (256, 512, 1024) | ||
| @testset "@groupreduce" begin |
There was a problem hiding this comment.
Suggested change
| @testset "@groupreduce" begin | |
| return @testset "@groupreduce" begin |
Comment on lines
+1
to
+2
| @kernel cpu=false function groupreduce_1!(y, x, op, neutral) | ||
| i = @index(Global) |
There was a problem hiding this comment.
Suggested change
| @kernel cpu=false function groupreduce_1!(y, x, op, neutral) | |
| i = @index(Global) | |
| @kernel cpu = false function groupreduce_1!(y, x, op, neutral) | |
| @kernel cpu = false function groupreduce_2!(y, x, op, neutral, ::Val{groupsize}) where {groupsize} |
| groupsizes = "$backend" == "oneAPIBackend" ? | ||
| (256,) : | ||
| (256, 512, 1024) | ||
| @testset "@groupreduce" begin |
There was a problem hiding this comment.
Suggested change
| @testset "@groupreduce" begin | |
| return @testset "@groupreduce" begin |
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.
Implement reduction API. Supports two types of algorithms:
groupsize, no bank conflict, no divergence.shlf_downwithin warps: uses shmem of length32, reduction within warps storing results in shmem, followed by final warp reduction using values stored in shmem. Backends are required to only implementshlf_downintrinsic which AMDGPU/CUDA/Metal have (no sure about other backends).KA.__supports_warp_reduction().