Skip to content

Commit 5588a30

Browse files
adhamsishakedregev
adhamsi
and
shakedregev
authored
SpGEMM for HyKKT (#366)
* main class and header * begin cpu * cpu implementation: no test yet * handle csc vs csr and implement minimal test * hip implemenation * cuda implementation * hip fix signature * fix cmakelists * minor fixes * Apply pre-commmit fixes * symbolic test * reuse test * fix for hip for reuse: broken * Apply pre-commmit fixes * compiler warning * reuse test working * reuse test cuda * documentation * Apply pre-commmit fixes * fix build error * Apply pre-commmit fixes * rename add functions to load * modify symbolic test for different case of D sparsity * Apply pre-commmit fixes * doc update * straggling unsaved changes * fixed typo --------- Co-authored-by: adhamsi <adhamsi@users.noreply.github.com> Co-authored-by: shakedregev <shakedvregev@gmail.com>
1 parent 825882c commit 5588a30

File tree

15 files changed

+1569
-3
lines changed

15 files changed

+1569
-3
lines changed

‎resolve/hykkt/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
add_subdirectory(permutation)
22
add_subdirectory(ruiz)
33
add_subdirectory(cholesky)
4+
add_subdirectory(spgemm)

‎resolve/hykkt/README.md‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
# HyKKT
2+
3+
## Description
4+
HyKKT (pronounced as _hiked_) is a package for solving systems of equations and unknowns resulting from an iterative solution of an optimization
5+
problem, for example optimal power flow, which uses hardware accelerators (GPUs) efficiently.
6+
7+
The HyKKT package contains a linear solver tailored for Karush Kuhn Tucker (KKT) linear systems and
8+
deployment on hardware accelerator hardware such as GPUs. The solver requires
9+
all blocks of the $4\times 4$ block system:
10+
11+
```math
12+
\begin{bmatrix}
13+
H + D_x & 0 & J_c^T & J_d^T \\
14+
0 & D_s & 0 & -I \\
15+
J_c & 0 & 0 & 0 \\
16+
J_d & -I & 0 & 0
17+
\end{bmatrix}
18+
\begin{bmatrix}
19+
\Delta x \\ \Delta s \\ \Delta y_c \\ \Delta y_d
20+
\end{bmatrix} =
21+
\begin{bmatrix}
22+
\tilde{r}_x \\ r_s \\ r_c \\ r_d
23+
\end{bmatrix}
24+
```
25+
26+
separately and solves the system to a desired
27+
numerical precision exactly via block reduction and conjugate gradient on the
28+
schur complement. Please see the [HyKKT paper](https://www.tandfonline.com/doi/abs/10.1080/10556788.2022.2124990) for mathematical details.
29+
30+
## Integration within ReSolve
31+
32+
HyKKT code within ReSolve is not yet fully integrated. The code is currently experimental, but is actively in development.
33+
The main issue is the requirement that the user pass matrix and vector blocks, rather than the constructed KKT system.
34+
This is currently not available within ReSolve as it is specific to HyKKT and KKT matrices.
35+
TODO:
36+
- Work with AMD to fix bug in Cholesky factorization.
37+
- Allow the representation of a matrix or vector by blocks within ReSolve.
38+
- Schur Compliment Conjugate Gradient within HyKKT.
39+
- Complete HyKKT integration within ReSolve.
40+
41+
42+
## Dependencies
43+
44+
Most HyKKT dependencies are identical to ReSolve. Exceptions are:
45+
- HyKKT does not require KLU
46+
- For AMD builds, HyKKT requires HIP/ROCm >= 6.4 due to a known bug in SpGEMM.

‎resolve/hykkt/spgemm/CMakeLists.txt‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#[[
2+
3+
@brief Build SpGEMM module for HyKKT solver
4+
5+
@author Adham Ibrahim <ibrahimas@ornl.gov>
6+
7+
]]
8+
9+
# Add source files for the SpGEMM module
10+
set(SPGEMM_SRC
11+
SpGEMM.cpp
12+
SpGEMMCpu.cpp
13+
)
14+
set(SPGEMM_HEADER_INSTALL
15+
SpGEMM.hpp
16+
SpGEMMImpl.hpp
17+
SpGEMMCpu.hpp
18+
SpGEMMCuda.hpp
19+
SpGEMMHip.hpp
20+
)
21+
22+
set(SPGEMM_CUDA_SRC
23+
SpGEMMCuda.cpp
24+
)
25+
26+
set(SPGEMM_ROCM_SRC
27+
SpGEMMHip.cpp
28+
)
29+
30+
# Build shared library ReSolve::hykkt
31+
add_library(resolve_hykkt_spgemm SHARED ${SPGEMM_SRC})
32+
33+
target_link_libraries(resolve_hykkt_spgemm PRIVATE resolve_logger ${suitesparse_cholmod})
34+
target_include_directories(resolve_hykkt_spgemm PUBLIC ${SUITESPARSE_INCLUDE_DIR})
35+
36+
if (RESOLVE_USE_CUDA)
37+
target_sources(resolve_hykkt_spgemm PRIVATE ${SPGEMM_CUDA_SRC})
38+
target_link_libraries(resolve_hykkt_spgemm PUBLIC resolve_backend_cuda)
39+
endif()
40+
41+
if (RESOLVE_USE_HIP)
42+
target_sources(resolve_hykkt_spgemm PRIVATE ${SPGEMM_ROCM_SRC})
43+
target_link_libraries(resolve_hykkt_spgemm PUBLIC resolve_backend_hip)
44+
endif()
45+
46+
# Link to dummy device backend if GPU support is not enabled
47+
if (NOT RESOLVE_USE_GPU)
48+
target_link_libraries(resolve_hykkt_spgemm PUBLIC resolve_backend_cpu)
49+
endif()
50+
51+
target_include_directories(resolve_hykkt_spgemm INTERFACE
52+
$<BUILD_INTERFACE:${CMAKE_SOURCE_DIR}>
53+
$<INSTALL_INTERFACE:include>
54+
)
55+
56+
install(FILES ${SPGEMM_HEADER_INSTALL} DESTINATION include/resolve/hykkt)

