fix: better MLIR codegen for simple getindex/setindex (#1200)

* fix: emit slice instead of gather

* fix: emit dus instead of scatter

* fix: diag indexing

* fix: indexing

* fix: promotion

* fix: ops

* fix: index

* chore: run formatter

* fix: diag indices

* fix: we don't need diag specialization

* fix: contiguous check

* fix: generalize subarray indexing behavior

* fix: try fixing ambiguity

* fix: more bad dispatches

* fix: fixup

* fix: some ambiguities

* fix: more cleanup

* fix: materialize subarray

* fix: more fixes

* Update src/TracedRArray.jl


* fix: remove unwanted overlay

* Update test/indexing.jl


* fix: fill! compile

Author: avik-pal

src/ConcreteRArray.jl

@@ -547,25 +547,22 @@ function Base.fill!(a::ConcretePJRTArray{T,N}, val) where {T,N}
         return a
     end
 
-    idxs = ntuple(Returns(Colon()), N)
-    fn = compile(mysetindex!, (a, val, idxs...))
-    fn(a, val, idxs...)
+    fn = compile(fill!, (a, val))
+    fn(a, val)
     return a
 end
 
 function Base.fill!(a::ConcreteIFRTArray{T,N}, val) where {T,N}
     isempty(a) && throw("Cannot setindex! to empty buffer")
 
-    idxs = ntuple(Returns(Colon()), N)
-    fn = compile(mysetindex!, (a, val, idxs...))
-    fn(a, val, idxs...)
+    fn = compile(fill!, (a, val))
+    fn(a, val)
     return a
 end
 
 function Base.fill!(x::Union{AnyConcreteIFRTArray,AnyConcretePJRTArray}, val)
-    idxs = ntuple(Returns(Colon()), ndims(x))
-    fn = compile(mysetindex!, (x, val, idxs...))
-    fn(x, val, idxs...)
+    fn = compile(fill!, (x, val))
+    fn(x, val)
     return x
 end
 

src/Ops.jl

@@ -532,7 +532,7 @@ end
     limit_indices = limit_indices
     rsize = limit_indices .- start_indices
     @assert all(rsize .> 0) "Invalid slice dimensions"
-    strides = isnothing(strides) ? [1, size(x)[1:(end - 1)]...] : strides
+    strides = isnothing(strides) ? ones(Int64, N) : strides
     res = MLIR.IR.result(
         stablehlo.slice(
             x.mlir_data;
@@ -1699,7 +1699,8 @@ instead.
 @noinline function scatter_setindex(
     dest::TracedRArray{T,N},
     scatter_indices::TracedRArray{Int64,2},
-    updates::TracedRArray{T2,1},
+    updates::TracedRArray{T2,1};
+    location=mlir_stacktrace("scatter_setindex", @__FILE__, @__LINE__),
 ) where {T,N,T2}
     @assert length(updates) == size(scatter_indices, 1)
     @assert size(scatter_indices, 2) == N
@@ -1716,14 +1717,39 @@ instead.
     push!(block, return_op)
     pushfirst!(update_computation, block)
 
-    #! format: off
-    update_window_dims = Int64[]
-    inserted_window_dims = collect(Int64, 0:(N - 1))
-    input_batching_dims = Int64[]
-    scatter_indices_batching_dims = Int64[]
-    scatter_dims_to_operand_dims = collect(Int64, 0:(N - 1))
-    index_vector_dim = Int64(1)
+    return scatter(
+        [dest],
+        scatter_indices,
+        [updates];
+        update_computation,
+        update_window_dims=Int64[],
+        inserted_window_dims=collect(Int64, 0:(N - 1)),
+        input_batching_dims=Int64[],
+        scatter_indices_batching_dims=Int64[],
+        scatter_dims_to_operand_dims=collect(Int64, 0:(N - 1)),
+        index_vector_dim=Int64(1),
+        location,
+    )[1]
+end
+
+@noinline function scatter(
+    dest::Vector{TracedRArray{T,N}},
+    scatter_indices::TracedRArray{Int64,2},
+    updates::Vector{<:TracedRArray{T}};
+    update_computation::MLIR.IR.Region,
+    update_window_dims::Vector{Int64},
+    inserted_window_dims::Vector{Int64},
+    input_batching_dims::Vector{Int64},
+    scatter_indices_batching_dims::Vector{Int64},
+    scatter_dims_to_operand_dims::Vector{Int64},
+    index_vector_dim::Int64,
+    location=mlir_stacktrace("scatter", @__FILE__, @__LINE__),
+) where {T,N}
+    scatter_indices = subtract(
+        scatter_indices, fill(Int64(1), size(scatter_indices)); location
+    )
 
+    #! format: off
     scatter_dimension_numbers = MLIR.API.stablehloScatterDimensionNumbersGet(
         MLIR.IR.context(),
         length(update_window_dims), update_window_dims,
@@ -1735,21 +1761,22 @@ instead.
     )
     #! format: on
 
-    return TracedRArray{T,N}(
-        (),
-        MLIR.IR.result(
-            MLIR.Dialects.stablehlo.scatter(
-                [dest.mlir_data],
-                scatter_indices.mlir_data,
-                [updates.mlir_data];
-                result_0=[mlir_type(TracedRArray{T,N}, size(dest))],
-                update_computation,
-                scatter_dimension_numbers,
-            ),
-            1,
-        ),
-        size(dest),
+    dest_values = [d.mlir_data for d in dest]
+    update_values = [u.mlir_data for u in updates]
+    scatter_op = stablehlo.scatter(
+        dest_values,
+        scatter_indices.mlir_data,
+        update_values;
+        update_computation,
+        scatter_dimension_numbers,
+        result_0=[mlir_type(TracedRArray{T,N}, size(d)) for d in dest],
+        location,
     )
+
+    return [
+        TracedRArray{T,N}((), MLIR.IR.result(scatter_op, i), size(dest[i])) for
+        i in eachindex(dest)
+    ]
 end
 
 """
@@ -1760,7 +1787,9 @@ specified by `gather_indices`. If the indices are contiguous it is recommended t
 use [`MLIR.Dialects.stablehlo.dynamic_slice`](@ref) instead.
 """
 @noinline function gather_getindex(
-    src::TracedRArray{T,N}, gather_indices::TracedRArray{Int64,2}
+    src::TracedRArray{T,N},
+    gather_indices::TracedRArray{Int64,2};
+    location=mlir_stacktrace("gather_getindex", @__FILE__, @__LINE__),
 ) where {T,N}
     @assert size(gather_indices, 2) == N
 
@@ -1770,14 +1799,42 @@ use [`MLIR.Dialects.stablehlo.dynamic_slice`](@ref) instead.
         )
     end
 
-    #! format: off
-    offset_dims = Int64[1]
-    collapsed_slice_dims = collect(Int64, 0:(N - 2))
-    operand_batching_dims = Int64[]
-    start_indices_batching_dims = Int64[]
-    start_index_map = collect(Int64, 0:(N - 1))
-    index_vector_dim = Int64(1)
+    return reshape(
+        gather(
+            src,
+            gather_indices;
+            offset_dims=Int64[1],
+            collapsed_slice_dims=collect(Int64, 0:(N - 2)),
+            operand_batching_dims=Int64[],
+            start_indices_batching_dims=Int64[],
+            start_index_map=collect(Int64, 0:(N - 1)),
+            index_vector_dim=Int64(1),
+            slice_sizes=ones(Int64, N),
+            indices_are_sorted=false,
+            location,
+        ),
+        [size(gather_indices, 1)],
+    )
+end
+
+@noinline function gather(
+    src::TracedRArray{T,N},
+    gather_indices::TracedRArray{Int64,2};
+    offset_dims::Vector{Int64},
+    collapsed_slice_dims::Vector{Int64},
+    operand_batching_dims::Vector{Int64},
+    start_indices_batching_dims::Vector{Int64},
+    start_index_map::Vector{Int64},
+    index_vector_dim::Int64,
+    slice_sizes::Vector{Int64},
+    indices_are_sorted::Bool=false,
+    location=mlir_stacktrace("gather", @__FILE__, @__LINE__),
+) where {T,N}
+    gather_indices = subtract(
+        gather_indices, fill(Int64(1), size(gather_indices)); location
+    )
 
+    #! format: off
     dimension_numbers = MLIR.API.stablehloGatherDimensionNumbersGet(
         MLIR.IR.context(),
         Int64(length(offset_dims)), offset_dims,
@@ -1789,20 +1846,18 @@ use [`MLIR.Dialects.stablehlo.dynamic_slice`](@ref) instead.
     )
     #! format: on
 
-    return reshape(
-        TracedRArray{T}(
-            MLIR.IR.result(
-                MLIR.Dialects.stablehlo.gather(
-                    src.mlir_data,
-                    gather_indices.mlir_data;
-                    dimension_numbers,
-                    slice_sizes=Base.fill(Int64(1), N),
-                    indices_are_sorted=false,
-                ),
-                1,
+    return TracedRArray{T}(
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.gather(
+                src.mlir_data,
+                gather_indices.mlir_data;
+                dimension_numbers,
+                slice_sizes,
+                indices_are_sorted,
+                location,
             ),
+            1,
         ),
-        size(gather_indices, 1),
     )
 end
 

src/Overlay.jl

@@ -8,14 +8,6 @@
     return f
 end
 
-@reactant_overlay @noinline function Base.setindex!(
-    a::AnyTracedRArray{T,N}, v, indices::Vararg{Any,N}
-) where {T,N}
-    ancestor_indices = TracedUtils.get_ancestor_indices(a, indices...)
-    (Base.inferencebarrier(setindex!))(Reactant.ancestor(a), v, ancestor_indices...)
-    return a
-end
-
 # Enzyme.jl overlays
 @reactant_overlay @noinline function Enzyme.autodiff_deferred(
     rmode::Enzyme.Mode, f::FA, rt::Type{A}, args::Vararg{Annotation,Nargs}