JuliaGPU / JuliaGPU/KernelAbstractions.jl
v0.10, CPU backend: `@device_override __validindex(ctx)` on `::Any` makes `__validindex` non-extensible
Nobody has claimed this yet.
- 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)withBase.Experimental.@overlayinsrc/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 runsI in __ndrange(ctx)with__ndrange(ctx) === nothing, and the compiler turns the guaranteediterate(::Nothing)MethodError into thejl_f_throw_methoderrorthat GPUCompiler rejects.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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