Added better explantaion on merged implemenntation

Author: gabrielpreviato

src/conv.jl

@@ -166,13 +166,16 @@ end
 
 # First, we will define mappings from the generic API names to our accelerated backend
 # implementations. For homogeneous-datatype 1, 2 and 3d convolutions, we default to using
-# im2col + GEMM.  Do so in a loop, here:
+# im2col + GEMM.
+# But we always support a fallback, non-accelerated path, where we use the direct, but
+# slow, implementations. These should not typically be used, hence the `@warn`,
 
 # These are the GEMM types we will accelerate with `im2col`
 const G = Union{[x[2] for x in gemm_datatype_mappings]...}
 
 for (front_name, backend, signature) in (
-    # This maps from public, front-facing name, to internal backend name
+    # This maps from public, front-facing name, to internal backend name, given the function signature and the where clause
+    # (frontend, backend, (out Array signature, in1 Array signature, in2 Array signature, (parametric Types)))
     (:conv, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
     (:conv, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
 )
@@ -213,7 +216,8 @@ end
 
 # im2col-accelerated function forwarding definition
 for (front_name, backend, signature) in (
-    # This maps from public, front-facing name, to internal backend name
+    # This maps from public, front-facing name, to internal backend name, given the function signature and the where clause
+    # (frontend, backend, (out Array signature, in1 Array signature, in2 Array signature, (parametric Types)))
     (:∇conv_data, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
     (:∇conv_data, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
 )
@@ -256,13 +260,13 @@ for (front_name, backend, signature) in (
 end
 
 for (front_name, backend, signature) in (
-    # This maps from public, front-facing name, to internal backend name
+    # This maps from public, front-facing name, to internal backend name, given the function signature and the where clause
+    # (frontend, backend, (out Array signature, in1 Array signature, in2 Array signature, (parametric Types)))
     (:∇conv_filter, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
     (:∇conv_filter, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
 )
     # We only define 3d conv primitives, we reshape lower down to get 1d and 2d convolution
     @eval begin
-        # println($(Symbol(["$(i)" for i in "$(signature[5])"]...))...)
         function $(Symbol("$(front_name)!"))(
                         out::AbstractArray{$(signature[1][1]), $(signature[1][2])},
                         in1::AbstractArray{$(signature[2][1]), $(signature[1][2])},
@@ -298,98 +302,37 @@ for (front_name, backend, signature) in (
 end
 
 
-for (front_name, backend) in (
-        # This maps from public, front-facing name, to internal backend name
-        :depthwiseconv          => :im2col,
-        :∇depthwiseconv_data    => :im2col,
-        :∇depthwiseconv_filter  => :im2col,
-    )
+for (front_name, backend, signature) in (
+    # This maps from public, front-facing name, to internal backend name, given the function signature and the where clause
+    # (frontend, backend, (out Array signature, in1 Array signature, in2 Array signature, (parametric Types)))
+    (:depthwiseconv, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
+    (:depthwiseconv, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
+    
+    (:∇depthwiseconv_data, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
+    (:∇depthwiseconv_data, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
+    
+    (:∇depthwiseconv_filter, :im2col, ((:T, 5), (:T, 5), (:T, 5), :C, (:(T <: G), :(C <: ConvDims)))),
+    (:∇depthwiseconv_filter, :direct, ((:yT, :N), (:T1, :N), (:T2, :N), :C, (:yT, :T1, :T2, :N, :(C <: ConvDims)))),
+)
 
     # We only define 3d conv primitives, we reshape lower down to get 1d and 2d convolution
     @eval begin
         # im2col-accelerated function forwarding definition
         function $(Symbol("$(front_name)!"))(
-                        out::AbstractArray{T,5}, in1::AbstractArray{T,5},
-                        in2::AbstractArray{T,5}, cdims::C; kwargs...) where {T <: $G, C <: ConvDims}
-            $(Symbol("$(front_name)_$(backend)!"))(out, in1, in2, cdims; kwargs...)
-        end
-    end
-end
-
-# We always support a fallback, non-accelerated path, where we use the direct, but
-# slow, implementations.  These should not typically be used, hence the `@warn`,
-# but let's go ahead and define them first:
-for front_name in (:depthwiseconv, :∇depthwiseconv_data, :∇depthwiseconv_filter)
-    @eval begin
-        function $(Symbol("$(front_name)!"))(
-                        y::AbstractArray{yT,N}, in1::AbstractArray{T1,N},
-                        in2::AbstractArray{T2,N}, cdims::ConvDims;
-                        kwargs...) where {yT, T1, T2, N}
-            if yT == Float64  # warn for Float32 + accidental Float64, but don't print warning for ForwardDiff.Dual
+                        out::AbstractArray{$(signature[1][1]), $(signature[1][2])},
+                        in1::AbstractArray{$(signature[2][1]), $(signature[1][2])},
+                        in2::AbstractArray{$(signature[3][1]), $(signature[1][2])},
+                        cdims::$(signature[4]),
+                        kwargs...) where {$(signature[5]...)}
+            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
+                        "You probably don't want this; check your datatypes.") yT T1 T2 maxlog=1
             end
-            $(Symbol("$(front_name)_direct!"))(y, in1, in2, cdims; kwargs...)
+            $(Symbol("$(front_name)_$(backend)!"))(out, in1, in2, cdims; kwargs...)
         end
     end
 end
 
-# # direct function forwarding definition
-# function ∇conv_data!(out::AbstractArray{yT,N}, in1::AbstractArray{T1,N},
-#                      in2::AbstractArray{T2,N}, cdims::C; kwargs...) where {yT, T1, T2, N, C <: ConvDims}
-#     if 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
-#     end
-
-#     dx_cs = Iterators.partition(1:size(out, 4),
-#                 channels_in(cdims) ÷ groupcount(cdims))
-#     w_cs = Iterators.partition(1:size(in2, 5),
-#                 channels_out(cdims) ÷ groupcount(cdims))
-#     dy_cs = Iterators.partition(1:size(in1, 4),
-#                 channels_out(cdims) ÷ groupcount(cdims))
-#     cdims2 = basetype(C)(cdims,
-#             G = 1,
-#             C_in = channels_in(cdims) ÷ groupcount(cdims),
-#             C_out = channels_out(cdims) ÷ groupcount(cdims))
-
-#     Threads.@sync for (xc, yc, wc) in zip(dx_cs, dy_cs, w_cs)
-#         dxv = @view out[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
-#         dyv = @view in1[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-#         wv = @view in2[ntuple(i -> i == 5  ? wc : Colon(), 5)...]
-#         Threads.@spawn ∇conv_data_direct!(dxv, dyv, wv, cdims2; kwargs...)
-#     end
-
-#     return out
-# end
-
-# function ∇conv_filter!(out::AbstractArray{yT,N}, in1::AbstractArray{T1,N},
-#                        in2::AbstractArray{T2,N}, cdims::C; kwargs...) where {yT, T1, T2, N, C <: ConvDims}
-#     if 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
-#     end
-#     dw_cs = Iterators.partition(1:size(out, 5),
-#                 channels_out(cdims) ÷ groupcount(cdims))
-#     dy_cs = Iterators.partition(1:size(in2, 4),
-#                 channels_out(cdims) ÷ groupcount(cdims))
-#     x_cs = Iterators.partition(1:size(in1, 4),
-#                 channels_in(cdims) ÷ groupcount(cdims))
-#     cdims2 = basetype(C)(cdims,
-#             G = 1,
-#             C_in = channels_in(cdims) ÷ groupcount(cdims),
-#             C_out = channels_out(cdims) ÷ groupcount(cdims))
-
-#     Threads.@sync for (wc, xc, yc) in zip(dw_cs, x_cs, dy_cs)
-#         x = @view in1[ntuple(i -> i == 4 ? xc : Colon(), 5)...]
-#         dy = @view in2[ntuple(i -> i == 4 ? yc : Colon(), 5)...]
-#         dw = @view out[ntuple(i -> i == 5 ? yc : Colon(), 5)...]
-#         Threads.@spawn ∇conv_filter_direct!(dw, x, dy, cdims2; kwargs...)
-#     end
-
-#     return out
-# end
-
 for Dims in [:DenseConvDims, :DepthwiseConvDims, :PoolDims]
     @eval @non_differentiable $Dims(::Any...)
 end