Actual source code: ex69.c

  1: static char help[] = "Tests MatCreateDenseCUDA(), MatDenseCUDAPlaceArray(), MatDenseCUDAReplaceArray(), MatDenseCUDAResetArray()\n";

  3: #include <petscmat.h>
  4: #include <petscpkg_version.h>

  6: static PetscErrorCode MatMult_S(Mat S, Vec x, Vec y)
  7: {
  8:   Mat A;

 10:   PetscFunctionBeginUser;
 11:   PetscCall(MatShellGetContext(S, &A));
 12:   PetscCall(MatMult(A, x, y));
 13:   PetscFunctionReturn(PETSC_SUCCESS);
 14: }

 16: static PetscBool test_cusparse_transgen = PETSC_FALSE;

 18: static PetscErrorCode MatMultTranspose_S(Mat S, Vec x, Vec y)
 19: {
 20:   Mat A;

 22:   PetscFunctionBeginUser;
 23:   PetscCall(MatShellGetContext(S, &A));
 24:   PetscCall(MatMultTranspose(A, x, y));

 26:   /* alternate transgen true and false to test code logic */
 27:   PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, test_cusparse_transgen));
 28:   test_cusparse_transgen = (PetscBool)!test_cusparse_transgen;
 29:   PetscFunctionReturn(PETSC_SUCCESS);
 30: }

 32: int main(int argc, char **argv)
 33: {
 34:   Mat          A, B, C, S;
 35:   Vec          t, v;
 36:   PetscScalar *vv, *aa;
 37:   // We met a mysterious cudaErrorMisalignedAddress error on some systems with cuda-12.0,1 but not
 38:   // with prior cuda-11.2,3,7,8 versions. Making nloc an even number somehow 'fixes' the problem.
 39:   // See more at https://gitlab.com/petsc/petsc/-/merge_requests/6225
 40: #if PETSC_PKG_CUDA_VERSION_GE(12, 0, 0)
 41:   PetscInt n = 32;
 42: #else
 43:   PetscInt n = 30;
 44: #endif
 45:   PetscInt  k = 6, l = 0, i, Istart, Iend, nloc, bs, test = 1;
 46:   PetscBool flg, reset, use_shell = PETSC_FALSE;
 47:   VecType   vtype;

 49:   PetscFunctionBeginUser;
 50:   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
 51:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
 52:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-k", &k, NULL));
 53:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-l", &l, NULL));
 54:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-test", &test, NULL));
 55:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-use_shell", &use_shell, NULL));
 56:   PetscCheck(k >= 0, PETSC_COMM_WORLD, PETSC_ERR_USER, "k %" PetscInt_FMT " must be positive", k);
 57:   PetscCheck(l >= 0, PETSC_COMM_WORLD, PETSC_ERR_USER, "l %" PetscInt_FMT " must be positive", l);
 58:   PetscCheck(l <= k, PETSC_COMM_WORLD, PETSC_ERR_USER, "l %" PetscInt_FMT " must be smaller or equal than k %" PetscInt_FMT, l, k);

 60:   /* sparse matrix */
 61:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 62:   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, n, n));
 63:   PetscCall(MatSetType(A, MATAIJCUSPARSE));
 64:   PetscCall(MatSetOptionsPrefix(A, "A_"));
 65:   PetscCall(MatSetFromOptions(A));
 66:   PetscCall(MatSetUp(A));

 68:   /* test special case for SeqAIJCUSPARSE to generate explicit transpose (not default) */
 69:   PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, PETSC_TRUE));

 71:   PetscCall(MatGetOwnershipRange(A, &Istart, &Iend));
 72:   for (i = Istart; i < Iend; i++) {
 73:     if (i > 0) PetscCall(MatSetValue(A, i, i - 1, -1.0, INSERT_VALUES));
 74:     if (i < n - 1) PetscCall(MatSetValue(A, i, i + 1, -1.0, INSERT_VALUES));
 75:     PetscCall(MatSetValue(A, i, i, 2.0, INSERT_VALUES));
 76:   }
 77:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
 78:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));

 80:   /* template vector */
 81:   PetscCall(MatCreateVecs(A, NULL, &t));
 82:   PetscCall(VecGetType(t, &vtype));

 84:   /* long vector, contains the stacked columns of an nxk dense matrix */
 85:   PetscCall(VecGetLocalSize(t, &nloc));
 86:   PetscCall(VecGetBlockSize(t, &bs));
 87:   PetscCall(VecCreate(PetscObjectComm((PetscObject)t), &v));
 88:   PetscCall(VecSetType(v, vtype));
 89:   PetscCall(VecSetSizes(v, k * nloc, k * n));
 90:   PetscCall(VecSetBlockSize(v, bs));
 91:   PetscCall(VecSetRandom(v, NULL));

 93:   /* dense matrix that contains the columns of v */
 94:   PetscCall(VecCUDAGetArray(v, &vv));

 96:   /* test few cases for MatDenseCUDA handling pointers */
 97:   switch (test) {
 98:   case 1:
 99:     PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, vv, &B)); /* pass a pointer to avoid allocation of storage */
100:     PetscCall(MatDenseCUDAReplaceArray(B, NULL));                                                         /* replace with a null pointer, the value after BVRestoreMat */
101:     PetscCall(MatDenseCUDAPlaceArray(B, vv + l * nloc));                                                  /* set the actual pointer */
102:     reset = PETSC_TRUE;
103:     break;
104:   case 2:
105:     PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &B));
106:     PetscCall(MatDenseCUDAPlaceArray(B, vv + l * nloc)); /* set the actual pointer */
107:     reset = PETSC_TRUE;
108:     break;
109:   default:
110:     PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, vv + l * nloc, &B));
111:     reset = PETSC_FALSE;
112:     break;
113:   }

115:   /* Test MatMatMult */
116:   if (use_shell) {
117:     /* we could have called the general converter below, but we explicitly set the operations
118:        ourselves to test MatProductSymbolic_X_Dense, MatProductNumeric_X_Dense code */
119:     /* PetscCall(MatConvert(A,MATSHELL,MAT_INITIAL_MATRIX,&S)); */
120:     PetscCall(MatCreateShell(PetscObjectComm((PetscObject)v), nloc, nloc, n, n, A, &S));
121:     PetscCall(MatShellSetOperation(S, MATOP_MULT, (void (*)(void))MatMult_S));
122:     PetscCall(MatShellSetOperation(S, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_S));
123:     PetscCall(MatShellSetVecType(S, vtype));
124:   } else {
125:     PetscCall(PetscObjectReference((PetscObject)A));
126:     S = A;
127:   }

129:   PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &C));

131:   /* test MatMatMult */
132:   PetscCall(MatProductCreateWithMat(S, B, NULL, C));
133:   PetscCall(MatProductSetType(C, MATPRODUCT_AB));
134:   PetscCall(MatProductSetFromOptions(C));
135:   PetscCall(MatProductSymbolic(C));
136:   PetscCall(MatProductNumeric(C));
137:   PetscCall(MatMatMultEqual(S, B, C, 10, &flg));
138:   if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatMatMult\n"));

140:   /* test MatTransposeMatMult */
141:   PetscCall(MatProductCreateWithMat(S, B, NULL, C));
142:   PetscCall(MatProductSetType(C, MATPRODUCT_AtB));
143:   PetscCall(MatProductSetFromOptions(C));
144:   PetscCall(MatProductSymbolic(C));
145:   PetscCall(MatProductNumeric(C));
146:   PetscCall(MatTransposeMatMultEqual(S, B, C, 10, &flg));
147:   if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatTransposeMatMult\n"));

149:   PetscCall(MatDestroy(&C));
150:   PetscCall(MatDestroy(&S));

152:   /* finished using B */
153:   PetscCall(MatDenseGetArrayAndMemType(B, &aa, NULL));
154:   PetscCheck(vv == aa - l * nloc, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong array");
155:   PetscCall(MatDenseRestoreArrayAndMemType(B, &aa));
156:   if (reset) PetscCall(MatDenseCUDAResetArray(B));
157:   PetscCall(VecCUDARestoreArray(v, &vv));

159:   if (test == 1) {
160:     PetscCall(MatDenseGetArrayAndMemType(B, &aa, NULL));
161:     PetscCheck(!aa, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Expected a null pointer");
162:     PetscCall(MatDenseRestoreArrayAndMemType(B, &aa));
163:   }

165:   /* free work space */
166:   PetscCall(MatDestroy(&B));
167:   PetscCall(MatDestroy(&A));
168:   PetscCall(VecDestroy(&t));
169:   PetscCall(VecDestroy(&v));
170:   PetscCall(PetscFinalize());
171:   return 0;
172: }

174: /*TEST

176:   build:
177:     requires: cuda

179:   test:
180:     requires: cuda
181:     suffix: 1
182:     nsize: {{1 2}}
183:     args: -A_mat_type {{aij aijcusparse}} -test {{0 1 2}} -k 6 -l {{0 5}} -use_shell {{0 1}}

185: TEST*/