解压和内存拷贝可以并行，但是效果一般，可以调节并行和串行

Co-authored-by: Copilot <copilot@github.com>

src/sort/phase_1.cpp

@@ -26,7 +26,9 @@ void phase1Pipeline() {
     const size_t kReadBufSize = 4L * 1024 * 1024 * phase1Arg.numThread;  // 平均每线程4M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略）
 
     phase1Arg.uncompressBufBytes = nsgv::gSortArg.MAX_MEM * 0.9; // 比最大内存参数小点
-    phase1Arg.threadUncompressWrap.Resize(phase1Arg.numThread); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
+    for (int i = 0; i<phase1Arg.UNCOMPRESS_BUF_NUM; ++i) {
+        phase1Arg.threadUncompressWrap[i].Resize(phase1Arg.numThread); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
+    }
     phase1Arg.singleThreadMemBytes = kReadBufSize * BAM_COMPRESS_RATIAO;
     spdlog::info("max mem: {}, uncompress mem: {}, single thread mem: {}", nsgv::gSortArg.MAX_MEM, phase1Arg.uncompressBufBytes, phase1Arg.singleThreadMemBytes);
     for (int i = 0; i < phase1Arg.READ_BUF_NUM; ++i) {
@@ -38,11 +40,12 @@ void phase1Pipeline() {
 
     PROF_G_BEG(mid_all);
     /* create threads */
-    pthread_t tidArr[2]; // 2-stage pipeline
+    pthread_t tidArr[3]; // 2-stage pipeline
     pthread_create(&tidArr[0], NULL, phase1ReadFile, &phase1Arg);
     pthread_create(&tidArr[1], NULL, phase1Uncompress, &phase1Arg);
+    pthread_create(&tidArr[2], NULL, phase1MemCopy, &phase1Arg);
 
-    for (int i = 0; i < 2; ++i) pthread_join(tidArr[i], NULL);
+    for (int i = 0; i < 3; ++i) pthread_join(tidArr[i], NULL);
 
     spdlog::info("all bams num: {}", phase1Arg.bamNum);
 

src/sort/phase_1.h

@@ -98,7 +98,8 @@ struct ThreadUncompressWrap {
 /* 第一阶段的多线程流水线参数 */
 struct Phase1PipelineArg {
     static const int READ_BUF_NUM = 2; // 读入的buf数量
-    static const int UNCOMPRESS_BUF_NUM = 1; // 解压的buf数量, 只有一个
+    static const int UNCOMPRESS_BUF_NUM = 1; // 解压的buf数量
+
     static const int COMPRESS_BUF_NUM = 1;  // 压缩的buf数量, 只有一个
     static const int WRITE_BUF_NUM = 2;  // 写入文件buf数量
 
@@ -114,10 +115,15 @@ struct Phase1PipelineArg {
     uint64_t memCopyOrder = 0; // 串行拷贝解压数据到uncompressData的轮次编号，与上边的uncompressOrder对应
 
     volatile int readFinish = 0;
+    volatile int uncompressFinish = 0;
+
     yarn::lock_t* readSig;
     yarn::lock_t* uncompressSig;
+
     ReadBuffer readData[READ_BUF_NUM];  // 用来读如数据，双缓冲
-    ThreadUncompressWrap threadUncompressWrap;        // 每个thread一个，用来保存解压后的block数据
+    ThreadUncompressWrap threadUncompressWrap[UNCOMPRESS_BUF_NUM];  // 每个thread一个，用来保存解压后的block数据
+    DataBuffer lastRoundBuf;                                            // 上一轮不完整的数据放在这里
+
     UncompressBlockBuffer uncompressData;  // 所有线程共用一个，串行往这里添加解压后的block数据
     BamArr allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据
 
@@ -141,6 +147,7 @@ struct Phase1PipelineArg {
     Phase1PipelineArg() {
         readSig = yarn::NEW_LOCK(0);
         uncompressSig = yarn::NEW_LOCK(0);
+        lastRoundBuf.AllocMem(SINGLE_BLOCK_SIZE);
     }
 };
 

src/sort/phase_1_uncompress.cpp

@@ -138,18 +138,20 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) {
 
     Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
     ReadBuffer& readData = p.readData[p.uncompressOrder % p.READ_BUF_NUM];
+    ThreadUncompressWrap& uncompressWrap = p.threadUncompressWrap[p.uncompressOrder % p.UNCOMPRESS_BUF_NUM];
 
-    tid = idx; // 静态分配任务，此时用idx代替tid
+    tid = idx;  // 静态分配任务，此时用idx代替tid
     int startIdx = START_IDX(idx, p.numThread, readData.startAddrArr.size());
     int stopIdx = STOP_IDX(idx, p.numThread, readData.startAddrArr.size());
 
-    auto &blockBuf = p.threadUncompressWrap.threadUncompressDataArr[tid].blockBuf;
+    auto& blockBuf = uncompressWrap.threadUncompressDataArr[tid].blockBuf;
-    auto &bamArr = p.threadUncompressWrap.threadUncompressDataArr[tid].bamArr;
+    auto& bamArr = uncompressWrap.threadUncompressDataArr[tid].bamArr;
 
     // 开辟足够的内存
     if (stopIdx - startIdx > blockBuf.maxLen / SINGLE_BLOCK_SIZE) {
         blockBuf.ReAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE);
     }
+    uncompressWrap.threadUncompressDataArr[tid].blockNum = stopIdx - startIdx;
 
     // 解压block
     for (int i = startIdx; i < stopIdx; ++i) {
@@ -192,7 +194,7 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) {
 // 处理相邻线程的block数据，可能有bam跨越这两个线程的block（GATK的bam）
 static void handleAdjacentThreadBlock(Phase1PipelineArg& p) {
     auto& uncompressData = p.uncompressData;
-    auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr;
+    auto& threadUncompressDataArr = p.threadUncompressWrap[p.uncompressOrder % p.UNCOMPRESS_BUF_NUM].threadUncompressDataArr;
     size_t offset = 0;  // 当前线程对应的全局数据的起始偏移量
     size_t bamOffset = p.allBams.Size(); // 当前线程解析的bam在全局数据中的偏移量
 
@@ -203,6 +205,7 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) {
         auto& bamArr = threadUncompressDataArr[tid].bamArr;
         auto& firstBam = threadUncompressDataArr[tid].firstBam;
         auto& lastBamBuf = threadUncompressDataArr[tid].lastBamBuf;
+        auto& lastRoundBuf = p.lastRoundBuf;
 
         bool hasLastData = false; // 上一个block里有不完整bam数据
         int lastDataLen = 0; // 上一个block里不完整bam数据的长度
@@ -210,11 +213,11 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) {
 
         if (tid == 0) { // 第一个线程
             // 检查一下bam的定位是否正确
-            hasLastData = uncompressData.usedBufSize != uncompressData.lastEndPos;
+            hasLastData = lastRoundBuf.curLen > 0;
             if (hasLastData || blockBuf.readPos > 0) {
-                lastDataLen = uncompressData.usedBufSize - uncompressData.lastEndPos;
+                lastDataLen = lastRoundBuf.curLen;
                 leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据
-                lastBamBuf.MemCopy(uncompressData.dataBuf + uncompressData.lastEndPos, lastDataLen);
+                lastBamBuf.MemCopy(lastRoundBuf.data, lastDataLen);
                 lastBamBuf.MemCopy(blockBuf.data, leftDataLen);
             }
         } else {
@@ -250,69 +253,26 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) {
         }
         offset += threadUncompressDataArr[tid].blockBuf.curLen;
         bamOffset += threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
-    }
+        if (tid == p.numThread - 1) {  // 最后一个线程
-}
+            lastRoundBuf.Clear();
-
+            int lastBlockLeftDataLen = blockBuf.curLen - blockBuf.lastPos;
-static void mtMemCopy(void* data, long idx, int tid) {
+            if (lastBlockLeftDataLen > 0) {
-    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
+                lastRoundBuf.MemCopy(blockBuf.data + blockBuf.lastPos, lastBlockLeftDataLen);
-    tid = idx; // 静态分配任务，此时用idx代替tid
+            }
-
+        }
-    auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr;  // 每个thread一个，用来保存解压后的block数据
-    auto &uncompressData = p.uncompressData;              // 所有线程共用一个，串行往这里添加解压后的block数据
-
-    // 拷贝bam未解析数据到全局的uncompressData里
-    memcpy(uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset, threadUncompressDataArr[tid].blockBuf.data,
-           threadUncompressDataArr[tid].blockBuf.curLen);
-           
-    // 拷贝解析的bam到全局数据里
-    size_t i = 0;
-    for (; i < threadUncompressDataArr[tid].firstBam.Size(); ++i) {
-        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].firstBam.arr[i];
-    }
-    for (size_t j = 0; j < threadUncompressDataArr[tid].bamArr.Size(); ++i, ++j) {
-        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].bamArr.arr[j];
-    }
-
-    if (tid == p.numThread - 1) { // 最后一个线程，更新全局uncompressData的usedBufSize
-        uncompressData.usedBufSize += threadUncompressDataArr[tid].memOffset + threadUncompressDataArr[tid].blockBuf.curLen;
-        uncompressData.lastEndPos = uncompressData.usedBufSize - (threadUncompressDataArr[tid].blockBuf.curLen - threadUncompressDataArr[tid].blockBuf.lastPos);
-        p.allBams.curIdx += threadUncompressDataArr[tid].bamOffset + threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
     }
 }
 
 /* 将gz block进行解压，并进行线程内排序 */
 static void doPhase1Uncompress(Phase1PipelineArg& p, int finish = 0) {
     PROF_G_BEG(uncompress);
-    uint64_t blockNum = p.readData[p.uncompressOrder % p.READ_BUF_NUM].startAddrArr.size();
-    p.blockNum += blockNum;
 
     kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread);
-    PROF_G_END(uncompress);
 
-    // 并行拷贝所有blocks
+    // 处理相邻线程的block数据，可能有bam跨越这两个线程的block（GATK的bam）
-    PROF_G_BEG(mem_copy);
     handleAdjacentThreadBlock(p);
-    p.allBams.Add(p.threadUncompressWrap.GetTotalBamNum());
-    p.bamNum += p.threadUncompressWrap.GetTotalBamNum();
-    kt_for(p.numThread, mtMemCopy, &p, p.numThread);
-    PROF_G_END(mem_copy);
-
-    PROF_G_BEG(parse_block);
-    PROF_G_END(parse_block);
-
-    if (true) { // 缓冲区满了
-        spdlog::info("blocks num: {}, uncompressed: {}, bam num: {}, all bam num: {}, zero start blocks: {}", blockNum, p.blockNum,
-                     p.threadUncompressWrap.GetTotalBamNum(), p.bamNum, p.zeroStartBlockNum);
-        // p.uncompressData.Clear();
-        p.uncompressData.NextRound();
-        // spdlog::info("last data - 0: {}", p.uncompressData.usedBufSize - p.uncompressData.lastEndPos);
-        p.threadUncompressWrap.ResetBlockArr();
-        //for (size_t i = 0; i < p.allBams.Size(); ++i) {
-        //    fprintf(gfp[0], "%d-%ld\n", p.allBams.arr[i].tid, p.allBams.arr[i].pos);
-        //}
-        p.allBams.Clear();
-    }
 
+    PROF_G_END(uncompress);
 }
 
 /* phase1Uncompress step-2 解压线程 */
@@ -323,14 +283,18 @@ void* phase1Uncompress(void* data) {
     while (true) {
         // previous dependency
         yarn::DEPENDENCY_NOT_TO_BE(p.readSig, 0);
+        yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, p.UNCOMPRESS_BUF_NUM);
 
         if (p.readFinish) {
             while (p.uncompressOrder < p.readOrder) {
+                yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, p.UNCOMPRESS_BUF_NUM);
                 doPhase1Uncompress(p, 1);
-                p.uncompressOrder += 1;
+                yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder);
             }
+            yarn::SIGNAL_FINISH(p.uncompressSig, p.uncompressFinish);
             break;
         }
+#if 0
         if (parseFirstBlock) {
             parseFirstBlock = 0;
             // 计算bam的平均长度，以及第一个block里的bam个数，用来指导后续的解压和排序
@@ -364,13 +328,96 @@ void* phase1Uncompress(void* data) {
             p.uncompressData.avgBamNumPerBlock = bamNum;
             spdlog::info("avg bam size: {}, avg bam num per block: {}, max bam len: {}, max seq len: {}", p.uncompressData.avgBamSize, p.uncompressData.avgBamNumPerBlock, p.maxBamLen, p.maxSeqLen);
         }
+#endif
         doPhase1Uncompress(p);
 
         // update status
         yarn::CONSUME_SIGNAL(p.readSig);
-        p.uncompressOrder += 1;
+        yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder);
     }
 
     spdlog::info("uncompress order: {}", p.uncompressOrder);
     return nullptr;
