An optimization problem on GPU works in Julia v1.8, but does not work in v1.9


I’m doing an optimization problem and trying to work on CUDA. I can run my program in Julia v1.8.2 and v1.8.5. However, I get some errors when running on the GPU in Julia v1.9.0 and v1.9.1.

Here is a minimum example:

using CUDA, Distributions, Optim

x = rand(truncated(Normal(0, 1); lower=0.0), 100)

logL(x, p1, p2) = logpdf(truncated(Normal(p1, exp(p2)); lower=0.0), x)

obj(x, p1, p2) = -mean(logL.(CuArray(x), p1, p2))

init = ones(2)

H = TwiceDifferentiable(vars -> obj(x, vars[1], vars[2]), init)

opt = optimize(H, init, Optim.Options(iterations = 10, show_trace=true))

And here are errors when running the last line of the above example in Julia v1.9.0 and v1.9.1:

ERROR: InvalidIRError: compiling MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::CUDA.CuKernelContext, ::CuDeviceVector{Float64, 1}, ::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(logL), Tuple{Base.Broadcast.Extruded{CuDeviceVector{Float64, 1}, Tuple{Bool}, Tuple{Int64}}, Float64, Float64}}, ::Int64) resulted in invalid LLVM IR

Reason: unsupported dynamic function invocation (call to var"#setprecision#25"(kws::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(setprecision), f::Function, ::Type{T}, prec::Integer) where T @ Base.MPFR mpfr.jl:969)

The above errors point out Reason: unsupported dynamic function. I’m not sure whether it is a bug of Julia v1.9.x or not. How can I solve the problem?

Not necessarily a bug, but possibly an upstream change that introduced GPU-incompatible code. You should have included the backtrace to setprecision, but seeing MPFR in there, you’re attempting to use BigFloat which are GPU incompatible (as they rely on a binary C library).

I try to check the type problem via @device_code_warntype, Cthulhu.jl and its macro @descend_code_typed. They show all variables are stable in work.

After that, I guess the problem may arise from the broadcast operator, so I merge logL and obj into a function as follows:

obj(x, p1, p2) = -mean(logpdf.(truncated(Normal(p1, exp(p2)); lower=0.0), CuArray(x)))

In this way, Julia v1.9.0 and v1.9.1 can work.

Although combining functions can avoid the above problem, I wonder whether there exists any way that splitting functions works fine, or if this is a problem because the way of broadcast identification is changed.