add allowslow

Author: mcabbott

src/NNlib.jl

@@ -18,6 +18,15 @@ using Statistics: mean
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
+"""
+    allowslow(::Bool)
+
+By default, NNlib will print warnings the first time various slow fallback paths are taken.
+Calling `allowslow(false)` will instead make these into errors.
+"""
+allowslow(flag::Bool) = (SLOWERROR[] = !flag; nothing)
+const SLOWERROR = Ref(true)
+
 # Include APIs
 include("dim_helpers.jl")
 export ConvDims, DenseConvDims, PoolDims, DepthwiseConvDims

src/batched/batchedmul.jl

@@ -274,7 +274,10 @@ for (TA, fA) in _BATCHED_LIST, (TB, fB) in _BATCHED_LIST
 
         size(A, 3) == size(C, 3) || size(A, 3) == 1 || throw(DimensionMismatch("batch size mismatch: A != C"))
         size(B, 3) == size(C, 3) || size(B, 3) == 1 || throw(DimensionMismatch("batch size mismatch: B != C"))
-        @debug "calling fallback method for batched_mul!" typeof(A) size(A) typeof(B) size(B) typeof(C)
+        @warn "calling fallback method for batched_mul!" typeof(A) size(A) typeof(B) size(B) typeof(C) maxlog=1
+        if SLOWERROR[]
+            error("calling fallback method for batched_mul!")
+        end
 
         Abase, Bbase = _unbatch(A), _unbatch(B)
         sA, oA = size(A,3) == 1 ? (0,1) : (1,0)

src/conv.jl

@@ -191,6 +191,7 @@ for (front_name, backend, signature) in (
             if $(string(backend)) == "direct" && yT == Float64  # warn for Float32 + accidental Float64, but don't print warning for ForwardDiff.Dual
                 @warn string("Slow fallback implementation invoked for ", $(string(front_name)), "!  ",
                         "You probably don't want this; check your datatypes.") yT T1 T2 maxlog=1
+                SLOWERROR[] && error(string("calling slow fallback method for ", $(string(front_name))))
             end
 
             x_cs = Iterators.partition(1:size(in1, 4),
@@ -232,6 +233,7 @@ for (front_name, backend, signature) in (
             if $(string(backend)) == "direct" && yT == Float64  # warn for Float32 + accidental Float64, but don't print warning for ForwardDiff.Dual
                 @warn string("Slow fallback implementation invoked for ", $(string(front_name)), "!  ",
                         "You probably don't want this; check your datatypes.") yT T1 T2 maxlog=1
+                SLOWERROR[] && error(string("calling slow fallback method for ", $(string(front_name))))
             end
 
 
@@ -275,6 +277,7 @@ for (front_name, backend, signature) in (
             if $(string(backend)) == "direct" && yT == Float64  # warn for Float32 + accidental Float64, but don't print warning for ForwardDiff.Dual
                 @warn string("Slow fallback implementation invoked for ", $(string(front_name)), "!  ",
                         "You probably don't want this; check your datatypes.") yT T1 T2 maxlog=1
+                SLOWERROR[] && error(string("calling slow fallback method for ", $(string(front_name))))
             end
 
             dw_cs = Iterators.partition(1:size(out, 5),
@@ -326,6 +329,7 @@ for (front_name, backend, signature) in (
             if $(string(backend)) == "direct" && yT == Float64  # warn for Float32 + accidental Float64, but don't print warning for ForwardDiff.Dual
                 @warn string("Slow fallback implementation invoked for ", $(string(front_name)), "!  ",
                         "You probably don't want this; check your datatypes.") yT T1 T2 maxlog=1
+                SLOWERROR[] && error(string("calling slow fallback method for ", $(string(front_name))))
             end
             $(Symbol("$(front_name)_$(backend)!"))(out, in1, in2, cdims; kwargs...)
         end

test/batchedmul.jl

@@ -303,3 +303,14 @@ FiniteDifferences.to_vec(x::BatchedTranspose) = FiniteDifferences.to_vec(collect
 
     gradtest(batched_vec, randn(rng, M, P, B), randn(rng, P))
 end
+
+@testset "warning / error" begin
+    prev = NNlib.SLOWERROR[]
+    NNlib.allowslow(true)
+    A = rand(1:99, 3,4,7)
+    B = rand(1:99, 4,5,7)
+    @test batched_mul(A, B) isa Array  # no error!
+    NNlib.allowslow(false)
+    @test_throws Exception batched_mul(A, B)
+    NNlib.SLOWERROR[] = prev
+end