Integrate NNPACK convolutions into tests a bit more

Author: staticfloat

src/conv.jl

@@ -159,7 +159,7 @@ end
 if is_nnpack_available()
     function conv(x::Array{xT, 4}, w::Array{wT, 4},
                   cdims::DenseConvDims{2, K, C_in, C_out, (1, 1), P, (1, 1), F};
-                  kwargs...) where {xT, wT, K, C_in, C_out, S, P, F}
+                  kwargs...) where {xT, wT, K, C_in, C_out, P, F}
         return conv_nnpack(x, w, cdims; kwargs...)
     end
 end

src/nnpack/interface.jl

@@ -52,11 +52,19 @@ end
 
 
 """
-    check_supported_operation(x::Array, pdims::PoolDims)
+    nnpack_supported_operation(cdims::ConvDims)
+    nnpack_supported_operation(pdims::PoolDims)
 
-Returns `true` if nnpack supports the pooling operation for the given input.
+Returns `true` if nnpack supports the convolution/pooling operation for the given parameters.
 """
-function check_supported_operation(x::Array{T, 4}, pdims::PoolDims{2, K, S, P, (1, 1)}) where {T, K, S, P}
-    val = size(x)[1:2] .+ (P[1] + P[2], P[3] + P[4]) .- K
+function nnpack_supported_operation(pdims::PoolDims{2, K, S, P, (1, 1)}) where {K, S, P}
+    val = input_size(pdims)[1:2] .+ (P[1] + P[2], P[3] + P[4]) .- K
     return val .% S == (0, 0) ? true : false
 end
+
+function nnpack_supported_operation(cdims::ConvDims{2, K, (1, 1), P, (1, 1)}) where {K, S, P}
+    return true
+end
+
+# Return false for everything else
+nnpack_supported_operation(dims) = false

src/pooling.jl

@@ -132,7 +132,7 @@ end
 # Use NNPACK if it is available and operation is supported
 if is_nnpack_available()
     function maxpool(x::Array{T, 4}, pdims::PoolDims{2, K, S, P, (1, 1)}; kwargs...) where {T, K, S, P}
-        func = check_supported_operation(x, pdims) ? maxpool_nnpack : maxpool_direct
+        func = nnpack_supported_operation(pdims) ? maxpool_nnpack : maxpool_direct
         return func(x, pdims; kwargs...)
     end
 end

test/conv.jl

@@ -274,10 +274,13 @@ conv_answer_dict = Dict(
             # A "drop channels and batch dimension" helper
             ddims(x) = dropdims(x, dims=(rank+1, rank+2))
 
-            for conv in (NNlib.conv, NNlib.conv_im2col, NNlib.conv_direct)
+            for conv in (NNlib.conv, NNlib.conv_im2col, NNlib.conv_direct, NNlib.conv_nnpack)
+                if conv == NNlib.conv_nnpack && !NNlib.nnpack_supported_operation(DenseConvDims(x, w))
+                    continue
+                end
                 @testset "$(conv)" begin
-                    # First, your basic convolution with no parameters
                     cdims = DenseConvDims(x, w)
+                    # First, your basic convolution with no parameters
                     @test isapprox(ddims(conv(x, w, cdims)), y_plain, rtol = 1.0e-7)
 
                     # Next, test convolution on views and alternate datatypes: