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