‎resolve/hykkt/spgemm/SpGEMM.cpp‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/**
2+
* @file SpGEMM.cpp
3+
* @author Adham Ibrahim (ibrahimas@ornl.gov)
4+
*/
5+
6+
#include "SpGEMM.hpp"
7+
8+
#include "SpGEMMCpu.hpp"
9+
#ifdef RESOLVE_USE_CUDA
10+
#include "SpGEMMCuda.hpp"
11+
#elif defined(RESOLVE_USE_HIP)
12+
#include "SpGEMMHip.hpp"
13+
#endif
14+
15+
namespace ReSolve
16+
{
17+
using real_type = ReSolve::real_type;
18+
using out = ReSolve::io::Logger;
19+
20+
namespace hykkt
21+
{
22+
/**
23+
* Constructor for SpGEMM
24+
* @param memspace[in] - Memory space for computation.
25+
* @param alpha[in] - Scalar multiplier for the product.
26+
* @param beta[in] - Scalar multiplier for the sum.
27+
*/
28+
SpGEMM::SpGEMM(memory::MemorySpace memspace, real_type alpha, real_type beta)
29+
: memspace_(memspace)
30+
{
31+
if (memspace_ == memory::HOST)
32+
{
33+
impl_ = new SpGEMMCpu(alpha, beta);
34+
}
35+
else
36+
{
37+
#ifdef RESOLVE_USE_CUDA
38+
impl_ = new SpGEMMCuda(alpha, beta);
39+
#elif defined(RESOLVE_USE_HIP)
40+
impl_ = new SpGEMMHip(alpha, beta);
41+
#else
42+
out::error() << "No GPU support enabled, and memory space set to DEVICE.\n";
43+
exit(1);
44+
#endif
45+
}
46+
}
47+
48+
/**
49+
* Destructor for SpGEMM
50+
*/
51+
SpGEMM::~SpGEMM()
52+
{
53+
delete impl_;
54+
}
55+
56+
/**
57+
* Loads the two matrices for the product
58+
* @param A[in] - Pointer to CSR matrix
59+
* @param B[in] - Pointer to CSR matrix
60+
*/
61+
void SpGEMM::loadProductMatrices(matrix::Csr* A, matrix::Csr* B)
62+
{
63+
impl_->loadProductMatrices(A, B);
64+
}
65+
66+
/**
67+
* Loads the sum matrix for the operation
68+
* @param D[in] - Pointer to CSR matrix
69+
*/
70+
void SpGEMM::loadSumMatrix(matrix::Csr* D)
71+
{
72+
impl_->loadSumMatrix(D);
73+
}
74+
75+
/**
76+
* Loads the result matrix
77+
* @param E[in] - Pointer to pointer to CSR matrix
78+
*/
79+
void SpGEMM::loadResultMatrix(matrix::Csr** E_ptr)
80+
{
81+
impl_->loadResultMatrix(E_ptr);
82+
}
83+
84+
/**
85+
* Computes the result of the SpGEMM operation
86+
* E = alpha * A * B + beta * D
87+
*/
88+
void SpGEMM::compute()
89+
{
90+
impl_->compute();
91+
}
92+
} // namespace hykkt
93+
} // namespace ReSolve

