Inference of reductions with a non-concrete `alg` recurses through GPUArrays 12 and never finishes

オープン
#161 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
5/5
見積もり時間
1週間以上
初心者へのやさしさ
8/100
issue の種類
バグ
明瞭さ
おおむね明確
活発さ
活発
技術スタック
julia

調査の方向性

The cycle runs through _mapreduce_setup / _Plan and the host fallback at src/reduce/mapreduce_1d_gpu.jl:116, so start by reproducing the code_typed call with a non-concrete alg global. Done means the abstract-alg case infers in a bounded time and the constant-alg path is unchanged. The fix needs a design decision about function barriers across several entry points, so get maintainer direction before writing code.

索引モデルが issue の本文から書いたものです。

説明

With GPUArrays 12 (the KernelAbstractions 0.10 stack), inferring a reduction whose alg is not inferred concretely never finishes in practice. Base's reductions of GPU arrays now go through AK.mapreduce, and AK's reduction path can call back into them, so inference recurses through AK and GPUArrays with a type that grows one level per round.

using CUDACore
import AcceleratedKernels as AK

global ALG::AK.Algorithm = AK.Auto()      # typed or untyped global: any non-concrete `alg`
redmin(s) = AK.reduce((x, y) -> x < y ? x : y, s; alg=ALG,
                      init=typemax(eltype(s)), neutral=typemax(eltype(s)))

code_typed(redmin, (CuArray{Int32, 1, CUDACore.DeviceMemory},))
info: inference of kwcall(NamedTuple{(:alg, :init, :neutral), ...<:Tuple{AcceleratedKernels.Algorithm, Int32, Int32}}, typeof(AcceleratedKernels.reduce), ...) exceeding 2500 frames (may be slow).
info: ... exceeding 5000 frames (may be slow).
info: ... exceeding 10000 frames (may be slow).

Timings (Julia 1.12.7, AK main at 1b807a3):

stack alg=ALG::AK.Algorithm alg=AK.Auto()
CUDA (RTX 5080): KA 0.10.0, GPUArrays 12.1.0, CUDACore main > 7 min (killed) ~4 s
Metal (M1): KA 0.9.45, GPUArrays 11.5.16, Metal 1.11.1 2.7 s
Metal (M1): KA 0.10.0, GPUArrays 12.1.0, Metal main > 420 s

The same thing makes test/generic/reduce.jl hang for over an hour in its first testset on the KA 0.10 stack, because it passes alg=REDUCE_ALG, a non-const global. CI passes today only because it still resolves KA 0.9 / GPUArrays 11.

Cause

I logged every inference frame with a custom AbstractInterpreter. The constant-alg call takes about 10k frames. The abstract one passes 65k frames in 4 minutes, and the depth of its signatures keeps growing (past 450). The cycle:

  1. _resolve_reduce(alg::Algorithm, ...) infers fine, as Union{BlockReduce, CPUThreads.Partitioned}. But _mapreduce_setup builds _Plan(backend, a, sizes) from it, and merging the two _Plan types widens it to _Plan{CUDABackend, A, @NamedTuple{}, S} where {A, S}. So in _mapreduce_run, s.plan.alg and the buffers from _buffers(s.plan, ...) are not concrete, and partials reaches mapreduce_1d_gpu as Any.
  2. In mapreduce_1d_gpu, Base.reduce(op_host, Vector(p); init=neutral) (src/reduce/mapreduce_1d_gpu.jl:116) then infers its argument only as AbstractArray. That call matches GPUArrays' Base.mapreduce(f, op, ::AnyGPUArray), which on GPUArrays 12 is GPUArrays._mapreduce → AK.mapreduce.
  3. That nested AK.mapreduce has op = _AccOp{Int32, <user op>}. _lanefuncs wraps it again (_AccOp{Int32, _AccOp{...}}), and the chain goes back to step 2 with a deeper type each time. Julia's recursion limiting doesn't stop it, because each round goes through several different methods and only partly widened signatures.

When inference finally stops, finish! for the cycle spends the rest of the time comparing these deep types in ijl_compress_ir and store_backedges (egal_types recursion more than 200 levels deep in every sample).

With GPUArrays 11, step 2 lands in GPUArrays' own reduction kernels instead of AK, so there is no cycle. That is why the KA version seems to matter.

Things I checked:

  • Typing the global as AK.Auto, Union{AK.Auto, AK.BlockReduce} or Union{AK.Auto, AK.CPUThreads.Partitioned} infers in about 0.5 s. Julia union-splits the keyword call at the call site, so each path is concrete. Only a fully abstract alg hits the problem.
  • Resolving the algorithm before the call, so that AK sees Union{BlockReduce, CPUThreads.Partitioned}, also infers in 0.2 s:
    a = AK._resolve_reduce(ALG, get_backend(s), eltype(s), :)
    AK.reduce(op, s; alg=a, init=..., neutral=...)
    

Suggested fix

  • Add a function barrier after algorithm resolution. mapreduce / mapreducedim! (and the other entry points that use the same _resolve_* → _Plan pattern: accumulate, sort, findall, predicates) would resolve alg first and pass the resolved algorithm into an inner function. That way the call is union-split and _Plan stays concrete per algorithm.
  • Separately, make the host fallbacks in the reduction code unable to dispatch back into GPU reductions. For example, Base.reduce(op_host, Vector(p)::Vector; ...) (or an explicit Array{eltype(p)}(p)), and similarly for the Base.mapreduce(f_host, op_host, h_src; ...) path after _host_copy. Now that GPUArrays sends Base reductions to AK, any AK host fallback whose argument isn't inferred as a host array can loop back into AK.

Versions: Julia 1.12.7, AcceleratedKernels main (1b807a3), KernelAbstractions 0.10.0, GPUArrays 12.1.0, CUDACore 6.4.2 (CUDA.jl main) / Metal.jl main.

主要言語
Julia
スター
204
フォーク
20
平均マージ
2日 19時間
マージ済み PR(30日)
41

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

JuliaGPU/AcceleratedKernels.jl のほかの issue

JuliaGPU/AcceleratedKernels.jl の issue をすべて見る

似ている issue

Julia の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。