diff --git a/lib/MadNLPGPU/src/KKT/sparse.jl b/lib/MadNLPGPU/src/KKT/sparse.jl index 47d229cc..36c1e575 100644 --- a/lib/MadNLPGPU/src/KKT/sparse.jl +++ b/lib/MadNLPGPU/src/KKT/sparse.jl @@ -112,12 +112,14 @@ function MadNLP.get_sparse_condensed_ext( jt_map, hess_map, ) where {T,VT<:CuVector{T}} - hess_com_ptr = map((i, j) -> (i, j), hess_map, 1:length(hess_map)) + zvals = CuVector{Int}(1:length(hess_map)) + hess_com_ptr = map((i, j) -> (i, j), hess_map, zvals) if length(hess_com_ptr) > 0 # otherwise error is thrown sort!(hess_com_ptr) end - jt_csc_ptr = map((i, j) -> (i, j), jt_map, 1:length(jt_map)) + jvals = CuVector{Int}(1:length(jt_map)) + jt_csc_ptr = map((i, j) -> (i, j), jt_map, jvals) if length(jt_csc_ptr) > 0 # otherwise error is thrown sort!(jt_csc_ptr) end @@ -320,7 +322,8 @@ end function MadNLP.coo_to_csc( coo::MadNLP.SparseMatrixCOO{T,I,VT,VI}, ) where {T,I,VT<:CuArray,VI<:CuArray} - coord = map((i, j, k) -> ((i, j), k), coo.I, coo.J, 1:length(coo.I)) + zvals = CuVector{Int}(1:length(coo.I)) + coord = map((i, j, k) -> ((i, j), k), coo.I, coo.J, zvals) if length(coord) > 0 sort!(coord, lt = (((i, j), k), ((n, m), l)) -> (j, i) < (m, n)) end @@ -505,7 +508,8 @@ function MadNLP._set_con_scale_sparse!( jac_I, jac_buffer, ) where {T,VT<:CuVector{T}} - inds = map((i, j) -> (i, j), jac_I, 1:length(jac_I)) + ind_jac = CuVector{Int}(1:length(jac_I)) + inds = map((i, j) -> (i, j), jac_I, ind_jac) if !isempty(inds) sort!(inds) end diff --git a/src/KKT/Sparse/condensed.jl b/src/KKT/Sparse/condensed.jl index b4697b5d..694231c7 100644 --- a/src/KKT/Sparse/condensed.jl +++ b/src/KKT/Sparse/condensed.jl @@ -221,15 +221,18 @@ nzval(H) = H.nzval n = size(H,2) + # N.B.: we should ensure that zind is allocated on the proper device. + zind = similar(nzval(H), Int, size(H, 2)) + zind .= 1:size(H, 2) map!( i->(-1,i,0), @view(sym[1:n]), - 1:size(H,2) + zind, ) map!( i->(i,i), @view(sym2[1:n]), - 1:size(H,2) + zind, ) _build_condensed_aug_symbolic_hess(