feat: generalize Ops.scatter to handle tracing (#1392)

* feat: generalize Ops.scatter to handle tracing

* Update src/TracedRArray.jl

Author: avik-pal

src/Ops.jl

@@ -1678,23 +1678,11 @@ instead.
     @assert length(updates) == size(scatter_indices, 1)
     @assert size(scatter_indices, 2) == N
 
-    updates = convert(TracedRArray{T,1}, updates)
-
-    update_computation = MLIR.IR.Region()
-    block = MLIR.IR.Block(
-        [mlir_type(TracedRNumber{T}), 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)
-
     return scatter(
+        (a, b) -> b,
         [dest],
         scatter_indices,
-        [updates];
-        update_computation,
+        [convert(TracedRArray{T,1}, updates)];
         update_window_dims=Int64[],
         inserted_window_dims=collect(Int64, 1:N),
         input_batching_dims=Int64[],
@@ -1705,6 +1693,36 @@ instead.
     )[1]
 end
 
+@noinline function scatter(
+    f::F,
+    dest::Vector{TracedRArray{T,N}},
+    scatter_indices::TracedRArray{Int64},
+    updates::Vector{<:TracedRArray{T}};
+    location=mlir_stacktrace("scatter", @__FILE__, @__LINE__),
+    kwargs...,
+) where {F,T,N}
+    sample_inputs = (
+        Reactant.TracedUtils.promote_to(TracedRNumber{T}, zero(T)),
+        Reactant.TracedUtils.promote_to(TracedRNumber{T}, zero(T)),
+    )
+
+    compiled_fn =
+        Reactant.TracedUtils.make_mlir_fn(
+            f,
+            sample_inputs,
+            (),
+            "update_computation",
+            false;
+            args_in_result=:result,
+            return_dialect=:stablehlo,
+        ).f
+    update_computation = MLIR.IR.Region()
+    MLIR.API.mlirRegionTakeBody(update_computation, MLIR.IR.region(compiled_fn, 1))
+    MLIR.IR.rmfromparent!(compiled_fn)
+
+    return scatter(dest, scatter_indices, updates; update_computation, location, kwargs...)
+end
+
 @noinline function scatter(
     dest::Vector{TracedRArray{T,N}},
     scatter_indices::TracedRArray{Int64},

src/TracedRArray.jl

@@ -439,22 +439,11 @@ function Base.setindex!(a::TracedRArray{T,N}, v, indices::Vararg{Any,N}) where {
         ]
         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(
+            (xᵢ, xⱼ) -> xⱼ,
             [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[],