Implement sametype_transpose and use it for A -> At (#35)

Author: gdalle

src/CoolPDLP.jl

@@ -50,6 +50,8 @@ end
 
 include("MOI_wrapper.jl")
 
+@public sametype_transpose
+
 export GPUSparseMatrixCOO, GPUSparseMatrixCSR, GPUSparseMatrixELL
 
 export MILP, nbvar, nbvar_int, nbvar_cont, nbcons, nbcons_eq, nbcons_ineq

src/problems/milp.jl

@@ -10,6 +10,7 @@ Represent a Mixed Integer Linear Program in "cuPDLPx form":
 
     MILP(;
         c, lv, uv, A, lc, uc,
+        At=sametype_transpose(A),
         [D1, D2, int_var, var_names, dataset, name, path]
     )
 
@@ -21,6 +22,7 @@ struct MILP{
         T <: Number,
         V <: DenseVector{T},
         M <: AbstractMatrix{T},
+        Mt <: AbstractMatrix{T},
         Vb <: DenseVector{Bool},
     }
     "objective vector"
@@ -32,7 +34,7 @@ struct MILP{
     "constraint matrix"
     A::M
     "transposed constraint matrix"
-    At::M
+    At::Mt
     "constraint lower bound"
     lc::V
     "constraint upper bound"
@@ -57,7 +59,7 @@ struct MILP{
             lv,
             uv,
             A,
-            At = convert(typeof(A), transpose(A)),
+            At = sametype_transpose(A),
             lc,
             uc,
             D1 = Diagonal(one!(similar(lc))),
@@ -77,13 +79,15 @@ struct MILP{
 
         T = Base.promote_eltype(c, lv, uv, A, At, lc, uc, D1, D2)
         V = promote_type(typeof(c), typeof(lv), typeof(uv), typeof(lc), typeof(uc))
-        M = promote_type(typeof(A), typeof(At))
+        M = typeof(A)
+        Mt = typeof(At)
         Vb = typeof(int_var)
 
         if (
                 !isconcretetype(T) ||
                     !isconcretetype(V) ||
                     !isconcretetype(M) ||
+                    !isconcretetype(Mt) ||
                     !isconcretetype(Vb)
             )
             throw(ArgumentError("Abstract type parameter"))
@@ -95,7 +99,7 @@ struct MILP{
             name = splitext(splitpath(path)[end])[1]
         end
 
-        return new{T, V, M, Vb}(
+        return new{T, V, M, Mt, Vb}(
             c,
             lv,
             uv,
@@ -136,14 +140,14 @@ function MILP(qps::QPSData; kwargs...)
     )
 end
 
-function Base.show(io::IO, milp::MILP{T, V, M}) where {T, V, M}
+function Base.show(io::IO, milp::MILP{T, V, M, Mt}) where {T, V, M, Mt}
     return print(
         io, """
         MILP instance $(milp.name) from dataset $(milp.dataset):
         - types:
           - values $T
           - vectors $V
-          - matrices $M
+          - matrices $(M == Mt ? M : (M, Mt))
         - variables: $(nbvar(milp))
           - $(nbvar_cont(milp)) continuous
           - $(nbvar_int(milp)) integer

src/utils/linalg.jl

@@ -1,3 +1,12 @@
+"""
+    sametype_transpose(A::AbstractMatrix)
+
+Return a matrix of the same type of `A` containing `transpose(A)` (as opposed to a `Transpose{...}` wrapper).
+
+The default implementation is just `convert(typeof(A), transpose(A))` but it may need to be overloaded for certain matrix types.
+"""
+sametype_transpose(A::AbstractMatrix) = convert(typeof(A), transpose(A))
+
 zero!(x::AbstractArray) = fill!(x, zero(eltype(x)))
 one!(x::AbstractArray) = fill!(x, one(eltype(x)))
 

src/utils/mat_coo.jl

@@ -47,11 +47,19 @@ function Adapt.adapt_structure(to, A::GPUSparseMatrixCOO)
     )
 end
 
+function SparseArrays.SparseMatrixCSC(A::GPUSparseMatrixCOO)
+    return sparse(Vector(A.rowval), Vector(A.colval), Vector(A.nzval), A.m, A.n)
+end
+
 function GPUSparseMatrixCOO(A::SparseMatrixCSC{T, Ti}) where {T, Ti}
     rowval, colval, nzval = findnz(A)
     return GPUSparseMatrixCOO(A.m, A.n, rowval, colval, nzval)
 end
 
+function sametype_transpose(A::GPUSparseMatrixCOO)
+    return GPUSparseMatrixCOO(A.n, A.m, A.colval, A.rowval, A.nzval)
+end
+
 @kernel function spmv_coo!(
         c::DenseVector{T},
         A_rowval::DenseVector{Ti},

src/utils/mat_csr.jl

@@ -56,8 +56,21 @@ function Adapt.adapt_structure(to, A::GPUSparseMatrixCSR)
 end
 
 function GPUSparseMatrixCSR(A::SparseMatrixCSC{T, Ti}) where {T, Ti}
-    At = sparse(transpose(A))
-    return GPUSparseMatrixCSR(At.n, At.m, At.colptr, At.rowval, At.nzval)
+    At_csc = SparseMatrixCSC(transpose(A))
+    return GPUSparseMatrixCSR(At_csc.n, At_csc.m, At_csc.colptr, At_csc.rowval, At_csc.nzval)
+end
+
+function SparseArrays.SparseMatrixCSC(A::GPUSparseMatrixCSR)
+    At_csc = SparseMatrixCSC(A.n, A.m, Vector(A.rowptr), Vector(A.colval), Vector(A.nzval))
+    return SparseMatrixCSC(transpose(At_csc))
+end
+
+function sametype_transpose(A::GPUSparseMatrixCSR)
+    A_csc = SparseMatrixCSC(A)
+    return adapt(
+        get_backend(A),
+        GPUSparseMatrixCSR(A_csc.n, A_csc.m, A_csc.colptr, A_csc.rowval, A_csc.nzval)
+    )
 end
 
 @kernel function spmv_csr!(

src/utils/mat_ell.jl

@@ -69,6 +69,19 @@ function GPUSparseMatrixELL(A::SparseMatrixCSC{T, Ti}) where {T, Ti}
     return GPUSparseMatrixELL(m, n, colval, nzval)
 end
 
+function SparseArrays.SparseMatrixCSC(A::GPUSparseMatrixELL)
+    I = Matrix(map(ind -> Tuple(ind)[1], CartesianIndices(A.colval)))
+    J = Matrix(A.colval)
+    V = Matrix(A.nzval)
+    inds = J .> 0
+    return sparse(I[inds], J[inds], V[inds], A.m, A.n)
+end
+
+function sametype_transpose(A::GPUSparseMatrixELL)
+    At = SparseMatrixCSC(transpose(SparseMatrixCSC(A)))
+    return adapt(get_backend(A), GPUSparseMatrixELL(At))
+end
+
 @kernel function spmv_ell!(
         c::DenseVector{T},
         A_colval::AbstractMatrix{Ti},

test/problems/modify.jl

@@ -19,9 +19,21 @@ end
 
 @testset "Change backend" begin
     milp_flexible = CoolPDLP.set_matrix_type(GPUSparseMatrixCSR, milp)
-    @test milp_flexible isa MILP{Float64, Vector{Float64}, GPUSparseMatrixCSR{Float64, Int, Vector{Float64}, Vector{Int}}, Vector{Bool}}
+    @test milp_flexible isa MILP{
+        Float64,
+        Vector{Float64},
+        GPUSparseMatrixCSR{Float64, Int, Vector{Float64}, Vector{Int}},
+        GPUSparseMatrixCSR{Float64, Int, Vector{Float64}, Vector{Int}},
+        Vector{Bool},
+    }
     milp_gpu = adapt(JLBackend(), milp_flexible)
-    @test milp_gpu isa MILP{Float64, JLVector{Float64}, GPUSparseMatrixCSR{Float64, Int, JLVector{Float64}, JLVector{Int}}, JLVector{Bool}}
+    @test milp_gpu isa MILP{
+        Float64,
+        JLVector{Float64},
+        GPUSparseMatrixCSR{Float64, Int, JLVector{Float64}, JLVector{Int}},
+        GPUSparseMatrixCSR{Float64, Int, JLVector{Float64}, JLVector{Int}},
+        JLVector{Bool},
+    }
 
     sol_gpu = adapt(JLBackend(), sol)
     @test sol_gpu isa PrimalDualSolution{Float64, JLVector{Float64}}

test/utils/linalg.jl

@@ -2,6 +2,7 @@ using CoolPDLP
 using LinearAlgebra
 using SparseArrays
 using Test
+using Random: Xoshiro
 
 @testset "Simple projections" begin
     for x in randn(100)
@@ -67,3 +68,10 @@ end
         end
     end
 end
+
+@testset "Same-type transpose" begin
+    for A in Any[rand(10, 10), sprand(10, 10, 0.6)]
+        @test CoolPDLP.sametype_transpose(A) == transpose(A)
+        @test CoolPDLP.sametype_transpose(A) isa typeof(A)
+    end
+end

test/utils/matrices.jl

@@ -19,16 +19,20 @@ c_candidates = [rand(size(A, 1)) for A in A_candidates];
 
 function test_sparse_matrix(::Type{M}; A, b, c, α, β) where {M}
     A_jl = adapt(JLBackend(), M(A))
+    At_jl = adapt(JLBackend(), M(sparse(transpose(A))))
     b_jl, c_jl = jl(b), jl(c)
     @test @allowscalar Matrix(A_jl) == A
+    @test @allowscalar SparseMatrixCSC(A_jl) == A
     @test nnz(A_jl) == nnz(A)
     @test get_backend(A_jl) isa JLBackend
     @test mul!(copy(c_jl), A_jl, b_jl, α, β) ≈ α * A * b + β * c
+    @test @allowscalar Matrix(CoolPDLP.sametype_transpose(A_jl)) == transpose(A)
+    @test typeof(CoolPDLP.sametype_transpose(A_jl)) == typeof(At_jl)
     return nothing
 end
 
 @testset for M in (GPUSparseMatrixCOO, GPUSparseMatrixCSR, GPUSparseMatrixELL)
-    for (A, b, c) in zip(A_candidates, b_candidates, c_candidates)
+    for (A, b, c) in collect(zip(A_candidates, b_candidates, c_candidates))
         test_sparse_matrix(M; A, b, c, α, β)
     end
 end