Merge pull request #153 from FluxML/sf/nnpack_fix

Fix asymmetric padding arguments to NNPACK's `conv()`

Author: DhairyaLGandhi

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/impl/depthwiseconv_direct.jl

@@ -20,7 +20,7 @@ See the docstring for `conv_direct!()` for more on the optional parameters.
 """
 function depthwiseconv_direct!(y::AbstractArray{yT,5}, x::AbstractArray{xT,5},
                       w::AbstractArray{wT,5}, cdims::DepthwiseConvDims;
-                      alpha::yT = yT(1), beta = false) where {yT, xT, wT}
+                      alpha::yT=yT(1), beta=false) where {yT, xT, wT}
     check_dims(size(x), size(w), size(y), cdims)
 
     width, height, depth = input_size(cdims)
@@ -135,7 +135,7 @@ for each batch and channel independently.
 function ∇depthwiseconv_data_direct!(
                 dx::AbstractArray{xT,5}, dy::AbstractArray{yT,5},
                 w::AbstractArray{wT,5}, cdims::DepthwiseConvDims;
-                alpha::xT=xT(1), beta::xT=xT(0)) where {xT, yT, wT}
+                alpha::xT=xT(1), beta=false) where {xT, yT, wT}
     # We do a separate convolution for each channel in x
     @inbounds for cidx in 1:channels_in(cdims)
         # For this batch and in-channel, we have a normal transposed convolution
@@ -168,7 +168,7 @@ Calculate the gradient imposed upon `w` in the depthwise convolution `y = x * w`
 function ∇depthwiseconv_filter_direct!(
                 dw::AbstractArray{wT,5}, x::AbstractArray{xT,5},
                 dy::AbstractArray{yT,5}, cdims::DepthwiseConvDims;
-                alpha::wT=wT(1),beta::wT=wT(0)) where {xT, yT, wT}
+                alpha::wT=wT(1),beta=false) where {xT, yT, wT}
     # We do a separate convolution for each channel in x
     @inbounds for cidx in 1:channels_in(cdims)
         # For this batch and in-channel, we have a normal transposed convolution

src/impl/depthwiseconv_im2col.jl

@@ -14,7 +14,7 @@ function depthwiseconv_im2col!(
                 y::AbstractArray{T,5}, x::AbstractArray{T,5},
                 w::AbstractArray{T,5}, cdims::DepthwiseConvDims;
                 col::AbstractArray{T,2} = similar(x, im2col_dims(cdims)),
-                alpha=T(1), beta=T(0)) where T
+                alpha::T=T(1), beta::T=T(0)) where T
     check_dims(size(x), size(w), size(y), cdims)
 
     # This functions exactly the same as conv_im2col!(), except that we shard the
@@ -56,7 +56,7 @@ function ∇depthwiseconv_filter_im2col!(
                 dw::AbstractArray{T,5}, x::AbstractArray{T,5},
                 dy::AbstractArray{T,5}, cdims::DepthwiseConvDims;
                 col::AbstractArray{T,2} = similar(dw, im2col_dims(cdims)),
-                alpha=T(1), beta=T(0)) where T
+                alpha::T=T(1), beta::T=T(0)) where T
     check_dims(size(x), size(dw), size(dy), cdims)
 
     M = prod(kernel_size(cdims))
@@ -96,7 +96,7 @@ function ∇depthwiseconv_data_im2col!(
                 dx::AbstractArray{T,5}, dy::AbstractArray{T,5},
                 w::AbstractArray{T,5}, cdims::DepthwiseConvDims;
                 col::AbstractArray{T,2} = similar(dx, im2col_dims(cdims)),
-                alpha=T(1), beta=T(0)) where T
+                alpha::T=T(1), beta::T=T(0)) where T
     check_dims(size(dx), size(w), size(dy), cdims)
 
     M = prod(output_size(cdims))

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/nnpack/libnnpack.jl

@@ -127,7 +127,7 @@ end
 function nnp_convolution_output(y::Array{Float32,4}, x::Array{Float32,4}, w::Array{Float32,4}, b::Array{Float32,1}; algo::nnp_convolution_algorithm = UInt32(0), workspace_buffer = nothing, workspace_size = 0, padding = 0, stride = 1, threadpool = C_NULL, profile = nothing)
     input_size = nnp_size(Csize_t.((size(x,1), size(x,2)))...)
     kernel_size = nnp_size(Csize_t.((size(w,1),size(w,2)))...)
-    input_padding = nnp_padding(Csize_t(padding[2]), Csize_t(padding[1]), Csize_t(padding[2]), Csize_t(padding[1]))
+    input_padding = nnp_padding(Csize_t(padding[3]), Csize_t(padding[2]), Csize_t(padding[4]), Csize_t(padding[1]))
     profile = profile == nothing ? nnp_profile() : profile
     workspace_buffer = workspace_buffer === nothing ? C_NULL : workspace_buffer
     nnp_convolution_output(UInt32(algo), size(x,4), size(x,3), size(w,4), input_size, input_padding, kernel_size, x, w, b, y, workspace_buffer, workspace_size, UInt32(0), C_NULL, threadpool, profile)

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: