GPU Support for Integration

[electronsandstuff]

Hey folks, I am running into another problem related to integrals. I can't seem to get GPU support working correctly with them. The script I wrote runs correctly when the equations do not include an integral, but fails at compilation when they do. Here is a small example that has this problem.

using NeuralPDE, Lux, Optimization, OptimizationOptimJL, Random
using ModelingToolkit, Logging, DomainSets, OptimizationOptimisers

function generate_eqs()
    # Setup all parameters and operators for vlasov equation
    @parameters s r vᵣ
    @variables f(..) Eself(..)
    Ds = Differential(s)
    Dr = Differential(r)
    Ii = Integral(vᵣ in DomainSets.ClosedInterval(-1.0, 1.0))

    # Define equation
    eqs = [
        Ds(f(s, r, vᵣ)) + vᵣ*Dr(f(s, r, vᵣ)) ~ 0,
        Eself(s, r) ~ Ii(f(s, r, vᵣ))  # Commenting out this line fixes issue
    ]

    # Boundary conditions
    σr = 0.1
    σvᵣ = 0.3
    bcs = [
        # Initial distribution
        f(0, r, vᵣ) ~ exp(-r^2/2/σr^2)/sqrt(2*π)/σr * exp(-vᵣ^2/2/σvᵣ^2)/sqrt(2*π)/σvᵣ,
    ]

    # Domains
    domains = [
        s ∈ Interval(0.0, 1.0),
        r ∈ Interval(-1.0, 1.0),
        vᵣ ∈ Interval(-1.0, 1.0),
    ]

    # Return as named tuple
    @named pde_system = PDESystem(eqs, bcs, domains, [s, r, vᵣ], [f(s, r, vᵣ), Eself(s, r)])
    pde_system
end

function make_discretization()
    # Neural network
    idim = 32
    chains = [
        Lux.Chain(Dense(3,idim,Lux.σ),Dense(idim,idim,Lux.σ),Dense(idim,idim,Lux.σ),Dense(idim,1)),  # For f
        Lux.Chain(Dense(2,idim,Lux.σ),Dense(idim,idim,Lux.σ),Dense(idim,idim,Lux.σ),Dense(idim,1))  # For Eself
    ]

    # Get initial params
    names = :f, :Eself
    init_params = Lux.initialparameters.(Random.default_rng(), chains)
    init_params = NamedTuple{names}(init_params)
    init_params = Lux.ComponentArray(init_params)
    init_params = Float64.(Lux.gpu(init_params))

    # Discretization
    strategy = QuasiRandomTraining(128)
    PhysicsInformedNN(chains, strategy, init_params=init_params)
end

# Set up training for PINN
pde_system = generate_eqs()
discretization = make_discretization()

# Try to run the optimization (fails to compile GPU kernel)
prob = discretize(pde_system, discretization)
opt = Adam(0.03)
res = Optimization.solve(prob, opt, maxiters=32)

The stack trace is pretty long, but the top part is copied below and points out the issue as occurring in the methods related to numeric integration in the loss function.

ERROR: GPU compilation of kernel #broadcast_kernel#28(CUDA.CuKernelContext, SubArray{Float64, 1, CUDA.CuDeviceVector{Float64, 1}, Tuple{CUDA.CuDeviceVector{Int64, 1}}, false}, Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(identity), Tuple{Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}}}}, Int64) failed
KernelError: passing and using non-bitstype argument

Argument 4 to your kernel function is of type Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(identity), Tuple{Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}}}}, which is not isbits:
  .args is of type Tuple{Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}}} which is not isbits.
    .1 is of type Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}} which is not isbits.
      .x is of type Vector{Float64} which is not isbits.


