Allow non-Int64 indices in scatter (#543)

* Fix typo in version comparison

* Allow non-Int64 indices in scatter

* Disable Enzyme for AMDGPU

* Refactor

Author: pxl-th

src/scatter.jl

@@ -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(_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])
@@ -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(_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 _convert_i64(x::Int) = x
+@inline _convert_i64(x::Integer) = Int(x)
+@inline _convert_i64(x) = x
+
 """
     NNlib.scatter(op, src, idx; [init, dstsize])
 

test/runtests.jl

@@ -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)
 

test/scatter.jl

@@ -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
@@ -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)) :
@@ -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