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-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 : 802 : 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 : 802 : std::unordered_set< std::size_t > >& nodecommmap ) :
76 : 802 : ConjugateGradients( std::move(std::get<0>(system)),
77 : 802 : std::move(std::get<1>(system)),
78 : 802 : std::move(std::get<2>(system)),
79 : 802 : gid, lid, nodecommmap ) {}
80 : :
81 : : #if defined(__clang__)
82 : : #pragma clang diagnostic pop
83 : : #endif
84 : :
85 : : //! Migrate constructor
86 : 1505 : explicit ConjugateGradients( CkMigrateMessage* m )
87 : 1505 : : 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 : 15562 : 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 : 677 : 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 : 4723 : void pup( PUP::er &p ) override {
168 : 4723 : p | m_A;
169 : 4723 : p | m_An;
170 : 4723 : p | m_x;
171 : 4723 : p | m_b;
172 : 4723 : p | m_pc;
173 : 4723 : p | m_gid;
174 : 4723 : p | m_lid;
175 : 4723 : p | m_nodeCommMap;
176 : 4723 : p | m_r;
177 : 4723 : p | m_z;
178 : 4723 : p | m_d;
179 : 4723 : p | m_rc;
180 : 4723 : p | m_nr;
181 : 4723 : p | m_na;
182 : 4723 : p | m_dirbc;
183 : 4723 : p | m_dirbcc;
184 : 4723 : p | m_nb;
185 : 4723 : p | m_p;
186 : 4723 : p | m_q;
187 : 4723 : p | m_qc;
188 : 4723 : p | m_nq;
189 : 4723 : p | m_nd;
190 : 4723 : p | m_initres;
191 : 4723 : p | m_solved;
192 : 4723 : p | m_normb;
193 : 4723 : p | m_it;
194 : 4723 : p | m_maxit;
195 : 4723 : p | m_finished;
196 : 4723 : p | m_verbose;
197 : 4723 : p | m_tol;
198 : 4723 : p | m_rho;
199 : 4723 : p | m_rho0;
200 : 4723 : p | m_alpha;
201 : 4723 : p | m_converged;
202 : 4723 : p | m_xc;
203 : 4723 : p | m_nx;
204 : 4723 : p | m_normr;
205 : 4723 : p | m_apply;
206 : 4723 : }
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::
|