+}
+
+// 多线程内存拷贝，静态分配任务，此时用idx代替tid，multi-thread memory copy uncompressed data to global buffer
+static void mtMemCopy(void* data, long idx, int tid) {
+    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
+    tid = idx;  // 静态分配任务，此时用idx代替tid
+
+    auto& threadUncompressDataArr =
+        p.threadUncompressWrap[p.memCopyOrder % p.UNCOMPRESS_BUF_NUM].threadUncompressDataArr;  // 每个thread一个，用来保存解压后的block数据
+    auto& uncompressData = p.uncompressData;                                                    // 所有线程共用一个，串行往这里添加解压后的block数据
+
+    // 拷贝bam未解析数据到全局的uncompressData里
+    memcpy(uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset, threadUncompressDataArr[tid].blockBuf.data,
+           threadUncompressDataArr[tid].blockBuf.curLen);
+
+    // 拷贝解析的bam到全局数据里
+    size_t i = 0;
+    for (; i < threadUncompressDataArr[tid].firstBam.Size(); ++i) {
+        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].firstBam.arr[i];
+    }
+    for (size_t j = 0; j < threadUncompressDataArr[tid].bamArr.Size(); ++i, ++j) {
+        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].bamArr.arr[j];
+    }
+
+    if (tid == p.numThread - 1) {  // 最后一个线程，更新全局uncompressData的usedBufSize
+        uncompressData.usedBufSize += threadUncompressDataArr[tid].memOffset + threadUncompressDataArr[tid].blockBuf.curLen;
+        uncompressData.lastEndPos =
+            uncompressData.usedBufSize - (threadUncompressDataArr[tid].blockBuf.curLen - threadUncompressDataArr[tid].blockBuf.lastPos);
+        // 更新全局bam数量和偏移
+        p.allBams.curIdx +=
+            threadUncompressDataArr[tid].bamOffset + threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
+    }
+}
+
+static void doMemCopy(Phase1PipelineArg& p) {
+    // 并行拷贝所有blocks
+    PROF_G_BEG(mem_copy);
+
+    auto& uncompressWrap = p.threadUncompressWrap[p.memCopyOrder % p.UNCOMPRESS_BUF_NUM];
+    p.allBams.Add(uncompressWrap.GetTotalBamNum());
+    p.bamNum += uncompressWrap.GetTotalBamNum();
+    p.blockNum += uncompressWrap.GetTotalBlockNum();
+
+    kt_for(p.numThread, mtMemCopy, &p, p.numThread);
+    PROF_G_END(mem_copy);
+
+    if (true) {  // 缓冲区满了
+        spdlog::info("block num: {}, all block num: {}, bam num: {}, all bam num: {}", uncompressWrap.GetTotalBlockNum(), p.blockNum, uncompressWrap.GetTotalBamNum(), p.bamNum);
+        // p.uncompressData.Clear();
+        p.uncompressData.NextRound();
+        // spdlog::info("last data - 0: {}", p.uncompressData.usedBufSize - p.uncompressData.lastEndPos);
+        uncompressWrap.ResetBlockArr();
+        // for (size_t i = 0; i < p.allBams.Size(); ++i) {
+        //     fprintf(gfp[0], "%d-%ld\n", p.allBams.arr[i].tid, p.allBams.arr[i].pos);
+        // }
+        p.allBams.Clear();
+    }
+}
+
+/* phase1Uncompress step-3 拷贝线程 */
+void* phase1MemCopy(void* data) {
+    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
+    /* 2. do the work */
+    while (true) {
+        // previous dependency
+        yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, 0);
+
+        if (p.uncompressFinish) {
+            while (p.memCopyOrder < p.uncompressOrder) {
+                doMemCopy(p);
+                p.memCopyOrder += 1;
+            }
+            break;
+        }
+        doMemCopy(p);
+        // update status
+        yarn::CONSUME_SIGNAL(p.uncompressSig);
+        p.memCopyOrder += 1;
+    }
+
+    spdlog::info("mem copy order: {}", p.memCopyOrder);
+    return nullptr;
 }

