`pad_zeros` is not type stable for certain usage patterns

[femtomc]

Using 0.8.5:

module TypeInstability

using InteractiveUtils
using StableRNGs
using NNlib

r = 128
s = 30
rng = StableRNG(3)
x = rand(rng, Float32, r * s * rand(rng, 25:100), 1, 1, 8)
@code_warntype NNlib.pad_zeros(x, (0, size(x)[1] % 2, 0, 0, 0, 0, 0, 0))

end # module

produces

julia> include("scratch/instability.jl")
Variables
  #self#::Core.Const(NNlib.pad_zeros)
  x::Array{Float32, 4}
  pad::NTuple{8, Int64}

Body::Any
1 ─ %1 = NNlib.:(var"#pad_zeros#386")(NNlib.:(:), #self#, x, pad)::Any
└──      return %1
Main.TypeInstability

---

julia> versioninfo()
Julia Version 1.6.5
Commit 9058264a69 (2021-12-19 12:30 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, skylake-avx512)
Environment:
  JULIA_VERSION = 1.6.5

[ToucheSir]

The root of the instability comes from

NNlib.jl/src/padding.jl

Line 115 in 9cd9013

function gen_pad(pad::NTuple{L,Int}, dims::NTuple{D,Int}, N) where {L,D}

. N should be a Val. Bad news is the original PR didn't comment on why this wasn't done, but good news is that this is an internal function and should be straightforward to patch.

[jesseylin]

Just ran into this issue recently. Is there any reason the padding functions cannot be implemented entirely in terms of gather/scatter, which do not seem to have this type instability?

[ToucheSir]

No particular reason. I don't think it makes sense to use gather and scatter for all of the padding functions, but feel free to drop a PR if you can come up with implementations that are faster and/or simpler than the current ones!