Loading...
Searching...
No Matches
gpu.hpp
Go to the documentation of this file.
1#pragma once
2
3/**
4 * GPUPlanewaveSubspaceManager — GPU implementation of subspace operations.
5 *
6 * Implements PlanewaveSubspaceManager using GPU BLAS for the large matrix
7 * operations (subspace Hamiltonian build, Ritz vector rotation, orthogonalization).
8 * Subspace diagonalization always runs on CPU via GPUEigensolver (which
9 * internally uses LAPACK — subspace matrices are small N×N).
10 *
11 * Constructor takes GPUBLAS& and GPUEigensolver& for delegation, plus
12 * the execution policy for band-group collectives.
13 *
14 * No CUDA headers. No Metal headers. Uses only vendor-neutral GPU abstractions.
15 */
16
17#include <algorithm>
18#include <cmath>
19#include <complex>
20#include <cstring>
21#include <memory>
22#include <vector>
23
24#include "io/logger.hpp"
25#include "layer_0_core/execution/gpu/device/init.hpp"
26#include "layer_0_core/execution/gpu/device/internal/buffer.hpp"
27#include "layer_0_core/execution/init.hpp"
28#include "layer_0_core/math/linear_algebra/blas.hpp"
29#include "layer_0_core/math/linear_algebra/gpu/blas/init.hpp"
30#include "layer_0_core/math/tolerances.hpp"
31#include "layer_0_core/types/matrix.hpp"
32#include "layer_4_computation/diagonalizer/planewave/internal/hamiltonian_applicator/gpu_kernels.hpp"
33#include "layer_4_computation/diagonalizer/planewave/internal/overlap/init.hpp"
34#include "layer_4_computation/diagonalizer/planewave/internal/subspace_manager/init.hpp"
35#include "layer_4_computation/diagonalizer/planewave/subspace/factory.hpp"
36
37namespace q3 {
38
39class GPUPlanewaveSubspaceManager : public PlanewaveSubspaceManager {
40 public:
41 /**
42 * Construct with GPU math wrappers and execution policy.
43 *
44 * @param gpu_blas GPU BLAS wrapper for GEMM operations.
45 * @param execution_policy Execution policy for band-group collectives.
46 * @param device_backend GPU device backend for buffer allocation.
47 */
48 GPUPlanewaveSubspaceManager(
49 execution::gpu::blas::GPUBLAS& gpu_blas,
50 const execution::ExecutionPolicy& execution_policy,
51 std::shared_ptr<execution::gpu::device::GPUDeviceBackend> device_backend)
52 : gpu_blas_{gpu_blas}
53 , execution_policy_{execution_policy}
54 , device_backend_{device_backend}
55 , subspace_diagonalizer_{SubspaceDiagonalizerFactory::create(execution_policy)} {}
56
57 /**
58 * S-orthogonalize and normalize candidates using GPU BLAS.
59 *
60 * Uses GPU GEMM for computing overlap coefficients and applying
61 * projections. Falls back to modified Gram-Schmidt if Cholesky fails.
62 */
63 void set_gamma_origin_is_local(bool origin_is_local) override {
64 gamma_origin_is_local_ = origin_is_local;
65 }
66
67 int orthogonalize_batch(int starting_vector_index,
68 int number_of_candidates,
69 int number_of_plane_waves,
70 bool is_gamma_only,
71 const OverlapOperator* overlap_operator,
72 Complex* basis_vectors,
73 Complex* overlap_basis_vectors) override {
74 if (number_of_candidates == 0) {
75 return 0;
76 }
77
78 int leading_dimension = std::max(number_of_plane_waves, 1);
79
80 // Apply overlap operator to candidates.
81 overlap_operator->apply_batch(
82 number_of_plane_waves,
83 number_of_candidates,
84 &basis_vectors[starting_vector_index * number_of_plane_waves],
85 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves]);
86
87 int number_of_existing = starting_vector_index;
88
89 // Project out existing basis via GPU GEMM.
90 if (number_of_existing > 0 && !is_gamma_only) {
91 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
92
93 // Compute overlap coefficients: coeffs = V† · S·candidates
94 std::vector<Complex> overlap_coefficients(number_of_existing * number_of_candidates,
95 Complex(0.0, 0.0));
96 execution::gpu::device::GPUDeviceBuffer<Complex> device_overlap_coefficients(
97 number_of_existing * number_of_candidates, device_backend_);
98
99 gpu_blas_.gemm('C',
100 'N',
101 number_of_existing,
102 number_of_candidates,
103 number_of_plane_waves,
104 alpha_one,
105 basis_vectors,
106 leading_dimension,
107 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves],
108 leading_dimension,
109 beta_zero,
110 device_overlap_coefficients.data(),
111 number_of_existing);
112
113 device_backend_->synchronize();
114
115 std::memcpy(overlap_coefficients.data(),
116 device_overlap_coefficients.data(),
117 overlap_coefficients.size() * sizeof(Complex));
118
119 execution_policy_.reduce_within_group(overlap_coefficients.data(),
120 number_of_existing * number_of_candidates);
121
122 std::memcpy(device_overlap_coefficients.data(),
123 overlap_coefficients.data(),
124 overlap_coefficients.size() * sizeof(Complex));
125
126 // Subtract projections: candidates -= V · coeffs
127 Complex alpha_neg_one(-1.0, 0.0), beta_one(1.0, 0.0);
128 gpu_blas_.gemm('N',
129 'N',
130 number_of_plane_waves,
131 number_of_candidates,
132 number_of_existing,
133 alpha_neg_one,
134 basis_vectors,
135 leading_dimension,
136 device_overlap_coefficients.data(),
137 number_of_existing,
138 beta_one,
139 &basis_vectors[starting_vector_index * number_of_plane_waves],
140 leading_dimension);
141
142 // Same for overlap vectors: S·candidates -= S·V · coeffs
143 gpu_blas_.gemm('N',
144 'N',
145 number_of_plane_waves,
146 number_of_candidates,
147 number_of_existing,
148 alpha_neg_one,
149 overlap_basis_vectors,
150 leading_dimension,
151 device_overlap_coefficients.data(),
152 number_of_existing,
153 beta_one,
154 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves],
155 leading_dimension);
156
157 device_backend_->synchronize();
158 } else if (number_of_existing > 0 && is_gamma_only) {
159 // Gamma-only: project out existing basis using GPU DGEMM.
160 int double_number_of_plane_waves = 2 * number_of_plane_waves;
161 int double_leading_dimension = 2 * leading_dimension;
162 double alpha_two = 2.0;
163 double d_beta_zero = 0.0;
164
165 auto* real_basis = reinterpret_cast<double*>(const_cast<Complex*>(basis_vectors));
166 auto* real_candidates_sv = reinterpret_cast<double*>(
167 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves]);
168
169 std::vector<double> overlap_coefficients(number_of_existing * number_of_candidates,
170 0.0);
171 execution::gpu::device::GPUDeviceBuffer<double> device_overlap_coefficients(
172 number_of_existing * number_of_candidates, device_backend_);
173
174 // coeffs = V† · S·candidates
175 gpu_blas_.dgemm('T',
176 'N',
177 number_of_existing,
178 number_of_candidates,
179 double_number_of_plane_waves,
180 alpha_two,
181 real_basis,
182 double_leading_dimension,
183 real_candidates_sv,
184 double_leading_dimension,
185 d_beta_zero,
186 device_overlap_coefficients.data(),
187 number_of_existing);
188
189 device_backend_->synchronize();
190
191 std::memcpy(overlap_coefficients.data(),
192 device_overlap_coefficients.data(),
193 overlap_coefficients.size() * sizeof(double));
194
195 // Subtract the double-counted G=0 component, where this rank holds it.
196 if (gamma_origin_is_local_) {
197 for (int candidate_index = 0; candidate_index < number_of_candidates;
198 ++candidate_index) {
199 for (int basis_index = 0; basis_index < number_of_existing; ++basis_index) {
200 Complex basis_g0 = basis_vectors[basis_index * number_of_plane_waves];
201 Complex sv_g0 =
202 overlap_basis_vectors[(starting_vector_index + candidate_index) *
203 number_of_plane_waves];
204 overlap_coefficients[candidate_index * number_of_existing + basis_index] -=
205 (std::conj(basis_g0) * sv_g0).real();
206 }
207 }
208 }
209
210 execution_policy_.reduce_within_group(overlap_coefficients.data(),
211 number_of_existing * number_of_candidates);
212
213 std::memcpy(device_overlap_coefficients.data(),
214 overlap_coefficients.data(),
215 overlap_coefficients.size() * sizeof(double));
216
217 auto* real_candidates = reinterpret_cast<double*>(
218 &basis_vectors[starting_vector_index * number_of_plane_waves]);
219 auto* real_overlap_basis = reinterpret_cast<double*>(overlap_basis_vectors);
220
221 double alpha_neg_one = -1.0;
222 double beta_one = 1.0;
223
224 // candidates -= V · coeffs
225 gpu_blas_.dgemm('N',
226 'N',
227 double_number_of_plane_waves,
228 number_of_candidates,
229 number_of_existing,
230 alpha_neg_one,
231 real_basis,
232 double_leading_dimension,
233 device_overlap_coefficients.data(),
234 number_of_existing,
235 beta_one,
236 real_candidates,
237 double_leading_dimension);
238
239 // S·candidates -= S·V · coeffs
240 gpu_blas_.dgemm('N',
241 'N',
242 double_number_of_plane_waves,
243 number_of_candidates,
244 number_of_existing,
245 alpha_neg_one,
246 real_overlap_basis,
247 double_leading_dimension,
248 device_overlap_coefficients.data(),
249 number_of_existing,
250 beta_one,
251 real_candidates_sv,
252 double_leading_dimension);
253
254 device_backend_->synchronize();
255 }
256
257 // Cholesky QR for inter-candidate orthogonalization.
258 if (!is_gamma_only) {
259 auto* cast_candidates = reinterpret_cast<std::complex<double>*>(
260 &basis_vectors[starting_vector_index * number_of_plane_waves]);
261 auto* cast_candidates_sv = reinterpret_cast<std::complex<double>*>(
262 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves]);
263
264 std::vector<Complex> cholesky_factor(number_of_candidates * number_of_candidates,
265 Complex(0.0, 0.0));
266 execution::gpu::device::GPUDeviceBuffer<Complex> device_cholesky_factor(
267 number_of_candidates * number_of_candidates, device_backend_);
268 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
269
270 gpu_blas_.gemm('C',
271 'N',
272 number_of_candidates,
273 number_of_candidates,
274 number_of_plane_waves,
275 alpha_one,
276 &basis_vectors[starting_vector_index * number_of_plane_waves],
277 leading_dimension,
278 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves],
279 leading_dimension,
280 beta_zero,
281 device_cholesky_factor.data(),
282 number_of_candidates);
283
284 device_backend_->synchronize();
285
286 std::memcpy(cholesky_factor.data(),
287 device_cholesky_factor.data(),
288 cholesky_factor.size() * sizeof(Complex));
289
290 execution_policy_.reduce_within_group(cholesky_factor.data(),
291 number_of_candidates * number_of_candidates);
292
293 int info = 0;
294 auto* cast_cholesky_factor =
295 reinterpret_cast<std::complex<double>*>(cholesky_factor.data());
296 q3::blas::zpotrf_(
297 "U", &number_of_candidates, cast_cholesky_factor, &number_of_candidates, &info);
298 if (info == 0) {
299 Complex z_alpha_one(1.0, 0.0);
300 auto* cast_alpha_one = reinterpret_cast<std::complex<double>*>(&z_alpha_one);
301 q3::blas::ztrsm_("R",
302 "U",
303 "N",
304 "N",
305 &number_of_plane_waves,
306 &number_of_candidates,
307 cast_alpha_one,
308 cast_cholesky_factor,
309 &number_of_candidates,
310 cast_candidates,
311 &leading_dimension);
312 q3::blas::ztrsm_("R",
313 "U",
314 "N",
315 "N",
316 &number_of_plane_waves,
317 &number_of_candidates,
318 cast_alpha_one,
319 cast_cholesky_factor,
320 &number_of_candidates,
321 cast_candidates_sv,
322 &leading_dimension);
323 return number_of_candidates;
324 }
325 } else {
326 // Gamma-only Cholesky: same as serial.
327 int double_number_of_plane_waves = 2 * number_of_plane_waves;
328 int double_leading_dimension = 2 * leading_dimension;
329 auto* real_candidates = reinterpret_cast<double*>(
330 &basis_vectors[starting_vector_index * number_of_plane_waves]);
331 auto* real_candidates_sv = reinterpret_cast<double*>(
332 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves]);
333
334 std::vector<double> cholesky_factor(number_of_candidates * number_of_candidates, 0.0);
335 execution::gpu::device::GPUDeviceBuffer<double> device_cholesky_factor(
336 number_of_candidates * number_of_candidates, device_backend_);
337 double alpha_two = 2.0;
338 double beta_zero = 0.0;
339 gpu_blas_.dgemm('T',
340 'N',
341 number_of_candidates,
342 number_of_candidates,
343 double_number_of_plane_waves,
344 alpha_two,
345 real_candidates,
346 double_leading_dimension,
347 real_candidates_sv,
348 double_leading_dimension,
349 beta_zero,
350 device_cholesky_factor.data(),
351 number_of_candidates);
352
353 device_backend_->synchronize();
354
355 std::memcpy(cholesky_factor.data(),
356 device_cholesky_factor.data(),
357 cholesky_factor.size() * sizeof(double));
358
359 if (gamma_origin_is_local_) {
360 for (int column_index = 0; column_index < number_of_candidates; ++column_index) {
361 for (int row_index = 0; row_index < number_of_candidates; ++row_index) {
362 Complex vi_g0 = basis_vectors[(starting_vector_index + row_index) *
363 number_of_plane_waves];
364 Complex svj_g0 =
365 overlap_basis_vectors[(starting_vector_index + column_index) *
366 number_of_plane_waves];
367 cholesky_factor[column_index * number_of_candidates + row_index] -=
368 (std::conj(vi_g0) * svj_g0).real();
369 }
370 }
371 }
372
373 execution_policy_.reduce_within_group(cholesky_factor.data(),
374 number_of_candidates * number_of_candidates);
375
376 int info = 0;
377 q3::blas::dpotrf_(
378 "U", &number_of_candidates, cholesky_factor.data(), &number_of_candidates, &info);
379 if (info == 0) {
380 double d_alpha_one = 1.0;
381 q3::blas::dtrsm_("R",
382 "U",
383 "N",
384 "N",
385 &double_number_of_plane_waves,
386 &number_of_candidates,
387 &d_alpha_one,
388 cholesky_factor.data(),
389 &number_of_candidates,
390 real_candidates,
391 &double_leading_dimension);
392 q3::blas::dtrsm_("R",
393 "U",
394 "N",
395 "N",
396 &double_number_of_plane_waves,
397 &number_of_candidates,
398 &d_alpha_one,
399 cholesky_factor.data(),
400 &number_of_candidates,
401 real_candidates_sv,
402 &double_leading_dimension);
403 return number_of_candidates;
404 }
405 }
406
407 // Fallback: Modified Gram-Schmidt (same as serial path).
408 int accepted_count = 0;
409 int one_increment = 1;
410 int double_leading_dimension = 2 * leading_dimension;
411 std::vector<double> batch_coefficients_buffer(number_of_candidates, 0.0);
412 std::vector<Complex> z_batch_coefficients_buffer(number_of_candidates, Complex(0.0, 0.0));
413
414 for (int candidate_index = 0; candidate_index < number_of_candidates; ++candidate_index) {
415 int source_index = starting_vector_index + candidate_index;
416 int destination_index = starting_vector_index + accepted_count;
417
418 if (destination_index != source_index) {
419 std::memcpy(&basis_vectors[destination_index * number_of_plane_waves],
420 &basis_vectors[source_index * number_of_plane_waves],
421 number_of_plane_waves * sizeof(Complex));
422 std::memcpy(&overlap_basis_vectors[destination_index * number_of_plane_waves],
423 &overlap_basis_vectors[source_index * number_of_plane_waves],
424 number_of_plane_waves * sizeof(Complex));
425 }
426
427 auto* current_vector = reinterpret_cast<std::complex<double>*>(
428 &basis_vectors[destination_index * number_of_plane_waves]);
429 auto* current_overlap = reinterpret_cast<std::complex<double>*>(
430 &overlap_basis_vectors[destination_index * number_of_plane_waves]);
431
432 int number_of_batch_existing = destination_index - starting_vector_index;
433 if (number_of_batch_existing > 0) {
434 if (is_gamma_only) {
435 int double_number_of_plane_waves = 2 * number_of_plane_waves;
436 double alpha_two = 2.0;
437 double d_beta_zero = 0.0;
438 int one_col = 1;
439
440 auto* real_batch_basis = reinterpret_cast<double*>(
441 &basis_vectors[starting_vector_index * number_of_plane_waves]);
442 auto* real_sv = reinterpret_cast<double*>(current_overlap);
443
444 q3::blas::dgemm_("T",
445 "N",
446 &number_of_batch_existing,
447 &one_col,
448 &double_number_of_plane_waves,
449 &alpha_two,
450 real_batch_basis,
451 &double_leading_dimension,
452 real_sv,
453 &double_leading_dimension,
454 &d_beta_zero,
455 batch_coefficients_buffer.data(),
456 &number_of_batch_existing);
457
458 if (gamma_origin_is_local_) {
459 for (int basis_index = 0; basis_index < number_of_batch_existing;
460 ++basis_index) {
461 Complex basis_g0 = basis_vectors[(starting_vector_index + basis_index) *
462 number_of_plane_waves];
463 Complex sv_g0 =
464 overlap_basis_vectors[destination_index * number_of_plane_waves];
465 batch_coefficients_buffer[basis_index] -=
466 (std::conj(basis_g0) * sv_g0).real();
467 }
468 }
469
470 execution_policy_.reduce_within_group(batch_coefficients_buffer.data(),
471 number_of_batch_existing);
472
473 double alpha_neg_one = -1.0;
474 double beta_one = 1.0;
475 q3::blas::dgemv_("N",
476 &double_number_of_plane_waves,
477 &number_of_batch_existing,
478 &alpha_neg_one,
479 real_batch_basis,
480 &double_leading_dimension,
481 batch_coefficients_buffer.data(),
482 &one_increment,
483 &beta_one,
484 reinterpret_cast<double*>(current_vector),
485 &one_increment);
486
487 auto* real_batch_overlap = reinterpret_cast<double*>(
488 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves]);
489 q3::blas::dgemv_("N",
490 &double_number_of_plane_waves,
491 &number_of_batch_existing,
492 &alpha_neg_one,
493 real_batch_overlap,
494 &double_leading_dimension,
495 batch_coefficients_buffer.data(),
496 &one_increment,
497 &beta_one,
498 reinterpret_cast<double*>(current_overlap),
499 &one_increment);
500 } else {
501 int one_col = 1;
502 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
503 auto* cast_batch_basis = reinterpret_cast<std::complex<double>*>(
504 &basis_vectors[starting_vector_index * number_of_plane_waves]);
505 auto* cast_alpha = reinterpret_cast<std::complex<double>*>(&alpha_one);
506 auto* cast_beta = reinterpret_cast<std::complex<double>*>(&beta_zero);
507 auto* cast_coeffs =
508 reinterpret_cast<std::complex<double>*>(z_batch_coefficients_buffer.data());
509
510 q3::blas::zgemm_("C",
511 "N",
512 &number_of_batch_existing,
513 &one_col,
514 &number_of_plane_waves,
515 cast_alpha,
516 cast_batch_basis,
517 &leading_dimension,
518 current_overlap,
519 &leading_dimension,
520 cast_beta,
521 cast_coeffs,
522 &number_of_batch_existing);
523
524 execution_policy_.reduce_within_group(z_batch_coefficients_buffer.data(),
525 number_of_batch_existing);
526
527 Complex alpha_neg_one(-1.0, 0.0), beta_one(1.0, 0.0);
528 auto* cast_alpha_neg = reinterpret_cast<std::complex<double>*>(&alpha_neg_one);
529 auto* cast_beta_one = reinterpret_cast<std::complex<double>*>(&beta_one);
530 q3::blas::zgemv_("N",
531 &number_of_plane_waves,
532 &number_of_batch_existing,
533 cast_alpha_neg,
534 cast_batch_basis,
535 &leading_dimension,
536 cast_coeffs,
537 &one_increment,
538 cast_beta_one,
539 current_vector,
540 &one_increment);
541
542 auto* cast_batch_overlap = reinterpret_cast<std::complex<double>*>(
543 &overlap_basis_vectors[starting_vector_index * number_of_plane_waves]);
544 q3::blas::zgemv_("N",
545 &number_of_plane_waves,
546 &number_of_batch_existing,
547 cast_alpha_neg,
548 cast_batch_overlap,
549 &leading_dimension,
550 cast_coeffs,
551 &one_increment,
552 cast_beta_one,
553 current_overlap,
554 &one_increment);
555 }
556 }
557
558 Real norm_squared = gamma_aware_real_dot(
559 is_gamma_only, number_of_plane_waves, current_vector, current_overlap);
560
561 if (norm_squared > math::tolerances::iterative_solver_norm_threshold) {
562 double inverse_norm = 1.0 / std::sqrt(norm_squared);
563 q3::blas::zdscal_(
564 &number_of_plane_waves, &inverse_norm, current_vector, &one_increment);
565 q3::blas::zdscal_(
566 &number_of_plane_waves, &inverse_norm, current_overlap, &one_increment);
567 accepted_count++;
568 }
569 }
570 return accepted_count;
571 }
572
573 /**
574 * Incrementally update subspace matrices via GPU GEMM.
575 */
576 void update_subspace_matrices(int start_column,
577 int new_columns,
578 int total_basis_size,
579 int maximum_basis_size,
580 int number_of_plane_waves,
581 bool is_gamma_only,
582 const Complex* basis_vectors,
583 const Complex* hamiltonian_basis_vectors,
584 const Complex* overlap_basis_vectors,
585 Complex* subspace_hamiltonian,
586 Complex* subspace_overlap) override {
587 int leading_dimension = std::max(number_of_plane_waves, 1);
588
589 if (is_gamma_only) {
590 // Gamma-only: use GPU DGEMM.
591 int double_number_of_plane_waves = 2 * number_of_plane_waves;
592 int double_leading_dimension = 2 * leading_dimension;
593 double alpha_two = 2.0;
594 double d_beta_zero = 0.0;
595
596 auto* real_basis = reinterpret_cast<const double*>(basis_vectors);
597 auto* real_h_new = reinterpret_cast<const double*>(
598 &hamiltonian_basis_vectors[start_column * number_of_plane_waves]);
599 auto* real_s_new = reinterpret_cast<const double*>(
600 &overlap_basis_vectors[start_column * number_of_plane_waves]);
601
602 std::vector<double> h_block_real(static_cast<size_t>(total_basis_size) * new_columns,
603 0.0);
604 std::vector<double> s_block_real(static_cast<size_t>(total_basis_size) * new_columns,
605 0.0);
606
607 execution::gpu::device::GPUDeviceBuffer<double> device_h_block_real(
608 total_basis_size * new_columns, device_backend_);
609 execution::gpu::device::GPUDeviceBuffer<double> device_s_block_real(
610 total_basis_size * new_columns, device_backend_);
611
612 gpu_blas_.dgemm('T',
613 'N',
614 total_basis_size,
615 new_columns,
616 double_number_of_plane_waves,
617 alpha_two,
618 real_basis,
619 double_leading_dimension,
620 real_h_new,
621 double_leading_dimension,
622 d_beta_zero,
623 device_h_block_real.data(),
624 total_basis_size);
625
626 gpu_blas_.dgemm('T',
627 'N',
628 total_basis_size,
629 new_columns,
630 double_number_of_plane_waves,
631 alpha_two,
632 real_basis,
633 double_leading_dimension,
634 real_s_new,
635 double_leading_dimension,
636 d_beta_zero,
637 device_s_block_real.data(),
638 total_basis_size);
639
640 device_backend_->synchronize();
641
642 std::memcpy(h_block_real.data(),
643 device_h_block_real.data(),
644 h_block_real.size() * sizeof(double));
645 std::memcpy(s_block_real.data(),
646 device_s_block_real.data(),
647 s_block_real.size() * sizeof(double));
648
649 for (int col = 0; col < new_columns; ++col) {
650 int source_col = start_column + col;
651 Complex h_right_g0 = hamiltonian_basis_vectors[source_col * number_of_plane_waves];
652 Complex s_right_g0 = overlap_basis_vectors[source_col * number_of_plane_waves];
653 for (int row = 0; row < total_basis_size; ++row) {
654 Complex left_g0 = basis_vectors[row * number_of_plane_waves];
655 int real_index = col * total_basis_size + row;
656 if (gamma_origin_is_local_) {
657 h_block_real[real_index] -= (std::conj(left_g0) * h_right_g0).real();
658 s_block_real[real_index] -= (std::conj(left_g0) * s_right_g0).real();
659 }
660
661 int sub_index = (start_column + col) * maximum_basis_size + row;
662 subspace_hamiltonian[sub_index] = Complex(h_block_real[real_index], 0.0);
663 subspace_overlap[sub_index] = Complex(s_block_real[real_index], 0.0);
664 }
665 }
666 } else {
667 // Non-gamma: use GPU GEMM.
668 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
669
670 gpu_blas_.gemm('C',
671 'N',
672 total_basis_size,
673 new_columns,
674 number_of_plane_waves,
675 alpha_one,
676 basis_vectors,
677 leading_dimension,
678 &hamiltonian_basis_vectors[start_column * number_of_plane_waves],
679 leading_dimension,
680 beta_zero,
681 &subspace_hamiltonian[start_column * maximum_basis_size],
682 maximum_basis_size);
683
684 gpu_blas_.gemm('C',
685 'N',
686 total_basis_size,
687 new_columns,
688 number_of_plane_waves,
689 alpha_one,
690 basis_vectors,
691 leading_dimension,
692 &overlap_basis_vectors[start_column * number_of_plane_waves],
693 leading_dimension,
694 beta_zero,
695 &subspace_overlap[start_column * maximum_basis_size],
696 maximum_basis_size);
697
698 device_backend_->synchronize();
699 }
700
701 int new_block_size = maximum_basis_size * new_columns;
702 int new_block_offset = start_column * maximum_basis_size;
703 execution_policy_.reduce_within_group(&subspace_hamiltonian[new_block_offset],
704 new_block_size);
705 execution_policy_.reduce_within_group(&subspace_overlap[new_block_offset], new_block_size);
706 }
707
708 /**
709 * Diagonalize subspace via CPU LAPACK (small matrix — GPU not beneficial).
710 */
711 int diagonalize_subspace(int basis_size,
712 int maximum_basis_size,
713 bool is_gamma_only,
714 const Complex* subspace_hamiltonian,
715 const Complex* subspace_overlap,
716 Real* eigenvalues,
717 Complex* eigenvector_matrix) override {
718 /**
719 * The solver overwrites both inputs, and the caller's matrices carry
720 * the next iteration's incremental update, so it works on copies. The
721 * Hamiltonian's copy is the eigenvector buffer itself -- the solver
722 * returns the eigenvectors where the matrix was -- and only the overlap
723 * needs a buffer of its own. Both are compacted from the caller's
724 * leading dimension to basis_size. That is two matrices of basis_size
725 * squared here rather than three, on every rank.
726 */
727 diagonalize_overlap_copy_workspace_.resize(static_cast<std::size_t>(basis_size) *
728 basis_size);
729 for (int column = 0; column < basis_size; ++column) {
730 for (int row = 0; row < basis_size; ++row) {
731 eigenvector_matrix[column * basis_size + row] =
732 subspace_hamiltonian[column * maximum_basis_size + row];
733 diagonalize_overlap_copy_workspace_[column * basis_size + row] =
734 subspace_overlap[column * maximum_basis_size + row];
735 }
736 }
737
738 /**
739 * Which ranks run the eigensolver, in three cases:
740 *
741 * - a collective solver (ScaLAPACK) must be entered by EVERY rank of
742 * the group — gating on the root would deadlock — and returns
743 * redundant results, so no broadcast follows;
744 * - replicated coefficient vectors (group_size() <= 1) mean every rank
745 * already holds the same matrices, so each solves its own copy and
746 * there is no group to broadcast over. This includes the fallback
747 * that disables band distribution AFTER this manager was built, where
748 * group_rank_index() keeps its pool-local value and is not
749 * necessarily 0;
750 * - distributed vectors with a non-collective solver: the group root
751 * solves and broadcasts.
752 *
753 * The two conditions must agree. Testing group_rank_index() == 0 to
754 * solve but group_size() > 1 to broadcast leaves the non-root ranks of a
755 * replicated group with neither a result nor a broadcast, and they then
756 * diverge from the root at the next collective.
757 */
758 const bool all_ranks_solve = subspace_diagonalizer_->is_collective();
759 const bool vectors_are_replicated = execution_policy_.group_size() <= 1;
760 const bool this_rank_solves =
761 all_ranks_solve || vectors_are_replicated || execution_policy_.group_rank_index() == 0;
762 int info = 0;
763 if (this_rank_solves) {
764 info = subspace_diagonalizer_->solve(basis_size,
765 eigenvector_matrix,
766 diagonalize_overlap_copy_workspace_.data(),
767 eigenvalues,
768 eigenvector_matrix,
769 is_gamma_only);
770 }
771
772 if (!all_ranks_solve && !vectors_are_replicated) {
773 execution_policy_.broadcast_from_group_root(eigenvector_matrix,
774 basis_size * basis_size);
775 execution_policy_.broadcast_from_group_root(eigenvalues, basis_size);
776 Real info_as_real = static_cast<Real>(info);
777 execution_policy_.broadcast_from_group_root(&info_as_real, 1);
778 info = static_cast<int>(info_as_real);
779 }
780
781 if (info != 0) {
782 log::warn(
783 "LAPACK eigensolver returned info={} in GPU Davidson subspace (basis_size={})",
784 info,
785 basis_size);
786 }
787 return info;
788 }
789
790 /**
791 * Collapse subspace via GPU GEMM rotation.
792 */
793 int build_correction_vectors(int basis_size,
794 int number_of_plane_waves,
795 int number_of_bands,
796 int maximum_candidates,
797 bool is_gamma_only,
798 const Complex* eigenvector_matrix,
799 int eigenvector_leading_dimension,
800 const Real* eigenvalues,
801 const std::vector<bool>& band_converged,
802 const Real* kinetic_energy_squared,
803 const Complex* hamiltonian_basis_vectors,
804 const Complex* overlap_basis_vectors,
805 Complex* correction_columns) override {
806 /**
807 * Same formulation as the CPU path: R = (HV)c - eps (SV)c, two GEMMs
808 * writing straight into the free basis columns, so no Ritz vectors and
809 * no residual array are held. The device path is single-process, so
810 * there is no band partition to apply.
811 */
812 std::vector<int> source_bands;
813 source_bands.reserve(static_cast<std::size_t>(number_of_bands));
814 for (int band_index = 0; band_index < number_of_bands; ++band_index) {
815 if (!band_converged[band_index] &&
816 static_cast<int>(source_bands.size()) < maximum_candidates) {
817 source_bands.push_back(band_index);
818 }
819 }
820 int number_of_candidates = static_cast<int>(source_bands.size());
821 if (number_of_candidates == 0) {
822 return 0;
823 }
824
825 execution::gpu::device::GPUDeviceBuffer<Complex> device_coefficients(
826 static_cast<std::size_t>(basis_size) * static_cast<std::size_t>(number_of_candidates),
827 device_backend_);
828 std::vector<Real> candidate_eigenvalues(static_cast<std::size_t>(number_of_candidates),
829 0.0);
830 for (int candidate_index = 0; candidate_index < number_of_candidates; ++candidate_index) {
831 int band_index = source_bands[static_cast<std::size_t>(candidate_index)];
832 candidate_eigenvalues[static_cast<std::size_t>(candidate_index)] =
833 eigenvalues[band_index];
834 std::memcpy(
835 device_coefficients.data() + static_cast<std::size_t>(candidate_index) * basis_size,
836 eigenvector_matrix + static_cast<std::size_t>(band_index) *
837 static_cast<std::size_t>(eigenvector_leading_dimension),
838 static_cast<std::size_t>(basis_size) * sizeof(Complex));
839 }
840
841 int leading_dimension = std::max(number_of_plane_waves, 1);
842 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0), beta_one(1.0, 0.0);
843
844 gpu_blas_.gemm('N',
845 'N',
846 number_of_plane_waves,
847 number_of_candidates,
848 basis_size,
849 alpha_one,
850 overlap_basis_vectors,
851 leading_dimension,
852 device_coefficients.data(),
853 basis_size,
854 beta_zero,
855 correction_columns,
856 leading_dimension);
857 device_backend_->synchronize();
858
859 for (int candidate_index = 0; candidate_index < number_of_candidates; ++candidate_index) {
860 Complex negative_eigenvalue(
861 -candidate_eigenvalues[static_cast<std::size_t>(candidate_index)], 0.0);
862 gpu_blas_.scale(number_of_plane_waves,
863 negative_eigenvalue,
864 correction_columns +
865 static_cast<std::size_t>(candidate_index) * number_of_plane_waves);
866 }
867 device_backend_->synchronize();
868
869 gpu_blas_.gemm('N',
870 'N',
871 number_of_plane_waves,
872 number_of_candidates,
873 basis_size,
874 alpha_one,
875 hamiltonian_basis_vectors,
876 leading_dimension,
877 device_coefficients.data(),
878 basis_size,
879 beta_one,
880 correction_columns,
881 leading_dimension);
882 device_backend_->synchronize();
883
884 /**
885 * The preconditioner kernel dereferences the eigenvalues on the
886 * device, so stage them there rather than hand over a host vector.
887 */
888 execution::gpu::device::GPUDeviceBuffer<Real> device_candidate_eigenvalues(
889 static_cast<std::size_t>(number_of_candidates), device_backend_);
890 std::memcpy(device_candidate_eigenvalues.data(),
891 candidate_eigenvalues.data(),
892 static_cast<std::size_t>(number_of_candidates) * sizeof(Real));
893 std::vector<bool> none_converged(static_cast<std::size_t>(number_of_candidates), false);
894 apply_preconditioner(number_of_candidates,
895 number_of_plane_waves,
896 kinetic_energy_squared,
897 device_candidate_eigenvalues.data(),
898 none_converged,
899 correction_columns);
900
901 enforce_gamma_g0_real(
902 is_gamma_only, number_of_plane_waves, number_of_candidates, correction_columns);
903 return number_of_candidates;
904 }
905
906 int collapse_subspace(int basis_size,
907 int number_of_plane_waves,
908 int number_of_bands,
909 int maximum_basis_size,
910 bool is_gamma_only,
911 bool is_identity_overlap,
912 const OverlapOperator* overlap_operator,
913 const Complex* subspace_hamiltonian,
914 const Complex* subspace_overlap,
915 Complex* basis_vectors,
916 Complex* hamiltonian_basis_vectors,
917 Complex* overlap_basis_vectors,
918 Complex* rotation_scratch) override {
919 std::vector<Real> eigenvalues_expanded(maximum_basis_size, 0.0);
920 std::vector<Complex> eigenvector_matrix(maximum_basis_size * maximum_basis_size);
921 /**
922 * As on the CPU path: the caller's subspace is complete through
923 * basis_size, so diagonalise it; rebuild from the vectors only as the
924 * fallback for the recovery collapse.
925 */
926 int info = -1;
927 if (subspace_hamiltonian != nullptr && subspace_overlap != nullptr) {
928 info = diagonalize_subspace(basis_size,
929 maximum_basis_size,
930 is_gamma_only,
931 subspace_hamiltonian,
932 subspace_overlap,
933 eigenvalues_expanded.data(),
934 eigenvector_matrix.data());
935 }
936 if (info != 0) {
937 build_subspace_matrices_and_diagonalize(basis_size,
938 number_of_plane_waves,
939 is_gamma_only,
940 overlap_operator,
941 basis_vectors,
942 hamiltonian_basis_vectors,
943 overlap_basis_vectors,
944 eigenvalues_expanded.data(),
945 eigenvector_matrix.data());
946 }
947
948 int leading_dimension = std::max(number_of_plane_waves, 1);
949 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
950 /**
951 * As on the CPU path: only number_of_bands columns are rotated and
952 * copied back, and the destination is the caller's scratch rather than
953 * three more device buffers the size of the whole basis.
954 */
955 std::size_t collapsed_size = static_cast<std::size_t>(number_of_bands) *
956 static_cast<std::size_t>(number_of_plane_waves);
957
958 // Rotate via GPU GEMM.
959 execution::gpu::device::GPUDeviceBuffer<Complex> device_eigenvector_matrix(
960 basis_size * basis_size, device_backend_);
961 std::memcpy(device_eigenvector_matrix.data(),
962 eigenvector_matrix.data(),
963 basis_size * basis_size * sizeof(Complex));
964
965 /**
966 * One scratch buffer serves all three rotations: each reads a
967 * different array than it writes, and the overlap needs no rotation at
968 * all when it aliases the basis.
969 */
970 auto rotate_into_scratch = [&](const Complex* source) {
971 gpu_blas_.gemm('N',
972 'N',
973 number_of_plane_waves,
974 number_of_bands,
975 basis_size,
976 alpha_one,
977 source,
978 leading_dimension,
979 device_eigenvector_matrix.data(),
980 basis_size,
981 beta_zero,
982 rotation_scratch,
983 leading_dimension);
984 device_backend_->synchronize();
985 };
986
987 rotate_into_scratch(basis_vectors);
988 std::memcpy(basis_vectors, rotation_scratch, collapsed_size * sizeof(Complex));
989
990 rotate_into_scratch(hamiltonian_basis_vectors);
991 std::memcpy(hamiltonian_basis_vectors, rotation_scratch, collapsed_size * sizeof(Complex));
992
993 /**
994 * Skip the overlap rotation only when the overlap array is the basis
995 * itself. An identity overlap is not enough: the CPU workspace aliases
996 * S|V> to |V> for it, so rotating the basis has already rotated S|V>,
997 * but the device workspace keeps S|V> in a buffer of its own. Keying
998 * the skip on the flag left that buffer unrotated after every collapse,
999 * so the next orthogonalisation and subspace update read stale S|V>,
1000 * and a run that collapsed often converged to the wrong answer -- 1.77
1001 * Ry low on the two-atom silicon deck at expansion factor 2, which
1002 * collapses every iteration. Factor 4 hid it: small decks converge
1003 * before the basis fills.
1004 */
1005 const bool overlap_needs_rotation =
1006 !is_identity_overlap || overlap_basis_vectors != basis_vectors;
1007 if (overlap_needs_rotation) {
1008 rotate_into_scratch(overlap_basis_vectors);
1009 std::memcpy(overlap_basis_vectors, rotation_scratch, collapsed_size * sizeof(Complex));
1010 }
1011
1012 enforce_gamma_g0_real(is_gamma_only, number_of_plane_waves, number_of_bands, basis_vectors);
1013 enforce_gamma_g0_real(
1014 is_gamma_only, number_of_plane_waves, number_of_bands, hamiltonian_basis_vectors);
1015 enforce_gamma_g0_real(
1016 is_gamma_only, number_of_plane_waves, number_of_bands, overlap_basis_vectors);
1017
1018 return number_of_bands;
1019 }
1020
1021 /**
1022 * Extract eigenvectors using GPU GEMM rotation.
1023 */
1024 void extract_eigenvectors(int basis_size,
1025 int number_of_plane_waves,
1026 int number_of_bands,
1027 int kpoint_index_state,
1028 bool is_gamma_only,
1029 const Complex* eigenvector_matrix,
1030 const Complex* basis_vectors,
1031 const Real* eigenvalues,
1032 Complex* rotated_scratch,
1033 PlanewaveKohnShamQuantumSystem& quantum_system) override {
1034 int leading_dimension = std::max(number_of_plane_waves, 1);
1035
1036 /**
1037 * The rotation lands in the caller's scratch, which on this path is a
1038 * device buffer already sized for it, so the extra device allocation
1039 * this used to make is gone. The host copy stays: band data is written
1040 * to the quantum system one band at a time from host memory.
1041 */
1042 std::vector<Complex> final_wavefunctions(static_cast<std::size_t>(number_of_bands) *
1043 static_cast<std::size_t>(number_of_plane_waves));
1044
1045 execution::gpu::device::GPUDeviceBuffer<Complex> device_eigenvector_matrix(
1046 basis_size * basis_size, device_backend_);
1047 std::memcpy(device_eigenvector_matrix.data(),
1048 eigenvector_matrix,
1049 basis_size * basis_size * sizeof(Complex));
1050
1051 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
1052
1053 gpu_blas_.gemm('N',
1054 'N',
1055 number_of_plane_waves,
1056 number_of_bands,
1057 basis_size,
1058 alpha_one,
1059 basis_vectors,
1060 leading_dimension,
1061 device_eigenvector_matrix.data(),
1062 basis_size,
1063 beta_zero,
1064 rotated_scratch,
1065 leading_dimension);
1066
1067 device_backend_->synchronize();
1068
1069 std::memcpy(final_wavefunctions.data(),
1070 rotated_scratch,
1071 final_wavefunctions.size() * sizeof(Complex));
1072
1073 enforce_gamma_g0_real(
1074 is_gamma_only, number_of_plane_waves, number_of_bands, final_wavefunctions.data());
1075
1076 for (int band_index = 0; band_index < number_of_bands; ++band_index) {
1077 quantum_system.set_eigenvalue(kpoint_index_state, band_index, eigenvalues[band_index]);
1078 Complex* band_data =
1079 quantum_system.wavefunction_band_data(kpoint_index_state, band_index);
1080 std::copy(&final_wavefunctions[band_index * number_of_plane_waves],
1081 &final_wavefunctions[(band_index + 1) * number_of_plane_waves],
1082 band_data);
1083 }
1084 }
1085
1086 void apply_preconditioner(int number_of_bands,
1087 int number_of_plane_waves,
1088 const Real* kinetic_energy_squared,
1089 const Real* eigenvalues,
1090 const std::vector<bool>& band_converged,
1091 Complex* residual_wavefunctions) override {
1092 // Copy band_converged to a temporary device buffer.
1093 execution::gpu::device::GPUDeviceBuffer<bool> device_band_converged(number_of_bands,
1094 device_backend_);
1095 /**
1096 * std::vector<bool> is bit-packed and has no contiguous bool storage,
1097 * so stage through a byte-per-flag vector for the device upload.
1098 */
1099 std::vector<char> host_converged(number_of_bands);
1100 for (int band = 0; band < number_of_bands; ++band) {
1101 host_converged[band] = band_converged[band] ? 1 : 0;
1102 }
1103 std::memcpy(
1104 device_band_converged.data(), host_converged.data(), number_of_bands * sizeof(bool));
1105
1106 // Copy kinetic_energy_squared to a temporary device buffer.
1107 execution::gpu::device::GPUDeviceBuffer<Real> device_kinetic_energy_squared(
1108 number_of_plane_waves, device_backend_);
1109 std::memcpy(device_kinetic_energy_squared.data(),
1110 kinetic_energy_squared,
1111 number_of_plane_waves * sizeof(Real));
1112
1113 diagonalizer::planewave::gpu::launch_apply_preconditioner(
1114 number_of_bands,
1115 number_of_plane_waves,
1116 device_kinetic_energy_squared.data(),
1117 eigenvalues,
1118 device_band_converged.data(),
1119 math::tolerances::preconditioner_denominator_floor,
1120 reinterpret_cast<std::complex<double>*>(residual_wavefunctions));
1121 device_backend_->synchronize();
1122 }
1123
1124 private:
1125 // ─── Helpers ──────────────────────────────────────────────────────
1126
1127 /**
1128 * Build full subspace matrices and diagonalize for collapse.
1129 */
1130 void build_subspace_matrices_and_diagonalize(int basis_size,
1131 int number_of_plane_waves,
1132 bool is_gamma_only,
1133 const OverlapOperator* overlap_operator,
1134 const Complex* basis_vectors,
1135 const Complex* hamiltonian_basis_vectors,
1136 const Complex* overlap_basis_vectors,
1137 Real* eigenvalues,
1138 Complex* eigenvector_matrix) {
1139
1140 std::vector<Complex> subspace_hamiltonian(basis_size * basis_size, Complex(0.0, 0.0));
1141 std::vector<Complex> subspace_overlap(basis_size * basis_size, Complex(0.0, 0.0));
1142
1143 int leading_dimension = std::max(number_of_plane_waves, 1);
1144 Complex alpha_one(1.0, 0.0), beta_zero(0.0, 0.0);
1145
1146 if (!is_gamma_only) {
1147 execution::gpu::device::GPUDeviceBuffer<Complex> device_subspace_hamiltonian(
1148 basis_size * basis_size, device_backend_);
1149 execution::gpu::device::GPUDeviceBuffer<Complex> device_subspace_overlap(
1150 basis_size * basis_size, device_backend_);
1151
1152 gpu_blas_.gemm('C',
1153 'N',
1154 basis_size,
1155 basis_size,
1156 number_of_plane_waves,
1157 alpha_one,
1158 basis_vectors,
1159 leading_dimension,
1160 hamiltonian_basis_vectors,
1161 leading_dimension,
1162 beta_zero,
1163 device_subspace_hamiltonian.data(),
1164 basis_size);
1165
1166 gpu_blas_.gemm('C',
1167 'N',
1168 basis_size,
1169 basis_size,
1170 number_of_plane_waves,
1171 alpha_one,
1172 basis_vectors,
1173 leading_dimension,
1174 overlap_basis_vectors,
1175 leading_dimension,
1176 beta_zero,
1177 device_subspace_overlap.data(),
1178 basis_size);
1179
1180 device_backend_->synchronize();
1181
1182 std::memcpy(subspace_hamiltonian.data(),
1183 device_subspace_hamiltonian.data(),
1184 subspace_hamiltonian.size() * sizeof(Complex));
1185 std::memcpy(subspace_overlap.data(),
1186 device_subspace_overlap.data(),
1187 subspace_overlap.size() * sizeof(Complex));
1188 } else {
1189 // Gamma-only: use GPU DGEMM.
1190 int double_number_of_plane_waves = 2 * number_of_plane_waves;
1191 int double_leading_dimension = 2 * leading_dimension;
1192 double alpha_two = 2.0;
1193 double d_beta_zero = 0.0;
1194
1195 auto* real_basis = reinterpret_cast<const double*>(basis_vectors);
1196 auto* real_h_basis = reinterpret_cast<const double*>(hamiltonian_basis_vectors);
1197 auto* real_s_basis = reinterpret_cast<const double*>(overlap_basis_vectors);
1198
1199 std::vector<double> h_real(static_cast<size_t>(basis_size) * basis_size, 0.0);
1200 std::vector<double> s_real(static_cast<size_t>(basis_size) * basis_size, 0.0);
1201
1202 execution::gpu::device::GPUDeviceBuffer<double> device_h_real(basis_size * basis_size,
1203 device_backend_);
1204 execution::gpu::device::GPUDeviceBuffer<double> device_s_real(basis_size * basis_size,
1205 device_backend_);
1206
1207 gpu_blas_.dgemm('T',
1208 'N',
1209 basis_size,
1210 basis_size,
1211 double_number_of_plane_waves,
1212 alpha_two,
1213 real_basis,
1214 double_leading_dimension,
1215 real_h_basis,
1216 double_leading_dimension,
1217 d_beta_zero,
1218 device_h_real.data(),
1219 basis_size);
1220
1221 gpu_blas_.dgemm('T',
1222 'N',
1223 basis_size,
1224 basis_size,
1225 double_number_of_plane_waves,
1226 alpha_two,
1227 real_basis,
1228 double_leading_dimension,
1229 real_s_basis,
1230 double_leading_dimension,
1231 d_beta_zero,
1232 device_s_real.data(),
1233 basis_size);
1234
1235 device_backend_->synchronize();
1236
1237 std::memcpy(h_real.data(), device_h_real.data(), h_real.size() * sizeof(double));
1238 std::memcpy(s_real.data(), device_s_real.data(), s_real.size() * sizeof(double));
1239
1240 for (int col = 0; col < basis_size; ++col) {
1241 Complex h_right_g0 = hamiltonian_basis_vectors[col * number_of_plane_waves];
1242 Complex s_right_g0 = overlap_basis_vectors[col * number_of_plane_waves];
1243 for (int row = 0; row < basis_size; ++row) {
1244 Complex left_g0 = basis_vectors[row * number_of_plane_waves];
1245 int element_index = col * basis_size + row;
1246 if (gamma_origin_is_local_) {
1247 h_real[element_index] -= (std::conj(left_g0) * h_right_g0).real();
1248 s_real[element_index] -= (std::conj(left_g0) * s_right_g0).real();
1249 }
1250 subspace_hamiltonian[element_index] = Complex(h_real[element_index], 0.0);
1251 subspace_overlap[element_index] = Complex(s_real[element_index], 0.0);
1252 }
1253 }
1254 }
1255
1256 execution_policy_.reduce_within_group(subspace_hamiltonian.data(), basis_size * basis_size);
1257 execution_policy_.reduce_within_group(subspace_overlap.data(), basis_size * basis_size);
1258
1259 /**
1260 * See diagonalize_subspace() for the three cases: collective solvers are
1261 * entered by every rank, replicated vectors are solved redundantly by
1262 * every rank, and only a distributed non-collective solve is root-only
1263 * plus a broadcast.
1264 */
1265 const bool all_ranks_solve = subspace_diagonalizer_->is_collective();
1266 const bool vectors_are_replicated = execution_policy_.group_size() <= 1;
1267 const bool this_rank_solves =
1268 all_ranks_solve || vectors_are_replicated || execution_policy_.group_rank_index() == 0;
1269 int info = 0;
1270 if (this_rank_solves) {
1271 info = subspace_diagonalizer_->solve(basis_size,
1272 subspace_hamiltonian.data(),
1273 subspace_overlap.data(),
1274 eigenvalues,
1275 eigenvector_matrix,
1276 is_gamma_only);
1277 }
1278 if (!all_ranks_solve && !vectors_are_replicated) {
1279 execution_policy_.broadcast_from_group_root(eigenvector_matrix,
1280 basis_size * basis_size);
1281 execution_policy_.broadcast_from_group_root(eigenvalues, basis_size);
1282 }
1283 static_cast<void>(info);
1284 }
1285
1286 /**
1287 * Enforce a real G=0 component for gamma-only calculations, on the rank
1288 * whose index 0 is the origin.
1289 */
1290 void enforce_gamma_g0_real(bool is_gamma_only,
1291 int number_of_plane_waves,
1292 int number_of_bands,
1293 Complex* vectors) const {
1294 if (!is_gamma_only || !gamma_origin_is_local_)
1295 return;
1296 for (int band = 0; band < number_of_bands; ++band) {
1297 Complex& g0 = vectors[band * number_of_plane_waves];
1298 g0 = Complex(g0.real(), 0.0);
1299 }
1300 }
1301
1302 /**
1303 * Compute gamma-aware real dot product.
1304 */
1305 Real gamma_aware_real_dot(bool is_gamma_only,
1306 int number_of_plane_waves,
1307 const std::complex<double>* vector_a,
1308 const std::complex<double>* vector_b) const {
1309 if (is_gamma_only) {
1310 double result = 0.0;
1311 for (int g_index = 0; g_index < number_of_plane_waves; ++g_index) {
1312 double re_a = vector_a[g_index].real();
1313 double im_a = vector_a[g_index].imag();
1314 double re_b = vector_b[g_index].real();
1315 double im_b = vector_b[g_index].imag();
1316 result += 2.0 * (re_a * re_b + im_a * im_b);
1317 }
1318 // G=0 correction, where this rank holds the origin.
1319 if (gamma_origin_is_local_) {
1320 result -= (vector_a[0].real() * vector_b[0].real() +
1321 vector_a[0].imag() * vector_b[0].imag());
1322 }
1323 return result;
1324 } else {
1325 double result = 0.0;
1326 for (int g_index = 0; g_index < number_of_plane_waves; ++g_index) {
1327 result += (std::conj(vector_a[g_index]) * vector_b[g_index]).real();
1328 }
1329 return result;
1330 }
1331 }
1332
1333 // ─── Member data ──────────────────────────────────────────────────
1334
1335 execution::gpu::blas::GPUBLAS& gpu_blas_;
1336 const execution::ExecutionPolicy& execution_policy_;
1337 bool gamma_origin_is_local_ = true;
1338 std::shared_ptr<execution::gpu::device::GPUDeviceBackend> device_backend_;
1339 std::unique_ptr<SubspaceDiagonalizer> subspace_diagonalizer_;
1340
1341 // ─── Workspaces ───────────────────────────────────────────────────
1342
1343 std::vector<Complex> diagonalize_overlap_copy_workspace_;
1344};
1345
1346} // namespace q3
q3::GPUPlanewaveSubspaceManager::set_gamma_origin_is_local
void set_gamma_origin_is_local(bool origin_is_local) override
Definition gpu.hpp:63
q3::GPUPlanewaveSubspaceManager::GPUPlanewaveSubspaceManager
GPUPlanewaveSubspaceManager(execution::gpu::blas::GPUBLAS &gpu_blas, const execution::ExecutionPolicy &execution_policy, std::shared_ptr< execution::gpu::device::GPUDeviceBackend > device_backend)
Definition gpu.hpp:48
q3::GPUPlanewaveSubspaceManager::orthogonalize_batch
int orthogonalize_batch(int starting_vector_index, int number_of_candidates, int number_of_plane_waves, bool is_gamma_only, const OverlapOperator *overlap_operator, Complex *basis_vectors, Complex *overlap_basis_vectors) override
Definition gpu.hpp:67
q3::GPUPlanewaveSubspaceManager::apply_preconditioner
void apply_preconditioner(int number_of_bands, int number_of_plane_waves, const Real *kinetic_energy_squared, const Real *eigenvalues, const std::vector< bool > &band_converged, Complex *residual_wavefunctions) override
Definition gpu.hpp:1086
q3::GPUPlanewaveSubspaceManager::collapse_subspace
int collapse_subspace(int basis_size, int number_of_plane_waves, int number_of_bands, int maximum_basis_size, bool is_gamma_only, bool is_identity_overlap, const OverlapOperator *overlap_operator, const Complex *subspace_hamiltonian, const Complex *subspace_overlap, Complex *basis_vectors, Complex *hamiltonian_basis_vectors, Complex *overlap_basis_vectors, Complex *rotation_scratch) override
Definition gpu.hpp:906
q3::GPUPlanewaveSubspaceManager::diagonalize_subspace
int diagonalize_subspace(int basis_size, int maximum_basis_size, bool is_gamma_only, const Complex *subspace_hamiltonian, const Complex *subspace_overlap, Real *eigenvalues, Complex *eigenvector_matrix) override
Definition gpu.hpp:711
q3::GPUPlanewaveSubspaceManager::extract_eigenvectors
void extract_eigenvectors(int basis_size, int number_of_plane_waves, int number_of_bands, int kpoint_index_state, bool is_gamma_only, const Complex *eigenvector_matrix, const Complex *basis_vectors, const Real *eigenvalues, Complex *rotated_scratch, PlanewaveKohnShamQuantumSystem &quantum_system) override
Definition gpu.hpp:1024
q3::GPUPlanewaveSubspaceManager::build_correction_vectors
int build_correction_vectors(int basis_size, int number_of_plane_waves, int number_of_bands, int maximum_candidates, bool is_gamma_only, const Complex *eigenvector_matrix, int eigenvector_leading_dimension, const Real *eigenvalues, const std::vector< bool > &band_converged, const Real *kinetic_energy_squared, const Complex *hamiltonian_basis_vectors, const Complex *overlap_basis_vectors, Complex *correction_columns) override
Definition gpu.hpp:793
q3::GPUPlanewaveSubspaceManager::update_subspace_matrices
void update_subspace_matrices(int start_column, int new_columns, int total_basis_size, int maximum_basis_size, int number_of_plane_waves, bool is_gamma_only, const Complex *basis_vectors, const Complex *hamiltonian_basis_vectors, const Complex *overlap_basis_vectors, Complex *subspace_hamiltonian, Complex *subspace_overlap) override
Definition gpu.hpp:576
q3::OverlapOperator::apply_batch
virtual void apply_batch(int number_of_plane_waves, int number_of_vectors, const Complex *wavefunctions, Complex *overlap_wavefunctions) const
Definition init.hpp:39
q3::PlanewaveKohnShamQuantumSystem::set_eigenvalue
void set_eigenvalue(int kpoint_state_index, int band_index, Real value)
Definition init.hpp:121
q3::PlanewaveKohnShamQuantumSystem::wavefunction_band_data
Complex * wavefunction_band_data(int kpoint_state_index, int band_index)
Definition init.hpp:107
q3::execution::ExecutionPolicy::group_size
virtual int group_size() const =0
q3::execution::ExecutionPolicy::broadcast_from_group_root
virtual void broadcast_from_group_root(ComplexValue *, int) const
Definition init.hpp:310
q3::execution::ExecutionPolicy::group_rank_index
virtual int group_rank_index() const =0
q3::execution::ExecutionPolicy::reduce_within_group
virtual void reduce_within_group(ComplexValue *data, int number_of_elements) const =0
q3::execution::gpu::blas::GPUBLAS::gemm
virtual void gemm(char transpose_a, char transpose_b, int m, int n, int k, Complex alpha, const Complex *matrix_a, int leading_dimension_a, const Complex *matrix_b, int leading_dimension_b, Complex beta, Complex *matrix_c, int leading_dimension_c) const =0
q3::execution::gpu::blas::GPUBLAS::dgemm
virtual void dgemm(char transpose_a, char transpose_b, int m, int n, int k, Real alpha, const Real *matrix_a, int leading_dimension_a, const Real *matrix_b, int leading_dimension_b, Real beta, Real *matrix_c, int leading_dimension_c) const =0
q3::execution::gpu::blas::GPUBLAS::scale
virtual void scale(int number_of_elements, Complex scaling_factor, Complex *vector_x) const =0
gpu_kernels.hpp
init.hpp
init.hpp
init.hpp
init.hpp
init.hpp
factory.hpp
q3::blas::dgemm_
void dgemm_(const char *transa, const char *transb, const int *m, const int *n, const int *k, const double *alpha, const double *a, const int *lda, const double *b, const int *ldb, const double *beta, double *c, const int *ldc)
q3::blas::dpotrf_
void dpotrf_(const char *uplo, const int *n, double *a, const int *lda, int *info)
q3::blas::dtrsm_
void dtrsm_(const char *side, const char *uplo, const char *transa, const char *diag, const int *m, const int *n, const double *alpha, const double *a, const int *lda, double *b, const int *ldb)
q3::blas::zgemm_
void zgemm_(const char *transa, const char *transb, const int *m, const int *n, const int *k, const std::complex< double > *alpha, const std::complex< double > *a, const int *lda, const std::complex< double > *b, const int *ldb, const std::complex< double > *beta, std::complex< double > *c, const int *ldc)
q3::blas::dgemv_
void dgemv_(const char *trans, const int *m, const int *n, const double *alpha, const double *a, const int *lda, const double *x, const int *incx, const double *beta, double *y, const int *incy)
q3::blas::zpotrf_
void zpotrf_(const char *uplo, const int *n, std::complex< double > *a, const int *lda, int *info)
q3::blas::ztrsm_
void ztrsm_(const char *side, const char *uplo, const char *transa, const char *diag, const int *m, const int *n, const std::complex< double > *alpha, const std::complex< double > *a, const int *lda, std::complex< double > *b, const int *ldb)
q3::blas::zdscal_
void zdscal_(const int *n, const double *alpha, std::complex< double > *x, const int *incx)
q3::blas::zgemv_
void zgemv_(const char *trans, const int *m, const int *n, const std::complex< double > *alpha, const std::complex< double > *a, const int *lda, const std::complex< double > *x, const int *incx, const std::complex< double > *beta, std::complex< double > *y, const int *incy)
q3::diagonalizer::planewave::gpu::launch_apply_preconditioner
void launch_apply_preconditioner(int number_of_bands, int number_of_plane_waves, const double *kinetic_energy_squared, const double *eigenvalues, const bool *band_converged, double floor_value, std::complex< double > *residual_wavefunctions)
Definition gpu_kernels_hip.cpp:525
q3::log::warn
void warn(spdlog::format_string_t< Args... > fmt, Args &&... args)
Definition logger.hpp:187
q3::log::info
void info(spdlog::format_string_t< Args... > fmt, Args &&... args)
Definition logger.hpp:183
q3::math::tolerances::iterative_solver_norm_threshold
constexpr double iterative_solver_norm_threshold
Definition tolerances.hpp:39
q3::math::tolerances::preconditioner_denominator_floor
constexpr double preconditioner_denominator_floor
Definition tolerances.hpp:50
q3::Real
double Real
Definition type_aliases.hpp:20
q3::Complex
std::complex< double > Complex
Definition type_aliases.hpp:21