Add distributed SpGEMM by pratikvn · Pull Request #2047 · ginkgo-project/ginkgo

pratikvn · 2026-07-08T17:59:36Z

This PR adds a simple 1D SpGEMM algorithm, a->spgemm(b,c) for the operation $C = A \times B$ It tries to reuse the existing components from the matrix setup. The phases are:

Merge: Combine into a global-column CSR matrix
Communicate: Exchange remote B rows
Assemble: Append the remote rows to local-B rows and remap A's global columns
Local SpGEMM: Delegate to the local device SpGEMM
Reassemble: Split the local C result into diag and off-diag matrix with a simplified version of separate_diag_off_diag.

Here is the performance for a relatively dense matrix (random sparsity with ~32 nnz per row):

MarcelKoch

I'm still in the process of looking through it. But right now it's a bit too complicated for me to follow.
I think it can be made simpler, by just using A and B in global indices. The compression only happens at the last step, after the local C matrix has been computed (in global indices).

MarcelKoch · 2026-07-09T07:30:16Z

+     * @param b  the right-hand operand of the product.
+     * @param c  a pre-created distributed Matrix (may be empty) to fill.
+     */
+    void spgemm(ptr_param<const Matrix> b, ptr_param<Matrix> c) const;


I think in the Csr this is called multiply. We should keep the same name here.

MarcelKoch · 2026-07-09T07:31:18Z

+    {
+        return partition_;
+    }


can you move this to the .cpp?

MarcelKoch · 2026-07-09T07:31:33Z

+    get_row_partition() const
+    {
+        return row_partition_;
+    }


MarcelKoch · 2026-07-09T07:33:37Z

+     * read_distributed. May be null if the matrix was not read.
+     */
+    std::shared_ptr<const Partition<local_index_type, global_index_type>>
+    get_row_partition() const


maybe we combine this with #2033 and store a column index map and a row index map. (The row index map would be trivial, since there are no non-local indices.)

MarcelKoch · 2026-07-09T07:35:44Z

+
+// Turns a per-rank counts vector into the corresponding exclusive-prefix-sum
+// offsets vector (as used for all_to_all_v send/recv displacement arrays).
+std::vector<int> counts_to_offsets(const std::vector<int>& counts)


why is this not using std::prefix_sum instead?

MarcelKoch · 2026-07-09T10:30:11Z

+    // A's imap_ gives the remote global column indices (= B rows) this rank
+    // needs and their owner ranks.


I'm wondering if this index communication step could be handled through the collective_communicator interface.
I think if you use create_inverse and then if we would have a function create_index_map from the collective communicator, you get the rows of B to send.

No, the index_map would need to have a create_inverse function. Although the result of that wouldn't be an index_map, since the send indices are again local indices. Maybe it should just be a create_send_indices(comm) -> array

You can also take a look at the row_gatherer implementation. It does the same in its constructor.

MarcelKoch · 2026-07-13T09:47:38Z

+    // A's imap_ gives the remote global column indices (= B rows) this rank
+    // needs and their owner ranks.


The whole communication seems overly complicated to me. I think it would be a lot simpler, if we communicated just the (device) matrix data. Then we only need to:

create send indices

essentially a row gather with the send indices

communicate all rows.

MarcelKoch · 2026-07-13T09:49:01Z

+        b_aug_row_ptrs[b_local_nrows + i + 1] =
+            static_cast<GlobalIndexType>(recv_nnz_counts[i]);
+    }
+    // Prefix-sum the remote-row counts into offsets.


why not use std::prefix_sum directly?

MarcelKoch · 2026-07-13T09:57:40Z

+    // Remap A's global columns to B_augmented row indices via A's imap_
+    // (combined index space); row_ptrs and values carry over from a_merged.


Is the remapping for both A and B even necessary? If we take both in global indexing, then there is only the unnecessary storage for the row_ptrs, right? The matrices can still be multiplied. I think we can live with that for now.

We cannot use global indices and compress at the end, because rocsparse does not support spgemm in int64.

TBH, that still seems like a better choice compared to the current implementation.

Sorry, but I dont understand what you mean. We need to have the local spgemm's in int32 because rocsparse does not support int64. That is the reason that I had to go with the current approach.

I mean to just not support it on AMD for now. At least until it's cleaned up more and we have found a better implementation (or support spgemm ourselves)

I think it is important that we support AMD GPUs. I will try to simplify the code, but we should still aim to support AMD GPUs.

Also doing the local SpGEMMs in int32 is signifcantly faster, and scales better as more ranks -> more local spgemms.

I just think that optimization is a low priority right now. We should just get it to work.

Turns out even CUDA also only supports int64 from version > 13. So without the current approach, distributed spgemm will fail in almost all cases. Our default is Matrix<vtype, int, long>, which will not have a supported distributed spgemm.

MarcelKoch · 2026-07-13T10:07:43Z

+    // The local spgemm runs with LocalIndexType (32-bit) indices, which every
+    // backend supports (rocSPARSE has no 64-bit spgemm). Compress B_augmented's
+    // global columns to a compact local space on the executor;
+    // b_aug_distinct_cols maps each compact index back to its global column for
+    // the reassemble below.


I think the compressing can be done using a new index map. Then we also don't need the additional compress_columns kernel.

Also, the remapping has to match the non-local indexing for A, right?

pratikvn self-assigned this Jul 8, 2026

ginkgo-bot added reg:build This is related to the build system. reg:testing This is related to testing. mod:all This touches all Ginkgo modules. labels Jul 8, 2026

pratikvn requested review from MarcelKoch and yhmtsai July 8, 2026 18:00

pratikvn added 4 commits July 12, 2026 14:33

add row_partition as Matrix member

5bb979f

add spgemm core and tests

1f91874

add device reassemble kernel for dist spgemm

8cff38a

use int32 for local spgemms

5a9c403

pratikvn force-pushed the feat/distr-spgemm branch from 35a5413 to 5a9c403 Compare July 12, 2026 12:34

skip hip <,long,long> tests

b1568bb

MarcelKoch reviewed Jul 13, 2026

View reviewed changes

		// A's imap_ gives the remote global column indices (= B rows) this rank
		// needs and their owner ranks.

		// Remap A's global columns to B_augmented row indices via A's imap_
		// (combined index space); row_ptrs and values carry over from a_merged.

Uh oh!

Conversation

pratikvn commented Jul 8, 2026

Uh oh!

MarcelKoch left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants