Line data Source code
1 : // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2 : // SPDX-License-Identifier: Apache-2.0
3 :
4 : #include "vector.hpp"
5 :
6 : #include <cstdint>
7 : #include <random>
8 : #include <mfem/general/forall.hpp>
9 : #include "linalg/hypre.hpp"
10 : #include "utils/omp.hpp"
11 :
12 : namespace palace
13 : {
14 :
15 123918 : ComplexVector::ComplexVector(int size) : xr(size), xi(size) {}
16 :
17 0 : ComplexVector::ComplexVector(const ComplexVector &y) : ComplexVector(y.Size())
18 : {
19 0 : UseDevice(y.UseDevice());
20 0 : Set(y);
21 0 : }
22 :
23 0 : ComplexVector::ComplexVector(const Vector &yr, const Vector &yi) : ComplexVector(yr.Size())
24 : {
25 : MFEM_ASSERT(yr.Size() == yi.Size(),
26 : "Mismatch in dimension of real and imaginary parts in ComplexVector!");
27 0 : UseDevice(yr.UseDevice() || yi.UseDevice());
28 0 : Set(yr, yi);
29 0 : }
30 :
31 0 : ComplexVector::ComplexVector(const std::complex<double> *py, int size, bool on_dev)
32 0 : : ComplexVector(size)
33 : {
34 0 : Set(py, size, on_dev);
35 0 : }
36 :
37 0 : ComplexVector::ComplexVector(Vector &y, int offset, int size)
38 : {
39 0 : MakeRef(y, offset, size);
40 0 : }
41 :
42 61938 : void ComplexVector::UseDevice(bool use_dev)
43 : {
44 : xr.UseDevice(use_dev);
45 : xi.UseDevice(use_dev);
46 61938 : }
47 :
48 48 : void ComplexVector::SetSize(int size)
49 : {
50 48 : xr.SetSize(size);
51 48 : xi.SetSize(size);
52 48 : }
53 :
54 0 : void ComplexVector::MakeRef(Vector &y, int offset, int size)
55 : {
56 : MFEM_ASSERT(y.Size() >= offset + 2 * size,
57 : "Insufficient storage for ComplexVector alias reference of the given size!");
58 0 : y.ReadWrite(); // Ensure memory is allocated on device before aliasing
59 : xr.MakeRef(y, offset, size);
60 0 : xi.MakeRef(y, offset + size, size);
61 0 : }
62 :
63 0 : void ComplexVector::Set(const ComplexVector &y)
64 : {
65 : MFEM_ASSERT(y.Size() == Size(),
66 : "Mismatch in dimension of provided parts in ComplexVector!");
67 0 : Real() = y.Real();
68 0 : Imag() = y.Imag();
69 0 : }
70 :
71 0 : void ComplexVector::Set(const Vector &yr, const Vector &yi)
72 : {
73 : MFEM_ASSERT(yr.Size() == yi.Size() && yr.Size() == Size(),
74 : "Mismatch in dimension of real and imaginary parts in ComplexVector!");
75 0 : Real() = yr;
76 0 : Imag() = yi;
77 0 : }
78 :
79 3 : void ComplexVector::Set(const std::complex<double> *py, int size, bool on_dev)
80 : {
81 : MFEM_ASSERT(size == Size(),
82 : "Mismatch in dimension for array of std::complex<double> in ComplexVector!");
83 3 : auto SetImpl = [this](const double *Y, const int N, bool use_dev)
84 : {
85 3 : auto *XR = Real().Write(use_dev);
86 3 : auto *XI = Imag().Write(use_dev);
87 : mfem::forall_switch(use_dev, N,
88 3 : [=] MFEM_HOST_DEVICE(int i)
89 : {
90 6 : XR[i] = Y[2 * i];
91 6 : XI[i] = Y[2 * i + 1];
92 : });
93 3 : };
94 : const bool use_dev = UseDevice();
95 3 : if (((!use_dev || !mfem::Device::Allows(mfem::Backend::DEVICE_MASK)) && !on_dev) ||
96 0 : (use_dev && mfem::Device::Allows(mfem::Backend::DEVICE_MASK) && on_dev))
97 : {
98 : // No copy (host pointer and not using device, or device pointer and using device).
99 3 : SetImpl(reinterpret_cast<const double *>(py), size, use_dev);
100 : }
101 0 : else if (!on_dev)
102 : {
103 : // Need copy from host to device (host pointer but using device).
104 0 : Vector y(2 * size);
105 : y.UseDevice(true);
106 : {
107 : auto *Y = y.HostWrite();
108 0 : PalacePragmaOmp(parallel for schedule(static))
109 : for (int i = 0; i < size; i++)
110 : {
111 : Y[2 * i] = py[i].real();
112 : Y[2 * i + 1] = py[i].imag();
113 : }
114 : }
115 0 : SetImpl(y.Read(use_dev), size, use_dev);
116 : }
117 : else
118 : {
119 0 : MFEM_ABORT("ComplexVector::Set using a device pointer is not implemented when MFEM is "
120 : "not configured to use the device!");
121 : }
122 3 : }
123 :
124 0 : void ComplexVector::Get(std::complex<double> *py, int size, bool on_dev) const
125 : {
126 : MFEM_ASSERT(size == Size(),
127 : "Mismatch in dimension for array of std::complex<double> in ComplexVector!");
128 0 : auto GetImpl = [this](double *Y, const int N, bool use_dev)
129 : {
130 0 : const auto *XR = Real().Read(use_dev);
131 0 : const auto *XI = Imag().Read(use_dev);
132 : mfem::forall_switch(use_dev, N,
133 0 : [=] MFEM_HOST_DEVICE(int i)
134 : {
135 0 : Y[2 * i] = XR[i];
136 0 : Y[2 * i + 1] = XI[i];
137 : });
138 0 : };
139 : const bool use_dev = UseDevice();
140 0 : if (((!use_dev || !mfem::Device::Allows(mfem::Backend::DEVICE_MASK)) && !on_dev) ||
141 0 : (use_dev && mfem::Device::Allows(mfem::Backend::DEVICE_MASK) && on_dev))
142 : {
143 : // No copy (host pointer and not using device, or device pointer and using device).
144 0 : GetImpl(reinterpret_cast<double *>(py), size, use_dev);
145 : }
146 0 : else if (!on_dev)
147 : {
148 : // Need copy from device to host (host pointer but using device).
149 0 : const auto *XR = Real().HostRead();
150 0 : const auto *XI = Imag().HostRead();
151 0 : PalacePragmaOmp(parallel for schedule(static))
152 : for (int i = 0; i < size; i++)
153 : {
154 : py[i].real(XR[i]);
155 : py[i].imag(XI[i]);
156 : }
157 : }
158 : else
159 : {
160 0 : MFEM_ABORT("ComplexVector::Get using a device pointer is not implemented when MFEM is "
161 : "not configured to use the device!");
162 : }
163 0 : }
164 :
165 18 : ComplexVector &ComplexVector::operator=(std::complex<double> s)
166 : {
167 18 : Real() = s.real();
168 18 : Imag() = s.imag();
169 18 : return *this;
170 : }
171 :
172 0 : void ComplexVector::SetBlocks(const std::vector<const ComplexVector *> &y,
173 : const std::vector<std::complex<double>> &s)
174 : {
175 : MFEM_ASSERT(s.empty() || y.size() == s.size(),
176 : "Mismatch in dimension of vector blocks and scaling coefficients!");
177 0 : auto *XR = Real().Write();
178 0 : auto *XI = Imag().Write();
179 : int offset = 0;
180 0 : for (std::size_t b = 0; b < y.size(); b++)
181 : {
182 0 : MFEM_VERIFY(y[b] && ((b < y.size() - 1 && offset + y[b]->Size() < Size()) ||
183 : (b == y.size() - 1 && offset + y[b]->Size() == Size())),
184 : "Mistmatch between sum of block dimensions and parent vector dimension!");
185 0 : const double sr = s.empty() ? 1.0 : s[b].real();
186 0 : const double si = s.empty() ? 0.0 : s[b].imag();
187 0 : const bool use_dev = UseDevice() || y[b]->UseDevice();
188 : const int N = y[b]->Size();
189 0 : const auto *YR = y[b]->Real().Read();
190 0 : const auto *YI = y[b]->Imag().Read();
191 : mfem::forall_switch(use_dev, N,
192 0 : [=] MFEM_HOST_DEVICE(int i)
193 : {
194 0 : XR[i] = sr * YR[i] - si * YI[i];
195 0 : XI[i] = si * YR[i] + sr * YI[i];
196 0 : });
197 0 : XR += N;
198 0 : XI += N;
199 0 : offset += N;
200 : }
201 0 : }
202 :
203 0 : ComplexVector &ComplexVector::operator*=(std::complex<double> s)
204 : {
205 : const double sr = s.real();
206 : const double si = s.imag();
207 0 : if (si == 0.0)
208 : {
209 0 : Real() *= sr;
210 0 : Imag() *= sr;
211 : }
212 : else
213 : {
214 : const bool use_dev = UseDevice();
215 : const int N = Size();
216 0 : auto *XR = Real().ReadWrite(use_dev);
217 0 : auto *XI = Imag().ReadWrite(use_dev);
218 : mfem::forall_switch(use_dev, N,
219 0 : [=] MFEM_HOST_DEVICE(int i)
220 : {
221 0 : const auto t = si * XR[i] + sr * XI[i];
222 0 : XR[i] = sr * XR[i] - si * XI[i];
223 0 : XI[i] = t;
224 : });
225 : }
226 0 : return *this;
227 : }
228 :
229 0 : void ComplexVector::Conj()
230 : {
231 0 : Imag() *= -1.0;
232 0 : }
233 :
234 3 : void ComplexVector::Abs()
235 : {
236 : const bool use_dev = UseDevice();
237 : const int N = Size();
238 3 : auto *XR = Real().ReadWrite(use_dev);
239 3 : auto *XI = Imag().ReadWrite(use_dev);
240 : mfem::forall_switch(use_dev, N,
241 3 : [=] MFEM_HOST_DEVICE(int i)
242 : {
243 118 : XR[i] = std::sqrt(XR[i] * XR[i] + XI[i] * XI[i]);
244 118 : XI[i] = 0.0;
245 118 : });
246 3 : }
247 :
248 0 : void ComplexVector::Reciprocal()
249 : {
250 : const bool use_dev = UseDevice();
251 : const int N = Size();
252 0 : auto *XR = Real().ReadWrite(use_dev);
253 0 : auto *XI = Imag().ReadWrite(use_dev);
254 : mfem::forall_switch(use_dev, N,
255 0 : [=] MFEM_HOST_DEVICE(int i)
256 : {
257 0 : const auto s = 1.0 / (XR[i] * XR[i] + XI[i] * XI[i]);
258 0 : XR[i] *= s;
259 0 : XI[i] *= -s;
260 : });
261 0 : }
262 :
263 0 : std::complex<double> ComplexVector::Dot(const ComplexVector &y) const
264 : {
265 0 : return {(Real() * y.Real()) + (Imag() * y.Imag()),
266 0 : (this == &y) ? 0.0 : ((Imag() * y.Real()) - (Real() * y.Imag()))};
267 : }
268 :
269 0 : std::complex<double> ComplexVector::TransposeDot(const ComplexVector &y) const
270 : {
271 0 : return {(Real() * y.Real()) - (Imag() * y.Imag()),
272 0 : (this == &y) ? (2.0 * (Imag() * y.Real()))
273 0 : : ((Imag() * y.Real()) + (Real() * y.Imag()))};
274 : }
275 :
276 9 : void ComplexVector::AXPY(std::complex<double> alpha, const ComplexVector &x)
277 : {
278 9 : AXPY(alpha, x.Real(), x.Imag(), Real(), Imag());
279 9 : }
280 :
281 9 : void ComplexVector::AXPY(std::complex<double> alpha, const Vector &xr, const Vector &xi,
282 : Vector &yr, Vector &yi)
283 : {
284 9 : const bool use_dev = yr.UseDevice() || xr.UseDevice();
285 : const int N = yr.Size();
286 : const double ar = alpha.real();
287 : const double ai = alpha.imag();
288 9 : const auto *XR = xr.Read(use_dev);
289 9 : const auto *XI = xi.Read(use_dev);
290 9 : auto *YR = yr.ReadWrite(use_dev);
291 9 : auto *YI = yi.ReadWrite(use_dev);
292 9 : if (ai == 0.0)
293 : {
294 : mfem::forall_switch(use_dev, N,
295 3 : [=] MFEM_HOST_DEVICE(int i)
296 : {
297 118 : YR[i] += ar * XR[i];
298 118 : YI[i] += ar * XI[i];
299 118 : });
300 : }
301 : else
302 : {
303 : mfem::forall_switch(use_dev, N,
304 6 : [=] MFEM_HOST_DEVICE(int i)
305 : {
306 236 : const auto t = ai * XR[i] + ar * XI[i];
307 236 : YR[i] += ar * XR[i] - ai * XI[i];
308 236 : YI[i] += t;
309 236 : });
310 : }
311 9 : }
312 :
313 0 : void ComplexVector::AXPBY(std::complex<double> alpha, const ComplexVector &x,
314 : std::complex<double> beta)
315 : {
316 0 : AXPBY(alpha, x.Real(), x.Imag(), beta, Real(), Imag());
317 0 : }
318 :
319 0 : void ComplexVector::AXPBY(std::complex<double> alpha, const Vector &xr, const Vector &xi,
320 : std::complex<double> beta, Vector &yr, Vector &yi)
321 : {
322 0 : const bool use_dev = yr.UseDevice() || xr.UseDevice();
323 : const int N = yr.Size();
324 : const double ar = alpha.real();
325 : const double ai = alpha.imag();
326 0 : const auto *XR = xr.Read(use_dev);
327 0 : const auto *XI = xi.Read(use_dev);
328 : if (beta == 0.0)
329 : {
330 0 : auto *YR = yr.Write(use_dev);
331 0 : auto *YI = yi.Write(use_dev);
332 0 : if (ai == 0.0)
333 : {
334 : mfem::forall_switch(use_dev, N,
335 0 : [=] MFEM_HOST_DEVICE(int i)
336 : {
337 0 : YR[i] = ar * XR[i];
338 0 : YI[i] = ar * XI[i];
339 : });
340 : }
341 : else
342 : {
343 : mfem::forall_switch(use_dev, N,
344 0 : [=] MFEM_HOST_DEVICE(int i)
345 : {
346 0 : const auto t = ai * XR[i] + ar * XI[i];
347 0 : YR[i] = ar * XR[i] - ai * XI[i];
348 0 : YI[i] = t;
349 0 : });
350 : }
351 : }
352 : else
353 : {
354 : const double br = beta.real();
355 : const double bi = beta.imag();
356 0 : auto *YR = yr.ReadWrite(use_dev);
357 0 : auto *YI = yi.ReadWrite(use_dev);
358 0 : if (ai == 0.0 && bi == 0.0)
359 : {
360 : mfem::forall_switch(use_dev, N,
361 0 : [=] MFEM_HOST_DEVICE(int i)
362 : {
363 0 : YR[i] = ar * XR[i] + br * YR[i];
364 0 : YI[i] = ar * XI[i] + br * YI[i];
365 0 : });
366 0 : }
367 : else
368 : {
369 : mfem::forall_switch(use_dev, N,
370 0 : [=] MFEM_HOST_DEVICE(int i)
371 : {
372 0 : const auto t =
373 0 : ai * XR[i] + ar * XI[i] + bi * YR[i] + br * YI[i];
374 0 : YR[i] = ar * XR[i] - ai * XI[i] + br * YR[i] - bi * YI[i];
375 0 : YI[i] = t;
376 0 : });
377 : }
378 : }
379 0 : }
380 :
381 0 : void ComplexVector::AXPBYPCZ(std::complex<double> alpha, const ComplexVector &x,
382 : std::complex<double> beta, const ComplexVector &y,
383 : std::complex<double> gamma)
384 : {
385 0 : AXPBYPCZ(alpha, x.Real(), x.Imag(), beta, y.Real(), y.Imag(), gamma, Real(), Imag());
386 0 : }
387 :
388 0 : void ComplexVector::AXPBYPCZ(std::complex<double> alpha, const Vector &xr, const Vector &xi,
389 : std::complex<double> beta, const Vector &yr, const Vector &yi,
390 : std::complex<double> gamma, Vector &zr, Vector &zi)
391 : {
392 0 : const bool use_dev = zr.UseDevice() || xr.UseDevice() || yr.UseDevice();
393 : const int N = zr.Size();
394 : const double ar = alpha.real();
395 : const double ai = alpha.imag();
396 : const double br = beta.real();
397 : const double bi = beta.imag();
398 0 : const auto *XR = xr.Read(use_dev);
399 0 : const auto *XI = xi.Read(use_dev);
400 0 : const auto *YR = yr.Read(use_dev);
401 0 : const auto *YI = yi.Read(use_dev);
402 : if (gamma == 0.0)
403 : {
404 0 : auto *ZR = zr.Write(use_dev);
405 0 : auto *ZI = zi.Write(use_dev);
406 0 : if (ai == 0.0 && bi == 0.0)
407 : {
408 : mfem::forall_switch(use_dev, N,
409 0 : [=] MFEM_HOST_DEVICE(int i)
410 : {
411 0 : ZR[i] = ar * XR[i] + br * YR[i];
412 0 : ZI[i] = ar * XI[i] + br * YI[i];
413 0 : });
414 0 : }
415 : else
416 : {
417 : mfem::forall_switch(use_dev, N,
418 0 : [=] MFEM_HOST_DEVICE(int i)
419 : {
420 0 : const auto t =
421 0 : ai * XR[i] + ar * XI[i] + bi * YR[i] + br * YI[i];
422 0 : ZR[i] = ar * XR[i] - ai * XI[i] + br * YR[i] - bi * YI[i];
423 0 : ZI[i] = t;
424 0 : });
425 : }
426 : }
427 : else
428 : {
429 : const double gr = gamma.real();
430 : const double gi = gamma.imag();
431 0 : auto *ZR = zr.ReadWrite(use_dev);
432 0 : auto *ZI = zi.ReadWrite(use_dev);
433 0 : if (ai == 0.0 && bi == 0.0 && gi == 0.0)
434 : {
435 : mfem::forall_switch(use_dev, N,
436 0 : [=] MFEM_HOST_DEVICE(int i)
437 : {
438 0 : ZR[i] = ar * XR[i] + br * YR[i] + gr * ZR[i];
439 0 : ZI[i] = ar * XI[i] + br * YI[i] + gr * ZI[i];
440 0 : });
441 0 : }
442 : else
443 : {
444 : mfem::forall_switch(use_dev, N,
445 0 : [=] MFEM_HOST_DEVICE(int i)
446 : {
447 0 : const auto t = ai * XR[i] + ar * XI[i] + bi * YR[i] +
448 0 : br * YI[i] + gi * ZR[i] + gr * ZI[i];
449 0 : ZR[i] = ar * XR[i] - ai * XI[i] + br * YR[i] - bi * YI[i] +
450 0 : gr * ZR[i] - gi * ZI[i];
451 0 : ZI[i] = t;
452 0 : });
453 : }
454 : }
455 0 : }
456 :
457 : namespace linalg
458 : {
459 :
460 : template <>
461 0 : void SetSubVector(Vector &x, const mfem::Array<int> &rows, double s)
462 : {
463 0 : const bool use_dev = x.UseDevice();
464 : const int N = rows.Size();
465 : const double sr = s;
466 0 : const auto *idx = rows.Read(use_dev);
467 0 : auto *X = x.ReadWrite(use_dev);
468 : mfem::forall_switch(use_dev, N,
469 0 : [=] MFEM_HOST_DEVICE(int i)
470 : {
471 0 : const auto id = idx[i];
472 0 : X[id] = sr;
473 : });
474 0 : }
475 :
476 : template <>
477 0 : void SetSubVector(ComplexVector &x, const mfem::Array<int> &rows, double s)
478 : {
479 : const bool use_dev = x.UseDevice();
480 : const int N = rows.Size();
481 : const double sr = s;
482 0 : const auto *idx = rows.Read(use_dev);
483 0 : auto *XR = x.Real().ReadWrite(use_dev);
484 0 : auto *XI = x.Imag().ReadWrite(use_dev);
485 : mfem::forall_switch(use_dev, N,
486 0 : [=] MFEM_HOST_DEVICE(int i)
487 : {
488 0 : const int id = idx[i];
489 0 : XR[id] = sr;
490 0 : XI[id] = 0.0;
491 : });
492 0 : }
493 :
494 : template <>
495 0 : void SetSubVector(Vector &x, const mfem::Array<int> &rows, const Vector &y)
496 : {
497 0 : const bool use_dev = x.UseDevice();
498 : const int N = rows.Size();
499 0 : const auto *idx = rows.Read(use_dev);
500 0 : const auto *Y = y.Read(use_dev);
501 0 : auto *X = x.ReadWrite(use_dev);
502 : mfem::forall_switch(use_dev, N,
503 0 : [=] MFEM_HOST_DEVICE(int i)
504 : {
505 0 : const int id = idx[i];
506 0 : X[id] = Y[id];
507 : });
508 0 : }
509 :
510 : template <>
511 0 : void SetSubVector(ComplexVector &x, const mfem::Array<int> &rows, const ComplexVector &y)
512 : {
513 : const bool use_dev = x.UseDevice();
514 : const int N = rows.Size();
515 0 : const auto *idx = rows.Read(use_dev);
516 0 : const auto *YR = y.Real().Read(use_dev);
517 0 : const auto *YI = y.Imag().Read(use_dev);
518 0 : auto *XR = x.Real().ReadWrite(use_dev);
519 0 : auto *XI = x.Imag().ReadWrite(use_dev);
520 : mfem::forall_switch(use_dev, N,
521 0 : [=] MFEM_HOST_DEVICE(int i)
522 : {
523 0 : const int id = idx[i];
524 0 : XR[id] = YR[id];
525 0 : XI[id] = YI[id];
526 : });
527 0 : }
528 :
529 : template <>
530 0 : void SetSubVector(Vector &x, int start, const Vector &y)
531 : {
532 0 : const bool use_dev = x.UseDevice();
533 : const int N = y.Size();
534 : MFEM_ASSERT(start >= 0 && start + N <= x.Size(), "Invalid range for SetSubVector!");
535 0 : const auto *Y = y.Read(use_dev);
536 0 : auto *X = x.ReadWrite(use_dev);
537 : mfem::forall_switch(use_dev, N,
538 0 : [=] MFEM_HOST_DEVICE(int i)
539 : {
540 0 : const int id = start + i;
541 0 : X[id] = Y[i];
542 : });
543 0 : }
544 :
545 : template <>
546 0 : void SetSubVector(ComplexVector &x, int start, const ComplexVector &y)
547 : {
548 : const bool use_dev = x.UseDevice();
549 : const int N = y.Size();
550 : MFEM_ASSERT(start >= 0 && start + N <= x.Size(), "Invalid range for SetSubVector!");
551 0 : const auto *YR = y.Real().Read(use_dev);
552 0 : const auto *YI = y.Imag().Read(use_dev);
553 0 : auto *XR = x.Real().ReadWrite(use_dev);
554 0 : auto *XI = x.Imag().ReadWrite(use_dev);
555 : mfem::forall_switch(use_dev, N,
556 0 : [=] MFEM_HOST_DEVICE(int i)
557 : {
558 0 : const int id = start + i;
559 0 : XR[id] = YR[i];
560 0 : XI[id] = YI[i];
561 : });
562 0 : }
563 :
564 : template <>
565 0 : void SetSubVector(Vector &x, int start, int end, double s)
566 : {
567 0 : const bool use_dev = x.UseDevice();
568 : MFEM_ASSERT(start >= 0 && end <= x.Size() && start <= end,
569 : "Invalid range for SetSubVector!");
570 0 : const int N = end - start;
571 : const double sr = s;
572 0 : auto *X = x.ReadWrite(use_dev) + start;
573 0 : mfem::forall_switch(use_dev, N, [=] MFEM_HOST_DEVICE(int i) { X[i] = sr; });
574 0 : }
575 :
576 : template <>
577 0 : void SetSubVector(ComplexVector &x, int start, int end, double s)
578 : {
579 : const bool use_dev = x.UseDevice();
580 : MFEM_ASSERT(start >= 0 && end <= x.Size() && start <= end,
581 : "Invalid range for SetSubVector!");
582 0 : const int N = end - start;
583 : const double sr = s;
584 0 : auto *XR = x.Real().ReadWrite(use_dev) + start;
585 0 : auto *XI = x.Imag().ReadWrite(use_dev) + start;
586 : mfem::forall_switch(use_dev, N,
587 0 : [=] MFEM_HOST_DEVICE(int i)
588 : {
589 0 : XR[i] = sr;
590 0 : XI[i] = 0.0;
591 : });
592 0 : }
593 :
594 : template <>
595 0 : void SetRandom(MPI_Comm comm, Vector &x, int seed)
596 : {
597 0 : if (seed == 0)
598 : {
599 : std::vector<std::uint32_t> seeds(1);
600 0 : std::seed_seq seed_gen{Mpi::Rank(comm)};
601 0 : seed_gen.generate(seeds.begin(), seeds.end());
602 0 : seed = static_cast<int>(seeds[0]);
603 : }
604 0 : x.Randomize(seed); // On host always
605 0 : }
606 :
607 : template <>
608 0 : void SetRandomReal(MPI_Comm comm, Vector &x, int seed)
609 : {
610 0 : SetRandom(comm, x, seed);
611 0 : }
612 :
613 : template <>
614 0 : void SetRandomSign(MPI_Comm comm, Vector &x, int seed)
615 : {
616 0 : SetRandom(comm, x, seed);
617 0 : const bool use_dev = x.UseDevice();
618 : const int N = x.Size();
619 0 : auto *X = x.ReadWrite(use_dev);
620 0 : mfem::forall_switch(use_dev, N, [=] MFEM_HOST_DEVICE(int i)
621 0 : { X[i] = (X[i] > 0.0) ? 1.0 : ((X[i] < 0.0) ? -1.0 : 0.0); });
622 0 : }
623 :
624 : template <>
625 0 : void SetRandom(MPI_Comm comm, ComplexVector &x, int seed)
626 : {
627 0 : if (seed == 0)
628 : {
629 0 : std::vector<std::uint32_t> seeds(2);
630 0 : std::seed_seq seed_gen{2 * Mpi::Rank(comm), 2 * Mpi::Rank(comm) + 1};
631 0 : seed_gen.generate(seeds.begin(), seeds.end());
632 0 : SetRandom(comm, x.Real(), static_cast<int>(seeds[0]));
633 0 : SetRandom(comm, x.Imag(), static_cast<int>(seeds[1]));
634 : }
635 : else
636 : {
637 0 : SetRandom(comm, x.Real(), seed);
638 0 : SetRandom(comm, x.Imag(), seed);
639 : }
640 0 : }
641 :
642 : template <>
643 0 : void SetRandomReal(MPI_Comm comm, ComplexVector &x, int seed)
644 : {
645 0 : SetRandom(comm, x.Real(), seed);
646 0 : x.Imag() = 0.0;
647 0 : }
648 :
649 : template <>
650 0 : void SetRandomSign(MPI_Comm comm, ComplexVector &x, int seed)
651 : {
652 0 : SetRandom(comm, x, seed);
653 : const bool use_dev = x.UseDevice();
654 : const int N = x.Size();
655 0 : auto *XR = x.Real().ReadWrite(use_dev);
656 0 : auto *XI = x.Imag().ReadWrite(use_dev);
657 : mfem::forall_switch(use_dev, N,
658 0 : [=] MFEM_HOST_DEVICE(int i)
659 : {
660 0 : XR[i] = (XR[i] > 0.0) ? 1.0 : ((XR[i] < 0.0) ? -1.0 : 0.0);
661 0 : XI[i] = (XI[i] > 0.0) ? 1.0 : ((XI[i] < 0.0) ? -1.0 : 0.0);
662 : });
663 0 : }
664 :
665 36 : double LocalDot(const Vector &x, const Vector &y)
666 : {
667 42 : static hypre::HypreVector X, Y;
668 : MFEM_ASSERT(x.Size() == y.Size(), "Size mismatch for vector inner product!");
669 36 : X.Update(x);
670 36 : Y.Update(y);
671 36 : return hypre_SeqVectorInnerProd(X, Y);
672 : }
673 :
674 0 : std::complex<double> LocalDot(const ComplexVector &x, const ComplexVector &y)
675 : {
676 0 : if (&x == &y)
677 : {
678 0 : return {LocalDot(x.Real(), y.Real()) + LocalDot(x.Imag(), y.Imag()), 0.0};
679 : }
680 : else
681 : {
682 0 : return {LocalDot(x.Real(), y.Real()) + LocalDot(x.Imag(), y.Imag()),
683 0 : LocalDot(x.Imag(), y.Real()) - LocalDot(x.Real(), y.Imag())};
684 : }
685 : }
686 :
687 : // We implement LocalSum using Hypre instead of using MFEM's Sum because it is
688 : // more efficient on GPUs. TODO: Verify this
689 11 : double LocalSum(const Vector &x)
690 : {
691 19 : static hypre::HypreVector X;
692 11 : X.Update(x);
693 11 : return hypre_SeqVectorSumElts(X);
694 : }
695 :
696 3 : std::complex<double> LocalSum(const ComplexVector &x)
697 : {
698 3 : return {LocalSum(x.Real()), LocalSum(x.Imag())};
699 : }
700 :
701 : template <>
702 0 : void AXPY(double alpha, const Vector &x, Vector &y)
703 : {
704 0 : if (alpha == 1.0)
705 : {
706 0 : y += x;
707 : }
708 : else
709 : {
710 0 : y.Add(alpha, x);
711 : }
712 0 : }
713 :
714 : template <>
715 0 : void AXPY(double alpha, const ComplexVector &x, ComplexVector &y)
716 : {
717 0 : y.AXPY(alpha, x);
718 0 : }
719 :
720 : template <>
721 0 : void AXPY(std::complex<double> alpha, const ComplexVector &x, ComplexVector &y)
722 : {
723 0 : y.AXPY(alpha, x);
724 0 : }
725 :
726 : template <>
727 0 : void AXPBY(double alpha, const Vector &x, double beta, Vector &y)
728 : {
729 0 : add(alpha, x, beta, y, y);
730 0 : }
731 :
732 : template <>
733 0 : void AXPBY(std::complex<double> alpha, const ComplexVector &x, std::complex<double> beta,
734 : ComplexVector &y)
735 : {
736 0 : y.AXPBY(alpha, x, beta);
737 0 : }
738 :
739 : template <>
740 0 : void AXPBY(double alpha, const ComplexVector &x, double beta, ComplexVector &y)
741 : {
742 0 : y.AXPBY(alpha, x, beta);
743 0 : }
744 :
745 : template <>
746 0 : void AXPBYPCZ(double alpha, const Vector &x, double beta, const Vector &y, double gamma,
747 : Vector &z)
748 : {
749 0 : if (gamma == 0.0)
750 : {
751 0 : add(alpha, x, beta, y, z);
752 : }
753 : else
754 : {
755 0 : AXPBY(alpha, x, gamma, z);
756 0 : z.Add(beta, y);
757 : }
758 0 : }
759 :
760 : template <>
761 0 : void AXPBYPCZ(std::complex<double> alpha, const ComplexVector &x, std::complex<double> beta,
762 : const ComplexVector &y, std::complex<double> gamma, ComplexVector &z)
763 : {
764 0 : z.AXPBYPCZ(alpha, x, beta, y, gamma);
765 0 : }
766 :
767 : template <>
768 0 : void AXPBYPCZ(double alpha, const ComplexVector &x, double beta, const ComplexVector &y,
769 : double gamma, ComplexVector &z)
770 : {
771 0 : z.AXPBYPCZ(alpha, x, beta, y, gamma);
772 0 : }
773 :
774 2 : void Sqrt(Vector &x, double s)
775 : {
776 2 : const bool use_dev = x.UseDevice();
777 : const int N = x.Size();
778 2 : auto *X = x.ReadWrite(use_dev);
779 : mfem::forall_switch(use_dev, N,
780 9 : [=] MFEM_HOST_DEVICE(int i) { X[i] = std::sqrt(X[i] * s); });
781 2 : }
782 :
783 : } // namespace linalg
784 :
785 : } // namespace palace
|
