v0.10, CPU backend: `@device_override __validindex(ctx)` on `::Any` makes `__validindex` non-extensible

[giordano]

Reproducer of an issue observed in CliMA/Oceananigans.jl#5799

# Downstream extension of KernelAbstractions' iteration space, as done by Oceananigans.jl
# ("mapped kernels": the work items of a 1D kernel are looked up in a list of indices).
#
# KA 0.9: works on CPU().  KA 0.10 (POCL CPU backend): InvalidIRError, because the
# backend's `@device_override __validindex(ctx)` shadows the downstream `__validindex`
# method and falls back to `I in __ndrange(ctx)` with `__ndrange(ctx) === nothing`.

using KernelAbstractions
using KernelAbstractions: Kernel, CompilerMetadata, NDRange, StaticSize, DynamicSize, NDIteration
using KernelAbstractions: workitems, blocks, __iterspace, __groupindex, __dynamic_checkbounds
import KernelAbstractions: partition, expand, __validindex, workgroupsize
import Adapt

# A kernel function wrapped together with the list of indices it should visit.
struct MappedFunction{F, M} <: Function
    func      :: F
    index_map :: M
end

const MappedKernel{D} = Kernel{D, <:Any, <:Any, <:MappedFunction} where D

# `kernel.f` returns the wrapped function, so the backend compiles/launches that one.
@inline Base.getproperty(k::MappedKernel, prop::Symbol) = _prop(k, Val(prop))
@inline _prop(k, ::Val{prop}) where prop = getfield(k, prop)
@inline _prop(k, ::Val{:index_map}) = getfield(getfield(k, :f), :index_map)
@inline _prop(k, ::Val{:f})         = getfield(getfield(k, :f), :func)

mapped_kernel(k::Kernel{D, B, W}, map) where {D, B, W} =
    Kernel{D, B, W, typeof(MappedFunction(k.f, map))}(k.backend, MappedFunction(k.f, map))

# The index map takes the place of NDRange.workitems (a static workgroup size means
# `workitems(ndrange)` never reads that field).
struct IndexMap{T, N, M <: AbstractArray{T, N}} <: AbstractArray{T, N}
    index_map :: M
end
Base.size(m::IndexMap) = size(m.index_map)
Base.@propagate_inbounds Base.getindex(m::IndexMap, I...) = m.index_map[I...]
Adapt.adapt_structure(to, m::IndexMap) = IndexMap(Adapt.adapt(to, m.index_map))

const MappedNDRange{N, B, W} = NDRange{N, B, W, <:Any, <:IndexMap} where {N, B, W <: StaticSize}
const MappedCompilerMetadata{N, C} = CompilerMetadata{N, C, <:Any, <:Any, <:MappedNDRange} where {N, C}

Adapt.adapt_structure(to, nd::MappedNDRange{N, B, W}) where {N, B, W} =
    NDRange{N, B, W}(Adapt.adapt(to, nd.blocks), Adapt.adapt(to, nd.workitems))
Adapt.adapt_structure(to, cm::MappedCompilerMetadata{N, C}) where {N, C} =
    CompilerMetadata{N, C}(Adapt.adapt(to, cm.groupindex), Adapt.adapt(to, cm.ndrange), Adapt.adapt(to, cm.iterspace))

# The kernel is launched without an `ndrange`; the range is the length of the index map.
function partition(kernel::MappedKernel, inrange, ingroupsize)
    static_wg = workgroupsize(kernel)
    index_map = kernel.index_map
    blocks, _, _ = NDIteration.partition(length(index_map), NDIteration.get(static_wg))
    iterspace = NDRange{1, DynamicSize, static_wg}(CartesianIndices(blocks), IndexMap(index_map))
    return iterspace, NDIteration.DynamicCheck()  # last block must be bounds-checked
end

Base.@propagate_inbounds linear_expand(nd::MappedNDRange, gidx::Integer, idx::Integer) =
    (gidx - 1) * size(workitems(nd), 1) + idx

Base.@propagate_inbounds expand(nd::MappedNDRange, gidx::CartesianIndex{1}, idx::CartesianIndex{1}) =
    CartesianIndex(nd.workitems[linear_expand(nd, gidx.I[1], idx.I[1])])
Base.@propagate_inbounds expand(nd::MappedNDRange, gidx::Integer, idx::Integer) =
    CartesianIndex(nd.workitems[linear_expand(nd, gidx, idx)])

# Validity = the linear index does not run past the end of the index map.
# 2-arg form: called by the KA 0.9 CPU backend.
@inline function __validindex(ctx::MappedCompilerMetadata, idx::CartesianIndex)
    __dynamic_checkbounds(ctx) || return true
    return @inbounds linear_expand(__iterspace(ctx), __groupindex(ctx).I[1], idx.I[1]) ≤ length(__iterspace(ctx).workitems)
end
# 1-arg form: what the `@kernel` macro calls on GPU-style backends (including KA 0.10 CPU).
# NOT reached on KA 0.10 CPU: the backend's overlay `__validindex(ctx)` takes precedence.
@inline function __validindex(ctx::MappedCompilerMetadata)
    __dynamic_checkbounds(ctx) || return true
    return @inbounds linear_expand(__iterspace(ctx), KernelAbstractions.__index_Group_Linear(ctx), KernelAbstractions.__index_Local_Linear(ctx)) ≤ length(__iterspace(ctx).workitems)
end

@kernel function scatter!(a, v)
    I = @index(Global, Cartesian)
    @inbounds a[I] = v
end

a   = zeros(10)
map = [1, 3, 5, 6, 10]             # 5 items, workgroup of 4 → last group partially valid
k   = mapped_kernel(scatter!(CPU(), 4), map)
k(a, 1.0)
KernelAbstractions.synchronize(CPU())
println("KA ", pkgversion(KernelAbstractions), ": a = ", a)
@assert a == [1, 0, 1, 0, 1, 1, 0, 0, 0, 1]

Analysis by the bot:

Root cause, for the KA developers. In 0.9 the CPU backend's __validindex(ctx, idx) was an ordinary generic function, so a downstream package could add a more specific method for its own context type and dispatch picked it. In 0.10 the POCL backend defines __validindex(ctx) with Base.Experimental.@overlay in src/pocl/backend.jl:318. An overlay method wins over every method in the global table for any matching signature, regardless of specificity, and its signature is ::Any. The downstream method is therefore never considered, KA's fallback runs I in __ndrange(ctx) with __ndrange(ctx) === nothing, and the compiler turns the guaranteed iterate(::Nothing) MethodError into the jl_f_throw_methoderror that GPUCompiler rejects.

[vchuravy]

Yeah the bot is also a bit wrong, since they should now overload __validindex(ctx) on all platforms, but yeah this is messing with internals a bit too much.