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: MatSeqAIJCUSPARSE_CUPM<T, Policy>
136: Policy (C++11 traits class) requirements - all static methods:
138: // Device struct types
139: typedef ... mat_struct_type; // Mat_SeqAIJCUSPARSE / Mat_SeqAIJHIPSPARSE
140: typedef ... mult_struct_type; // ...MultStruct equivalent
142: // Storage-format constants (value of each format enumerator)
143: static int storage_format_csr();
144: static int storage_format_ell();
145: static int storage_format_hyb();
147: // Bookkeeping helpers (device-type specific)
148: static PetscErrorCode CopyToGPU(Mat);
149: static PetscErrorCode CopyFromGPU(Mat);
150: static PetscErrorCode InvalidateTranspose(Mat, PetscBool);
151: static PetscErrorCode ConvertFromSeqAIJ(Mat, MatType, MatReuse, Mat *);
152: static const char *mat_type_name; // "seqaijcusparse" / "seqaijhipsparse"
154: // Destruction helpers (device-type specific)
155: static PetscErrorCode Destroy(Mat);
156: static PetscErrorCode TriFactorsDestroy(void **);
158: // Compose-function keys that differ between CUDA and HIP
159: static const char *set_format_c; // "MatCUSPARSESetFormat_C" / "MatHIPSPARSESetFormat_C"
160: static const char *set_use_cpu_solve_c; // "MatCUSPARSESetUseCPUSolve_C" / "MatHIPSPARSESetUseCPUSolve_C"
161: static const char *product_seqdense_device_c; // "...seqdensecuda_C" / "...seqdensehip_C"
162: static const char *product_seqdense_c; // "...seqdense_C"
163: static const char *product_self_c; // "...seqaijcusparse_C" / "...seqaijhipsparse_C"
164: static const char *seq_convert_hypre_c; // "MatConvert_seqaijcusparse_hypre_C" / "_seqaijhipsparse_hypre_C"
166: // Vec device-array access (device-type specific)
167: static PetscErrorCode VecGetArrayRead (Vec, const PetscScalar **);
168: static PetscErrorCode VecRestoreArrayRead(Vec, const PetscScalar **);
169: static PetscErrorCode VecGetArrayWrite (Vec, PetscScalar **);
170: static PetscErrorCode VecRestoreArrayWrite(Vec, PetscScalar **);
171: ========================================================================== */
173: template <device::cupm::DeviceType T, typename Policy>
174: struct MatSeqAIJCUSPARSE_CUPM : device::cupm::impl::CUPMObject<T> {
175: PETSC_CUPMOBJECT_HEADER(T);
177: typedef typename Policy::mat_struct_type MatStructType;
178: typedef typename Policy::mult_struct_type MultStructType;
180: /* -------------------------------------------------------------------
181: Tier 1 - Trivial
182: ------------------------------------------------------------------- */
184: /* MatAssemblyEnd: delegation to SeqAIJ */
185: static PetscErrorCode AssemblyEnd(Mat A, MatAssemblyType mode) noexcept
186: {
187: PetscFunctionBegin;
188: PetscCall(MatAssemblyEnd_SeqAIJ(A, mode));
189: PetscFunctionReturn(PETSC_SUCCESS);
190: }
192: /* MatDuplicate */
193: static PetscErrorCode Duplicate(Mat A, MatDuplicateOption cpvalues, Mat *B) noexcept
194: {
195: PetscFunctionBegin;
196: PetscCall(MatDuplicate_SeqAIJ(A, cpvalues, B));
197: PetscCall(Policy::ConvertFromSeqAIJ(*B, Policy::mat_type_name, MAT_INPLACE_MATRIX, B));
198: PetscFunctionReturn(PETSC_SUCCESS);
199: }
201: /* MatGetCurrentMemType */
202: static PetscErrorCode GetCurrentMemType(PETSC_UNUSED Mat A, PetscMemType *m) noexcept
203: {
204: PetscFunctionBegin;
205: *m = PETSC_MEMTYPE_CUPM();
206: PetscFunctionReturn(PETSC_SUCCESS);
207: }
209: /* MatCOOStructDestroy: free device jmap and perm fields */
210: static PetscErrorCode COOStructDestroy(PetscCtxRt ctx) noexcept
211: {
212: MatCOOStruct_SeqAIJ *coo = *(MatCOOStruct_SeqAIJ **)ctx;
214: PetscFunctionBegin;
215: PetscCallCUPM(cupmFree(coo->perm));
216: PetscCallCUPM(cupmFree(coo->jmap));
217: PetscCall(PetscFree(coo));
218: PetscFunctionReturn(PETSC_SUCCESS);
219: }
221: /* -------------------------------------------------------------------
222: Tier 2 - Straightforward
223: ------------------------------------------------------------------- */
225: /* MatZeroEntries: fill device CSR values with zero */
226: static PetscErrorCode ZeroEntries(Mat A) noexcept
227: {
228: PetscBool gpu = PETSC_FALSE;
229: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
230: MatStructType *spptr;
232: PetscFunctionBegin;
233: if (A->factortype == MAT_FACTOR_NONE) {
234: spptr = (MatStructType *)A->spptr;
235: if (spptr->mat) {
236: CsrMatrix *matrix = (CsrMatrix *)spptr->mat->mat;
237: if (matrix->values) {
238: gpu = PETSC_TRUE;
239: PetscCallThrust(thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), (PetscScalar)0.));
240: }
241: }
242: if (spptr->matTranspose) {
243: CsrMatrix *matrix = (CsrMatrix *)spptr->matTranspose->mat;
244: if (matrix->values) PetscCallThrust(thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), (PetscScalar)0.));
245: }
246: }
247: if (gpu) A->offloadmask = PETSC_OFFLOAD_GPU;
248: else {
249: PetscCall(PetscArrayzero(a->a, a->i[A->rmap->n]));
250: A->offloadmask = PETSC_OFFLOAD_CPU;
251: }
252: PetscFunctionReturn(PETSC_SUCCESS);
253: }
255: /* MatScale: cupmBlasXscal on the device CSR values */
256: static PetscErrorCode Scale(Mat Y, PetscScalar a) noexcept
257: {
258: Mat_SeqAIJ *y = (Mat_SeqAIJ *)Y->data;
259: PetscScalar *ay = nullptr;
260: cupmBlasHandle_t blashandle;
261: PetscBLASInt one = 1, bnz = 1;
263: PetscFunctionBegin;
264: PetscCall(GetArray(Y, &ay));
265: PetscCall(GetHandles_(&blashandle));
266: PetscCall(PetscBLASIntCast(y->nz, &bnz));
267: PetscCall(PetscLogGpuTimeBegin());
268: PetscCallCUPMBLAS(cupmBlasXscal(blashandle, bnz, cupmScalarPtrCast(&a), cupmScalarPtrCast(ay), one));
269: PetscCall(PetscLogGpuFlops(bnz));
270: PetscCall(PetscLogGpuTimeEnd());
271: PetscCall(RestoreArray(Y, &ay));
272: PetscFunctionReturn(PETSC_SUCCESS);
273: }
275: /* MatDiagonalScale: Thrust-based left and right scaling of CSR values */
276: static PetscErrorCode DiagonalScale(Mat A, Vec ll, Vec rr) noexcept
277: {
278: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)A->data;
279: MatStructType *devstruct;
280: CsrMatrix *csr;
281: PetscScalar *av = nullptr;
282: PetscInt m, n, nz = aij->nz;
283: cupmStream_t stream;
285: PetscFunctionBegin;
286: PetscCall(GetHandles_(&stream));
287: PetscCall(PetscLogGpuTimeBegin());
288: PetscCall(GetArray(A, &av));
289: devstruct = (MatStructType *)A->spptr;
290: csr = (CsrMatrix *)devstruct->mat->mat;
291: if (ll) {
292: const PetscScalar *lv;
293: PetscCall(VecGetLocalSize(ll, &m));
294: PetscCheck(m == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector wrong length");
295: PetscCall(Policy::VecGetArrayRead(ll, &lv));
296: {
297: const PetscInt *cprow = devstruct->mat->cprowIndices ? devstruct->mat->cprowIndices->data().get() : NULL;
298: DiagonalScaleLeft_CSR_Functor functor = {csr->row_offsets->data().get(), av, lv, cprow};
299: PetscCallThrust(THRUST_CALL(thrust::for_each, stream, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(csr->num_rows), functor));
300: }
301: PetscCall(Policy::VecRestoreArrayRead(ll, &lv));
302: PetscCall(PetscLogGpuFlops(nz));
303: }
304: if (rr) {
305: const PetscScalar *rv;
306: PetscCall(VecGetLocalSize(rr, &n));
307: PetscCheck(n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vector wrong length");
308: PetscCall(Policy::VecGetArrayRead(rr, &rv));
309: #if PetscDefined(USING_NVCC) && CCCL_VERSION >= 3001000
310: 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>()));
311: #else
312: 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>()));
313: #endif
314: PetscCall(Policy::VecRestoreArrayRead(rr, &rv));
315: PetscCall(PetscLogGpuFlops(nz));
316: }
317: PetscCall(RestoreArray(A, &av));
318: PetscCall(PetscLogGpuTimeEnd());
319: PetscFunctionReturn(PETSC_SUCCESS);
320: }
322: /* MatSeqAIJGetIJ: return device CSR row-pointer and column-index arrays */
323: static PetscErrorCode GetIJ(Mat A, PetscBool compressed, const PetscInt **i, const PetscInt **j) noexcept
324: {
325: MatStructType *cusp = (MatStructType *)A->spptr;
326: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
327: CsrMatrix *csr;
329: PetscFunctionBegin;
331: if (!i || !j) PetscFunctionReturn(PETSC_SUCCESS);
332: PetscCheckTypeName(A, Policy::mat_type_name);
333: 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");
334: PetscCall(Policy::CopyToGPU(A));
335: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
336: csr = (CsrMatrix *)cusp->mat->mat;
337: if (i) {
338: if (!compressed && a->compressedrow.use) { /* need full row offset */
339: if (!cusp->rowoffsets_gpu) {
340: cusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
341: cusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
342: PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
343: }
344: *i = cusp->rowoffsets_gpu->data().get();
345: } else *i = csr->row_offsets->data().get();
346: }
347: if (j) *j = csr->column_indices->data().get();
348: PetscFunctionReturn(PETSC_SUCCESS);
349: }
351: /* MatSeqAIJRestoreIJ: nullify the pointers previously obtained with GetIJ */
352: static PetscErrorCode RestoreIJ(Mat A, PetscBool compressed, const PetscInt **i, const PetscInt **j) noexcept
353: {
354: PetscFunctionBegin;
356: PetscCheckTypeName(A, Policy::mat_type_name);
357: if (i) *i = NULL;
358: if (j) *j = NULL;
359: (void)compressed;
360: PetscFunctionReturn(PETSC_SUCCESS);
361: }
363: /* MatSetPreallocationCOO: copy COO bookkeeping struct to device */
364: static PetscErrorCode SetPreallocationCOO(Mat mat, PetscCount coo_n, PetscInt coo_i[], PetscInt coo_j[]) noexcept
365: {
366: PetscBool dev_ij = PETSC_FALSE;
367: PetscMemType mtype = PETSC_MEMTYPE_HOST;
368: PetscInt *i, *j;
369: PetscContainer container_h;
370: MatCOOStruct_SeqAIJ *coo_h, *coo_d;
372: PetscFunctionBegin;
373: PetscCall(PetscGetMemType(coo_i, &mtype));
374: if (PetscMemTypeDevice(mtype)) {
375: dev_ij = PETSC_TRUE;
376: PetscCall(PetscMalloc2(coo_n, &i, coo_n, &j));
377: PetscCallCUPM(cupmMemcpy(i, coo_i, coo_n * sizeof(PetscInt), cupmMemcpyDeviceToHost));
378: PetscCallCUPM(cupmMemcpy(j, coo_j, coo_n * sizeof(PetscInt), cupmMemcpyDeviceToHost));
379: } else {
380: i = coo_i;
381: j = coo_j;
382: }
383: PetscCall(MatSetPreallocationCOO_SeqAIJ(mat, coo_n, i, j));
384: if (dev_ij) PetscCall(PetscFree2(i, j));
385: mat->offloadmask = PETSC_OFFLOAD_CPU;
386: /* Create the GPU memory */
387: PetscCall(Policy::CopyToGPU(mat));
389: /* Copy the COO struct to device */
390: PetscCall(PetscObjectQuery((PetscObject)mat, "__PETSc_MatCOOStruct_Host", (PetscObject *)&container_h));
391: PetscCall(PetscContainerGetPointer(container_h, (void **)&coo_h));
392: PetscCall(PetscMalloc1(1, &coo_d));
393: *coo_d = *coo_h; /* shallow copy; device fields amended below */
394: PetscCallCUPM(cupmMalloc((void **)&coo_d->jmap, (coo_h->nz + 1) * sizeof(PetscCount)));
395: PetscCallCUPM(cupmMemcpy(coo_d->jmap, coo_h->jmap, (coo_h->nz + 1) * sizeof(PetscCount), cupmMemcpyHostToDevice));
396: PetscCallCUPM(cupmMalloc((void **)&coo_d->perm, coo_h->Atot * sizeof(PetscCount)));
397: PetscCallCUPM(cupmMemcpy(coo_d->perm, coo_h->perm, coo_h->Atot * sizeof(PetscCount), cupmMemcpyHostToDevice));
399: PetscCall(PetscObjectContainerCompose((PetscObject)mat, "__PETSc_MatCOOStruct_Device", coo_d, MatSeqAIJCUSPARSE_CUPM::COOStructDestroy));
400: PetscFunctionReturn(PETSC_SUCCESS);
401: }
403: /* MatSetValuesCOO: launch MatAddCOOValues kernel */
404: static PetscErrorCode SetValuesCOO(Mat A, const PetscScalar v[], InsertMode imode) noexcept
405: {
406: Mat_SeqAIJ *seq = (Mat_SeqAIJ *)A->data;
407: MatStructType *dev = (MatStructType *)A->spptr;
408: PetscCount Annz = seq->nz;
409: PetscMemType memtype;
410: const PetscScalar *v1 = v;
411: PetscScalar *Aa = nullptr;
412: PetscContainer container;
413: MatCOOStruct_SeqAIJ *coo;
414: cupmStream_t stream;
416: PetscFunctionBegin;
417: if (!dev->mat) PetscCall(Policy::CopyToGPU(A));
419: PetscCall(PetscObjectQuery((PetscObject)A, "__PETSc_MatCOOStruct_Device", (PetscObject *)&container));
420: PetscCall(PetscContainerGetPointer(container, (void **)&coo));
422: PetscCall(PetscGetMemType(v, &memtype));
423: if (PetscMemTypeHost(memtype)) { /* copy host values to device */
424: PetscCallCUPM(cupmMalloc((void **)&v1, coo->n * sizeof(PetscScalar)));
425: PetscCallCUPM(cupmMemcpy((void *)v1, v, coo->n * sizeof(PetscScalar), cupmMemcpyHostToDevice));
426: PetscCall(PetscLogCpuToGpu(coo->n * sizeof(PetscScalar)));
427: }
429: if (imode == INSERT_VALUES) PetscCall(GetArrayWrite(A, &Aa));
430: else PetscCall(GetArray(A, &Aa));
432: PetscCall(GetHandles_(&stream));
433: PetscCall(PetscLogGpuTimeBegin());
434: if (Annz) {
435: PetscCallCUPM(cupmLaunchKernel(MatAddCOOValues, (unsigned int)((Annz + 255) / 256), 256u, (size_t)0, stream, v1, Annz, coo->jmap, coo->perm, imode, Aa));
436: PetscCallCUPM(cupmGetLastError());
437: }
438: PetscCall(PetscLogGpuTimeEnd());
440: if (imode == INSERT_VALUES) PetscCall(RestoreArrayWrite(A, &Aa));
441: else PetscCall(RestoreArray(A, &Aa));
443: if (PetscMemTypeHost(memtype)) {
444: void *v1_device = (void *)v1;
445: PetscCallCUPM(cupmFree(v1_device));
446: }
447: PetscFunctionReturn(PETSC_SUCCESS);
448: }
450: /* MatSeqAIJCopySubArray: scatter-gather a sub-array of CSR values */
451: static PetscErrorCode CopySubArray(Mat A, PetscInt n, const PetscInt idx[], PetscScalar v[]) noexcept
452: {
453: const PetscScalar *av = nullptr;
454: PetscMemType mtype;
455: PetscBool dmem;
457: PetscFunctionBegin;
458: PetscCall(PetscCUPMGetMemType(v, &mtype));
459: dmem = PetscMemTypeDevice(mtype);
460: PetscCall(GetArrayRead(A, &av));
461: if (n && idx) {
462: THRUSTINTARRAY widx(n);
463: widx.assign(idx, idx + n);
464: PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
466: THRUSTARRAY *w = NULL;
467: thrust::device_ptr<PetscScalar> dv;
468: if (dmem) {
469: dv = thrust::device_pointer_cast(v);
470: } else {
471: w = new THRUSTARRAY(n);
472: dv = w->data();
473: }
474: {
475: thrust::device_ptr<const PetscScalar> dav = thrust::device_pointer_cast(av);
476: auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(dav, widx.begin()), dv));
477: auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(dav, widx.end()), dv + n));
478: PetscCallThrust(thrust::for_each(zibit, zieit, VecCUPMEquals{}));
479: }
480: if (w) PetscCallCUPM(cupmMemcpy(v, w->data().get(), n * sizeof(PetscScalar), cupmMemcpyDeviceToHost));
481: delete w;
482: } else {
483: PetscCallCUPM(cupmMemcpy(v, av, n * sizeof(PetscScalar), dmem ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost));
484: }
485: if (!dmem) PetscCall(PetscLogCpuToGpu(n * sizeof(PetscScalar)));
486: PetscCall(RestoreArrayRead(A, &av));
487: PetscFunctionReturn(PETSC_SUCCESS);
488: }
490: /* -------------------------------------------------------------------
491: Tier 3 - AXPY shared branches (SAME_NZ and DIFFERENT_NZ only).
492: The SUBSET_NZ branch calls cuSPARSE/hipSPARSE and stays in the caller.
493: ------------------------------------------------------------------- */
495: /* AXPY SAME_NONZERO_PATTERN branch: cupmBlasXaxpy */
496: static PetscErrorCode AXPY_SameNZ(Mat Y, PetscScalar a, Mat X) noexcept
497: {
498: Mat_SeqAIJ *x = (Mat_SeqAIJ *)X->data;
499: const PetscScalar *ax = nullptr;
500: PetscScalar *ay = nullptr;
501: cupmBlasHandle_t blashandle;
502: PetscBLASInt one = 1, bnz = 1;
504: PetscFunctionBegin;
505: PetscCall(GetArrayRead(X, &ax));
506: PetscCall(GetArray(Y, &ay));
507: PetscCall(GetHandles_(&blashandle));
508: PetscCall(PetscBLASIntCast(x->nz, &bnz));
509: PetscCall(PetscLogGpuTimeBegin());
510: PetscCallCUPMBLAS(cupmBlasXaxpy(blashandle, bnz, cupmScalarPtrCast(&a), cupmScalarPtrCast(ax), one, cupmScalarPtrCast(ay), one));
511: PetscCall(PetscLogGpuFlops(2.0 * bnz));
512: PetscCall(PetscLogGpuTimeEnd());
513: PetscCall(RestoreArrayRead(X, &ax));
514: PetscCall(RestoreArray(Y, &ay));
515: PetscFunctionReturn(PETSC_SUCCESS);
516: }
518: /* GetDiagonal: kernel-based extraction of the CSR diagonal */
519: static PetscErrorCode GetDiagonal(Mat A, Vec diag) noexcept
520: {
521: MatStructType *devstruct = (MatStructType *)A->spptr;
522: MultStructType *matstruct = (MultStructType *)devstruct->mat;
523: PetscScalar *darray;
524: cupmStream_t stream;
526: PetscFunctionBegin;
527: if (A->offloadmask == PETSC_OFFLOAD_BOTH || A->offloadmask == PETSC_OFFLOAD_GPU) {
528: PetscInt n = A->rmap->n;
529: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
531: PetscCheck(devstruct->format == (decltype(devstruct->format))Policy::storage_format_csr(), PETSC_COMM_SELF, PETSC_ERR_SUP, "Only CSR format supported");
532: if (n > 0) {
533: PetscCall(Policy::VecGetArrayWrite(diag, &darray));
534: PetscCall(GetHandles_(&stream));
535: 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));
536: PetscCallCUPM(cupmGetLastError());
537: PetscCall(Policy::VecRestoreArrayWrite(diag, &darray));
538: }
539: } else {
540: PetscCall(MatGetDiagonal_SeqAIJ(A, diag));
541: }
542: PetscFunctionReturn(PETSC_SUCCESS);
543: }
545: /* -------------------------------------------------------------------
546: Tier 4 - Device array access (moved here from vendor files so both
547: SeqAIJCUSPARSE and SeqAIJHIPSPARSE share one implementation).
548: ------------------------------------------------------------------- */
550: /* GetArrayRead: read-only access to device CSR value array */
551: static PetscErrorCode GetArrayRead(Mat A, const PetscScalar **a) noexcept
552: {
553: MatStructType *cusp = (MatStructType *)A->spptr;
554: CsrMatrix *csr;
556: PetscFunctionBegin;
558: PetscAssertPointer(a, 2);
559: PetscCheckTypeName(A, Policy::mat_type_name);
560: 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");
561: PetscCall(Policy::CopyToGPU(A));
562: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
563: csr = (CsrMatrix *)cusp->mat->mat;
564: PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing device memory");
565: *a = csr->values->data().get();
566: PetscFunctionReturn(PETSC_SUCCESS);
567: }
569: /* RestoreArrayRead: release read-only access obtained from GetArrayRead */
570: static PetscErrorCode RestoreArrayRead(Mat A, const PetscScalar **a) noexcept
571: {
572: PetscFunctionBegin;
574: PetscAssertPointer(a, 2);
575: PetscCheckTypeName(A, Policy::mat_type_name);
576: *a = NULL;
577: PetscFunctionReturn(PETSC_SUCCESS);
578: }
580: /* GetArray: read-write access to device CSR value array */
581: static PetscErrorCode GetArray(Mat A, PetscScalar **a) noexcept
582: {
583: MatStructType *cusp = (MatStructType *)A->spptr;
584: CsrMatrix *csr;
586: PetscFunctionBegin;
588: PetscAssertPointer(a, 2);
589: PetscCheckTypeName(A, Policy::mat_type_name);
590: 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");
591: PetscCall(Policy::CopyToGPU(A));
592: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
593: csr = (CsrMatrix *)cusp->mat->mat;
594: PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing device memory");
595: *a = csr->values->data().get();
596: A->offloadmask = PETSC_OFFLOAD_GPU;
597: PetscCall(Policy::InvalidateTranspose(A, PETSC_FALSE));
598: PetscFunctionReturn(PETSC_SUCCESS);
599: }
601: /* RestoreArray: restore read-write access obtained from GetArray */
602: static PetscErrorCode RestoreArray(Mat A, PetscScalar **a) noexcept
603: {
604: PetscFunctionBegin;
606: PetscAssertPointer(a, 2);
607: PetscCheckTypeName(A, Policy::mat_type_name);
608: PetscCall(PetscObjectStateIncrease((PetscObject)A));
609: *a = NULL;
610: PetscFunctionReturn(PETSC_SUCCESS);
611: }
613: /* GetArrayWrite: write-only access to device CSR value array (no host-to-device copy) */
614: static PetscErrorCode GetArrayWrite(Mat A, PetscScalar **a) noexcept
615: {
616: MatStructType *cusp = (MatStructType *)A->spptr;
617: CsrMatrix *csr;
619: PetscFunctionBegin;
621: PetscAssertPointer(a, 2);
622: PetscCheckTypeName(A, Policy::mat_type_name);
623: 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");
624: PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing MultStruct");
625: csr = (CsrMatrix *)cusp->mat->mat;
626: PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing device memory");
627: *a = csr->values->data().get();
628: A->offloadmask = PETSC_OFFLOAD_GPU;
629: PetscCall(Policy::InvalidateTranspose(A, PETSC_FALSE));
630: PetscFunctionReturn(PETSC_SUCCESS);
631: }
633: /* RestoreArrayWrite: restore write-only access obtained from GetArrayWrite */
634: static PetscErrorCode RestoreArrayWrite(Mat A, PetscScalar **a) noexcept
635: {
636: PetscFunctionBegin;
638: PetscAssertPointer(a, 2);
639: PetscCheckTypeName(A, Policy::mat_type_name);
640: PetscCall(PetscObjectStateIncrease((PetscObject)A));
641: *a = NULL;
642: PetscFunctionReturn(PETSC_SUCCESS);
643: }
645: /* SeqAIJGetArray: copy GPU-to-CPU then return host value array (ops->getarray) */
646: static PetscErrorCode SeqAIJGetArray(Mat A, PetscScalar *array[]) noexcept
647: {
648: PetscFunctionBegin;
649: PetscCall(Policy::CopyFromGPU(A));
650: *array = ((Mat_SeqAIJ *)A->data)->a;
651: PetscFunctionReturn(PETSC_SUCCESS);
652: }
654: /* SeqAIJRestoreArray: mark matrix data CPU-valid (ops->restorearray) */
655: static PetscErrorCode SeqAIJRestoreArray(Mat A, PetscScalar *array[]) noexcept
656: {
657: PetscFunctionBegin;
658: A->offloadmask = PETSC_OFFLOAD_CPU;
659: *array = NULL;
660: PetscFunctionReturn(PETSC_SUCCESS);
661: }
663: /* SeqAIJGetArrayRead: copy GPU-to-CPU then return host value array read-only (ops->getarrayread) */
664: static PetscErrorCode SeqAIJGetArrayRead(Mat A, const PetscScalar *array[]) noexcept
665: {
666: PetscFunctionBegin;
667: PetscCall(Policy::CopyFromGPU(A));
668: *array = ((Mat_SeqAIJ *)A->data)->a;
669: PetscFunctionReturn(PETSC_SUCCESS);
670: }
672: /* SeqAIJRestoreArrayRead: release read-only host array (ops->restorearrayread) */
673: static PetscErrorCode SeqAIJRestoreArrayRead(Mat /*A*/, const PetscScalar *array[]) noexcept
674: {
675: PetscFunctionBegin;
676: *array = NULL;
677: PetscFunctionReturn(PETSC_SUCCESS);
678: }
680: /* SeqAIJGetArrayWrite: return host value array for write-only access (ops->getarraywrite) */
681: static PetscErrorCode SeqAIJGetArrayWrite(Mat A, PetscScalar *array[]) noexcept
682: {
683: PetscFunctionBegin;
684: *array = ((Mat_SeqAIJ *)A->data)->a;
685: PetscFunctionReturn(PETSC_SUCCESS);
686: }
688: /* SeqAIJRestoreArrayWrite: mark matrix data CPU-valid after write (ops->restorearraywrite) */
689: static PetscErrorCode SeqAIJRestoreArrayWrite(Mat A, PetscScalar *array[]) noexcept
690: {
691: PetscFunctionBegin;
692: A->offloadmask = PETSC_OFFLOAD_CPU;
693: *array = NULL;
694: PetscFunctionReturn(PETSC_SUCCESS);
695: }
697: /* CreateSeqAIJ: allocate and preallocate a seq sparse matrix of this type */
698: static PetscErrorCode CreateSeqAIJ(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt nz, const PetscInt nnz[], Mat *A) noexcept
699: {
700: PetscFunctionBegin;
701: PetscCall(MatCreate(comm, A));
702: PetscCall(MatSetSizes(*A, m, n, m, n));
703: PetscCall(MatSetType(*A, Policy::mat_type_name));
704: PetscCall(MatSeqAIJSetPreallocation_SeqAIJ(*A, nz, (PetscInt *)nnz));
705: PetscFunctionReturn(PETSC_SUCCESS);
706: }
708: /* MatDestroy: free vendor-specific state, deregister composed functions */
709: static PetscErrorCode Destroy(Mat A) noexcept
710: {
711: PetscFunctionBegin;
712: if (A->factortype == MAT_FACTOR_NONE) PetscCall(Policy::Destroy(A));
713: else PetscCall(Policy::TriFactorsDestroy(&A->spptr));
714: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
715: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::set_format_c, NULL));
716: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::set_use_cpu_solve_c, NULL));
717: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::product_seqdense_device_c, NULL));
718: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::product_seqdense_c, NULL));
719: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::product_self_c, NULL));
720: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatFactorGetSolverType_C", NULL));
721: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
722: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
723: PetscCall(PetscObjectComposeFunction((PetscObject)A, Policy::seq_convert_hypre_c, NULL));
724: PetscCall(MatDestroy_SeqAIJ(A));
725: PetscFunctionReturn(PETSC_SUCCESS);
726: }
727: };
729: } // namespace impl
731: } // namespace cupm
733: } // namespace aij
735: } // namespace mat
737: } // namespace Petsc