1 static char help[] = "Tests MatCreateDenseCUDA(), MatDenseCUDAPlaceArray(), MatDenseCUDAReplaceArray(), MatDenseCUDAResetArray()\n"; 2 3 #include <petscmat.h> 4 5 static PetscErrorCode MatMult_S(Mat S, Vec x, Vec y) { 6 Mat A; 7 8 PetscFunctionBeginUser; 9 PetscCall(MatShellGetContext(S, &A)); 10 PetscCall(MatMult(A, x, y)); 11 PetscFunctionReturn(0); 12 } 13 14 static PetscBool test_cusparse_transgen = PETSC_FALSE; 15 16 static PetscErrorCode MatMultTranspose_S(Mat S, Vec x, Vec y) { 17 Mat A; 18 19 PetscFunctionBeginUser; 20 PetscCall(MatShellGetContext(S, &A)); 21 PetscCall(MatMultTranspose(A, x, y)); 22 23 /* alternate transgen true and false to test code logic */ 24 PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, test_cusparse_transgen)); 25 test_cusparse_transgen = (PetscBool)!test_cusparse_transgen; 26 PetscFunctionReturn(0); 27 } 28 29 int main(int argc, char **argv) { 30 Mat A, B, C, S; 31 Vec t, v; 32 PetscScalar *vv, *aa; 33 PetscInt n = 30, k = 6, l = 0, i, Istart, Iend, nloc, bs, test = 1; 34 PetscBool flg, reset, use_shell = PETSC_FALSE; 35 VecType vtype; 36 37 PetscFunctionBeginUser; 38 PetscCall(PetscInitialize(&argc, &argv, (char *)0, help)); 39 PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL)); 40 PetscCall(PetscOptionsGetInt(NULL, NULL, "-k", &k, NULL)); 41 PetscCall(PetscOptionsGetInt(NULL, NULL, "-l", &l, NULL)); 42 PetscCall(PetscOptionsGetInt(NULL, NULL, "-test", &test, NULL)); 43 PetscCall(PetscOptionsGetBool(NULL, NULL, "-use_shell", &use_shell, NULL)); 44 PetscCheck(k >= 0, PETSC_COMM_WORLD, PETSC_ERR_USER, "k %" PetscInt_FMT " must be positive", k); 45 PetscCheck(l >= 0, PETSC_COMM_WORLD, PETSC_ERR_USER, "l %" PetscInt_FMT " must be positive", l); 46 PetscCheck(l <= k, PETSC_COMM_WORLD, PETSC_ERR_USER, "l %" PetscInt_FMT " must be smaller or equal than k %" PetscInt_FMT, l, k); 47 48 /* sparse matrix */ 49 PetscCall(MatCreate(PETSC_COMM_WORLD, &A)); 50 PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, n, n)); 51 PetscCall(MatSetType(A, MATAIJCUSPARSE)); 52 PetscCall(MatSetOptionsPrefix(A, "A_")); 53 PetscCall(MatSetFromOptions(A)); 54 PetscCall(MatSetUp(A)); 55 56 /* test special case for SeqAIJCUSPARSE to generate explicit transpose (not default) */ 57 PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, PETSC_TRUE)); 58 59 PetscCall(MatGetOwnershipRange(A, &Istart, &Iend)); 60 for (i = Istart; i < Iend; i++) { 61 if (i > 0) PetscCall(MatSetValue(A, i, i - 1, -1.0, INSERT_VALUES)); 62 if (i < n - 1) PetscCall(MatSetValue(A, i, i + 1, -1.0, INSERT_VALUES)); 63 PetscCall(MatSetValue(A, i, i, 2.0, INSERT_VALUES)); 64 } 65 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 66 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 67 68 /* template vector */ 69 PetscCall(MatCreateVecs(A, NULL, &t)); 70 PetscCall(VecGetType(t, &vtype)); 71 72 /* long vector, contains the stacked columns of an nxk dense matrix */ 73 PetscCall(VecGetLocalSize(t, &nloc)); 74 PetscCall(VecGetBlockSize(t, &bs)); 75 PetscCall(VecCreate(PetscObjectComm((PetscObject)t), &v)); 76 PetscCall(VecSetType(v, vtype)); 77 PetscCall(VecSetSizes(v, k * nloc, k * n)); 78 PetscCall(VecSetBlockSize(v, bs)); 79 PetscCall(VecSetRandom(v, NULL)); 80 81 /* dense matrix that contains the columns of v */ 82 PetscCall(VecCUDAGetArray(v, &vv)); 83 84 /* test few cases for MatDenseCUDA handling pointers */ 85 switch (test) { 86 case 1: 87 PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, vv, &B)); /* pass a pointer to avoid allocation of storage */ 88 PetscCall(MatDenseCUDAReplaceArray(B, NULL)); /* replace with a null pointer, the value after BVRestoreMat */ 89 PetscCall(MatDenseCUDAPlaceArray(B, vv + l * nloc)); /* set the actual pointer */ 90 reset = PETSC_TRUE; 91 break; 92 case 2: 93 PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &B)); 94 PetscCall(MatDenseCUDAPlaceArray(B, vv + l * nloc)); /* set the actual pointer */ 95 reset = PETSC_TRUE; 96 break; 97 default: 98 PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, vv + l * nloc, &B)); 99 reset = PETSC_FALSE; 100 break; 101 } 102 PetscCall(VecCUDARestoreArray(v, &vv)); 103 104 /* Test MatMatMult */ 105 if (use_shell) { 106 /* we could have called the general convertor below, but we explicit set the operations 107 ourselves to test MatProductSymbolic_X_Dense, MatProductNumeric_X_Dense code */ 108 /* PetscCall(MatConvert(A,MATSHELL,MAT_INITIAL_MATRIX,&S)); */ 109 PetscCall(MatCreateShell(PetscObjectComm((PetscObject)v), nloc, nloc, n, n, A, &S)); 110 PetscCall(MatShellSetOperation(S, MATOP_MULT, (void (*)(void))MatMult_S)); 111 PetscCall(MatShellSetOperation(S, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_S)); 112 PetscCall(MatShellSetVecType(S, vtype)); 113 } else { 114 PetscCall(PetscObjectReference((PetscObject)A)); 115 S = A; 116 } 117 118 PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &C)); 119 120 /* test MatMatMult */ 121 PetscCall(MatProductCreateWithMat(S, B, NULL, C)); 122 PetscCall(MatProductSetType(C, MATPRODUCT_AB)); 123 PetscCall(MatProductSetFromOptions(C)); 124 PetscCall(MatProductSymbolic(C)); 125 PetscCall(MatProductNumeric(C)); 126 PetscCall(MatMatMultEqual(S, B, C, 10, &flg)); 127 if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatMatMult\n")); 128 129 /* test MatTransposeMatMult */ 130 PetscCall(MatProductCreateWithMat(S, B, NULL, C)); 131 PetscCall(MatProductSetType(C, MATPRODUCT_AtB)); 132 PetscCall(MatProductSetFromOptions(C)); 133 PetscCall(MatProductSymbolic(C)); 134 PetscCall(MatProductNumeric(C)); 135 PetscCall(MatTransposeMatMultEqual(S, B, C, 10, &flg)); 136 if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatTransposeMatMult\n")); 137 138 PetscCall(MatDestroy(&C)); 139 PetscCall(MatDestroy(&S)); 140 141 /* finished using B */ 142 PetscCall(MatDenseCUDAGetArray(B, &aa)); 143 PetscCheck(vv == aa - l * nloc, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong array"); 144 PetscCall(MatDenseCUDARestoreArray(B, &aa)); 145 if (reset) PetscCall(MatDenseCUDAResetArray(B)); 146 PetscCall(VecCUDARestoreArray(v, &vv)); 147 148 if (test == 1) { 149 PetscCall(MatDenseCUDAGetArray(B, &aa)); 150 PetscCheck(!aa, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Expected a null pointer"); 151 PetscCall(MatDenseCUDARestoreArray(B, &aa)); 152 } 153 154 /* free work space */ 155 PetscCall(MatDestroy(&B)); 156 PetscCall(MatDestroy(&A)); 157 PetscCall(VecDestroy(&t)); 158 PetscCall(VecDestroy(&v)); 159 PetscCall(PetscFinalize()); 160 return 0; 161 } 162 163 /*TEST 164 165 build: 166 requires: cuda 167 168 test: 169 requires: cuda 170 suffix: 1 171 nsize: {{1 2}} 172 args: -A_mat_type {{aij aijcusparse}} -test {{0 1 2}} -k 6 -l {{0 5}} -use_shell {{0 1}} 173 174 TEST*/ 175