-
-
Notifications
You must be signed in to change notification settings - Fork 13
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Update src/gather.jl Co-authored-by: Carlo Lucibello <carlo.lucibello@gmail.com> replace _check_dims as gather_check_dims fix update manifest.toml
- Loading branch information
Showing
5 changed files
with
157 additions
and
24 deletions.
There are no files selected for viewing
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| function gather_check_dims(X::AbstractArray{Tx,Nx}, | ||
| Y::AbstractArray{Ty,Ny}, | ||
| idx::AbstractArray{Tidx,Nidx}) where | ||
| {Tx,Ty,Tidx<:IntOrIntTuple,Nx,Ny,Nidx} | ||
| M = NNlib.typelength(Tidx) | ||
| dims = gather_check_dims(Nx, Ny, M, Nidx) | ||
| size(X)[1:dims] == size(Y)[1:dims] || throw(ArgumentError("Incompatible input shapes.")) | ||
| size(Y)[dims+1:end] == size(idx) || throw(ArgumentError("Incompatible input shapes.")) | ||
| return dims | ||
| end | ||
|
|
||
| function gather_check_dims(X::AbstractArray{Tx,Nx}, | ||
| Y::AbstractArray{Ty,Ny}, | ||
| idx::AbstractArray{CartesianIndex{M},Nidx}) where | ||
| {Tx,Ty,Nx,Ny,M,Nidx} | ||
| dims = gather_check_dims(Nx, Ny, M, Nidx) | ||
| size(X)[1:dims] == size(Y)[1:dims] || throw(ArgumentError("Incompatible input shapes.")) | ||
| size(Y)[dims+1:end] == size(idx) || throw(ArgumentError("Incompatible input shapes.")) | ||
| return dims | ||
| end | ||
|
|
||
| function gather_check_dims(Nx, Ny, M, Nidx) | ||
| @assert Nx - M == Ny - Nidx "Incompatible input shapes of (dst, src, idx) = ($Nx, $Ny, $Nidx)." | ||
| dims = Nx - M | ||
| dims < 0 && throw(ArgumentError("dims must be non-negative but got dims=$dims.")) | ||
| return dims | ||
| end | ||
|
|
||
| function gather_kernel!(dst, src, idx, max_idx, max_dims_idx, dims_size) | ||
| index = threadIdx().x + (blockIdx().x - 1) * blockDim().x | ||
|
|
||
| @inbounds if index <= max_idx | ||
| j, k = divrem(index-1, max_dims_idx) | ||
| dims_i = CartesianIndices(dims_size)[k+1] | ||
| dst[index] = src[dims_i, idx[j+1]...] | ||
| end | ||
| return nothing | ||
| end | ||
|
|
||
| function NNlib.gather!(dst::AnyCuArray, src::AnyCuArray, idx::AnyCuArray) | ||
| dims = gather_check_dims(src, dst, idx) | ||
| dims_size = size(src)[1:dims] | ||
| max_dims_idx = prod(dims_size) | ||
| max_idx = max_dims_idx * length(idx) | ||
| args = dst, src, idx, max_idx, max_dims_idx, dims_size | ||
|
|
||
| kernel = @cuda launch=false gather_kernel!(args...) | ||
| config = launch_configuration(kernel.fun; max_threads=256) | ||
| threads = min(max_idx, config.threads) | ||
| blocks = cld(max_idx, threads) | ||
| kernel(args...; threads=threads, blocks=blocks) | ||
| return dst | ||
| end |
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| @testset "gather" begin | ||
| T = Float32 | ||
| CT = CuArray{Float32} | ||
|
|
||
| ## 1d src, 2d index of ints -> 2d output | ||
| src = CT([3, 4, 5, 6, 7]) | ||
| index = cu([1 2 3 4; | ||
| 4 2 1 3; | ||
| 3 5 5 3]) | ||
| output = CT([3 4 5 6; | ||
| 6 4 3 5; | ||
| 5 7 7 5]) | ||
|
|
||
| y = NNlib.gather(src, index) | ||
| @test y isa CuArray{Float32,2} | ||
| @test size(y) == size(index) | ||
| gputest(src -> NNlib.gather(src, index), src, checkgrad=false) | ||
| @test NNlib.gather!(CUDA.zeros(T, size(index)...), src, index) == output | ||
| @test_throws ArgumentError NNlib.gather!(zeros(T, 3, 5), src, index) | ||
|
|
||
| ## 1d src, 3d index of ints -> 3d output | ||
| src = CT([3, 4, 5, 6, 7]) | ||
| index = cu([1 2 3 4; | ||
| 4 2 1 3; | ||
| 3 5 5 3][:,:,1:1]) | ||
| output = CT([3 4 5 6; | ||
| 6 4 3 5; | ||
| 5 7 7 5][:,:,1:1]) | ||
|
|
||
| y = NNlib.gather(src, index) | ||
| @test y isa CuArray{Float32,3} | ||
| @test size(y) == size(index) | ||
| gputest(src -> NNlib.gather(src, index), src, checkgrad=false) | ||
|
|
||
|
|
||
| ## 2d src, 2d index of ints -> 3d output | ||
| src = CT([3 5 7 | ||
| 4 6 8]) | ||
| index = cu([1 2 3; | ||
| 2 2 1; | ||
| 3 1 3]) | ||
|
|
||
| output = zeros(T, 2, 3, 3) | ||
|
|
||
| output[:,:,1] = [3 5 7 | ||
| 4 6 8] | ||
|
|
||
| output[:,:,2] = [5 5 3 | ||
| 6 6 4] | ||
|
|
||
| output[:,:,3] = [7 3 7 | ||
| 8 4 8] | ||
|
|
||
| y = NNlib.gather(src, index) | ||
| M = NNlib.typelength(eltype(index)) | ||
| Nsrc = ndims(src) | ||
| @test y isa CuArray{Float32,3} | ||
| @test size(y) == (size(src)[1:Nsrc-M]..., size(index)...) | ||
| gputest(src -> NNlib.gather(src, index), src, checkgrad=false) | ||
| end |
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,4 +18,5 @@ if CUDA.has_cuda() | |
| include("softmax.jl") | ||
| include("batchnorm.jl") | ||
| include("scatter.jl") | ||
| include("gather.jl") | ||
| end | ||