Branch data Line data Source code
1 : : // *****************************************************************************
2 : : /*!
3 : : \file src/LinearSolver/ConjugateGradients.hpp
4 : : \copyright 2012-2015 J. Bakosi,
5 : : 2016-2018 Los Alamos National Security, LLC.,
6 : : 2019-2021 Triad National Security, LLC.,
7 : : 2022-2024 J. Bakosi
8 : : All rights reserved. See the LICENSE file for details.
9 : : \brief Charm++ chare array for distributed conjugate gradients
10 : : \details Charm++ chare array for asynchronous distributed
11 : : conjugate gradients linear solver.
12 : :
13 : : There are a potentially large number of ConjugateGradients Charm++ chares.
14 : : Each ConjugateGradient chare gets a chunk of the full load, due to partiting
15 : : the mesh, on which the solve is performed.
16 : :
17 : : The implementation uses the Charm++ runtime system and is fully
18 : : asynchronous, overlapping computation and communication. The algorithm
19 : : utilizes the structured dagger (SDAG) Charm++ functionality.
20 : : */
21 : : // *****************************************************************************
22 : : #pragma once
23 : :
24 : : #include "Types.hpp"
25 : : #include "CSR.hpp"
26 : :
27 : : #include "NoWarning/conjugategradients.decl.h"
28 : :
29 : : namespace tk {
30 : :
31 : : //! \brief ConjugateGradients Charm++ chare array used to perform a distributed
32 : : //! linear solve with the conjugate gradients algorithm
33 : : class ConjugateGradients : public CBase_ConjugateGradients {
34 : :
35 : : public:
36 : : #if defined(__clang__)
37 : : #pragma clang diagnostic push
38 : : #pragma clang diagnostic ignored "-Wunused-parameter"
39 : : #elif defined(STRICT_GNUC)
40 : : #pragma GCC diagnostic push
41 : : #pragma GCC diagnostic ignored "-Wunused-parameter"
42 : : #endif
43 : : // Include Charm++ SDAG code. See http://charm.cs.illinois.edu/manuals/html/
44 : : // charm++/manual.html, Sec. "Structured Control Flow: Structured Dagger".
45 : : ConjugateGradients_SDAG_CODE
46 : : #if defined(__clang__)
47 : : #pragma clang diagnostic pop
48 : : #elif defined(STRICT_GNUC)
49 : : #pragma GCC diagnostic pop
50 : : #endif
51 : :
52 : : //! Constructor
53 : : explicit ConjugateGradients(
54 : : const CSR& A,
55 : : const std::vector< tk::real >& x,
56 : : const std::vector< tk::real >& b,
57 : : const std::vector< std::size_t >& gid = {},
58 : : const std::unordered_map< std::size_t, std::size_t >& lid = {},
59 : : const std::unordered_map< int,
60 : : std::unordered_set< std::size_t > >& nodecommmap = {} );
61 : :
62 : : #if defined(__clang__)
63 : : #pragma clang diagnostic push
64 : : #pragma clang diagnostic ignored "-Wundefined-func-template"
65 : : #endif
66 : :
67 : : //! Constructor taking a tuple of {A,x,b} by rvalue reference
68 : 703 : explicit ConjugateGradients(
69 : : std::tuple< tk::CSR,
70 : : std::vector< tk::real >,
71 : : std::vector< tk::real > >&& system,
72 : : const std::vector< std::size_t >& gid,
73 : : const std::unordered_map< std::size_t, std::size_t >& lid,
74 : : const std::unordered_map< int,
75 : 703 : std::unordered_set< std::size_t > >& nodecommmap ) :
76 : 703 : ConjugateGradients( std::move(std::get<0>(system)),
77 : 703 : std::move(std::get<1>(system)),
78 : 703 : std::move(std::get<2>(system)),
79 : 703 : gid, lid, nodecommmap ) {}
80 : :
81 : : #if defined(__clang__)
82 : : #pragma clang diagnostic pop
83 : : #endif
84 : :
85 : : //! Migrate constructor
86 : 1235 : explicit ConjugateGradients( CkMigrateMessage* m )
87 : 1235 : : CBase_ConjugateGradients( m ) {}
88 : :
89 : : //! Solve linear system
90 : : void solve( std::size_t maxit,
91 : : tk::real tol,
92 : : int pe,
93 : : uint64_t verbose,
94 : : CkCallback c );
95 : :
96 : : //! Initialize linear solve: set initial guess and boundary conditions
97 : : void init( const std::vector< tk::real >& x,
98 : : const std::vector< tk::real >& b,
99 : : const std::vector< tk::real >& neubc,
100 : : const std::unordered_map< std::size_t,
101 : : std::vector< std::pair< int, tk::real > > >& dirbc,
102 : : const std::string& pc,
103 : : CkCallback cb );
104 : :
105 : : //! Setup solver
106 : : void setup( CkCallback c );
107 : :
108 : : //! Compute the norm of the right hand side
109 : : void normb( tk::real n );
110 : :
111 : : //! Compute rho = (r,z)
112 : : void rho( tk::real r );
113 : :
114 : : //! Receive contributions to r = b - A * x on chare-boundaries
115 : : void comres( const std::vector< std::size_t >& gid,
116 : : const std::vector< std::vector< tk::real > >& rc );
117 : :
118 : : //! Receive contributions to boundary conditions and rhs on chare-boundaries
119 : : void combc( const std::map< std::size_t,
120 : : std::vector< std::pair< int, tk::real > > >& dbc,
121 : : const std::vector< std::size_t >& gid,
122 : : const std::vector< std::vector< tk::real > >& qc );
123 : :
124 : : //! \brief Receive contributions to rhs with Dirichlet BCs applied on
125 : : //! chare-boundaries
126 : : void comr( const std::vector< std::size_t >& gid,
127 : : const std::vector< std::vector< tk::real > >& rc );
128 : :
129 : : //! Receive contributions to preconditioner chare-boundaries
130 : : void comd( const std::vector< std::size_t >& gid,
131 : : const std::vector< std::vector< tk::real > >& qc );
132 : :
133 : : //! Receive contributions to q = A * p on chare-boundaries
134 : : void comq( const std::vector< std::size_t >& gid,
135 : : const std::vector< std::vector< tk::real > >& qc );
136 : :
137 : : //! Receive contributions to final solution on chare-boundaries
138 : : void comx( const std::vector< std::size_t >& gid,
139 : : const std::vector< std::vector< tk::real > >& xc );
140 : :
141 : : //! Compute the dot product (p,q)
142 : : void pq( tk::real d );
143 : :
144 : : //! Compute the norm of the residual (r,r)
145 : : void normres( tk::real r );
146 : :
147 : : //! Compute the dot product (r,z)
148 : : void rz( tk::real rz );
149 : :
150 : : //! Access solution
151 : 12364 : const std::vector< tk::real >& solution() const { return m_x; }
152 : :
153 : : //! Return convergence flag
154 : : bool converged() const { return m_converged; }
155 : :
156 : : //! Return number of iterations taken
157 : 583 : std::size_t it() const { return m_it; }
158 : :
159 : : //! Non-const-ref access to lhs matrix
160 : 40 : tk::CSR& lhs() { return m_A; }
161 : :
162 : : /** @name Pack/unpack (Charm++ serialization) routines */
163 : : ///@{
164 : : //! \brief Pack/Unpack serialize member function
165 : : //! \param[in,out] p Charm++'s PUP::er serializer object reference
166 : 3814 : void pup( PUP::er &p ) override {
167 : 3814 : p | m_A;
168 : 3814 : p | m_An;
169 : 3814 : p | m_x;
170 : 3814 : p | m_b;
171 : 3814 : p | m_pc;
172 : 3814 : p | m_gid;
173 : 3814 : p | m_lid;
174 : 3814 : p | m_nodeCommMap;
175 : 3814 : p | m_r;
176 : 3814 : p | m_z;
177 : 3814 : p | m_d;
178 : 3814 : p | m_rc;
179 : 3814 : p | m_nr;
180 : 3814 : p | m_na;
181 : 3814 : p | m_dirbc;
182 : 3814 : p | m_dirbcc;
183 : 3814 : p | m_nb;
184 : 3814 : p | m_p;
185 : 3814 : p | m_q;
186 : 3814 : p | m_qc;
187 : 3814 : p | m_nq;
188 : 3814 : p | m_nd;
189 : 3814 : p | m_initres;
190 : 3814 : p | m_solved;
191 : 3814 : p | m_normb;
192 : 3814 : p | m_it;
193 : 3814 : p | m_maxit;
194 : 3814 : p | m_finished;
195 : 3814 : p | m_verbose;
196 : 3814 : p | m_tol;
197 : 3814 : p | m_rho;
198 : 3814 : p | m_rho0;
199 : 3814 : p | m_alpha;
200 : 3814 : p | m_converged;
201 : 3814 : p | m_xc;
202 : 3814 : p | m_nx;
203 : 3814 : p | m_normr;
204 : 3814 : }
205 : : //! \brief Pack/Unpack serialize operator|
206 : : //! \param[in,out] p Charm++'s PUP::er serializer object reference
207 : : //! \param[in,out] c ConjugateGradients object reference
208 : : friend void operator|( PUP::er& p, ConjugateGradients& c ) { c.pup(p); }
209 : : ///@}
210 : :
211 : : private:
212 : : //! Sparse matrix
213 : : CSR m_A;
214 : : //! Sparse matrix before boundary conditions
215 : : CSR m_An;
216 : : //! Solution/unknown
217 : : std::vector< tk::real > m_x;
218 : : //! Right hand side
219 : : std::vector< tk::real > m_b;
220 : : //! Preconditioner to use
221 : : std::string m_pc;
222 : : //! Global node IDs
223 : : std::vector< std::size_t > m_gid;
224 : : //! Local node IDs associated to global ones
225 : : std::unordered_map< std::size_t, std::size_t > m_lid;
226 : : //! Global mesh node IDs shared with other chares associated to chare IDs
227 : : std::unordered_map< int, std::unordered_set< std::size_t > > m_nodeCommMap;
228 : : //! Auxiliary vector for CG solve
229 : : std::vector< tk::real > m_r;
230 : : //! Receive buffer for communication of r = b - A * x
231 : : std::unordered_map< std::size_t, std::vector< tk::real > > m_rc;
232 : : //! Auxiliary vector for preconditioned CG solve
233 : : std::vector< tk::real > m_z;
234 : : //! Jacobi preconditioner
235 : : std::vector< tk::real > m_d;
236 : : //! Counter for assembling m_r
237 : : std::size_t m_nr;
238 : : //! Counter for assembling m_r (rhs with BCs applied)
239 : : std::size_t m_na;
240 : : //! Dirichlet boundary conditions
241 : : std::map< std::size_t, std::vector< std::pair<int,tk::real> > > m_dirbc;
242 : : //! Dirichlet boundary conditions communication buffer
243 : : std::map< std::size_t, std::vector< std::pair<int,tk::real> > > m_dirbcc;
244 : : //! Counter for assembling boundary conditions
245 : : std::size_t m_nb;
246 : : //! Auxiliary vector for CG solve
247 : : std::vector< tk::real > m_p;
248 : : //! Auxiliary vector for CG solve
249 : : std::vector< tk::real > m_q;
250 : : //! Receive buffer for communication of q = A * p
251 : : std::unordered_map< std::size_t, std::vector< tk::real > > m_qc;
252 : : //! Counter for assembling m_q
253 : : std::size_t m_nq;
254 : : //! Counter for assembling the preconditioner
255 : : std::size_t m_nd;
256 : : //! Charm++ callback to continue with when the setup is complete
257 : : CkCallback m_initres;
258 : : //! Charm++ callback to continue with when the solve is complete
259 : : CkCallback m_solved;
260 : : //! L2 norm of the right hand side
261 : : tk::real m_normb;
262 : : //! Iteration count
263 : : std::size_t m_it;
264 : : //! Max iteration count
265 : : std::size_t m_maxit;
266 : : //! True if finished
267 : : bool m_finished;
268 : : //! Verbose output
269 : : uint64_t m_verbose;
270 : : //! Stop tolerance
271 : : tk::real m_tol;
272 : : //! Helper scalar for CG algorithm
273 : : tk::real m_rho;
274 : : //! Helper scalar for CG algorithm
275 : : tk::real m_rho0;
276 : : //! Helper scalar for CG algorithm
277 : : tk::real m_alpha;
278 : : //! Convergence flag: true if linear smoother converged to tolerance
279 : : bool m_converged;
280 : : //! Receive buffer for solution
281 : : std::unordered_map< std::size_t, std::vector< tk::real > > m_xc;
282 : : //! Counter for assembling the solution on chare boundaries
283 : : std::size_t m_nx;
284 : : //! Norm of the residual
285 : : tk::real m_normr;
286 : :
287 : : //! Initiate computationa of dot product of two vectors
288 : : void dot( const std::vector< tk::real >& a,
289 : : const std::vector< tk::real >& b,
290 : : CkCallback c );
291 : :
292 : : //! Initiate A * x for computing the residual, r = b - A * x
293 : : void residual();
294 : :
295 : : //! Finish computing the initial residual, r = b - A * x
296 : : void initres();
297 : :
298 : : //! Setup preconditioner
299 : : void pc();
300 : :
301 : : //! Apply boundary conditions
302 : : void apply( CkCallback cb );
303 : :
304 : : //! Finish computing rhs with applied BCs
305 : : void r( CkCallback cb );
306 : :
307 : : //! Initiate computing q = A * p
308 : : void qAp();
309 : :
310 : : //! Finish computing q = A * p
311 : : void q();
312 : :
313 : : //! Start next linear solver iteration
314 : : void next();
315 : :
316 : : //! Assemble solution on chare boundaries and decide what's next
317 : : void x();
318 : : };
319 : :
320 : : } // tk::
|