JuliaGPU / JuliaGPU/KernelAbstractions.jl

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

Open
#757 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

wontfix
Dominant language
Julia
Stars
523
Forks
88
Avg merge
1d 11h
Merged PRs (30d)
25

Description

Reproducer of an issue observed in https://github.com/CliMA/Oceananigans.jl/pull/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.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Run the supplied Julia reproducer with KernelAbstractions 0.10 and inspect the POCL backend overlay at src/pocl/backend.jl:318, alongside the downstream __validindex methods shown in the issue. Confirm the CPU backend no longer prevents the mapped context method from being considered, and that the reproducer's final assertion passes without InvalidIRError.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend-api-design, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
58/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.