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