From ce82b3a6f38cd9aa6c3fc1032a8a1d124255caf0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20Pawela?= Date: Wed, 20 Mar 2024 15:50:40 +0100 Subject: [PATCH] change CUDA sparse to CSR, rename function createing sparse matrices --- src/contractions/site.jl | 17 ++++------------- src/contractions/sparse.jl | 14 +++++++------- src/contractions/virtual.jl | 4 ++-- 3 files changed, 13 insertions(+), 22 deletions(-) diff --git a/src/contractions/site.jl b/src/contractions/site.jl index 5aadbd5..bae7c6c 100644 --- a/src/contractions/site.jl +++ b/src/contractions/site.jl @@ -63,7 +63,7 @@ function contract_sparse_with_three( le = @view loc_exp[from:to] outp .*= reshape(le, 1, 1, :) outpp = reshape(outp, s1 * s4, :) - ipr, rf, rt = SparseCSC(R, lp, kout, device; from, to) + ipr, rf, rt = sparse(R, lp, kout, device; from, to) @inbounds out[rf:rt, :, :] .+= reshape(ipr * outpp', :, s1, s4) from = to + 1 end @@ -155,7 +155,7 @@ function update_reduced_env_right( outp = dropdims(Bp ⊠ REp, dims = 2) outp .*= reshape(le .* Kp, 1, :) - ipr, rf, rt = SparseCSC(R, M.lp, k1, device; from, to) + ipr, rf, rt = sparse(R, M.lp, k1, device; from, to) @inbounds out[rf:rt, :] .+= ipr * outp' from = to + 1 end @@ -169,16 +169,7 @@ function contract_tensors43(M::SiteTensor{R,4}, B::Tensor{R,3}) where {R<:Real} sm1, sm2, sm3 = size.(Ref(M.lp), M.projs[1:3]) @inbounds Bp = B[:, :, p4] .* reshape(M.loc_exp, 1, 1, :) @cast Bp[(x, y), z] := Bp[x, y, z] - ip123 = SparseCSC(R, M.lp, M.projs[1], M.projs[2], M.projs[3], device) - @show nnz(ip123) - @show size(ip123) - @show size(Bp) - @show size(Bp') - @show typeof(Bp) - @show typeof(ip123) - @show typeof(Bp') - @show (sm1, sm2, sm3, sb1, sb2) - @show size(ip123 * Bp') + ip123 = sparse(R, M.lp, M.projs[1], M.projs[2], M.projs[3], device) out = reshape(ip123 * Bp', sm1, sm2, sm3, sb1, sb2) out = permutedims(out, (4, 1, 5, 3, 2)) reshape(out, sb1 * sm1, sb2 * sm3, sm2) @@ -198,7 +189,7 @@ function corner_matrix( @cast outp[(x, y), z] := outp[x, y, z] sm1, sm2 = maximum(projs[1]), maximum(projs[2]) @inbounds p12 = projs[1] .+ (projs[2] .- 1) .* sm1 - ip12 = SparseCSC(R, p12; mp = sm1 * sm2) + ip12 = sparse(R, p12; mp = sm1 * sm2) out = reshape(ip12 * outp', sm1, maximum(projs[2]), size(B, 1), size(C, 2)) permutedims(out, (3, 1, 4, 2)) end diff --git a/src/contractions/sparse.jl b/src/contractions/sparse.jl index fc58547..e0bc0d6 100644 --- a/src/contractions/sparse.jl +++ b/src/contractions/sparse.jl @@ -10,17 +10,17 @@ end =# # TODO shouldn't we have CSR format instead? -function SparseCSC(::Type{R}, p::CuArray{Int64,1}; mp = nothing) where {R<:Real} +function SparseArrays.sparse(::Type{R}, p::CuArray{Int64,1}; mp = nothing) where {R<:Real} n = length(p) if isnothing(mp) mp = maximum(p) end cn = CuArray(1:n+1) # aux_cusparse(R, n) co = CUDA.ones(R, n) - CuSparseMatrixCSC(cn, p, co, (mp, n)) + CuSparseMatrixCSR(CuSparseMatrixCSC(cn, p, co, (mp, n))) # TODO: Change when CUDA.jl is fixed end -function SparseCSC(::Type{R}, p::Vector{Int64}; mp = nothing) where {R<:Real} +function SparseArrays.sparse(::Type{R}, p::Vector{Int64}; mp = nothing) where {R<:Real} n = length(p) if isnothing(mp) mp = maximum(p) @@ -30,7 +30,7 @@ function SparseCSC(::Type{R}, p::Vector{Int64}; mp = nothing) where {R<:Real} sparse(p, cn, co, mp, n) end -@memoize Dict function SparseCSC( +@memoize Dict function SparseArrays.sparse( ::Type{T}, lp::PoolOfProjectors, k1::R, @@ -47,10 +47,10 @@ end if device == :GPU p = CuArray(p) end - SparseCSC(T, p; mp = s1 * s2 * s3) + sparse(T, p; mp = s1 * s2 * s3) end -@memoize Dict function SparseCSC( +@memoize Dict function SparseArrays.sparse( ::Type{R}, lp::PoolOfProjectors, k::Int, @@ -65,6 +65,6 @@ end if device == :GPU pp = CuArray(pp) end - ipr = SparseCSC(R, pp .- (rf - 1)) + ipr = sparse(R, pp .- (rf - 1)) (ipr, rf, rt) end diff --git a/src/contractions/virtual.jl b/src/contractions/virtual.jl index ebaa4a3..78d6f57 100644 --- a/src/contractions/virtual.jl +++ b/src/contractions/virtual.jl @@ -397,12 +397,12 @@ function contract_tensors43(M::VirtualTensor{R,4}, B::Tensor{R,3}) where {R<:Rea B = reshape(B, (slb, srb, slcb, srcb)) - pls = SparseCSC(R, M.lp, p_lb, p_l, p_lt, :CPU) + pls = sparse(R, M.lp, p_lb, p_l, p_lt, :CPU) pls = typeof(B) <: CuArray ? CuArray(pls) : Array(pls) pls = reshape(pls, (slcb, slc, slct * slcp)) pls = permutedims(pls, (3, 1, 2)) # [(slct, slcp), lcb, lc] - prs = SparseCSC(R, M.lp, p_rb, p_r, p_rt, :CPU) + prs = sparse(R, M.lp, p_rb, p_r, p_rt, :CPU) prs = typeof(B) <: CuArray ? CuArray(prs) : Array(prs) prs = reshape(prs, (srcb, src, srct * srcp)) prs = permutedims(prs, (3, 1, 2)) # [(rct, rcp), rcb, rc]