Block-Structured AMR Software Framework
Loading...
Searching...
No Matches
AMReX_SpGEMM.H
Go to the documentation of this file.
1#ifndef AMREX_SPGEMM_H_
2#define AMREX_SPGEMM_H_
3#include <AMReX_Config.H>
4
5#include <AMReX_Algorithm.H>
6#include <AMReX_GpuComplex.H>
7#include <AMReX_OpenMP.H>
8#include <AMReX_SpMatrix.H>
9
10#include <algorithm>
11#include <limits>
12
13namespace amrex::detail {
14
16template <typename T, template<typename> class V>
17CSR<T,V> spgemm_empty (Long nrows)
18{
19 CSR<T,V> C;
20 C.row_offset.resize(nrows+1);
21 auto* p = C.row_offset.data();
22 ParallelForOMP(nrows+1, [=] AMREX_GPU_DEVICE (Long i) { p[i] = 0; });
23 Gpu::streamSynchronize();
24 return C;
25}
26
27#if !defined(AMREX_USE_GPU)
28
29// Gustavson's algorithm. Pass 1 counts, pass 2 accumulates into a dense
30// per-thread accumulator. Output rows are sorted.
31template <typename T, template<typename> class V>
32CSR<T,V> spgemm_local_cpu (Long nrows, Long ncols,
33 CsrView<T const> const& A, CsrView<T const> const& B)
34{
35 CSR<T,V> C;
36 C.row_offset.resize(nrows+1);
37 Long* AMREX_RESTRICT crow = C.row_offset.data();
38 AMREX_ASSUME(crow != nullptr); // gcc -Wnull-dereference false positive
39 crow[0] = 0;
40
41#ifdef AMREX_USE_OMP
42#pragma omp parallel
43#endif
44 {
45 Vector<Long> marker(ncols, -1);
46#ifdef AMREX_USE_OMP
47#pragma omp for schedule(dynamic,64)
48#endif
49 for (Long i = 0; i < nrows; ++i) {
50 Long cnt = 0;
51 for (Long ap = A.row_offset[i]; ap < A.row_offset[i+1]; ++ap) {
52 Long const k = A.col_index[ap];
53 for (Long bp = B.row_offset[k]; bp < B.row_offset[k+1]; ++bp) {
54 Long const j = B.col_index[bp];
55 if (marker[j] != i) {
56 marker[j] = i;
57 ++cnt;
58 }
59 }
60 }
61 crow[i+1] = cnt;
62 }
63 }
64
65 for (Long i = 0; i < nrows; ++i) { crow[i+1] += crow[i]; }
66 C.nnz = crow[nrows];
67 C.col_index.resize(C.nnz);
68 C.mat.resize(C.nnz);
69 Long* AMREX_RESTRICT ccol = C.col_index.data();
70 T* AMREX_RESTRICT cmat = C.mat.data();
71
72#ifdef AMREX_USE_OMP
73#pragma omp parallel
74#endif
75 {
76 Vector<Long> marker(ncols, -1);
77 Vector<T> acc(ncols);
78 Vector<Long> cols;
79#ifdef AMREX_USE_OMP
80#pragma omp for schedule(dynamic,64)
81#endif
82 for (Long i = 0; i < nrows; ++i) {
83 cols.clear();
84 for (Long ap = A.row_offset[i]; ap < A.row_offset[i+1]; ++ap) {
85 Long const k = A.col_index[ap];
86 T const a = A.mat[ap];
87 for (Long bp = B.row_offset[k]; bp < B.row_offset[k+1]; ++bp) {
88 Long const j = B.col_index[bp];
89 if (marker[j] != i) {
90 marker[j] = i;
91 cols.push_back(j);
92 acc[j] = a * B.mat[bp];
93 } else {
94 acc[j] += a * B.mat[bp];
95 }
96 }
97 }
98 std::sort(cols.begin(), cols.end());
99 Long p = crow[i];
100 for (Long j : cols) {
101 ccol[p] = j;
102 cmat[p] = acc[j];
103 ++p;
104 }
105 }
106 }
107
108 return C;
109}
110
111#elif defined(AMREX_USE_CUDA)
112
113template <typename T, template<typename> class V>
114CSR<T,V> spgemm_local_cusparse (Long nrows, Long ncols,
115 CsrView<T const> const& A, CsrView<T const> const& B)
116{
117 cusparseHandle_t handle;
118 AMREX_CUSPARSE_SAFE_CALL(cusparseCreate(&handle));
119 AMREX_CUSPARSE_SAFE_CALL(cusparseSetStream(handle, Gpu::gpuStream()));
120
121 cudaDataType data_type;
122 if constexpr (std::is_same_v<T,float>) {
123 data_type = CUDA_R_32F;
124 } else if constexpr (std::is_same_v<T,double>) {
125 data_type = CUDA_R_64F;
126 } else if constexpr (std::is_same_v<T,GpuComplex<float>>) {
127 data_type = CUDA_C_32F;
128 } else if constexpr (std::is_same_v<T,GpuComplex<double>>) {
129 data_type = CUDA_C_64F;
130 } else {
131 amrex::Abort("SpGEMM: unsupported data type");
132 }
133
134 CSR<T,V> C;
135 C.row_offset.resize(nrows+1);
136
137#if (CUDART_VERSION >= 13000)
138 // 64-bit indices are supported by cusparseSpGEMM since CUDA 13.
139 constexpr cusparseIndexType_t index_type = CUSPARSE_INDEX_64I;
140 void* rowA = (void*)A.row_offset;
141 void* colA = (void*)A.col_index;
142 void* rowB = (void*)B.row_offset;
143 void* colB = (void*)B.col_index;
144 void* rowC = (void*)C.row_offset.data();
145#else
146 AMREX_ALWAYS_ASSERT(nrows < Long(std::numeric_limits<int>::max()) &&
147 ncols < Long(std::numeric_limits<int>::max()));
148 constexpr cusparseIndexType_t index_type = CUSPARSE_INDEX_32I;
149 CsrIndex<int,V> ciA, ciB, ciC;
150 ciA.copyFrom(A);
151 ciB.copyFrom(B);
152 ciC.row_offset.resize(nrows+1);
153 void* rowA = (void*)ciA.row_offset.data();
154 void* colA = (void*)ciA.col_index.data();
155 void* rowB = (void*)ciB.row_offset.data();
156 void* colB = (void*)ciB.col_index.data();
157 void* rowC = (void*)ciC.row_offset.data();
158#endif
159
160 cusparseSpMatDescr_t mat_A, mat_B, mat_C;
161 AMREX_CUSPARSE_SAFE_CALL
162 (cusparseCreateCsr(&mat_A, nrows, B.nrows, A.nnz, rowA, colA, (void*)A.mat,
163 index_type, index_type, CUSPARSE_INDEX_BASE_ZERO, data_type));
164 AMREX_CUSPARSE_SAFE_CALL
165 (cusparseCreateCsr(&mat_B, B.nrows, ncols, B.nnz, rowB, colB, (void*)B.mat,
166 index_type, index_type, CUSPARSE_INDEX_BASE_ZERO, data_type));
167 AMREX_CUSPARSE_SAFE_CALL
168 (cusparseCreateCsr(&mat_C, nrows, ncols, 0, rowC, nullptr, nullptr,
169 index_type, index_type, CUSPARSE_INDEX_BASE_ZERO, data_type));
170
171 cusparseSpGEMMDescr_t spgemm_descr;
172 AMREX_CUSPARSE_SAFE_CALL(cusparseSpGEMM_createDescr(&spgemm_descr));
173
174 T alpha = T(1);
175 T beta = T(0);
176 cusparseOperation_t op = CUSPARSE_OPERATION_NON_TRANSPOSE;
177
178 std::size_t buffer_size1 = 0;
179 AMREX_CUSPARSE_SAFE_CALL
180 (cusparseSpGEMM_workEstimation(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_C,
181 data_type, CUSPARSE_SPGEMM_DEFAULT, spgemm_descr,
182 &buffer_size1, nullptr));
183 auto* buffer1 = (void*)The_Arena()->alloc(buffer_size1);
184 AMREX_CUSPARSE_SAFE_CALL
185 (cusparseSpGEMM_workEstimation(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_C,
186 data_type, CUSPARSE_SPGEMM_DEFAULT, spgemm_descr,
187 &buffer_size1, buffer1));
188
189 std::size_t buffer_size2 = 0;
190 AMREX_CUSPARSE_SAFE_CALL
191 (cusparseSpGEMM_compute(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_C,
192 data_type, CUSPARSE_SPGEMM_DEFAULT, spgemm_descr,
193 &buffer_size2, nullptr));
194 auto* buffer2 = (void*)The_Arena()->alloc(buffer_size2);
195 AMREX_CUSPARSE_SAFE_CALL
196 (cusparseSpGEMM_compute(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_C,
197 data_type, CUSPARSE_SPGEMM_DEFAULT, spgemm_descr,
198 &buffer_size2, buffer2));
199
200 std::int64_t c_nrows, c_ncols, c_nnz;
201 AMREX_CUSPARSE_SAFE_CALL(cusparseSpMatGetSize(mat_C, &c_nrows, &c_ncols, &c_nnz));
202 AMREX_ALWAYS_ASSERT(c_nrows == nrows && c_ncols == ncols);
203#if (CUDART_VERSION < 13000)
204 AMREX_ALWAYS_ASSERT(c_nnz < std::int64_t(std::numeric_limits<int>::max()));
205#endif
206
207 C.mat.resize(c_nnz);
208 C.col_index.resize(c_nnz);
209 C.nnz = c_nnz;
210#if (CUDART_VERSION >= 13000)
211 void* colC = (void*)C.col_index.data();
212#else
213 ciC.col_index.resize(c_nnz);
214 void* colC = (void*)ciC.col_index.data();
215#endif
216 AMREX_CUSPARSE_SAFE_CALL(cusparseCsrSetPointers(mat_C, rowC, colC, (void*)C.mat.data()));
217
218 // cuSPARSE guarantees sorted column indices in C.
219 AMREX_CUSPARSE_SAFE_CALL
220 (cusparseSpGEMM_copy(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_C,
221 data_type, CUSPARSE_SPGEMM_DEFAULT, spgemm_descr));
222
223#if (CUDART_VERSION < 13000)
224 ciC.copyTo(C);
225#endif
226
227 Gpu::streamSynchronize();
228 AMREX_CUSPARSE_SAFE_CALL(cusparseSpGEMM_destroyDescr(spgemm_descr));
229 AMREX_CUSPARSE_SAFE_CALL(cusparseDestroySpMat(mat_A));
230 AMREX_CUSPARSE_SAFE_CALL(cusparseDestroySpMat(mat_B));
231 AMREX_CUSPARSE_SAFE_CALL(cusparseDestroySpMat(mat_C));
232 AMREX_CUSPARSE_SAFE_CALL(cusparseDestroy(handle));
233 The_Arena()->free(buffer1);
234 The_Arena()->free(buffer2);
235
236 return C;
237}
238
239#elif defined(AMREX_USE_HIP)
240
241template <typename T, template<typename> class V>
242CSR<T,V> spgemm_local_rocsparse (Long nrows, Long ncols,
243 CsrView<T const> const& A, CsrView<T const> const& B)
244{
245 static_assert(sizeof(Long) == 8);
246
247 rocsparse_handle handle;
248 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_create_handle(&handle));
249 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_set_stream(handle, Gpu::gpuStream()));
250
251 rocsparse_datatype data_type;
252 if constexpr (std::is_same_v<T,float>) {
253 data_type = rocsparse_datatype_f32_r;
254 } else if constexpr (std::is_same_v<T,double>) {
255 data_type = rocsparse_datatype_f64_r;
256 } else if constexpr (std::is_same_v<T,GpuComplex<float>>) {
257 data_type = rocsparse_datatype_f32_c;
258 } else if constexpr (std::is_same_v<T,GpuComplex<double>>) {
259 data_type = rocsparse_datatype_f64_c;
260 } else {
261 amrex::Abort("SpGEMM: unsupported data type");
262 }
263
264 constexpr rocsparse_indextype index_type = rocsparse_indextype_i64;
265 constexpr rocsparse_index_base index_base = rocsparse_index_base_zero;
266
267 CSR<T,V> C;
268 C.row_offset.resize(nrows+1);
269
270 rocsparse_spmat_descr mat_A, mat_B, mat_C, mat_D;
271 AMREX_ROCSPARSE_SAFE_CALL
272 (rocsparse_create_csr_descr(&mat_A, nrows, B.nrows, A.nnz,
273 (void*)A.row_offset, (void*)A.col_index, (void*)A.mat,
274 index_type, index_type, index_base, data_type));
275 AMREX_ROCSPARSE_SAFE_CALL
276 (rocsparse_create_csr_descr(&mat_B, B.nrows, ncols, B.nnz,
277 (void*)B.row_offset, (void*)B.col_index, (void*)B.mat,
278 index_type, index_type, index_base, data_type));
279 AMREX_ROCSPARSE_SAFE_CALL
280 (rocsparse_create_csr_descr(&mat_C, nrows, ncols, 0,
281 (void*)C.row_offset.data(), nullptr, nullptr,
282 index_type, index_type, index_base, data_type));
283 // D is unused because beta is zero, but a valid descriptor is required.
284 AMREX_ROCSPARSE_SAFE_CALL
285 (rocsparse_create_csr_descr(&mat_D, 0, 0, 0, nullptr, nullptr, nullptr,
286 index_type, index_type, index_base, data_type));
287
288 T alpha = T(1);
289 T beta = T(0);
290 auto const op = rocsparse_operation_none;
291 auto const alg = rocsparse_spgemm_alg_default;
292
293 std::size_t buffer_size = 0;
294 AMREX_ROCSPARSE_SAFE_CALL
295 (rocsparse_spgemm(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_D, mat_C,
296 data_type, alg, rocsparse_spgemm_stage_buffer_size,
297 &buffer_size, nullptr));
298 auto* buffer = (void*)The_Arena()->alloc(buffer_size);
299
300 // This stage fills C's row offsets.
301 AMREX_ROCSPARSE_SAFE_CALL
302 (rocsparse_spgemm(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_D, mat_C,
303 data_type, alg, rocsparse_spgemm_stage_nnz,
304 &buffer_size, buffer));
305
306 std::int64_t c_nrows, c_ncols, c_nnz;
307 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_spmat_get_size(mat_C, &c_nrows, &c_ncols, &c_nnz));
308 AMREX_ALWAYS_ASSERT(c_nrows == nrows && c_ncols == ncols);
309
310 C.mat.resize(c_nnz);
311 C.col_index.resize(c_nnz);
312 C.nnz = c_nnz;
313 AMREX_ROCSPARSE_SAFE_CALL
314 (rocsparse_csr_set_pointers(mat_C, (void*)C.row_offset.data(),
315 (void*)C.col_index.data(), (void*)C.mat.data()));
316
317 AMREX_ROCSPARSE_SAFE_CALL
318 (rocsparse_spgemm(handle, op, op, &alpha, mat_A, mat_B, &beta, mat_D, mat_C,
319 data_type, alg, rocsparse_spgemm_stage_compute,
320 &buffer_size, buffer));
321
322 Gpu::streamSynchronize();
323 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_destroy_spmat_descr(mat_A));
324 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_destroy_spmat_descr(mat_B));
325 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_destroy_spmat_descr(mat_C));
326 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_destroy_spmat_descr(mat_D));
327 AMREX_ROCSPARSE_SAFE_CALL(rocsparse_destroy_handle(handle));
328 The_Arena()->free(buffer);
329
330 C.sort(); // rocSPARSE does not promise sorted rows for the generic API.
331
332 return C;
333}
334
335#elif defined(AMREX_USE_SYCL)
336
337template <typename T, template<typename> class V>
338CSR<T,V> spgemm_local_onemkl (Long nrows, Long ncols,
339 CsrView<T const> const& A, CsrView<T const> const& B)
340{
341 auto& q = Gpu::Device::streamQueue();
342
343 CSR<T,V> C;
344 C.row_offset.resize(nrows+1);
345 // oneMKL wants valid pointers for C before the nnz is known.
346 V<Long> dummy_col(1);
347 V<T> dummy_mat(1);
348
349 mkl::sparse::matrix_handle_t hA{}, hB{}, hC{};
350 mkl::sparse::init_matrix_handle(&hA);
351 mkl::sparse::init_matrix_handle(&hB);
352 mkl::sparse::init_matrix_handle(&hC);
353
354#if defined(INTEL_MKL_VERSION) && (INTEL_MKL_VERSION < 20250300)
355 mkl::sparse::set_csr_data(q, hA, nrows, B.nrows, mkl::index_base::zero,
356 (Long*)A.row_offset, (Long*)A.col_index, (T*)A.mat);
357 mkl::sparse::set_csr_data(q, hB, B.nrows, ncols, mkl::index_base::zero,
358 (Long*)B.row_offset, (Long*)B.col_index, (T*)B.mat);
359 mkl::sparse::set_csr_data(q, hC, nrows, ncols, mkl::index_base::zero,
360 C.row_offset.data(), dummy_col.data(), dummy_mat.data());
361#else
362 mkl::sparse::set_csr_data(q, hA, nrows, B.nrows, A.nnz, mkl::index_base::zero,
363 (Long*)A.row_offset, (Long*)A.col_index, (T*)A.mat);
364 mkl::sparse::set_csr_data(q, hB, B.nrows, ncols, B.nnz, mkl::index_base::zero,
365 (Long*)B.row_offset, (Long*)B.col_index, (T*)B.mat);
366 mkl::sparse::set_csr_data(q, hC, nrows, ncols, Long(0), mkl::index_base::zero,
367 C.row_offset.data(), dummy_col.data(), dummy_mat.data());
368#endif
369
370 mkl::sparse::matmat_descr_t descr = nullptr;
371 mkl::sparse::init_matmat_descr(&descr);
372 mkl::sparse::set_matmat_data(descr,
373 mkl::sparse::matrix_view_descr::general,
374 mkl::transpose::nontrans,
375 mkl::sparse::matrix_view_descr::general,
376 mkl::transpose::nontrans,
377 mkl::sparse::matrix_view_descr::general);
378
379 using req = mkl::sparse::matmat_request;
380 auto* size_buf = (std::int64_t*)The_Pinned_Arena()->alloc(sizeof(std::int64_t));
381
382 mkl::sparse::matmat(q, hA, hB, hC, req::get_work_estimation_buf_size, descr,
383 size_buf, nullptr, {}).wait();
384 auto* buffer1 = (void*)The_Arena()->alloc(std::size_t(*size_buf));
385 mkl::sparse::matmat(q, hA, hB, hC, req::work_estimation, descr,
386 size_buf, buffer1, {}).wait();
387
388 mkl::sparse::matmat(q, hA, hB, hC, req::get_compute_buf_size, descr,
389 size_buf, nullptr, {}).wait();
390 auto* buffer2 = (void*)The_Arena()->alloc(std::size_t(*size_buf));
391 mkl::sparse::matmat(q, hA, hB, hC, req::compute, descr,
392 size_buf, buffer2, {}).wait();
393
394 mkl::sparse::matmat(q, hA, hB, hC, req::get_nnz, descr,
395 size_buf, nullptr, {}).wait();
396 Long const c_nnz = *size_buf;
397
398 C.mat.resize(c_nnz);
399 C.col_index.resize(c_nnz);
400 C.nnz = c_nnz;
401#if defined(INTEL_MKL_VERSION) && (INTEL_MKL_VERSION < 20250300)
402 mkl::sparse::set_csr_data(q, hC, nrows, ncols, mkl::index_base::zero,
403 C.row_offset.data(), C.col_index.data(), C.mat.data());
404#else
405 mkl::sparse::set_csr_data(q, hC, nrows, ncols, c_nnz, mkl::index_base::zero,
406 C.row_offset.data(), C.col_index.data(), C.mat.data());
407#endif
408
409 mkl::sparse::matmat(q, hA, hB, hC, req::finalize, descr,
410 size_buf, nullptr, {}).wait();
411
412 mkl::sparse::release_matmat_descr(&descr);
413 mkl::sparse::release_matrix_handle(q, &hA);
414 mkl::sparse::release_matrix_handle(q, &hB);
415 auto ev = mkl::sparse::release_matrix_handle(q, &hC);
416 ev.wait();
417 Gpu::streamSynchronize();
418 The_Arena()->free(buffer1);
419 The_Arena()->free(buffer2);
420 The_Pinned_Arena()->free(size_buf);
421
422 C.sort(); // oneMKL does not sort the output.
423
424 return C;
425}
426
427#endif
428
435template <typename T, template<typename> class V>
436CSR<T,V> spgemm_local (Long nrows, Long ncols,
437 CsrView<T const> const& A, CsrView<T const> const& B)
438{
439 AMREX_ASSERT(A.nrows == nrows);
440
441 if (nrows <= 0 || ncols <= 0 || A.nnz <= 0 || B.nnz <= 0 || B.nrows <= 0) {
442 return spgemm_empty<T,V>(nrows);
443 }
444
445#if !defined(AMREX_USE_GPU)
446 return spgemm_local_cpu<T,V>(nrows, ncols, A, B);
447#elif defined(AMREX_USE_CUDA)
448 return spgemm_local_cusparse<T,V>(nrows, ncols, A, B);
449#elif defined(AMREX_USE_HIP)
450 return spgemm_local_rocsparse<T,V>(nrows, ncols, A, B);
451#elif defined(AMREX_USE_SYCL)
452 return spgemm_local_onemkl<T,V>(nrows, ncols, A, B);
453#endif
454}
455
456#ifdef AMREX_USE_MPI
457
458// Row i of the result is row i of X0 followed by row i of X1, with column
459// indices transformed by map0 and map1. Both views must have full row offsets.
460template <typename T, template<typename> class V, typename F0, typename F1>
461CSR<T,V> concat_csr_cols (Long nrows, CsrView<T const> const& X0, CsrView<T const> const& X1,
462 F0 const& map0, F1 const& map1)
463{
464 CSR<T,V> C;
465 C.resize(nrows, X0.nnz + X1.nnz);
466 Long* AMREX_RESTRICT crow = C.row_offset.data();
467 Long* AMREX_RESTRICT ccol = C.col_index.data();
468 T* AMREX_RESTRICT cmat = C.mat.data();
469 ParallelForOMP(nrows+1, [=] AMREX_GPU_DEVICE (Long i)
470 {
471 crow[i] = X0.row_offset[i] + X1.row_offset[i];
472 if (i < nrows) {
473 Long p = crow[i];
474 for (Long q = X0.row_offset[i]; q < X0.row_offset[i+1]; ++q) {
475 ccol[p] = map0(X0.col_index[q]);
476 cmat[p] = X0.mat[q];
477 ++p;
478 }
479 for (Long q = X1.row_offset[i]; q < X1.row_offset[i+1]; ++q) {
480 ccol[p] = map1(X1.col_index[q]);
481 cmat[p] = X1.mat[q];
482 ++p;
483 }
484 }
485 });
486 return C;
487}
488
489// Append rows with global column indices, sorted within each row, to Bh
490// whose first rows are already in the compact column space. Each row is
491// rotated so that the local block [c0,c1) comes first, which keeps it
492// sorted in the compact space [local | remote_union].
493template <typename T, template<typename> class V>
494void append_ext_rows (CSR<T,V>& Bh, Long const* ext_row_offset, Long const* ext_col,
495 T const* ext_mat, Long n_ext, Long nnz_ext,
496 Long c0, Long c1, Long const* ru, Long nru)
497{
498 Long const nb = Bh.nrows();
499 Long const nnz_b = Bh.nnz;
500 Bh.mat.resize(nnz_b + nnz_ext);
501 Bh.col_index.resize(nnz_b + nnz_ext);
502 Bh.row_offset.resize(nb + n_ext + 1);
503 Bh.nnz = nnz_b + nnz_ext;
504 Long* AMREX_RESTRICT brow = Bh.row_offset.data();
505 Long* AMREX_RESTRICT bcol = Bh.col_index.data();
506 T* AMREX_RESTRICT bmat = Bh.mat.data();
507 Long const nlocal = c1 - c0;
508 ParallelForOMP(n_ext, [=] AMREX_GPU_DEVICE (Long r)
509 {
510 Long const b = ext_row_offset[r];
511 Long const e = ext_row_offset[r+1];
512 brow[nb+r+1] = nnz_b + e;
513 Long const p0 = amrex::lower_bound(ext_col+b, ext_col+e, c0) - ext_col;
514 Long const p1 = amrex::lower_bound(ext_col+b, ext_col+e, c1) - ext_col;
515 Long p = nnz_b + b;
516 for (Long q = p0; q < p1; ++q) {
517 bcol[p] = ext_col[q] - c0;
518 bmat[p] = ext_mat[q];
519 ++p;
520 }
521 for (Long q = b; q < e; ++q) {
522 if (q < p0 || q >= p1) {
523 bcol[p] = nlocal + (amrex::lower_bound(ru, ru+nru, ext_col[q]) - ru);
524 bmat[p] = ext_mat[q];
525 ++p;
526 }
527 }
528 });
529}
530
531#endif
532
533}
534
535namespace amrex {
536
552template <typename T, template <typename> class Allocator>
553SpMatrix<T,Allocator>
554SpGEMM (SpMatrix<T,Allocator> const& A, SpMatrix<T,Allocator> const& B,
555 AlgPartition const& col_partition)
556{
557 using SpMat = SpMatrix<T,Allocator>;
558 using csr_type = typename SpMat::csr_type;
559
560 auto& Am = const_cast<SpMat&>(A);
561 auto& Bm = const_cast<SpMat&>(B);
562 Long const nrows = Am.numLocalRows();
563
564#ifdef AMREX_USE_MPI
565 using LongVec = typename SpMat::template container_type<Long>;
566
567 Am.setColumnPartition(Bm.partition());
568 Bm.setColumnPartition(col_partition);
569
570 Long const nb = Bm.numLocalRows();
571 Long const c0 = col_partition.globalRowBegin();
572 Long const c1 = col_partition.globalRowEnd();
573 Long const nlocal = c1 - c0;
574
575 // Rows of B needed by A's off-diagonal columns, global column indices.
576 typename SpMat::RemoteRowsMM ext;
577 if (! detail::spmat_comm_is_local(Am.partition(), Bm.partition())) {
578 ext = Am.fetch_remote_rows_mm(Bm);
579 }
580
581 // Compact column space: [local columns | sorted remote columns].
582 // TODO: this serial host loop over ext.nnz limits scaling on GPUs.
583 Vector<Long> ru_h = Bm.m_remote_cols_v;
584 for (Long i = 0; i < ext.nnz; ++i) {
585 auto g = ext.col_index[i];
586 if (g < c0 || g >= c1) { ru_h.push_back(g); }
587 }
588 RemoveDuplicates(ru_h);
589 Long const nru = Long(ru_h.size());
590 LongVec ru_d(nru);
591 Gpu::copyAsync(Gpu::hostToDevice, ru_h.begin(), ru_h.end(), ru_d.begin());
592 Long const* ru = ru_d.data();
593 Long const ncols_hat = nlocal + nru;
594
595 csr_type Ah = detail::concat_csr_cols<T,SpMat::template container_type>
596 (nrows, Am.m_csr.const_view(), Am.remote_full_const_view(),
597 [=] AMREX_GPU_DEVICE (Long c) { return c; },
598 [=] AMREX_GPU_DEVICE (Long j) { return nb + j; });
599
600#ifdef AMREX_USE_GPU
601 Long const* b_rcols = Bm.m_remote_cols_dv.data();
602#else
603 Long const* b_rcols = Bm.m_remote_cols_v.data();
604#endif
605 csr_type Bh = detail::concat_csr_cols<T,SpMat::template container_type>
606 (nb, Bm.m_csr.const_view(), Bm.remote_full_const_view(),
607 [=] AMREX_GPU_DEVICE (Long c) { return c; },
608 [=] AMREX_GPU_DEVICE (Long j) {
609 return nlocal + (amrex::lower_bound(ru, ru+nru, b_rcols[j]) - ru); });
610
611 if (ext.nrows > 0) {
612 LongVec ext_col_d(ext.nnz);
613 Gpu::copyAsync(Gpu::hostToDevice, ext.col_index, ext.col_index+ext.nnz,
614 ext_col_d.begin());
615 detail::append_ext_rows(Bh, ext.row_offset.data(), ext_col_d.data(), ext.mat,
616 ext.nrows, ext.nnz, c0, c1, ru, nru);
617 Gpu::streamSynchronize();
618 }
619 ext.clear();
620
621 csr_type Ch = detail::spgemm_local<T,SpMat::template container_type>
622 (nrows, ncols_hat, Ah.const_view(), Bh.const_view());
623 Ah = csr_type{};
624 Bh = csr_type{};
625
626 // Back to global column indices. Rows are sorted within the local and
627 // remote blocks, which is all split_csr needs.
628 {
629 auto* pc = Ch.col_index.data();
630 ParallelForOMP(Ch.nnz, [=] AMREX_GPU_DEVICE (Long i) {
631 auto c = pc[i];
632 pc[i] = (c < nlocal) ? c + c0 : ru[c - nlocal];
633 });
634 Gpu::streamSynchronize();
635 }
636
637 SpMat C(Am.partition(), std::move(Ch));
638 C.setColumnPartition(col_partition);
639 return C;
640
641#else
642
643 Am.setColumnPartition(Bm.partition());
644 Bm.setColumnPartition(col_partition);
645
646 Long const ncols = col_partition.numGlobalRows();
647 csr_type Ch = detail::spgemm_local<T,SpMat::template container_type>
648 (nrows, ncols, Am.m_csr.const_view(), Bm.m_csr.const_view());
649 SpMat C(Am.partition(), std::move(Ch));
650 C.setColumnPartition(col_partition);
651 return C;
652
653#endif
654}
655
656}
657
658#endif
AMReX_Algorithm.H
General-purpose algorithm utilities available on both host and device.
AMREX_ASSERT
#define AMREX_ASSERT(EX)
Definition AMReX_BLassert.H:38
AMREX_ALWAYS_ASSERT
#define AMREX_ALWAYS_ASSERT(EX)
Definition AMReX_BLassert.H:50
AMREX_ASSUME
#define AMREX_ASSUME(ASSUMPTION)
Definition AMReX_Extension.H:287
AMREX_RESTRICT
#define AMREX_RESTRICT
Definition AMReX_Extension.H:37
AMReX_GpuComplex.H
AMREX_CUSPARSE_SAFE_CALL
#define AMREX_CUSPARSE_SAFE_CALL(call)
Definition AMReX_GpuError.H:101
AMREX_GPU_DEVICE
#define AMREX_GPU_DEVICE
Definition AMReX_GpuQualifiers.H:18
beta
GpuArray< Real, 3 > beta
Definition AMReX_MLEBNodeFDLaplacian.cpp:1099
AMReX_OpenMP.H
AMReX_SpMatrix.H
amrex::AlgPartition
Definition AMReX_AlgPartition.H:21
amrex::AlgPartition::numGlobalRows
Long numGlobalRows() const
Total number of rows covered by the partition.
Definition AMReX_AlgPartition.H:50
amrex::AlgPartition::globalRowEnd
Long globalRowEnd() const
Exclusive global index end on this process.
Definition AMReX_AlgPartition.H:62
amrex::AlgPartition::globalRowBegin
Long globalRowBegin() const
Inclusive global index begin on this process.
Definition AMReX_AlgPartition.H:57
amrex::Arena::free
virtual void free(void *pt)=0
Free a previously allocated block pointed to by pt.
amrex::Arena::alloc
virtual void * alloc(std::size_t sz)=0
Allocate sz bytes from this arena.
amrex::SpMatrix
Distributed CSR matrix that manages storage and GPU-friendly partitions.
Definition AMReX_SpMatrix.H:63
amrex::SpMatrix::numLocalRows
Long numLocalRows() const
Number of rows owned by this rank.
Definition AMReX_SpMatrix.H:191
amrex::Vector
This class is a thin wrapper around std::vector. Unlike vector, Vector::operator[] provides bound che...
Definition AMReX_Vector.H:29
amrex::Vector::size
Long size() const noexcept
Definition AMReX_Vector.H:54
amrex::Long
amrex_long Long
Definition AMReX_INT.H:30
amrex::ParallelForOMP
void ParallelForOMP(T n, L const &f) noexcept
Performance-portable kernel launch function with optional OpenMP threading.
Definition AMReX_GpuLaunch.H:328
amrex::The_Pinned_Arena
Arena * The_Pinned_Arena()
Definition AMReX_Arena.cpp:855
amrex::The_Arena
Arena * The_Arena()
Definition AMReX_Arena.cpp:815
amrex::lower_bound
__host__ __device__ ItType lower_bound(ItType first, ItType last, const ValType &val)
Return an iterator to the first element not less than a given value.
Definition AMReX_Algorithm.H:298
amrex::Gpu::copyAsync
void copyAsync(HostToDevice, InIter begin, InIter end, OutIter result) noexcept
A host-to-device copy routine. Note this is just a wrapper around memcpy, so it assumes contiguous st...
Definition AMReX_GpuContainers.H:228
amrex::Gpu::hostToDevice
static constexpr HostToDevice hostToDevice
Definition AMReX_GpuContainers.H:105
amrex::Gpu::streamSynchronize
void streamSynchronize() noexcept
Definition AMReX_GpuDevice.H:310
amrex::Gpu::gpuStream
gpuStream_t gpuStream() noexcept
Definition AMReX_GpuDevice.H:291
amrex
Definition AMReX_Amr.cpp:50
amrex::Abort
void Abort(const std::string &msg)
Print a fatal-error message to stderr and abort execution.
Definition AMReX.cpp:242
amrex::RemoveDuplicates
void RemoveDuplicates(Vector< T > &vec)
Definition AMReX_Vector.H:210
amrex::SpGEMM
SpMatrix< T, Allocator > SpGEMM(SpMatrix< T, Allocator > const &A, SpMatrix< T, Allocator > const &B, AlgPartition const &col_partition)
Sparse matrix-matrix multiplication, C = A * B.
Definition AMReX_SpGEMM.H:554