1 static char help[] = "Tests MatCreateDenseCUDA(), MatDenseCUDAPlaceArray(), MatDenseCUDAReplaceArray(), MatDenseCUDAResetArray()\n"; 2 3 #include <petscmat.h> 4 #include <petscpkg_version.h> 5 6 static PetscErrorCode MatMult_S(Mat S, Vec x, Vec y) 7 { 8 Mat A; 9 10 PetscFunctionBeginUser; 11 PetscCall(MatShellGetContext(S, &A)); 12 PetscCall(MatMult(A, x, y)); 13 PetscFunctionReturn(PETSC_SUCCESS); 14 } 15 16 static PetscBool test_cusparse_transgen = PETSC_FALSE; 17 18 static PetscErrorCode MatMultTranspose_S(Mat S, Vec x, Vec y) 19 { 20 Mat A; 21 22 PetscFunctionBeginUser; 23 PetscCall(MatShellGetContext(S, &A)); 24 PetscCall(MatMultTranspose(A, x, y)); 25 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 } 31 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; 48 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); 59 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)); 67 68 /* test special case for SeqAIJCUSPARSE to generate explicit transpose (not default) */ 69 PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, PETSC_TRUE)); 70 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)); 79 80 /* template vector */ 81 PetscCall(MatCreateVecs(A, NULL, &t)); 82 PetscCall(VecGetType(t, &vtype)); 83 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)); 92 93 /* dense matrix that contains the columns of v */ 94 PetscCall(VecCUDAGetArray(v, &vv)); 95 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 } 114 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 } 128 129 PetscCall(MatCreateDenseCUDA(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &C)); 130 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")); 139 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")); 148 149 PetscCall(MatDestroy(&C)); 150 PetscCall(MatDestroy(&S)); 151 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)); 158 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 } 164 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 } 173 174 /*TEST 175 176 build: 177 requires: cuda 178 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}} 184 185 TEST*/ 186