Scalar indexing is disallowed

[gautierronan]

Bug Description

Hi, I'm trying to run a simple example on GPU, and running into a Scalar indexing is disallowed. I might be doing something wrong, as I'm new to the library, and to julia overall.

Here's my example:

using QuantumToolbox
using CUDA
using CUDA.CUSPARSE
CUDA.allowscalar(false)

# parameters
kappa_2 = 1.0
g_cnot = 0.3
nbar = 4.0
num_tsave = 100
N = 32

# save times
alpha = sqrt(nbar)
gate_time = pi / (4 * alpha * g_cnot)
tlist = range(0, gate_time, num_tsave)

# operators
ac = cu(tensor(destroy(N), qeye(N)))
at = cu(tensor(qeye(N), destroy(N)))
i = cu(tensor(qeye(N), qeye(N)))
H = g_cnot * (ac + ac') * (at' * at - nbar * i)
c_ops = [sqrt(kappa_2) * (ac * ac - nbar * i)]

# initial state
plus = normalize(coherent(N, alpha) + coherent(N, -alpha))
psi0 = cu(tensor(plus, plus))

# run mesolve
sol = mesolve(H, psi0, tlist, c_ops)

and the stack trace:

ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/indexing.jl:50 [inlined]
  [6] scalar_getindex
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/indexing.jl:36 [inlined]
  [7] _getindex
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/indexing.jl:19 [inlined]
  [8] getindex
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/indexing.jl:17 [inlined]
  [9] getindex
    @ ~/.julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/adjtrans.jl:334 [inlined]
 [10] _getindex
    @ ./abstractarray.jl:1358 [inlined]
 [11] getindex
    @ ./abstractarray.jl:1312 [inlined]
 [12] iterate
    @ ./abstractarray.jl:1209 [inlined]
 [13] iterate
    @ ./abstractarray.jl:1207 [inlined]
 [14] SparseMatrixCSC{ComplexF64, Int64}(M::Adjoint{ComplexF64, CuArray{ComplexF64, 2, CUDA.DeviceMemory}})
    @ SparseArrays ~/.julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:891
 [15] _sparsem
    @ ~/.julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/sparseconvert.jl:90 [inlined]
 [16] SparseMatrixCSC
    @ ~/.julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/sparseconvert.jl:20 [inlined]
 [17] convert
    @ ~/.julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:980 [inlined]
 [18] sparse
    @ ~/.julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:1016 [inlined]
 [19] _sprepost(A::CuArray{ComplexF64, 2, CUDA.DeviceMemory}, B::Adjoint{ComplexF64, CuArray{…}})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:13
 [20] _lindblad_dissipator(O::CuArray{ComplexF64, 2, CUDA.DeviceMemory}, Id::Diagonal{Bool, Vector{Bool}})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:52
 [21] lindblad_dissipator(O::QuantumObject{…}, Id_cache::Diagonal{…})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:135
 [22] #72
    @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:187 [inlined]
 [23] mapreduce_first
    @ ./reduce.jl:421 [inlined]
 [24] _mapreduce(f::QuantumToolbox.var"#72#76"{Diagonal{…}}, op::typeof(+), ::IndexLinear, A::Vector{QuantumObject{…}})
    @ Base ./reduce.jl:432
 [25] _mapreduce_dim
    @ ./reducedim.jl:337 [inlined]
 [26] mapreduce
    @ ./reducedim.jl:329 [inlined]
 [27] _sum_lindblad_dissipators(c_ops::Vector{QuantumObject{…}}, Id_cache::Diagonal{Bool, Vector{…}})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:187
 [28] liouvillian(H::QuantumObject{…}, c_ops::Vector{…}, Id_cache::Diagonal{…})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:167
 [29] liouvillian
    @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/superoperators.jl:165 [inlined]
 [30] _mesolve_make_L_QobjEvo(H::QuantumObject{…}, c_ops::Vector{…})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/time_evolution/mesolve.jl:3
 [31] mesolveProblem(H::QuantumObject{…}, ψ0::QuantumObject{…}, tlist::StepRangeLen{…}, c_ops::Vector{…}; e_ops::Nothing, params::SciMLBase.NullParameters, progress_bar::Val{…}, inplace::Val{…}, kwargs::@Kwargs{…})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/time_evolution/mesolve.jl:75
 [32] mesolveProblem
    @ ~/.julia/packages/QuantumToolbox/YmsAS/src/time_evolution/mesolve.jl:56 [inlined]
 [33] mesolve(H::QuantumObject{…}, ψ0::QuantumObject{…}, tlist::StepRangeLen{…}, c_ops::Vector{…}; alg::OrdinaryDiffEqTsit5.Tsit5{…}, e_ops::Nothing, params::SciMLBase.NullParameters, progress_bar::Val{…}, inplace::Val{…}, kwargs::@Kwargs{})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/time_evolution/mesolve.jl:161
 [34] mesolve(H::QuantumObject{…}, ψ0::QuantumObject{…}, tlist::StepRangeLen{…}, c_ops::Vector{…})
    @ QuantumToolbox ~/.julia/packages/QuantumToolbox/YmsAS/src/time_evolution/mesolve.jl:146
 [35] top-level scope
    @ REPL[75]:1
Some type information was truncated. Use `show(err)` to see complete types.

Code to Reproduce the Bug

Code Output

Expected Behaviour

N/A

Your Environment

QuantumToolbox.jl: Quantum Toolbox in Julia
≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡
Copyright © QuTiP team 2022 and later.
Current admin team:
    Alberto Mercurio and Yi-Te Huang

Package information:
====================================
Julia              Ver. 1.11.4
QuantumToolbox     Ver. 0.29.1
SciMLOperators     Ver. 0.3.13
LinearSolve        Ver. 3.7.2
OrdinaryDiffEqCore Ver. 1.21.0

System information:
====================================
OS       : Linux (x86_64-linux-gnu)
CPU      : 24 × AMD EPYC 9334 32-Core Processor
Memory   : 235.941 GB
WORD_SIZE: 64
LIBM     : libopenlibm
LLVM     : libLLVM-16.0.6 (ORCJIT, znver4)
BLAS     : libopenblas64_.so (ilp64)
Threads  : 1 (on 24 virtual cores)

Additional Context

No response

[ytdHuang]

@gautierronan

Hello, welcome to the Julia world !

I just looked through your sample code. I think it is indeed a small bug !

This line you were running:

i = cu(tensor(qeye(N), qeye(N)))

It was a Diagonal-type matrix before the CUDA matrix conversion (cu).

But after the cu conversion, it becomes a CuArray (dense GPU) matrix. This makes H and the elements in c_ops become dense and somehow falls into scalar indexing in mesolve.

To make your code work now before we fix this, you can manually convert dense GPU array i into a sparse one, namely

i = to_sparse(cu(tensor(qeye(N), qeye(N))))

We will fix this bug in the next patch or minor release.

[albertomercurio]

The main problem is that cu on a diagonal term returns a dense array. This doesn't happen to Diagonal matrices with CUDA vectors so it is something on our side. I will make a PR fixing this, such that cu(qeye(N)) still remains a diagonal matrix.

[ytdHuang]

Currently waiting this PR JuliaGPU/CUDA.jl#2729 to be merged and CUDA.jl bump a patch release.