JuliaDiff / JuliaDiff/ChainRules.jl

Rules that use scalar indexing are not GPU-compatible

Open
#617 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

GPU
Dominant language
Julia
Stars
475
Forks
98
PR merge metrics
No merged PRs in 30d

Description

Some rules use scalar indexing, breaking GPU compatibility, e.g., https://github.com/JuliaDiff/ChainRules.jl/blob/3b3791f10bc88c41f004fbb9eb229745d1764593/src/rulesets/LinearAlgebra/norm.jl#L187

One solution would be to use @allowscalar from GPUArrays, but one concern about adding that dependency is the loading time (h/t @mcabbott).

Another is to be clever and replace all scalar indexing by size-1 views, such as:

@inbounds @views ∂x[yind:yind] .= sign.(x[yind:yind]) .* Δu

For people unfamiliar with GPU programming and @allowscalar: The reason scalar setindex/getindex is disallowed for GPU arrays is that sequential processing of GPU arrays from loopy CPU code defeats the purpose of using the GPU and leads to terrible performance. However, scalar setindex/getindex is of course legitimate in cases like the above, where you're not in a loop and only intend to set or retrieve that single element in the array. @allowscalar from GPUArrays is how you indicate that an exception should be made in any given instance, e.g.,

@inbounds @allowscalar ∂x[yind] = sign(x[yind]) * Δu

See https://cuda.juliagpu.org/stable/usage/workflow/#UsageWorkflowScalar for more about this mechanism


MWE:

using ChainRules
using CUDA
using LinearAlgebra

CUDA.allowscalar(false)
x = CUDA.ones(3)
nx, pb = ChainRules.rrule(norm, x, Inf)
ChainRules.unthunk(pb(1f0)[2])
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] assertscalar(op::String)
    @ GPUArrays ~/.julia/packages/GPUArrays/Zecv7/src/host/indexing.jl:53
  [3] getindex
    @ ~/.julia/packages/GPUArrays/Zecv7/src/host/indexing.jl:86 [inlined]
  [4] findprev(testf::ChainRules.var"#1770#1771"{Float32}, A::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, start::Int64)
    @ Base ./array.jl:2151
  [5] findlast
    @ ./array.jl:2199 [inlined]
  [6] _normInf_back(x::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, y::Float32, Δy::Float32)
    @ ChainRules ~/.julia/packages/ChainRules/SrdPq/src/rulesets/LinearAlgebra/norm.jl:185
  [7] (::ChainRules.var"#1742#1746"{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float64, Float32})()
    @ ChainRules ~/.julia/packages/ChainRules/SrdPq/src/rulesets/LinearAlgebra/norm.jl:49
  [8] unthunk
    @ ~/.julia/packages/ChainRulesCore/RbX5a/src/tangent_types/thunks.jl:195 [inlined]
  [9] unthunk(x::ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk{ChainRules.var"#1742#1746"{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float64, Float32}}, ChainRules.var"#1741#1745"{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float64, Float32}})
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/RbX5a/src/tangent_types/thunks.jl:222
 [10] top-level scope
    @ REPL[329]:1
 [11] top-level scope

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

Start with src/rulesets/LinearAlgebra/norm.jl around the _normInf_back implementation and run the CUDA MWE with scalar indexing disabled. Search the rulesets for similar scalar indexing and compare the proposed @allowscalar and size-1 view approaches. Done means the affected rules work with GPU arrays without triggering the scalar-indexing error.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.