perf: restore old batched_mul

Project.toml

@@ -65,7 +65,7 @@ ChainRulesCore = "1.24"
 Compat = "4.15.0"
 CpuId = "0.3"
 DispatchDoctor = "0.4.12"
-Enzyme = "0.13.1"
+Enzyme = "0.13.12"
 EnzymeCore = "0.8.1"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36"

src/impl/batched_mul.jl

@@ -61,9 +61,7 @@ function batched_matmul_cpu!(z::AbstractArray{zT, 3}, x::AbstractArray{xT, 3},
         batched_matmul_loopvec_impl!(z, x, y)
         return
     end
-    # Avoid an Enzyme segfault https://github.com/EnzymeAD/Enzyme.jl/issues/1983
-    fallback_batched_matmul!(z, LoopedArrayOp(), x, y)
-    # NNlib.batched_mul!(z, x, y) # XXX: restore once the enzyme segfault is fixed
+    NNlib.batched_mul!(z, x, y)
     return
 end
 
@@ -80,10 +78,9 @@ end
 function fallback_batched_matmul!(
         z::AbstractArray{zT, 3}, dev, x::AbstractArray{xT, 3},
         y::AbstractArray{yT, 3}) where {zT, xT, yT}
-    # XXX: bring back once the enzyme segfault is fixed
-    # @warn "Using fallback Batched Matrix Multiply routine for $(dev) with A: size = \
-    #        $(size(x)) eltype = $(xT) and B: size = $(size(y)) eltype = $(yT). This may be \
-    #        slow." maxlog=1
+    @warn "Using fallback Batched Matrix Multiply routine for $(dev) with A: size = \
+           $(size(x)) eltype = $(xT) and B: size = $(size(y)) eltype = $(yT). This may be \
+           slow." maxlog=1
 
     if (size(x, 3) != size(y, 3) && size(x, 3) != 1 && size(y, 3) != 1) ||
        (size(x, 2) != size(y, 1))