diff --git a/examples/fluids/src/mat-ceed.c b/examples/fluids/src/mat-ceed.c index bf2f52b006..5c5a8760db 100644 --- a/examples/fluids/src/mat-ceed.c +++ b/examples/fluids/src/mat-ceed.c @@ -209,6 +209,58 @@ static PetscErrorCode MatGetBlockDiagonal_Ceed(Mat mat_ceed, Mat *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); +} + /** @brief Get on-process diagonal block of `MATCEED` @@ -472,6 +524,8 @@ PetscErrorCode MatCreateCeed(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperato 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(MatShellSetOperation(*mat, MATOP_GET_BLOCK_DIAGONAL, (void (*)(void))MatGetBlockDiagonal_Ceed)); PetscCall(MatShellSetOperation(*mat, MATOP_GET_VBLOCK_DIAGONAL, (void (*)(void))MatGetVariableBlockDiagonal_Ceed)); PetscCall(MatShellSetVecType(*mat, vec_type)); @@ -538,6 +592,8 @@ PetscErrorCode MatCeedCopy(Mat mat_ceed, Mat mat_other) { 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)); PetscCall(MatShellSetOperation(mat_other, MATOP_GET_BLOCK_DIAGONAL, (void (*)(void))MatGetBlockDiagonal_Ceed)); PetscCall(MatShellSetOperation(mat_other, MATOP_GET_VBLOCK_DIAGONAL, (void (*)(void))MatGetVariableBlockDiagonal_Ceed)); {