1: #include "../matmpidensecupm.hpp"

  3: using namespace Petsc::mat::cupm;
  4: using Petsc::device::cupm::DeviceType;

  6: static constexpr impl::MatDense_MPI_CUPM<DeviceType::HIP> mat_cupm{};

  8: /*MC
  9:   MATDENSEHIP - "densehip" - A matrix type to be used for dense matrices on GPUs.

 11:   This matrix type is identical to `MATSEQDENSEHIP` when constructed with a single process
 12:   communicator, and `MATMPIDENSEHIP` otherwise.

 14:   Options Database Key:
 15: . -mat_type densehip - sets the matrix type to `MATDENSEHIP` during a call to
 16:                         `MatSetFromOptions()`

 18:   Level: beginner

 20: .seealso: [](ch_matrices), `Mat`, `MATSEQDENSEHIP`, `MATMPIDENSEHIP`, `MATSEQDENSECUDA`,
 21: `MATMPIDENSECUDA`, `MATDENSE`
 22: M*/

 24: /*MC
 25:   MATMPIDENSEHIP - "mpidensehip" - A matrix type to be used for distributed dense matrices on
 26:   GPUs.

 28:   Options Database Key:
 29: . -mat_type mpidensehip - sets the matrix type to `MATMPIDENSEHIP` during a call to
 30:                            `MatSetFromOptions()`

 32:   Level: beginner

 34: .seealso: [](ch_matrices), `Mat`, `MATDENSEHIP`, `MATMPIDENSE`, `MATSEQDENSE`,
 35: `MATSEQDENSEHIP`, `MATSEQDENSECUDA`
 36: M*/
 37: PETSC_INTERN PetscErrorCode MatCreate_MPIDenseHIP(Mat A)
 38: {
 39:   PetscFunctionBegin;
 40:   PetscCall(mat_cupm.Create(A));
 41:   PetscFunctionReturn(PETSC_SUCCESS);
 42: }

 44: PetscErrorCode MatConvert_MPIDense_MPIDenseHIP(Mat A, MatType type, MatReuse reuse, Mat *ret)
 45: {
 46:   PetscFunctionBegin;
 47:   PetscCall(mat_cupm.Convert_MPIDense_MPIDenseCUPM(A, type, reuse, ret));
 48:   PetscFunctionReturn(PETSC_SUCCESS);
 49: }

 51: /*@C
 52:   MatCreateDenseHIP - Creates a matrix in `MATDENSEHIP` format using HIP.

 54:   Collective

 56:   Input Parameters:
 57: + comm - MPI communicator
 58: . m    - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
 59: . n    - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
 60: . M    - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
 61: . N    - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
 62: - data - optional location of GPU matrix data. Pass `NULL` to have PETSc to control matrix
 63:          memory allocation.

 65:   Output Parameter:
 66: . A - the matrix

 68:   Level: intermediate

 70: .seealso: `MATDENSEHIP`, `MatCreate()`, `MatCreateDense()`
 71: @*/
 72: PetscErrorCode MatCreateDenseHIP(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar *data, Mat *A)
 73: {
 74:   PetscFunctionBegin;
 75:   PetscCall(MatCreateDenseCUPM<DeviceType::HIP>(comm, m, n, M, N, data, A));
 76:   PetscFunctionReturn(PETSC_SUCCESS);
 77: }

 79: /*@C
 80:   MatDenseHIPPlaceArray - Allows one to replace the GPU array in a `MATDENSEHIP` matrix with an
 81:   array provided by the user. This is useful to avoid copying an array into a matrix.

 83:   Not Collective

 85:   Input Parameters:
 86: + mat   - the matrix
 87: - array - the array in column major order

 89:   Level: developer

 91:   Note:
 92:   You can return to the original array with a call to `MatDenseHIPResetArray()`. The user is
 93:   responsible for freeing this array; it will not be freed when the matrix is destroyed. The
 94:   array must have been allocated with `hipMalloc()`.

 96: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPResetArray()`,
 97: `MatDenseHIPReplaceArray()`
 98: @*/
 99: PetscErrorCode MatDenseHIPPlaceArray(Mat mat, const PetscScalar *array)
100: {
101:   PetscFunctionBegin;
102:   PetscCall(MatDenseCUPMPlaceArray<DeviceType::HIP>(mat, array));
103:   PetscFunctionReturn(PETSC_SUCCESS);
104: }

