added fold/unfold and gpu tests

Project.toml

@@ -13,7 +13,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 Adapt = "3.3"
 CUDA = "3.11"
-NNlib = "0.8.9"
+NNlib = "0.8.11"
 julia = "1.6"
 
 [extras]

src/NNlibCUDA.jl

@@ -12,6 +12,7 @@ include("activations.jl")
 include("batchedadjtrans.jl")
 include("batchedmul.jl")
 include("ctc.jl")
+include("fold.jl")
 include("scatter.jl")
 include("gather.jl")
 include("utils.jl")

src/fold.jl

@@ -0,0 +1,111 @@
+
+function unfold_kernel!(col::AbstractArray{T}, x, col_size, input_size, output_size, kernel_size, flipkernel, stride, pad_lo, dilation, max_idx) where {T}
+    index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+
+    @inbounds if index <= max_idx
+        i, kw, kh, kd, c, b = CartesianIndices(col_size)[index].I # col indices
+        w, h, d = CartesianIndices(output_size)[i].I # x indices
+
+        # project 
+        w, h, d = @. ((w, h, d) - 1)*stride - pad_lo + 1 + ((kw, kh, kd) - 1)*dilation
+
+        if !flipkernel
+            kw, kh, kd = kernel_size .- (kw, kh, kd) .+ 1
+        end
+
+        # check out of bounds
+        if !all(checkindex.(Bool, UnitRange.(1, input_size), (w, h, d)))
+            col[i, kw, kh, kd, c, b] = T(0)
+            return nothing
+        end
+
+        xval::T = x[w, h, d, c, b]
+        col[i, kw, kh, kd, c, b] = xval
+    end
+
+    return nothing
+end
+
+function fold_kernel!(x::AbstractArray{T}, col, col_size, input_size, output_size, kernel_size, flipkernel, stride, pad_lo, dilation, max_idx) where {T}
+    index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+
+    @inbounds if index <= max_idx
+        i, kw, kh, kd, c, b = CartesianIndices(col_size)[index].I # col indices
+        w, h, d = CartesianIndices(output_size)[i].I # x indices
+
+        # project 
+        w, h, d = @. ((w, h, d) - 1)*stride - pad_lo + 1 + ((kw, kh, kd) - 1)*dilation
+
+        # check out of bounds
+        if !all(checkindex.(Bool, UnitRange.(1, input_size), (w, h, d)))
+            return nothing
+        end
+
+        if !flipkernel
+            kw, kh, kd = kernel_size .- (kw, kh, kd) .+ 1
+        end
+
+        cval::T = col[i, kw, kh, kd, c, b]
+        CUDA.@atomic x[w, h, d, c, b] += cval
+    end
+
+    return nothing
+end
+
+function NNlib.unfold!(col::AnyCuArray{cT,3}, x::AnyCuArray{xT,5}, cdims::NNlib.DenseConvDims) where {cT, xT}
+    if NNlib.spatial_dims(cdims) != 3
+        throw(DimensionMismatch("unfold!() only accepts 3d convoluitional inputs"))
+    end
+
+    input_size = NNlib.input_size(cdims)
+    C_in = NNlib.channels_in(cdims)
+    kernel_size = NNlib.kernel_size(cdims)
+    pad_w_lo, pad_w_hi, pad_h_lo, pad_h_hi, pad_d_lo, pad_d_hi = NNlib.padding(cdims)
+    pad_lo = (pad_w_lo, pad_h_lo, pad_d_lo)
+    dilation = NNlib.dilation(cdims)
+    stride = NNlib.stride(cdims)
+    output_size = NNlib.output_size(cdims)
+    flipkernel = NNlib.flipkernel(cdims) 
+
+    col_reshaped = reshape(col, (prod(output_size), kernel_size..., C_in, :))
+
+    max_idx = prod(size(col))
+    args = col_reshaped, x, size(col_reshaped), input_size, output_size, kernel_size, flipkernel, stride, pad_lo, dilation, max_idx
+    kernel = @cuda launch=false unfold_kernel!(args...)
+    config = launch_configuration(kernel.fun; max_threads=256)
+    threads = min(max_idx, config.threads)
+    blocks = cld(max_idx, threads)
+    kernel(args...; threads=threads, blocks=blocks)
+    return col
+end
+
+function NNlib.fold!(x::AnyCuArray{xT,5}, col::AnyCuArray{cT,3}, cdims::NNlib.DenseConvDims) where {xT, cT}
+    if NNlib.spatial_dims(cdims) != 3
+        throw(DimensionMismatch("fold!() only accepts 3d convoluitional inputs"))
+    end
+
+    # going to accumulate into x
+    fill!(x, xT(0))
+
+    input_size = NNlib.input_size(cdims)
+    C_in = NNlib.channels_in(cdims)
+    kernel_size = NNlib.kernel_size(cdims)
+    pad_w_lo, pad_w_hi, pad_h_lo, pad_h_hi, pad_d_lo, pad_d_hi = NNlib.padding(cdims)
+    pad_lo = (pad_w_lo, pad_h_lo, pad_d_lo)
+    dilation = NNlib.dilation(cdims)
+    stride = NNlib.stride(cdims)
+    output_size = NNlib.output_size(cdims)
+    flipkernel = NNlib.flipkernel(cdims) 
+
+    col_reshaped = reshape(col, (prod(output_size), kernel_size..., C_in, :))
+
+    max_idx = prod(size(col))
+    args = x, col_reshaped, size(col_reshaped), input_size, output_size, kernel_size, flipkernel, stride, pad_lo, dilation, max_idx
+    kernel = @cuda launch=false fold_kernel!(args...)
+    config = launch_configuration(kernel.fun; max_threads=256)
+    threads = min(max_idx, config.threads)
+    blocks = cld(max_idx, threads)
+    kernel(args...; threads=threads, blocks=blocks)
+    return x
+end
+

test/fold.jl

@@ -0,0 +1,36 @@
+
+@testset "fold" begin
+    # Test for agreement between CPU/GPU versions, across a variety of kwargs
+    options = Dict{Any, Any}.((
+        (), (:dilation => 2), (:flipkernel => true), (:stride => 2),
+        (:padding => 1),
+        (:padding => (1,0)),
+        (:padding => (0,1)),
+        (:padding => (2,3)),
+    ))
+
+    C_in = 3
+    C_out = 4
+    batch_size = 1
+
+    @testset "spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
+        for opts in options
+            if :padding in keys(opts)
+                padding = opts[:padding]
+                if 1 < length(padding) && length(padding) != 2spatial_rank
+                    opts[:padding] = ntuple(i -> padding[mod1(i,2)] .+ 2div(i-1,2), 2spatial_rank)   
+                end
+            end
+
+            x = rand(Float64, fill(8, spatial_rank)..., C_in, batch_size)
+            w = rand(Float64, fill(2, spatial_rank)..., C_in, C_out)
+            cdims = DenseConvDims(x, w; opts...)
+            y = NNlib.unfold(x, cdims)
+
+            # test equivalence of fold/unfold across GPU/CPU
+            gputest(x -> NNlib.unfold(x, cdims), x) 
+            gputest(y -> NNlib.fold(y, size(x), cdims), y) 
+        end
+    end
+end
+

test/runtests.jl

@@ -15,6 +15,7 @@ include("batchedmul.jl")
 include("upsample.jl")
 include("conv.jl")
 include("ctc.jl")
+include("fold.jl")
 include("pooling.jl")
 include("softmax.jl")
 include("batchnorm.jl")