‎resolve/hykkt/spgemm/SpGEMM.hpp‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/**
2+
* @file SpGEMM.hpp
3+
* @author Adham Ibrahim (ibrahimas@ornl.gov)
4+
*/
5+
6+
#pragma once
7+
#include "SpGEMMImpl.hpp"
8+
#include <resolve/Common.hpp>
9+
#include <resolve/MemoryUtils.hpp>
10+
#include <resolve/matrix/Csr.hpp>
11+
12+
namespace ReSolve
13+
{
14+
using real_type = ReSolve::real_type;
15+
16+
namespace hykkt
17+
{
18+
class SpGEMM
19+
{
20+
public:
21+
// computes E = alpha * A * B + beta * D
22+
SpGEMM(memory::MemorySpace memspace, real_type alpha, real_type beta);
23+
~SpGEMM();
24+
25+
void loadProductMatrices(matrix::Csr* A, matrix::Csr* B);
26+
void loadSumMatrix(matrix::Csr* D);
27+
void loadResultMatrix(matrix::Csr** E_ptr);
28+
29+
void compute();
30+
31+
private:
32+
memory::MemorySpace memspace_;
33+
34+
SpGEMMImpl* impl_;
35+
};
36+
} // namespace hykkt
37+
} // namespace ReSolve