106: /*@C
107:   MatDenseHIPResetArray - Resets the matrix array to that it previously had before the call to
108:   `MatDenseHIPPlaceArray()`

110:   Not Collective

112:   Input Parameter:
113: . mat - the matrix

115:   Level: developer

117:   Note:
118:   You can only call this after a call to `MatDenseHIPPlaceArray()`

120: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPPlaceArray()`
121: @*/
122: PetscErrorCode MatDenseHIPResetArray(Mat mat)
123: {
124:   PetscFunctionBegin;
125:   PetscCall(MatDenseCUPMResetArray<DeviceType::HIP>(mat));
126:   PetscFunctionReturn(PETSC_SUCCESS);
127: }

129: /*@C
130:   MatDenseHIPReplaceArray - Allows one to replace the GPU array in a `MATDENSEHIP` matrix
131:   with an array provided by the user. This is useful to avoid copying an array into a matrix.

133:   Not Collective

135:   Input Parameters:
136: + mat   - the matrix
137: - array - the array in column major order

139:   Level: developer

141:   Note:
142:   This permanently replaces the GPU array and frees the memory associated with the old GPU
143:   array. The memory passed in CANNOT be freed by the user. It will be freed when the matrix is
144:   destroyed. The array should respect the matrix leading dimension.

146: .seealso: `MatDenseHIPGetArray()`, `MatDenseHIPPlaceArray()`, `MatDenseHIPResetArray()`
147: @*/
148: PetscErrorCode MatDenseHIPReplaceArray(Mat mat, const PetscScalar *array)
149: {
150:   PetscFunctionBegin;
151:   PetscCall(MatDenseCUPMReplaceArray<DeviceType::HIP>(mat, array));
152:   PetscFunctionReturn(PETSC_SUCCESS);
153: }

155: /*@C
156:   MatDenseHIPGetArrayWrite - Provides write access to the HIP buffer inside a `MATDENSEHIP`
157:   matrix.

159:   Not Collective

161:   Input Parameter:
162: . A - the matrix

164:   Output Parameter:
165: . a - the GPU array in column major order

167:   Level: developer

169:   Notes:
170:   The data on the GPU may not be updated due to operations done on the CPU. If you need updated
171:   data, use `MatDenseHIPGetArray()`.

173:   The array must be restored with `MatDenseHIPRestoreArrayWrite()` when no longer needed.

175: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPRestoreArray()`,
176: `MatDenseHIPRestoreArrayWrite()`, `MatDenseHIPGetArrayRead()`,
177: `MatDenseHIPRestoreArrayRead()`
178: @*/
179: PetscErrorCode MatDenseHIPGetArrayWrite(Mat A, PetscScalar **a)
180: {
181:   PetscFunctionBegin;
182:   PetscCall(MatDenseCUPMGetArrayWrite<DeviceType::HIP>(A, a));
183:   PetscFunctionReturn(PETSC_SUCCESS);
184: }

186: /*@C
187:   MatDenseHIPRestoreArrayWrite - Restore write access to the HIP buffer inside a
188:   `MATDENSEHIP` matrix previously obtained with `MatDenseHIPGetArrayWrite()`.

190:   Not Collective

192:   Input Parameters:
193: + A - the matrix
194: - a - the GPU array in column major order

196:   Level: developer

198: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPRestoreArray()`,
199: `MatDenseHIPGetArrayWrite()`, `MatDenseHIPRestoreArrayRead()`, `MatDenseHIPGetArrayRead()`
200: @*/
201: PetscErrorCode MatDenseHIPRestoreArrayWrite(Mat A, PetscScalar **a)
202: {
203:   PetscFunctionBegin;
204:   PetscCall(MatDenseCUPMRestoreArrayWrite<DeviceType::HIP>(A, a));
205:   PetscFunctionReturn(PETSC_SUCCESS);
206: }

208: /*@C
209:   MatDenseHIPGetArrayRead - Provides read-only access to the HIP buffer inside a
210:   `MATDENSEHIP` matrix. The array must be restored with `MatDenseHIPRestoreArrayRead()` when
211:   no longer needed.

213:   Not Collective

215:   Input Parameter:
216: . A - the matrix

218:   Output Parameter:
219: . a - the GPU array in column major order

221:   Level: developer

223:   Note:
224:   Data may be copied to the GPU due to operations done on the CPU. If you need write only
225:   access, use `MatDenseHIPGetArrayWrite()`.

227: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPRestoreArray()`,
228: `MatDenseHIPRestoreArrayWrite()`, `MatDenseHIPGetArrayWrite()`,
229: `MatDenseHIPRestoreArrayRead()`
230: @*/
231: PetscErrorCode MatDenseHIPGetArrayRead(Mat A, const PetscScalar **a)
232: {
233:   PetscFunctionBegin;
234:   PetscCall(MatDenseCUPMGetArrayRead<DeviceType::HIP>(A, a));
235:   PetscFunctionReturn(PETSC_SUCCESS);
236: }

