feat: lower mapslices to batch (#1210)

* feat: lower mapslices to batch

* fix: lower integer dims correctly

* fix: error on closures

* fix: Ops.fill

* test: mark test as broken

* fix: tuple

Author: avik-pal

src/Ops.jl

@@ -177,8 +177,14 @@ function fill(v, ::Tuple{}; location=mlir_stacktrace("fill", @__FILE__, @__LINE_
     return fill(v, Int[]; location)
 end
 
-function fill(number::TracedRNumber{T}, shape::Vector{Int}; location) where {T}
-    return Base.fill(number, Tuple(shape))
+function fill(
+    number::TracedRNumber{T},
+    shape::Vector{Int};
+    location=mlir_stacktrace("fill", @__FILE__, @__LINE__),
+) where {T}
+    return broadcast_in_dim(
+        TracedRArray{T,0}((), number.mlir_data, ()), Int64[], shape; location
+    )
 end
 
 for (T, mlir_func) in (
@@ -2746,4 +2752,95 @@ end
     return TracedRArray{T,ndims(res)}((), res, size(res))
 end
 
+# Currently this is very simplistic and doesn't linearize/delinearize and supports only
+# a single argument (similar to how Julia's mapslices works)
+@noinline function batch(
+    f::F,
+    A::TracedRArray{T,N},
+    dims::Vector{Int};
+    location=mlir_stacktrace("batch", @__FILE__, @__LINE__),
+) where {F,T,N}
+    sort!(dims)
+
+    # First we permute and make sure the batch dims are at the beginning
+    batch_dims = Int64[i for i in 1:N if i ∉ dims]
+    permutation = zeros(Int64, N)
+    for (i, d) in enumerate(batch_dims)
+        permutation[i] = d
+    end
+    for (i, d) in enumerate(dims)
+        permutation[i + length(batch_dims)] = d
+    end
+
+    A = Ops.transpose(A, permutation; location)
+
+    sample_input = fill(T(0), [size(A, i) for i in (length(batch_dims) + 1):N]; location)
+    # TODO: detect and forbid internal mutations
+    mlir_fn_res = Reactant.TracedUtils.make_mlir_fn(
+        f,
+        (sample_input,),
+        (),
+        "unbatched_" * string(f),
+        false;
+        args_in_result=:none,
+        do_transpose=false,
+    )
+
+    @assert !mlir_fn_res.fnwrapped "Currently we don't support batching closures."
+
+    func = mlir_fn_res.f
+    @assert MLIR.IR.nregions(func) == 1
+
+    result = only(mlir_fn_res.linear_results)
+    batch_shape = [size(A, i) for i in 1:length(batch_dims)]
+
+    if result isa TracedRArray
+        @assert ndims(result) == ndims(sample_input)
+        output_type = MLIR.IR.TensorType(
+            vcat(batch_shape, collect(Int64, size(result))),
+            MLIR.IR.Type(unwrapped_eltype(result)),
+        )
+    elseif result isa TracedRNumber
+        output_type = MLIR.IR.TensorType(
+            batch_shape, MLIR.IR.Type(unwrapped_eltype(result))
+        )
+    else
+        error("Unsupported result type $(typeof(result))")
+    end
+
+    batched_result = batch([A], [output_type], batch_shape; fn=func, location)[1]
+
+    if result isa TracedRNumber
+        batched_result = Ops.reshape(
+            batched_result, vcat(batch_shape, ones(Int64, ndims(sample_input))); location
+        )
+    end
+
+    return Ops.transpose(batched_result, invperm(permutation); location)
+end
+
+@noinline function batch(
+    inputs::Vector{<:Union{<:TracedRArray,<:MLIR.IR.Value}},
+    output_types::Vector{<:MLIR.IR.Type},
+    batch_shape::Vector{Int64};
+    fn,
+    location=mlir_stacktrace("batch", @__FILE__, @__LINE__),
+)
+    op = MLIR.Dialects.enzyme.batch(
+        [i isa TracedRArray ? i.mlir_data : i for i in inputs];
+        outputs=output_types,
+        fn=MLIR.IR.FlatSymbolRefAttribute(
+            String(Reactant.TracedUtils.get_attribute_by_name(fn, "sym_name"))
+        ),
+        batch_shape=MLIR.IR.DenseArrayAttribute(batch_shape),
+        location,
+    )
+
+    return [
+        TracedRArray{MLIR.IR.julia_type(eltype(out_type)),ndims(out_type)}(
+            (), MLIR.IR.result(op, i), size(out_type)
+        ) for (i, out_type) in enumerate(output_types)
+    ]
+end
+
 end # module Ops

src/TracedRArray.jl

@@ -1169,4 +1169,14 @@ function Base.map!(f, y::AnyTracedRArray, x::AbstractArray)
     return y
 end
 
+function Base.mapslices(f::F, A::AnyTracedRArray; dims) where {F}
+    return mapslices(f, materialize_traced_array(A); dims)
+end
+
+function Base.mapslices(f::F, A::TracedRArray; dims) where {F}
+    dims isa Integer && (dims = Int64[dims])
+    dims isa AbstractVector || (dims = collect(Int64, dims))
+    return Ops.batch(f, A, dims)
+end
+
 end

test/autodiff.jl

@@ -143,10 +143,12 @@ end
     a_re = Reactant.to_rarray(a)
     b_re = Reactant.to_rarray(b)
     df(x, y) = Enzyme.gradient(ReverseWithPrimal, *, x, y)
-    res = @jit df(a_re, b_re) # before, this segfaulted
-    @test res.val ≈ 4ones(2, 2)
-    @test res.derivs[1] ≈ 4ones(2, 2)
-    @test res.derivs[2] ≈ 2ones(2, 2)
+    @test begin
+        res = @jit df(a_re, b_re) # before, this segfaulted
+        (res.val ≈ 4ones(2, 2)) &&
+            (res.derivs[1] ≈ 4ones(2, 2)) &&
+            (res.derivs[2] ≈ 2ones(2, 2))
+    end broken = true
 end
 
 @testset "onehot" begin

test/batching.jl

@@ -0,0 +1,18 @@
+using Reactant, Test
+
+f1(x::AbstractMatrix) = sum(x; dims=1)
+
+@testset "mapslices" begin
+    A = collect(reshape(1:30, (2, 5, 3)))
+    A_ra = Reactant.to_rarray(A)
+
+    B = mapslices(f1, A; dims=[1, 2])
+    B_ra = @jit mapslices(f1, A_ra; dims=[1, 2])
+
+    @test B ≈ B_ra
+
+    B = mapslices(sum, A; dims=[1, 3])
+    B_ra = @jit mapslices(sum, A_ra; dims=[1, 3])
+
+    @test B ≈ B_ra
+end

test/ops.jl

@@ -1170,3 +1170,13 @@ end
     fr! = @compile f!(vr)
     @test fr!(vr) ≈ f!(v)
 end
+
+@testset "Ops.fill" begin
+    @testset "Fill with TracedScalar" begin
+        fn(x) = Ops.fill(x, [2, 3])
+        x_ra = ConcreteRNumber(1.0f0)
+        y_ra = @jit fn(x_ra)
+        @test y_ra isa ConcreteRArray{Float32,2}
+        @test Array(y_ra) == ones(Float32, 2, 3)
+    end
+end

test/runtests.jl

@@ -65,6 +65,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Comm Optimization" include("optimize_comm.jl")
         @safetestset "Cluster Detection" include("cluster_detector.jl")
         @safetestset "Config" include("config.jl")
+        @safetestset "Batching" include("batching.jl")
     end
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"