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-2025 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 804 : 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 804 : std::unordered_set< std::size_t > >& nodecommmap ) :
76 804 : ConjugateGradients( std::move(std::get<0>(system)),
77 804 : std::move(std::get<1>(system)),
78 804 : std::move(std::get<2>(system)),
79 804 : gid, lid, nodecommmap ) {}
80 :
81 : #if defined(__clang__)
82 : #pragma clang diagnostic pop
83 : #endif
84 :
85 : //! Migrate constructor
86 1323 : explicit ConjugateGradients( CkMigrateMessage* m )
87 1323 : : 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 : bool apply,
103 : const std::string& pc,
104 : CkCallback cb );
105 :
106 : //! Setup solver
107 : void setup( CkCallback c );
108 :
109 : //! Compute the norm of the right hand side
110 : void normb( tk::real n );
111 :
112 : //! Compute rho = (r,z)
113 : void rho( tk::real r );
114 :
115 : //! Receive contributions to r = b - A * x on chare-boundaries
116 : void comres( const std::vector< std::size_t >& gid,
117 : const std::vector< std::vector< tk::real > >& rc );
118 :
119 : //! Receive contributions to boundary conditions and rhs on chare-boundaries
120 : void combc( const std::map< std::size_t,
121 : std::vector< std::pair< int, tk::real > > >& dbc,
122 : const std::vector< std::size_t >& gid,
123 : const std::vector< std::vector< tk::real > >& qc );
124 :
125 : //! \brief Receive contributions to rhs with Dirichlet BCs applied on
126 : //! chare-boundaries
127 : void comr( const std::vector< std::size_t >& gid,
128 : const std::vector< std::vector< tk::real > >& rc );
129 :
130 : //! Receive contributions to preconditioner chare-boundaries
131 : void comd( const std::vector< std::size_t >& gid,
132 : const std::vector< std::vector< tk::real > >& qc );
133 :
134 : //! Receive contributions to q = A * p on chare-boundaries
135 : void comq( const std::vector< std::size_t >& gid,
136 : const std::vector< std::vector< tk::real > >& qc );
137 :
138 : //! Receive contributions to final solution on chare-boundaries
139 : void comx( const std::vector< std::size_t >& gid,
140 : const std::vector< std::vector< tk::real > >& xc );
141 :
142 : //! Compute the dot product (p,q)
143 : void pq( tk::real d );
144 :
145 : //! Compute the norm of the residual (r,r)
146 : void normres( tk::real r );
147 :
148 : //! Compute the dot product (r,z)
149 : void rz( tk::real rz );
150 :
151 : //! Access solution
152 15686 : const std::vector< tk::real >& solution() const { return m_x; }
153 :
154 : //! Return convergence flag
155 : bool converged() const { return m_converged; }
156 :
157 : //! Return number of iterations taken
158 721 : std::size_t it() const { return m_it; }
159 :
160 : //! Non-const-ref access to lhs matrix
161 40 : tk::CSR& lhs() { return m_A; }
162 :
163 : /** @name Pack/unpack (Charm++ serialization) routines */
164 : ///@{
165 : //! \brief Pack/Unpack serialize member function
166 : //! \param[in,out] p Charm++'s PUP::er serializer object reference
167 4179 : void pup( PUP::er &p ) override {
168 4179 : p | m_A;
169 4179 : p | m_An;
170 4179 : p | m_x;
171 4179 : p | m_b;
172 4179 : p | m_pc;
173 4179 : p | m_gid;
174 4179 : p | m_lid;
175 4179 : p | m_nodeCommMap;
176 4179 : p | m_r;
177 4179 : p | m_z;
178 4179 : p | m_d;
179 4179 : p | m_rc;
180 4179 : p | m_nr;
181 4179 : p | m_na;
182 4179 : p | m_dirbc;
183 4179 : p | m_dirbcc;
184 4179 : p | m_nb;
185 4179 : p | m_p;
186 4179 : p | m_q;
187 4179 : p | m_qc;
188 4179 : p | m_nq;
189 4179 : p | m_nd;
190 4179 : p | m_initres;
191 4179 : p | m_solved;
192 4179 : p | m_normb;
193 4179 : p | m_it;
194 4179 : p | m_maxit;
195 4179 : p | m_finished;
196 4179 : p | m_verbose;
197 4179 : p | m_tol;
198 4179 : p | m_rho;
199 4179 : p | m_rho0;
200 4179 : p | m_alpha;
201 4179 : p | m_converged;
202 4179 : p | m_xc;
203 4179 : p | m_nx;
204 4179 : p | m_normr;
205 4179 : p | m_apply;
206 4179 : }
207 : //! \brief Pack/Unpack serialize operator|
208 : //! \param[in,out] p Charm++'s PUP::er serializer object reference
209 : //! \param[in,out] c ConjugateGradients object reference
210 : friend void operator|( PUP::er& p, ConjugateGradients& c ) { c.pup(p); }
211 : ///@}
212 :
213 : private:
214 : //! Sparse matrix
215 : CSR m_A;
216 : //! Sparse matrix before boundary conditions
217 : CSR m_An;
218 : //! Solution/unknown
219 : std::vector< tk::real > m_x;
220 : //! Right hand side
221 : std::vector< tk::real > m_b;
222 : //! Preconditioner to use
223 : std::string m_pc;
224 : //! Global node IDs
225 : std::vector< std::size_t > m_gid;
226 : //! Local node IDs associated to global ones
227 : std::unordered_map< std::size_t, std::size_t > m_lid;
228 : //! Global mesh node IDs shared with other chares associated to chare IDs
229 : std::unordered_map< int, std::unordered_set< std::size_t > > m_nodeCommMap;
230 : //! Auxiliary vector for CG solve
231 : std::vector< tk::real > m_r;
232 : //! Receive buffer for communication of r = b - A * x
233 : std::unordered_map< std::size_t, std::vector< tk::real > > m_rc;
234 : //! Auxiliary vector for preconditioned CG solve
235 : std::vector< tk::real > m_z;
236 : //! Jacobi preconditioner
237 : std::vector< tk::real > m_d;
238 : //! Counter for assembling m_r
239 : std::size_t m_nr;
240 : //! Counter for assembling m_r (rhs with BCs applied)
241 : std::size_t m_na;
242 : //! Dirichlet boundary conditions
243 : std::map< std::size_t, std::vector< std::pair<int,tk::real> > > m_dirbc;
244 : //! Dirichlet boundary conditions communication buffer
245 : std::map< std::size_t, std::vector< std::pair<int,tk::real> > > m_dirbcc;
246 : //! Counter for assembling boundary conditions
247 : std::size_t m_nb;
248 : //! Auxiliary vector for CG solve
249 : std::vector< tk::real > m_p;
250 : //! Auxiliary vector for CG solve
251 : std::vector< tk::real > m_q;
252 : //! Receive buffer for communication of q = A * p
253 : std::unordered_map< std::size_t, std::vector< tk::real > > m_qc;
254 : //! Counter for assembling m_q
255 : std::size_t m_nq;
256 : //! Counter for assembling the preconditioner
257 : std::size_t m_nd;
258 : //! Charm++ callback to continue with when the setup is complete
259 : CkCallback m_initres;
260 : //! Charm++ callback to continue with when the solve is complete
261 : CkCallback m_solved;
262 : //! L2 norm of the right hand side
263 : tk::real m_normb;
264 : //! Iteration count
265 : std::size_t m_it;
266 : //! Max iteration count
267 : std::size_t m_maxit;
268 : //! True if finished
269 : bool m_finished;
270 : //! Verbose output
271 : uint64_t m_verbose;
272 : //! Stop tolerance
273 : tk::real m_tol;
274 : //! Helper scalar for CG algorithm
275 : tk::real m_rho;
276 : //! Helper scalar for CG algorithm
277 : tk::real m_rho0;
278 : //! Helper scalar for CG algorithm
279 : tk::real m_alpha;
280 : //! Convergence flag: true if linear smoother converged to tolerance
281 : bool m_converged;
282 : //! Receive buffer for solution
283 : std::unordered_map< std::size_t, std::vector< tk::real > > m_xc;
284 : //! Counter for assembling the solution on chare boundaries
285 : std::size_t m_nx;
286 : //! Norm of the residual
287 : tk::real m_normr;
288 : //! True to apply boundary conditions
289 : bool m_apply;
290 :
291 : //! Initiate computationa of dot product of two vectors
292 : void dot( const std::vector< tk::real >& a,
293 : const std::vector< tk::real >& b,
294 : CkCallback c );
295 :
296 : //! Initiate A * x for computing the residual, r = b - A * x
297 : void residual();
298 :
299 : //! Finish computing the initial residual, r = b - A * x
300 : void initres();
301 :
302 : //! Setup preconditioner
303 : void pc();
304 :
305 : //! Apply boundary conditions
306 : void apply( CkCallback cb );
307 :
308 : //! Finish computing rhs with applied BCs
309 : void r( CkCallback cb );
310 :
311 : //! Initiate computing q = A * p
312 : void qAp();
313 :
314 : //! Finish computing q = A * p
315 : void q();
316 :
317 : //! Start next linear solver iteration
318 : void next();
319 :
320 : //! Assemble solution on chare boundaries and decide what's next
321 : void x();
322 : };
323 :
324 : } // tk::
|
