From ed9ed3dea34fef5853dd3f23fdf30a7807a29b2c Mon Sep 17 00:00:00 2001 From: James Wright Date: Tue, 12 Mar 2024 09:15:33 -0600 Subject: [PATCH 1/7] fluids: Make MatShell functions static This is done primarily to allow MatCeed to coexist with the context operators for right now. This has the nasty side-effect of forcing `Mat` objects to be NULL when they're passed into CreateMatShell_Ceed, but that will soon be replaced with MatCeed, so it is only temporary for the transition in this PR --- examples/fluids/Makefile | 1 + examples/fluids/include/petsc_ops.h | 2 -- examples/fluids/src/differential_filter.c | 2 +- examples/fluids/src/grid_anisotropy_tensor.c | 2 +- examples/fluids/src/petsc_ops.c | 16 ++++++++++++---- examples/fluids/src/setupts.c | 7 ++----- examples/fluids/src/turb_spanstats.c | 2 +- .../fluids/src/velocity_gradient_projection.c | 2 +- 8 files changed, 19 insertions(+), 15 deletions(-) diff --git a/examples/fluids/Makefile b/examples/fluids/Makefile index c8c3956929..f5a1d7c8c0 100644 --- a/examples/fluids/Makefile +++ b/examples/fluids/Makefile @@ -41,6 +41,7 @@ AFLAGS ?= -fsanitize=address #-fsanitize=undefined -fno-omit-frame-pointer CFLAGS += $(if $(ASAN),$(AFLAGS)) FFLAGS += $(if $(ASAN),$(AFLAGS)) LDFLAGS += $(if $(ASAN),$(AFLAGS)) +CPPFLAGS += -I./include OBJDIR := build SRCDIR := src diff --git a/examples/fluids/include/petsc_ops.h b/examples/fluids/include/petsc_ops.h index 879ee0d447..95038baf39 100644 --- a/examples/fluids/include/petsc_ops.h +++ b/examples/fluids/include/petsc_ops.h @@ -23,8 +23,6 @@ struct OperatorApplyContext_ { PetscErrorCode OperatorApplyContextCreate(DM dm_x, DM dm_y, Ceed ceed, CeedOperator op_apply, CeedVector x_ceed, CeedVector y_ceed, Vec X_loc, Vec Y_loc, OperatorApplyContext *op_apply_ctx); PetscErrorCode OperatorApplyContextDestroy(OperatorApplyContext op_apply_ctx); -PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D); -PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y); PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat); PetscErrorCode DMGetGlobalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type); diff --git a/examples/fluids/src/differential_filter.c b/examples/fluids/src/differential_filter.c index 873ceef907..e9c1e06e21 100644 --- a/examples/fluids/src/differential_filter.c +++ b/examples/fluids/src/differential_filter.c @@ -80,7 +80,7 @@ PetscErrorCode DifferentialFilterCreateOperators(Ceed ceed, User user, CeedData { // Setup LHS Operator and KSP for the differential filtering solve CeedOperator op_lhs; OperatorApplyContext mat_ctx; - Mat mat_lhs; + Mat mat_lhs = NULL; CeedInt num_comp_qd; PetscInt dim, num_comp_grid_aniso; CeedElemRestriction elem_restr_grid_aniso; diff --git a/examples/fluids/src/grid_anisotropy_tensor.c b/examples/fluids/src/grid_anisotropy_tensor.c index e397bf1f14..428843232c 100644 --- a/examples/fluids/src/grid_anisotropy_tensor.c +++ b/examples/fluids/src/grid_anisotropy_tensor.c @@ -73,7 +73,7 @@ PetscErrorCode GridAnisotropyTensorProjectionSetupApply(Ceed ceed, User user, Ce PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", *elem_restr_grid_aniso, basis_grid_aniso, CEED_VECTOR_ACTIVE)); { // -- Setup KSP for L^2 projection - Mat mat_mass; + Mat mat_mass = NULL; OperatorApplyContext mass_matop_ctx; PetscCall(OperatorApplyContextCreate(grid_aniso_proj->dm, grid_aniso_proj->dm, ceed, op_mass, NULL, NULL, NULL, NULL, &mass_matop_ctx)); diff --git a/examples/fluids/src/petsc_ops.c b/examples/fluids/src/petsc_ops.c index bf6ca5c424..f9fc2f801e 100644 --- a/examples/fluids/src/petsc_ops.c +++ b/examples/fluids/src/petsc_ops.c @@ -428,7 +428,7 @@ PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorAp // ----------------------------------------------------------------------------- // Wraps the libCEED operator for a MatShell // ----------------------------------------------------------------------------- -PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { +static PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { OperatorApplyContext ctx; PetscFunctionBeginUser; @@ -440,7 +440,7 @@ PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { // ----------------------------------------------------------------------------- // Returns the computed diagonal of the operator // ----------------------------------------------------------------------------- -PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { +static PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { OperatorApplyContext ctx; Vec Y_loc; PetscMemType mem_type; @@ -474,7 +474,7 @@ PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { * @brief Create PETSc MatShell object for the corresponding OperatorApplyContext * * @param[in] ctx Context that does the action of the operator - * @param[out] mat MatShell for the operator + * @param[out] mat MatShell for the operator, must be NULL if mat should be created */ PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { MPI_Comm comm_x = PetscObjectComm((PetscObject)(ctx->dm_x)); @@ -488,7 +488,15 @@ PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { PetscCall(DMGetGlobalVectorInfo(ctx->dm_x, &X_loc_size, &X_size, &X_vec_type)); PetscCall(DMGetGlobalVectorInfo(ctx->dm_y, &Y_loc_size, &Y_size, &Y_vec_type)); - PetscCall(MatCreateShell(comm_x, Y_loc_size, X_loc_size, Y_size, X_size, ctx, mat)); + if (*mat) { + PetscBool mat_is_shell; + const char *type; + + PetscCall(PetscObjectTypeCompare((PetscObject)*mat, MATSHELL, &mat_is_shell)); + PetscObjectGetType((PetscObject)*mat, &type); + PetscCheck(mat_is_shell, comm_x, PETSC_ERR_ARG_WRONGSTATE, "Mat must be of type SHELL, instead got: %s", type); + PetscCall(MatShellSetContext(*mat, ctx)); + } else PetscCall(MatCreateShell(comm_x, Y_loc_size, X_loc_size, Y_size, X_size, ctx, mat)); PetscCall(MatShellSetContextDestroy(*mat, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy)); PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (void (*)(void))MatMult_Ceed)); PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed)); diff --git a/examples/fluids/src/setupts.c b/examples/fluids/src/setupts.c index df518c422c..d49978ac27 100644 --- a/examples/fluids/src/setupts.c +++ b/examples/fluids/src/setupts.c @@ -35,7 +35,7 @@ PetscErrorCode CreateKSPMassOperator(User user, CeedData ceed_data) { PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE)); { // -- Setup KSP for mass operator - Mat mat_mass; + Mat mat_mass = NULL; Vec Zeros_loc; MPI_Comm comm = PetscObjectComm((PetscObject)dm); @@ -261,10 +261,7 @@ PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, PetscReal OperatorApplyContext op_ijacobian_ctx; OperatorApplyContextCreate(user->dm, user->dm, user->ceed, user->op_ijacobian, user->q_ceed, user->g_ceed, user->Q_dot_loc, NULL, &op_ijacobian_ctx); - PetscCall(MatShellSetContext(J, op_ijacobian_ctx)); - PetscCall(MatShellSetContextDestroy(J, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy)); - PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))MatMult_Ceed)); - PetscCall(MatShellSetOperation(J, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed)); + PetscCall(CreateMatShell_Ceed(op_ijacobian_ctx, &J)); PetscCall(MatSetUp(J)); } if (!J_pre_is_shell) { diff --git a/examples/fluids/src/turb_spanstats.c b/examples/fluids/src/turb_spanstats.c index 0ee268566e..28bc7c0775 100644 --- a/examples/fluids/src/turb_spanstats.c +++ b/examples/fluids/src/turb_spanstats.c @@ -333,7 +333,7 @@ PetscErrorCode SetupL2ProjectionStats(Ceed ceed, User user, CeedData ceed_data, { // Setup KSP for L^2 projection OperatorApplyContext M_ctx; - Mat mat_mass; + Mat mat_mass = NULL; KSP ksp; PetscCall(OperatorApplyContextCreate(user->spanstats.dm, user->spanstats.dm, user->ceed, op_mass, NULL, NULL, NULL, NULL, &M_ctx)); diff --git a/examples/fluids/src/velocity_gradient_projection.c b/examples/fluids/src/velocity_gradient_projection.c index 8c503a1022..5cc2572f15 100644 --- a/examples/fluids/src/velocity_gradient_projection.c +++ b/examples/fluids/src/velocity_gradient_projection.c @@ -102,7 +102,7 @@ PetscErrorCode VelocityGradientProjectionSetup(Ceed ceed, User user, CeedData ce PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", elem_restr_grad_velo, basis_grad_velo, CEED_VECTOR_ACTIVE)); { // -- Setup KSP for L^2 projection with lumped mass operator - Mat mat_mass; + Mat mat_mass = NULL; OperatorApplyContext mass_matop_ctx; MPI_Comm comm = PetscObjectComm((PetscObject)grad_velo_proj->dm); From c8564c30f55def119ddd4dd977f1131cd364f231 Mon Sep 17 00:00:00 2001 From: James Wright Date: Tue, 12 Mar 2024 09:22:24 -0600 Subject: [PATCH 2/7] fluids: Add MatCeed Code from Ratel --- examples/fluids/include/mat-ceed-impl.h | 69 ++ examples/fluids/include/mat-ceed.h | 41 + examples/fluids/src/mat-ceed.c | 1376 +++++++++++++++++++++++ 3 files changed, 1486 insertions(+) create mode 100644 examples/fluids/include/mat-ceed-impl.h create mode 100644 examples/fluids/include/mat-ceed.h create mode 100644 examples/fluids/src/mat-ceed.c diff --git a/examples/fluids/include/mat-ceed-impl.h b/examples/fluids/include/mat-ceed-impl.h new file mode 100644 index 0000000000..8342ec6f15 --- /dev/null +++ b/examples/fluids/include/mat-ceed-impl.h @@ -0,0 +1,69 @@ +#ifndef MAT_CEED_IMPL_H +#define MAT_CEED_IMPL_H + +#include +#include +#include + +#if defined(__clang_analyzer__) +#define MATCEED_EXTERN extern +#elif defined(__cplusplus) +#define MATCEED_EXTERN extern "C" +#else +#define MATCEED_EXTERN extern +#endif + +#if defined(__clang_analyzer__) +#define MATCEED_INTERN +#else +#define MATCEED_INTERN MATCEED_EXTERN __attribute__((visibility("hidden"))) +#endif + +/** + @brief Calls a libCEED function and then checks the resulting error code. + If the error code is non-zero, then a PETSc error is set with the libCEED error message. +**/ +#define PetscCeedCall(ceed_, ...) \ + do { \ + int ierr_q_ = __VA_ARGS__; \ + if (ierr_q_ != CEED_ERROR_SUCCESS) { \ + const char *error_message; \ + CeedGetErrorMessage(ceed_, &error_message); \ + SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "%s", error_message); \ + } \ + } while (0) + +// MatCeed context for applying composite CeedOperator on a DM +typedef struct MatCeedContext_private *MatCeedContext; +struct MatCeedContext_private { + Ceed ceed; + char *name, *internal_mat_type; + PetscMemType mem_type; + PetscInt ref_count, num_mats_assembled_full, num_mats_assembled_pbd; + PetscBool is_destroyed, is_ceed_pbd_valid, is_ceed_vpbd_valid; + PetscLogEvent log_event_mult, log_event_mult_transpose; + DM dm_x, dm_y; + Mat *mats_assembled_full, *mats_assembled_pbd, mat_assembled_full_internal, mat_assembled_pbd_internal; + Vec X_loc, Y_loc_transpose; + CeedVector x_loc, y_loc, coo_values_full, coo_values_pbd; + CeedOperator op_mult, op_mult_transpose; + PetscLogDouble flops_mult, flops_mult_transpose; +}; + +// Context data +MATCEED_INTERN PetscErrorCode MatCeedContextCreate(DM dm_x, DM dm_y, Vec X_loc, Vec Y_loc_transpose, CeedOperator op_mult, + CeedOperator op_mult_transpose, PetscLogEvent log_event_mult, + PetscLogEvent log_event_mult_transpose, MatCeedContext *ctx); +MATCEED_INTERN PetscErrorCode MatCeedContextReference(MatCeedContext ctx); +MATCEED_INTERN PetscErrorCode MatCeedContextReferenceCopy(MatCeedContext ctx, MatCeedContext *ctx_copy); +MATCEED_INTERN PetscErrorCode MatCeedContextDestroy(MatCeedContext ctx); + +// Mat Ceed +MATCEED_INTERN PetscErrorCode MatGetDiagonal_Ceed(Mat A, Vec D); +MATCEED_INTERN PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y); +MATCEED_INTERN PetscErrorCode MatMultTranspose_Ceed(Mat A, Vec Y, Vec X); + +extern PetscClassId MATCEED_CLASSID; +extern PetscLogEvent MATCEED_MULT, MATCEED_MULT_TRANSPOSE; + +#endif // MAT_CEED_IMPL_H diff --git a/examples/fluids/include/mat-ceed.h b/examples/fluids/include/mat-ceed.h new file mode 100644 index 0000000000..c999a6d9b4 --- /dev/null +++ b/examples/fluids/include/mat-ceed.h @@ -0,0 +1,41 @@ +#ifndef MAT_CEED_H +#define MAT_CEED_H + +#include +#include +#include + +#define MATCEED "ceed" + +#if defined(__clang_analyzer__) +#define MATCEED_EXTERN extern +#elif defined(__cplusplus) +#define MATCEED_EXTERN extern "C" +#else +#define MATCEED_EXTERN extern +#endif + +#if defined(__clang_analyzer__) +#define MATCEED_INTERN +#else +#define MATCEED_INTERN MATCEED_EXTERN __attribute__((visibility("hidden"))) +#endif + +// Context data +MATCEED_INTERN PetscErrorCode MatCeedCreate(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperator op_mult_transpose, Mat *mat); +MATCEED_INTERN PetscErrorCode MatCeedCopy(Mat mat_ceed, Mat mat_other); +MATCEED_INTERN PetscErrorCode MatCeedAssembleCOO(Mat mat_ceed, Mat mat_coo); +MATCEED_INTERN PetscErrorCode MatCeedSetContext(Mat mat, PetscErrorCode (*f)(void *), void *ctx); +MATCEED_INTERN PetscErrorCode MatCeedGetContext(Mat mat, void *ctx); +MATCEED_INTERN PetscErrorCode MatCeedSetInnerMatType(Mat mat, MatType type); +MATCEED_INTERN PetscErrorCode MatCeedGetInnerMatType(Mat mat, MatType *type); +MATCEED_INTERN PetscErrorCode MatCeedSetOperation(Mat mat, MatOperation op, void (*g)(void)); +MATCEED_INTERN PetscErrorCode MatCeedSetLocalVectors(Mat mat, Vec X_loc, Vec Y_loc_transpose); +MATCEED_INTERN PetscErrorCode MatCeedGetLocalVectors(Mat mat, Vec *X_loc, Vec *Y_loc_transpose); +MATCEED_INTERN PetscErrorCode MatCeedRestoreLocalVectors(Mat mat, Vec *X_loc, Vec *Y_loc_transpose); +MATCEED_INTERN PetscErrorCode MatCeedGetCeedOperators(Mat mat, CeedOperator *op_mult, CeedOperator *op_mult_transpose); +MATCEED_INTERN PetscErrorCode MatCeedRestoreCeedOperators(Mat mat, CeedOperator *op_mult, CeedOperator *op_mult_transpose); +MATCEED_INTERN PetscErrorCode MatCeedSetLogEvents(Mat mat, PetscLogEvent log_event_mult, PetscLogEvent log_event_mult_transpose); +MATCEED_INTERN PetscErrorCode MatCeedGetLogEvents(Mat mat, PetscLogEvent *log_event_mult, PetscLogEvent *log_event_mult_transpose); + +#endif // MAT_CEED_H diff --git a/examples/fluids/src/mat-ceed.c b/examples/fluids/src/mat-ceed.c new file mode 100644 index 0000000000..35e8bb41a2 --- /dev/null +++ b/examples/fluids/src/mat-ceed.c @@ -0,0 +1,1376 @@ +/// @file +/// MatCeed and it's related operators + +#include +#include +#include +#include +#include +#include +#include + +PetscClassId MATCEED_CLASSID; +PetscLogEvent MATCEED_MULT, MATCEED_MULT_TRANSPOSE; + +/** + @brief Register MATCEED log events. + + Not collective across MPI processes. + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatCeedRegisterLogEvents() { + static bool registered = false; + + PetscFunctionBeginUser; + if (registered) PetscFunctionReturn(PETSC_SUCCESS); + PetscCall(PetscClassIdRegister("MATCEED", &MATCEED_CLASSID)); + PetscCall(PetscLogEventRegister("MATCEED Mult", MATCEED_CLASSID, &MATCEED_MULT)); + PetscCall(PetscLogEventRegister("MATCEED Mult Transpose", MATCEED_CLASSID, &MATCEED_MULT_TRANSPOSE)); + registered = true; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Translate PetscMemType to CeedMemType + + @param[in] mem_type PetscMemType + + @return Equivalent CeedMemType +**/ +static inline CeedMemType MemTypeP2C(PetscMemType mem_type) { return PetscMemTypeDevice(mem_type) ? CEED_MEM_DEVICE : CEED_MEM_HOST; } + +/** + @brief Translate array of `CeedInt` to `PetscInt`. + If the types differ, `array_ceed` is freed with `free()` and `array_petsc` is allocated with `malloc()`. + Caller is responsible for freeing `array_petsc` with `free()`. + + Not collective across MPI processes. + + @param[in] num_entries Number of array entries + @param[in,out] array_ceed Array of `CeedInt` + @param[out] array_petsc Array of `PetscInt` + + @return An error code: 0 - success, otherwise - failure +**/ +static inline PetscErrorCode IntArrayC2P(PetscInt num_entries, CeedInt **array_ceed, PetscInt **array_petsc) { + const CeedInt int_c = 0; + const PetscInt int_p = 0; + + PetscFunctionBeginUser; + if (sizeof(int_c) == sizeof(int_p)) { + *array_petsc = (PetscInt *)*array_ceed; + } else { + *array_petsc = malloc(num_entries * sizeof(PetscInt)); + for (PetscInt i = 0; i < num_entries; i++) (*array_petsc)[i] = (*array_ceed)[i]; + free(*array_ceed); + } + *array_ceed = NULL; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Transfer array from PETSc `Vec` to `CeedVector`. + + Collective across MPI processes. + + @param[in] ceed libCEED context + @param[in] X_petsc PETSc `Vec` + @param[out] mem_type PETSc `MemType` + @param[out] x_ceed `CeedVector` + + @return An error code: 0 - success, otherwise - failure +**/ +static inline PetscErrorCode VecP2C(Ceed ceed, Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { + PetscScalar *x; + + PetscFunctionBeginUser; + PetscCall(VecGetArrayAndMemType(X_petsc, &x, mem_type)); + PetscCeedCall(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Transfer array from `CeedVector` to PETSc `Vec`. + + Collective across MPI processes. + + @param[in] ceed libCEED context + @param[in] x_ceed `CeedVector` + @param[in] mem_type PETSc `MemType` + @param[out] X_petsc PETSc `Vec` + + @return An error code: 0 - success, otherwise - failure +**/ +static inline PetscErrorCode VecC2P(Ceed ceed, CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { + PetscScalar *x; + + PetscFunctionBeginUser; + PetscCeedCall(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); + PetscCall(VecRestoreArrayAndMemType(X_petsc, &x)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Transfer read only array from PETSc `Vec` to `CeedVector`. + + Collective across MPI processes. + + @param[in] ceed libCEED context + @param[in] X_petsc PETSc `Vec` + @param[out] mem_type PETSc `MemType` + @param[out] x_ceed `CeedVector` + + @return An error code: 0 - success, otherwise - failure +**/ +static inline PetscErrorCode VecReadP2C(Ceed ceed, Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { + PetscScalar *x; + + PetscFunctionBeginUser; + PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, mem_type)); + PetscCeedCall(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Transfer read only array from `CeedVector` to PETSc `Vec`. + + Collective across MPI processes. + + @param[in] ceed libCEED context + @param[in] x_ceed `CeedVector` + @param[in] mem_type PETSc `MemType` + @param[out] X_petsc PETSc `Vec` + + @return An error code: 0 - success, otherwise - failure +**/ +static inline PetscErrorCode VecReadC2P(Ceed ceed, CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { + PetscScalar *x; + + PetscFunctionBeginUser; + PetscCeedCall(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); + PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Setup inner `Mat` for `PC` operations not directly supported by libCEED. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to setup + @param[out] mat_inner Inner `Mat` + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatCeedSetupInnerMat(Mat mat_ceed, Mat *mat_inner) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + + PetscCheck(ctx->dm_x == ctx->dm_y, PetscObjectComm((PetscObject)mat_ceed), PETSC_ERR_SUP, "PC only supported for MATCEED on a single DM"); + + // Check cl mat type + { + PetscBool is_internal_mat_type_cl = PETSC_FALSE; + char internal_mat_type_cl[64]; + + // Check for specific CL inner mat type for this Mat + { + const char *mat_ceed_prefix = NULL; + + PetscCall(MatGetOptionsPrefix(mat_ceed, &mat_ceed_prefix)); + PetscOptionsBegin(PetscObjectComm((PetscObject)mat_ceed), mat_ceed_prefix, "", NULL); + PetscCall(PetscOptionsFList("-ceed_inner_mat_type", "MATCEED inner assembled MatType for PC support", NULL, MatList, internal_mat_type_cl, + internal_mat_type_cl, sizeof(internal_mat_type_cl), &is_internal_mat_type_cl)); + PetscOptionsEnd(); + if (is_internal_mat_type_cl) { + PetscCall(PetscFree(ctx->internal_mat_type)); + PetscCall(PetscStrallocpy(internal_mat_type_cl, &ctx->internal_mat_type)); + } + } + } + + // Create sparse matrix + { + MatType dm_mat_type, dm_mat_type_copy; + + PetscCall(DMGetMatType(ctx->dm_x, &dm_mat_type)); + PetscCall(PetscStrallocpy(dm_mat_type, (char **)&dm_mat_type_copy)); + PetscCall(DMSetMatType(ctx->dm_x, ctx->internal_mat_type)); + PetscCall(DMCreateMatrix(ctx->dm_x, mat_inner)); + PetscCall(DMSetMatType(ctx->dm_x, dm_mat_type_copy)); + PetscCall(PetscFree(dm_mat_type_copy)); + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Assemble the point block diagonal of a `MATCEED` into a `MATAIJ` or similar. + The `mat_coo` preallocation is set to match the sparsity pattern of `mat_ceed`. + The caller is responsible for assuring the global and local sizes are compatible, otherwise this function will fail. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to assemble + @param[in,out] mat_coo `MATAIJ` or similar to assemble into + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatCeedAssemblePointBlockDiagonalCOO(Mat mat_ceed, Mat mat_coo) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + + // Check if COO pattern set + { + PetscInt index = -1; + + for (PetscInt i = 0; i < ctx->num_mats_assembled_pbd; i++) { + if (ctx->mats_assembled_pbd[i] == mat_coo) index = i; + } + if (index == -1) { + PetscInt *rows_petsc = NULL, *cols_petsc = NULL; + CeedInt *rows_ceed, *cols_ceed; + PetscCount num_entries; + PetscLogStage stage_amg_setup; + + // -- Assemble sparsity pattern if mat hasn't been assembled before + PetscCall(PetscLogStageGetId("MATCEED Assembly Setup", &stage_amg_setup)); + if (stage_amg_setup == -1) { + PetscCall(PetscLogStageRegister("MATCEED Assembly Setup", &stage_amg_setup)); + } + PetscCall(PetscLogStagePush(stage_amg_setup)); + PetscCeedCall(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonalSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed)); + PetscCall(IntArrayC2P(num_entries, &rows_ceed, &rows_petsc)); + PetscCall(IntArrayC2P(num_entries, &cols_ceed, &cols_petsc)); + PetscCall(MatSetPreallocationCOOLocal(mat_coo, num_entries, rows_petsc, cols_petsc)); + free(rows_petsc); + free(cols_petsc); + if (!ctx->coo_values_pbd) PetscCeedCall(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_pbd)); + PetscCall(PetscRealloc(++ctx->num_mats_assembled_pbd * sizeof(Mat), &ctx->mats_assembled_pbd)); + ctx->mats_assembled_pbd[ctx->num_mats_assembled_pbd - 1] = mat_coo; + PetscCall(PetscLogStagePop()); + } + } + + // Assemble mat_ceed + PetscCall(MatAssemblyBegin(mat_coo, MAT_FINAL_ASSEMBLY)); + { + const CeedScalar *values; + MatType mat_type; + CeedMemType mem_type = CEED_MEM_HOST; + PetscBool is_spd, is_spd_known; + + PetscCall(MatGetType(mat_coo, &mat_type)); + if (strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE; + else if (strstr(mat_type, "kokkos")) mem_type = CEED_MEM_DEVICE; + else mem_type = CEED_MEM_HOST; + + PetscCeedCall(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonal(ctx->op_mult, ctx->coo_values_pbd, CEED_REQUEST_IMMEDIATE)); + PetscCeedCall(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_pbd, mem_type, &values)); + PetscCall(MatSetValuesCOO(mat_coo, values, INSERT_VALUES)); + PetscCall(MatIsSPDKnown(mat_ceed, &is_spd_known, &is_spd)); + if (is_spd_known) PetscCall(MatSetOption(mat_coo, MAT_SPD, is_spd)); + PetscCeedCall(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_pbd, &values)); + } + PetscCall(MatAssemblyEnd(mat_coo, MAT_FINAL_ASSEMBLY)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Assemble inner `Mat` for diagonal `PC` operations + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to invert + @param[in] use_ceed_pbd Boolean flag to use libCEED PBD assembly + @param[out] mat_inner Inner `Mat` for diagonal operations + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatCeedAssembleInnerBlockDiagonalMat(Mat mat_ceed, PetscBool use_ceed_pbd, Mat *mat_inner) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + if (use_ceed_pbd) { + // Check if COO pattern set + if (!ctx->mat_assembled_pbd_internal) PetscCall(MatCeedSetupInnerMat(mat_ceed, &ctx->mat_assembled_pbd_internal)); + + // Assemble mat_assembled_full_internal + PetscCall(MatCeedAssemblePointBlockDiagonalCOO(mat_ceed, ctx->mat_assembled_pbd_internal)); + if (mat_inner) *mat_inner = ctx->mat_assembled_pbd_internal; + } else { + // Check if COO pattern set + if (!ctx->mat_assembled_full_internal) PetscCall(MatCeedSetupInnerMat(mat_ceed, &ctx->mat_assembled_full_internal)); + + // Assemble mat_assembled_full_internal + PetscCall(MatCeedAssembleCOO(mat_ceed, ctx->mat_assembled_full_internal)); + if (mat_inner) *mat_inner = ctx->mat_assembled_full_internal; + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Get `MATCEED` diagonal block for Jacobi. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to invert + @param[out] mat_block The diagonal block matrix + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatGetDiagonalBlock_Ceed(Mat mat_ceed, Mat *mat_block) { + Mat mat_inner = NULL; + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + + // Assemble inner mat if needed + PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_pbd_valid, &mat_inner)); + + // Get block diagonal + PetscCall(MatGetDiagonalBlock(mat_inner, mat_block)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Invert `MATCEED` diagonal block for Jacobi. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to invert + @param[out] values The block inverses in column major order + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatInvertBlockDiagonal_Ceed(Mat mat_ceed, const PetscScalar **values) { + Mat mat_inner = NULL; + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + + // Assemble inner mat if needed + PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_pbd_valid, &mat_inner)); + + // Invert PB diagonal + PetscCall(MatInvertBlockDiagonal(mat_inner, values)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Invert `MATCEED` variable diagonal block for Jacobi. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to invert + @param[in] num_blocks The number of blocks on the process + @param[in] block_sizes The size of each block on the process + @param[out] values The block inverses in column major order + + @return An error code: 0 - success, otherwise - failure +**/ +static PetscErrorCode MatInvertVariableBlockDiagonal_Ceed(Mat mat_ceed, PetscInt num_blocks, const PetscInt *block_sizes, PetscScalar *values) { + Mat mat_inner = NULL; + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + + // Assemble inner mat if needed + PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_vpbd_valid, &mat_inner)); + + // Invert PB diagonal + PetscCall(MatInvertVariableBlockDiagonal(mat_inner, num_blocks, block_sizes, values)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +// ----------------------------------------------------------------------------- +// MatCeed +// ----------------------------------------------------------------------------- + +/** + @brief Create PETSc `Mat` from libCEED operators. + + Collective across MPI processes. + + @param[in] dm_x Input `DM` + @param[in] dm_y Output `DM` + @param[in] op_mult `CeedOperator` for forward evaluation + @param[in] op_mult_transpose `CeedOperator` for transpose evaluation + @param[out] mat New MatCeed + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedCreate(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperator op_mult_transpose, Mat *mat) { + PetscInt X_l_size, X_g_size, Y_l_size, Y_g_size; + VecType vec_type; + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatCeedRegisterLogEvents()); + + // Collect context data + PetscCall(DMGetVecType(dm_x, &vec_type)); + { + Vec X; + + PetscCall(DMGetGlobalVector(dm_x, &X)); + PetscCall(VecGetSize(X, &X_g_size)); + PetscCall(VecGetLocalSize(X, &X_l_size)); + PetscCall(DMRestoreGlobalVector(dm_x, &X)); + } + if (dm_y) { + Vec Y; + + PetscCall(DMGetGlobalVector(dm_y, &Y)); + PetscCall(VecGetSize(Y, &Y_g_size)); + PetscCall(VecGetLocalSize(Y, &Y_l_size)); + PetscCall(DMRestoreGlobalVector(dm_y, &Y)); + } else { + dm_y = dm_x; + Y_g_size = X_g_size; + Y_l_size = X_l_size; + } + // Create context + { + Vec X_loc, Y_loc_transpose = NULL; + + PetscCall(DMCreateLocalVector(dm_x, &X_loc)); + PetscCall(VecZeroEntries(X_loc)); + if (op_mult_transpose) { + PetscCall(DMCreateLocalVector(dm_y, &Y_loc_transpose)); + PetscCall(VecZeroEntries(Y_loc_transpose)); + } + PetscCall(MatCeedContextCreate(dm_x, dm_y, X_loc, Y_loc_transpose, op_mult, op_mult_transpose, MATCEED_MULT, MATCEED_MULT_TRANSPOSE, &ctx)); + PetscCall(VecDestroy(&X_loc)); + PetscCall(VecDestroy(&Y_loc_transpose)); + } + + // Create mat + PetscCall(MatCreateShell(PetscObjectComm((PetscObject)dm_x), Y_l_size, X_l_size, Y_g_size, X_g_size, ctx, mat)); + PetscCall(PetscObjectChangeTypeName((PetscObject)*mat, MATCEED)); + // -- Set block and variable block sizes + if (dm_x == dm_y) { + MatType dm_mat_type, dm_mat_type_copy; + Mat temp_mat; + + PetscCall(DMGetMatType(dm_x, &dm_mat_type)); + PetscCall(PetscStrallocpy(dm_mat_type, (char **)&dm_mat_type_copy)); + PetscCall(DMSetMatType(dm_x, MATAIJ)); + PetscCall(DMCreateMatrix(dm_x, &temp_mat)); + PetscCall(DMSetMatType(dm_x, dm_mat_type_copy)); + PetscCall(PetscFree(dm_mat_type_copy)); + + { + PetscInt block_size, num_blocks, max_vblock_size = PETSC_INT_MAX; + const PetscInt *vblock_sizes; + + // -- Get block sizes + PetscCall(MatGetBlockSize(temp_mat, &block_size)); + PetscCall(MatGetVariableBlockSizes(temp_mat, &num_blocks, &vblock_sizes)); + { + PetscInt local_min_max[2] = {0}, global_min_max[2] = {0, PETSC_INT_MAX}; + + for (PetscInt i = 0; i < num_blocks; i++) local_min_max[1] = PetscMax(local_min_max[1], vblock_sizes[i]); + PetscCall(PetscGlobalMinMaxInt(PetscObjectComm((PetscObject)dm_x), local_min_max, global_min_max)); + max_vblock_size = global_min_max[1]; + } + + // -- Copy block sizes + if (block_size > 1) PetscCall(MatSetBlockSize(*mat, block_size)); + if (num_blocks) PetscCall(MatSetVariableBlockSizes(*mat, num_blocks, (PetscInt *)vblock_sizes)); + + // -- Check libCEED compatibility + { + bool is_composite; + + ctx->is_ceed_pbd_valid = PETSC_TRUE; + ctx->is_ceed_vpbd_valid = PETSC_TRUE; + PetscCeedCall(ctx->ceed, CeedOperatorIsComposite(op_mult, &is_composite)); + if (is_composite) { + CeedInt num_sub_operators; + CeedOperator *sub_operators; + + PetscCeedCall(ctx->ceed, CeedCompositeOperatorGetNumSub(op_mult, &num_sub_operators)); + PetscCeedCall(ctx->ceed, CeedCompositeOperatorGetSubList(op_mult, &sub_operators)); + for (CeedInt i = 0; i < num_sub_operators; i++) { + CeedInt num_bases, num_comp; + CeedBasis *active_bases; + CeedOperatorAssemblyData assembly_data; + + PetscCeedCall(ctx->ceed, CeedOperatorGetOperatorAssemblyData(sub_operators[i], &assembly_data)); + PetscCeedCall(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL)); + PetscCeedCall(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp)); + if (num_bases > 1) { + ctx->is_ceed_pbd_valid = PETSC_FALSE; + ctx->is_ceed_vpbd_valid = PETSC_FALSE; + } + if (num_comp != block_size) ctx->is_ceed_pbd_valid = PETSC_FALSE; + if (num_comp < max_vblock_size) ctx->is_ceed_vpbd_valid = PETSC_FALSE; + } + } else { + // LCOV_EXCL_START + CeedInt num_bases, num_comp; + CeedBasis *active_bases; + CeedOperatorAssemblyData assembly_data; + + PetscCeedCall(ctx->ceed, CeedOperatorGetOperatorAssemblyData(op_mult, &assembly_data)); + PetscCeedCall(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL)); + PetscCeedCall(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp)); + if (num_bases > 1) { + ctx->is_ceed_pbd_valid = PETSC_FALSE; + ctx->is_ceed_vpbd_valid = PETSC_FALSE; + } + if (num_comp != block_size) ctx->is_ceed_pbd_valid = PETSC_FALSE; + if (num_comp < max_vblock_size) ctx->is_ceed_vpbd_valid = PETSC_FALSE; + // LCOV_EXCL_STOP + } + { + PetscInt local_is_valid[2], global_is_valid[2]; + + local_is_valid[0] = local_is_valid[1] = ctx->is_ceed_pbd_valid; + PetscCall(PetscGlobalMinMaxInt(PetscObjectComm((PetscObject)dm_x), local_is_valid, global_is_valid)); + ctx->is_ceed_pbd_valid = global_is_valid[0]; + local_is_valid[0] = local_is_valid[1] = ctx->is_ceed_vpbd_valid; + PetscCall(PetscGlobalMinMaxInt(PetscObjectComm((PetscObject)dm_x), local_is_valid, global_is_valid)); + ctx->is_ceed_vpbd_valid = global_is_valid[0]; + } + } + } + PetscCall(MatDestroy(&temp_mat)); + } + // -- Set internal mat type + { + VecType vec_type; + MatType internal_mat_type = MATAIJ; + + PetscCall(VecGetType(ctx->X_loc, &vec_type)); + if (strstr(vec_type, VECCUDA)) internal_mat_type = MATAIJCUSPARSE; + else if (strstr(vec_type, VECKOKKOS)) internal_mat_type = MATAIJKOKKOS; + else internal_mat_type = MATAIJ; + PetscCall(PetscStrallocpy(internal_mat_type, &ctx->internal_mat_type)); + } + // -- Set mat operations + PetscCall(MatShellSetContextDestroy(*mat, (PetscErrorCode(*)(void *))MatCeedContextDestroy)); + PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (void (*)(void))MatMult_Ceed)); + if (op_mult_transpose) PetscCall(MatShellSetOperation(*mat, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_Ceed)); + PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_Ceed)); + PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL_BLOCK, (void (*)(void))MatGetDiagonalBlock_Ceed)); + PetscCall(MatShellSetOperation(*mat, MATOP_INVERT_BLOCK_DIAGONAL, (void (*)(void))MatInvertBlockDiagonal_Ceed)); + PetscCall(MatShellSetOperation(*mat, MATOP_INVERT_VBLOCK_DIAGONAL, (void (*)(void))MatInvertVariableBlockDiagonal_Ceed)); + PetscCall(MatShellSetVecType(*mat, vec_type)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Copy `MATCEED` into a compatible `Mat` with type `MatShell` or `MATCEED`. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to copy from + @param[out] mat_other `MatShell` or `MATCEED` to copy into + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedCopy(Mat mat_ceed, Mat mat_other) { + PetscFunctionBeginUser; + PetscCall(MatCeedRegisterLogEvents()); + + // Check type compatibility + { + MatType mat_type_ceed, mat_type_other; + + PetscCall(MatGetType(mat_ceed, &mat_type_ceed)); + PetscCheck(!strcmp(mat_type_ceed, MATCEED), PETSC_COMM_SELF, PETSC_ERR_LIB, "mat_ceed must have type " MATCEED); + PetscCall(MatGetType(mat_ceed, &mat_type_other)); + PetscCheck(!strcmp(mat_type_other, MATCEED) || !strcmp(mat_type_other, MATSHELL), PETSC_COMM_SELF, PETSC_ERR_LIB, + "mat_other must have type " MATCEED " or " MATSHELL); + } + + // Check dimension compatibility + { + PetscInt X_l_ceed_size, X_g_ceed_size, Y_l_ceed_size, Y_g_ceed_size, X_l_other_size, X_g_other_size, Y_l_other_size, Y_g_other_size; + + PetscCall(MatGetSize(mat_ceed, &Y_g_ceed_size, &X_g_ceed_size)); + PetscCall(MatGetLocalSize(mat_ceed, &Y_l_ceed_size, &X_l_ceed_size)); + PetscCall(MatGetSize(mat_ceed, &Y_g_other_size, &X_g_other_size)); + PetscCall(MatGetLocalSize(mat_ceed, &Y_l_other_size, &X_l_other_size)); + PetscCheck((Y_g_ceed_size == Y_g_other_size) && (X_g_ceed_size == X_g_other_size) && (Y_l_ceed_size == Y_l_other_size) && + (X_l_ceed_size == X_l_other_size), + PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, + "mat_ceed and mat_other must have compatible sizes; found mat_ceed (Global: %" PetscInt_FMT ", %" PetscInt_FMT + "; Local: %" PetscInt_FMT ", %" PetscInt_FMT ") mat_other (Global: %" PetscInt_FMT ", %" PetscInt_FMT "; Local: %" PetscInt_FMT + ", %" PetscInt_FMT ")", + Y_g_ceed_size, X_g_ceed_size, Y_l_ceed_size, X_l_ceed_size, Y_g_other_size, X_g_other_size, Y_l_other_size, X_l_other_size); + } + + // Convert + { + VecType vec_type; + MatCeedContext ctx; + + PetscCall(PetscObjectChangeTypeName((PetscObject)mat_other, MATCEED)); + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + PetscCall(MatCeedContextReference(ctx)); + PetscCall(MatShellSetContext(mat_other, ctx)); + PetscCall(MatShellSetContextDestroy(mat_other, (PetscErrorCode(*)(void *))MatCeedContextDestroy)); + PetscCall(MatShellSetOperation(mat_other, MATOP_MULT, (void (*)(void))MatMult_Ceed)); + if (ctx->op_mult_transpose) PetscCall(MatShellSetOperation(mat_other, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_Ceed)); + PetscCall(MatShellSetOperation(mat_other, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_Ceed)); + PetscCall(MatShellSetOperation(mat_other, MATOP_GET_DIAGONAL_BLOCK, (void (*)(void))MatGetDiagonalBlock_Ceed)); + PetscCall(MatShellSetOperation(mat_other, MATOP_INVERT_BLOCK_DIAGONAL, (void (*)(void))MatInvertBlockDiagonal_Ceed)); + PetscCall(MatShellSetOperation(mat_other, MATOP_INVERT_VBLOCK_DIAGONAL, (void (*)(void))MatInvertVariableBlockDiagonal_Ceed)); + { + PetscInt block_size; + + PetscCall(MatGetBlockSize(mat_ceed, &block_size)); + if (block_size > 1) PetscCall(MatSetBlockSize(mat_other, block_size)); + } + { + PetscInt num_blocks; + const PetscInt *block_sizes; + + PetscCall(MatGetVariableBlockSizes(mat_ceed, &num_blocks, &block_sizes)); + if (num_blocks) PetscCall(MatSetVariableBlockSizes(mat_other, num_blocks, (PetscInt *)block_sizes)); + } + PetscCall(DMGetVecType(ctx->dm_x, &vec_type)); + PetscCall(MatShellSetVecType(mat_other, vec_type)); + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Assemble a `MATCEED` into a `MATAIJ` or similar. + The `mat_coo` preallocation is set to match the sparsity pattern of `mat_ceed`. + The caller is responsible for assuring the global and local sizes are compatible, otherwise this function will fail. + + Collective across MPI processes. + + @param[in] mat_ceed `MATCEED` to assemble + @param[in,out] mat_coo `MATAIJ` or similar to assemble into + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedAssembleCOO(Mat mat_ceed, Mat mat_coo) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat_ceed, &ctx)); + + // Check if COO pattern set + { + PetscInt index = -1; + + for (PetscInt i = 0; i < ctx->num_mats_assembled_full; i++) { + if (ctx->mats_assembled_full[i] == mat_coo) index = i; + } + if (index == -1) { + PetscInt *rows_petsc = NULL, *cols_petsc = NULL; + CeedInt *rows_ceed, *cols_ceed; + PetscCount num_entries; + PetscLogStage stage_amg_setup; + + // -- Assemble sparsity pattern if mat hasn't been assembled before + PetscCall(PetscLogStageGetId("MATCEED Assembly Setup", &stage_amg_setup)); + if (stage_amg_setup == -1) { + PetscCall(PetscLogStageRegister("MATCEED Assembly Setup", &stage_amg_setup)); + } + PetscCall(PetscLogStagePush(stage_amg_setup)); + PetscCeedCall(ctx->ceed, CeedOperatorLinearAssembleSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed)); + PetscCall(IntArrayC2P(num_entries, &rows_ceed, &rows_petsc)); + PetscCall(IntArrayC2P(num_entries, &cols_ceed, &cols_petsc)); + PetscCall(MatSetPreallocationCOOLocal(mat_coo, num_entries, rows_petsc, cols_petsc)); + free(rows_petsc); + free(cols_petsc); + if (!ctx->coo_values_full) PetscCeedCall(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_full)); + PetscCall(PetscRealloc(++ctx->num_mats_assembled_full * sizeof(Mat), &ctx->mats_assembled_full)); + ctx->mats_assembled_full[ctx->num_mats_assembled_full - 1] = mat_coo; + PetscCall(PetscLogStagePop()); + } + } + + // Assemble mat_ceed + PetscCall(MatAssemblyBegin(mat_coo, MAT_FINAL_ASSEMBLY)); + { + const CeedScalar *values; + MatType mat_type; + CeedMemType mem_type = CEED_MEM_HOST; + PetscBool is_spd, is_spd_known; + + PetscCall(MatGetType(mat_coo, &mat_type)); + if (strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE; + else if (strstr(mat_type, "kokkos")) mem_type = CEED_MEM_DEVICE; + else mem_type = CEED_MEM_HOST; + + PetscCeedCall(ctx->ceed, CeedOperatorLinearAssemble(ctx->op_mult, ctx->coo_values_full)); + PetscCeedCall(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_full, mem_type, &values)); + PetscCall(MatSetValuesCOO(mat_coo, values, INSERT_VALUES)); + PetscCall(MatIsSPDKnown(mat_ceed, &is_spd_known, &is_spd)); + if (is_spd_known) PetscCall(MatSetOption(mat_coo, MAT_SPD, is_spd)); + PetscCeedCall(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_full, &values)); + } + PetscCall(MatAssemblyEnd(mat_coo, MAT_FINAL_ASSEMBLY)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Set user context for a `MATCEED`. + + Collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[in] f The context destroy function, or NULL + @param[in] ctx User context, or NULL to unset + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedSetContext(Mat mat, PetscErrorCode (*f)(void *), void *ctx) { + PetscContainer user_ctx = NULL; + + PetscFunctionBeginUser; + if (ctx) { + PetscCall(PetscContainerCreate(PetscObjectComm((PetscObject)mat), &user_ctx)); + PetscCall(PetscContainerSetPointer(user_ctx, ctx)); + PetscCall(PetscContainerSetUserDestroy(user_ctx, f)); + } + PetscCall(PetscObjectCompose((PetscObject)mat, "MatCeed user context", (PetscObject)user_ctx)); + PetscCall(PetscContainerDestroy(&user_ctx)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Retrieve the user context for a `MATCEED`. + + Collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[in] ctx User context + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedGetContext(Mat mat, void *ctx) { + PetscContainer user_ctx; + + PetscFunctionBeginUser; + PetscCall(PetscObjectQuery((PetscObject)mat, "MatCeed user context", (PetscObject *)&user_ctx)); + if (user_ctx) PetscCall(PetscContainerGetPointer(user_ctx, (void **)ctx)); + else *(void **)ctx = NULL; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Sets the inner matrix type as a string from the `MATCEED`. + + Collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[in] type Inner `MatType` to set + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedSetInnerMatType(Mat mat, MatType type) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + // Check if same + { + size_t len_old, len_new; + PetscBool is_same = PETSC_FALSE; + + PetscCall(PetscStrlen(ctx->internal_mat_type, &len_old)); + PetscCall(PetscStrlen(type, &len_new)); + if (len_old == len_new) PetscCall(PetscStrncmp(ctx->internal_mat_type, type, len_old, &is_same)); + if (is_same) PetscFunctionReturn(PETSC_SUCCESS); + } + // Clean up old mats in different format + // LCOV_EXCL_START + if (ctx->mat_assembled_full_internal) { + for (PetscInt i = 0; i < ctx->num_mats_assembled_full; i++) { + if (ctx->mats_assembled_full[i] == ctx->mat_assembled_full_internal) { + for (PetscInt j = i + 1; j < ctx->num_mats_assembled_full; j++) { + ctx->mats_assembled_full[j - 1] = ctx->mats_assembled_full[j]; + } + ctx->num_mats_assembled_full--; + // Note: we'll realloc this array again, so no need to shrink the allocation + PetscCall(MatDestroy(&ctx->mat_assembled_full_internal)); + } + } + } + if (ctx->mat_assembled_pbd_internal) { + for (PetscInt i = 0; i < ctx->num_mats_assembled_pbd; i++) { + if (ctx->mats_assembled_pbd[i] == ctx->mat_assembled_pbd_internal) { + for (PetscInt j = i + 1; j < ctx->num_mats_assembled_pbd; j++) { + ctx->mats_assembled_pbd[j - 1] = ctx->mats_assembled_pbd[j]; + } + // Note: we'll realloc this array again, so no need to shrink the allocation + ctx->num_mats_assembled_pbd--; + PetscCall(MatDestroy(&ctx->mat_assembled_pbd_internal)); + } + } + } + PetscCall(PetscFree(ctx->internal_mat_type)); + PetscCall(PetscStrallocpy(type, &ctx->internal_mat_type)); + PetscFunctionReturn(PETSC_SUCCESS); + // LCOV_EXCL_STOP +} + +/** + @brief Gets the inner matrix type as a string from the `MATCEED`. + + Collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[in] type Inner `MatType` + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedGetInnerMatType(Mat mat, MatType *type) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + *type = ctx->internal_mat_type; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Set a user defined matrix operation for a `MATCEED` matrix. + + Within each user-defined routine, the user should call `MatCeedGetContext()` to obtain the user-defined context that was set by +`MatCeedSetContext()`. + + Collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[in] op Name of the `MatOperation` + @param[in] g Function that provides the operation + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedSetOperation(Mat mat, MatOperation op, void (*g)(void)) { + PetscFunctionBeginUser; + PetscCall(MatShellSetOperation(mat, op, g)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Set input local vectors for `MATCEED` `MatMult()` and `MatMultTranspose()` operations. + + Not collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[in] X_loc Input PETSc local vector, or NULL + @param[in] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedSetLocalVectors(Mat mat, Vec X_loc, Vec Y_loc_transpose) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + if (X_loc) { + PetscInt len_old, len_new; + + PetscCall(VecGetSize(ctx->X_loc, &len_old)); + PetscCall(VecGetSize(X_loc, &len_new)); + PetscCheck(len_old == len_new, PETSC_COMM_SELF, PETSC_ERR_LIB, "new X_loc length %" PetscInt_FMT " should match old X_loc length %" PetscInt_FMT, + len_new, len_old); + PetscCall(VecDestroy(&ctx->X_loc)); + ctx->X_loc = X_loc; + PetscCall(PetscObjectReference((PetscObject)X_loc)); + } + if (Y_loc_transpose) { + PetscInt len_old, len_new; + + PetscCall(VecGetSize(ctx->Y_loc_transpose, &len_old)); + PetscCall(VecGetSize(Y_loc_transpose, &len_new)); + PetscCheck(len_old == len_new, PETSC_COMM_SELF, PETSC_ERR_LIB, + "new Y_loc_transpose length %" PetscInt_FMT " should match old Y_loc_transpose length %" PetscInt_FMT, len_new, len_old); + PetscCall(VecDestroy(&ctx->Y_loc_transpose)); + ctx->Y_loc_transpose = Y_loc_transpose; + PetscCall(PetscObjectReference((PetscObject)Y_loc_transpose)); + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Get input local vectors for `MATCEED` `MatMult()` and `MatMultTranspose()` operations. + + Not collective across MPI processes. + + @param[in,out] mat `MATCEED` + @param[out] X_loc Input PETSc local vector, or NULL + @param[out] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedGetLocalVectors(Mat mat, Vec *X_loc, Vec *Y_loc_transpose) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + if (X_loc) { + *X_loc = ctx->X_loc; + PetscCall(PetscObjectReference((PetscObject)*X_loc)); + } + if (Y_loc_transpose) { + *Y_loc_transpose = ctx->Y_loc_transpose; + PetscCall(PetscObjectReference((PetscObject)*Y_loc_transpose)); + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Restore input local vectors for `MATCEED` `MatMult()` and `MatMultTranspose()` operations. + + Not collective across MPI processes. + + @param[in,out] mat MatCeed + @param[out] X_loc Input PETSc local vector, or NULL + @param[out] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedRestoreLocalVectors(Mat mat, Vec *X_loc, Vec *Y_loc_transpose) { + PetscFunctionBeginUser; + if (X_loc) PetscCall(VecDestroy(X_loc)); + if (Y_loc_transpose) PetscCall(VecDestroy(Y_loc_transpose)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Get libCEED `CeedOperator` for `MATCEED` `MatMult()` and `MatMultTranspose()` operations. + + Not collective across MPI processes. + + @param[in,out] mat MatCeed + @param[out] op_mult libCEED `CeedOperator` for `MatMult()`, or NULL + @param[out] op_mult_transpose libCEED `CeedOperator` for `MatMultTranspose()`, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedGetCeedOperators(Mat mat, CeedOperator *op_mult, CeedOperator *op_mult_transpose) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + if (op_mult) { + *op_mult = NULL; + PetscCeedCall(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult, op_mult)); + } + if (op_mult_transpose) { + *op_mult_transpose = NULL; + PetscCeedCall(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult_transpose, op_mult_transpose)); + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Restore libCEED `CeedOperator` for `MATCEED` `MatMult()` and `MatMultTranspose()` operations. + + Not collective across MPI processes. + + @param[in,out] mat MatCeed + @param[out] op_mult libCEED `CeedOperator` for `MatMult()`, or NULL + @param[out] op_mult_transpose libCEED `CeedOperator` for `MatMultTranspose()`, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedRestoreCeedOperators(Mat mat, CeedOperator *op_mult, CeedOperator *op_mult_transpose) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + if (op_mult) PetscCeedCall(ctx->ceed, CeedOperatorDestroy(op_mult)); + if (op_mult_transpose) PetscCeedCall(ctx->ceed, CeedOperatorDestroy(op_mult_transpose)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Set `PetscLogEvent` for `MATCEED` `MatMult()` and `MatMultTranspose()` operators. + + Not collective across MPI processes. + + @param[in,out] mat MatCeed + @param[out] log_event_mult `PetscLogEvent` for forward evaluation, or NULL + @param[out] log_event_mult_transpose `PetscLogEvent` for transpose evaluation, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedSetLogEvents(Mat mat, PetscLogEvent log_event_mult, PetscLogEvent log_event_mult_transpose) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + if (log_event_mult) ctx->log_event_mult = log_event_mult; + if (log_event_mult_transpose) ctx->log_event_mult_transpose = log_event_mult_transpose; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Get `PetscLogEvent` for `MATCEED` `MatMult()` and `MatMultTranspose()` operators. + + Not collective across MPI processes. + + @param[in,out] mat MatCeed + @param[out] log_event_mult `PetscLogEvent` for forward evaluation, or NULL + @param[out] log_event_mult_transpose `PetscLogEvent` for transpose evaluation, or NULL + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedGetLogEvents(Mat mat, PetscLogEvent *log_event_mult, PetscLogEvent *log_event_mult_transpose) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(mat, &ctx)); + if (log_event_mult) *log_event_mult = ctx->log_event_mult; + if (log_event_mult_transpose) *log_event_mult_transpose = ctx->log_event_mult_transpose; + PetscFunctionReturn(PETSC_SUCCESS); +} + +// ----------------------------------------------------------------------------- +// Operator context data +// ----------------------------------------------------------------------------- + +/** + @brief Setup context data for operator application. + + Collective across MPI processes. + + @param[in] dm_x Input `DM` + @param[in] dm_y Output `DM` + @param[in] X_loc Input PETSc local vector, or NULL + @param[in] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL + @param[in] op_mult `CeedOperator` for forward evaluation + @param[in] op_mult_transpose `CeedOperator` for transpose evaluation + @param[in] log_event_mult `PetscLogEvent` for forward evaluation + @param[in] log_event_mult_transpose `PetscLogEvent` for transpose evaluation + @param[out] ctx Context data for operator evaluation + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedContextCreate(DM dm_x, DM dm_y, Vec X_loc, Vec Y_loc_transpose, CeedOperator op_mult, CeedOperator op_mult_transpose, + PetscLogEvent log_event_mult, PetscLogEvent log_event_mult_transpose, MatCeedContext *ctx) { + CeedSize x_loc_len, y_loc_len; + + PetscFunctionBeginUser; + + // Allocate + PetscCall(PetscNew(ctx)); + (*ctx)->ref_count = 1; + + // Logging + (*ctx)->log_event_mult = log_event_mult; + (*ctx)->log_event_mult_transpose = log_event_mult_transpose; + + // PETSc objects + PetscCall(PetscObjectReference((PetscObject)dm_x)); + (*ctx)->dm_x = dm_x; + PetscCall(PetscObjectReference((PetscObject)dm_y)); + (*ctx)->dm_y = dm_y; + if (X_loc) PetscCall(PetscObjectReference((PetscObject)X_loc)); + (*ctx)->X_loc = X_loc; + if (Y_loc_transpose) PetscCall(PetscObjectReference((PetscObject)Y_loc_transpose)); + (*ctx)->Y_loc_transpose = Y_loc_transpose; + + // Memtype + { + const PetscScalar *x; + Vec X; + + PetscCall(DMGetLocalVector(dm_x, &X)); + PetscCall(VecGetArrayReadAndMemType(X, &x, &(*ctx)->mem_type)); + PetscCall(VecRestoreArrayReadAndMemType(X, &x)); + PetscCall(DMRestoreLocalVector(dm_x, &X)); + } + + // libCEED objects + PetscCheck(CeedOperatorGetCeed(op_mult, &(*ctx)->ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB, + "retrieving Ceed context object failed"); + PetscCeedCall((*ctx)->ceed, CeedReference((*ctx)->ceed)); + PetscCeedCall((*ctx)->ceed, CeedOperatorGetActiveVectorLengths(op_mult, &x_loc_len, &y_loc_len)); + PetscCeedCall((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult, &(*ctx)->op_mult)); + if (op_mult_transpose) PetscCeedCall((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult_transpose, &(*ctx)->op_mult_transpose)); + PetscCeedCall((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, x_loc_len, &(*ctx)->x_loc)); + PetscCeedCall((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, y_loc_len, &(*ctx)->y_loc)); + + // Flop counting + { + CeedSize ceed_flops_estimate = 0; + + PetscCeedCall((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult, &ceed_flops_estimate)); + (*ctx)->flops_mult = ceed_flops_estimate; + if (op_mult_transpose) { + PetscCeedCall((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult_transpose, &ceed_flops_estimate)); + (*ctx)->flops_mult_transpose = ceed_flops_estimate; + } + } + + // Check sizes + if (x_loc_len > 0 || y_loc_len > 0) { + CeedSize ctx_x_loc_len, ctx_y_loc_len; + PetscInt X_loc_len, dm_x_loc_len, Y_loc_len, dm_y_loc_len; + Vec dm_X_loc, dm_Y_loc; + + // -- Input + PetscCall(DMGetLocalVector(dm_x, &dm_X_loc)); + PetscCall(VecGetLocalSize(dm_X_loc, &dm_x_loc_len)); + PetscCall(DMRestoreLocalVector(dm_x, &dm_X_loc)); + if (X_loc) PetscCall(VecGetLocalSize(X_loc, &X_loc_len)); + PetscCeedCall((*ctx)->ceed, CeedVectorGetLength((*ctx)->x_loc, &ctx_x_loc_len)); + if (X_loc) PetscCheck(X_loc_len == dm_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "X_loc must match dm_x dimensions"); + PetscCheck(x_loc_len == dm_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "op must match dm_x dimensions"); + PetscCheck(x_loc_len == ctx_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "x_loc must match op dimensions"); + + // -- Output + PetscCall(DMGetLocalVector(dm_y, &dm_Y_loc)); + PetscCall(VecGetLocalSize(dm_Y_loc, &dm_y_loc_len)); + PetscCall(DMRestoreLocalVector(dm_y, &dm_Y_loc)); + PetscCeedCall((*ctx)->ceed, CeedVectorGetLength((*ctx)->y_loc, &ctx_y_loc_len)); + PetscCheck(ctx_y_loc_len == dm_y_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "op must match dm_y dimensions"); + + // -- Transpose + if (Y_loc_transpose) { + PetscCall(VecGetLocalSize(Y_loc_transpose, &Y_loc_len)); + PetscCheck(Y_loc_len == dm_y_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "Y_loc_transpose must match dm_y dimensions"); + } + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Increment reference counter for `MATCEED` context. + + Not collective across MPI processes. + + @param[in,out] ctx Context data + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedContextReference(MatCeedContext ctx) { + PetscFunctionBeginUser; + ctx->ref_count++; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Copy reference for `MATCEED`. + Note: If `ctx_copy` is non-null, it is assumed to be a valid pointer to a `MatCeedContext`. + + Not collective across MPI processes. + + @param[in] ctx Context data + @param[out] ctx_copy Copy of pointer to context data + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedContextReferenceCopy(MatCeedContext ctx, MatCeedContext *ctx_copy) { + PetscFunctionBeginUser; + PetscCall(MatCeedContextReference(ctx)); + PetscCall(MatCeedContextDestroy(*ctx_copy)); + *ctx_copy = ctx; + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Destroy context data for operator application. + + Collective across MPI processes. + + @param[in,out] ctx Context data for operator evaluation + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatCeedContextDestroy(MatCeedContext ctx) { + PetscFunctionBeginUser; + if (!ctx || --ctx->ref_count > 0) PetscFunctionReturn(PETSC_SUCCESS); + + // PETSc objects + PetscCall(DMDestroy(&ctx->dm_x)); + PetscCall(DMDestroy(&ctx->dm_y)); + PetscCall(VecDestroy(&ctx->X_loc)); + PetscCall(VecDestroy(&ctx->Y_loc_transpose)); + PetscCall(MatDestroy(&ctx->mat_assembled_full_internal)); + PetscCall(MatDestroy(&ctx->mat_assembled_pbd_internal)); + PetscCall(PetscFree(ctx->internal_mat_type)); + PetscCall(PetscFree(ctx->mats_assembled_full)); + PetscCall(PetscFree(ctx->mats_assembled_pbd)); + + // libCEED objects + PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->x_loc)); + PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->y_loc)); + PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->coo_values_full)); + PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->coo_values_pbd)); + PetscCeedCall(ctx->ceed, CeedOperatorDestroy(&ctx->op_mult)); + PetscCeedCall(ctx->ceed, CeedOperatorDestroy(&ctx->op_mult_transpose)); + PetscCheck(CeedDestroy(&ctx->ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB, "destroying libCEED context object failed"); + + // Deallocate + ctx->is_destroyed = PETSC_TRUE; // Flag as destroyed in case someone has stale ref + PetscCall(PetscFree(ctx)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Compute the diagonal of an operator via libCEED. + + Collective across MPI processes. + + @param[in] A `MATCEED` + @param[out] D Vector holding operator diagonal + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatGetDiagonal_Ceed(Mat A, Vec D) { + PetscMemType mem_type; + Vec D_loc; + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(A, &ctx)); + + // Place PETSc vector in libCEED vector + PetscCall(DMGetLocalVector(ctx->dm_x, &D_loc)); + PetscCall(VecP2C(ctx->ceed, D_loc, &mem_type, ctx->x_loc)); + + // Compute Diagonal + PetscCeedCall(ctx->ceed, CeedOperatorLinearAssembleDiagonal(ctx->op_mult, ctx->x_loc, CEED_REQUEST_IMMEDIATE)); + + // Restore PETSc vector + PetscCall(VecC2P(ctx->ceed, ctx->x_loc, mem_type, D_loc)); + + // Local-to-Global + PetscCall(VecZeroEntries(D)); + PetscCall(DMLocalToGlobal(ctx->dm_x, D_loc, ADD_VALUES, D)); + PetscCall(DMRestoreLocalVector(ctx->dm_x, &D_loc)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Compute `A X = Y` for a `MATCEED`. + + Collective across MPI processes. + + @param[in] A `MATCEED` + @param[in] X Input PETSc vector + @param[out] Y Output PETSc vector + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(A, &ctx)); + PetscCall(PetscLogEventBegin(ctx->log_event_mult, A, X, Y, 0)); + + { + PetscMemType x_mem_type, y_mem_type; + Vec X_loc = ctx->X_loc, Y_loc; + + // Get local vectors + if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); + PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); + + // Global-to-local + PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc)); + + // Setup libCEED vectors + PetscCall(VecReadP2C(ctx->ceed, X_loc, &x_mem_type, ctx->x_loc)); + PetscCall(VecZeroEntries(Y_loc)); + PetscCall(VecP2C(ctx->ceed, Y_loc, &y_mem_type, ctx->y_loc)); + + // Apply libCEED operator + PetscCall(PetscLogGpuTimeBegin()); + PetscCeedCall(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult, ctx->x_loc, ctx->y_loc, CEED_REQUEST_IMMEDIATE)); + PetscCall(PetscLogGpuTimeEnd()); + + // Restore PETSc vectors + PetscCall(VecReadC2P(ctx->ceed, ctx->x_loc, x_mem_type, X_loc)); + PetscCall(VecC2P(ctx->ceed, ctx->y_loc, y_mem_type, Y_loc)); + + // Local-to-global + PetscCall(VecZeroEntries(Y)); + PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y)); + + // Restore local vectors, as needed + if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); + PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); + } + + // Log flops + if (PetscMemTypeDevice(ctx->mem_type)) PetscCall(PetscLogGpuFlops(ctx->flops_mult)); + else PetscCall(PetscLogFlops(ctx->flops_mult)); + + PetscCall(PetscLogEventEnd(ctx->log_event_mult, A, X, Y, 0)); + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + @brief Compute `A^T Y = X` for a `MATCEED`. + + Collective across MPI processes. + + @param[in] A `MATCEED` + @param[in] Y Input PETSc vector + @param[out] X Output PETSc vector + + @return An error code: 0 - success, otherwise - failure +**/ +PetscErrorCode MatMultTranspose_Ceed(Mat A, Vec Y, Vec X) { + MatCeedContext ctx; + + PetscFunctionBeginUser; + PetscCall(MatShellGetContext(A, &ctx)); + PetscCall(PetscLogEventBegin(ctx->log_event_mult_transpose, A, Y, X, 0)); + + { + PetscMemType x_mem_type, y_mem_type; + Vec X_loc, Y_loc = ctx->Y_loc_transpose; + + // Get local vectors + if (!ctx->Y_loc_transpose) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); + PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); + + // Global-to-local + PetscCall(DMGlobalToLocal(ctx->dm_y, Y, INSERT_VALUES, Y_loc)); + + // Setup libCEED vectors + PetscCall(VecReadP2C(ctx->ceed, Y_loc, &y_mem_type, ctx->y_loc)); + PetscCall(VecZeroEntries(X_loc)); + PetscCall(VecP2C(ctx->ceed, X_loc, &x_mem_type, ctx->x_loc)); + + // Apply libCEED operator + PetscCall(PetscLogGpuTimeBegin()); + PetscCeedCall(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult_transpose, ctx->y_loc, ctx->x_loc, CEED_REQUEST_IMMEDIATE)); + PetscCall(PetscLogGpuTimeEnd()); + + // Restore PETSc vectors + PetscCall(VecReadC2P(ctx->ceed, ctx->y_loc, y_mem_type, Y_loc)); + PetscCall(VecC2P(ctx->ceed, ctx->x_loc, x_mem_type, X_loc)); + + // Local-to-global + PetscCall(VecZeroEntries(X)); + PetscCall(DMLocalToGlobal(ctx->dm_x, X_loc, ADD_VALUES, X)); + + // Restore local vectors, as needed + if (!ctx->Y_loc_transpose) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); + PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); + } + + // Log flops + if (PetscMemTypeDevice(ctx->mem_type)) PetscCall(PetscLogGpuFlops(ctx->flops_mult_transpose)); + else PetscCall(PetscLogFlops(ctx->flops_mult_transpose)); + + PetscCall(PetscLogEventEnd(ctx->log_event_mult_transpose, A, Y, X, 0)); + PetscFunctionReturn(PETSC_SUCCESS); +} From 91c97f41bd25cf2fbada21392acdb50722cb486e Mon Sep 17 00:00:00 2001 From: James Wright Date: Tue, 12 Mar 2024 09:47:34 -0600 Subject: [PATCH 3/7] fluids: Use MatCeed for Jacobian --- examples/fluids/navierstokes.c | 6 +- examples/fluids/navierstokes.h | 10 +-- examples/fluids/src/cloptions.c | 2 - examples/fluids/src/setupdm.c | 2 + examples/fluids/src/setuplibceed.c | 19 ++++- examples/fluids/src/setupts.c | 120 +++++++++-------------------- 6 files changed, 62 insertions(+), 97 deletions(-) diff --git a/examples/fluids/navierstokes.c b/examples/fluids/navierstokes.c index a4ffd7c5e6..c1f440d741 100644 --- a/examples/fluids/navierstokes.c +++ b/examples/fluids/navierstokes.c @@ -22,7 +22,7 @@ //TESTARGS(name="Blasius, SGS DataDriven Sequential") -ceed {ceed_resource} -options_file examples/fluids/tests-output/blasius_stgtest.yaml -sgs_model_type data_driven -sgs_model_dd_leakyrelu_alpha 0.3 -sgs_model_dd_parameter_dir examples/fluids/dd_sgs_data -ts_dt 2e-9 -state_var primitive -ksp_rtol 1e-12 -snes_rtol 1e-12 -stg_mean_only -stg_fluctuating_IC -test_type solver -compare_final_state_atol 1e-10 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-blasius-sgs-data-driven.bin -sgs_model_dd_use_fused false //TESTARGS(name="Advection, skew") -ceed {ceed_resource} -test_type solver -options_file examples/fluids/advection.yaml -ts_max_steps 0 -wind_type translation -wind_translation -0.5547002,0.83205029,0 -advection_ic_type skew -dm_plex_box_faces 2,1,1 -compare_final_state_atol 1e-10 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-adv-skew.bin //TESTARGS(name="Advection, rotation, cosine") -ceed {ceed_resource} -test_type solver -options_file examples/fluids/advection.yaml -ts_max_steps 0 -advection_ic_type cosine_hill -dm_plex_box_faces 2,1,1 -compare_final_state_atol 1e-10 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-adv-rotation-cosine.bin -//TESTARGS(name="Gaussian Wave, using MatShell") -ceed {ceed_resource} -test_type solver -options_file examples/fluids/gaussianwave.yaml -compare_final_state_atol 1e-8 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-gaussianwave-shell.bin -dm_plex_box_faces 2,2,1 -ts_max_steps 5 -degree 3 -amat_type shell -pmat_pbdiagonal -pc_type vpbjacobi +//TESTARGS(name="Gaussian Wave, using MatShell") -ceed {ceed_resource} -test_type solver -options_file examples/fluids/gaussianwave.yaml -compare_final_state_atol 1e-8 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-gaussianwave-shell.bin -dm_plex_box_faces 2,2,1 -ts_max_steps 5 -degree 3 -amat_type shell -pc_type vpbjacobi //TESTARGS(name="Taylor-Green Vortex IC") -ceed {ceed_resource} -problem taylor_green -test_type solver -dm_plex_dim 3 -dm_plex_box_faces 6,6,6 -ts_max_steps 0 -compare_final_state_atol 1e-12 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-taylor-green-IC.bin //TESTARGS(name="Blasius, SGS DataDriven Fused") -ceed {ceed_resource} -options_file examples/fluids/tests-output/blasius_stgtest.yaml -sgs_model_type data_driven -sgs_model_dd_leakyrelu_alpha 0.3 -sgs_model_dd_parameter_dir examples/fluids/dd_sgs_data -ts_dt 2e-9 -state_var primitive -ksp_rtol 1e-12 -snes_rtol 1e-12 -stg_mean_only -stg_fluctuating_IC -test_type solver -compare_final_state_atol 1e-10 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-blasius-sgs-data-driven.bin //TESTARGS(name="Blasius, Anisotropic Differential Filter") -ceed {ceed_resource} -test_type diff_filter -options_file examples/fluids/tests-output/blasius_test.yaml -compare_final_state_atol 5e-10 -compare_final_state_filename examples/fluids/tests-output/fluids-navierstokes-blasius_diff_filter_aniso_vandriest.bin -diff_filter_monitor -ts_max_steps 0 -state_var primitive -diff_filter_friction_length 1e-5 -diff_filter_wall_damping_function van_driest -diff_filter_ksp_rtol 1e-8 -diff_filter_grid_based_width -diff_filter_width_scaling 1,0.7,1 @@ -264,8 +264,6 @@ int main(int argc, char **argv) { PetscCallCeed(ceed, CeedVectorDestroy(&user->q_ceed)); PetscCallCeed(ceed, CeedVectorDestroy(&user->q_dot_ceed)); PetscCallCeed(ceed, CeedVectorDestroy(&user->g_ceed)); - PetscCallCeed(ceed, CeedVectorDestroy(&user->coo_values_amat)); - PetscCallCeed(ceed, CeedVectorDestroy(&user->coo_values_pmat)); // -- Bases PetscCallCeed(ceed, CeedBasisDestroy(&ceed_data->basis_q)); @@ -315,7 +313,6 @@ int main(int argc, char **argv) { PetscCall(OperatorApplyContextDestroy(user->op_rhs_ctx)); PetscCall(OperatorApplyContextDestroy(user->op_strong_bc_ctx)); PetscCallCeed(ceed, CeedOperatorDestroy(&user->op_ifunction)); - PetscCallCeed(ceed, CeedOperatorDestroy(&user->op_ijacobian)); // -- Ceed PetscCheck(CeedDestroy(&ceed) == CEED_ERROR_SUCCESS, comm, PETSC_ERR_LIB, "Destroying Ceed object failed"); @@ -339,6 +336,7 @@ int main(int argc, char **argv) { // -- Matrices PetscCall(MatDestroy(&user->interp_viz)); + PetscCall(MatDestroy(&user->mat_ijacobian)); // -- DM PetscCall(DMDestroy(&dm)); diff --git a/examples/fluids/navierstokes.h b/examples/fluids/navierstokes.h index b344992175..003ff4bcfd 100644 --- a/examples/fluids/navierstokes.h +++ b/examples/fluids/navierstokes.h @@ -9,6 +9,7 @@ #define libceed_fluids_examples_navier_stokes_h #include +#include #include #include @@ -123,8 +124,7 @@ struct AppCtx_private { PetscInt degree; PetscInt q_extra; // Solver arguments - MatType amat_type; - PetscBool pmat_pbdiagonal; + MatType amat_type; // Post-processing arguments PetscInt checkpoint_interval; PetscInt viz_refine; @@ -245,11 +245,11 @@ struct User_private { Vec Q_loc, Q_dot_loc; Physics phys; AppCtx app_ctx; - CeedVector q_ceed, q_dot_ceed, g_ceed, coo_values_amat, coo_values_pmat, x_ceed; - CeedOperator op_rhs_vol, op_ifunction_vol, op_ifunction, op_ijacobian; + CeedVector q_ceed, q_dot_ceed, g_ceed, x_ceed; + CeedOperator op_rhs_vol, op_ifunction_vol, op_ifunction; + Mat mat_ijacobian; KSP mass_ksp; OperatorApplyContext op_rhs_ctx, op_strong_bc_ctx; - bool matrices_set_up; CeedScalar time_bc_set; SpanStatsData spanstats; NodalProjectionData grad_velo_proj; diff --git a/examples/fluids/src/cloptions.c b/examples/fluids/src/cloptions.c index cbb5e7841d..571d8dbe30 100644 --- a/examples/fluids/src/cloptions.c +++ b/examples/fluids/src/cloptions.c @@ -102,8 +102,6 @@ PetscErrorCode ProcessCommandLineOptions(MPI_Comm comm, AppCtx app_ctx, SimpleBC sizeof(amat_type), &option_set)); if (option_set) PetscCall(PetscStrallocpy(amat_type, (char **)&app_ctx->amat_type)); } - PetscCall(PetscOptionsBool("-pmat_pbdiagonal", "Assemble only point-block diagonal for Pmat", NULL, app_ctx->pmat_pbdiagonal, - &app_ctx->pmat_pbdiagonal, NULL)); // Provide default ceed resource if not specified if (!ceed_flag) { diff --git a/examples/fluids/src/setupdm.c b/examples/fluids/src/setupdm.c index d59ff1619b..b88e7f9715 100644 --- a/examples/fluids/src/setupdm.c +++ b/examples/fluids/src/setupdm.c @@ -32,6 +32,8 @@ PetscErrorCode CreateDM(MPI_Comm comm, ProblemData *problem, MatType mat_type, V // Set Tensor elements PetscCall(PetscOptionsSetValue(NULL, "-dm_plex_simplex", "0")); PetscCall(PetscOptionsSetValue(NULL, "-dm_sparse_localize", "0")); + PetscCall(PetscOptionsSetValue(NULL, "-dm_blocking_type", "field_node")); + // Set CL options PetscCall(DMSetFromOptions(*dm)); PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view")); diff --git a/examples/fluids/src/setuplibceed.c b/examples/fluids/src/setuplibceed.c index bee3cf87f0..bbde161ef7 100644 --- a/examples/fluids/src/setuplibceed.c +++ b/examples/fluids/src/setuplibceed.c @@ -139,6 +139,7 @@ PetscErrorCode SetupBCQFunctions(Ceed ceed, PetscInt dim_sur, PetscInt num_comp_ if (apply_bc.qfunction) { PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, apply_bc.qfunction, apply_bc.qfunction_loc, qf_apply_bc)); PetscCallCeed(ceed, CeedQFunctionSetContext(*qf_apply_bc, apply_bc.qfunction_context)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(*qf_apply_bc, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf_apply_bc, "q", num_comp_q, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf_apply_bc, "Grad_q", num_comp_q * dim_sur, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf_apply_bc, "surface qdata", q_data_size_sur, CEED_EVAL_NONE)); @@ -149,6 +150,7 @@ PetscErrorCode SetupBCQFunctions(Ceed ceed, PetscInt dim_sur, PetscInt num_comp_ if (apply_bc_jacobian.qfunction) { PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, apply_bc_jacobian.qfunction, apply_bc_jacobian.qfunction_loc, qf_apply_bc_jacobian)); PetscCallCeed(ceed, CeedQFunctionSetContext(*qf_apply_bc_jacobian, apply_bc_jacobian.qfunction_context)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(*qf_apply_bc_jacobian, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf_apply_bc_jacobian, "dq", num_comp_q, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf_apply_bc_jacobian, "Grad_dq", num_comp_q * dim_sur, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf_apply_bc_jacobian, "surface qdata", q_data_size_sur, CEED_EVAL_NONE)); @@ -213,6 +215,7 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App if (problem->setup_vol.qfunction_context) { PetscCallCeed(ceed, CeedQFunctionSetContext(ceed_data->qf_setup_vol, problem->setup_vol.qfunction_context)); } + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(ceed_data->qf_setup_vol, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_setup_vol, "dx", num_comp_x * dim, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_setup_vol, "weight", 1, CEED_EVAL_WEIGHT)); PetscCallCeed(ceed, CeedQFunctionAddOutput(ceed_data->qf_setup_vol, "qdata", q_data_size_vol, CEED_EVAL_NONE)); @@ -220,6 +223,7 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App // -- Create QFunction for ICs PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, problem->ics.qfunction, problem->ics.qfunction_loc, &ceed_data->qf_ics)); PetscCallCeed(ceed, CeedQFunctionSetContext(ceed_data->qf_ics, problem->ics.qfunction_context)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(ceed_data->qf_ics, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_ics, "x", num_comp_x, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_ics, "dx", num_comp_x * dim, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddOutput(ceed_data->qf_ics, "q0", num_comp_q, CEED_EVAL_NONE)); @@ -229,6 +233,7 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App PetscCallCeed( ceed, CeedQFunctionCreateInterior(ceed, 1, problem->apply_vol_rhs.qfunction, problem->apply_vol_rhs.qfunction_loc, &ceed_data->qf_rhs_vol)); PetscCallCeed(ceed, CeedQFunctionSetContext(ceed_data->qf_rhs_vol, problem->apply_vol_rhs.qfunction_context)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(ceed_data->qf_rhs_vol, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_rhs_vol, "q", num_comp_q, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_rhs_vol, "Grad_q", num_comp_q * dim, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_rhs_vol, "qdata", q_data_size_vol, CEED_EVAL_NONE)); @@ -242,6 +247,7 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, problem->apply_vol_ifunction.qfunction, problem->apply_vol_ifunction.qfunction_loc, &ceed_data->qf_ifunction_vol)); PetscCallCeed(ceed, CeedQFunctionSetContext(ceed_data->qf_ifunction_vol, problem->apply_vol_ifunction.qfunction_context)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(ceed_data->qf_ifunction_vol, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_ifunction_vol, "q", num_comp_q, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_ifunction_vol, "Grad_q", num_comp_q * dim, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_ifunction_vol, "q dot", num_comp_q, CEED_EVAL_INTERP)); @@ -257,6 +263,7 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, problem->apply_vol_ijacobian.qfunction, problem->apply_vol_ijacobian.qfunction_loc, &qf_ijacobian_vol)); PetscCallCeed(ceed, CeedQFunctionSetContext(qf_ijacobian_vol, problem->apply_vol_ijacobian.qfunction_context)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(qf_ijacobian_vol, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(qf_ijacobian_vol, "dq", num_comp_q, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(qf_ijacobian_vol, "Grad_dq", num_comp_q * dim, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(qf_ijacobian_vol, "qdata", q_data_size_vol, CEED_EVAL_NONE)); @@ -380,6 +387,7 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App if (problem->setup_sur.qfunction_context) { PetscCallCeed(ceed, CeedQFunctionSetContext(ceed_data->qf_setup_sur, problem->setup_sur.qfunction_context)); } + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(ceed_data->qf_setup_sur, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_setup_sur, "dx", num_comp_x * dim_sur, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(ceed_data->qf_setup_sur, "weight", 1, CEED_EVAL_WEIGHT)); PetscCallCeed(ceed, CeedQFunctionAddOutput(ceed_data->qf_setup_sur, "surface qdata", q_data_size_sur, CEED_EVAL_NONE)); @@ -407,10 +415,15 @@ PetscErrorCode SetupLibceed(Ceed ceed, CeedData ceed_data, DM dm, User user, App PetscCall(OperatorApplyContextCreate(dm, dm, ceed, op_rhs, user->q_ceed, user->g_ceed, user->Q_loc, NULL, &user->op_rhs_ctx)); PetscCallCeed(ceed, CeedOperatorDestroy(&op_rhs)); } else { // IFunction + CeedOperator op_ijacobian = NULL; + PetscCall(CreateOperatorForDomain(ceed, dm, bc, ceed_data, user->phys, user->op_ifunction_vol, op_ijacobian_vol, height, P_sur, Q_sur, - q_data_size_sur, jac_data_size_sur, &user->op_ifunction, op_ijacobian_vol ? &user->op_ijacobian : NULL)); - if (user->op_ijacobian) { - PetscCallCeed(ceed, CeedOperatorGetContextFieldLabel(user->op_ijacobian, "ijacobian time shift", &user->phys->ijacobian_time_shift_label)); + q_data_size_sur, jac_data_size_sur, &user->op_ifunction, op_ijacobian_vol ? &op_ijacobian : NULL)); + if (op_ijacobian) { + PetscCall(MatCeedCreate(user->dm, user->dm, op_ijacobian, NULL, &user->mat_ijacobian)); + PetscCall(MatCeedSetLocalVectors(user->mat_ijacobian, user->Q_dot_loc, NULL)); + PetscCallCeed(ceed, CeedOperatorGetContextFieldLabel(op_ijacobian, "ijacobian time shift", &user->phys->ijacobian_time_shift_label)); + PetscCallCeed(ceed, CeedOperatorDestroy(&op_ijacobian)); } if (problem->use_strong_bc_ceed) PetscCall(SetupStrongBC_Ceed(ceed, ceed_data, dm, user, problem, bc)); if (app_ctx->sgs_model_type == SGS_MODEL_DATA_DRIVEN) PetscCall(SgsDDSetup(ceed, user, ceed_data, problem)); diff --git a/examples/fluids/src/setupts.c b/examples/fluids/src/setupts.c index d49978ac27..24815c00aa 100644 --- a/examples/fluids/src/setupts.c +++ b/examples/fluids/src/setupts.c @@ -198,88 +198,33 @@ PetscErrorCode IFunction_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, Vec G, void *u PetscFunctionReturn(PETSC_SUCCESS); } -static PetscErrorCode FormPreallocation(User user, PetscBool pbdiagonal, Mat J, CeedVector *coo_values) { - PetscCount ncoo; - PetscInt *rows_petsc, *cols_petsc; - CeedInt *rows_ceed, *cols_ceed; - - PetscFunctionBeginUser; - if (pbdiagonal) { - PetscCallCeed(user->ceed, CeedOperatorLinearAssemblePointBlockDiagonalSymbolic(user->op_ijacobian, &ncoo, &rows_ceed, &cols_ceed)); - } else { - PetscCallCeed(user->ceed, CeedOperatorLinearAssembleSymbolic(user->op_ijacobian, &ncoo, &rows_ceed, &cols_ceed)); - } - PetscCall(IntArrayC2P(ncoo, &rows_ceed, &rows_petsc)); - PetscCall(IntArrayC2P(ncoo, &cols_ceed, &cols_petsc)); - PetscCall(MatSetPreallocationCOOLocal(J, ncoo, rows_petsc, cols_petsc)); - free(rows_petsc); - free(cols_petsc); - PetscCallCeed(user->ceed, CeedVectorCreate(user->ceed, ncoo, coo_values)); - PetscFunctionReturn(PETSC_SUCCESS); -} - -static PetscErrorCode FormSetValues(User user, PetscBool pbdiagonal, Mat J, CeedVector coo_values) { - CeedMemType mem_type = CEED_MEM_HOST; - const PetscScalar *values; - MatType mat_type; - - PetscFunctionBeginUser; - PetscCall(MatGetType(J, &mat_type)); - if (strstr(mat_type, "kokkos") || strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE; - if (pbdiagonal) { - PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssemblePointBlockDiagonal, J, 0, 0, 0)); - PetscCall(PetscLogGpuTimeBegin()); - PetscCallCeed(user->ceed, CeedOperatorLinearAssemblePointBlockDiagonal(user->op_ijacobian, coo_values, CEED_REQUEST_IMMEDIATE)); - PetscCall(PetscLogGpuTimeEnd()); - PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssemblePointBlockDiagonal, J, 0, 0, 0)); - } else { - PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssemble, J, 0, 0, 0)); - PetscCall(PetscLogGpuTimeBegin()); - PetscCallCeed(user->ceed, CeedOperatorLinearAssemble(user->op_ijacobian, coo_values)); - PetscCall(PetscLogGpuTimeEnd()); - PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssemble, J, 0, 0, 0)); - } - PetscCallCeed(user->ceed, CeedVectorGetArrayRead(coo_values, mem_type, &values)); - PetscCall(MatSetValuesCOO(J, values, INSERT_VALUES)); - PetscCallCeed(user->ceed, CeedVectorRestoreArrayRead(coo_values, &values)); - PetscFunctionReturn(PETSC_SUCCESS); -} - PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, PetscReal shift, Mat J, Mat J_pre, void *user_data) { User user = *(User *)user_data; Ceed ceed = user->ceed; - PetscBool J_is_shell, J_is_mffd, J_pre_is_shell; + PetscBool J_is_matceed, J_is_mffd, J_pre_is_matceed, J_pre_is_mffd; PetscFunctionBeginUser; - if (user->phys->ijacobian_time_shift_label) - PetscCallCeed(ceed, CeedOperatorSetContextDouble(user->op_ijacobian, user->phys->ijacobian_time_shift_label, &shift)); PetscCall(PetscObjectTypeCompare((PetscObject)J, MATMFFD, &J_is_mffd)); - PetscCall(PetscObjectTypeCompare((PetscObject)J, MATSHELL, &J_is_shell)); - PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATSHELL, &J_pre_is_shell)); - if (!user->matrices_set_up) { - if (J_is_shell) { - OperatorApplyContext op_ijacobian_ctx; - OperatorApplyContextCreate(user->dm, user->dm, user->ceed, user->op_ijacobian, user->q_ceed, user->g_ceed, user->Q_dot_loc, NULL, - &op_ijacobian_ctx); - PetscCall(CreateMatShell_Ceed(op_ijacobian_ctx, &J)); - PetscCall(MatSetUp(J)); - } - if (!J_pre_is_shell) { - PetscCall(FormPreallocation(user, user->app_ctx->pmat_pbdiagonal, J_pre, &user->coo_values_pmat)); - } - if (J != J_pre && !J_is_shell && !J_is_mffd) { - PetscCall(FormPreallocation(user, PETSC_FALSE, J, &user->coo_values_amat)); - } - user->matrices_set_up = true; - } - if (!J_pre_is_shell) { - PetscCall(FormSetValues(user, user->app_ctx->pmat_pbdiagonal, J_pre, user->coo_values_pmat)); + PetscCall(PetscObjectTypeCompare((PetscObject)J, MATCEED, &J_is_matceed)); + PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATMFFD, &J_pre_is_mffd)); + PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATCEED, &J_pre_is_matceed)); + if (user->phys->ijacobian_time_shift_label) { + CeedOperator op_ijacobian; + + PetscCall(MatCeedGetCeedOperators(user->mat_ijacobian, &op_ijacobian, NULL)); + PetscCallCeed(ceed, CeedOperatorSetContextDouble(op_ijacobian, user->phys->ijacobian_time_shift_label, &shift)); } - if (user->coo_values_amat) { - PetscCall(FormSetValues(user, PETSC_FALSE, J, user->coo_values_amat)); - } else if (J_is_mffd) { + + if (J_is_matceed || J_is_mffd) { PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY)); PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY)); + } else PetscCall(MatCeedAssembleCOO(user->mat_ijacobian, J)); + + if (J_pre_is_matceed && J != J_pre) { + PetscCall(MatAssemblyBegin(J_pre, MAT_FINAL_ASSEMBLY)); + PetscCall(MatAssemblyEnd(J_pre, MAT_FINAL_ASSEMBLY)); + } else if (!J_pre_is_matceed && !J_pre_is_mffd && J != J_pre) { + PetscCall(MatCeedAssembleCOO(user->mat_ijacobian, J_pre)); } PetscFunctionReturn(PETSC_SUCCESS); } @@ -419,17 +364,8 @@ PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys, Vec *Q } else { // Implicit integrators can fall back to using an RHSFunction PetscCall(TSSetRHSFunction(*ts, NULL, RHS_NS, &user)); } - if (user->op_ijacobian) { + if (user->mat_ijacobian) { PetscCall(DMTSSetIJacobian(dm, FormIJacobian_NS, &user)); - if (app_ctx->amat_type) { - Mat Pmat, Amat; - PetscCall(DMCreateMatrix(dm, &Pmat)); - PetscCall(DMSetMatType(dm, app_ctx->amat_type)); - PetscCall(DMCreateMatrix(dm, &Amat)); - PetscCall(TSSetIJacobian(*ts, Amat, Pmat, NULL, NULL)); - PetscCall(MatDestroy(&Amat)); - PetscCall(MatDestroy(&Pmat)); - } } } else { PetscCheck(user->op_rhs_ctx, comm, PETSC_ERR_ARG_NULL, "Problem does not provide RHSFunction"); @@ -444,6 +380,24 @@ PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys, Vec *Q PetscCall(TSGetAdapt(*ts, &adapt)); PetscCall(TSAdaptSetStepLimits(adapt, 1.e-12 * user->units->second, 1.e2 * user->units->second)); PetscCall(TSSetFromOptions(*ts)); + if (user->mat_ijacobian) { + if (app_ctx->amat_type && !strcmp(app_ctx->amat_type, MATSHELL)) { + PetscBool use_matceed_pmat; + SNES snes; + KSP ksp; + PC pc; + PCType pc_type; + Mat Pmat; + + PetscCall(TSGetSNES(*ts, &snes)); + PetscCall(SNESGetKSP(snes, &ksp)); + PetscCall(KSPGetPC(ksp, &pc)); + PetscCall(PCGetType(pc, &pc_type)); + PetscCall(PetscStrcmpAny(pc_type, &use_matceed_pmat, PCJACOBI, PCVPBJACOBI, PCPBJACOBI, "")); + Pmat = use_matceed_pmat ? user->mat_ijacobian : NULL; + PetscCall(TSSetIJacobian(*ts, user->mat_ijacobian, Pmat, NULL, NULL)); + } + } user->time_bc_set = -1.0; // require all BCs be updated if (app_ctx->cont_steps) { // continue from previous timestep data PetscInt count; From 7f2a93030e798875971a84e5c64071653e0488c9 Mon Sep 17 00:00:00 2001 From: James Wright Date: Tue, 12 Mar 2024 11:14:12 -0600 Subject: [PATCH 4/7] fluids: Add KSPSetFromOptions_WithMatCeed --- examples/fluids/include/petsc_ops.h | 4 + examples/fluids/src/differential_filter.c | 20 +++-- examples/fluids/src/grid_anisotropy_tensor.c | 9 +-- examples/fluids/src/misc.c | 1 + examples/fluids/src/petsc_ops.c | 79 +++++++++++++++++++ examples/fluids/src/setupts.c | 37 ++++----- examples/fluids/src/turb_spanstats.c | 11 +-- .../fluids/src/velocity_gradient_projection.c | 11 +-- 8 files changed, 119 insertions(+), 53 deletions(-) diff --git a/examples/fluids/include/petsc_ops.h b/examples/fluids/include/petsc_ops.h index 95038baf39..5f1d890ee8 100644 --- a/examples/fluids/include/petsc_ops.h +++ b/examples/fluids/include/petsc_ops.h @@ -10,6 +10,7 @@ #include #include +#include typedef struct OperatorApplyContext_ *OperatorApplyContext; struct OperatorApplyContext_ { @@ -44,4 +45,7 @@ PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorAp PetscErrorCode DMGetLocalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type); PetscErrorCode DMGetGlobalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type); + +PetscErrorCode CreateSolveOperatorsFromMatCeed(KSP ksp, Mat mat_ceed, PetscBool assemble, Mat *Amat, Mat *Pmat); +PetscErrorCode KSPSetFromOptions_WithMatCeed(KSP ksp, Mat mat_ceed); #endif // petsc_ops_h diff --git a/examples/fluids/src/differential_filter.c b/examples/fluids/src/differential_filter.c index e9c1e06e21..0cc3eafb66 100644 --- a/examples/fluids/src/differential_filter.c +++ b/examples/fluids/src/differential_filter.c @@ -78,13 +78,12 @@ PetscErrorCode DifferentialFilterCreateOperators(Ceed ceed, User user, CeedData } { // Setup LHS Operator and KSP for the differential filtering solve - CeedOperator op_lhs; - OperatorApplyContext mat_ctx; - Mat mat_lhs = NULL; - CeedInt num_comp_qd; - PetscInt dim, num_comp_grid_aniso; - CeedElemRestriction elem_restr_grid_aniso; - CeedVector grid_aniso_ceed; + CeedOperator op_lhs; + Mat mat_lhs; + CeedInt num_comp_qd; + PetscInt dim, num_comp_grid_aniso; + CeedElemRestriction elem_restr_grid_aniso; + CeedVector grid_aniso_ceed; PetscCall(DMGetDimension(user->dm, &dim)); PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_qd_i, &num_comp_qd)); @@ -119,6 +118,7 @@ PetscErrorCode DifferentialFilterCreateOperators(Ceed ceed, User user, CeedData } PetscCallCeed(ceed, CeedQFunctionSetContext(qf_lhs, diff_filter_qfctx)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(qf_lhs, 0)); PetscCallCeed(ceed, CeedQFunctionAddInput(qf_lhs, "q", num_comp_filter, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(qf_lhs, "Grad_q", num_comp_filter * dim, CEED_EVAL_GRAD)); PetscCallCeed(ceed, CeedQFunctionAddInput(qf_lhs, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE)); @@ -149,8 +149,7 @@ PetscErrorCode DifferentialFilterCreateOperators(Ceed ceed, User user, CeedData PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_lhs)); PetscCallCeed(ceed, CeedOperatorDestroy(&op_lhs_sub)); } - PetscCall(OperatorApplyContextCreate(dm_filter, dm_filter, ceed, op_lhs, NULL, NULL, NULL, NULL, &mat_ctx)); - PetscCall(CreateMatShell_Ceed(mat_ctx, &mat_lhs)); + PetscCall(MatCeedCreate(dm_filter, dm_filter, op_lhs, NULL, &mat_lhs)); PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm_filter), &diff_filter->ksp)); PetscCall(KSPSetOptionsPrefix(diff_filter->ksp, "diff_filter_")); @@ -163,8 +162,7 @@ PetscErrorCode DifferentialFilterCreateOperators(Ceed ceed, User user, CeedData PetscCall(KSPSetNormType(diff_filter->ksp, KSP_NORM_NATURAL)); PetscCall(KSPSetTolerances(diff_filter->ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT)); } - PetscCall(KSPSetOperators(diff_filter->ksp, mat_lhs, mat_lhs)); - PetscCall(KSPSetFromOptions(diff_filter->ksp)); + PetscCall(KSPSetFromOptions_WithMatCeed(diff_filter->ksp, mat_lhs)); PetscCall(MatDestroy(&mat_lhs)); PetscCallCeed(ceed, CeedOperatorDestroy(&op_lhs)); diff --git a/examples/fluids/src/grid_anisotropy_tensor.c b/examples/fluids/src/grid_anisotropy_tensor.c index 428843232c..9972b53873 100644 --- a/examples/fluids/src/grid_anisotropy_tensor.c +++ b/examples/fluids/src/grid_anisotropy_tensor.c @@ -73,11 +73,9 @@ PetscErrorCode GridAnisotropyTensorProjectionSetupApply(Ceed ceed, User user, Ce PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", *elem_restr_grid_aniso, basis_grid_aniso, CEED_VECTOR_ACTIVE)); { // -- Setup KSP for L^2 projection - Mat mat_mass = NULL; - OperatorApplyContext mass_matop_ctx; + Mat mat_mass; - PetscCall(OperatorApplyContextCreate(grid_aniso_proj->dm, grid_aniso_proj->dm, ceed, op_mass, NULL, NULL, NULL, NULL, &mass_matop_ctx)); - PetscCall(CreateMatShell_Ceed(mass_matop_ctx, &mat_mass)); + PetscCall(MatCeedCreate(grid_aniso_proj->dm, grid_aniso_proj->dm, op_mass, NULL, &mat_mass)); PetscCall(KSPCreate(comm, &ksp)); PetscCall(KSPSetOptionsPrefix(ksp, "grid_anisotropy_tensor_projection_")); @@ -90,8 +88,7 @@ PetscErrorCode GridAnisotropyTensorProjectionSetupApply(Ceed ceed, User user, Ce PetscCall(KSPSetNormType(ksp, KSP_NORM_NATURAL)); PetscCall(KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT)); } - PetscCall(KSPSetOperators(ksp, mat_mass, mat_mass)); - PetscCall(KSPSetFromOptions(ksp)); + PetscCall(KSPSetFromOptions_WithMatCeed(ksp, mat_mass)); PetscCall(MatDestroy(&mat_mass)); } diff --git a/examples/fluids/src/misc.c b/examples/fluids/src/misc.c index 17abe040b7..131e2b4f9a 100644 --- a/examples/fluids/src/misc.c +++ b/examples/fluids/src/misc.c @@ -243,6 +243,7 @@ PetscErrorCode CreateMassQFunction(Ceed ceed, CeedInt N, CeedInt q_data_size, Ce PetscCallCeed(ceed, CeedQFunctionAddInput(*qf, "u", N, CEED_EVAL_INTERP)); PetscCallCeed(ceed, CeedQFunctionAddInput(*qf, "qdata", q_data_size, CEED_EVAL_NONE)); PetscCallCeed(ceed, CeedQFunctionAddOutput(*qf, "v", N, CEED_EVAL_INTERP)); + PetscCallCeed(ceed, CeedQFunctionSetUserFlopsEstimate(*qf, N)); PetscFunctionReturn(PETSC_SUCCESS); } diff --git a/examples/fluids/src/petsc_ops.c b/examples/fluids/src/petsc_ops.c index f9fc2f801e..42d96f1082 100644 --- a/examples/fluids/src/petsc_ops.c +++ b/examples/fluids/src/petsc_ops.c @@ -507,4 +507,83 @@ PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { PetscFunctionReturn(PETSC_SUCCESS); } +/** + * @brief Return Mats for KSP solve + * + * Uses command-line flag with `ksp`'s prefix to determine if mat_ceed should be used directly or whether it should be assembled. + * If `Amat` is to be assembled, then `Pmat` is equal to `Amat`. + * + * If `Amat` uses `mat_ceed`, then `Pmat` is either assembled or uses `mat_ceed` based on the preconditioner choice in `ksp`. + * + * @param[in] ksp `KSP` object for used for solving + * @param[in] mat_ceed `MATCEED` for the linear operator + * @param[in] assemble Whether to assemble `Amat` and `Pmat` if they are not `mat_ceed` + * @param[out] Amat `Mat` to be used for the solver `Amat` + * @param[out] Pmat `Mat` to be used for the solver `Pmat` + */ +PetscErrorCode CreateSolveOperatorsFromMatCeed(KSP ksp, Mat mat_ceed, PetscBool assemble, Mat *Amat, Mat *Pmat) { + PetscBool use_matceed_pmat, assemble_amat = PETSC_FALSE; + MatType mat_ceed_inner_type; + + PetscFunctionBeginUser; + PetscCall(MatCeedGetInnerMatType(mat_ceed, &mat_ceed_inner_type)); + { // Determine if Amat should be MATCEED or assembled + const char *ksp_prefix = NULL; + + PetscCall(KSPGetOptionsPrefix(ksp, &ksp_prefix)); + PetscOptionsBegin(PetscObjectComm((PetscObject)mat_ceed), ksp_prefix, "", NULL); + PetscCall(PetscOptionsBool("-matceed_assemble_amat", "Assemble the A matrix for KSP solve", NULL, assemble_amat, &assemble_amat, NULL)); + PetscOptionsEnd(); + } + + if (assemble_amat) { + PetscCall(MatConvert(mat_ceed, mat_ceed_inner_type, MAT_INITIAL_MATRIX, Amat)); + if (assemble) PetscCall(MatCeedAssembleCOO(mat_ceed, *Amat)); + + PetscCall(PetscObjectReference((PetscObject)*Amat)); + *Pmat = *Amat; + PetscFunctionReturn(PETSC_SUCCESS); + } else { + PetscCall(PetscObjectReference((PetscObject)mat_ceed)); + *Amat = mat_ceed; + } + + { // Determine if Pmat should be MATCEED or assembled + PC pc; + PCType pc_type; + + PetscCall(KSPGetPC(ksp, &pc)); + PetscCall(PCGetType(pc, &pc_type)); + PetscCall(PetscStrcmpAny(pc_type, &use_matceed_pmat, PCJACOBI, PCVPBJACOBI, PCPBJACOBI, "")); + } + + if (use_matceed_pmat) { + PetscCall(PetscObjectReference((PetscObject)mat_ceed)); + *Pmat = mat_ceed; + } else { + PetscCall(MatConvert(mat_ceed, mat_ceed_inner_type, MAT_INITIAL_MATRIX, Pmat)); + if (assemble) PetscCall(MatCeedAssembleCOO(mat_ceed, *Pmat)); + } + PetscFunctionReturn(PETSC_SUCCESS); +} + +/** + * @brief Runs KSPSetFromOptions and sets Operators based on mat_ceed + * + * See CreateSolveOperatorsFromMatCeed for details on how the KSPSolve operators are set. + * + * @param[in] ksp `KSP` of the solve + * @param[in] mat_ceed `MatCeed` linear operator to solve for + */ +PetscErrorCode KSPSetFromOptions_WithMatCeed(KSP ksp, Mat mat_ceed) { + Mat Amat, Pmat; + + PetscFunctionBeginUser; + PetscCall(KSPSetFromOptions(ksp)); + PetscCall(CreateSolveOperatorsFromMatCeed(ksp, mat_ceed, PETSC_TRUE, &Amat, &Pmat)); + PetscCall(KSPSetOperators(ksp, Amat, Pmat)); + PetscCall(MatDestroy(&Amat)); + PetscCall(MatDestroy(&Pmat)); + PetscFunctionReturn(PETSC_SUCCESS); +} // ----------------------------------------------------------------------------- diff --git a/examples/fluids/src/setupts.c b/examples/fluids/src/setupts.c index 24815c00aa..8157cbc297 100644 --- a/examples/fluids/src/setupts.c +++ b/examples/fluids/src/setupts.c @@ -17,12 +17,11 @@ // @brief Create KSP to solve the inverse mass operator for explicit time stepping schemes PetscErrorCode CreateKSPMassOperator(User user, CeedData ceed_data) { - Ceed ceed = user->ceed; - DM dm = user->dm; - CeedQFunction qf_mass; - CeedOperator op_mass; - OperatorApplyContext mass_matop_ctx; - CeedInt num_comp_q, q_data_size; + Ceed ceed = user->ceed; + DM dm = user->dm; + CeedQFunction qf_mass; + CeedOperator op_mass; + CeedInt num_comp_q, q_data_size; PetscFunctionBeginUser; PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q)); @@ -35,14 +34,14 @@ PetscErrorCode CreateKSPMassOperator(User user, CeedData ceed_data) { PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE)); { // -- Setup KSP for mass operator - Mat mat_mass = NULL; + Mat mat_mass; Vec Zeros_loc; MPI_Comm comm = PetscObjectComm((PetscObject)dm); PetscCall(DMCreateLocalVector(dm, &Zeros_loc)); PetscCall(VecZeroEntries(Zeros_loc)); - PetscCall(OperatorApplyContextCreate(dm, dm, ceed, op_mass, NULL, NULL, Zeros_loc, NULL, &mass_matop_ctx)); - PetscCall(CreateMatShell_Ceed(mass_matop_ctx, &mat_mass)); + PetscCall(MatCeedCreate(dm, dm, op_mass, NULL, &mat_mass)); + PetscCall(MatCeedSetLocalVectors(mat_mass, Zeros_loc, NULL)); PetscCall(KSPCreate(comm, &user->mass_ksp)); PetscCall(KSPSetOptionsPrefix(user->mass_ksp, "mass_")); @@ -53,13 +52,11 @@ PetscErrorCode CreateKSPMassOperator(User user, CeedData ceed_data) { PetscCall(PCJacobiSetType(pc, PC_JACOBI_ROWSUM)); PetscCall(KSPSetType(user->mass_ksp, KSPPREONLY)); } - PetscCall(KSPSetOperators(user->mass_ksp, mat_mass, mat_mass)); - PetscCall(KSPSetFromOptions(user->mass_ksp)); + PetscCall(KSPSetFromOptions_WithMatCeed(user->mass_ksp, mat_mass)); PetscCall(VecDestroy(&Zeros_loc)); PetscCall(MatDestroy(&mat_mass)); } - // Cleanup PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_mass)); PetscCallCeed(ceed, CeedOperatorDestroy(&op_mass)); PetscFunctionReturn(PETSC_SUCCESS); @@ -382,20 +379,16 @@ PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys, Vec *Q PetscCall(TSSetFromOptions(*ts)); if (user->mat_ijacobian) { if (app_ctx->amat_type && !strcmp(app_ctx->amat_type, MATSHELL)) { - PetscBool use_matceed_pmat; - SNES snes; - KSP ksp; - PC pc; - PCType pc_type; - Mat Pmat; + SNES snes; + KSP ksp; + Mat Pmat, Amat; PetscCall(TSGetSNES(*ts, &snes)); PetscCall(SNESGetKSP(snes, &ksp)); - PetscCall(KSPGetPC(ksp, &pc)); - PetscCall(PCGetType(pc, &pc_type)); - PetscCall(PetscStrcmpAny(pc_type, &use_matceed_pmat, PCJACOBI, PCVPBJACOBI, PCPBJACOBI, "")); - Pmat = use_matceed_pmat ? user->mat_ijacobian : NULL; + PetscCall(CreateSolveOperatorsFromMatCeed(ksp, user->mat_ijacobian, PETSC_FALSE, &Amat, &Pmat)); PetscCall(TSSetIJacobian(*ts, user->mat_ijacobian, Pmat, NULL, NULL)); + PetscCall(MatDestroy(&Amat)); + PetscCall(MatDestroy(&Pmat)); } } user->time_bc_set = -1.0; // require all BCs be updated diff --git a/examples/fluids/src/turb_spanstats.c b/examples/fluids/src/turb_spanstats.c index 28bc7c0775..e7cf8f2161 100644 --- a/examples/fluids/src/turb_spanstats.c +++ b/examples/fluids/src/turb_spanstats.c @@ -332,12 +332,10 @@ PetscErrorCode SetupL2ProjectionStats(Ceed ceed, User user, CeedData ceed_data, PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", stats_data->elem_restr_parent_stats, stats_data->basis_stats, CEED_VECTOR_ACTIVE)); { // Setup KSP for L^2 projection - OperatorApplyContext M_ctx; - Mat mat_mass = NULL; - KSP ksp; + Mat mat_mass; + KSP ksp; - PetscCall(OperatorApplyContextCreate(user->spanstats.dm, user->spanstats.dm, user->ceed, op_mass, NULL, NULL, NULL, NULL, &M_ctx)); - PetscCall(CreateMatShell_Ceed(M_ctx, &mat_mass)); + PetscCall(MatCeedCreate(user->spanstats.dm, user->spanstats.dm, op_mass, NULL, &mat_mass)); PetscCall(KSPCreate(comm, &ksp)); PetscCall(KSPSetOptionsPrefix(ksp, "turbulence_spanstats_")); @@ -350,8 +348,7 @@ PetscErrorCode SetupL2ProjectionStats(Ceed ceed, User user, CeedData ceed_data, PetscCall(KSPSetNormType(ksp, KSP_NORM_NATURAL)); PetscCall(KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT)); } - PetscCall(KSPSetOperators(ksp, mat_mass, mat_mass)); - PetscCall(KSPSetFromOptions(ksp)); + PetscCall(KSPSetFromOptions_WithMatCeed(ksp, mat_mass)); user->spanstats.ksp = ksp; PetscCall(MatDestroy(&mat_mass)); } diff --git a/examples/fluids/src/velocity_gradient_projection.c b/examples/fluids/src/velocity_gradient_projection.c index 5cc2572f15..b4287a4a44 100644 --- a/examples/fluids/src/velocity_gradient_projection.c +++ b/examples/fluids/src/velocity_gradient_projection.c @@ -102,12 +102,10 @@ PetscErrorCode VelocityGradientProjectionSetup(Ceed ceed, User user, CeedData ce PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", elem_restr_grad_velo, basis_grad_velo, CEED_VECTOR_ACTIVE)); { // -- Setup KSP for L^2 projection with lumped mass operator - Mat mat_mass = NULL; - OperatorApplyContext mass_matop_ctx; - MPI_Comm comm = PetscObjectComm((PetscObject)grad_velo_proj->dm); + Mat mat_mass; + MPI_Comm comm = PetscObjectComm((PetscObject)grad_velo_proj->dm); - PetscCall(OperatorApplyContextCreate(grad_velo_proj->dm, grad_velo_proj->dm, ceed, op_mass, NULL, NULL, NULL, NULL, &mass_matop_ctx)); - PetscCall(CreateMatShell_Ceed(mass_matop_ctx, &mat_mass)); + PetscCall(MatCeedCreate(grad_velo_proj->dm, grad_velo_proj->dm, op_mass, NULL, &mat_mass)); PetscCall(KSPCreate(comm, &grad_velo_proj->ksp)); PetscCall(KSPSetOptionsPrefix(grad_velo_proj->ksp, "velocity_gradient_projection_")); @@ -118,8 +116,7 @@ PetscErrorCode VelocityGradientProjectionSetup(Ceed ceed, User user, CeedData ce PetscCall(PCJacobiSetType(pc, PC_JACOBI_ROWSUM)); PetscCall(KSPSetType(grad_velo_proj->ksp, KSPPREONLY)); } - PetscCall(KSPSetOperators(grad_velo_proj->ksp, mat_mass, mat_mass)); - PetscCall(KSPSetFromOptions(grad_velo_proj->ksp)); + PetscCall(KSPSetFromOptions_WithMatCeed(grad_velo_proj->ksp, mat_mass)); PetscCall(MatDestroy(&mat_mass)); } From 5da97eb481848ec884009d838cdc5c6d54668dcb Mon Sep 17 00:00:00 2001 From: James Wright Date: Wed, 13 Mar 2024 08:18:36 -0600 Subject: [PATCH 5/7] fluids: Remove old CreateMatShell_Ceed This is fully replaced by MatCeed --- examples/fluids/include/petsc_ops.h | 1 - examples/fluids/src/petsc_ops.c | 82 ----------------------------- 2 files changed, 83 deletions(-) diff --git a/examples/fluids/include/petsc_ops.h b/examples/fluids/include/petsc_ops.h index 5f1d890ee8..c1ec4b3a8e 100644 --- a/examples/fluids/include/petsc_ops.h +++ b/examples/fluids/include/petsc_ops.h @@ -24,7 +24,6 @@ struct OperatorApplyContext_ { PetscErrorCode OperatorApplyContextCreate(DM dm_x, DM dm_y, Ceed ceed, CeedOperator op_apply, CeedVector x_ceed, CeedVector y_ceed, Vec X_loc, Vec Y_loc, OperatorApplyContext *op_apply_ctx); PetscErrorCode OperatorApplyContextDestroy(OperatorApplyContext op_apply_ctx); -PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat); PetscErrorCode DMGetGlobalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type); PetscErrorCode DMGetLocalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type); diff --git a/examples/fluids/src/petsc_ops.c b/examples/fluids/src/petsc_ops.c index 42d96f1082..1603619ff6 100644 --- a/examples/fluids/src/petsc_ops.c +++ b/examples/fluids/src/petsc_ops.c @@ -425,88 +425,6 @@ PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorAp PetscFunctionReturn(PETSC_SUCCESS); } -// ----------------------------------------------------------------------------- -// Wraps the libCEED operator for a MatShell -// ----------------------------------------------------------------------------- -static PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { - OperatorApplyContext ctx; - - PetscFunctionBeginUser; - PetscCall(MatShellGetContext(A, &ctx)); - PetscCall(ApplyCeedOperatorGlobalToGlobal(X, Y, ctx)); - PetscFunctionReturn(PETSC_SUCCESS); -}; - -// ----------------------------------------------------------------------------- -// Returns the computed diagonal of the operator -// ----------------------------------------------------------------------------- -static PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { - OperatorApplyContext ctx; - Vec Y_loc; - PetscMemType mem_type; - - PetscFunctionBeginUser; - PetscCall(MatShellGetContext(A, &ctx)); - Ceed ceed = ctx->ceed; - if (ctx->Y_loc) Y_loc = ctx->Y_loc; - else PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); - - // -- Place PETSc vector in libCEED vector - PetscCall(VecP2C(Y_loc, &mem_type, ctx->y_ceed)); - - // -- Compute Diagonal - PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssembleDiagonal, A, D, 0, 0)); - PetscCall(PetscLogGpuTimeBegin()); - PetscCallCeed(ceed, CeedOperatorLinearAssembleDiagonal(ctx->op, ctx->y_ceed, CEED_REQUEST_IMMEDIATE)); - PetscCall(PetscLogGpuTimeEnd()); - PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssembleDiagonal, A, D, 0, 0)); - - // -- Local-to-Global - PetscCall(VecC2P(ctx->y_ceed, mem_type, Y_loc)); - PetscCall(VecZeroEntries(D)); - PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, D)); - - if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); - PetscFunctionReturn(PETSC_SUCCESS); -}; - -/** - * @brief Create PETSc MatShell object for the corresponding OperatorApplyContext - * - * @param[in] ctx Context that does the action of the operator - * @param[out] mat MatShell for the operator, must be NULL if mat should be created - */ -PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { - MPI_Comm comm_x = PetscObjectComm((PetscObject)(ctx->dm_x)); - MPI_Comm comm_y = PetscObjectComm((PetscObject)(ctx->dm_y)); - PetscInt X_loc_size, X_size, Y_size, Y_loc_size; - VecType X_vec_type, Y_vec_type; - - PetscFunctionBeginUser; - PetscCheck(comm_x == comm_y, PETSC_COMM_WORLD, PETSC_ERR_ARG_NOTSAMECOMM, "Input and output DM must have the same comm"); - - PetscCall(DMGetGlobalVectorInfo(ctx->dm_x, &X_loc_size, &X_size, &X_vec_type)); - PetscCall(DMGetGlobalVectorInfo(ctx->dm_y, &Y_loc_size, &Y_size, &Y_vec_type)); - - if (*mat) { - PetscBool mat_is_shell; - const char *type; - - PetscCall(PetscObjectTypeCompare((PetscObject)*mat, MATSHELL, &mat_is_shell)); - PetscObjectGetType((PetscObject)*mat, &type); - PetscCheck(mat_is_shell, comm_x, PETSC_ERR_ARG_WRONGSTATE, "Mat must be of type SHELL, instead got: %s", type); - PetscCall(MatShellSetContext(*mat, ctx)); - } else PetscCall(MatCreateShell(comm_x, Y_loc_size, X_loc_size, Y_size, X_size, ctx, mat)); - PetscCall(MatShellSetContextDestroy(*mat, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy)); - PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (void (*)(void))MatMult_Ceed)); - PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed)); - - PetscCheck(X_vec_type == Y_vec_type, PETSC_COMM_WORLD, PETSC_ERR_ARG_NOTSAMETYPE, "Vec_type of ctx->dm_x (%s) and ctx->dm_y (%s) must be the same", - X_vec_type, Y_vec_type); - PetscCall(MatShellSetVecType(*mat, X_vec_type)); - PetscFunctionReturn(PETSC_SUCCESS); -} - /** * @brief Return Mats for KSP solve * From d231d939bc5d5ae8ddde621a24e0fa81942b7b00 Mon Sep 17 00:00:00 2001 From: James Wright Date: Wed, 13 Mar 2024 08:28:00 -0600 Subject: [PATCH 6/7] fluids: Add -pmat_pbdiagonal warning --- examples/fluids/src/cloptions.c | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/examples/fluids/src/cloptions.c b/examples/fluids/src/cloptions.c index 571d8dbe30..f44347ec0a 100644 --- a/examples/fluids/src/cloptions.c +++ b/examples/fluids/src/cloptions.c @@ -102,6 +102,11 @@ PetscErrorCode ProcessCommandLineOptions(MPI_Comm comm, AppCtx app_ctx, SimpleBC sizeof(amat_type), &option_set)); if (option_set) PetscCall(PetscStrallocpy(amat_type, (char **)&app_ctx->amat_type)); } + { + PetscBool option_set; + PetscCall(PetscOptionsHasName(NULL, NULL, "-pmat_pbdiagonal", &option_set)); + if (option_set) PetscCall(PetscPrintf(comm, "Warning! -pmat_pbdiagonal no longer used. Pmat assembly determined from -pc_type setting\n")); + } // Provide default ceed resource if not specified if (!ceed_flag) { From 54831c5f47ae63ad598bd81054291277b4e7e451 Mon Sep 17 00:00:00 2001 From: James Wright Date: Wed, 13 Mar 2024 09:50:15 -0600 Subject: [PATCH 7/7] fluids(MatCeed): PetscCeedCall -> PetscCallCeed --- examples/fluids/include/mat-ceed-impl.h | 2 +- examples/fluids/src/mat-ceed.c | 92 ++++++++++++------------- 2 files changed, 47 insertions(+), 47 deletions(-) diff --git a/examples/fluids/include/mat-ceed-impl.h b/examples/fluids/include/mat-ceed-impl.h index 8342ec6f15..0069693bd9 100644 --- a/examples/fluids/include/mat-ceed-impl.h +++ b/examples/fluids/include/mat-ceed-impl.h @@ -23,7 +23,7 @@ @brief Calls a libCEED function and then checks the resulting error code. If the error code is non-zero, then a PETSc error is set with the libCEED error message. **/ -#define PetscCeedCall(ceed_, ...) \ +#define PetscCallCeed(ceed_, ...) \ do { \ int ierr_q_ = __VA_ARGS__; \ if (ierr_q_ != CEED_ERROR_SUCCESS) { \ diff --git a/examples/fluids/src/mat-ceed.c b/examples/fluids/src/mat-ceed.c index 35e8bb41a2..d03ed235db 100644 --- a/examples/fluids/src/mat-ceed.c +++ b/examples/fluids/src/mat-ceed.c @@ -86,7 +86,7 @@ static inline PetscErrorCode VecP2C(Ceed ceed, Vec X_petsc, PetscMemType *mem_ty PetscFunctionBeginUser; PetscCall(VecGetArrayAndMemType(X_petsc, &x, mem_type)); - PetscCeedCall(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); + PetscCallCeed(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); PetscFunctionReturn(PETSC_SUCCESS); } @@ -106,7 +106,7 @@ static inline PetscErrorCode VecC2P(Ceed ceed, CeedVector x_ceed, PetscMemType m PetscScalar *x; PetscFunctionBeginUser; - PetscCeedCall(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); + PetscCallCeed(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); PetscCall(VecRestoreArrayAndMemType(X_petsc, &x)); PetscFunctionReturn(PETSC_SUCCESS); } @@ -128,7 +128,7 @@ static inline PetscErrorCode VecReadP2C(Ceed ceed, Vec X_petsc, PetscMemType *me PetscFunctionBeginUser; PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, mem_type)); - PetscCeedCall(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); + PetscCallCeed(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); PetscFunctionReturn(PETSC_SUCCESS); } @@ -148,7 +148,7 @@ static inline PetscErrorCode VecReadC2P(Ceed ceed, CeedVector x_ceed, PetscMemTy PetscScalar *x; PetscFunctionBeginUser; - PetscCeedCall(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); + PetscCallCeed(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); PetscFunctionReturn(PETSC_SUCCESS); } @@ -243,13 +243,13 @@ static PetscErrorCode MatCeedAssemblePointBlockDiagonalCOO(Mat mat_ceed, Mat mat PetscCall(PetscLogStageRegister("MATCEED Assembly Setup", &stage_amg_setup)); } PetscCall(PetscLogStagePush(stage_amg_setup)); - PetscCeedCall(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonalSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed)); + PetscCallCeed(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonalSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed)); PetscCall(IntArrayC2P(num_entries, &rows_ceed, &rows_petsc)); PetscCall(IntArrayC2P(num_entries, &cols_ceed, &cols_petsc)); PetscCall(MatSetPreallocationCOOLocal(mat_coo, num_entries, rows_petsc, cols_petsc)); free(rows_petsc); free(cols_petsc); - if (!ctx->coo_values_pbd) PetscCeedCall(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_pbd)); + if (!ctx->coo_values_pbd) PetscCallCeed(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_pbd)); PetscCall(PetscRealloc(++ctx->num_mats_assembled_pbd * sizeof(Mat), &ctx->mats_assembled_pbd)); ctx->mats_assembled_pbd[ctx->num_mats_assembled_pbd - 1] = mat_coo; PetscCall(PetscLogStagePop()); @@ -269,12 +269,12 @@ static PetscErrorCode MatCeedAssemblePointBlockDiagonalCOO(Mat mat_ceed, Mat mat else if (strstr(mat_type, "kokkos")) mem_type = CEED_MEM_DEVICE; else mem_type = CEED_MEM_HOST; - PetscCeedCall(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonal(ctx->op_mult, ctx->coo_values_pbd, CEED_REQUEST_IMMEDIATE)); - PetscCeedCall(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_pbd, mem_type, &values)); + PetscCallCeed(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonal(ctx->op_mult, ctx->coo_values_pbd, CEED_REQUEST_IMMEDIATE)); + PetscCallCeed(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_pbd, mem_type, &values)); PetscCall(MatSetValuesCOO(mat_coo, values, INSERT_VALUES)); PetscCall(MatIsSPDKnown(mat_ceed, &is_spd_known, &is_spd)); if (is_spd_known) PetscCall(MatSetOption(mat_coo, MAT_SPD, is_spd)); - PetscCeedCall(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_pbd, &values)); + PetscCallCeed(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_pbd, &values)); } PetscCall(MatAssemblyEnd(mat_coo, MAT_FINAL_ASSEMBLY)); PetscFunctionReturn(PETSC_SUCCESS); @@ -493,21 +493,21 @@ PetscErrorCode MatCeedCreate(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperato ctx->is_ceed_pbd_valid = PETSC_TRUE; ctx->is_ceed_vpbd_valid = PETSC_TRUE; - PetscCeedCall(ctx->ceed, CeedOperatorIsComposite(op_mult, &is_composite)); + PetscCallCeed(ctx->ceed, CeedOperatorIsComposite(op_mult, &is_composite)); if (is_composite) { CeedInt num_sub_operators; CeedOperator *sub_operators; - PetscCeedCall(ctx->ceed, CeedCompositeOperatorGetNumSub(op_mult, &num_sub_operators)); - PetscCeedCall(ctx->ceed, CeedCompositeOperatorGetSubList(op_mult, &sub_operators)); + PetscCallCeed(ctx->ceed, CeedCompositeOperatorGetNumSub(op_mult, &num_sub_operators)); + PetscCallCeed(ctx->ceed, CeedCompositeOperatorGetSubList(op_mult, &sub_operators)); for (CeedInt i = 0; i < num_sub_operators; i++) { CeedInt num_bases, num_comp; CeedBasis *active_bases; CeedOperatorAssemblyData assembly_data; - PetscCeedCall(ctx->ceed, CeedOperatorGetOperatorAssemblyData(sub_operators[i], &assembly_data)); - PetscCeedCall(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL)); - PetscCeedCall(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp)); + PetscCallCeed(ctx->ceed, CeedOperatorGetOperatorAssemblyData(sub_operators[i], &assembly_data)); + PetscCallCeed(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL)); + PetscCallCeed(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp)); if (num_bases > 1) { ctx->is_ceed_pbd_valid = PETSC_FALSE; ctx->is_ceed_vpbd_valid = PETSC_FALSE; @@ -521,9 +521,9 @@ PetscErrorCode MatCeedCreate(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperato CeedBasis *active_bases; CeedOperatorAssemblyData assembly_data; - PetscCeedCall(ctx->ceed, CeedOperatorGetOperatorAssemblyData(op_mult, &assembly_data)); - PetscCeedCall(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL)); - PetscCeedCall(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp)); + PetscCallCeed(ctx->ceed, CeedOperatorGetOperatorAssemblyData(op_mult, &assembly_data)); + PetscCallCeed(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL)); + PetscCallCeed(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp)); if (num_bases > 1) { ctx->is_ceed_pbd_valid = PETSC_FALSE; ctx->is_ceed_vpbd_valid = PETSC_FALSE; @@ -683,13 +683,13 @@ PetscErrorCode MatCeedAssembleCOO(Mat mat_ceed, Mat mat_coo) { PetscCall(PetscLogStageRegister("MATCEED Assembly Setup", &stage_amg_setup)); } PetscCall(PetscLogStagePush(stage_amg_setup)); - PetscCeedCall(ctx->ceed, CeedOperatorLinearAssembleSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed)); + PetscCallCeed(ctx->ceed, CeedOperatorLinearAssembleSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed)); PetscCall(IntArrayC2P(num_entries, &rows_ceed, &rows_petsc)); PetscCall(IntArrayC2P(num_entries, &cols_ceed, &cols_petsc)); PetscCall(MatSetPreallocationCOOLocal(mat_coo, num_entries, rows_petsc, cols_petsc)); free(rows_petsc); free(cols_petsc); - if (!ctx->coo_values_full) PetscCeedCall(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_full)); + if (!ctx->coo_values_full) PetscCallCeed(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_full)); PetscCall(PetscRealloc(++ctx->num_mats_assembled_full * sizeof(Mat), &ctx->mats_assembled_full)); ctx->mats_assembled_full[ctx->num_mats_assembled_full - 1] = mat_coo; PetscCall(PetscLogStagePop()); @@ -709,12 +709,12 @@ PetscErrorCode MatCeedAssembleCOO(Mat mat_ceed, Mat mat_coo) { else if (strstr(mat_type, "kokkos")) mem_type = CEED_MEM_DEVICE; else mem_type = CEED_MEM_HOST; - PetscCeedCall(ctx->ceed, CeedOperatorLinearAssemble(ctx->op_mult, ctx->coo_values_full)); - PetscCeedCall(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_full, mem_type, &values)); + PetscCallCeed(ctx->ceed, CeedOperatorLinearAssemble(ctx->op_mult, ctx->coo_values_full)); + PetscCallCeed(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_full, mem_type, &values)); PetscCall(MatSetValuesCOO(mat_coo, values, INSERT_VALUES)); PetscCall(MatIsSPDKnown(mat_ceed, &is_spd_known, &is_spd)); if (is_spd_known) PetscCall(MatSetOption(mat_coo, MAT_SPD, is_spd)); - PetscCeedCall(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_full, &values)); + PetscCallCeed(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_full, &values)); } PetscCall(MatAssemblyEnd(mat_coo, MAT_FINAL_ASSEMBLY)); PetscFunctionReturn(PETSC_SUCCESS); @@ -965,11 +965,11 @@ PetscErrorCode MatCeedGetCeedOperators(Mat mat, CeedOperator *op_mult, CeedOpera PetscCall(MatShellGetContext(mat, &ctx)); if (op_mult) { *op_mult = NULL; - PetscCeedCall(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult, op_mult)); + PetscCallCeed(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult, op_mult)); } if (op_mult_transpose) { *op_mult_transpose = NULL; - PetscCeedCall(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult_transpose, op_mult_transpose)); + PetscCallCeed(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult_transpose, op_mult_transpose)); } PetscFunctionReturn(PETSC_SUCCESS); } @@ -990,8 +990,8 @@ PetscErrorCode MatCeedRestoreCeedOperators(Mat mat, CeedOperator *op_mult, CeedO PetscFunctionBeginUser; PetscCall(MatShellGetContext(mat, &ctx)); - if (op_mult) PetscCeedCall(ctx->ceed, CeedOperatorDestroy(op_mult)); - if (op_mult_transpose) PetscCeedCall(ctx->ceed, CeedOperatorDestroy(op_mult_transpose)); + if (op_mult) PetscCallCeed(ctx->ceed, CeedOperatorDestroy(op_mult)); + if (op_mult_transpose) PetscCallCeed(ctx->ceed, CeedOperatorDestroy(op_mult_transpose)); PetscFunctionReturn(PETSC_SUCCESS); } @@ -1096,21 +1096,21 @@ PetscErrorCode MatCeedContextCreate(DM dm_x, DM dm_y, Vec X_loc, Vec Y_loc_trans // libCEED objects PetscCheck(CeedOperatorGetCeed(op_mult, &(*ctx)->ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB, "retrieving Ceed context object failed"); - PetscCeedCall((*ctx)->ceed, CeedReference((*ctx)->ceed)); - PetscCeedCall((*ctx)->ceed, CeedOperatorGetActiveVectorLengths(op_mult, &x_loc_len, &y_loc_len)); - PetscCeedCall((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult, &(*ctx)->op_mult)); - if (op_mult_transpose) PetscCeedCall((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult_transpose, &(*ctx)->op_mult_transpose)); - PetscCeedCall((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, x_loc_len, &(*ctx)->x_loc)); - PetscCeedCall((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, y_loc_len, &(*ctx)->y_loc)); + PetscCallCeed((*ctx)->ceed, CeedReference((*ctx)->ceed)); + PetscCallCeed((*ctx)->ceed, CeedOperatorGetActiveVectorLengths(op_mult, &x_loc_len, &y_loc_len)); + PetscCallCeed((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult, &(*ctx)->op_mult)); + if (op_mult_transpose) PetscCallCeed((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult_transpose, &(*ctx)->op_mult_transpose)); + PetscCallCeed((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, x_loc_len, &(*ctx)->x_loc)); + PetscCallCeed((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, y_loc_len, &(*ctx)->y_loc)); // Flop counting { CeedSize ceed_flops_estimate = 0; - PetscCeedCall((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult, &ceed_flops_estimate)); + PetscCallCeed((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult, &ceed_flops_estimate)); (*ctx)->flops_mult = ceed_flops_estimate; if (op_mult_transpose) { - PetscCeedCall((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult_transpose, &ceed_flops_estimate)); + PetscCallCeed((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult_transpose, &ceed_flops_estimate)); (*ctx)->flops_mult_transpose = ceed_flops_estimate; } } @@ -1126,7 +1126,7 @@ PetscErrorCode MatCeedContextCreate(DM dm_x, DM dm_y, Vec X_loc, Vec Y_loc_trans PetscCall(VecGetLocalSize(dm_X_loc, &dm_x_loc_len)); PetscCall(DMRestoreLocalVector(dm_x, &dm_X_loc)); if (X_loc) PetscCall(VecGetLocalSize(X_loc, &X_loc_len)); - PetscCeedCall((*ctx)->ceed, CeedVectorGetLength((*ctx)->x_loc, &ctx_x_loc_len)); + PetscCallCeed((*ctx)->ceed, CeedVectorGetLength((*ctx)->x_loc, &ctx_x_loc_len)); if (X_loc) PetscCheck(X_loc_len == dm_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "X_loc must match dm_x dimensions"); PetscCheck(x_loc_len == dm_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "op must match dm_x dimensions"); PetscCheck(x_loc_len == ctx_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "x_loc must match op dimensions"); @@ -1135,7 +1135,7 @@ PetscErrorCode MatCeedContextCreate(DM dm_x, DM dm_y, Vec X_loc, Vec Y_loc_trans PetscCall(DMGetLocalVector(dm_y, &dm_Y_loc)); PetscCall(VecGetLocalSize(dm_Y_loc, &dm_y_loc_len)); PetscCall(DMRestoreLocalVector(dm_y, &dm_Y_loc)); - PetscCeedCall((*ctx)->ceed, CeedVectorGetLength((*ctx)->y_loc, &ctx_y_loc_len)); + PetscCallCeed((*ctx)->ceed, CeedVectorGetLength((*ctx)->y_loc, &ctx_y_loc_len)); PetscCheck(ctx_y_loc_len == dm_y_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "op must match dm_y dimensions"); // -- Transpose @@ -1206,12 +1206,12 @@ PetscErrorCode MatCeedContextDestroy(MatCeedContext ctx) { PetscCall(PetscFree(ctx->mats_assembled_pbd)); // libCEED objects - PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->x_loc)); - PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->y_loc)); - PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->coo_values_full)); - PetscCeedCall(ctx->ceed, CeedVectorDestroy(&ctx->coo_values_pbd)); - PetscCeedCall(ctx->ceed, CeedOperatorDestroy(&ctx->op_mult)); - PetscCeedCall(ctx->ceed, CeedOperatorDestroy(&ctx->op_mult_transpose)); + PetscCallCeed(ctx->ceed, CeedVectorDestroy(&ctx->x_loc)); + PetscCallCeed(ctx->ceed, CeedVectorDestroy(&ctx->y_loc)); + PetscCallCeed(ctx->ceed, CeedVectorDestroy(&ctx->coo_values_full)); + PetscCallCeed(ctx->ceed, CeedVectorDestroy(&ctx->coo_values_pbd)); + PetscCallCeed(ctx->ceed, CeedOperatorDestroy(&ctx->op_mult)); + PetscCallCeed(ctx->ceed, CeedOperatorDestroy(&ctx->op_mult_transpose)); PetscCheck(CeedDestroy(&ctx->ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB, "destroying libCEED context object failed"); // Deallocate @@ -1243,7 +1243,7 @@ PetscErrorCode MatGetDiagonal_Ceed(Mat A, Vec D) { PetscCall(VecP2C(ctx->ceed, D_loc, &mem_type, ctx->x_loc)); // Compute Diagonal - PetscCeedCall(ctx->ceed, CeedOperatorLinearAssembleDiagonal(ctx->op_mult, ctx->x_loc, CEED_REQUEST_IMMEDIATE)); + PetscCallCeed(ctx->ceed, CeedOperatorLinearAssembleDiagonal(ctx->op_mult, ctx->x_loc, CEED_REQUEST_IMMEDIATE)); // Restore PETSc vector PetscCall(VecC2P(ctx->ceed, ctx->x_loc, mem_type, D_loc)); @@ -1291,7 +1291,7 @@ PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { // Apply libCEED operator PetscCall(PetscLogGpuTimeBegin()); - PetscCeedCall(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult, ctx->x_loc, ctx->y_loc, CEED_REQUEST_IMMEDIATE)); + PetscCallCeed(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult, ctx->x_loc, ctx->y_loc, CEED_REQUEST_IMMEDIATE)); PetscCall(PetscLogGpuTimeEnd()); // Restore PETSc vectors @@ -1351,7 +1351,7 @@ PetscErrorCode MatMultTranspose_Ceed(Mat A, Vec Y, Vec X) { // Apply libCEED operator PetscCall(PetscLogGpuTimeBegin()); - PetscCeedCall(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult_transpose, ctx->y_loc, ctx->x_loc, CEED_REQUEST_IMMEDIATE)); + PetscCallCeed(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult_transpose, ctx->y_loc, ctx->x_loc, CEED_REQUEST_IMMEDIATE)); PetscCall(PetscLogGpuTimeEnd()); // Restore PETSc vectors