From ad7790b3cb7e484bffdc58b29a60b2a1b03e8ee0 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 17 Oct 2024 18:31:03 -0400
Subject: [PATCH] fix: patch more enzyme issues

---
 ext/LuxLibLoopVectorizationExt.jl   |  2 +-
 src/impl/batched_mul.jl             | 15 +++++++++++----
 src/utils.jl                        |  7 +++++++
 test/common_ops/activation_tests.jl |  2 +-
 test/common_ops/bias_act_tests.jl   | 11 +++++------
 5 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/ext/LuxLibLoopVectorizationExt.jl b/ext/LuxLibLoopVectorizationExt.jl
index 87a912be..14061903 100644
--- a/ext/LuxLibLoopVectorizationExt.jl
+++ b/ext/LuxLibLoopVectorizationExt.jl
@@ -49,7 +49,7 @@ end
 
 # batched matmul
 function LuxLib.Impl.batched_matmul_loopvec_impl!(
-        z::AbstractArray{zT, 3}, x::AbstractArray{xT, 3},
+        ::True, z::AbstractArray{zT, 3}, x::AbstractArray{xT, 3},
         y::AbstractArray{yT, 3}, α::Number=true, β::Number=false) where {zT, xT, yT}
     if size(x, 3) == size(y, 3)
         @batch for L in axes(z, 3)
diff --git a/src/impl/batched_mul.jl b/src/impl/batched_mul.jl
index 911e7927..7517b82c 100644
--- a/src/impl/batched_mul.jl
+++ b/src/impl/batched_mul.jl
@@ -54,14 +54,21 @@ function batched_matmul!(z::AbstractArray{zT, 3}, ::LoopedArrayOp,
     return
 end
 
-function batched_matmul_cpu!(z::AbstractArray{zT, 3},
-        x::AbstractArray{xT, 3}, y::AbstractArray{yT, 3}) where {zT, xT, yT}
+function batched_matmul_cpu!(
+        z::AbstractArray{zT, 3}, x::AbstractArray{xT, 3}, y::AbstractArray{yT, 3},
+        α::Number=true, β::Number=false) where {zT, xT, yT}
     if can_loopvec_args(batchview(z, 1), batchview(x, 1), batchview(y, 1)) &&
        !unsafe_known(explicit_blas_loaded())
-        batched_matmul_loopvec_impl!(z, x, y)
+        batched_matmul_loopvec_impl!is_extension_loaded(
+            Val(:LoopVectorization), z, x, y, α, β)
         return
     end
-    NNlib.batched_mul!(z, x, y)
+    NNlib.batched_mul!(z, x, y, α, β)
+    return
+end
+
+function batched_matmul_loopvec_impl!(_, z, x, y, α, β)
+    NNlib.batched_mul!(z, x, y, α, β)
     return
 end
 
diff --git a/src/utils.jl b/src/utils.jl
index c96e4611..eaa60f08 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -18,6 +18,9 @@ const KA = KernelAbstractions
 
 is_extension_loaded(::Val) = False()
 
+CRC.@non_differentiable is_extension_loaded(::Any...)
+EnzymeRules.inactive_noinl(::typeof(is_extension_loaded), ::Any...) = nothing
+
 # Simple Operations -- no rrules needed
 ofeltype_array(::Type{T}, x::AbstractArray{T}) where {T} = x
 function ofeltype_array(
@@ -328,4 +331,8 @@ end
 
 @inline can_loopvec_args_check(::False, args...) = false
 
+CRC.@non_differentiable can_loopvec_args_check(::Any...)
+
+EnzymeRules.inactive_noinl(::typeof(can_loopvec_args_check), ::Any...) = nothing
+
 end
diff --git a/test/common_ops/activation_tests.jl b/test/common_ops/activation_tests.jl
index 2045f20f..e2b80e71 100644
--- a/test/common_ops/activation_tests.jl
+++ b/test/common_ops/activation_tests.jl
@@ -36,7 +36,7 @@
             @jet apply_act_fast2(f, x)
 
             @test @inferred(Zygote.gradient(apply_act, f, x)) isa Any
-            if f !== lisht || (f === lisht && T == Float32 && !ongpu)
+            if f !== lisht
                 @test @inferred(Zygote.gradient(apply_act_fast, f, x)) isa Any
             end
             @test @inferred(Zygote.gradient(apply_act_fast2, f, x)) isa Any
diff --git a/test/common_ops/bias_act_tests.jl b/test/common_ops/bias_act_tests.jl
index 1429c9b2..a7499654 100644
--- a/test/common_ops/bias_act_tests.jl
+++ b/test/common_ops/bias_act_tests.jl
@@ -44,12 +44,11 @@
             @jet bias_act_loss2(act, x, b)
             @jet bias_act_loss3(act, x, b)
 
-            if (act !== lisht || (act === lisht && T == Float32 && !ongpu)) && T != Float16
-                @test @inferred(Zygote.gradient(bias_act_loss2, act, x, b)) isa Any
-                @test @inferred(Zygote.gradient(bias_act_loss3, act, x, b)) isa Any
-            elseif T != Float16
-                @test_broken @inferred(Zygote.gradient(bias_act_loss2, act, x, b)) isa Any
-                @test_broken @inferred(Zygote.gradient(bias_act_loss3, act, x, b)) isa Any
+            if act !== lisht
+                @test @inferred(Zygote.gradient(bias_act_loss2, act, x, b)) isa Any broken=(T !=
+                                                                                            Float16)
+                @test @inferred(Zygote.gradient(bias_act_loss3, act, x, b)) isa Any broken=(T !=
+                                                                                            Float16)
             end
 
             @test_gradients(__Fix1(bias_act_loss1, act), x, b; atol, rtol,