feat: support vector mode AD (#519)

* feat: support vector mode AD

* Update test/autodiff.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update Interpreter.jl

* test: new JLL

* fix: codegen for vector mode

* fix: bump version

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: William Moses <[email protected]>

Author: 3 people

Project.toml

@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.9"
-Reactant_jll = "0.0.151"
+Reactant_jll = "0.0.152"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Interpreter.jl

@@ -192,7 +192,11 @@ function push_acts!(ad_inputs, x::BatchDuplicated, path, reverse)
         predims = size(x.val)
         cval = MLIR.IR.result(
             MLIR.Dialects.stablehlo.concatenate(
-                [Ops.reshape(v, Int64[1, predims...]) for v in x.dval]; dimension=Int64(0)
+                [
+                    TracedUtils.get_mlir_data(Ops.reshape(v, Int64[1, predims...])) for
+                    v in x.dval
+                ];
+                dimension=Int64(0),
             ),
         )
         tval = TracedRArray{ET,length(predims) + 1}((), cval, (length(x.dval), predims...))
@@ -244,12 +248,6 @@ function overload_autodiff(
     width = Enzyme.same_or_one(1, args...)
     if width == 0
         throw(ErrorException("Cannot differentiate with a batch size of 0"))
-    elseif width != 1
-        throw(
-            ErrorException(
-                "EnzymeMLIR does not presently support width=$width, please rewrite your code to not use BatchDuplicated and/or call gradient(; chunk=1)",
-            ),
-        )
     end
 
     primf = f.val
@@ -389,9 +387,10 @@ function overload_autodiff(
     fname = TracedUtils.get_attribute_by_name(func2, "sym_name")
     fname = MLIR.IR.FlatSymbolRefAttribute(Base.String(fname))
     res = (reverse ? MLIR.Dialects.enzyme.autodiff : MLIR.Dialects.enzyme.fwddiff)(
-        [TracedUtils.transpose_val(v) for v in ad_inputs];
+        [TracedUtils.transpose_val(v; keep_first_intact=width > 1) for v in ad_inputs];
         outputs=outtys,
         fn=fname,
+        width,
         activity=MLIR.IR.Attribute([act_attr(a) for a in activity]),
         ret_activity=MLIR.IR.Attribute([act_attr(a) for a in ret_activity]),
     )
@@ -434,8 +433,11 @@ function overload_autodiff(
                             push!(starts, 0)
                             push!(limits, v)
                         end
-                        sval = Ops.slice(sval, starts, limits)
-                        TracedUtils.set!(dresult[i], path[2:end], sval)
+                        sval = Ops.slice(TracedRArray(tval), starts, limits)
+                        sval = Ops.reshape(sval, collect(Int64, sz))
+                        TracedUtils.set!(
+                            dresult[i], path[2:end], TracedUtils.get_mlir_data(sval)
+                        )
                     end
                 end
                 residx += 1

src/TracedUtils.jl

@@ -186,11 +186,15 @@ Base.@nospecializeinfer function transpose_ty(
 end
 
 Base.@nospecializeinfer function transpose_val(
-    @nospecialize(val::MLIR.IR.Value)
+    @nospecialize(val::MLIR.IR.Value); keep_first_intact::Bool=false
 )::MLIR.IR.Value
     val_size = size(MLIR.IR.type(val))
     val_size == () && return val
-    attr = MLIR.IR.DenseArrayAttribute(Int64[reverse(0:(length(val_size) - 1))...])
+    if keep_first_intact
+        attr = MLIR.IR.DenseArrayAttribute(Int64[0, reverse(1:(length(val_size) - 1))...])
+    else
+        attr = MLIR.IR.DenseArrayAttribute(Int64[reverse(0:(length(val_size) - 1))...])
+    end
     return MLIR.IR.result(MLIR.Dialects.stablehlo.transpose(val; permutation=attr), 1)
 end
 

test/autodiff.jl

@@ -156,3 +156,17 @@ end
     hlo = @code_hlo optimize = false Enzyme.onehot(x)
     @test !contains("stablehlo.constant", repr(hlo))
 end
+
+fn(x) = sum(abs2, x)
+vector_forward_ad(x) = Enzyme.autodiff(Forward, fn, BatchDuplicated(x, Enzyme.onehot(x)))
+
+@testset "Vector Mode AD" begin
+    x = Reactant.to_rarray(reshape(collect(Float32, 1:4), 2, 2))
+    res = @jit vector_forward_ad(x)
+    res_enz = vector_forward_ad(Array(x))
+
+    @test res[1][1] ≈ res_enz[1][1]
+    @test res[1][2] ≈ res_enz[1][2]
+    @test res[1][3] ≈ res_enz[1][3]
+    @test res[1][4] ≈ res_enz[1][4]
+end