change CUDA sparse to CSR, rename function createing sparse matrices

euro-hpc-pl · Mar 20, 2024 · ce82b3a · ce82b3a
1 parent 9cc37a8
commit ce82b3a
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 22 deletions.
diff --git 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
@@ -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
@@ -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]