Disable at-nospecialize for GPU codegen
- Dominant language
- Julia
- Stars
- 1.4k
- Forks
- 281
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 30
Description
Now that we have a GPU runtime library that can allocate and box, I tried to get rid of this hack: https://github.com/JuliaGPU/CUDAnative.jl/blob/53368d48b6405ee962e54d4c3b9f90e3eb623310/src/compiler/irgen.jl#L229-L236
Turns out we still need it, as the argument to `throw` often is the value returned by the `BoundsError` constructor, which has a `@nospecialize` resulting in a `jl_invoke`. LLVM obviously can't remove this function by itself, so we end up with GPU-incompatible code. Just try to `cu([1]) .+ cu([2])` with `--check-bounds=yes` (and the above hack disabled, of course).
We could try and redefine `BoundsError `or `throw_boundserror` as it occurs for GPU code, e.g.:
```julia
using LinearAlgebra
for (W, _) in (:AT => (A,mut)->mut(A), Adapt.wrappers...)
@eval begin
# BoundsError has a @nospecialize constructor resulting in `invoke`
Base.throw_boundserror(A::$W where {AT <: CuDeviceArray}, I) = throw(nothing)
end
end
```
But then we still miss cases thrown from the broadcasting code (and I couldn't find a way to dispatch on `throw_boundserror(Broadcasted{...CuDeviceArray...}`. It would be best to just get rid of `@nospecialize` for GPU code altogether.
Not sure whether that would need to happen at the inference/optimizer/codegen level though. Monkey-patching `MethodInstance`s from within CUDAnative's `emit_function` hook didn't seem to work.
@Keno did you do something similar for XLA, since you worked on more exhaustive inference there?
Are those interfaces public already?
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.