feat: NNlib.scatter (#1395)

avik-pal · JulianTrommer · web-flow · commit 115ac2585a3f · 2025-06-12T09:16:39.000-04:00
* First version of NNlib.scatter code &amp; tests

* Support for higher scatter dims  + refactoring

* Added support for mean in NNlib.scatter

* refactor: reuse scatter impl

* fix: emit better scatter

---------

Co-authored-by: Julian Trommer &lt;julian.trommer@uni-a.de&gt;
diff --git a/Project.toml b/Project.toml
@@ -52,7 +52,7 @@ ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = ["CUDA", "GPUCompiler", "KernelAbstractions", "LLVM"]
 ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
-ReactantNNlibExt = "NNlib"
+ReactantNNlibExt = ["NNlib", "Statistics"]
 ReactantOffsetArraysExt = "OffsetArrays"
 ReactantOneHotArraysExt = "OneHotArrays"
 ReactantPythonCallExt = "PythonCall"
diff --git a/ext/ReactantNNlibExt/Implementations.jl b/ext/ReactantNNlibExt/Implementations.jl
@@ -487,3 +487,90 @@ function NNlib.upsample_linear_kernel!(
     copyto!(y, upsample_linear(x, size(y)[1:(end - 2)], ratios..., align_corners))
     return y
 end
+
+# Scatter
+function NNlib.scatter(
+    op::OP, src::AnyTracedRArray{T}, idx::AbstractArray; init=nothing, dstsize=nothing
+) where {OP,T}
+    dims = ndims(src) - ndims(idx)
+    dstsz = if isnothing(dstsize)
+        (size(src)[1:dims]..., NNlib.maximum_dims(idx)...)
+    else
+        dstsize
+    end
+    if any(d -> d isa TracedRNumber, dstsz)
+        throw(
+            ArgumentError(
+                "dstsize must be specified when idx is a TracedRArray or contains a TracedRNumber.",
+            ),
+        )
+    end
+    xinit = isnothing(init) ? NNlib.scatter_empty(op, T) : init
+    dst = Ops.fill(xinit, dstsz)
+
+    NNlib.scatter!(op, dst, src, idx)
+    return dst
+end
+
+function NNlib.scatter!(
+    op::OP, dst::AnyTracedRArray, src::AnyTracedRArray, idx::AbstractArray
+) where {OP}
+    dims = NNlib.scatter_dims(dst, src, idx)
+    res = _nnlib_scatter_impl(op, dst, src, _stack_indices(idx), dims)
+    set_mlir_data!(dst, get_mlir_data(res))
+    return dst
+end
+
+function NNlib.scatter!(
+    op::OP, dst::AnyTracedRArray, src::AnyTracedRArray, idx::AbstractArray{<:Number}
+) where {OP}
+    dims = NNlib.scatter_dims(dst, src, idx)
+    res = _nnlib_scatter_impl(op, dst, src, reshape(idx, 1, size(idx)...), dims)
+    set_mlir_data!(dst, get_mlir_data(res))
+    return dst
+end
+
+for AT in (AbstractArray, AbstractArray{<:Number})
+    @eval function NNlib.scatter!(
+        ::typeof(mean), dst::AnyTracedRArray, src::AnyTracedRArray, idx::$AT
+    )
+        Ns = NNlib.scatter!(+, zero(dst), one.(src), idx)
+        dst_ = NNlib.scatter!(+, zero(dst), src, idx)
+        res = dst .+ NNlib.safe_div.(dst_, Ns)
+        set_mlir_data!(dst, get_mlir_data(res))
+        return dst
+    end
+end
+
+function _nnlib_scatter_impl(
+    op::OP,
+    dst::AnyTracedRArray{T},
+    src::AnyTracedRArray{T},
+    idx::AbstractArray,
+    n_dims::Int,
+) where {OP,T}
+    scatter_indices = TracedUtils.promote_to(TracedRArray{Int,ndims(idx)}, idx)
+    n_idxs = size(scatter_indices, 1)
+    return Ops.scatter(
+        op,
+        [dst],
+        scatter_indices,
+        [src];
+        update_window_dims=collect(Int64, 1:n_dims),
+        inserted_window_dims=collect(Int64, (n_dims + 1):ndims(dst)),
+        input_batching_dims=Int64[],
+        scatter_indices_batching_dims=Int64[],
+        scatter_dims_to_operand_dims=collect(Int64, (ndims(dst) - n_idxs + 1):ndims(dst)),
+        index_vector_dim=Int64(1),
+    )[1]
+end
+
+function NNlib.maximum_dims(dims::AnyTracedRArray{<:Integer})
+    return (maximum(dims),)
+end
+function NNlib.maximum_dims(dims::AnyTracedRArray{NTuple{N,T}}) where {N,T}
+    return ntuple(i -> maximum(x -> x[i], dims), N)
+end
+function NNlib.maximum_dims(dims::AnyTracedRArray{CartesianIndex{N}}) where {N}
+    return ntuple(i -> maximum(x -> x[i], dims), N)
+end
diff --git a/ext/ReactantNNlibExt/ReactantNNlibExt.jl b/ext/ReactantNNlibExt/ReactantNNlibExt.jl
@@ -10,6 +10,7 @@ using Reactant.TracedUtils:
 
 using ReactantCore: @trace
 using LinearAlgebra: LinearAlgebra, triu
+using Statistics: mean
 
 include("Overlay.jl")
 include("Ops.jl")
diff --git a/src/Compiler.jl b/src/Compiler.jl
@@ -838,7 +838,8 @@ function optimization_passes(;
         "divide_sqrt_to_multiply_rsqrt<16>",
         "associative_binary_op_reordering<1>",
         "transpose_broadcast_in_dim_to_broadcast_in_dim<16>",
-        "scatter_indices_are_unique",
+        # XXX: needs upstream fix
+        # "scatter_indices_are_unique",
         "replace_neg_add_with_subtract",
         "binop_const_simplify",
         "not_select_simplify",
diff --git a/src/Ops.jl b/src/Ops.jl
@@ -1734,6 +1734,8 @@ end
     scatter_indices_batching_dims::Vector{Int64},
     scatter_dims_to_operand_dims::Vector{Int64},
     index_vector_dim::Int64,
+    unique_indices::Union{Bool,Nothing}=nothing,
+    indices_are_sorted::Union{Bool,Nothing}=nothing,
     location=mlir_stacktrace("scatter", @__FILE__, @__LINE__),
 ) where {T,N}
     scatter_indices = subtract(
@@ -1768,6 +1770,8 @@ end
         update_computation,
         scatter_dimension_numbers,
         result_0=[mlir_type(TracedRArray{T,N}, size(d)) for d in dest],
+        unique_indices,
+        indices_are_sorted,
         location,
     )
 
diff --git a/test/nn/nnlib.jl b/test/nn/nnlib.jl
@@ -1,4 +1,5 @@
 using NNlib, Reactant, Enzyme
+using Statistics
 
 @testset "Activation Functions" begin
     sumabs2(f, x) = sum(abs2, f.(x))
@@ -381,6 +382,255 @@ end
     end
 end
 
+# Adapted from https://github.com/FluxML/NNlib.jl/blob/1468582c4db5f18149cc8fff6fb4633c5debe5c5/test/testsuite/scatter.jl#L108
+@testset "NNlib scatter" begin
+    function test_scatter(dsts, srcs, idxs, res; dims)
+        @testset "scatter Float32 $op" for op in (+, -, max, min, *, /, mean)
+            for idx in values(idxs), dim in dims
+                dst = copy(dsts[dim])
+                target_y = res[(op, dim, true)]
+                src = srcs[(dim, true)]
+                if op == /
+                    src = src .* 2.0f0
+                end
+
+                y1 = @jit(
+                    NNlib.scatter!(
+                        op, Reactant.to_rarray(dst), Reactant.to_rarray(src), idx
+                    )
+                )
+                @test y1 ≈ target_y
+                @test y1 isa ConcreteRArray{Float32,ndims(dst)}
+                @test size(y1) == size(dsts[dim])
+                dst = copy(dsts[dim])
+                y2 = @jit(
+                    NNlib.scatter!(
+                        op,
+                        Reactant.to_rarray(dst),
+                        Reactant.to_rarray(src),
+                        Reactant.to_rarray(idx),
+                    )
+                )
+                @test y2 ≈ target_y
+                @test y2 isa ConcreteRArray{Float32,ndims(dst)}
+                @test size(y2) == size(dsts[dim])
+
+                target_y = res[(op, dim, false)]
+                src = srcs[(dim, false)]
+                if op == /
+                    src = src .* 2.0f0
+                end
+
+                y3 = @jit(NNlib.scatter(op, Reactant.to_rarray(src), idx))
+                @test y3 ≈ target_y
+                @test y3 isa ConcreteRArray{Float32,ndims(dst)}
+                @test size(y3) == size(dsts[dim])
+                y4 = @jit(
+                    NNlib.scatter(
+                        op,
+                        Reactant.to_rarray(src),
+                        Reactant.to_rarray(idx);
+                        dstsize=size(dsts[dim]),
+                    )
+                )
+                @test y4 ≈ target_y
+                @test y4 isa ConcreteRArray{Float32,ndims(dst)}
+                @test size(y4) == size(dsts[dim])
+
+                ridx = Reactant.to_rarray(idx)
+                if ridx isa Reactant.AbstractConcreteArray
+                    @test_throws ArgumentError @jit(
+                        NNlib.scatter(op, Reactant.to_rarray(src), ridx)
+                    )
+                else
+                    y5 = @jit(NNlib.scatter(op, Reactant.to_rarray(src), ridx))
+                    @test y5 ≈ target_y
+                    @test y5 isa ConcreteRArray{Float32,ndims(dst)}
+                    @test size(y5) == size(dsts[dim])
+                end
+            end
+        end
+    end
+
+    @testset "scatter 1d src, 1d index => 1d output" begin
+        #! format: off
+        dsts = Dict(
+            0 => Float32[3, 4, 5, 6, 7]
+        )
+
+        srcs = Dict(
+            (0, true) => ones(Float32, 5),
+            (0, false) => collect(Float32, 1:5),
+        )
+
+        idxs = Dict(
+            :int => [4, 2, 1, 5, 3],
+            :tup => [(4,), (2,), (1,), (5,), (3,)],
+            :car => CartesianIndex.([(4,), (2,), (1,), (5,), (3,)]),
+        )
+
+        res = Dict(
+            (+, 0, true) => Float32[4, 5, 6, 7, 8],
+            (+, 0, false) => Float32[3, 2, 5, 1, 4],
+
+            (-, 0, true) => Float32[2, 3, 4, 5, 6],
+            (-, 0, false) => Float32[-3, -2, -5, -1, -4],
+
+            (max, 0, true) => Float32[3, 4, 5, 6, 7],
+            (max, 0, false) => Float32[3, 2, 5, 1, 4],
+
+            (min, 0, true) => Float32[1, 1, 1, 1, 1],
+            (min, 0, false) => Float32[3, 2, 5, 1, 4],
+
+            (*, 0, true) => Float32[3, 4, 5, 6, 7],
+            (*, 0, false) => Float32[3, 2, 5, 1, 4],
+
+            (/, 0, true) => Float32[1.5, 2.0, 2.5, 3.0, 3.5],
+            (/, 0, false) => Float32[1//6, 1//4, 1//10, 1//2, 1//8],
+
+            (mean, 0, true) => Float32[4, 5, 6, 7, 8],
+            (mean, 0, false) => Float32[3, 2, 5, 1, 4],
+        )
+        #! format: on
+        test_scatter(dsts, srcs, idxs, res; dims=[0])
+    end
+
+    @testset "scatter 2d src, 1d index => 2d output" begin
+        #! format: off
+        dsts = Dict(
+                0 => Float32[3 3 4 4 5
+                            5 5 6 6 7]
+        )
+
+        srcs = Dict(
+            (0, true) => ones(Float32, 2, 5),
+            (0, false) => ones(Float32, 2) * collect(1:5)',
+        )
+
+        idxs = Dict(
+            :int => [4, 2, 1, 5, 3],
+            :tup => [(4,), (2,), (1,), (5,), (3,)],
+            :car => CartesianIndex.([(4,), (2,), (1,), (5,), (3,)]),
+        )
+
+        res = Dict(
+            (+, 0, true) => Float32[4 4 5 5 6;
+                                    6 6 7 7 8],
+            (+, 0, false) => Float32[3 2 5 1 4;
+                                        3 2 5 1 4],
+
+            (-, 0, true) => Float32[2 2 3 3 4;
+                                    4 4 5 5 6],
+            (-, 0, false) => Float32[-3 -2 -5 -1 -4;
+                                        -3 -2 -5 -1 -4],
+
+            (max, 0, true) => Float32[3 3 4 4 5;
+                                        5 5 6 6 7],
+            (max, 0, false) => Float32[3 2 5 1 4;
+                                        3 2 5 1 4],
+
+            (min, 0, true) => Float32[1 1 1 1 1;
+                                        1 1 1 1 1],
+            (min, 0, false) => Float32[3 2 5 1 4;
+                                        3 2 5 1 4],
+
+            (*, 0, true) => Float32[3 3 4 4 5;
+                                    5 5 6 6 7],
+            (*, 0, false) => Float32[3 2 5 1 4;
+                                        3 2 5 1 4],
+
+            (/, 0, true) => Float32[1.5 1.5 2.0 2.0 2.5;
+                                    2.5 2.5 3.0 3.0 3.5],
+            (/, 0, false) => Float32[1//6 1//4 1//10 1//2 1//8;
+                                        1//6 1//4 1//10 1//2 1//8],
+
+            (mean, 0, true) => Float32[4 4 5 5 6;
+                                        6 6 7 7 8],
+            (mean, 0, false) => Float32[3 2 5 1 4;
+                                        3 2 5 1 4],
+        )
+        #! format: on
+        test_scatter(dsts, srcs, idxs, res; dims=[0])
+    end
+
+    @testset "scatter 2d+3d src, 2d index => 1d+2d output" begin
+        #! format: off
+        dsts = Dict(
+            0 => Float32[3, 4, 5, 6, 7],
+            1 => Float32[3 3 4 4 5;
+                         5 5 6 6 7],
+        )
+
+        srcs = Dict(
+            (0, true) => ones(Float32, 3, 4),
+            (0, false) => ones(Float32, 3) * collect(1:4)',
+            (1, true) => ones(Float32, 2, 3, 4),
+            (1, false) => Float32[1, 2] .* reshape(ones(Float32, 3) * collect(1:4)', 1,3,4),
+        )
+
+        idxs = Dict(
+            :int => [1 2 3 4;
+                     4 2 1 3;
+                     3 5 5 3],
+            :tup => [(1,) (2,) (3,) (4,);
+                     (4,) (2,) (1,) (3,);
+                     (3,) (5,) (5,) (3,)],
+            :car => CartesianIndex.(
+                    [(1,) (2,) (3,) (4,);
+                     (4,) (2,) (1,) (3,);
+                     (3,) (5,) (5,) (3,)]),
+        )
+
+        res = Dict(
+            (+, 0, true) => Float32[5, 6, 9, 8, 9],
+            (+, 1, true) => Float32[5 5 8 6 7;
+                                    7 7 10 8 9],
+            (+, 0, false) => Float32[4, 4, 12, 5, 5],
+            (+, 1, false) => Float32[4 4 12 5 5;
+                                     8 8 24 10 10],
+            (-, 0, true) => Float32[1, 2, 1, 4, 5],
+            (-, 1, true) => Float32[1 1 0 2 3;
+                                    3 3 2 4 5],
+            (-, 0, false) => Float32[-4, -4, -12, -5, -5],
+            (-, 1, false) => Float32[-4 -4 -12 -5 -5;
+                                     -8 -8 -24 -10 -10],
+            (max, 0, true) => Float32[3, 4, 5, 6, 7],
+            (max, 1, true) => Float32[3 3 4 4 5;
+                                      5 5 6 6 7],
+            (max, 0, false) => Float32[3, 2, 4, 4, 3],
+            (max, 1, false) => Float32[3 2 4 4 3;
+                                       6 4 8 8 6],
+            (min, 0, true) => Float32[1, 1, 1, 1, 1],
+            (min, 1, true) => Float32[1 1 1 1 1;
+                                      1 1 1 1 1],
+            (min, 0, false) => Float32[1, 2, 1, 1, 2],
+            (min, 1, false) => Float32[1 2 1 1 2;
+                                       2 4 2 2 4],
+            (*, 0, true) => Float32[3, 4, 5, 6, 7],
+            (*, 1, true) => Float32[3 3 4 4 5;
+                                    5 5 6 6 7],
+            (*, 0, false) => Float32[3, 4, 48, 4, 6],
+            (*, 1, false) => Float32[3 4 48 4 6;
+                                     12 16 768 16 24],
+            (/, 0, true) => Float32[0.75, 1., 0.3125, 1.5, 1.75],
+            (/, 1, true) => Float32[0.75 0.75 0.25 1. 1.25;
+                                    1.25 1.25 0.375 1.5 1.75],
+            (/, 0, false) => Float32[1//12, 1//16, 1//768, 1//16, 1//24],
+            (/, 1, false) => Float32[1//12 1//16 1//768 1//16 1//24;
+                                     1//48 1//64 1//12288 1//64 1//96],
+            (mean, 0, true) => Float32[4., 5., 6., 7., 8.],
+            (mean, 1, true) => Float32[4. 4. 5. 5. 6.;
+                                       6. 6. 7. 7. 8.],
+            (mean, 0, false) => Float32[2, 2, 3, 2.5, 2.5],
+            (mean, 1, false) => Float32[2. 2. 3. 2.5 2.5;
+                                        4. 4. 6. 5. 5.],
+        )
+        #! format: on
+
+        test_scatter(dsts, srcs, idxs, res; dims=[0, 1])
+    end
+end
+
 @testset "∇conv(D = $ndim)" for ndim in 1:3
     x_spatial_dim = 4
     batch_size = 2
