Actual source code: baijlibxsmm.c
1: #include <../src/mat/impls/baij/seq/baij.h>
2: #include <libxsmm/libxsmm.h>
4: typedef struct {
5: libxsmm_gemmfunction kernel;
6: PetscInt n, ldb, ldc;
7: } SeqBAIJLIBXSMM_SeqDense;
9: PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqBAIJLIBXSMM(Mat, MatType, MatReuse, Mat *);
11: static PetscErrorCode MatDuplicate_SeqBAIJLIBXSMM(Mat A, MatDuplicateOption op, Mat *B)
12: {
13: PetscFunctionBegin;
14: PetscCall(MatDuplicate_SeqBAIJ(A, op, B));
15: PetscCall(MatConvert_SeqBAIJ_SeqBAIJLIBXSMM(*B, MATSEQBAIJLIBXSMM, MAT_INPLACE_MATRIX, B));
16: PetscFunctionReturn(PETSC_SUCCESS);
17: }
19: static PetscErrorCode MatProductDataDestroy_SeqBAIJLIBXSMM(PetscCtxRt ctx)
20: {
21: SeqBAIJLIBXSMM_SeqDense *data = *(SeqBAIJLIBXSMM_SeqDense **)ctx;
23: PetscFunctionBegin;
24: PetscCall(PetscFree(data));
25: PetscFunctionReturn(PETSC_SUCCESS);
26: }
28: static PetscErrorCode MatProductNumeric_SeqBAIJLIBXSMM_SeqDense(Mat C)
29: {
30: Mat A = C->product->A, B = C->product->B;
31: Mat_SeqBAIJ *baij = (Mat_SeqBAIJ *)A->data;
32: SeqBAIJLIBXSMM_SeqDense *data = (SeqBAIJLIBXSMM_SeqDense *)C->product->data;
33: const PetscScalar *b;
34: PetscScalar *c;
35: const MatScalar *values = baij->a;
36: const PetscInt *cols = baij->j, *rows = baij->i;
37: PetscInt bs = A->rmap->bs, ldb, ldc;
38: libxsmm_gemm_param param = {0};
40: PetscFunctionBegin;
41: MatCheckProduct(C, 1);
42: PetscCheck(data && (data->kernel || !data->n), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing LIBXSMM product data");
43: PetscCall(MatDenseGetLDA(B, &ldb));
44: PetscCall(MatDenseGetLDA(C, &ldc));
45: PetscCheck(ldb == data->ldb && ldc == data->ldc, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Cannot reuse dense matrices with leading dimensions (%" PetscInt_FMT ",%" PetscInt_FMT ") instead of (%" PetscInt_FMT ",%" PetscInt_FMT ")", ldb, ldc, data->ldb,
46: data->ldc);
47: PetscCall(MatZeroEntries(C));
48: if (!data->n) PetscFunctionReturn(PETSC_SUCCESS);
49: PetscCall(MatDenseGetArrayRead(B, &b));
50: PetscCall(MatDenseGetArray(C, &c));
51: for (PetscInt i = 0; i < baij->mbs; ++i) {
52: for (PetscInt j = rows[i]; j < rows[i + 1]; ++j) {
53: param.a.primary = (void *)(values + j * bs * bs);
54: param.b.primary = (void *)(b + cols[j] * bs);
55: param.c.primary = c + i * bs;
56: PetscCallExternalVoid("LIBXSMM JIT kernel", data->kernel(¶m));
57: }
58: }
59: PetscCall(MatDenseRestoreArrayRead(B, &b));
60: PetscCall(MatDenseRestoreArray(C, &c));
61: PetscCall(PetscLogFlops(2.0 * baij->nz * bs * bs * data->n));
62: PetscFunctionReturn(PETSC_SUCCESS);
63: }
65: static PetscErrorCode MatProductSymbolic_SeqBAIJLIBXSMM_SeqDense(Mat C)
66: {
67: Mat A = C->product->A, B = C->product->B;
68: SeqBAIJLIBXSMM_SeqDense *data;
69: PetscInt m = A->rmap->n, n = B->cmap->n, bs = A->rmap->bs, ldb, ldc;
70: PetscBLASInt bn, bbs, bldb, bldc;
71: libxsmm_gemm_shape shape;
72: libxsmm_datatype datatype = PetscDefined(USE_REAL_SINGLE) ? LIBXSMM_DATATYPE_F32 : LIBXSMM_DATATYPE_F64;
74: PetscFunctionBegin;
75: MatCheckProduct(C, 1);
76: PetscCheck(A->cmap->n == B->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, %" PetscInt_FMT " != %" PetscInt_FMT, A->cmap->n, B->rmap->n);
77: PetscCall(MatSetSizes(C, m, n, m, n));
78: PetscCall(MatSetBlockSizesFromMats(C, A, B));
79: PetscCall(MatSetType(C, MATSEQDENSE));
80: PetscCall(MatSetVecType(C, B->defaultvectype));
81: PetscCall(MatSetUp(C));
82: PetscCall(MatDenseGetLDA(B, &ldb));
83: PetscCall(MatDenseGetLDA(C, &ldc));
84: PetscCheck((PetscInt64)PetscMax(ldb, ldc) * ((PetscInt64)n + 1) * sizeof(PetscScalar) <= PETSC_INT32_MAX, PETSC_COMM_SELF, PETSC_ERR_SUP, "LIBXSMM JIT cannot address dense matrices with leading dimensions (%" PetscInt_FMT ",%" PetscInt_FMT ") and %" PetscInt_FMT " columns", ldb, ldc, n);
85: PetscCall(PetscBLASIntCast(n, &bn));
86: PetscCall(PetscBLASIntCast(bs, &bbs));
87: PetscCall(PetscBLASIntCast(ldb, &bldb));
88: PetscCall(PetscBLASIntCast(ldc, &bldc));
89: PetscCall(PetscNew(&data));
90: data->n = n;
91: data->ldb = ldb;
92: data->ldc = ldc;
93: if (n) {
94: PetscCallExternalVoid("libxsmm_create_gemm_shape", shape = libxsmm_create_gemm_shape(bbs, bn, bbs, bbs, bldb, bldc, datatype, datatype, datatype, datatype));
95: PetscCallExternalVoid("libxsmm_dispatch_gemm", data->kernel = libxsmm_dispatch_gemm(shape, LIBXSMM_GEMM_FLAG_NONE, LIBXSMM_GEMM_PREFETCH_NONE));
96: PetscCheck(data->kernel, PETSC_COMM_SELF, PETSC_ERR_SUP, "LIBXSMM cannot generate a kernel for block size %" PetscInt_FMT " and %" PetscInt_FMT " dense columns", bs, n);
97: }
98: C->product->data = data;
99: C->product->destroy = MatProductDataDestroy_SeqBAIJLIBXSMM;
100: C->ops->productnumeric = MatProductNumeric_SeqBAIJLIBXSMM_SeqDense;
101: PetscFunctionReturn(PETSC_SUCCESS);
102: }
104: static PetscErrorCode MatProductSetFromOptions_SeqBAIJLIBXSMM_SeqDense(Mat C)
105: {
106: PetscFunctionBegin;
107: MatCheckProduct(C, 1);
108: if (C->product->type == MATPRODUCT_AB) C->ops->productsymbolic = MatProductSymbolic_SeqBAIJLIBXSMM_SeqDense;
109: PetscFunctionReturn(PETSC_SUCCESS);
110: }
112: PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJLIBXSMM_SeqBAIJ(Mat A, MatType type, MatReuse reuse, Mat *newmat)
113: {
114: Mat B = *newmat;
116: PetscFunctionBegin;
117: if (reuse == MAT_INITIAL_MATRIX) PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &B));
118: B->ops->duplicate = MatDuplicate_SeqBAIJ;
119: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_seqbaijlibxsmm_seqdense_C", NULL));
120: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqbaijlibxsmm_seqbaij_C", NULL));
121: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQBAIJ));
122: *newmat = B;
123: PetscFunctionReturn(PETSC_SUCCESS);
124: }
126: PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqBAIJLIBXSMM(Mat A, MatType type, MatReuse reuse, Mat *newmat)
127: {
128: Mat B = *newmat;
129: PetscBool sametype;
131: PetscFunctionBegin;
132: if (reuse == MAT_INITIAL_MATRIX) PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &B));
133: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQBAIJLIBXSMM, &sametype));
134: if (!sametype) {
135: B->ops->duplicate = MatDuplicate_SeqBAIJLIBXSMM;
136: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQBAIJLIBXSMM));
137: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_seqbaijlibxsmm_seqdense_C", MatProductSetFromOptions_SeqBAIJLIBXSMM_SeqDense));
138: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqbaijlibxsmm_seqbaij_C", MatConvert_SeqBAIJLIBXSMM_SeqBAIJ));
139: }
140: *newmat = B;
141: PetscFunctionReturn(PETSC_SUCCESS);
142: }
144: /*MC
145: MATSEQBAIJLIBXSMM - "seqbaijlibxsmm" - A sequential block sparse matrix that uses LIBXSMM kernels for products with `MATSEQDENSE` matrices
147: Options Database Key:
148: . -mat_type seqbaijlibxsmm - sets the matrix type to `MATSEQBAIJLIBXSMM` during a call to `MatSetFromOptions()`
150: Level: beginner
152: Notes:
153: This matrix type is available when PETSc is configured with `--download-libxsmm` or `--with-libxsmm-dir=directory`.
155: It has the same storage format and supports the same operations as `MATSEQBAIJ`.
157: .seealso: [](ch_matrices), `Mat`, `MATBAIJLIBXSMM`, `MATMPIBAIJLIBXSMM`, `MATSEQBAIJ`, `MatMatMult()`
158: M*/
159: PETSC_EXTERN PetscErrorCode MatCreate_SeqBAIJLIBXSMM(Mat A)
160: {
161: PetscFunctionBegin;
162: PetscCall(MatSetType(A, MATSEQBAIJ));
163: PetscCall(MatConvert_SeqBAIJ_SeqBAIJLIBXSMM(A, MATSEQBAIJLIBXSMM, MAT_INPLACE_MATRIX, &A));
164: PetscFunctionReturn(PETSC_SUCCESS);
165: }