From 898edf4f64260f5b323b2939a82ab6325e9e8530 Mon Sep 17 00:00:00 2001
From: murfalo <git@murfalo.com>
Date: Tue, 11 Aug 2026 19:27:41 -0400
Subject: [PATCH] Balanced `makepaddedseqdb` across threads

Avoids pathological load imbalance on length-skewed databases like NT by
partitioning work by total sequence length instead of sequence count
---
 src/commons/DBReader.cpp     |  4 +++-
 src/util/makepaddedseqdb.cpp | 13 ++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/commons/DBReader.cpp b/src/commons/DBReader.cpp
index af7aefa81..b20b2198e 100644
--- a/src/commons/DBReader.cpp
+++ b/src/commons/DBReader.cpp
@@ -1347,7 +1347,9 @@ void DBReader<T>::decomposeDomainByAminoAcid(size_t worldRank, size_t worldSize,
             sumCharsAssignedToCurrRank = 0;
             currentRank++;
         }
-        sumCharsAssignedToCurrRank += index[i].length;
+        // Add the sequence length at position i (after any sorting)
+        // This is identical to index[i].length when the reader is unsorted
+        sumCharsAssignedToCurrRank += getEntryLen(i);
         entriesPerWorker[currentRank] += 1;
     }
 
diff --git a/src/util/makepaddedseqdb.cpp b/src/util/makepaddedseqdb.cpp
index da8919dc5..9d8fa6045 100644
--- a/src/util/makepaddedseqdb.cpp
+++ b/src/util/makepaddedseqdb.cpp
@@ -36,8 +36,10 @@ int makepaddedseqdb(int argc, const char **argv, const Command &command) {
 #pragma omp parallel
 {
     unsigned int thread_idx = 0;
+    unsigned int thread_count = 1;
 #ifdef OPENMP
     thread_idx = static_cast<unsigned int>(omp_get_thread_num());
+    thread_count = static_cast<unsigned int>(omp_get_num_threads());
 #endif
     Masker masker(subMat);
     std::string result;
@@ -55,8 +57,13 @@ int makepaddedseqdb(int argc, const char **argv, const Command &command) {
         charSequence = (unsigned char*)malloc(charSeqBufferSize * sizeof(char));
     }
 
-#pragma omp for schedule(static)
-    for (size_t i = 0; i < dbr.getSize(); i++) {
+    size_t rangeStart = 0, rangeSize = 0;
+    // Processing cost scales with sequence length, so balance total bytes instead
+    dbr.decomposeDomainByAminoAcid(thread_count - 1 - thread_idx, thread_count, &rangeStart, &rangeSize);
+    // The reader processes longest first, but the loop below processes shortest first
+    // This keeps the GPU database sorted by increasing sequence length
+    rangeStart = dbr.getSize() - rangeStart - rangeSize;
+    for (size_t i = rangeStart; i < rangeStart + rangeSize; i++) {
         progress.updateProgress();
 
         if (firstIt == SIZE_MAX) {
@@ -151,4 +158,4 @@ int makepaddedseqdb(int argc, const char **argv, const Command &command) {
     }
     dbr.close();
     return EXIT_SUCCESS;
-}
\ No newline at end of file
+}