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