JuliaDiff / JuliaDiff/ChainRules.jl
Rules that use scalar indexing are not GPU-compatible
Nobody has claimed this yet.
- 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
- 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
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