Actual source code: ex319.c
1: static char help[] = "Tests MatMatMult() and MatTransposeMatMult() of an MPIAIJ-family matrix with an MPIDENSE-family matrix, including the batched path, against a host reference\n\n";
3: // Contributed by: Steven Dargaville
5: #include <petscmat.h>
6: #include <petsc/private/petscimpl.h>
8: /* The product of the matrices whose types are set with -A_mat_type and -B_mat_type is compared against
9: the product of host MATAIJ and MATDENSE copies holding the same values */
10: static PetscErrorCode CheckProduct(Mat C, Mat Ch, const char *label)
11: {
12: Mat Ct;
13: PetscReal nrm, err;
15: PetscFunctionBeginUser;
16: /* a host matrix is duplicated, a device one is copied back to the host */
17: PetscCall(MatConvert(C, MATDENSE, MAT_INITIAL_MATRIX, &Ct));
18: PetscCall(MatNorm(Ch, NORM_INFINITY, &nrm));
19: PetscCall(MatAXPY(Ct, -1.0, Ch, SAME_NONZERO_PATTERN));
20: PetscCall(MatNorm(Ct, NORM_INFINITY, &err));
21: PetscCheck(err <= 100.0 * PETSC_SMALL * PetscMax(nrm, 1.0), PETSC_COMM_WORLD, PETSC_ERR_PLIB, "%s: error %g relative to %g", label, (double)err, (double)nrm);
22: PetscCall(MatDestroy(&Ct));
23: PetscFunctionReturn(PETSC_SUCCESS);
24: }
26: // C = A * B, or C = A^T * B if transpose is true
27: static PetscErrorCode Product(Mat A, Mat B, MatReuse reuse, PetscBool transpose, Mat *C)
28: {
29: PetscFunctionBeginUser;
30: if (transpose) PetscCall(MatTransposeMatMult(A, B, reuse, PETSC_DETERMINE, C));
31: else PetscCall(MatMatMult(A, B, reuse, PETSC_DETERMINE, C));
32: PetscFunctionReturn(PETSC_SUCCESS);
33: }
35: // Ch comes from the same batching and reduction code as C, so also check it against matrix-vector products
36: static PetscErrorCode CheckTransposeReference(Mat Ah, Mat Bh, Mat Ch, const char *label)
37: {
38: PetscBool flg;
40: PetscFunctionBeginUser;
41: PetscCall(MatTransposeMatMultEqual(Ah, Bh, Ch, 10, &flg));
42: PetscCheck(flg, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "%s: wrong host reference", label);
43: PetscFunctionReturn(PETSC_SUCCESS);
44: }
46: // Reuse Ch = Ah^T * Bh with an output whose leading dimension changes, on one process only and then on all of them
47: static PetscErrorCode CheckChangedLDA(Mat Ah, Mat Bh, Mat Ch)
48: {
49: Mat C;
50: PetscScalar *store;
51: PetscInt m, n, M, N;
52: PetscMPIInt rank;
54: PetscFunctionBeginUser;
55: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
56: PetscCall(MatGetLocalSize(Ch, &m, &n));
57: PetscCall(MatGetSize(Ch, &M, &N));
58: PetscCall(PetscMalloc1((m + 3) * N, &store));
59: PetscCall(MatCreateDense(PETSC_COMM_WORLD, m, n, M, N, store, &C));
60: PetscCall(MatDenseSetLDA(C, m + 1));
61: PetscCall(MatProductCreateWithMat(Ah, Bh, NULL, C));
62: PetscCall(MatProductSetType(C, MATPRODUCT_AtB));
63: PetscCall(MatProductSetFromOptions(C));
64: PetscCall(MatProductSymbolic(C));
65: PetscCall(MatProductNumeric(C));
66: PetscCall(CheckProduct(C, Ch, "lda m + 1"));
67: if (rank == 0) PetscCall(MatDenseSetLDA(C, m + 3));
68: PetscCall(MatProductNumeric(C));
69: PetscCall(CheckProduct(C, Ch, "lda m + 3 on the first process"));
70: PetscCall(MatDenseSetLDA(C, m + 2));
71: PetscCall(MatProductNumeric(C));
72: PetscCall(CheckProduct(C, Ch, "lda m + 2"));
73: PetscCall(MatDenseSetLDA(C, m + 1));
74: PetscCall(MatProductNumeric(C));
75: PetscCall(CheckProduct(C, Ch, "lda m + 1 again"));
76: PetscCall(MatDestroy(&C));
77: PetscCall(PetscFree(store));
78: PetscFunctionReturn(PETSC_SUCCESS);
79: }
81: int main(int argc, char **args)
82: {
83: Mat A, B, C, Ah, Bh, Ch;
84: char atype[256], btype[256];
85: PetscScalar vals[4];
86: PetscInt m = 20, n = 7, rstart, rend, ncols, cols[4];
87: PetscBool block_diagonal = PETSC_FALSE, change_lda = PETSC_FALSE, check_copies = PETSC_FALSE, tridiagonal = PETSC_FALSE, transpose = PETSC_FALSE;
88: PetscLogEvent event;
89: PetscMPIInt size;
91: PetscFunctionBeginUser;
92: PetscCall(PetscInitialize(&argc, &args, NULL, help));
93: PetscCall(PetscStrncpy(atype, MATAIJ, sizeof(atype)));
94: PetscCall(PetscStrncpy(btype, MATDENSE, sizeof(btype)));
95: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "MatMatMult() test options", "Mat");
96: PetscCall(PetscOptionsInt("-m", "Number of rows and columns of A", NULL, m, &m, NULL));
97: PetscCall(PetscOptionsInt("-n", "Number of columns of B", NULL, n, &n, NULL));
98: PetscCall(PetscOptionsFList("-A_mat_type", "Type of A", "MatSetType", MatList, atype, atype, sizeof(atype), NULL));
99: PetscCall(PetscOptionsFList("-B_mat_type", "Type of B", "MatSetType", MatList, btype, btype, sizeof(btype), NULL));
100: PetscCall(PetscOptionsBool("-block_diagonal", "Only keep the columns of A owned by this process, so the off-diagonal blocks have no columns", NULL, block_diagonal, &block_diagonal, NULL));
101: PetscCall(PetscOptionsBool("-tridiagonal", "Make A tridiagonal, so that only the first and last row of each off-diagonal block are nonempty and the block is stored in compressed-row format", NULL, tridiagonal, &tridiagonal, NULL));
102: PetscCall(PetscOptionsBool("-transpose", "Compute A^T * B with MatTransposeMatMult() instead of A * B", NULL, transpose, &transpose, NULL));
103: PetscCall(PetscOptionsBool("-change_lda", "Reuse A^T * B with an output whose leading dimension changes", NULL, change_lda, &change_lda, NULL));
104: PetscCall(PetscOptionsBool("-check_copies", "Check that a product of device matrices does not copy between the host and the device", NULL, check_copies, &check_copies, NULL));
105: PetscOptionsEnd();
106: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
107: /* PetscSF stages device buffers through the host when GPU-aware MPI is disabled. */
108: if (check_copies && size > 1 && !use_gpu_aware_mpi) {
109: PetscCall(PetscInfo(NULL, "Skipping copy checks because GPU-aware MPI is disabled; numerical checks remain enabled\n"));
110: check_copies = PETSC_FALSE;
111: }
112: if (check_copies) PetscCall(PetscLogDefaultBegin());
113: PetscCall(PetscLogEventRegister("ProductCheck", MAT_CLASSID, &event));
115: /* host reference matrices, deliberately not calling MatSetFromOptions() so they keep their host types */
116: PetscCall(MatCreate(PETSC_COMM_WORLD, &Ah));
117: PetscCall(MatSetSizes(Ah, PETSC_DECIDE, PETSC_DECIDE, m, m));
118: PetscCall(MatSetType(Ah, MATAIJ));
119: PetscCall(MatSeqAIJSetPreallocation(Ah, 4, NULL));
120: PetscCall(MatMPIAIJSetPreallocation(Ah, 4, NULL, 4, NULL));
121: PetscCall(MatGetOwnershipRange(Ah, &rstart, &rend));
122: for (PetscInt i = rstart; i < rend; i++) {
123: if (tridiagonal) {
124: ncols = 0;
125: if (i > 0) {
126: cols[ncols] = i - 1;
127: vals[ncols] = -1.0;
128: ncols++;
129: }
130: cols[ncols] = i;
131: vals[ncols] = 2.0;
132: ncols++;
133: if (i < m - 1) {
134: cols[ncols] = i + 1;
135: vals[ncols] = -1.0;
136: ncols++;
137: }
138: } else {
139: ncols = 4;
140: cols[0] = i;
141: cols[1] = (i + 1) % m;
142: cols[2] = (i + m / 2) % m;
143: cols[3] = (i * 7 + 3) % m;
144: for (PetscInt k = 0; k < ncols; k++) vals[k] = 1.0 + 0.1 * (PetscReal)(i + cols[k]);
145: }
146: /* duplicate columns accumulate with ADD_VALUES */
147: for (PetscInt k = 0; k < ncols; k++) {
148: if (block_diagonal && (cols[k] < rstart || cols[k] >= rend)) continue;
149: PetscCall(MatSetValue(Ah, i, cols[k], vals[k], ADD_VALUES));
150: }
151: }
152: PetscCall(MatAssemblyBegin(Ah, MAT_FINAL_ASSEMBLY));
153: PetscCall(MatAssemblyEnd(Ah, MAT_FINAL_ASSEMBLY));
155: PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, m, n, NULL, &Bh));
156: PetscCall(MatGetOwnershipRange(Bh, &rstart, &rend));
157: for (PetscInt i = rstart; i < rend; i++) {
158: for (PetscInt j = 0; j < n; j++) PetscCall(MatSetValue(Bh, i, j, 0.5 + PetscSinReal((PetscReal)(i + 3 * j)), INSERT_VALUES));
159: }
160: PetscCall(MatAssemblyBegin(Bh, MAT_FINAL_ASSEMBLY));
161: PetscCall(MatAssemblyEnd(Bh, MAT_FINAL_ASSEMBLY));
163: /* the matrices under test hold the same values with the requested types */
164: PetscCall(MatDuplicate(Ah, MAT_COPY_VALUES, &A));
165: PetscCall(MatConvert(A, atype, MAT_INPLACE_MATRIX, &A));
166: PetscCall(MatDuplicate(Bh, MAT_COPY_VALUES, &B));
167: PetscCall(MatConvert(B, btype, MAT_INPLACE_MATRIX, &B));
169: PetscCall(Product(A, B, MAT_INITIAL_MATRIX, transpose, &C));
170: PetscCall(Product(Ah, Bh, MAT_INITIAL_MATRIX, transpose, &Ch));
171: PetscCall(CheckProduct(C, Ch, "MAT_INITIAL_MATRIX"));
172: if (transpose) PetscCall(CheckTransposeReference(Ah, Bh, Ch, "MAT_INITIAL_MATRIX"));
174: /* change the values of both pairs identically and reuse the products */
175: PetscCall(MatScale(A, 2.0));
176: PetscCall(MatShift(A, 1.0));
177: PetscCall(MatScale(B, -0.5));
178: PetscCall(MatScale(Ah, 2.0));
179: PetscCall(MatShift(Ah, 1.0));
180: PetscCall(MatScale(Bh, -0.5));
182: PetscCall(Product(A, B, MAT_REUSE_MATRIX, transpose, &C));
183: PetscCall(Product(Ah, Bh, MAT_REUSE_MATRIX, transpose, &Ch));
184: PetscCall(CheckProduct(C, Ch, "MAT_REUSE_MATRIX"));
185: if (transpose) PetscCall(CheckTransposeReference(Ah, Bh, Ch, "MAT_REUSE_MATRIX"));
186: if (change_lda) PetscCall(CheckChangedLDA(Ah, Bh, Ch));
188: /* the reuse above brings the modified values to the device (MatShift() may run on the host), so any copy logged
189: in a further reuse is made by the product itself */
190: PetscCall(PetscLogEventBegin(event, 0, 0, 0, 0));
191: PetscCall(Product(A, B, MAT_REUSE_MATRIX, transpose, &C));
192: PetscCall(PetscLogEventEnd(event, 0, 0, 0, 0));
193: PetscCall(CheckProduct(C, Ch, "MAT_REUSE_MATRIX again"));
195: #if PetscDefined(HAVE_DEVICE)
196: /* a product of device matrices must run entirely on the device, in either direction and with any number of processes */
197: if (check_copies) {
198: PetscEventPerfInfo info;
200: PetscCall(PetscLogEventGetPerfInfo(PETSC_DETERMINE, event, &info));
201: PetscCheck(info.GpuToCpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected GPU to CPU copies (%g bytes) in the product", info.GpuToCpuCount, info.GpuToCpuSize);
202: PetscCheck(info.CpuToGpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected CPU to GPU copies (%g bytes) in the product", info.CpuToGpuCount, info.CpuToGpuSize);
203: }
204: #endif
206: PetscCall(MatDestroy(&A));
207: PetscCall(MatDestroy(&B));
208: PetscCall(MatDestroy(&C));
209: PetscCall(MatDestroy(&Ah));
210: PetscCall(MatDestroy(&Bh));
211: PetscCall(MatDestroy(&Ch));
212: PetscCall(PetscFinalize());
213: return 0;
214: }
216: /*TEST
218: testset:
219: output_file: output/empty.out
221: test:
222: suffix: host
223: nsize: {{1 3}}
224: args: -block_diagonal {{0 1}} -transpose {{0 1}}
226: test:
227: suffix: host_batch
228: nsize: 3
229: args: -matproduct_batch_size {{2 3 4}} -tridiagonal {{0 1}} -transpose {{0 1}}
231: test:
232: suffix: host_change_lda
233: nsize: {{1 3}}
234: args: -transpose -change_lda
236: test:
237: suffix: host_change_lda_batch
238: nsize: 3
239: args: -transpose -change_lda -matproduct_batch_size {{2 3}}
241: # the local block of B is a host MATSEQDENSE here, so the scatter of the off-process rows
242: # does not go through device memory and -use_gpu_aware_mpi would only duplicate the runs
243: test:
244: suffix: kokkos
245: requires: kokkos_kernels
246: nsize: {{1 3}}
247: args: -A_mat_type aijkokkos -transpose {{0 1}}
249: test:
250: suffix: kokkos_batch
251: requires: kokkos_kernels
252: nsize: 3
253: args: -A_mat_type aijkokkos -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -transpose {{0 1}}
255: # -check_copies in parallel requires GPU-aware MPI: without it PetscSF stages the device buffers of the
256: # scatter through the host, and those copies are logged inside the product. The plain device tests run
257: # with that staging, which is valid on every build, and the _copies_par tests cover the GPU-aware path
259: testset:
260: requires: cuda kokkos_kernels
261: args: -A_mat_type aijkokkos -B_mat_type densecuda
262: output_file: output/empty.out
264: test:
265: suffix: kokkos_cuda
266: nsize: 3
267: args: -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0 -transpose {{0 1}}
269: test:
270: suffix: kokkos_cuda_copies
271: nsize: 1
272: requires: defined(PETSC_USE_LOG)
273: args: -check_copies -transpose {{0 1}}
275: test:
276: suffix: kokkos_cuda_copies_par
277: nsize: 3
278: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
279: args: -check_copies -matproduct_batch_size {{0 3}} -transpose {{0 1}}
281: testset:
282: requires: cuda
283: args: -A_mat_type aijcusparse
284: output_file: output/empty.out
286: test:
287: suffix: cuda
288: nsize: 3
289: args: -B_mat_type {{dense densecuda}} -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0 -transpose {{0 1}}
291: test:
292: suffix: cuda_copies
293: nsize: 1
294: requires: defined(PETSC_USE_LOG)
295: args: -B_mat_type densecuda -check_copies -transpose {{0 1}}
297: test:
298: suffix: cuda_copies_par
299: nsize: 3
300: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
301: args: -B_mat_type densecuda -check_copies -matproduct_batch_size {{0 3}} -transpose {{0 1}}
303: testset:
304: requires: hip kokkos_kernels
305: args: -A_mat_type aijkokkos -B_mat_type densehip
306: output_file: output/empty.out
308: test:
309: suffix: kokkos_hip
310: nsize: 3
311: args: -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0 -transpose {{0 1}}
313: test:
314: suffix: kokkos_hip_copies
315: nsize: 1
316: requires: defined(PETSC_USE_LOG)
317: args: -check_copies -transpose {{0 1}}
319: test:
320: suffix: kokkos_hip_copies_par
321: nsize: 3
322: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
323: args: -check_copies -matproduct_batch_size {{0 3}} -transpose {{0 1}}
325: testset:
326: requires: hip
327: args: -A_mat_type aijhipsparse
328: output_file: output/empty.out
330: test:
331: suffix: hip
332: nsize: 3
333: args: -B_mat_type {{dense densehip}} -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0 -transpose {{0 1}}
335: test:
336: suffix: hip_copies
337: nsize: 1
338: requires: defined(PETSC_USE_LOG)
339: args: -B_mat_type densehip -check_copies -transpose {{0 1}}
341: test:
342: suffix: hip_copies_par
343: nsize: 3
344: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
345: args: -B_mat_type densehip -check_copies -matproduct_batch_size {{0 3}} -transpose {{0 1}}
347: TEST*/