Error: GPU compilation of kernel

[khorrami1]

There is an error related to OperetorKernel. When I use the simple network, model2, there is no error in the calculation of the gradient.

Julia Env:
[587475ba] Flux v0.13.13
[ea5c82af] NeuralOperators v0.4.5

CUDA.versioninfo()
CUDA toolkit 11.7, artifact installation
NVIDIA driver 531.18.0, for CUDA 12.1
CUDA driver 12.1

Libraries:

• CUBLAS: 11.10.1
• CURAND: 10.2.10
• CUFFT: 10.7.1
• CUSOLVER: 11.3.5
• CUSPARSE: 11.7.3
• CUPTI: 17.0.0
• NVML: 12.0.0+531.18
• CUDNN: 8.30.2 (for CUDA 11.5.0)
• CUTENSOR: 1.4.0 (for CUDA 11.5.0)

Toolchain:

• Julia: 1.8.5
• LLVM: 13.0.1
• PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0, 7.1, 7.2
• Device capability support: sm_35, sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86

1 device:
0: NVIDIA GeForce RTX 2070 (sm_75, 5.844 GiB / 8.000 GiB available)

using Flux, NeuralOperators

device = Flux.gpu

ch = [10,10,10,10,10]
modes = (10,10)
σ1=gelu
σ2=sigmoid

model = Chain(
        Dense(3, ch[1]),
        OperatorKernel(ch[1]=>ch[2], modes, FourierTransform, σ1, permuted=false),
        OperatorKernel(ch[2]=>ch[3], modes, FourierTransform, σ1, permuted=false),
        OperatorKernel(ch[3]=>ch[4], modes, FourierTransform, σ1, permuted=false),
        OperatorKernel(ch[4]=>ch[5], modes, FourierTransform, σ1, permuted=false),
        Dense(ch[5], 1, σ2; bias=false),
)|>device


model2 = Chain(
        Dense(3, ch[1]),
        Dense(ch[1], 1, σ2; bias=false),
)|>device

X = rand(Float32, 3, 64,64,2)|>device
Y = model(X)

loss2(x, y) = Flux.Losses.mae(model2(x), y)

loss(x,y) = Flux.Losses.mae(model(x), y)
loss(X,Y)

ps = Flux.params(model)

gs = Flux.gradient(()->loss2(X,Y), ps)

gs = Flux.gradient(()->loss(X,Y), ps)

[khorrami1]

I think there is a conflict between packages. There were no errors until I updated all the packages in my environment:

(@v1.8) pkg> st
Status ~/.julia/environments/v1.8/Project.toml
[fbb218c0] BSON v0.3.6
[4c6ed407] BayesianOptimization v0.2.5
⌅ [052768ef] CUDA v3.13.1
[31a5f54b] Debugger v0.7.8
[79314e9b] FESolvers v0.2.0
[7a1cc6ca] FFTW v1.6.0
[c061ca5d] Ferrite v0.3.12
[fd21fc07] FerriteAssembly v0.2.0
[4f95f4f8] FerriteGmsh v1.0.0
[f0450e97] FerriteNeumann v0.1.1
[e7aa30d1] FerriteProblems v0.1.4
[59d0093e] FerriteViz v0.2.0
⌃ [587475ba] Flux v0.13.13
[e9467ef8] GLMakie v0.8.2
[705231aa] Gmsh v0.2.2
⌃ [3025c34a] GridapGmsh v0.5.0
[f67ccb44] HDF5 v0.16.14
[7073ff75] IJulia v1.24.0
⌃ [b2108857] Lux v0.4.37
[f1d291b0] MLUtils v0.4.1
[d7e62f90] MaterialModels v0.1.0 https://github.com/kimauth/MaterialModels.jl#main
[af893363] MaterialModelsBase v0.1.3
⌃ [961ee093] ModelingToolkit v8.46.1
[872c559c] NNlib v0.8.19
[15e1cf62] NPZ v0.4.3
[ea5c82af] NeuralOperators v0.4.5
[315f7962] NeuralPDE v5.4.0
[3bd65402] Optimisers v0.2.15
[7f7a1694] Optimization v3.12.1
[36348300] OptimizationOptimJL v0.1.5
[91a5bcdd] Plots v1.38.8
[6038ab10] Rotations v1.4.0
[48a634ad] Tensors v1.13.1
[276b4fcb] WGLMakie v0.8.6
[29a6e085] Wavelets v0.9.5
[e88e6eb3] Zygote v0.6.57
[10745b16] Statistics

[khorrami1]

I got this error message:
GPU compilation of kernel #broadcast_kernel#28(CUDA.CuKernelContext, CUDA.CuDeviceArray{ComplexF32, 4, 1}, Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{4}, NTuple{4, Base.OneTo{Int64}}, typeof(*), Tuple{Base.Broadcast.Extruded{Array{Float32, 4}, NTuple{4, Bool}, NTuple{4, Int64}}, Base.Broadcast.Extruded{CUDA.CuDeviceArray{ComplexF32, 4, 1}, NTuple{4, Bool}, NTuple{4, Int64}}}}, Int64) failed
KernelError: passing and using non-bitstype argument

Argument 4 to your kernel function is of type Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{4}, NTuple{4, Base.OneTo{Int64}}, typeof(*), Tuple{Base.Broadcast.Extruded{Array{Float32, 4}, NTuple{4, Bool}, NTuple{4, Int64}}, Base.Broadcast.Extruded{CUDA.CuDeviceArray{ComplexF32, 4, 1}, NTuple{4, Bool}, NTuple{4, Int64}}}}, which is not isbits:
.args is of type Tuple{Base.Broadcast.Extruded{Array{Float32, 4}, NTuple{4, Bool}, NTuple{4, Int64}}, Base.Broadcast.Extruded{CUDA.CuDeviceArray{ComplexF32, 4, 1}, NTuple{4, Bool}, NTuple{4, Int64}}} which is not isbits.
.1 is of type Base.Broadcast.Extruded{Array{Float32, 4}, NTuple{4, Bool}, NTuple{4, Int64}} which is not isbits.
.x is of type Array{Float32, 4} which is not isbits.

The package might need to be updated.

[khorrami1]

Finally, I can resolve this issue by adding a new registry in my environment:

] add https://github.com/ziyiyin97/AbstractFFTs.jl.git#master