From 701334031c165c06ebb3dbee0490ab6420df2db7 Mon Sep 17 00:00:00 2001 From: Hans Pabst Date: Thu, 25 Jun 2026 11:23:04 +0200 Subject: [PATCH] Code comments on top of #5353 (#5457) --- src/dbm/dbm_matrix.c | 4 ++++ src/dbm/dbm_multiply.c | 20 ++++++++++++++++++-- src/dbm/dbm_multiply_comm.c | 10 ++++++++-- src/dbm/dbm_multiply_gpu.c | 3 +++ src/dbm/dbm_multiply_gpu.h | 3 +++ src/dbm/dbm_multiply_opencl.c | 6 +++--- src/dbm/dbm_shard.c | 5 +++++ src/offload/offload_mempool.c | 2 ++ 8 files changed, 46 insertions(+), 7 deletions(-) diff --git a/src/dbm/dbm_matrix.c b/src/dbm/dbm_matrix.c index 9c8fb78b7d..407e2ffe9d 100644 --- a/src/dbm/dbm_matrix.c +++ b/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)) { diff --git a/src/dbm/dbm_multiply.c b/src/dbm/dbm_multiply.c index a619f5a11d..f0373ffb18 100644 --- a/src/dbm/dbm_multiply.c +++ b/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; diff --git a/src/dbm/dbm_multiply_comm.c b/src/dbm/dbm_multiply_comm.c index 61080d1bcc..27751087d7 100644 --- a/src/dbm/dbm_multiply_comm.c +++ b/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); diff --git a/src/dbm/dbm_multiply_gpu.c b/src/dbm/dbm_multiply_gpu.c index 5aaa8bdc50..e89385d68b 100644 --- a/src/dbm/dbm_multiply_gpu.c +++ b/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); diff --git a/src/dbm/dbm_multiply_gpu.h b/src/dbm/dbm_multiply_gpu.h index bd1edc5bfe..72735b76af 100644 --- a/src/dbm/dbm_multiply_gpu.h +++ b/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; diff --git a/src/dbm/dbm_multiply_opencl.c b/src/dbm/dbm_multiply_opencl.c index e74fa83c1b..230ac1925c 100644 --- a/src/dbm/dbm_multiply_opencl.c +++ b/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; diff --git a/src/dbm/dbm_shard.c b/src/dbm/dbm_shard.c index 72fd4107de..4e7fbccb9f 100644 --- a/src/dbm/dbm_shard.c +++ b/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]; diff --git a/src/offload/offload_mempool.c b/src/offload/offload_mempool.c index 4ef0580c8e..8267279db7 100644 --- a/src/offload/offload_mempool.c +++ b/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),