add allowed_setindex method for CUDA

Author: Fe-r-oz

ext/QuantumToolboxCUDAExt.jl

@@ -2,8 +2,8 @@ module QuantumToolboxCUDAExt
 
 using QuantumToolbox
 using QuantumToolbox: makeVal, getVal
-import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize
-import CUDA: cu, CuArray, allowscalar, @allowscalar, has_cuda, AbstractCuArray
+import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize, allowed_setindex!
+import CUDA: cu, CuArray, allowscalar, @allowscalar, has_cuda
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
 import SparseArrays: SparseVector, SparseMatrixCSC, sparse, spzeros
 import CUDA.Adapt: adapt
@@ -104,11 +104,11 @@ QuantumToolbox.to_dense(::Type{T}, A::AbstractCuSparseArray) where {T<:Number} =
 
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
-_sparse_similar(A::CuSparseMatrixCSC, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
     CuSparseMatrixCSC(sparse(I, J, V, m, n))
-_sparse_similar(A::CuSparseMatrixCSC, m::Int, n::Int) = CuSparseMatrixCSC(spzeros(eltype(A), m, n))
-_sparse_similar(A::CuSparseMatrixCSR, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, m::Int, n::Int) = CuSparseMatrixCSC(spzeros(eltype(A), m, n))
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
     CuSparseMatrixCSR(sparse(I, J, V, m, n))
-_sparse_similar(A::CuSparseMatrixCSR, m::Int, n::Int) = CuSparseMatrixCSR(spzeros(eltype(A), m, n))
-allowed_setindex!(A::AbstractCuSparseArray, v, I...) = @allowscalar A[I...] = v
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, m::Int, n::Int) = CuSparseMatrixCSR(spzeros(eltype(A), m, n))
+QuantumToolbox.allowed_setindex!(A::AbstractCuSparseArray, v, I...) = @allowscalar A[I...] = v
 end

src/steadystate.jl

@@ -397,7 +397,7 @@ function _steadystate_fourier(
     v0 = similar(L_0_mat, n_fourier * N)
     fill!(v0, zero(T))
     target_idx = n_max*N + 1
-    allowed_setindex!(v0, weight, target_idx)
+    QuantumToolbox.allowed_setindex!(v0, weight, target_idx)
 
     (haskey(kwargs, :Pl) || haskey(kwargs, :Pr)) && error("The use of preconditioners must be defined in the solver.")
     if !isnothing(solver.Pl)

src/utilities.jl

@@ -200,3 +200,5 @@ _convert_eltype_wordsize(::Type{T}, ::Val{64}) where {T<:AbstractFloat} = Float6
 _convert_eltype_wordsize(::Type{T}, ::Val{32}) where {T<:AbstractFloat} = Float32
 _convert_eltype_wordsize(::Type{Complex{T}}, ::Val{64}) where {T<:Union{Int,AbstractFloat}} = ComplexF64
 _convert_eltype_wordsize(::Type{Complex{T}}, ::Val{32}) where {T<:Union{Int,AbstractFloat}} = ComplexF32
+
+function allowed_setindex! end