fixup

Author: Fe-r-oz

ext/QuantumToolboxCUDAExt.jl

@@ -2,12 +2,14 @@ module QuantumToolboxCUDAExt
 
 using QuantumToolbox
 using QuantumToolbox: makeVal, getVal
-import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize
-import CUDA: cu, CuArray, allowscalar
+import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize, _safe_setindex!
+import CUDA: cu, CuArray, allowscalar, @allowscalar, has_cuda
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
 import SparseArrays: SparseVector, SparseMatrixCSC, sparse
 import CUDA.Adapt: adapt
 
+export _safe_setindex!
+
 allowscalar(false)
 
 @doc raw"""
@@ -108,10 +110,17 @@ function QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, rows, cols, vals,
     cpu_sparse = sparse(rows, cols, vals, m, n)
     return CuSparseMatrixCSC(cpu_sparse)
 end
-
 function QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, rows, cols, vals, m::Int64, n::Int64)
     cpu_sparse = sparse(rows, cols, vals, m, n)
     return CuSparseMatrixCSR(cpu_sparse)
 end
 
+function _safe_setindex!(A, val, idx)
+    if has_cuda() && isa(A, CuArray)
+        @allowscalar A[idx] = val
+    else
+        A[idx] = val
+    end
+end
+
 end

src/steadystate.jl

@@ -417,7 +417,7 @@ function _steadystate_fourier(
 
     v0 = _dense_similar(L_0_mat, n_fourier * N)
     fill!(v0, 0)
-    @allowscalar v0[n_max*N+1] = weight
+    _safe_setindex!(v0, weight, n_max * N + 1)
 
     if !isnothing(solver.Pl)
         kwargs = merge((; kwargs...), (Pl = solver.Pl(M),))

src/utilities.jl

@@ -135,6 +135,7 @@ _dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
 
 _sparse_similar(A::AbstractArray, args...) = sparse(args...)
 _sparse_similar(A::AbstractArray, m::Int, n::Int) = spzeros(eltype(A), m, n)
+function _safe_setindex! end
 
 _Ginibre_ensemble(n::Int, rank::Int = n) = randn(ComplexF64, n, rank) / sqrt(n)