How to integrate ginkgo with my hydrodynamic model as an external library?

[lj-cug]

Dear Developers,
ginkgo is a powerful linear solver that can be ported into different processors including Nvidia, Intel and AMD etc. I have integrated PETSc, AmgX with my hydrodynamic code, and ginkgo also attracted me. Before using ginkgo, I have some questions to consult with developers, and hope you can guide me:
(1) I can see the distributed multi-grid algorithm implemention in examples code, and I want to know how to inject PETSc Mat A and Vector b into ginkgo's A and b?
Because my hydrodynamic code uses PETSc to integrate the external linear algebraic library AmgX (Nvidia) like the project AmgXWrapper? In setA.cpp of AmgXWrapper, the folowing function can easily upload my assembly PETSc-CSR matrix into AmgX following (Vector b operation is alike):

AMGX_matrix_upload_distributed(
AmgXA, nGlobalRows, nLocalRows, row[nLocalRows],
1, 1, row.data(), col.data(), data.data(),
nullptr, dist);

(2) Does the MPI-parallelized ginkgo (the distributed solvers) achieve the production-level application? And the application in multi-GPU clusters?

(3) Can the example code in external-lib-interfacing be used to implement the above-mentioned AmgXWrapper ?

Sincerely,
Li Jian
China University of Geoscience

[MarcelKoch]

Hi @lj-cug, I've converted your issue into a discussion since I think that fits better the topic of your questions.
Regarding your questions:

1. Converting a Petsc matrix into a Ginkgo matrix is highly non-trivial. To my knowledge, Petsc doesn't expose internal data, so you essentially have to extract all matrix entries from Petsc and copy them over to Ginkgo. The code you linked to might be a good reference. You should be able to replace the std::vector with gko::array and then pass (or better move) the arrays to the constructor of your Ginkgo matrix.
2. Yes, the distributed solvers are production-ready. You can run them on any number of GPUs. In terms of preconditioners, you have two options. 1) Use our AMG, which might require some tuning to find good parameters for your application, and 2) use the Schwarz preconditioner, where you can choose any preconditioner (e.g. Jacobi, Gauss-Seidel, etc.) that is applied only to the local matrix.
3. The external-lib-interfacing example might not be a good starting point for you. It has a lot of deal.II specifics, which obscure the actual interfacing parts. Maybe the paper https://arxiv.org/abs/2509.16081 can help you. It also has references to other applications interfacing with Ginkgo.

Feel free to reach out if you encounter any issues.

[lj-cug]

Thank you very much. I will compare AmgXWrapper code with ginkgo matrix and vector operations. Keep in touch.

[lj-cug]

PETSc_AmgXWrapper code

PetscErrorCode AmgXSolver::setA(const Mat &A)
{
PetscFunctionBeginUser;

Mat                 localA;
IS                  devIS;
PetscInt            nGlobalRows,
                    nLocalRows;
PetscBool           usesOffsets;

std::vector<PetscInt>       row;
std::vector<PetscInt64>     col;
std::vector<PetscScalar>    data;
std::vector<PetscInt>       partData;

// get number of rows in global matrix
PetscCall(MatGetSize(A, &nGlobalRows, nullptr));

// get the row indices of redistributed matrix owned by processes in gpuWorld
PetscCall(getDevIS(A, devIS));

// get sequential local portion of redistributed matrix
PetscCall(getLocalA(A, devIS, localA));

// get compressed row layout of the local Mat
PetscCall(getLocalMatRawData(localA, nLocalRows, row, col, data));

// destroy local matrix
PetscCall(destroyLocalA(A, localA));

// get a partition vector required by AmgX
PetscCall(getPartData(devIS, nGlobalRows, partData, usesOffsets));

// we use 64 bit column indices
std::vector<PetscInt64> offsets;
offsets.assign(partData.begin(), partData.end());

... ...  // Some operations to transfer PETSc-related matrix's data to AmgX's space

PetscFunctionReturn(0);

}

ginkgo's example distributed-multigrid-preconditioned-solver

// If have already the partition in my code, is it not necessary to create the partition like that in the example code ? I want to inject the PETSc's local matrix and vector into Ginkgo's distributed matrix and vector.
auto partition = gko::share(part_type::build_from_global_size_uniform(
exec->get_master(), comm.size(),
static_cast(num_rows)));

// Assemble the matrix A and fill the right-hand-side b and unknown vector x
// In real HPC case, the global matrix and vector will be very large?
for (int i = range_start; i < range_end; i++) {

// Use .emplace_back to create the following matrix and vector data
A_data ...;
b_data ...;
x_data ...;
}

// Create the distributed Matrix and Vector on Host
auto A_host = gko::share(dist_mtx::create(exec->get_master(), comm));
auto x_host = dist_vec::create(exec->get_master(), comm);
auto b_host = dist_vec::create(exec->get_master(), comm);

// Transfer matrix and vector data into the corresponding ginkgo array
** A_data is the global Matrix (it's dense mode in ginkgo), if I have local Matrix like the PETSc's matrix, how can I transfer the local matrix into A_host of ginkgo? **
** I cannot see what's the dimension of A_host like the operation A_data.size = {num_rows, num_rows}; ??? **
A_host->read_distributed(A_data, partition);
b_host->read_distributed(b_data, partition);
x_host->read_distributed(x_data, partition);

[lj-cug]

(1) Actually I want to transfer the local matrix and vector directly into ginkgo's distributed matrix and vector. In the example code of distributed-multigrid-preconditioned-solver, I think the calling of read_distributed has the same operations like PETSc_AmgXWrapper.
Also AmgXWrapper used
AMGX_distribution_set_partition_data(dist, AMGX_DIST_PARTITION_OFFSETS, offsets.data());
to set the partition like ginkgo

I am not sure how to use some kind of mode for gko::distributed::Partition to do what like AMGX_distribution_set_partition_data ?

(2) And how to upload PETSc data into ginkgo like following ?
AMGX_matrix_upload_distributed(
AmgXA, nGlobalRows, nLocalRows, row[nLocalRows],
1, 1, row.data(), col.data(), data.data(),
nullptr, dist);

(3) I can see gko::matrix::Csr<ValueType, IndexType>* matrix of creating CSR formatted matrix in Ginkgo. Can distributed CSR matrix be used in Ginkgo? This only has serial example code "poission-solver", it's better to have a distributed version. AmgX library uses default CSR matrix.

[MarcelKoch]

1. Can you clarify what you mean by transferring the local matrix directly into Ginkgo? Where does the matrix/vector come from? Are they stored in some internal format or as PETSc matrices/vectors?

2 + 3) In Ginkgo we have gko::experimental::distributed::Matrix and gko::experimental::distributed::Vector, which are the classes relevant to you. By default the distributed matrix uses the CSR format locally, but you can change that. The constructor for creating a distributed matrix from existing data is:

std::unique_ptr<Matrix> Matrix::create(
    std::shared_ptr<const Executor> exec, mpi::communicator comm,
    index_map<local_index_type, global_index_type> imap,
    std::shared_ptr<LinOp> local_linop,
    std::shared_ptr<LinOp> non_local_linop);

That means, you will need to define a Ginkgo linop (e.g. CSR matrix) with the local block (rows and columns local to the process) and a linop for the non-local block (rows local to the process and columns not local to the process). Additionally you need to define an index map, which defines a mapping for the non-local columns to their global column indices.
How you create these objects depends on 1).

[lj-cug]

Following your guides, I firstly write a code to transfer PETSc CSR Matrix to Ginkgo distributed matrix. Could you please help to check the code before I compile and check the code in my code, especially the Ginkgo operations ?

