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*/