Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow non-Int64 indices in scatter #543

Merged
merged 4 commits into from
Oct 4, 2023
Merged

Allow non-Int64 indices in scatter #543

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ce88867
Fix typo in version comparison
pxl-th Oct 2, 2023
1b00dd1
Allow non-Int64 indices in scatter
pxl-th Oct 2, 2023
3571306
Disable Enzyme for AMDGPU
pxl-th Oct 2, 2023
103ed32
Refactor
pxl-th Oct 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions src/scatter.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ end

@kernel function _scatter!(op::OP, dst, src, idxs) where OP
i = @index(Global)
@inbounds idx = Tuple(idxs[i])
@inbounds idx = Tuple(_atomix_convert_i64(idxs[i]))
@inbounds Atomix.modify!(Atomix.IndexableRef(dst, idx), op, src[i])
# FIXME `@atomic` macro silently fails to perform atomic op below
# @atomic dst[idx...] = op(dst[idx...], src[i])
Expand All @@ -119,14 +119,20 @@ end
) where OP
i = @index(Global)
j, k = divrem(i - 1, max_dims_idx)
@inbounds idx = (Tuple(dim_ids[k + 1])..., Tuple(idxs[j + 1])...)
@inbounds idx = (Tuple(dim_ids[k + 1])..., Tuple(_atomix_convert_i64(idxs[j + 1]))...)
@inbounds Atomix.modify!(Atomix.IndexableRef(dst, idx), op, src[i])
# FIXME
# FIXME `@atomic` macro silently fails to perform atomic op below
# dim_i = Tuple(dim_ids[k + 1])
# idx = idxs[j + 1]
# @atomic dst[dim_i..., idx...] = op(dst[dim_i..., idx...], src[i])
end

# Allow non-Int64 indices by converting them to Int64 when index eltype <: Integer.
# All other index types (tuples, cartesian indices) must be in Int64 already.
@inline _atomix_convert_i64(x::Int64) = x
@inline _atomix_convert_i64(x::Integer) = Int(x)
@inline _atomix_convert_i64(x) = x
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does the _atomix_ here mean?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess I just wanted to "signal" that it is related to Atomix.IndexableRef initially. Renamed to _convert_i64.

Copy link
Member Author

@pxl-th pxl-th Oct 3, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, I forgot to mention that casting to Int is only needed for atomics, where we take Atomix.pointer(::Atomix.IndexableRef) which works only with Int type.


"""
NNlib.scatter(op, src, idx; [init, dstsize])

Expand Down
4 changes: 3 additions & 1 deletion test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@ using Adapt
using KernelAbstractions
import ReverseDiff as RD # used in `pooling.jl`

const Test_Enzyme = VERSION <= v"1.10" && !Sys.iswindows()
const Test_Enzyme = VERSION <= v"1.10-" && !Sys.iswindows() &&
# TODO Enzyme is not working properly with AMDGPU yet.
get(ENV, "NNLIB_TEST_AMDGPU", "false") != "true"

DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib, UnicodePlots); recursive=true)

Expand Down
68 changes: 34 additions & 34 deletions test/scatter.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,37 +69,36 @@ res = Dict(
)

function test_scatter(device, types, ops; pt, ops_skip_types)
for T in types
for T in types, IT in (Int8, Int64)
PT = promote_type(T, pt)
@testset "$T" begin
for op in ops
skip_types = get(ops_skip_types, op, [])
@testset "$op" begin
for idx = values(idxs), dims = [0, 1]
idx = device(idx)
dst = device(dsts[dims])

mutated = true
target_y = res[(op, dims, mutated)]
src = device(srcs[(dims, mutated)])
if op == /
src = src .* T(2)
end

@test cpu(scatter!(op, T.(dst), T.(src), idx)) == T.(target_y)
@test cpu(scatter!(op, T.(dst), src, idx)) == PT.(target_y)
if op == /
@test cpu(scatter!(op, T.(dst), T.(src), idx)) == PT.(target_y)
else
@test cpu(scatter!(op, copy(dst), T.(src), idx)) == PT.(target_y)
end

if T ∉ skip_types
mutated = false
src = device(srcs[(dims, mutated)])
@test cpu(scatter(op, T.(src), idx)) == T.(res[(op, dims, mutated)])
end
end
@testset "eltype $T - idx eltype $IT - $op" for op in ops
skip_types = get(ops_skip_types, op, [])
for idx = values(idxs), dims = [0, 1]
# Tests with indices of different types.
eltype(idx) == Int && (idx = IT.(idx);)

idx = device(idx)
dst = device(dsts[dims])

mutated = true
target_y = res[(op, dims, mutated)]
src = device(srcs[(dims, mutated)])
if op == /
src = src .* T(2)
end

@test cpu(scatter!(op, T.(dst), T.(src), idx)) == T.(target_y)
@test cpu(scatter!(op, T.(dst), src, idx)) == PT.(target_y)
if op == /
@test cpu(scatter!(op, T.(dst), T.(src), idx)) == PT.(target_y)
else
@test cpu(scatter!(op, copy(dst), T.(src), idx)) == PT.(target_y)
end

if T ∉ skip_types
mutated = false
src = device(srcs[(dims, mutated)])
@test cpu(scatter(op, T.(src), idx)) == T.(res[(op, dims, mutated)])
end
end
end
Expand Down Expand Up @@ -174,14 +173,14 @@ function scatter_testsuite(Backend)
else
(+, -, mean, max, min)
end
for op in ops, i in (0, 1)
for op in ops, i in (0, 1), IT in (Int8, Int64)
PT = ( # If not CPU and CUDA -> use Int64 for min/max.
Backend != CPU &&
Symbol(Backend) != :CUDABackend &&
(op == max || op == min)) ? Int64 : T

src = device(srcs[(i, true)])
idx = device(idxs[:int])
idx = device(IT.(idxs[:int]))
dst = device(PT.(dsts[i]))
Backend == CPU ?
gradtest_fn(x -> scatter!(op, copy(x), src, idx), dst; fdm=fdm(op)) :
Expand All @@ -195,19 +194,20 @@ function scatter_testsuite(Backend)
else
(+, -, mean, max, min)
end
for op in ops, i in (0, 1)
for op in ops, i in (0, 1), IT in (Int8, Int64)
PT = ( # If not CPU and CUDA -> use Int64 for min/max.
Backend != CPU &&
Symbol(Backend) != :CUDABackend &&
(op == max || op == min)) ? Int64 : T
src = PT.(device(srcs[(i, false)]))
idx = device(idxs[:int])
idx = device(IT.(idxs[:int]))
Backend == CPU ?
gradtest_fn(xs -> scatter(op, xs, idx), src; fdm=fdm(op)) :
gradtest_fn((xs, i) -> scatter(op, xs, i), src, idx)
end
end


@static if Test_Enzyme

@testset "EnzymeRules" begin
Expand Down
Loading