Disable at-nospecialize for GPU codegen

[maleadt]

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.:

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 MethodInstances 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?

[timholy]

Can you intercept the CodeInfo object and rewrite the invokes? In the case of successful type inference you can do method lookup yourself since you have the types of the arguments to the invoke.

JuliaDebug/ASTInterpreter2.jl#37 pulls tricks like that for the purpose of optimizing a runtime interpreter. Not sure if that's a useful model, however.

[maleadt]

Can you intercept the CodeInfo object and rewrite the invokes?

Interesting, I was instead waiting for something like the TPU hacks from kf/tpu3 to enable custom inference params a la https://github.com/JuliaGPU/CUDAnative.jl/compare/vc/crazy

So no, we don't do anything like that right now, but go straight from the method instance to its LLVM IR: https://github.com/JuliaGPU/CUDAnative.jl/blob/9b524d5e1da8e59715af5e37f303cdffc450f17b/src/compiler/irgen.jl#L124-L135

Is it possible to take and rewrite CodeInfos recursively without eg. breaking 265?

[vchuravy]

We need the inferred codeinfo objects to be able to rewrite the invokes, right?.
I think it might be feasible, but replacing the codeinfo object like linfo.inferred = ci seems hacky (even though Jameson said it might be alright...).

https://github.com/vchuravy/Cthulhu.jl/blob/8e59d4d4d4278024adcd1348ac1baa5d8554e79c/src/Cthulhu.jl#L167-L180 is another way to directly interact with inference.

[timholy]

Is it possible to take and rewrite CodeInfos recursively without eg. breaking 265?

If you can't/shouldn't do it "in place," you can always copy the MethodInstance and then tweak whatever's in its inferred field. You might need to do shenanigans like this: https://github.com/JuliaDebug/ASTInterpreter2.jl/blob/624483df083293408d674ff4f3af3f0becdcc9e7/src/ASTInterpreter2.jl#L558-L567