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