feat: better codegen with gather/scatter (#1283)

* feat: don't flatten indices for getindex

* fix: repeats

* fix: indexing consistency

* fix: avoid reshape

* feat: better codegen for NNlib.gather

* revert: reshape change

* feat: strided indexing

* fix: better diag codegen

* fix: missing singletons

* fix: more checks

* fix: lower to dynamic_slice

* fix: temporarily disable slice_elementwise

* fix: one-elem cartesian index

* feat: setindex as scatter

* chore: bump jll

Author: avik-pal

ext/ReactantNNlibExt/Implementations.jl

@@ -430,64 +430,43 @@ function NNlib.pad_constant(
 end
 
 # Gather
-# XXX: reevaluate this manual optimization once
-#      https://github.com/EnzymeAD/Enzyme-JAX/issues/164 is handled
-function NNlib.gather!(
-    dst::AnyTracedRArray{T1,2},
-    src::AnyTracedRArray{T2,2},
-    idxs::Union{AbstractUnitRange{<:Number}},
-) where {T1,T2}
-    set_mlir_data!(dst, get_mlir_data(src[:, idxs]))
+function NNlib.gather!(dst::AnyTracedRArray, src::AnyTracedRArray, idxs::AbstractArray)
+    n_dims = NNlib.scatter_dims(src, dst, idxs)
+    res = _nnlib_gather_impl(src, _stack_indices(idxs), n_dims)
+    set_mlir_data!(dst, get_mlir_data(res))
     return dst
 end
 
 function NNlib.gather!(
-    dst::AnyTracedRArray{T1,2}, src::AnyTracedRArray{T2,2}, idxs::AbstractVector{<:Number}
-) where {T1,T2}
-    dims = NNlib.scatter_dims(src, dst, idxs)
-    @assert dims == 1  # scatter_dims lets us do some size checks so we call that function
-    idxs = get_mlir_data(TracedUtils.promote_to(TracedRArray{Int,1}, idxs) .- 1)
-    slice_sizes = get_mlir_data(
-        TracedUtils.promote_to(TracedRArray{Int,1}, [size(src, 1), 1])
-    )
-
-    #! format: off
-    dimension_numbers = MLIR.API.stablehloGatherDimensionNumbersGet(
-        MLIR.IR.context(),
-        Int64(1), Int64[0],
-        Int64(1), Int64[1],
-        Int64(0), Int64[],
-        Int64(0), Int64[],
-        Int64(1), Int64[1],
-        Int64(1)
-    )
-    #! format: on
-
-    res = MLIR.IR.result(
-        Reactant.MLIR.Dialects.stablehlo.dynamic_gather(
-            get_mlir_data(src), idxs, slice_sizes; dimension_numbers
-        ),
-        1,
-    )
-    set_mlir_data!(dst, res)
+    dst::AnyTracedRArray, src::AnyTracedRArray, idxs::AbstractArray{<:Number}
+)
+    n_dims = NNlib.scatter_dims(src, dst, idxs)
+    res = _nnlib_gather_impl(src, reshape(idxs, 1, size(idxs)...), n_dims)
+    set_mlir_data!(dst, get_mlir_data(res))
     return dst
 end
 
-# XXX: For performance to use `stablehlo.dynamic_gather` or atleast use traced loop
-#      instead of unrolling the loop (the case for AbstractArray can just use
-#      `stablehlo.gather`). See above for the special case implementation that is optimized.
-function NNlib.gather!(dst::AnyTracedRArray, src::AnyTracedRArray, idxs::AbstractArray)
-    @warn "Using fallback implementation of `gather!` for using `stablehlo.dynamic_slice`. \
-           This case is not optimized and will be slow." maxlog = 1
-    dims = NNlib.scatter_dims(src, dst, idxs)
-    colons = ntuple(Returns(Colon()), dims)
-    start_sizes = ntuple(Base.Fix1(size, src), dims)
-    results = map(CartesianIndices(idxs)) do k
-        res = @allowscalar src[colons..., Tuple(idxs[k])...]
-        res isa TracedRNumber && (res = TracedUtils.broadcast_to_size(res, (1,)))
-        return reshape(res, start_sizes..., :)
+_stack_indices(idxs::AbstractArray) = stack(idxs)
+function _stack_indices(idxs::AbstractArray{<:CartesianIndex})
+    stacked_idxs = similar(idxs, Int, length(first(idxs)), size(idxs)...)
+    for k in CartesianIndices(idxs)
+        stacked_idxs[:, k.I...] .= idxs[k].I
     end
-    res = reshape(cat(results...; dims=(dims + 1)), size(dst))
-    set_mlir_data!(dst, get_mlir_data(res))
-    return dst
+    return stacked_idxs
+end
+
+function _nnlib_gather_impl(src::AnyTracedRArray, idxs::AbstractArray, n_dims::Int)
+    idxs = TracedUtils.promote_to(TracedRArray{Int,ndims(idxs)}, idxs)
+    n_idxs = size(idxs, 1)
+    return Ops.gather(
+        src,
+        idxs;
+        offset_dims=collect(Int64, 1:n_dims),
+        collapsed_slice_dims=collect(Int64, (n_dims + 1):ndims(src)),
+        operand_batching_dims=Int64[],
+        start_indices_batching_dims=Int64[],
+        start_index_map=collect(Int64, (ndims(src) - n_idxs + 1):ndims(src)),
+        index_vector_dim=1,
+        slice_sizes=Int64[size(src)[1:n_dims]..., ones(Int64, ndims(src) - n_dims)...],
+    )
 end

src/Ops.jl

@@ -532,16 +532,24 @@ end
 
 @noinline function slice(
     x::TracedRArray{T,N},
-    start_indices,
-    limit_indices;
-    strides=nothing,
+    start_indices::Vector{<:Integer},
+    limit_indices::Vector{<:Integer};
+    strides::Union{Nothing,Vector{<:Integer}}=nothing,
     location=mlir_stacktrace("slice", @__FILE__, @__LINE__),
 ) where {T,N}
     start_indices = start_indices .- 1
     limit_indices = limit_indices
-    rsize = limit_indices .- start_indices
-    @assert all(rsize .> 0) "Invalid slice dimensions"
+    @assert all(Base.Fix2(≥, 0), start_indices) "Invalid start indices: $(start_indices)"
+    @assert all(s < l for (s, l) in zip(start_indices, limit_indices)) "Invalid slice indices: $(start_indices), $(limit_indices)"
+
     strides = isnothing(strides) ? ones(Int64, N) : strides
+    @assert all(s > 0 for s in strides) "Invalid strides: $(strides)"
+    rsize = [
+        length((start + 1):st:stop) for
+        (start, stop, st) in zip(start_indices, limit_indices, strides)
+    ]
+    @assert all(rsize .> 0) "Invalid slice dimensions"
+
     res = MLIR.IR.result(
         stablehlo.slice(
             x.mlir_data;
@@ -1732,18 +1740,18 @@ instead.
         [updates];
         update_computation,
         update_window_dims=Int64[],
-        inserted_window_dims=collect(Int64, 0:(N - 1)),
+        inserted_window_dims=collect(Int64, 1:N),
         input_batching_dims=Int64[],
         scatter_indices_batching_dims=Int64[],
-        scatter_dims_to_operand_dims=collect(Int64, 0:(N - 1)),
-        index_vector_dim=Int64(1),
+        scatter_dims_to_operand_dims=collect(Int64, 1:N),
+        index_vector_dim=Int64(2),
         location,
     )[1]
 end
 
 @noinline function scatter(
     dest::Vector{TracedRArray{T,N}},
-    scatter_indices::TracedRArray{Int64,2},
+    scatter_indices::TracedRArray{Int64},
     updates::Vector{<:TracedRArray{T}};
     update_computation::MLIR.IR.Region,
     update_window_dims::Vector{Int64},
@@ -1758,6 +1766,13 @@ end
         scatter_indices, fill(Int64(1), size(scatter_indices)); location
     )
 
+    update_window_dims = update_window_dims .- 1
+    inserted_window_dims = inserted_window_dims .- 1
+    input_batching_dims = input_batching_dims .- 1
+    scatter_indices_batching_dims = scatter_indices_batching_dims .- 1
+    scatter_dims_to_operand_dims = scatter_dims_to_operand_dims .- 1
+    index_vector_dim -= 1
+
     #! format: off
     scatter_dimension_numbers = MLIR.API.stablehloScatterDimensionNumbersGet(
         MLIR.IR.context(),
@@ -1813,11 +1828,11 @@ use [`MLIR.Dialects.stablehlo.dynamic_slice`](@ref) instead.
             src,
             gather_indices;
             offset_dims=Int64[1],
-            collapsed_slice_dims=collect(Int64, 0:(N - 2)),
+            collapsed_slice_dims=collect(Int64, 1:(N - 1)),
             operand_batching_dims=Int64[],
             start_indices_batching_dims=Int64[],
-            start_index_map=collect(Int64, 0:(N - 1)),
-            index_vector_dim=Int64(1),
+            start_index_map=collect(Int64, 1:N),
+            index_vector_dim=Int64(2),
             slice_sizes=ones(Int64, N),
             indices_are_sorted=false,
             location,
@@ -1828,7 +1843,7 @@ end
 
 @noinline function gather(
     src::TracedRArray{T,N},
-    gather_indices::TracedRArray{Int64,2};
+    gather_indices::TracedRArray{Int64};
     offset_dims::Vector{Int64},
     collapsed_slice_dims::Vector{Int64},
     operand_batching_dims::Vector{Int64},
@@ -1843,6 +1858,13 @@ end
         gather_indices, fill(Int64(1), size(gather_indices)); location
     )
 
+    offset_dims = offset_dims .- 1
+    start_indices_batching_dims = start_indices_batching_dims .- 1
+    start_index_map = start_index_map .- 1
+    operand_batching_dims = operand_batching_dims .- 1
+    collapsed_slice_dims = collapsed_slice_dims .- 1
+    index_vector_dim -= 1
+
     #! format: off
     dimension_numbers = MLIR.API.stablehloGatherDimensionNumbersGet(
         MLIR.IR.context(),

src/TracedRArray.jl

@@ -140,20 +140,28 @@ function Base.getindex(a::TracedRArray{T,1}, indices::CartesianIndex{1}) where {
     return _getindex_cartesian(a, indices)
 end
 
-_isone(x) = isone(x)
-_isone(::CartesianIndex) = false
-
-__contiguous_indices(::Base.LogicalIndex) = false
-__contiguous_indices(x) = all(_isone, diff(x))
-
 function _getindex_linear(a::TracedRArray{T,N}, indices::AbstractArray) where {T,N}
-    if !(indices isa Reactant.TracedType) && __contiguous_indices(vec(indices))
-        a_flat = materialize_traced_array(vec(a))
-        indices_flat = vec(indices)
-        return Ops.reshape(
-            Ops.dynamic_slice(a_flat, [first(indices_flat)], [length(indices_flat)]),
-            collect(size(indices)),
-        )
+    if !(indices isa Reactant.TracedType)
+        if length(indices) == 1 && first(indices) isa CartesianIndex
+            # fast-path else we will end up with a gather
+            return TracedUtils.broadcast_to_size(
+                @allowscalar(_getindex_cartesian(a, first(indices))), (1,)
+            )
+        end
+        stride = TracedUtils._get_slice_stride(vec(indices))
+        if stride > 0
+            a_flat = materialize_traced_array(vec(a))
+            indices_flat = vec(indices)
+            return Ops.reshape(
+                Ops.slice(
+                    a_flat,
+                    Int64[first(indices_flat)],
+                    Int64[last(indices_flat)];
+                    strides=Int64[stride],
+                ),
+                collect(Int64, size(indices)),
+            )
+        end
     end
 
     if !(indices isa TracedRArray)
@@ -180,14 +188,21 @@ function Base.getindex(a::TracedRArray{T,N}, indices::Vararg{Any,N}) where {T,N}
     indices = Base.to_indices(a, indices)
 
     use_gather_getindex = false
+    use_dynamic_slice = false
+    strides = Int64[]
     for idxs in indices
-        idxs isa Number && continue
+        if idxs isa Number
+            idxs isa TracedRNumber && (use_dynamic_slice = true)
+            push!(strides, 1)
+            continue
+        end
         if idxs isa Reactant.TracedType
             use_gather_getindex = true
             break
         end
-        contiguous = __contiguous_indices(vec(idxs))
-        if typeof(contiguous) <: Bool && !contiguous
+        stride = TracedUtils._get_slice_stride(vec(idxs))
+        push!(strides, stride)
+        if stride ≤ 0 || (use_dynamic_slice && stride != 1)
             use_gather_getindex = true
             break
         end
@@ -200,18 +215,37 @@ function Base.getindex(a::TracedRArray{T,N}, indices::Vararg{Any,N}) where {T,N}
             error("Boolean indexing with TracedRArrays isn't fully supported yet.")
         end
 
-        indices, integer_indices, result_size, preddim_result_size, _ = TracedUtils.traced_indices(
-            indices...
-        )
-        res = Ops.reshape(
-            Ops.gather_getindex(a, generate_index_list(indices...)), preddim_result_size
+        gather_dims = TracedUtils.indices_to_gather_dims(indices...)
+
+        return Ops.reshape(
+            Ops.transpose(
+                Ops.reshape(
+                    Ops.gather(
+                        a,
+                        gather_dims.start_indices;
+                        gather_dims.offset_dims,
+                        gather_dims.collapsed_slice_dims,
+                        operand_batching_dims=Int64[],
+                        start_indices_batching_dims=Int64[],
+                        gather_dims.start_index_map,
+                        gather_dims.index_vector_dim,
+                        gather_dims.slice_sizes,
+                    ),
+                    gather_dims.gather_reshape_shape,
+                ),
+                gather_dims.permutation,
+            ),
+            gather_dims.result_shape,
         )
-        isempty(integer_indices) ||
-            (res = materialize_traced_array(dropdims(res; dims=integer_indices)))
-        return Ops.reshape(res, result_size)
     end
 
-    x = Ops.dynamic_slice(a, [first.(indices)...], [length.(indices)...])
+    if use_dynamic_slice
+        @assert all(isone, strides) "This should not happen, please report a bug"
+        x = Ops.dynamic_slice(a, [first.(indices)...], [length.(indices)...])
+    else
+        x = Ops.slice(a, [first.(indices)...], [last.(indices)...]; strides)
+    end
+
     ddims = findall(indices) do idx
         return idx isa Integer || idx isa TracedRNumber{<:Integer}
     end
@@ -313,7 +347,7 @@ function Base.setindex!(a::TracedRArray{T,N}, v, index::CartesianIndex{N}) where
 end
 
 function _setindex_linear!(a::TracedRArray{T,N}, v, indices::AbstractArray) where {T,N}
-    if !(indices isa Reactant.TracedType) && __contiguous_indices(vec(indices))
+    if !(indices isa Reactant.TracedType) && TracedUtils.__contiguous_indices(vec(indices))
         res = Ops.reshape(
             Ops.dynamic_update_slice(
                 materialize_traced_array(vec(a)),
@@ -371,7 +405,7 @@ function Base.setindex!(a::TracedRArray{T,N}, v, indices::Vararg{Any,N}) where {
             use_scatter_setindex = true
             break
         end
-        contiguous = __contiguous_indices(idxs)
+        contiguous = TracedUtils.__contiguous_indices(idxs)
         if typeof(contiguous) <: Bool && !contiguous
             use_scatter_setindex = true
             break
@@ -384,9 +418,44 @@ function Base.setindex!(a::TracedRArray{T,N}, v, indices::Vararg{Any,N}) where {
         if any(i -> unwrapped_eltype(i) <: Bool, indices)
             error("Boolean indexing with TracedRArrays isn't fully supported yet.")
         end
-        indices_list = map(Base.Fix1(TracedUtils.promote_to, TracedRArray{Int,1}), indices)
-        indices_list = generate_index_list(indices_list...)
-        res = Ops.scatter_setindex(a, indices_list, Ops.reshape(v, length(v)))
+
+        gather_dims = TracedUtils.indices_to_gather_dims(indices...)
+
+        v = Ops.convert(
+            TracedRArray{T,ndims(v)},
+            TracedUtils.promote_to(TracedRArray{unwrapped_eltype(v),ndims(v)}, v),
+        )
+
+        updates = Ops.transpose(v, invperm(gather_dims.permutation))
+        n_collapsed = length(gather_dims.collapsed_slice_dims)
+        updates_shape = Int64[
+            prod(size(updates)[1:n_collapsed]), size(updates)[(n_collapsed + 1):end]...
+        ]
+        updates = Ops.reshape(updates, updates_shape)
+
+        # simply set the 2nd block argument as a result
+        update_computation = MLIR.IR.Region()
+        block = MLIR.IR.Block(
+            [Ops.mlir_type(TracedRNumber{T}), Ops.mlir_type(TracedRNumber{T})],
+            [MLIR.IR.Location(), MLIR.IR.Location()],
+        )
+        return_op = MLIR.Dialects.stablehlo.return_([MLIR.IR.argument(block, 2)])
+        MLIR.IR.rmfromparent!(return_op)
+        push!(block, return_op)
+        pushfirst!(update_computation, block)
+
+        res = Ops.scatter(
+            [a],
+            gather_dims.start_indices,
+            [updates];
+            update_computation,
+            update_window_dims=gather_dims.offset_dims,
+            inserted_window_dims=gather_dims.collapsed_slice_dims,
+            input_batching_dims=Int64[],
+            scatter_indices_batching_dims=Int64[],
+            scatter_dims_to_operand_dims=gather_dims.start_index_map,
+            index_vector_dim=gather_dims.index_vector_dim,
+        )[1]
         set_mlir_data!(a, get_mlir_data(res))
         return v
     end