238: /*@C
239:   MatDenseHIPRestoreArrayRead - Restore read-only access to the HIP buffer inside a
240:   `MATDENSEHIP` matrix previously obtained with a call to `MatDenseHIPGetArrayRead()`.

242:   Not Collective

244:   Input Parameters:
245: + A - the matrix
246: - a - the GPU array in column major order

248:   Level: developer

250:   Note:
251:   Data can be copied to the GPU due to operations done on the CPU. If you need write only
252:   access, use `MatDenseHIPGetArrayWrite()`.

254: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPRestoreArray()`,
255: `MatDenseHIPRestoreArrayWrite()`, `MatDenseHIPGetArrayWrite()`, `MatDenseHIPGetArrayRead()`
256: @*/
257: PetscErrorCode MatDenseHIPRestoreArrayRead(Mat A, const PetscScalar **a)
258: {
259:   PetscFunctionBegin;
260:   PetscCall(MatDenseCUPMRestoreArrayRead<DeviceType::HIP>(A, a));
261:   PetscFunctionReturn(PETSC_SUCCESS);
262: }

264: /*@C
265:   MatDenseHIPGetArray - Provides access to the HIP buffer inside a `MATDENSEHIP` matrix. The
266:   array must be restored with `MatDenseHIPRestoreArray()` when no longer needed.

268:   Not Collective

270:   Input Parameter:
271: . A - the matrix

273:   Output Parameter:
274: . a - the GPU array in column major order

276:   Level: developer

278:   Note:
279:   Data can be copied to the GPU due to operations done on the CPU. If you need write only
280:   access, use `MatDenseHIPGetArrayWrite()`. For read-only access, use
281:   `MatDenseHIPGetArrayRead()`.

283: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArrayRead()`, `MatDenseHIPRestoreArray()`,
284: `MatDenseHIPRestoreArrayWrite()`, `MatDenseHIPGetArrayWrite()`,
285: `MatDenseHIPRestoreArrayRead()`
286: @*/
287: PetscErrorCode MatDenseHIPGetArray(Mat A, PetscScalar **a)
288: {
289:   PetscFunctionBegin;
290:   PetscCall(MatDenseCUPMGetArray<DeviceType::HIP>(A, a));
291:   PetscFunctionReturn(PETSC_SUCCESS);
292: }

294: /*@C
295:   MatDenseHIPRestoreArray - Restore access to the HIP buffer inside a `MATDENSEHIP` matrix
296:   previously obtained with `MatDenseHIPGetArray()`.

298:   Not Collective

300:   Level: developer

302:   Input Parameters:
303: + A - the matrix
304: - a - the GPU array in column major order

306: .seealso: `MATDENSEHIP`, `MatDenseHIPGetArray()`, `MatDenseHIPRestoreArrayWrite()`,
307: `MatDenseHIPGetArrayWrite()`, `MatDenseHIPRestoreArrayRead()`, `MatDenseHIPGetArrayRead()`
308: @*/
309: PetscErrorCode MatDenseHIPRestoreArray(Mat A, PetscScalar **a)
310: {
311:   PetscFunctionBegin;
312:   PetscCall(MatDenseCUPMRestoreArray<DeviceType::HIP>(A, a));
313:   PetscFunctionReturn(PETSC_SUCCESS);
314: }

316: /*@C
317:   MatDenseHIPSetPreallocation - Set the device array used for storing the matrix elements of a
318:   `MATDENSEHIP` matrix

320:   Collective

322:   Input Parameters:
323: + A            - the matrix
324: - device_array - the array (or `NULL`)

326:   Level: intermediate

328: .seealso: [](ch_matrices), `Mat`, `MATDENSEHIP`, `MatCreate()`, `MatCreateDenseHIP()`,
329: `MatSetValues()`, `MatDenseSetLDA()`
330: @*/
331: PetscErrorCode MatDenseHIPSetPreallocation(Mat A, PetscScalar *device_array)
332: {
333:   PetscFunctionBegin;
334:   PetscCall(MatDenseCUPMSetPreallocation<DeviceType::HIP>(A, device_array));
335:   PetscFunctionReturn(PETSC_SUCCESS);
336: }