xref: /petsc/src/mat/tests/ex69.c (revision 28b400f66ebc7ae0049166a2294dfcd3df27e64b)
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