feat: upsampling functions (#1387)

* test: pixel_shuffle

* feat: upsample linear

* Update test/nn/nnlib.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* feat: upsample bilinear & trilinear

* test: upsampling functions

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

ext/ReactantNNlibExt/Implementations.jl

@@ -474,3 +474,16 @@ function _nnlib_gather_impl(src::AnyTracedRArray, idxs::AbstractArray, n_dims::I
         slice_sizes=Int64[size(src)[1:n_dims]..., ones(Int64, ndims(src) - n_dims)...],
     )
 end
+
+function NNlib.upsample_linear_kernel!(
+    y::AnyTracedRArray{T,N}, x::AnyTracedRArray{T,N}; align_corners::Bool=true
+) where {T,N}
+    wT = real(Reactant.unwrapped_eltype(T))
+    ratios = if align_corners
+        ntuple(i -> wT((size(x, i) - 1) / (size(y, i) - 1)), N - 2)
+    else
+        ntuple(i -> wT(size(x, i) / size(y, i)), N - 2)
+    end
+    copyto!(y, upsample_linear(x, size(y)[1:(end - 2)], ratios..., align_corners))
+    return y
+end

ext/ReactantNNlibExt/Ops.jl

@@ -25,3 +25,105 @@ function reduce_window(
         base_dilations=ones(Int, N),
     )[1]
 end
+
+function upsample_linear(
+    x::AnyTracedRArray{T,3}, out_size::Tuple{Int}, rwidth, align_corners::Bool
+) where {T}
+    W, _, _ = size(x)
+
+    out_idxs = Ops.iota(Int32, [out_size[1]]; iota_dimension=1)
+    iw0, iw1, w0_λ, w1_λ = source_idx_and_λ(rwidth, out_idxs, align_corners, W)
+
+    x0 = x[iw0, :, :]
+    x1 = x[iw1, :, :]
+
+    return w0_λ .* x0 .+ w1_λ .* x1
+end
+
+function upsample_linear(
+    x::AnyTracedRArray{T,4}, out_size::Tuple{Int,Int}, rwidth, rheight, align_corners::Bool
+) where {T}
+    W, H, _, _ = size(x)
+
+    out_width = Ops.iota(Int32, [out_size[1]]; iota_dimension=1)
+    out_height = Ops.iota(Int32, [out_size[2]]; iota_dimension=1)
+
+    iw0, iw1, w0_λ, w1_λ = source_idx_and_λ(rwidth, out_width, align_corners, W)
+    ih0, ih1, h0_λ, h1_λ = source_idx_and_λ(rheight, out_height, align_corners, H)
+
+    w0_λ, w1_λ = reshape(w0_λ, (:, 1, 1, 1)), reshape(w1_λ, (:, 1, 1, 1))
+    h0_λ, h1_λ = reshape(h0_λ, (1, :, 1, 1)), reshape(h1_λ, (1, :, 1, 1))
+
+    x00 = x[iw0, ih0, :, :]
+    x10 = x[iw1, ih0, :, :]
+    x01 = x[iw0, ih1, :, :]
+    x11 = x[iw1, ih1, :, :]
+
+    return h0_λ .* (w0_λ .* x00 .+ w1_λ .* x10) .+ h1_λ .* (w0_λ .* x01 .+ w1_λ .* x11)
+end
+
+function upsample_linear(
+    x::AnyTracedRArray{T,5},
+    out_size::Tuple{Int,Int,Int},
+    rwidth,
+    rheight,
+    rdepth,
+    align_corners::Bool,
+) where {T}
+    W, H, D, _, _ = size(x)
+
+    out_width = Ops.iota(Int32, [out_size[1]]; iota_dimension=1)
+    out_height = Ops.iota(Int32, [out_size[2]]; iota_dimension=1)
+    out_depth = Ops.iota(Int32, [out_size[3]]; iota_dimension=1)
+
+    iw0, iw1, w0_λ, w1_λ = source_idx_and_λ(rwidth, out_width, align_corners, W)
+    ih0, ih1, h0_λ, h1_λ = source_idx_and_λ(rheight, out_height, align_corners, H)
+    id0, id1, d0_λ, d1_λ = source_idx_and_λ(rdepth, out_depth, align_corners, D)
+
+    w0_λ = reshape(w0_λ, (:, 1, 1, 1))
+    w1_λ = reshape(w1_λ, (:, 1, 1, 1))
+    h0_λ = reshape(h0_λ, (1, :, 1, 1))
+    h1_λ = reshape(h1_λ, (1, :, 1, 1))
+    d0_λ = reshape(d0_λ, (1, 1, :, 1))
+    d1_λ = reshape(d1_λ, (1, 1, :, 1))
+
+    x000 = x[iw0, ih0, id0, :, :]
+    x100 = x[iw1, ih0, id0, :, :]
+    x010 = x[iw0, ih1, id0, :, :]
+    x110 = x[iw1, ih1, id0, :, :]
+
+    x001 = x[iw0, ih0, id1, :, :]
+    x101 = x[iw1, ih0, id1, :, :]
+    x011 = x[iw0, ih1, id1, :, :]
+    x111 = x[iw1, ih1, id1, :, :]
+
+    return (
+        (
+            d0_λ .* (
+                h0_λ .* (w0_λ .* x000 .+ w1_λ .* x100) .+
+                h1_λ .* (w0_λ .* x010 .+ w1_λ .* x110)
+            )
+        ) .+ (
+            d1_λ .* (
+                h0_λ .* (w0_λ .* x001 .+ w1_λ .* x101) .+
+                h1_λ .* (w0_λ .* x011 .+ w1_λ .* x111)
+            )
+        )
+    )
+end
+
+@inline function source_idx_and_λ(
+    ratio::T, out_idx::AbstractVector, align::Bool, in_width::Int
+) where {T}
+    real_index = ifelse(
+        align, ratio .* out_idx, max.(zero(T), ratio .* (out_idx .+ T(0.5)) .- T(0.5))
+    )
+
+    iw0 = Base.Fix1(floor, Int).(real_index)
+    offset = ifelse.(iw0 .< in_width - 1, 1, 0)
+    iw1 = iw0 .+ offset .+ 1
+
+    w1lambda = real_index .- iw0
+    w0lambda = one(T) .- w1lambda
+    return iw0 .+ 1, iw1, w0lambda, w1lambda
+end

test/nn/nnlib.jl

@@ -416,7 +416,53 @@ end
     x = randn(Float32, 4, 4, 3, 2)
     x_ra = Reactant.to_rarray(x)
 
-    @test @jit(NNlib.upsample_nearest(x_ra, (2, 2))) ≈ NNlib.upsample_nearest(x, (2, 2))
+    @testset "Nearest" begin
+        @test @jit(NNlib.upsample_nearest(x_ra, (2, 2))) ≈ NNlib.upsample_nearest(x, (2, 2))
+    end
+
+    @testset "Linear" begin
+        x = randn(Float32, 4, 3, 2)
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(NNlib.upsample_linear(x_ra, (2,))) ≈ NNlib.upsample_linear(x, (2,))
+
+        @test @jit(NNlib.upsample_linear(x_ra, (2,); align_corners=false)) ≈
+            NNlib.upsample_linear(x, (2,); align_corners=false)
+    end
+
+    @testset "Bi-Linear" begin
+        x = randn(Float32, 4, 4, 3, 2)
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(NNlib.upsample_bilinear(x_ra, (2, 2))) ≈
+            NNlib.upsample_bilinear(x, (2, 2))
+
+        @test @jit(NNlib.upsample_bilinear(x_ra, (2, 2); align_corners=false)) ≈
+            NNlib.upsample_bilinear(x, (2, 2); align_corners=false)
+    end
+
+    @testset "Tri-Linear" begin
+        x = randn(Float32, 4, 4, 4, 3, 2)
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(NNlib.upsample_trilinear(x_ra, (2, 2, 2))) ≈
+            NNlib.upsample_trilinear(x, (2, 2, 2))
+
+        @test @jit(NNlib.upsample_trilinear(x_ra, (2, 2, 2); align_corners=false)) ≈
+            NNlib.upsample_trilinear(x, (2, 2, 2); align_corners=false)
+    end
+end
+
+@testset "Pixel shuffle" begin
+    x = [10i + j + channel / 10 for i in 1:2, j in 1:3, channel in 1:4, batch in 1:1]
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(NNlib.pixel_shuffle(x_ra, 2)) ≈ NNlib.pixel_shuffle(x, 2)
+
+    y = [i + channel / 10 for i in 1:3, channel in 1:6, batch in 1:1]
+    y_ra = Reactant.to_rarray(y)
+
+    @test @jit(NNlib.pixel_shuffle(y_ra, 2)) ≈ NNlib.pixel_shuffle(y, 2)
 end
 
 @testset "softmax/logsoftmax reshaped input" begin