feat: add a preference to disable loop vectorization

Project.toml

@@ -19,6 +19,7 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -78,6 +79,7 @@ MLDataDevices = "1.2"
 Markdown = "1.10"
 NNlib = "0.9.24"
 Octavian = "0.3.28"
+Preferences = "1.4.3"
 Polyester = "0.7.15"
 Random = "1.10"
 Reexport = "1"

src/LuxLib.jl

@@ -1,6 +1,7 @@
 module LuxLib
 
 using Compat: @compat
+using Preferences: @load_preference
 using Reexport: @reexport
 using Static: Static, known
 
@@ -15,6 +16,8 @@ const Numeric = Union{AbstractArray{<:T}, T} where {T <: Number}
 const ∂∅ = NoTangent()
 const CRC = ChainRulesCore
 
+const DISABLE_LOOP_VECTORIZATION = @load_preference("disable_loop_vectorization", false)
+
 include("utils.jl")
 include("traits.jl")
 include("impl/Impl.jl")

src/traits.jl

@@ -80,6 +80,7 @@ using ChainRulesCore: ChainRulesCore
 using Hwloc: Hwloc
 using Static: static, False, True
 
+using ..LuxLib: DISABLE_LOOP_VECTORIZATION
 using ..Utils: is_extension_loaded, safe_minimum
 
 const CRC = ChainRulesCore
@@ -130,9 +131,13 @@ end
 
 CRC.@non_differentiable explicit_blas_loaded()
 
-function use_octavian()
-    return is_extension_loaded(Val(:Octavian)) & is_x86_64() &
-           (INTEL_HARDWARE | AMD_RYZEN_HARDWARE)
+@static if DISABLE_LOOP_VECTORIZATION
+    use_octavian() = False()
+else
+    function use_octavian()
+        return is_extension_loaded(Val(:Octavian)) & is_x86_64() &
+               (INTEL_HARDWARE | AMD_RYZEN_HARDWARE)
+    end
 end
 
 CRC.@non_differentiable use_octavian()

src/utils.jl

@@ -11,7 +11,7 @@ using NNlib: NNlib
 using Static: Static, StaticBool, False, True, static
 using StaticArraysCore: SVector, SMatrix
 
-using ..LuxLib: Optional, ∂∅
+using ..LuxLib: Optional, ∂∅, DISABLE_LOOP_VECTORIZATION
 
 const CRC = ChainRulesCore
 const KA = KernelAbstractions
@@ -325,8 +325,12 @@ end
 
 CRC.@non_differentiable static_training_mode_check(::Any...)
 
-@inline function can_loopvec_args(args...)
-    return can_loopvec_args_check(is_extension_loaded(Val(:LoopVectorization)), args...)
+@static if DISABLE_LOOP_VECTORIZATION
+    @inline can_loopvec_args(args...) = false
+else
+    @inline function can_loopvec_args(args...)
+        return can_loopvec_args_check(is_extension_loaded(Val(:LoopVectorization)), args...)
+    end
 end
 
 @inline can_loopvec_args_check(::False, args...) = false