Clean the code

Author: Avik Pal

src/NNlib.jl

@@ -13,8 +13,15 @@ include("linalg.jl")
 include("conv.jl")
 include("cubroadcast.jl")
 
-if Sys.islinux()
+try
+    global ENABLE_NNPACK = parse(UInt64, ENV["ENABLE_NNPACK"])
+catch
+    global ENABLE_NNPACK = 1
+end
+
+if Sys.islinux() && ENABLE_NNPACK == 1
     include("nnpack/NNPACK.jl")
+    include("backends.jl")
 end
 
 end # module

src/backends.jl

@@ -0,0 +1,68 @@
+function nnpack_supported_operation(x::AbstractArray{<:Real, 4}, k, pad, stride, dilation)
+    fallback = false
+    # NNPACK does not support dilated convolutions
+    dilation == 1 || dilation == (1, 1) || (fallback = true)
+    # Expand the pad and stride to have same dimensions as k
+    pad_, stride_ = expand(Val{length(k)}, pad), expand(Val{length(k)}, stride)
+    (size(x, 1) - k[1] + 2 * pad_[1]) % stride_[1] == 0 && (size(x, 2) - k[2] + 2 * pad_[2]) % stride_[2] == 0 || (fallback = true)
+    # Return the pad_ and stride_ as well
+    return pad_, stride_, fallback
+end
+
+function nnpack_speed_check(x::AbstractArray{<:Real, 4}, k, pad, stride, dilation)
+    # Add heurestics here to determine whether or not to use NNPACK
+    # For now just return true
+    return true
+end
+
+# NNPACK supports only Float32 operations. So Float64 will have it default behaviour
+
+# Pooling
+function maxpool!(y::A, x::A, k; pad = map(_ -> 0, k), stride = k) where A<:Array{Float32, 4}
+    pad_, stride_, use_default = nnpack_supported_operation(x, k, pad, stride, 1)
+    use_nnpack = !use_default
+    # Only use NNPACK if we get speed improvement
+    use_nnpack && (use_nnpack = nnpack_speed_check(x, k, pad, stride, 1))
+    if use_nnpack
+        nnpack_max_pooling!(y, x, k, pad = pad_, stride = stride_)
+    else
+        maxpool_cpu!(y, x, k, pad = pad_, stride = stride_)
+    end
+end
+
+# Convolutions
+function conv!(y::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
+    k = (size(w, 1), size(w, 2))
+    pad_, stride_, use_default = nnpack_supported_operation(x, k, pad, stride, 1)
+    use_nnpack = !use_default
+    use_nnpack && (use_nnpack = nnpack_speed_check(x, k, pad, stride, 1))
+    if use_nnpack
+        nnpack_convolution_forward!(y, x, w, zeros(Float32, size(y, 3)), algo = algo, pad = pad, stride = stride, flipkernel = flipkernel)
+    else
+        conv2d!(y, x, w, padding = pad_, stride = stride_, dilation = dilation, mode = flipkernel)
+    end
+end
+
+function ∇conv_data!(dx::A, dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
+    k = (size(w, 1), size(w, 2))
+    pad_, stride_, use_default = nnpack_supported_operation(x, k, pad, stride, 1)
+    use_nnpack = !use_default
+    use_nnpack && (use_nnpack = nnpack_speed_check(x, k, pad, stride, 1))
+    if use_nnpack
+        nnpack_convolution_backward_data!(dx, x, dy, w, pad = pad_, stride = stride_, algo = algo, flipkernel = flipkernel)
+    else
+        conv2d_grad_x!(dx, x, w, dy, padding = pad_, stride = stride_, dilation = dilation, mode = flipkernel)
+    end
+end
+
+function ∇conv_filter!(dw::A, dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
+    k = (size(w, 1), size(w, 2))
+    pad_, stride_, use_default = nnpack_supported_operation(x, k, pad, stride, 1)
+    use_nnpack = !use_default
+    use_nnpack && (use_nnpack = nnpack_speed_check(x, k, pad, stride, 1))
+    if use_nnpack
+        nnpack_convolution_backward_filter!(dw, x, dy, w, pad = pad_, stride = stride_, algo = algo, flipkernel = flipkernel)
+    else
+        conv2d_grad_w!(dw, x, w, dy, padding = pad_, stride = stride_, dilation = dilation, mode = flipkernel)
+    end
+end

src/nnpack/NNPACK.jl

@@ -9,27 +9,20 @@ end
 include(depsjl_path)
 
 const nnlib_interface_path = joinpath(dirname(@__FILE__), "interface.jl")
-# const shared_threadpool = Ref(C_NULL)
+const shared_threadpool = Ref(C_NULL)
 
 @init begin
     check_deps()
+    status = nnp_initialize()
+    if status == nnp_status_unsupported_hardware
+        @warn "HARDWARE is unsupported by NNPACK so falling back to default NNlib"
+    else
+        include(nnlib_interface_path)
+    end
     try
-        global ENABLE_NNPACK = parse(UInt64, ENV["ENABLE_NNPACK"])
+        global NNPACK_CPU_THREADS = parse(UInt64, ENV["NNPACK_CPU_THREADS"])
     catch
-        global ENABLE_NNPACK = 1
-    end
-    if ENABLE_NNPACK == 1
-        status = nnp_initialize()
-        if status == nnp_status_unsupported_hardware
-            @warn "HARDWARE is unsupported by NNPACK so falling back to default NNlib"
-        else
-            include(nnlib_interface_path)
-        end
-        try
-            global NNPACK_CPU_THREADS = parse(UInt64, ENV["NNPACK_CPU_THREADS"])
-        catch
-            global NNPACK_CPU_THREADS = 4
-        end
-        global shared_threadpool = pthreadpool_create(NNPACK_CPU_THREADS)
+        global NNPACK_CPU_THREADS = 4
     end
+    shared_threadpool = pthreadpool_create(NNPACK_CPU_THREADS)
 end

src/nnpack/interface.jl

@@ -1,52 +1,29 @@
 flipweight(w::Array{<:Any,4}) = w[end:-1:1,end:-1:1,:,:]
 
-function check_support(x, k, pad, stride, dilation = 1)
-    fallback = false
-    dilation == 1 || dilation == (1, 1) || (fallback = true)
-    pad_, stride_ = expand(Val{length(k)}, pad), expand(Val{length(k)}, stride)
-    ((size(x, 1) - k[1] + 2 * pad_[1]) % stride_[1] == 0 && (size(x, 2) - k[2] + 2 * pad_[2]) % stride_[2] == 0) || (fallback = true)
-    return pad_, stride_, fallback
-end
-
 softmax!(y::A, x::A) where A<:AbstractVecOrMat{Float32} = nnp_softmax_output(x, y)
 
-function maxpool!(y::A, x::A, k; pad = map(_->0,k), stride = k) where A<:Array{Float32, 4}
-    pad_, stride_, fallback = check_support(x, k, pad, stride)
-    if fallback
-        maxpool_cpu!(y, x, k, pad = pad_, stride = stride_)
-    else
-        nnp_max_pooling_output(x, y, k, padding = expand(Val{length(k)}, pad), stride = expand(Val{length(k)}, stride))
-    end
+nnpack_max_pooling!(y::A, x::A, k; pad = 0, stride = 1) where A<:Array{Float32, 4} =
+    nnp_max_pooling_output(y, x, k, padding = pad, stride = stride)
+
+function nnpack_convolution_forward!(y::A1, x::A1, w::A1, b::A2; pad = 0, stride = 1, algo = UInt32(0),
+                                     flipkernel = 0) where {A1<:Array{Float32, 4}, A2<:Array{Float32, 1}}
+    flipkernel == 0 && (w .= flipweight(w))        
+    # Use nnp_convolution_inference if the batch size is 1.
+    # The wrapper for nnp_convolution_inference is not present so use nnp_convolution_output for now
+    nnp_convolution_output(y, x, w, b, algo = algo, padding = pad, stride = stride)
 end
 
-function conv!(y::A1, x::A1, w::A1; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A1<:Array{Float32, 4}
-    pad_, stride_, fallback = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
+function nnpack_convolution_backward_data!(dx::A, dy::A, x::A, w::A; pad = 0, stride = 1,
+                                           algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
     flipkernel == 0 && (w .= flipweight(w))
-    if fallback
-        conv2d!(y, x, w, padding = pad, stride = stride, dilation = dilation, mode = 1)
-    else
-        nnp_convolution_output(y, x, w, zeros(Float32, size(y, 3)), algo = algo, padding = pad, stride = stride)
-    end
+    nnp_convolution_input_gradient(dx, x, dy, w, padding = pad, stride = stride, algo = algo)
 end
 
-function ∇conv_data!(dx::A, dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
-    pad_, stride_, fallback = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
-    if fallback
-        conv2d_grad_x!(dx, x, w, dy, padding = pad_, stride = stride_, dilation = dilation)
-    else
-        flipkernel == 0 && (w .= flipweight(w))
-        nnp_convolution_input_gradient(dx, x, dy, w, padding = pad, stride = stride, algo = algo)
-    end
+function nnpack_convolution_backward_filter!(dw::A, dy::A, x::A, w::A; pad = 0, stride = 1,
+                                             algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
+    flipkernel == 0 && (w .= flipweight(w))
+    nnp_convolution_kernel_gradient(dw, x, dy, w, padding = pad, stride = stride, algo = algo)
+    flipkernel && (dw .= flipkernel(dw))
+    dw
 end
 
-function ∇conv_filter!(dw::A, dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:Array{Float32, 4}
-    pad_, stride_, fallback = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
-    if fallback
-        conv2d_grad_w!(dw, x, w, dy, padding = pad_, stride = stride_, dilation = dilation)
-    else
-        flipkernel == 0 && (w .= flipweight(w))
-        nnp_convolution_kernel_gradient(dw, x, dy, w, padding = pad, stride = stride, algo = algo)
-        flipkernel && (dw .= flipkernel(dw))
-        dw
-    end
-end

src/nnpack/libnnpack.jl

@@ -81,7 +81,7 @@ function nnp_max_pooling_output(batch_size, channels, input_size, input_padding,
     @check ccall((:nnp_max_pooling_output, libnnpack), nnp_status, (Csize_t, Csize_t, nnp_size, nnp_padding, nnp_size, nnp_size, Ptr{Cfloat}, Ptr{Cfloat}, pthreadpool_t), batch_size, channels, input_size, input_padding, pooling_size, pooling_stride, input, output, threadpool)
 end
 
-function nnp_max_pooling_output(x::Array{Float32,4}, y::Array{Float32,4}, kernel::Tuple; padding = 0, stride = 1, threadpool = shared_threadpool[])
+function nnp_max_pooling_output(y::Array{Float32,4}, x::Array{Float32,4}, kernel::Tuple; padding = 0, stride = 1, threadpool = shared_threadpool[])
     input_size = nnp_size(Csize_t.((size(x, 1), size(x, 2)))...)
     pooling_size = nnp_size(Csize_t.(kernel)...)
     input_padding = nnp_padding(Csize_t(padding[2]), Csize_t(padding[1]), Csize_t(padding[2]), Csize_t(padding[1]))

test/conv.jl

@@ -5,26 +5,17 @@ using NNlib: conv, crosscor, ∇conv_filter, ∇conv_data, ∇maxpool, maxpool,
     x = reshape(Float32[1:16;], 4, 4, 1, 1)
     w = reshape(Float32[1:9;], 3, 3, 1, 1)
 
-    # Fix these tests
-    #=
-    @test dropdims(conv(x, w), dims = (3,4)) == [
-        29 79 129;
-        39 89 139;
-        49 99 149;
-        59 109 159.]
+    @test dropdims(conv(x, w), dims = (3,4)) == Float32.([
+        192 372;
+        237 417])
 
-    @test dropdims(conv(view(x, :, :, :, :), w), dims = (3,4)) == [
-        29 79 129;
-        39 89 139;
-        49 99 149;
-        59 109 159.]
-
-    @test dropdims(crosscor(x, w), dims = (3,4)) == [
-	 51  101  151;
-	 61  111  161;
- 	 71  121  171;
-	 81  131  181.]
-    =#
+    @test dropdims(conv(view(x, :, :, :, :), w), dims = (3,4)) == Float32.([
+        192 372;
+        237 417])
+
+    @test dropdims(crosscor(x, w), dims = (3,4)) == Float32.([
+	 348.0 528.0;
+         393.0 573.0])
 
     @test dropdims(conv(x, w, pad=1), dims=(3,4)) ≈ Float32.([
         29   99  207  263