Actual source code: aijcupm.hpp
1: #pragma once
3: /* Shared CUPM (CUDA/HIP) implementations for SeqAIJCUSPARSE and SeqAIJHIPSPARSE
4: that do not depend on the cuSPARSE/hipSPARSE library proper.
6: Include ordering requirement: the vendor-specific impl header
7: (cusparsematimpl.h or hipsparsematimpl.h) must be included before this
8: header so that CsrMatrix, THRUSTINTARRAY*, THRUSTARRAY and all device-specific
9: struct types are visible when this header is processed.
11: Instantiated by:
12: aijcusparse.cu (DeviceType::CUDA, using MatSeqAIJCUSPARSE_Policy)
13: aijhipsparse.hip.cxx (DeviceType::HIP, using MatSeqAIJHIPSPARSE_Policy) */
15: #include <petsc/private/cupmobject.hpp>
16: #include <petsc/private/cupmblasinterface.hpp>
17: #include <petsc/private/matimpl.h>
18: #include <../src/sys/objects/device/impls/cupm/cupmthrustutility.hpp>
19: #include <../src/mat/impls/aij/seq/aij.h>
21: #include <thrust/device_ptr.h>
22: #include <thrust/iterator/counting_iterator.h>
23: #include <thrust/iterator/permutation_iterator.h>
24: #include <thrust/functional.h>
25: #include <thrust/fill.h>
26: #include <thrust/tuple.h>
27: #include <thrust/transform.h>
28: #include <thrust/for_each.h>
29: #include <thrust/equal.h>
31: /* Forward declaration of SeqAIJ fallback function used inside template */
32: PETSC_INTERN PetscErrorCode MatGetDiagonal_SeqAIJ(Mat, Vec);
34: namespace Petsc
35: {
37: namespace mat
38: {
40: namespace aij
41: {
43: namespace cupm
44: {
46: namespace impl
47: {
49: // expand row_offsets to uncompressed row indices coo_i[]
50: struct Csr2coo {
51: const PetscInt *row_offsets;
52: PetscInt *coo_i;
54: Csr2coo(const PetscInt *roff, PetscInt *cooi) : row_offsets(roff), coo_i(cooi) { }
56: PETSC_HOSTDEVICE_INLINE_DECL void operator()(PetscInt i) const
57: {
58: for (PetscInt j = row_offsets[i]; j < row_offsets[i + 1]; ++j) coo_i[j] = i;
59: }
60: };
62: struct PetscIntToCInt {
63: // Caller should check overflow
64: PETSC_HOSTDEVICE_INLINE_DECL int operator()(PetscInt i) const { return static_cast<int>(i); }
65: };
67: /* --------------------------------------------------------------------------
68: Shared device functor: left-scale CSR rows.
69: cprow[i] gives the logical row index for compressed row i; NULL = identity.
70: -------------------------------------------------------------------------- */
71: struct DiagonalScaleLeft_CSR_Functor {
72: const PetscInt *row_ptr;
73: PetscScalar *val_ptr;
74: const PetscScalar *lv_ptr;
75: const PetscInt *cprow;
77: PETSC_HOSTDEVICE_INLINE_DECL void operator()(PetscInt i) const
78: {
79: const PetscInt row = cprow ? cprow[i] : i;
80: const PetscScalar s = lv_ptr[row];
81: for (PetscInt j = row_ptr[i]; j < row_ptr[i + 1]; j++) val_ptr[j] *= s;
82: }
83: };
85: /* --------------------------------------------------------------------------
86: Shared device functor: get<1>(t) = get<0>(t).
87: Replaces the identical VecCUDAEquals / VecHIPEquals structs.
88: -------------------------------------------------------------------------- */
89: struct VecCUPMEquals {
90: template <typename Tuple>
91: PETSC_HOSTDEVICE_INLINE_DECL void operator()(Tuple t) const
92: {
93: thrust::get<1>(t) = thrust::get<0>(t);
94: }
95: };
97: /* --------------------------------------------------------------------------
98: Shared __global__ kernel: accumulate COO values into a CSR array.
99: __global__ is valid for both nvcc and hipcc; the body is identical.
100: -------------------------------------------------------------------------- */
101: __global__ static void MatAddCOOValues(const PetscScalar kv[], PetscCount nnz, const PetscCount jmap[], const PetscCount perm[], InsertMode imode, PetscScalar a[])
102: {
103: PetscCount i = blockIdx.x * blockDim.x + threadIdx.x;
104: const PetscCount grid_size = gridDim.x * blockDim.x;
105: for (; i < nnz; i += grid_size) {
106: PetscScalar sum = 0.0;
107: for (PetscCount k = jmap[i]; k < jmap[i + 1]; k++) sum += kv[perm[k]];
108: a[i] = (imode == INSERT_VALUES ? (PetscScalar)0.0 : a[i]) + sum;
109: }
110: }
112: /* --------------------------------------------------------------------------
113: Shared __global__ kernel: extract the CSR diagonal.
114: -------------------------------------------------------------------------- */
115: __global__ void GetDiagonal_CSR(const PetscInt *row, const PetscInt *col, const PetscScalar *val, const PetscInt len, PetscScalar *diag)
116: {
117: const size_t x = blockIdx.x * blockDim.x + threadIdx.x;
119: if (x < (size_t)len) {
120: const PetscInt rowx = row[x], num_non0_row = row[x + 1] - rowx;
121: PetscScalar d = 0.0;
123: for (PetscInt i = 0; i < num_non0_row; i++) {
124: if (col[i + rowx] == (PetscInt)x) {
125: d = val[i + rowx];
126: break;
127: }
128: }
129: diag[x] = d;
130: }
131: }
133: /* --------------------------------------------------------------------------
134: Shared __global__ kernel: sum the absolute values of each CSR row.
135: With compressed rows, cprow maps CSR row i to its row in the matrix.
136: -------------------------------------------------------------------------- */
137: __global__ void GetRowSumAbs_CSR(const PetscInt *rowptr, const PetscScalar *val, const PetscInt *cprow, const PetscInt m, PetscScalar *sum)
138: {
139: const size_t i = blockIdx.x * blockDim.x + threadIdx.x;
141: if (i < (size_t)m) {
142: PetscReal s = 0.0;
144: for (PetscInt k = rowptr[i]; k < rowptr[i + 1]; k++) s += PetscAbsScalar(val[k]);
145: sum[cprow ? cprow[i] : (PetscInt)i] = s;
146: }
147: }
149: /* ==========================================================================
150: MatSeqAIJCUSPARSE_CUPM<T, Policy>
152: Policy (C++11 traits class) requirements - all static methods:
154: // Device struct types
155: typedef ... mat_struct_type; // Mat_SeqAIJCUSPARSE / Mat_SeqAIJHIPSPARSE
156: typedef ... mult_struct_type; // ...MultStruct equivalent
158: // Storage-format constants (value of each format enumerator)
159: static int storage_format_csr();
160: static int storage_format_ell();
161: static int storage_format_hyb();
163: // Bookkeeping helpers (device-type specific)
164: static PetscErrorCode CopyToGPU(Mat);
165: static PetscErrorCode CopyFromGPU(Mat);
166: static PetscErrorCode InvalidateTranspose(Mat, PetscBool);
167: static PetscErrorCode ConvertFromSeqAIJ(Mat, MatType, MatReuse, Mat *);
168: static const char *mat_type_name; // "seqaijcusparse" / "seqaijhipsparse"
170: // Destruction helpers (device-type specific)
171: static PetscErrorCode Destroy(Mat);
172: static PetscErrorCode TriFactorsDestroy(void **);
174: // Compose-function keys that differ between CUDA and HIP
175: static const char *set_format_c; // "MatCUSPARSESetFormat_C" / "MatHIPSPARSESetFormat_C"
176: static const char *set_use_cpu_solve_c; // "MatCUSPARSESetUseCPUSolve_C" / "MatHIPSPARSESetUseCPUSolve_C"
177: static const char *product_seqdense_device_c; // "...seqdensecuda_C" / "...seqdensehip_C"
178: static const char *product_seqdense_c; // "...seqdense_C"
179: static const char *product_self_c; // "...seqaijcusparse_C" / "...seqaijhipsparse_C"
180: static const char *seq_convert_hypre_c; // "MatConvert_seqaijcusparse_hypre_C" / "_seqaijhipsparse_hypre_C"
182: // Vec device-array access (device-type specific)
183: static PetscErrorCode VecGetArrayRead (Vec, const PetscScalar **);
184: static PetscErrorCode VecRestoreArrayRead(Vec, const PetscScalar **);
185: static PetscErrorCode VecGetArrayWrite (Vec, PetscScalar **);
186: static PetscErrorCode VecRestoreArrayWrite(Vec, PetscScalar **);
187: ========================================================================== */
189: template <device::cupm::DeviceType T, typename Policy>
190: struct MatSeqAIJCUSPARSE_CUPM : device::cupm::impl::CUPMObject<T> {
191: PETSC_CUPMOBJECT_HEADER(T);
193: typedef typename Policy::mat_struct_type MatStructType;
194: typedef typename Policy::mult_struct_type MultStructType;
196: /* -------------------------------------------------------------------
197: Tier 1 - Trivial
198: ------------------------------------------------------------------- */
200: /* MatAssemblyEnd: delegation to SeqAIJ */
201: static PetscErrorCode AssemblyEnd(Mat A, MatAssemblyType mode) noexcept
202: {
203: PetscFunctionBegin;
204: PetscCall(MatAssemblyEnd_SeqAIJ(A, mode));
205: PetscFunctionReturn(PETSC_SUCCESS);
206: }
208: /* MatDuplicate */
209: static PetscErrorCode Duplicate(Mat A, MatDuplicateOption cpvalues, Mat *B) noexcept
210: {
211: PetscFunctionBegin;
212: PetscCall(MatDuplicate_SeqAIJ(A, cpvalues, B));
213: PetscCall(Policy::ConvertFromSeqAIJ(*B, Policy::mat_type_name, MAT_INPLACE_MATRIX, B));
214: PetscFunctionReturn(PETSC_SUCCESS);
215: }
217: /* MatGetCurrentMemType */
218: static PetscErrorCode GetCurrentMemType(PETSC_UNUSED Mat A, PetscMemType *m) noexcept
219: {
220: PetscFunctionBegin;
221: *m = PETSC_MEMTYPE_CUPM();
222: PetscFunctionReturn(PETSC_SUCCESS);
223: }
225: /* MatCOOStructDestroy: free device jmap and perm fields */
226: static PetscErrorCode COOStructDestroy(PetscCtxRt ctx) noexcept
227: {
228: MatCOOStruct_SeqAIJ *coo = *(MatCOOStruct_SeqAIJ **)ctx;
230: PetscFunctionBegin;
231: PetscCallCUPM(cupmFree(coo->perm));
232: PetscCallCUPM(cupmFree(coo->jmap));
233: PetscCall(PetscFree(coo));
234: PetscFunctionReturn(PETSC_SUCCESS);
235: }
237: /* -------------------------------------------------------------------
238: Tier 2 - Straightforward
239: ------------------------------------------------------------------- */
241: /* MatZeroEntries: fill device CSR values with zero */
242: static PetscErrorCode ZeroEntries(Mat A) noexcept
243: {
244: PetscBool gpu = PETSC_FALSE;
245: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
246: MatStructType *spptr;
248: PetscFunctionBegin;
249: if (A->factortype == MAT_FACTOR_NONE) {
250: spptr = (MatStructType *)A->spptr;
251: if (spptr->mat) {
252: CsrMatrix *matrix = (CsrMatrix *)spptr->mat->mat;
253: if (matrix->values) {
254: gpu = PETSC_TRUE;
255: PetscCallThrust(thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), (PetscScalar)0.));
256: }
257: }
258: if (spptr->matTranspose) {
259: CsrMatrix *matrix = (CsrMatrix *)spptr->matTranspose->mat;
260: if (matrix->values) PetscCallThrust(thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), (PetscScalar)0.));
261: }
262: }
263: if (gpu) A->offloadmask = PETSC_OFFLOAD_GPU;
264: else {
265: PetscCall(PetscArrayzero(a->a, a->i[A->rmap->n]));
266: A->offloadmask = PETSC_OFFLOAD_CPU;
267: }
268: PetscFunctionReturn(PETSC_SUCCESS);
269: }
271: /* MatScale: cupmBlasXscal on the device CSR values */
272: static PetscErrorCode Scale(Mat Y, PetscScalar a) noexcept
273: {
274: Mat_SeqAIJ *y = (Mat_SeqAIJ *)Y->data;
275: PetscScalar *ay = nullptr;
276: cupmBlasHandle_t blashandle;
277: PetscBLASInt one = 1, bnz = 1;
279: PetscFunctionBegin;
280: PetscCall(GetArray(Y, &ay));
281: PetscCall(GetHandles_(&blashandle));
282: PetscCall(PetscBLASIntCast(y->nz, &bnz));
283: PetscCall(PetscLogGpuTimeBegin());
284: PetscCallCUPMBLAS(cupmBlasXscal(blashandle, bnz, cupmScalarPtrCast(&a), cupmScalarPtrCast(ay), one));
285: PetscCall(PetscLogGpuFlops(bnz));
286: PetscCall(PetscLogGpuTimeEnd());
287: PetscCall(RestoreArray(Y, &ay));
288: PetscFunctionReturn(PETSC_SUCCESS);
289: }
291: /* MatDiagonalScale: Thrust-based left and right scaling of CSR values */
292: static PetscErrorCode DiagonalScale(Mat A, Vec ll, Vec rr) noexcept
293: {
294: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)A->data;
295: MatStructType *devstruct;
296: CsrMatrix *csr;
297: PetscScalar *av = nullptr;
298: PetscInt m, n, nz = aij->nz;
299: cupmStream_t stream;
301: PetscFunctionBegin;
302: PetscCall(GetHandles_(&stream));
303: PetscCall(PetscLogGpuTimeBegin());
304: PetscCall(GetArray(A, &av));
305: devstruct = (MatStructType *)A->spptr;
306: csr = (CsrMatrix *)devstruct->mat->mat;
307: if (ll) {
308: const PetscScalar *lv;
309: PetscCall(VecGetLocalSize(ll, &m));
310: PetscCheck(m == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector wrong length");
311: PetscCall(Policy::VecGetArrayRead(ll, &lv));
312: {
313: const PetscInt *cprow = devstruct->mat->cprowIndices ? devstruct->mat->cprowIndices->data().get() : NULL;
314: DiagonalScaleLeft_CSR_Functor functor = {csr->row_offsets->data().get(), av, lv, cprow};
315: PetscCallThrust(THRUST_CALL(thrust::for_each, stream, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(csr->num_rows), functor));
316: }
317: PetscCall(Policy::VecRestoreArrayRead(ll, &lv));
318: PetscCall(PetscLogGpuFlops(nz));
319: }
320: if (rr) {
321: const PetscScalar *rv;
322: PetscCall(VecGetLocalSize(rr, &n));
323: PetscCheck(n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vector wrong length");
324: PetscCall(Policy::VecGetArrayRead(rr, &rv));
325: #if PetscDefined(USING_NVCC) && CCCL_VERSION >= 3001000
326: PetscCallThrust(THRUST_CALL(thrust::transform, stream, csr->values->begin(), csr->values->end(), thrust::make_permutation_iterator(thrust::device_pointer_cast(rv), csr->column_indices->begin()), csr->values->begin(), cuda::std::multiplies<PetscScalar>()));
327: #else
328: PetscCallThrust(THRUST_CALL(thrust::transform, stream, csr->values->begin(), csr->values->end(), thrust::make_permutation_iterator(thrust::device_pointer_cast(rv), csr->column_indices->begin()), csr->values->begin(), thrust::multiplies<PetscScalar>()));
329: #endif
330: PetscCall(Policy::VecRestoreArrayRead(rr, &rv));
331: PetscCall(PetscLogGpuFlops(nz));
332: }
333: PetscCall(RestoreArray(A, &av));
334: PetscCall(PetscLogGpuTimeEnd());
335: PetscFunctionReturn(PETSC_SUCCESS);
336: }
338: /* MatSeqAIJGetIJ: return device CSR row-pointer and column-index arrays */
339: static PetscErrorCode GetIJ(Mat A, PetscBool compressed, const PetscInt **i, const PetscInt **j) noexcept
340: {
341: MatStructType *cusp = (MatStructType *)A->spptr;
342: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
343: CsrMatrix *csr;
345: PetscFunctionBegin;
347: if (!i || !j) PetscFunctionReturn(PETSC_SUCCESS);
348: PetscCheckTypeName(A, Policy::mat_type_name);
349: PetscCheck(cusp->format != (decltype(cusp->format))Policy::storage_format_ell() && cusp->format != (decltype(cusp->format))Policy::storage_format_hyb(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
350: PetscCall(Policy::CopyToGPU(A));
351: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
352: csr = (CsrMatrix *)cusp->mat->mat;
353: if (i) {
354: if (!compressed && a->compressedrow.use) { /* need full row offset */
355: if (!cusp->rowoffsets_gpu) {
356: cusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
357: cusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
358: PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
359: }
360: *i = cusp->rowoffsets_gpu->data().get();
361: } else *i = csr->row_offsets->data().get();
362: }
363: if (j) *j = csr->column_indices->data().get();
364: PetscFunctionReturn(PETSC_SUCCESS);
365: }
367: /* MatSeqAIJRestoreIJ: nullify the pointers previously obtained with GetIJ */
368: static PetscErrorCode RestoreIJ(Mat A, PetscBool compressed, const PetscInt **i, const PetscInt **j) noexcept
369: {
370: PetscFunctionBegin;
372: PetscCheckTypeName(A, Policy::mat_type_name);
373: if (i) *i = NULL;
374: if (j) *j = NULL;
375: (void)compressed;
376: PetscFunctionReturn(PETSC_SUCCESS);
377: }
379: /* MatSetPreallocationCOO: copy COO bookkeeping struct to device */
380: static PetscErrorCode SetPreallocationCOO(Mat mat, PetscCount coo_n, PetscInt coo_i[], PetscInt coo_j[]) noexcept
381: {
382: PetscBool dev_ij = PETSC_FALSE;
383: PetscMemType mtype = PETSC_MEMTYPE_HOST;
384: PetscInt *i, *j;
385: PetscContainer container_h;
386: MatCOOStruct_SeqAIJ *coo_h, *coo_d;
388: PetscFunctionBegin;
389: PetscCall(PetscGetMemType(coo_i, &mtype));
390: if (PetscMemTypeDevice(mtype)) {
391: dev_ij = PETSC_TRUE;
392: PetscCall(PetscMalloc2(coo_n, &i, coo_n, &j));
393: PetscCallCUPM(cupmMemcpy(i, coo_i, coo_n * sizeof(PetscInt), cupmMemcpyDeviceToHost));
394: PetscCallCUPM(cupmMemcpy(j, coo_j, coo_n * sizeof(PetscInt), cupmMemcpyDeviceToHost));
395: } else {
396: i = coo_i;
397: j = coo_j;
398: }
399: PetscCall(MatSetPreallocationCOO_SeqAIJ(mat, coo_n, i, j));
400: if (dev_ij) PetscCall(PetscFree2(i, j));
401: mat->offloadmask = PETSC_OFFLOAD_CPU;
402: /* Create the GPU memory */
403: PetscCall(Policy::CopyToGPU(mat));
405: /* Copy the COO struct to device */
406: PetscCall(PetscObjectQuery((PetscObject)mat, "__PETSc_MatCOOStruct_Host", (PetscObject *)&container_h));
407: PetscCall(PetscContainerGetPointer(container_h, (void **)&coo_h));
408: PetscCall(PetscMalloc1(1, &coo_d));
409: *coo_d = *coo_h; /* shallow copy; device fields amended below */
410: PetscCallCUPM(cupmMalloc((void **)&coo_d->jmap, (coo_h->nz + 1) * sizeof(PetscCount)));
411: PetscCallCUPM(cupmMemcpy(coo_d->jmap, coo_h->jmap, (coo_h->nz + 1) * sizeof(PetscCount), cupmMemcpyHostToDevice));
412: PetscCallCUPM(cupmMalloc((void **)&coo_d->perm, coo_h->Atot * sizeof(PetscCount)));
413: PetscCallCUPM(cupmMemcpy(coo_d->perm, coo_h->perm, coo_h->Atot * sizeof(PetscCount), cupmMemcpyHostToDevice));
415: PetscCall(PetscObjectContainerCompose((PetscObject)mat, "__PETSc_MatCOOStruct_Device", coo_d, MatSeqAIJCUSPARSE_CUPM::COOStructDestroy));
416: PetscFunctionReturn(PETSC_SUCCESS);
417: }
419: /* MatSetValuesCOO: launch MatAddCOOValues kernel */
420: static PetscErrorCode SetValuesCOO(Mat A, const PetscScalar v[], InsertMode imode) noexcept
421: {
422: Mat_SeqAIJ *seq = (Mat_SeqAIJ *)A->data;
423: MatStructType *dev = (MatStructType *)A->spptr;
424: PetscCount Annz = seq->nz;
425: PetscMemType memtype;
426: const PetscScalar *v1 = v;
427: PetscScalar *Aa = nullptr;
428: PetscContainer container;
429: MatCOOStruct_SeqAIJ *coo;
430: cupmStream_t stream;
432: PetscFunctionBegin;
433: PetscCall(MatSeqAIJCheckCOONonzeroState(A));
434: if (!dev->mat) PetscCall(Policy::CopyToGPU(A));
436: PetscCall(PetscObjectQuery((PetscObject)A, "__PETSc_MatCOOStruct_Device", (PetscObject *)&container));
437: PetscCall(PetscContainerGetPointer(container, (void **)&coo));
439: PetscCall(PetscGetMemType(v, &memtype));
440: if (PetscMemTypeHost(memtype)) { /* copy host values to device */
441: PetscCallCUPM(cupmMalloc((void **)&v1, coo->n * sizeof(PetscScalar)));
442: PetscCallCUPM(cupmMemcpy((void *)v1, v, coo->n * sizeof(PetscScalar), cupmMemcpyHostToDevice));
443: PetscCall(PetscLogCpuToGpu(coo->n * sizeof(PetscScalar)));
444: }
446: if (imode == INSERT_VALUES) PetscCall(GetArrayWrite(A, &Aa));
447: else PetscCall(GetArray(A, &Aa));
449: PetscCall(GetHandles_(&stream));
450: PetscCall(PetscLogGpuTimeBegin());
451: if (Annz) {
452: PetscCallCUPM(cupmLaunchKernel(MatAddCOOValues, (unsigned int)((Annz + 255) / 256), 256u, (size_t)0, stream, v1, Annz, coo->jmap, coo->perm, imode, Aa));
453: PetscCallCUPM(cupmGetLastError());
454: }
455: PetscCall(PetscLogGpuTimeEnd());
457: if (imode == INSERT_VALUES) PetscCall(RestoreArrayWrite(A, &Aa));
458: else PetscCall(RestoreArray(A, &Aa));
460: if (PetscMemTypeHost(memtype)) {
461: void *v1_device = (void *)v1;
462: PetscCallCUPM(cupmFree(v1_device));
463: }
464: PetscFunctionReturn(PETSC_SUCCESS);
465: }
467: /* MatSeqAIJCopySubArray: scatter-gather a sub-array of CSR values */
468: static PetscErrorCode CopySubArray(Mat A, PetscInt n, const PetscInt idx[], PetscScalar v[]) noexcept
469: {
470: const PetscScalar *av = nullptr;
471: PetscMemType mtype;
472: PetscBool dmem;
474: PetscFunctionBegin;
475: PetscCall(PetscCUPMGetMemType(v, &mtype));
476: dmem = PetscMemTypeDevice(mtype);
477: PetscCall(GetArrayRead(A, &av));
478: if (n && idx) {
479: THRUSTINTARRAY widx(n);
480: widx.assign(idx, idx + n);
481: PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
483: THRUSTARRAY *w = NULL;
484: thrust::device_ptr<PetscScalar> dv;
485: if (dmem) {
486: dv = thrust::device_pointer_cast(v);
487: } else {
488: w = new THRUSTARRAY(n);
489: dv = w->data();
490: }
491: {
492: thrust::device_ptr<const PetscScalar> dav = thrust::device_pointer_cast(av);
493: auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(dav, widx.begin()), dv));
494: auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(dav, widx.end()), dv + n));
495: PetscCallThrust(thrust::for_each(zibit, zieit, VecCUPMEquals{}));
496: }
497: if (w) PetscCallCUPM(cupmMemcpy(v, w->data().get(), n * sizeof(PetscScalar), cupmMemcpyDeviceToHost));
498: delete w;
499: } else {
500: PetscCallCUPM(cupmMemcpy(v, av, n * sizeof(PetscScalar), dmem ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost));
501: }
502: if (!dmem) PetscCall(PetscLogCpuToGpu(n * sizeof(PetscScalar)));
503: PetscCall(RestoreArrayRead(A, &av));
504: PetscFunctionReturn(PETSC_SUCCESS);
505: }
507: /* -------------------------------------------------------------------
508: Tier 3 - AXPY shared branches (SAME_NZ and DIFFERENT_NZ only).
509: The SUBSET_NZ branch calls cuSPARSE/hipSPARSE and stays in the caller.
510: ------------------------------------------------------------------- */
512: /* AXPY SAME_NONZERO_PATTERN branch: cupmBlasXaxpy */
513: static PetscErrorCode AXPY_SameNZ(Mat Y, PetscScalar a, Mat X) noexcept
514: {
515: Mat_SeqAIJ *x = (Mat_SeqAIJ *)X->data;
516: const PetscScalar *ax = nullptr;
517: PetscScalar *ay = nullptr;
518: cupmBlasHandle_t blashandle;
519: PetscBLASInt one = 1, bnz = 1;
521: PetscFunctionBegin;
522: PetscCall(GetArrayRead(X, &ax));
523: PetscCall(GetArray(Y, &ay));
524: PetscCall(GetHandles_(&blashandle));
525: PetscCall(PetscBLASIntCast(x->nz, &bnz));
526: PetscCall(PetscLogGpuTimeBegin());
527: PetscCallCUPMBLAS(cupmBlasXaxpy(blashandle, bnz, cupmScalarPtrCast(&a), cupmScalarPtrCast(ax), one, cupmScalarPtrCast(ay), one));
528: PetscCall(PetscLogGpuFlops(2.0 * bnz));
529: PetscCall(PetscLogGpuTimeEnd());
530: PetscCall(RestoreArrayRead(X, &ax));
531: PetscCall(RestoreArray(Y, &ay));
532: PetscFunctionReturn(PETSC_SUCCESS);
533: }
535: /* GetDiagonal: kernel-based extraction of the CSR diagonal */
536: static PetscErrorCode GetDiagonal(Mat A, Vec diag) noexcept
537: {
538: MatStructType *devstruct = (MatStructType *)A->spptr;
539: MultStructType *matstruct = (MultStructType *)devstruct->mat;
540: PetscScalar *darray;
541: cupmStream_t stream;
543: PetscFunctionBegin;
544: if (A->offloadmask == PETSC_OFFLOAD_BOTH || A->offloadmask == PETSC_OFFLOAD_GPU) {
545: PetscInt n = A->rmap->n;
546: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
548: PetscCheck(devstruct->format == (decltype(devstruct->format))Policy::storage_format_csr(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Only CSR format supported");
549: if (n > 0) {
550: PetscCall(Policy::VecGetArrayWrite(diag, &darray));
551: PetscCall(GetHandles_(&stream));
552: PetscCallCUPM(cupmLaunchKernel(GetDiagonal_CSR, (unsigned int)((n + 255) / 256), 256u, (size_t)0, stream, mat->row_offsets->data().get(), mat->column_indices->data().get(), mat->values->data().get(), n, darray));
553: PetscCallCUPM(cupmGetLastError());
554: PetscCall(Policy::VecRestoreArrayWrite(diag, &darray));
555: }
556: } else {
557: PetscCall(MatGetDiagonal_SeqAIJ(A, diag));
558: }
559: PetscFunctionReturn(PETSC_SUCCESS);
560: }
562: /* GetRowSumAbs: kernel-based sum of the absolute values of each CSR row */
563: static PetscErrorCode GetRowSumAbs(Mat A, Vec v) noexcept
564: {
565: PetscInt n;
567: PetscFunctionBegin;
568: PetscCheck(!A->factortype, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
569: PetscCall(VecGetLocalSize(v, &n));
570: PetscCheck(n == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming matrix and vector");
571: if (A->offloadmask == PETSC_OFFLOAD_BOTH || A->offloadmask == PETSC_OFFLOAD_GPU) {
572: MatStructType *devstruct = (MatStructType *)A->spptr;
573: MultStructType *matstruct = (MultStructType *)devstruct->mat;
574: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
575: const PetscInt m = mat->num_rows;
576: const PetscInt *cprow = matstruct->cprowIndices ? matstruct->cprowIndices->data().get() : nullptr;
577: PetscScalar *varray;
578: cupmStream_t stream;
580: PetscCheck(devstruct->format == (decltype(devstruct->format))Policy::storage_format_csr(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Only CSR format supported");
581: if (n > 0) {
582: PetscCall(Policy::VecGetArrayWrite(v, &varray));
583: PetscCall(GetHandles_(&stream));
584: // compressed rows leave the empty rows unwritten by the kernel
585: if (matstruct->cprowIndices) PetscCallCUPM(cupmMemsetAsync(varray, 0, n * sizeof(*varray), stream));
586: if (m > 0) {
587: PetscCallCUPM(cupmLaunchKernel(GetRowSumAbs_CSR, (unsigned int)((m + 255) / 256), 256u, (size_t)0, stream, mat->row_offsets->data().get(), mat->values->data().get(), cprow, m, varray));
588: PetscCallCUPM(cupmGetLastError());
589: }
590: PetscCall(Policy::VecRestoreArrayWrite(v, &varray));
591: PetscCall(PetscLogGpuFlops(mat->num_entries));
592: }
593: } else PetscCall(MatGetRowSumAbs_SeqAIJ(A, v));
594: PetscFunctionReturn(PETSC_SUCCESS);
595: }
597: /* -------------------------------------------------------------------
598: Tier 4 - Device array access (moved here from vendor files so both
599: SeqAIJCUSPARSE and SeqAIJHIPSPARSE share one implementation).
600: ------------------------------------------------------------------- */
602: /* GetArrayRead: read-only access to device CSR value array */
603: static PetscErrorCode GetArrayRead(Mat A, const PetscScalar **a) noexcept
604: {
605: MatStructType *cusp = (MatStructType *)A->spptr;
606: CsrMatrix *csr;
608: PetscFunctionBegin;
610: PetscAssertPointer(a, 2);
611: PetscCheckTypeName(A, Policy::mat_type_name);
612: PetscCheck(cusp->format != (decltype(cusp->format))Policy::storage_format_ell() && cusp->format != (decltype(cusp->format))Policy::storage_format_hyb(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
613: PetscCall(Policy::CopyToGPU(A));
614: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
615: csr = (CsrMatrix *)cusp->mat->mat;
616: PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing device memory");
617: *a = csr->values->data().get();
618: PetscFunctionReturn(PETSC_SUCCESS);
619: }
621: /* RestoreArrayRead: release read-only access obtained from GetArrayRead */
622: static PetscErrorCode RestoreArrayRead(Mat A, const PetscScalar **a) noexcept
623: {
624: PetscFunctionBegin;
626: PetscAssertPointer(a, 2);
627: PetscCheckTypeName(A, Policy::mat_type_name);
628: *a = NULL;
629: PetscFunctionReturn(PETSC_SUCCESS);
630: }
632: /* GetArray: read-write access to device CSR value array */
633: static PetscErrorCode GetArray(Mat A, PetscScalar **a) noexcept
634: {
635: MatStructType *cusp = (MatStructType *)A->spptr;
636: CsrMatrix *csr;
638: PetscFunctionBegin;
640: PetscAssertPointer(a, 2);
641: PetscCheckTypeName(A, Policy::mat_type_name);
642: PetscCheck(cusp->format != (decltype(cusp->format))Policy::storage_format_ell() && cusp->format != (decltype(cusp->format))Policy::storage_format_hyb(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
643: PetscCall(Policy::CopyToGPU(A));
644: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
645: csr = (CsrMatrix *)cusp->mat->mat;
646: PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing device memory");
647: *a = csr->values->data().get();
648: A->offloadmask = PETSC_OFFLOAD_GPU;
649: PetscCall(Policy::InvalidateTranspose(A, PETSC_FALSE));
650: PetscFunctionReturn(PETSC_SUCCESS);
651: }
653: /* RestoreArray: restore read-write access obtained from GetArray */
654: static PetscErrorCode RestoreArray(Mat A, PetscScalar **a) noexcept
655: {
656: PetscFunctionBegin;
658: PetscAssertPointer(a, 2);
659: PetscCheckTypeName(A, Policy::mat_type_name);
660: PetscCall(PetscObjectStateIncrease((PetscObject)A));
661: *a = NULL;
662: PetscFunctionReturn(PETSC_SUCCESS);
663: }
665: /* GetArrayWrite: write-only access to device CSR value array (no host-to-device copy) */
666: static PetscErrorCode GetArrayWrite(Mat A, PetscScalar **a) noexcept
667: {
668: MatStructType *cusp = (MatStructType *)A->spptr;
669: CsrMatrix *csr;
671: PetscFunctionBegin;
673: PetscAssertPointer(a, 2);
674: PetscCheckTypeName(A, Policy::mat_type_name);
675: PetscCheck(cusp->format != (decltype(cusp->format))Policy::storage_format_ell() && cusp->format != (decltype(cusp->format))Policy::storage_format_hyb(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
676: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
677: csr = (CsrMatrix *)cusp->mat->mat;
678: PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing device memory");
679: *a = csr->values->data().get();
680: A->offloadmask = PETSC_OFFLOAD_GPU;
681: PetscCall(Policy::InvalidateTranspose(A, PETSC_FALSE));
682: PetscFunctionReturn(PETSC_SUCCESS);
683: }
685: /* RestoreArrayWrite: restore write-only access obtained from GetArrayWrite */
686: static PetscErrorCode RestoreArrayWrite(Mat A, PetscScalar **a) noexcept
687: {
688: PetscFunctionBegin;
690: PetscAssertPointer(a, 2);
691: PetscCheckTypeName(A, Policy::mat_type_name);
692: PetscCall(PetscObjectStateIncrease((PetscObject)A));
693: *a = NULL;
694: PetscFunctionReturn(PETSC_SUCCESS);
695: }
697: /* SeqAIJGetArray: copy GPU-to-CPU then return host value array (ops->getarray) */
698: static PetscErrorCode SeqAIJGetArray(Mat A, PetscScalar *array[]) noexcept
699: {
700: PetscFunctionBegin;
701: PetscCall(Policy::CopyFromGPU(A));
702: *array = ((Mat_SeqAIJ *)A->data)->a;
703: PetscFunctionReturn(PETSC_SUCCESS);
704: }
706: /* SeqAIJRestoreArray: mark matrix data CPU-valid (ops->restorearray) */
707: static PetscErrorCode SeqAIJRestoreArray(Mat A, PetscScalar *array[]) noexcept
708: {
709: PetscFunctionBegin;
710: A->offloadmask = PETSC_OFFLOAD_CPU;
711: *array = NULL;
712: PetscFunctionReturn(PETSC_SUCCESS);
713: }
715: /* SeqAIJGetArrayRead: copy GPU-to-CPU then return host value array read-only (ops->getarrayread) */
716: static PetscErrorCode SeqAIJGetArrayRead(Mat A, const PetscScalar *array[]) noexcept
717: {
718: PetscFunctionBegin;
719: PetscCall(Policy::CopyFromGPU(A));
720: *array = ((Mat_SeqAIJ *)A->data)->a;
721: PetscFunctionReturn(PETSC_SUCCESS);
722: }
724: /* SeqAIJRestoreArrayRead: release read-only host array (ops->restorearrayread) */
725: static PetscErrorCode SeqAIJRestoreArrayRead(Mat /*A*/, const PetscScalar *array[]) noexcept
726: {
727: PetscFunctionBegin;
728: *array = NULL;
729: PetscFunctionReturn(PETSC_SUCCESS);
730: }
732: /* SeqAIJGetArrayWrite: return host value array for write-only access (ops->getarraywrite) */
733: static PetscErrorCode SeqAIJGetArrayWrite(Mat A, PetscScalar *array[]) noexcept
734: {
735: PetscFunctionBegin;
736: *array = ((Mat_SeqAIJ *)A->data)->a;
737: PetscFunctionReturn(PETSC_SUCCESS);
738: }
740: /* SeqAIJRestoreArrayWrite: mark matrix data CPU-valid after write (ops->restorearraywrite) */
741: static PetscErrorCode SeqAIJRestoreArrayWrite(Mat A, PetscScalar *array[]) noexcept
742: {
743: PetscFunctionBegin;
744: A->offloadmask = PETSC_OFFLOAD_CPU;
745: *array = NULL;
746: PetscFunctionReturn(PETSC_SUCCESS);
747: }
749: /* CreateSeqAIJ: allocate and preallocate a seq sparse matrix of this type */
750: static PetscErrorCode CreateSeqAIJ(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt nz, const PetscInt nnz[], Mat *A) noexcept
751: {
752: PetscFunctionBegin;
753: PetscCall(MatCreate(comm, A));
754: PetscCall(MatSetSizes(*A, m, n, m, n));
755: PetscCall(MatSetType(*A, Policy::mat_type_name));
756: PetscCall(MatSeqAIJSetPreallocation_SeqAIJ(*A, nz, (PetscInt *)nnz));
757: PetscFunctionReturn(PETSC_SUCCESS);
758: }
760: /* MatDestroy: free vendor-specific state, deregister composed functions */
761: static PetscErrorCode Destroy(Mat A) noexcept
762: {
763: PetscFunctionBegin;
764: if (A->factortype == MAT_FACTOR_NONE) PetscCall(Policy::Destroy(A));
765: else PetscCall(Policy::TriFactorsDestroy(&A->spptr));
766: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
767: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::set_format_c, NULL));
768: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::set_use_cpu_solve_c, NULL));
769: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::product_seqdense_device_c, NULL));
770: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::product_seqdense_c, NULL));
771: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::product_self_c, NULL));
772: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatFactorGetSolverType_C", NULL));
773: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
774: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
775: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::seq_convert_hypre_c, NULL));
776: PetscCall(MatDestroy_SeqAIJ(A));
777: PetscFunctionReturn(PETSC_SUCCESS);
778: }
779: };
781: } // namespace impl
783: } // namespace cupm
785: } // namespace aij
787: } // namespace mat
789: } // namespace Petsc