Actual source code: baijlibxsmm.c
1: #include <../src/mat/impls/baij/seq/baij.h>
2: #include <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(MatSetUp(C));
81: PetscCall(MatDenseGetLDA(B, &ldb));
82: PetscCall(MatDenseGetLDA(C, &ldc));
83: 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);
84: PetscCall(PetscBLASIntCast(n, &bn));
85: PetscCall(PetscBLASIntCast(bs, &bbs));
86: PetscCall(PetscBLASIntCast(ldb, &bldb));
87: PetscCall(PetscBLASIntCast(ldc, &bldc));
88: PetscCall(PetscNew(&data));
89: data->n = n;
90: data->ldb = ldb;
91: data->ldc = ldc;
92: if (n) {
93: PetscCallExternalVoid("libxsmm_create_gemm_shape", shape = libxsmm_create_gemm_shape(bbs, bn, bbs, bbs, bldb, bldc, datatype, datatype, datatype, datatype));
94: PetscCallExternalVoid("libxsmm_dispatch_gemm", data->kernel = libxsmm_dispatch_gemm(shape, LIBXSMM_GEMM_FLAG_NONE, LIBXSMM_GEMM_PREFETCH_NONE));
95: 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);
96: }
97: C->product->data = data;
98: C->product->destroy = MatProductDataDestroy_SeqBAIJLIBXSMM;
99: C->ops->productnumeric = MatProductNumeric_SeqBAIJLIBXSMM_SeqDense;
100: PetscFunctionReturn(PETSC_SUCCESS);
101: }
103: static PetscErrorCode MatProductSetFromOptions_SeqBAIJLIBXSMM_SeqDense(Mat C)
104: {
105: PetscFunctionBegin;
106: MatCheckProduct(C, 1);
107: if (C->product->type == MATPRODUCT_AB) C->ops->productsymbolic = MatProductSymbolic_SeqBAIJLIBXSMM_SeqDense;
108: PetscFunctionReturn(PETSC_SUCCESS);
109: }
111: PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJLIBXSMM_SeqBAIJ(Mat A, MatType type, MatReuse reuse, Mat *newmat)
112: {
113: Mat B = *newmat;
115: PetscFunctionBegin;
116: if (reuse == MAT_INITIAL_MATRIX) PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &B));
117: B->ops->duplicate = MatDuplicate_SeqBAIJ;
118: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_seqbaijlibxsmm_seqdense_C", NULL));
119: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqbaijlibxsmm_seqbaij_C", NULL));
120: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQBAIJ));
121: *newmat = B;
122: PetscFunctionReturn(PETSC_SUCCESS);
123: }
125: PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqBAIJLIBXSMM(Mat A, MatType type, MatReuse reuse, Mat *newmat)
126: {
127: Mat B = *newmat;
128: PetscBool sametype;
130: PetscFunctionBegin;
131: if (reuse == MAT_INITIAL_MATRIX) PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &B));
132: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQBAIJLIBXSMM, &sametype));
133: if (!sametype) {
134: B->ops->duplicate = MatDuplicate_SeqBAIJLIBXSMM;
135: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQBAIJLIBXSMM));
136: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_seqbaijlibxsmm_seqdense_C", MatProductSetFromOptions_SeqBAIJLIBXSMM_SeqDense));
137: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqbaijlibxsmm_seqbaij_C", MatConvert_SeqBAIJLIBXSMM_SeqBAIJ));
138: }
139: *newmat = B;
140: PetscFunctionReturn(PETSC_SUCCESS);
141: }
143: /*MC
144: MATSEQBAIJLIBXSMM - "seqbaijlibxsmm" - A sequential block sparse matrix that uses LIBXSMM kernels for products with `MATSEQDENSE` matrices
146: Options Database Key:
147: . -mat_type seqbaijlibxsmm - sets the matrix type to `MATSEQBAIJLIBXSMM` during a call to `MatSetFromOptions()`
149: Level: beginner
151: Notes:
152: This matrix type is available when PETSc is configured with `--download-libxsmm` or `--with-libxsmm-dir=directory`.
154: It has the same storage format and supports the same operations as `MATSEQBAIJ`.
156: .seealso: [](ch_matrices), `Mat`, `MATBAIJLIBXSMM`, `MATMPIBAIJLIBXSMM`, `MATSEQBAIJ`, `MatMatMult()`
157: M*/
158: PETSC_EXTERN PetscErrorCode MatCreate_SeqBAIJLIBXSMM(Mat A)
159: {
160: PetscFunctionBegin;
161: PetscCall(MatSetType(A, MATSEQBAIJ));
162: PetscCall(MatConvert_SeqBAIJ_SeqBAIJLIBXSMM(A, MATSEQBAIJLIBXSMM, MAT_INPLACE_MATRIX, &A));
163: PetscFunctionReturn(PETSC_SUCCESS);
164: }