src/sort/phase_1_uncompress.h

@@ -10,4 +10,6 @@
 #pragma once
 
 /* phase1Uncompress step-2 解压线程 */
-void* phase1Uncompress(void* data);
+void* phase1Uncompress(void* data);
+
+void* phase1MemCopy(void* data);

src/sort/sam_io.h

@@ -40,9 +40,11 @@ struct DataBuffer {
     }
 
     void MemCopy(uint8_t *src, size_t len) {
-        ReAllocMem(curLen + len);
+        if (len > 0) {
-        memcpy(&data[curLen], src, len);
+            ReAllocMem(curLen + len);
-        curLen += len;
+            memcpy(&data[curLen], src, len);
+            curLen += len;
+        }
     }
 
     void Clear() {

src/sort/sort.cpp

@@ -662,10 +662,10 @@ static void samSortFirstPipe() {
 int doSort() {
 
 #if 1
-    gfp[0] = fopen("f0.txt", "w");
+    gfp[0] = fopen("output/f0.txt", "w");
-    gfp[1] = fopen("f1.txt", "w");
+    gfp[1] = fopen("output/f1.txt", "w");
-    gfp[2] = fopen("f2.txt", "w");
+    gfp[2] = fopen("output/f2.txt", "w");
-    gfp[3] = fopen("f3.txt", "w");
+    gfp[3] = fopen("output/f3.txt", "w");
 #endif
 
 #if 1