template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
void petsc_to_ginkgo_distributed(
    Mat petsc_mat, 
    MPI_Comm petsc_comm,
    std::shared_ptr <gko::Executor> exec,
    std::shared_ptr <gko::experimental::distributed::Matrix<
        ValueType, LocalIndexType, GlobalIndexType>>& dist_matrix)
{
    // 1. Make sure to assembly Matrix
    MatAssemblyBegin(petsc_mat, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(petsc_mat, MAT_FINAL_ASSEMBLY);
    
    // 2. Get the local and global information of Mat
    PetscInt local_rows, global_rows, nGlobalCols;
    MatGetLocalSize(petsc_mat, &local_rows, nullptr);   // local dimension of Mat A
    MatGetSize(petsc_mat, &global_rows, &nGlobalCols);  // global dimension of Mat A
    
    PetscInt local_start, local_end;
    MatGetOwnershipRange(petsc_mat, &local_start, &local_end);
    
    // 3. Get CSR formatted matrix data
    const PetscInt* petsc_row_ptr = nullptr;
    const PetscInt* petsc_col_idx = nullptr;
    const PetscScalar* petsc_vals = nullptr;
    PetscBool done = PETSC_FALSE;
    
    PetscInt nrows;
    PetscBool is_compressed;
    
    MatGetRowIJ(petsc_mat, 
                0,              // begin from index-0
                PETSC_FALSE,
                PETSC_FALSE, 
                &nrows, 
                &petsc_row_ptr, 
                &petsc_col_idx, 
                &is_compressed);
    
    MatSeqAIJGetArrayRead(petsc_mat, &petsc_vals);
    if (nrows != local_rows) {
        // throw exception ...
    }
    
    // 4. calculate local nnz
    PetscInt local_nnz = petsc_row_ptr[local_rows] - petsc_row_ptr[0];
    
    // 5. Transfer to Ginkgo arrray ...
    auto gko_values = gko::array<ValueType>(exec, local_nnz);
    auto gko_col_idx = gko::array<IndexType>(exec, local_nnz);
    auto gko_row_ptr = gko::array<IndexType>(exec, nLocalRows + 1);
    
    auto host_exec = exec->get_master();
    bool need_copy = (exec != host_exec);
    
    if (need_copy) {
        auto host_values = gko::array<ValueType>(host_exec, local_nnz);
        auto host_col_idx = gko::array<IndexType>(host_exec, local_nnz);
        auto host_row_ptr = gko::array<IndexType>(host_exec, nLocalRows + 1);
        
        auto h_vals = host_values.get_data();
        auto h_cols = host_col_idx.get_data();
        auto h_rows = host_row_ptr.get_data();
        
        for (PetscInt i = 0; i < local_nnz; ++i) {
            h_vals[i] = static_cast<ValueType>(vals[i]);
            h_cols[i] = static_cast<IndexType>(col_idx[i] + local_start);
        }        
        for (PetscInt i = 0; i <= nLocalRows; ++i) {
            h_rows[i] = static_cast<IndexType>(row_ptr[i] - row_ptr[0]);
        }

        gko_values = gko::array<ValueType>(exec, host_values);
        gko_col_idx = gko::array<IndexType>(exec, host_col_idx);
        gko_row_ptr = gko::array<IndexType>(exec, host_row_ptr);
    } else {
        auto gko_vals = gko_values.get_data();
        auto gko_cols = gko_col_idx.get_data();
        auto gko_rows = gko_row_ptr.get_data();
        
        for (PetscInt i = 0; i < local_nnz; ++i) {
            gko_vals[i] = static_cast<ValueType>(vals[i]);
            gko_cols[i] = static_cast<IndexType>(col_idx[i] + local_start);
        }
        
        for (PetscInt i = 0; i <= nLocalRows; ++i) {
            gko_rows[i] = static_cast<IndexType>(row_ptr[i] - row_ptr[0]);
        }
    }
    
	// 6. Create Ginkgo distributed matrix
    auto gko_comm = gko::experimental::mpi::communicator(petsc_comm);
    dist_matrix = DistMtx::create(
        exec,
        gko_comm,
        gko::dim<2>{static_cast<gko::size_type>(nGlobalRows),
                    static_cast<gko::size_type>(nGlobalCols)},
        std::move(gko_values),
        std::move(gko_row_ptr),
        std::move(gko_col_idx),
        static_cast<gko::size_type>(nLocalRows)
    );	
		
	// 7. Restore PETSc AIJ array
    MatRestoreRowIJ(petsc_mat, 0, PETSC_FALSE, PETSC_FALSE, 
                   &nrows, &petsc_row_ptr, &petsc_col_idx, &is_compressed);
    MatSeqAIJRestoreArrayRead(petsc_mat, &petsc_vals);
}

Thank you very much!

[MarcelKoch]

Hi @lj-cug, sorry for the delayed reply. Your point 6. is nearly correct for creating a non-distributed CSR matrix. The correct way would be:

csr_matrix = Csr::create(
    exec,
    gko::dim<2>{static_cast<gko::size_type>(nLocalRows),
                static_cast<gko::size_type>(nLocalCols)},
    std::move(gko_values),
    std::move(gko_col_idx),
    std::move(gko_row_ptr)
);	

Now, for the distributed case you can either:

1. extract all petsc entries into a gko::matrix_data (or gko::device_matrix_data) structure and create an empty matrix with DistMtx::create(exec, gko_comm) and then call ->read_distributed on that matrix. You will also need to provide a row partition for this.
2. create two matrices per process, one for the local diagonal block and one for the off-diagonal block (which we typically call local and non-local blocks). The documentation https://ginkgo-project.github.io/ginkgo-generated-documentation/doc/develop/classgko_1_1experimental_1_1distributed_1_1Matrix.html#details should clarify what these blocks mean. I don't know you would get these two blocks from PETSc. This is something you would need to consult its documentation for. In addition, you will have to provide an index_map. This requires to know how the column indices of the non-local block map into the global indices. Again, I don't know how to extract this from PETSc.

My advice is to go with 1., since you need to know a lot more about both Ginkgo and PETSc internals to make 2. work.

[lj-cug]

My hydrodynamics code was based on unstructured-mesh, so I think CSR formatted matrix is more suitable for me. Meanwhile gko::matrix_data is dense matrix, I have a question:
In the example code "distributed-solver", the data are structured. Is ginkgo more inclined to which one, dense matrix or CSR format ? especially for GPU.