Fix bugs for initializing equilibrium and preprocessing large socs

cvxgrp · Dec 7, 2024 · c752860 · c752860
1 parent bcba042
commit c752860
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 16 deletions.
diff --git a/src/gpucones/augment_socp.jl b/src/gpucones/augment_socp.jl
@@ -56,13 +56,18 @@ function augment_data(At0,
     return Atnew, bnew, conenew, augx_idx
 end
 
-function augment_A_b(cones,
-    P,
+function augment_A_b_soc(
+    cones::Vector{SupportedCone},
+    P::AbstractMatrix{T},
     q::Vector{T},
-    A,
-    b,
-    size_soc,
-    num_socs, last_sizes, soc_indices, soc_starts) where {T}
+    A::AbstractMatrix{T},
+    b::Vector{T},
+    size_soc::Int64,
+    num_socs::Vector{Int}, 
+    last_sizes::Vector{Int}, 
+    soc_indices::Vector{Int}, 
+    soc_starts::Vector{Int}
+) where {T}
 
     (m,n) = size(A)
 
@@ -115,7 +120,10 @@ function augment_A_b(cones,
     return Pnew,vcat(q,zeros(extra_dim)),SparseMatrixCSC(Atnew'), bnew, conesnew
 end
 
-function expand_soc(cones,size_soc)
+function expand_soc(
+    cones::Vector{SupportedCone},
+    size_soc::Int64
+)
     n_large_soc = 0
     soc_indices = sizehint!(Int[],length(cones))            #Indices of large second-order cones 
     soc_starts = sizehint!(Int[],length(cones))          #Starting index of each large second-order cone 

diff --git a/src/problemdata.jl b/src/problemdata.jl
@@ -29,13 +29,16 @@ function DefaultProblemData{T}(
 		(A, b, cones) = (A_new, b_new, cones_new)
 	end
 
+	use_gpu = settings.direct_solve_method === :cudss ? true : false
 	# Preprocess large second-order cones
-	if settings.direct_solve_method === :cudss 
+	if use_gpu
 		size_soc = GPUsocSize
 		num_socs, last_sizes, soc_indices, soc_starts = expand_soc(cones,size_soc)
-		P_new,q_new,A_new,b_new,cones_new =  augment_A_b(cones,P,q,A,b,size_soc,num_socs, last_sizes, soc_indices, soc_starts)
 
-		(P, q, A, b, cones) = (P_new, q_new, A_new, b_new, cones_new)
+		if (length(num_socs) > 0)
+			P_new,q_new,A_new,b_new,cones_new =  augment_A_b_soc(cones,P,q,A,b,size_soc,num_socs, last_sizes, soc_indices, soc_starts)
+			(P, q, A, b, cones) = (P_new, q_new, A_new, b_new, cones_new)
+		end
 	end
 
 	# chordal decomposition : return nothing if disabled or no decomp
@@ -69,7 +72,7 @@ function DefaultProblemData{T}(
 	#this ensures m is the *reduced* size m
 	(m,n) = size(A_new)
 
-	equilibration = DefaultEquilibration{T}(n,m)
+	equilibration = DefaultEquilibration{T}(n,m,use_gpu)
 
 	normq = norm(q_new, Inf)
 	normb = norm(b_new, Inf)

diff --git a/src/solution.jl b/src/solution.jl
@@ -39,7 +39,7 @@ function solution_post_process!(
 	elseif settings.direct_solve_method === :cudss && (length(solution.x) != length(variables.x))
 		lenx = length(solution.x)
 		@. solution.x = variables.x[1:lenx]		#extra entries are slack variables for large second-order cones
-		println("Solve an augmented problem that only returns value x")
+		# println("Solve an augmented problem that only returns value x")
 	else
 		@. solution.x = variables.x
 		@. solution.z = variables.z 

diff --git a/src/types.jl b/src/types.jl
@@ -150,6 +150,7 @@ struct DefaultEquilibration{T} <: AbstractEquilibration{T}
     function DefaultEquilibration{T}(
         n::Int64,
         m::Int64,
+        use_gpu::Bool
     ) where {T}
 
         #Left/Right diagonal scaling for problem data
@@ -160,10 +161,17 @@ struct DefaultEquilibration{T} <: AbstractEquilibration{T}
 
         c    = Ref(T(1.))
 
-        d_gpu = n > 0 ? unsafe_wrap(CuArray{T,1},d) : CUDA.zeros(T,0)
-        dinv_gpu = n > 0 ? unsafe_wrap(CuArray{T,1},dinv) : CUDA.zeros(T,0)
-        e_gpu = m > 0 ? unsafe_wrap(CuArray{T,1},e) : CUDA.zeros(T,0)
-        einv_gpu = m > 0 ? unsafe_wrap(CuArray{T,1},einv) : CUDA.zeros(T,0)
+        if use_gpu
+            d_gpu = n > 0 ? unsafe_wrap(CuArray{T,1},d) : CUDA.zeros(T,0)
+            dinv_gpu = n > 0 ? unsafe_wrap(CuArray{T,1},dinv) : CUDA.zeros(T,0)
+            e_gpu = m > 0 ? unsafe_wrap(CuArray{T,1},e) : CUDA.zeros(T,0)
+            einv_gpu = m > 0 ? unsafe_wrap(CuArray{T,1},einv) : CUDA.zeros(T,0)
+        else
+            d_gpu = Vector{Float64}()
+            dinv_gpu = Vector{Float64}()
+            e_gpu = Vector{Float64}()
+            einv_gpu = Vector{Float64}()
+        end
 
         new(d,dinv,e,einv,c,d_gpu,dinv_gpu,e_gpu,einv_gpu)
     end