‎resolve/hykkt/spgemm/SpGEMMCpu.cpp‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
/**
2+
* @file SpGEMMCpu.cpp
3+
* @author Adham Ibrahim (ibrahimas@ornl.gov)
4+
* @brief Implementation of SpGEMM using CHOLMOD for CPU
5+
*/
6+
7+
#include "SpGEMMCpu.hpp"
8+
9+
namespace ReSolve
10+
{
11+
using real_type = ReSolve::real_type;
12+
13+
namespace hykkt
14+
{
15+
/**
16+
* Constructor for SpGEMMCpu
17+
* @param alpha[in] - Scalar multiplier for the product.
18+
* @param beta[in] - Scalar multiplier for the sum.
19+
*/
20+
SpGEMMCpu::SpGEMMCpu(real_type alpha, real_type beta)
21+
: alpha_(alpha), beta_(beta)
22+
{
23+
cholmod_start(&Common_);
24+
25+
A_ = nullptr;
26+
B_ = nullptr;
27+
D_ = nullptr;
28+
}
29+
30+
/**
31+
* Destructor for SpGEMMCpu
32+
*/
33+
SpGEMMCpu::~SpGEMMCpu()
34+
{
35+
if (A_)
36+
{
37+
cholmod_free_sparse(&A_, &Common_);
38+
}
39+
if (B_)
40+
{
41+
cholmod_free_sparse(&B_, &Common_);
42+
}
43+
if (D_)
44+
{
45+
cholmod_free_sparse(&D_, &Common_);
46+
}
47+
cholmod_finish(&Common_);
48+
}
49+
50+
void SpGEMMCpu::loadProductMatrices(matrix::Csr* A, matrix::Csr* B)
51+
{
52+
if (!A_)
53+
{
54+
A_ = allocateCholmodType(A);
55+
}
56+
if (!B_)
57+
{
58+
B_ = allocateCholmodType(B);
59+
}
60+
copyDataToCholmodType(A, A_);
61+
copyDataToCholmodType(B, B_);
62+
}
63+
64+
void SpGEMMCpu::loadSumMatrix(matrix::Csr* D)
65+
{
66+
if (!D_)
67+
{
68+
D_ = allocateCholmodType(D);
69+
}
70+
copyDataToCholmodType(D, D_);
71+
}
72+
73+
void SpGEMMCpu::loadResultMatrix(matrix::Csr** E_ptr)
74+
{
75+
E_ptr_ = E_ptr;
76+
}
77+
78+
/**
79+
* @brief Computes the result of the SpGEMM operation
80+
*
81+
* Does not store partial results for reuse, as is this is not supported
82+
* by the CHOLMOD package.
83+
*/
84+
void SpGEMMCpu::compute()
85+
{
86+
cholmod_sparse* C_chol = cholmod_ssmult(B_, A_, 0, 1, 0, &Common_);
87+
cholmod_sparse* E_chol = cholmod_add(C_chol, D_, &alpha_, &beta_, 1, 0, &Common_);
88+
89+
if (!(*E_ptr_))
90+
{
91+
*E_ptr_ = new matrix::Csr((index_type) E_chol->nrow, (index_type) E_chol->ncol, (index_type) E_chol->nzmax);
92+
}
93+
else
94+
{
95+
(*E_ptr_)->destroyMatrixData(memory::HOST);
96+
}
97+
98+
// Previous data must be de-allocated and new data copied
99+
// Cholmod does not allow for reuse of arrays
100+
(*E_ptr_)->copyDataFrom(static_cast<index_type*>(E_chol->p),
101+
static_cast<index_type*>(E_chol->i),
102+
static_cast<real_type*>(E_chol->x),
103+
memory::HOST,
104+
memory::HOST);
105+
}
106+
107+
/**
108+
* @brief Allocates a CHOLMOD sparse matrix of the same size as the input CSR matrix
109+
* @param A[in] - Pointer to CSR matrix
110+
* @return Pointer to the allocated CHOLMOD sparse matrix
111+
*/
112+
cholmod_sparse* SpGEMMCpu::allocateCholmodType(matrix::Csr* A)
113+
{
114+
return cholmod_allocate_sparse((size_t) A->getNumRows(),
115+
(size_t) A->getNumColumns(),
116+
(size_t) A->getNnz(),
117+
1,
118+
1,
119+
0,
120+
CHOLMOD_REAL,
121+
&Common_);
122+
}
123+
124+
/**
125+
* Copies the data from a CSR matrix to a CHOLMOD sparse matrix
126+
* @param A[in] - Pointer to CSR matrix
127+
* @param A_chol[in] - Pointer to CHOLMOD sparse matrix
128+
*/
129+
void SpGEMMCpu::copyDataToCholmodType(matrix::Csr* A, cholmod_sparse* A_chol)
130+
{
131+
mem_.copyArrayHostToHost(
132+
static_cast<int*>(A_chol->p), A->getRowData(memory::HOST), A->getNumRows() + 1);
133+
mem_.copyArrayHostToHost(
134+
static_cast<int*>(A_chol->i), A->getColData(memory::HOST), A->getNnz());
135+
mem_.copyArrayHostToHost(
136+
static_cast<double*>(A_chol->x), A->getValues(memory::HOST), A->getNnz());
137+
}
138+
} // namespace hykkt
139+
} // namespace ReSolve

0 commit comments

Comments
 (0)