Actual source code: ex319.c
1: static char help[] = "Tests MatMatMult() 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>
7: /* The product of the matrices whose types are set with -A_mat_type and -B_mat_type is compared against
8: the product of host MATAIJ and MATDENSE copies holding the same values */
9: static PetscErrorCode CheckProduct(Mat C, Mat Ch, const char *label)
10: {
11: Mat Ct;
12: PetscReal nrm, err;
14: PetscFunctionBeginUser;
15: /* a host matrix is duplicated, a device one is copied back to the host */
16: PetscCall(MatConvert(C, MATDENSE, MAT_INITIAL_MATRIX, &Ct));
17: PetscCall(MatNorm(Ch, NORM_INFINITY, &nrm));
18: PetscCall(MatAXPY(Ct, -1.0, Ch, SAME_NONZERO_PATTERN));
19: PetscCall(MatNorm(Ct, NORM_INFINITY, &err));
20: 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);
21: PetscCall(MatDestroy(&Ct));
22: PetscFunctionReturn(PETSC_SUCCESS);
23: }
25: int main(int argc, char **args)
26: {
27: Mat A, B, C, Ah, Bh, Ch;
28: char atype[256], btype[256];
29: PetscScalar vals[4];
30: PetscInt m = 20, n = 7, rstart, rend, ncols, cols[4];
31: PetscBool block_diagonal = PETSC_FALSE, check_copies = PETSC_FALSE, tridiagonal = PETSC_FALSE;
32: PetscLogEvent event;
34: PetscFunctionBeginUser;
35: PetscCall(PetscInitialize(&argc, &args, NULL, help));
36: PetscCall(PetscStrncpy(atype, MATAIJ, sizeof(atype)));
37: PetscCall(PetscStrncpy(btype, MATDENSE, sizeof(btype)));
38: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "MatMatMult() test options", "Mat");
39: PetscCall(PetscOptionsInt("-m", "Number of rows and columns of A", NULL, m, &m, NULL));
40: PetscCall(PetscOptionsInt("-n", "Number of columns of B", NULL, n, &n, NULL));
41: PetscCall(PetscOptionsFList("-A_mat_type", "Type of A", "MatSetType", MatList, atype, atype, sizeof(atype), NULL));
42: PetscCall(PetscOptionsFList("-B_mat_type", "Type of B", "MatSetType", MatList, btype, btype, sizeof(btype), NULL));
43: 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));
44: 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));
45: 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));
46: PetscOptionsEnd();
47: if (check_copies) PetscCall(PetscLogDefaultBegin());
48: PetscCall(PetscLogEventRegister("ProductCheck", MAT_CLASSID, &event));
50: /* host reference matrices, deliberately not calling MatSetFromOptions() so they keep their host types */
51: PetscCall(MatCreate(PETSC_COMM_WORLD, &Ah));
52: PetscCall(MatSetSizes(Ah, PETSC_DECIDE, PETSC_DECIDE, m, m));
53: PetscCall(MatSetType(Ah, MATAIJ));
54: PetscCall(MatSeqAIJSetPreallocation(Ah, 4, NULL));
55: PetscCall(MatMPIAIJSetPreallocation(Ah, 4, NULL, 4, NULL));
56: PetscCall(MatGetOwnershipRange(Ah, &rstart, &rend));
57: for (PetscInt i = rstart; i < rend; i++) {
58: if (tridiagonal) {
59: ncols = 0;
60: if (i > 0) {
61: cols[ncols] = i - 1;
62: vals[ncols] = -1.0;
63: ncols++;
64: }
65: cols[ncols] = i;
66: vals[ncols] = 2.0;
67: ncols++;
68: if (i < m - 1) {
69: cols[ncols] = i + 1;
70: vals[ncols] = -1.0;
71: ncols++;
72: }
73: } else {
74: ncols = 4;
75: cols[0] = i;
76: cols[1] = (i + 1) % m;
77: cols[2] = (i + m / 2) % m;
78: cols[3] = (i * 7 + 3) % m;
79: for (PetscInt k = 0; k < ncols; k++) vals[k] = 1.0 + 0.1 * (PetscReal)(i + cols[k]);
80: }
81: /* duplicate columns accumulate with ADD_VALUES */
82: for (PetscInt k = 0; k < ncols; k++) {
83: if (block_diagonal && (cols[k] < rstart || cols[k] >= rend)) continue;
84: PetscCall(MatSetValue(Ah, i, cols[k], vals[k], ADD_VALUES));
85: }
86: }
87: PetscCall(MatAssemblyBegin(Ah, MAT_FINAL_ASSEMBLY));
88: PetscCall(MatAssemblyEnd(Ah, MAT_FINAL_ASSEMBLY));
90: PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, m, n, NULL, &Bh));
91: PetscCall(MatGetOwnershipRange(Bh, &rstart, &rend));
92: for (PetscInt i = rstart; i < rend; i++) {
93: for (PetscInt j = 0; j < n; j++) PetscCall(MatSetValue(Bh, i, j, 0.5 + PetscSinReal((PetscReal)(i + 3 * j)), INSERT_VALUES));
94: }
95: PetscCall(MatAssemblyBegin(Bh, MAT_FINAL_ASSEMBLY));
96: PetscCall(MatAssemblyEnd(Bh, MAT_FINAL_ASSEMBLY));
98: /* the matrices under test hold the same values with the requested types */
99: PetscCall(MatDuplicate(Ah, MAT_COPY_VALUES, &A));
100: PetscCall(MatConvert(A, atype, MAT_INPLACE_MATRIX, &A));
101: PetscCall(MatDuplicate(Bh, MAT_COPY_VALUES, &B));
102: PetscCall(MatConvert(B, btype, MAT_INPLACE_MATRIX, &B));
104: PetscCall(MatMatMult(A, B, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &C));
105: PetscCall(MatMatMult(Ah, Bh, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &Ch));
106: PetscCall(CheckProduct(C, Ch, "MAT_INITIAL_MATRIX"));
108: /* change the values of both pairs identically and reuse the products */
109: PetscCall(MatScale(A, 2.0));
110: PetscCall(MatShift(A, 1.0));
111: PetscCall(MatScale(B, -0.5));
112: PetscCall(MatScale(Ah, 2.0));
113: PetscCall(MatShift(Ah, 1.0));
114: PetscCall(MatScale(Bh, -0.5));
116: PetscCall(MatMatMult(A, B, MAT_REUSE_MATRIX, PETSC_DETERMINE, &C));
117: PetscCall(MatMatMult(Ah, Bh, MAT_REUSE_MATRIX, PETSC_DETERMINE, &Ch));
118: PetscCall(CheckProduct(C, Ch, "MAT_REUSE_MATRIX"));
120: /* the reuse above brings the modified values to the device (MatShift() may run on the host), so any copy logged
121: in a further reuse is made by the product itself */
122: PetscCall(PetscLogEventBegin(event, 0, 0, 0, 0));
123: PetscCall(MatMatMult(A, B, MAT_REUSE_MATRIX, PETSC_DETERMINE, &C));
124: PetscCall(PetscLogEventEnd(event, 0, 0, 0, 0));
125: PetscCall(CheckProduct(C, Ch, "MAT_REUSE_MATRIX again"));
127: #if PetscDefined(HAVE_DEVICE)
128: /* a product of device matrices must run entirely on the device, in either direction and with any number of processes */
129: if (check_copies) {
130: PetscEventPerfInfo info;
132: PetscCall(PetscLogEventGetPerfInfo(PETSC_DETERMINE, event, &info));
133: PetscCheck(info.GpuToCpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected GPU to CPU copies (%g bytes) in MatMatMult()", info.GpuToCpuCount, info.GpuToCpuSize);
134: PetscCheck(info.CpuToGpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected CPU to GPU copies (%g bytes) in MatMatMult()", info.CpuToGpuCount, info.CpuToGpuSize);
135: }
136: #endif
138: PetscCall(MatDestroy(&A));
139: PetscCall(MatDestroy(&B));
140: PetscCall(MatDestroy(&C));
141: PetscCall(MatDestroy(&Ah));
142: PetscCall(MatDestroy(&Bh));
143: PetscCall(MatDestroy(&Ch));
144: PetscCall(PetscFinalize());
145: return 0;
146: }
148: /*TEST
150: testset:
151: output_file: output/empty.out
153: test:
154: suffix: host
155: nsize: {{1 3}}
156: args: -block_diagonal {{0 1}}
158: test:
159: suffix: host_batch
160: nsize: 3
161: args: -matproduct_batch_size {{2 3 4}} -tridiagonal {{0 1}}
163: # the local block of B is a host MATSEQDENSE here, so the scatter of the off-process rows
164: # does not go through device memory and -use_gpu_aware_mpi would only duplicate the runs
165: test:
166: suffix: kokkos
167: requires: kokkos_kernels
168: nsize: {{1 3}}
169: args: -A_mat_type aijkokkos
171: test:
172: suffix: kokkos_batch
173: requires: kokkos_kernels
174: nsize: 3
175: args: -A_mat_type aijkokkos -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}}
177: # -check_copies in parallel requires GPU-aware MPI: without it PetscSF stages the device buffers of the
178: # scatter through the host, and those copies are logged inside MatMatMult(). The plain device tests run
179: # with that staging, which is valid on every build, and the _copies_par tests cover the GPU-aware path
181: testset:
182: requires: cuda kokkos_kernels
183: args: -A_mat_type aijkokkos -B_mat_type densecuda
184: output_file: output/empty.out
186: test:
187: suffix: kokkos_cuda
188: nsize: 3
189: args: -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0
191: test:
192: suffix: kokkos_cuda_copies
193: nsize: 1
194: requires: defined(PETSC_USE_LOG)
195: args: -check_copies
197: test:
198: suffix: kokkos_cuda_copies_par
199: nsize: 3
200: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
201: args: -check_copies -matproduct_batch_size {{0 3}}
203: testset:
204: requires: cuda
205: args: -A_mat_type aijcusparse
206: output_file: output/empty.out
208: test:
209: suffix: cuda
210: nsize: 3
211: args: -B_mat_type {{dense densecuda}} -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0
213: test:
214: suffix: cuda_copies
215: nsize: 1
216: requires: defined(PETSC_USE_LOG)
217: args: -B_mat_type densecuda -check_copies
219: test:
220: suffix: cuda_copies_par
221: nsize: 3
222: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
223: args: -B_mat_type densecuda -check_copies -matproduct_batch_size {{0 3}}
225: testset:
226: requires: hip kokkos_kernels
227: args: -A_mat_type aijkokkos -B_mat_type densehip
228: output_file: output/empty.out
230: test:
231: suffix: kokkos_hip
232: nsize: 3
233: args: -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0
235: test:
236: suffix: kokkos_hip_copies
237: nsize: 1
238: requires: defined(PETSC_USE_LOG)
239: args: -check_copies
241: test:
242: suffix: kokkos_hip_copies_par
243: nsize: 3
244: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
245: args: -check_copies -matproduct_batch_size {{0 3}}
247: testset:
248: requires: hip
249: args: -A_mat_type aijhipsparse
250: output_file: output/empty.out
252: test:
253: suffix: hip
254: nsize: 3
255: args: -B_mat_type {{dense densehip}} -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0
257: test:
258: suffix: hip_copies
259: nsize: 1
260: requires: defined(PETSC_USE_LOG)
261: args: -B_mat_type densehip -check_copies
263: test:
264: suffix: hip_copies_par
265: nsize: 3
266: requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
267: args: -B_mat_type densehip -check_copies -matproduct_batch_size {{0 3}}
269: TEST*/