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