Stacktrace:
  [1] (::Cubature.var"#17#18"{Bool, Bool, Int64, Float64, Float64, Int64, Int32, Ptr{Nothing}, Cubature.IntegrandData{IntegralsCubature.var"#3#15"{IntegralProblem{false, ComponentArrays.ComponentVector{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{ComponentArrays.Axis{(depvar = ViewAxis(1:4514, Axis(f = ViewAxis(1:2273, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:96, ShapedAxis((32, 3), NamedTuple())), 
bias = ViewAxis(97:128, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(129:1184, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1185:2240, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2241:2273, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))), Eself = ViewAxis(2274:4514, Axis(layer_1 = ViewAxis(1:96, Axis(weight = ViewAxis(1:64, ShapedAxis((32, 2), NamedTuple())), bias = ViewAxis(65:96, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(97:1152, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1153:2208, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2209:2241, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))))),)}}}, NeuralPDE.var"#integrand_#297"{CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, NeuralPDE.var"#12#13", Vector{NeuralPDE.Phi{Chain{NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), Tuple{Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), NTuple{4, NamedTuple{(), Tuple{}}}}}}, Vector{Int64}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#313"), :phi, :derivative, 
:integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0x12ccfe3d, 0x925aace2, 0x95a9ac30, 0x3d21bb37, 0x1b741398)}, typeof(NeuralPDE.numeric_derivative)}, Vector{Float64}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ComponentArrays.ComponentVector{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{ComponentArrays.Axis{(depvar = ViewAxis(1:4514, Axis(f = ViewAxis(1:2273, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:96, ShapedAxis((32, 3), NamedTuple())), bias = ViewAxis(97:128, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(129:1184, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1185:2240, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2241:2273, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))), Eself = ViewAxis(2274:4514, Axis(layer_1 = ViewAxis(1:96, Axis(weight = ViewAxis(1:64, ShapedAxis((32, 2), NamedTuple())), bias = ViewAxis(65:96, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(97:1152, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1153:2208, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2209:2241, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))))),)}}}}}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Int64})()
    @ Cubature C:\Users\chris\.julia\packages\Cubature\5zwuu\src\Cubature.jl:215
  [2] disable_sigint
    @ .\c.jl:473 [inlined]
  [3] cubature(xscalar::Bool, fscalar::Bool, vectorized::Bool, padaptive::Bool, fdim::Int64, f::IntegralsCubature.var"#3#15"{IntegralProblem{false, ComponentArrays.ComponentVector{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{ComponentArrays.Axis{(depvar = ViewAxis(1:4514, Axis(f = ViewAxis(1:2273, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:96, ShapedAxis((32, 3), NamedTuple())), bias = ViewAxis(97:128, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(129:1184, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1185:2240, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2241:2273, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))), Eself = ViewAxis(2274:4514, Axis(layer_1 = ViewAxis(1:96, Axis(weight = ViewAxis(1:64, ShapedAxis((32, 2), NamedTuple())), bias = ViewAxis(65:96, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(97:1152, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1153:2208, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2209:2241, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))))),)}}}, NeuralPDE.var"#integrand_#297"{CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, NeuralPDE.var"#12#13", Vector{NeuralPDE.Phi{Chain{NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), Tuple{Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), NTuple{4, NamedTuple{(), Tuple{}}}}}}, Vector{Int64}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#313"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0x12ccfe3d, 0x925aace2, 0x95a9ac30, 0x3d21bb37, 0x1b741398)}, typeof(NeuralPDE.numeric_derivative)}, Vector{Float64}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ComponentArrays.ComponentVector{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{ComponentArrays.Axis{(depvar = ViewAxis(1:4514, Axis(f = ViewAxis(1:2273, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:96, ShapedAxis((32, 3), NamedTuple())), bias = ViewAxis(97:128, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(129:1184, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1185:2240, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2241:2273, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))), Eself = ViewAxis(2274:4514, Axis(layer_1 = ViewAxis(1:96, Axis(weight = ViewAxis(1:64, ShapedAxis((32, 2), NamedTuple())), bias = ViewAxis(65:96, ShapedAxis((32, 1), NamedTuple())))), layer_2 = ViewAxis(97:1152, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_3 = ViewAxis(1153:2208, Axis(weight = ViewAxis(1:1024, ShapedAxis((32, 32), NamedTuple())), bias = ViewAxis(1025:1056, ShapedAxis((32, 1), NamedTuple())))), layer_4 = ViewAxis(2209:2241, Axis(weight = ViewAxis(1:32, ShapedAxis((1, 32), NamedTuple())), bias = ViewAxis(33:33, ShapedAxis((1, 1), NamedTuple())))))))),)}}}}, xmin_::Vector{Float64}, xmax_::Vector{Float64}, reqRelError::Float64, reqAbsError::Float64, maxEval::Int64, error_norm::Int32)
    @ Cubature C:\Users\chris\.julia\packages\Cubature\5zwuu\src\Cubature.jl:169
  [4] hcubature(f::Function, xmin::Vector{Float64}, xmax::Vector{Float64}; reltol::Float64, abstol::Float64, maxevals::Int64)
...

[ChrisRackauckas]

This is a duplicate of #267. It needs handling in Integrals.jl.