Code comments on top of #5353 (#5457)

src/dbm/dbm_matrix.c

@@ -570,6 +570,9 @@ double dbm_maxeps(const dbm_matrix_t *matrix_a, const dbm_matrix_t *matrix_b) {
   assert(matrix_a->dist == matrix_b->dist);
   assert(num_shards == dbm_get_num_shards(matrix_b));
 
+  // Handles asymmetric sparsity: blocks in A but not B are compared against
+  // zero, and blocks in B but not A are checked in a second pass. The
+  // reduction(max) avoids the previous data race on the shared epsilon.
 #pragma omp parallel for reduction(max : epsilon) DBM_OMP_SCHEDULE
   for (int ishard = 0; ishard < num_shards; ++ishard) {
     const dbm_shard_t *const shard_a = &matrix_a->shards[ishard];
@@ -594,6 +597,7 @@ double dbm_maxeps(const dbm_matrix_t *matrix_a, const dbm_matrix_t *matrix_b) {
         }
       }
     }
+    // Second pass: blocks present in B but absent from A (treat A as zero).
     for (int iblock = 0; iblock < shard_b->nblocks; ++iblock) {
       const dbm_block_t *const blk_b = &shard_b->blocks[iblock];
       if (NULL != dbm_shard_lookup(shard_a, blk_b->row, blk_b->col)) {

src/dbm/dbm_multiply.c

@@ -150,8 +150,22 @@ static void backend_stop(backend_context_t *ctx) {
 }
 
 /*******************************************************************************
- * \brief Private routine for multipling two packs.
- * \author Ole Schuett
+ * \brief Private routine for multiplying two packs (C += alpha * A * B).
+ *
+ * Blocks in each pack are grouped by shard (free_index % nshards) and sorted
+ * by sum_index within each group. The algorithm:
+ *  1. Builds shard-boundary lookup tables for A (rows) and B (cols).
+ *  2. For each (shard_row, shard_col) pair, determines the contiguous A and B
+ *     block ranges belonging to that shard.
+ *  3. Performs a merge-join over sum_index: advances A and B cursors in
+ *lockstep, caching the B sub-range for each sum_index so that multiple A blocks
+ *with the same sum_index reuse it without rescanning.
+ *  4. Applies a norm-based filter (alpha^2 * norm_a * norm_b < eps) for early
+ *     rejection before looking up or allocating the C block.
+ *  5. Accumulates matching pairs into a batched GEMM task list, flushing to the
+ *     backend (CPU or GPU) every DBM_MAX_BATCH_SIZE tasks.
+ *
+ * \author Ole Schuett and Hans Pabst
  ******************************************************************************/
 static void multiply_packs(const bool transa, const bool transb,
                            const double alpha, const dbm_pack_t *pack_a,
@@ -240,6 +254,8 @@ static void multiply_packs(const bool transa, const bool transb,
         }
 
         // Merge over sum_index (both ranges sorted by sum_index).
+        // Cache the B sub-range for each sum_index so that multiple A blocks
+        // sharing the same sum_index reuse it without re-scanning B.
         int i = iblock_start, j = jblock_start, last_sum_index = -1;
         int b_range_start = -1, b_range_end = -1;
 

src/dbm/dbm_multiply_comm.c

@@ -39,6 +39,8 @@ static int lcm(const int a, const int b) { return (a * b) / gcd(a, b); }
  * \author Hans Pabst
  ******************************************************************************/
 static int checked_byte_count(const int nelements, const size_t element_size) {
+  // MPI displacements/counts are int; unchecked multiplication could overflow
+  // silently and corrupt the alltoallv layout.
   assert(0 <= nelements);
   assert(element_size <= INT_MAX);
   assert(nelements <= INT_MAX / (int)element_size);
@@ -62,6 +64,8 @@ static inline int isum(const int n, const int input[n]) {
  * \author Ole Schuett and Hans Pabst
  ******************************************************************************/
 static inline void icumsum(const int n, const int input[n], int output[n]) {
+  // Exclusive prefix sum using carried accumulators to avoid a load from
+  // output[i-1] each iteration (removes loop-carried memory dependency).
   int oval = output[0] = 0, ival = input[0];
   for (int i = 1; i < n; i++) {
     output[i] = (oval += ival);
@@ -80,6 +84,8 @@ static void compute_data_recv_count(const int nranks,
                                     const int sum_index_sizes[],
                                     const dbm_pack_block_t blks_recv[],
                                     int data_recv_count[nranks]) {
+  // Derives per-rank data counts locally from the already-received block
+  // metadata, eliminating a collective alltoall for exchanging data amounts.
   memset(data_recv_count, 0, nranks * sizeof(int));
   for (int irank = 0; irank < nranks; irank++) {
     for (int i = 0; i < blks_recv_count[irank]; i++) {
@@ -450,7 +456,7 @@ static dbm_packed_matrix_t pack_matrix(const bool trans_matrix,
                           blks_recv, blks_recv_count_byte, blks_recv_displ_byte,
                           dist->comm);
 
-    // Compute data counts from the received block metadata.
+    // 3rd compute data counts from the received block metadata.
     int data_recv_count[nranks], data_recv_displ[nranks];
     compute_data_recv_count(nranks, blks_recv_count, blks_recv_displ,
                             free_index_sizes, sum_index_sizes, blks_recv,
@@ -469,7 +475,7 @@ static dbm_packed_matrix_t pack_matrix(const bool trans_matrix,
                             data_recv, data_recv_count, data_recv_displ,
                             dist->comm);
 
-    // Post-process received blocks and assemble them into a pack.
+    // 5th post-process received blocks and assemble them into a pack.
     postprocess_received_blocks(nranks, dist_indices->nshards, nblocks_recv,
                                 blks_recv_count, blks_recv_displ,
                                 data_recv_displ, blks_recv);

src/dbm/dbm_multiply_gpu.c

@@ -61,6 +61,9 @@ void dbm_multiply_gpu_start(const int max_batch_size, const int nshards,
  ******************************************************************************/
 static void upload_pack(const dbm_pack_t *pack_host, dbm_pack_gpu_t *pack_dev,
                         const offloadStream_t stream) {
+  // Reallocate only when the device buffer is too small; data_allocated tracks
+  // the device capacity independently from the host-side data_size that changes
+  // each tick (previous code compared against the wrong field).
   const size_t size = pack_host->data_size * sizeof(double);
   if (pack_dev->data_allocated < pack_host->data_size) {
     offload_mempool_device_free(pack_dev->data);

src/dbm/dbm_multiply_gpu.h

@@ -25,6 +25,9 @@ typedef struct {
   offloadEvent_t event;
 } dbm_shard_gpu_t;
 
+// Separate from dbm_pack_t because device allocation lifetime differs from the
+// host pack (which is reallocated each tick); tracks device capacity
+// explicitly.
 typedef struct {
   double *data; // on the device
   int data_allocated;

src/dbm/dbm_multiply_opencl.c

@@ -203,7 +203,7 @@ int dbm_multiply_opencl_launch_kernel(void *stream, double alpha, int ntasks,
       size_t work_size[] = {1, 1, 1}, ibatch = 0;
       const size_t work_tasks = ntasks;
       cl_kernel kernel = NULL;
-      int bk, use_blkrd = 0;
+      int bk, bk0, use_blkrd = 0;
       cl_int size;
       assert(NULL != str && NULL != str->queue);
       /* base init (once): env vars, flags, source, extensions */
@@ -346,8 +346,8 @@ int dbm_multiply_opencl_launch_kernel(void *stream, double alpha, int ntasks,
         task_complete = dbm_multiply_gpu_launch_info(&task, params_host, ntasks,
                                                      param_format, 0);
       }
-      bk = (0 < bk_max ? LIBXS_MIN(dbm_multiply_opencl_bk(task.max_k), bk_max)
-                       : dbm_multiply_opencl_bk(task.max_k));
+      bk0 = dbm_multiply_opencl_bk(task.max_k);
+      bk = (0 < bk_max ? LIBXS_MIN(bk0, bk_max) : bk0);
       { /* per-shape kernel lookup/compile */
         dbm_multiply_opencl_key_t key;
         cl_kernel *kptr;

src/dbm/dbm_shard.c

@@ -160,6 +160,9 @@ static void hashtable_insert(dbm_shard_t *shard, const int block_idx) {
   assert(0 <= block_idx && block_idx < shard->nblocks);
   const dbm_block_t *blk = &shard->blocks[block_idx];
   const unsigned int h = hash(blk->row, blk->col);
+  // Bounded probe: at most hashtable_size steps (load factor < 1 guarantees
+  // termination). Avoids unbounded scans when the slot variable wrapped around
+  // via the mask but the inner iteration continued past the table end.
   int slot = (shard->hashtable_prime * h) & hashtable_mask(shard);
   for (int i = 0; i < shard->hashtable_size; ++i) { // linear probing
     if (shard->hashtable[slot] == 0) {              // 0 means empty
@@ -177,6 +180,8 @@ static void hashtable_insert(dbm_shard_t *shard, const int block_idx) {
  ******************************************************************************/
 dbm_block_t *dbm_shard_lookup(const dbm_shard_t *shard, const int row,
                               const int col) {
+  // Bounded probe count prevents scanning the entire table on a miss when
+  // clusters exist (previous code could re-enter via slot wrap and re-scan).
   int slot = (shard->hashtable_prime * hash(row, col)) & hashtable_mask(shard);
   for (int i = 0; i < shard->hashtable_size; ++i) { // linear probing
     const int block_idx = shard->hashtable[slot];

src/offload/offload_mempool.c

@@ -402,6 +402,7 @@ void offload_mempool_stats_print(int fortran_comm,
                           " Number of allocations  Used [MiB]  Size [MiB]\n",
                           output_unit);
   }
+#if defined(__OFFLOAD)
   if (0 < memstats.device_mallocs) {
     cp_mpi_max_uint64(&memstats.device_peak, 1, comm);
     snprintf(buffer, sizeof(buffer),
@@ -412,6 +413,7 @@ void offload_mempool_stats_print(int fortran_comm,
              (uintptr_t)((memstats.device_peak + (512U << 10)) >> 20));
     OFFLOAD_MEMPOOL_PRINT(print_func, buffer, output_unit);
   }
+#endif
   if (0 < memstats.host_mallocs) {
     cp_mpi_max_uint64(&memstats.host_peak, 1, comm);
     snprintf(buffer, sizeof(buffer),