Constant Propagation on `Base.getindex`

[stev47]

Due to high compile times for very large static vectors constant propagation was changed in #32105 to bail out on Base.getindex in cases where the object is any non-singleton AbstractArray.

This can prevent one from harnessing constant propagation to provide type-stability in cases where the index is constant:

# constant propagation and getindex

# not inheriting from AbstractArray works
struct A <: AbstractArray{Int,1}
    # making A empty (i.e a singleton) works
    x::Int
end
# this does not help
#Base.issingletontype(::Type{<:A}) = true

Base.size(::A) = (1,)

my_getindex(a::A, i) = i > 1 ? missing : i
Base.getindex(a::A, i) = i > 1 ? missing : i

f(a) = my_getindex(a, 1)
g(a) = Base.getindex(a, 1)

a = A(0)

# this infers correctly to `Int` using constant propagation
@code_warntype f(a)
# this infers to `Union{Missing,Int}`
@code_warntype g(a)

A more concrete usecase is the package StaticKernels where the user is defining a kernel function with constant relative indexing on a moving Window array that efficiently handles array boundaries using the type system.
In order not to sacrifice the user interface the current workaround is to not inherit from AbstractArray even though it would make sense.

Can we provide a method for array types to opt-in to constant propagation for getindex or refine the current heuristic?

julia> versioninfo()
Julia Version 1.5.0-DEV.650
Commit dd738f9ee8* (2020-04-19 16:28 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: AMD A10-7860K Radeon R7, 12 Compute Cores 4C+8G
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, bdver3)
Environment:
  JULIA_PKG_SERVER = pkg.julialang.org

[JeffBezanson]

Base.issingletontype(::Type{<:A}) = true

😱

[JeffBezanson]

Can you give a bit more detail about the types involved and how the constants arise? For instance, in your example the x field is never used so the type actually could be a singleton.

[stev47]

Ah, sorry about that line ...

In StaticKernels the user can create a kernel by specifying a kernel function and apply it similar to a regular function using e.g.

# 1d finite difference kernel
k = Kernel{(0:1,)}(w -> w[1] - w[0])
map(k, rand(10), StaticKernels.ExtensionConstant(0))

So the kernel function is called with a statically sized moving window w that can be indexed using relative coordinates and returns an output value at that coordinate.
Traditionally boundaries (in the above simple example the highest index) need to be special cased or checked for by branching. Instead I use the static information on w to generate branch-free code for arbitrary user-defined kernel functions.
getindex on w is defined like this:

@propagate_inbounds @inline function Base.getindex(w::Window{<:Any,N}, wi::CartesianIndex{N}) where N
    # central piece to get efficient boundary handling.
    # we rely on the compiler to constant propagate this check away
    checkbounds(Bool, w, wi) || return getindex_extension(w, wi, extension(w.kernel))

    return getindex_parent(w, position(w) + wi)
end

If getindex does not constant propagate, then checkbounds will not compile away.

[stev47]

This seems to have been resolved by #43852 which refined the heuristic for constant propagation of getindex, I cannot reproduce the original example anymore.