diff --git a/src/sort/phase_1.cpp b/src/sort/phase_1.cpp index 6c7e626..860a1c1 100644 --- a/src/sort/phase_1.cpp +++ b/src/sort/phase_1.cpp @@ -17,6 +17,21 @@ #include "phase_1_write.h" #include "util/profiling.h" +std::string getFileNameWithoutExt(const std::string& filepath) { + // 1. 找到最后一个路径分隔符 + size_t sep = filepath.find_last_of("/\\"); + size_t nameStart = (sep == std::string::npos) ? 0 : sep + 1; + + // 2. 找到最后一个点（作为扩展名的分隔符） + size_t dot = filepath.find_last_of('.'); + + // 3. 仅当点位于文件名起始之后时，才认为是扩展名 + if (dot != std::string::npos && dot > nameStart) { + return filepath.substr(nameStart, dot - nameStart); + } + return filepath.substr(nameStart); +} + void phase1Pipeline() { #if 1 @@ -25,7 +40,9 @@ void phase1Pipeline() { phase1Arg.numThread = nsgv::gSortArg.NUM_THREADS; const size_t kReadBufSize = 4L * 1024 * 1024 * phase1Arg.numThread; // 平均每线程4M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略） - phase1Arg.uncompressBufBytes = nsgv::gSortArg.MAX_MEM * 0.9; // 比最大内存参数小点 + phase1Arg.midFileNamePrefix = getFileNameWithoutExt(nsgv::gSortArg.OUTPUT_FILE) + "_sort_mid."; + phase1Arg.maxMemBytes = nsgv::gSortArg.MAX_MEM; + phase1Arg.uncompressBufBytes = nsgv::gSortArg.MAX_MEM; // 比最大内存参数小点 for (int i = 0; i +#include + #include "const_val.h" #include "sam_io.h" -#include "util/yarn.h" #include "sort.h" +#include "util/yarn.h" + +using std::string; #define START_IDX(i, nt, nele) ((i) * (nele) / (nt)) #define STOP_IDX(i, nt, nele) (((i) + 1) * (nele) / (nt)) @@ -95,19 +99,32 @@ struct ThreadUncompressWrap { } }; +// 用于合并压缩的数据结构 +struct MergeCompressData { + vector blockDataArr; // 待压缩的数据 + vector compressDataArr; // 压缩后的数据 + + void Resize(int blockNum) { + blockDataArr.resize(blockNum); + compressDataArr.resize(blockNum); + } +}; + /* 第一阶段的多线程流水线参数 */ 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数量，1表示解压和拷贝是串行，2，表示并行 - static const int COMPRESS_BUF_NUM = 1; // 压缩的buf数量, 只有一个 - static const int WRITE_BUF_NUM = 2; // 写入文件buf数量 + static const int MERGE_BUF_NUM = 2; // 归并的buf数量 + static const int COMPRESS_BUF_NUM = 2; // 压缩的buf数量 // common parameters int numThread = 0; // 线程数 uint64_t singleThreadMemBytes = 0; // 单线程开辟的内存字节上限 uint64_t uncompressBufBytes = 0; // 总的解压缓冲区大小 uint64_t startBlockId = 0; // 当前轮次起始block id + uint64_t maxMemBytes = 0; // 最大可使用内存，参数 + int compressLevel = 1; // 压缩级别，参数 // for read-uncompress-parse uint64_t readOrder = 0; // 读取文件轮次编号，与下边的uncompressOrder对应 @@ -127,27 +144,47 @@ struct Phase1PipelineArg { UncompressBlockBuffer uncompressData; // 所有线程共用一个，串行往这里添加解压后的block数据 BamArr allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据 - // 判断bam是否有效的阈值 + // 判断bam是否有效的阈值，暂时没用了 int maxSeqLen = 0; // bam里seq的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件 int maxBamLen = 0; // bam的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件 + // 主要用来统计信息，测试用 uint64_t bamNum = 0; // 解压后的bam数量 uint64_t blockNum = 0; // 解压后的block数量 - int zeroStartBlockNum = 0; + ///////////////////////////////////////////////////// for sort-merge-compress-write // for merge-compress-write + int midFileOrder = 0; // 用来给中间文件编号 + int mergeBlocksThreshold = 0; // 进行归并排序的，单次block数量阈值，参数 + + uint64_t mergeOrder = 0; uint64_t compressOrder = 0; // 排序后压缩，这个和下边的writeOder对应，跟上边的order不相关 uint64_t writeOrder = 0; // 串行合并解压后的blocks，并解析每个bam的长度，达到内存阈值后，并行排序 + volatile int mergeFinish = 0; volatile int compressFinish = 0; + + yarn::lock_t* mergeSig; yarn::lock_t* compressSig; - yarn::lock_t* writeSig; + + MergeCompressData mergeCompressData[MERGE_BUF_NUM]; // 归并压缩用到的数据结构 + DataBuffer compressBuf[COMPRESS_BUF_NUM]; // 压缩用到的buf + string midFileNamePrefix; // 中间文件名前缀 + FILE* midFilePtr = nullptr; // 中间文件指针 Phase1PipelineArg() { readSig = yarn::NEW_LOCK(0); uncompressSig = yarn::NEW_LOCK(0); + mergeSig = yarn::NEW_LOCK(0); + compressSig = yarn::NEW_LOCK(0); lastRoundBuf.AllocMem(SINGLE_BLOCK_SIZE); + + mergeBlocksThreshold = 1000; // 大概64M的解压缩后的数据 + for(int i=0; i +#include + +#include +#include + +#include "common_data.h" +#include "const_val.h" +#include "phase_1.h" +#include "phase_1_compress.h" +#include "phase_1_write.h" +#include "sort.h" +#include "util/profiling.h" + +static void mtCompressBlock(void* data, long idx, int tid) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + MergeCompressData& mergeCompressData = p.mergeCompressData[p.compressOrder % p.COMPRESS_BUF_NUM]; + auto& blockData = mergeCompressData.blockDataArr[idx]; + auto& compressData = mergeCompressData.compressDataArr[idx]; + + // bgzfCompress(void* _dst, size_t* dlen, const void* src, size_t slen, int level) + compressData.ReAllocMem(SINGLE_BLOCK_SIZE); // 压缩后的block数据不会超过单个block的大小 + compressData.curLen = SINGLE_BLOCK_SIZE; + // spdlog::info("block bytes: {}", blockData.curLen); + bgzfCompress(compressData.data, &compressData.curLen, blockData.data, blockData.curLen, p.compressLevel); + //spdlog::info("block bytes: {}, compressed bytes: {}", blockData.curLen, compressData.curLen); +} + +static void doCompress(Phase1PipelineArg& p) { + PROF_G_BEG(compress); + DataBuffer& compressBuf = p.compressBuf[p.compressOrder % p.COMPRESS_BUF_NUM]; + MergeCompressData& mergeCompressData = p.mergeCompressData[p.compressOrder % p.MERGE_BUF_NUM]; + kt_for(p.numThread, mtCompressBlock, &p, mergeCompressData.blockDataArr.size()); + //kt_for(1, mtCompressBlock, &p, mergeCompressData.blockDataArr.size()); + + compressBuf.Clear(); + for (int i=0; i < mergeCompressData.blockDataArr.size(); ++i) { + compressBuf.MemCopy(mergeCompressData.compressDataArr[i].data, mergeCompressData.compressDataArr[i].curLen); + } + spdlog::info("compress bytes: {}", compressBuf.curLen); + PROF_G_END(compress); +} + +/* phase1Compress step- 压缩线程 */ +void* phase1Compress(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, 0); + yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM); + + if (p.mergeFinish) { + while (p.compressOrder < p.mergeOrder) { + yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM); + doCompress(p); + yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder); + } + yarn::SIGNAL_FINISH(p.compressSig, p.compressFinish); + break; + } + + doCompress(p); + + // update status + yarn::CONSUME_SIGNAL(p.mergeSig); + yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder); + } + + spdlog::info("End compress order: {}", p.compressOrder); + return nullptr; +} diff --git a/src/sort/phase_1_compress.h b/src/sort/phase_1_compress.h new file mode 100644 index 0000000..69b9e8b --- /dev/null +++ b/src/sort/phase_1_compress.h @@ -0,0 +1,12 @@ +/* + Description: 第一阶段的压缩 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/02 +*/ + +#pragma once + +void* phase1Compress(void* data); \ No newline at end of file diff --git a/src/sort/phase_1_sort.cpp b/src/sort/phase_1_sort.cpp new file mode 100644 index 0000000..c3e14fb --- /dev/null +++ b/src/sort/phase_1_sort.cpp @@ -0,0 +1,226 @@ +/* + Description: 第一阶段的排序 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/02 +*/ + +#include "phase_1_sort.h" + +#include +#include + +#include +#include + +#include "common_data.h" +#include "const_val.h" +#include "phase_1.h" +#include "phase_1_compress.h" +#include "phase_1_write.h" +#include "sort.h" +#include "util/profiling.h" + +/* bam 排序堆 */ +struct Phase1BamArrIdIdx { + size_t idx = 0; // 下一个待读入数据的idx + size_t endIdx = 0; + const OneBam* bam = nullptr; +}; + +struct Phase1BamPosGreaterThan { + bool operator()(const Phase1BamArrIdIdx& a, const Phase1BamArrIdIdx& b) const { return a.bam->pos > b.bam->pos; } +}; + +static uint8_t* gUncompressDataBuf = nullptr; // 全局的解压数据buf，排序的时候用来比较bam的坐标或者名字信息 +struct Phase1BamNameGreaterThan { + bool operator()(const Phase1BamArrIdIdx& a, const Phase1BamArrIdIdx& b) const { + int cmp = strncmp((char*)(gUncompressDataBuf + a.bam->offset + OneBam::QnameOffset), (char*)(gUncompressDataBuf + b.bam->offset + OneBam::QnameOffset), + std::min(a.bam->qnameLen, b.bam->qnameLen)); + if (cmp == 0) + return a.bam->qnameLen > b.bam->qnameLen; + return cmp > 0; + } +}; + +/* 用来排序 bam*/ +template +struct Phase1BamHeap { + Phase1PipelineArg* p; + priority_queue, GreaterThan> minHeap; + size_t popNum = 0; + + int Init(Phase1PipelineArg* _p) { + p = _p; + if (p == nullptr) { + return -1; + } + + for (int i = 0; i < p->numThread; ++i) { + size_t startIdx = START_IDX(i, p->numThread, p->allBams.Size()); + size_t stopIdx = STOP_IDX(i, p->numThread, p->allBams.Size()); + if (startIdx < stopIdx) { + minHeap.push({startIdx, stopIdx, &p->allBams.arr[startIdx]}); + } + } + return 0; + } + + const OneBam* Pop() { + const OneBam* ret = nullptr; + if (!minHeap.empty()) { + auto minVal = minHeap.top(); + minHeap.pop(); + ++popNum; + ret = minVal.bam; + if (minVal.idx + 1 < minVal.endIdx) { + minVal.idx++; + minVal.bam = &p->allBams.arr[minVal.idx]; + minHeap.push(minVal); + } + } + return ret; + } + + const OneBam* Top() { + const OneBam* ret = nullptr; + if (!minHeap.empty()) { + ret = minHeap.top().bam; + } + return ret; + } +}; + +// 线程内排序 +static void mtInThreadSort(void* data, long idx, int tid) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + tid = idx; + + // 先按照坐标排序 + int startIdx = START_IDX(tid, p.numThread, p.allBams.Size()); + int stopIdx = STOP_IDX(tid, p.numThread, p.allBams.Size()); + + if (nsgv::gSortArg.SORT_COORIDINATE) { + std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [](const OneBam& b1, const OneBam& b2) { + if (b1.tid == -1) + return false; + return b1.tid < b2.tid || (b1.tid == b2.tid && b1.pos < b2.pos) || (b1.tid == b2.tid && b1.pos == b2.pos && b1.qnameLen < b2.qnameLen); + }); + } else if (nsgv::gSortArg.QUERY_NAME_TYPE == nsmd::QueryNameType::PICARD) { + std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [&](const OneBam& b1, const OneBam& b2) { + int cmp = strncmp((char*)(p.uncompressData.dataBuf + b1.offset + OneBam::QnameOffset), + (char*)(p.uncompressData.dataBuf + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen)); + if (cmp == 0) + return b1.qnameLen < b2.qnameLen; + return cmp < 0; + }); + } +} + +void phase1Sort(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + PROF_G_BEG(sort); + kt_for(p.numThread, mtInThreadSort, &p, p.numThread); + PROF_G_END(sort); + + // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].tid, p.allBams.arr[1].tid, p.allBams.arr[2].tid, p.allBams.arr[3].tid, + // p.allBams.arr[4].tid); + // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].pos, p.allBams.arr[1].pos, p.allBams.arr[2].pos, p.allBams.arr[3].pos, + // p.allBams.arr[4].pos); + // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].offset, p.allBams.arr[1].offset, p.allBams.arr[2].offset, p.allBams.arr[3].offset, + // p.allBams.arr[4].offset); +} + +template +bool doMergeSort(Phase1PipelineArg& p, Phase1BamHeap& heap) { + auto &mergeCompressData = p.mergeCompressData[p.mergeOrder % p.MERGE_BUF_NUM]; + + bool finish = false; + const OneBam* bam = nullptr; + + int mergedBlockNum = 0; + size_t bamBytes = 0; + int singleBlockBytes = 0xff00; + mergeCompressData.blockDataArr[mergedBlockNum].Clear(); + while ((bam = heap.Top()) != nullptr) { + if (bamBytes + bam->wholeBamLen > singleBlockBytes) { + mergedBlockNum++; + if (mergedBlockNum >= p.mergeBlocksThreshold) { + break; + } + mergeCompressData.blockDataArr[mergedBlockNum].Clear(); // 清理，为添加bam数据做准备 + bamBytes = 0; + } + mergeCompressData.blockDataArr[mergedBlockNum].MemCopy(p.uncompressData.dataBuf + bam->offset, bam->wholeBamLen); + bamBytes += bam->wholeBamLen; // for test + heap.Pop(); + } + if (bam == nullptr) { + // 都处理完 + finish = true; + } + spdlog::info("mergedBlockNum: {}, bamBytes: {}", mergedBlockNum, bamBytes); + + return finish; +} + +void* phase1MergeSort(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + gUncompressDataBuf = p.uncompressData.dataBuf; // 主要用来比较bam的名字信息 + + // merge + Phase1BamHeap heap; + heap.Init(&p); + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM); + PROF_G_BEG(merge); + bool finish = doMergeSort(p, heap); + PROF_G_END(merge); + + if (finish) { + yarn::SIGNAL_FINISH(p.mergeSig, p.mergeFinish); + break; + } + + // update self status + yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder); + } + + spdlog::info("End merge sort order: {}", p.mergeOrder); + return nullptr; +} + +void ReInitMergePhase(Phase1PipelineArg& p) { + p.mergeOrder = 0; + p.compressOrder = 0; + p.writeOrder = 0; + p.mergeFinish = 0; + p.compressFinish = 0; + yarn::INIT_SIG(p.mergeSig); + yarn::INIT_SIG(p.compressSig); +} + +void phase1MergeCompress(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + ReInitMergePhase(p); + // 创建中间文件 + p.midFilePtr = fopen((p.midFileNamePrefix + std::to_string(p.midFileOrder)).c_str(), "wb"); + if (p.midFilePtr == nullptr) { + spdlog::error("failed to create mid file: {}", p.midFileNamePrefix + std::to_string(p.midFileOrder)); + exit(1); + } + + // 2-stage pipeline, merge-compress-write + pthread_t tidArr[3]; // 2-stage pipeline + pthread_create(&tidArr[0], NULL, phase1MergeSort, &p); + pthread_create(&tidArr[1], NULL, phase1Compress, &p); + pthread_create(&tidArr[2], NULL, phase1Write, &p); + + for (int i = 0; i < 3; ++i) pthread_join(tidArr[i], NULL); + p.midFileOrder += 1; // 写完一次中间文件 +1 + fclose(p.midFilePtr); +} \ No newline at end of file diff --git a/src/sort/phase_1_sort.h b/src/sort/phase_1_sort.h new file mode 100644 index 0000000..2d5a183 --- /dev/null +++ b/src/sort/phase_1_sort.h @@ -0,0 +1,15 @@ +/* + Description: 第一阶段的排序 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/02 +*/ + +#pragma once + +void phase1Sort(void* data); + + +void phase1MergeCompress(void* data); \ No newline at end of file diff --git a/src/sort/phase_1_uncompress.cpp b/src/sort/phase_1_uncompress.cpp index 403b5d7..209f94f 100644 --- a/src/sort/phase_1_uncompress.cpp +++ b/src/sort/phase_1_uncompress.cpp @@ -17,6 +17,7 @@ #include "common_data.h" #include "const_val.h" #include "phase_1.h" +#include "phase_1_sort.h" #include "sam_io.h" #include "sort.h" #include "util/profiling.h" @@ -31,8 +32,7 @@ int GetBamLen(uint8_t* dataAddr) { } void ParseBam(uint8_t* dataAddr, OneBam& bam) { - uint32_t bams = 0; - bam.bamLen = GetBamLen(dataAddr); + bam.wholeBamLen = GetBamLen(dataAddr) + 4; // bamLen不包含4字节的bam长度信息，所以要加上 dataAddr += 4; bam.tid = le_to_u32(dataAddr); bam.pos = le_to_i32(dataAddr + 4); @@ -41,9 +41,10 @@ void ParseBam(uint8_t* dataAddr, OneBam& bam) { } // 解析一个bam并放入bamArr -void ParseAddBam(uint8_t* dataAddr, BamArr& bamArr) { +void ParseAddBam(uint8_t* dataAddr, BamArr& bamArr, uint64_t offset = 0) { OneBam& bam = bamArr.Add(); ParseBam(dataAddr, bam); + bam.offset = offset; } // 返回解析bam的个数 @@ -56,9 +57,7 @@ size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, while (nextBamStart + 4 <= endOffset) { uint8_t* curAddr = dataAddr + nextBamStart; - memcpy(&bamLen, curAddr, 4); - if (nsgv::gIsBigEndian) - ed_swap_4p(&bamLen); + bamLen = GetBamLen(curAddr); nextBamStart += 4 + bamLen; if (nextBamStart == endOffset) { // 刚好解析到最后，说明这个block的内容都是完整bam lastPos = endOffset; // 继续解析当前的block @@ -66,13 +65,7 @@ size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, lastPos = nextBamStart - (4 + bamLen); // 记录最后一个不完整的bam起始位置 break; } - OneBam& bam = bamArr.Add(); - bam.bamLen = bamLen; - curAddr += 4; - bam.tid = le_to_u32(curAddr); - bam.pos = le_to_i32(curAddr + 4); - uint32_t x2 = le_to_u32(curAddr + 8); - bam.qnameLen = x2 & 0xff; + ParseAddBam(curAddr, bamArr, nextBamStart - (4 + bamLen)); ++bams; } if (nextBamStart < endOffset) { @@ -141,8 +134,8 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) { ThreadUncompressWrap& uncompressWrap = p.threadUncompressWrap[p.uncompressOrder % p.UNCOMPRESS_BUF_NUM]; tid = idx; // 静态分配任务，此时用idx代替tid - int startIdx = START_IDX(idx, p.numThread, readData.startAddrArr.size()); - int stopIdx = STOP_IDX(idx, p.numThread, readData.startAddrArr.size()); + int startIdx = START_IDX(tid, p.numThread, readData.startAddrArr.size()); + int stopIdx = STOP_IDX(tid, p.numThread, readData.startAddrArr.size()); auto& blockBuf = uncompressWrap.threadUncompressDataArr[tid].blockBuf; auto& bamArr = uncompressWrap.threadUncompressDataArr[tid].bamArr; @@ -196,7 +189,7 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) { auto& uncompressData = p.uncompressData; auto& threadUncompressDataArr = p.threadUncompressWrap[p.uncompressOrder % p.UNCOMPRESS_BUF_NUM].threadUncompressDataArr; size_t offset = 0; // 当前线程对应的全局数据的起始偏移量 - size_t bamOffset = p.allBams.Size(); // 当前线程解析的bam在全局数据中的偏移量 + size_t bamOffset = 0; // 当前线程解析的bam在全局数据中的偏移量 for (int tid = 0; tid < p.numThread; ++tid) { threadUncompressDataArr[tid].memOffset = offset; @@ -247,7 +240,7 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) { // 重新解析 bamArr.Clear(); ParseAddAllBams(blockBuf.data, leftDataLen, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos); - spdlog::error("bam len mismatch {}: {}, {}", tid, bamLen, claculatedBamLen); + spdlog::error("bam len mismatch {}: {}, {}, addition len: {}", tid, bamLen, claculatedBamLen, additionDataLen); } ParseAddBam(lastBamBuf.data, firstBam); } @@ -354,47 +347,66 @@ static void mtMemCopy(void* data, long idx, int tid) { 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]; + size_t i = p.allBams.Size(); // 当前线程解析的bam在全局数据中的起始偏移量, 这时curIdx还没更新 + uint64_t startOffset = uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset; // 当前线程解析的bam在全局数据中的起始偏移量 + for (size_t j = 0; j < threadUncompressDataArr[tid].firstBam.Size(); ++i, ++j) { + p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].firstBam.arr[j]; + p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset].offset += startOffset; // 更新bam的偏移量 } 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(); + p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset].offset += startOffset; // 更新bam的偏移量 } } static void doMemCopy(Phase1PipelineArg& p) { // 并行拷贝所有blocks - PROF_G_BEG(mem_copy); auto& uncompressWrap = p.threadUncompressWrap[p.memCopyOrder % p.UNCOMPRESS_BUF_NUM]; + auto& lastThreadData = uncompressWrap.threadUncompressDataArr[p.numThread - 1]; + size_t newDataLen = lastThreadData.memOffset + lastThreadData.blockBuf.curLen; + + // 判断缓存是否已满 + size_t totalDataLen = p.uncompressData.usedBufSize + p.allBams.Size() * sizeof(OneBam) + newDataLen; + if (totalDataLen > p.maxMemBytes) { + spdlog::info("block num: {}, all block num: {}, bam num: {}, all bam num: {}", uncompressWrap.GetTotalBlockNum(), p.blockNum, + uncompressWrap.GetTotalBamNum(), p.bamNum); + + // 测试一下bam的offset对不对 +#if 0 + auto &b = p.allBams.arr[p.allBams.Size() - 1]; // 最后一个bam + OneBam o; + ParseBam(p.uncompressData.dataBuf + b.offset, o); + spdlog::info("{}:{}, {}:{}, {}:{}, {}", b.tid, o.tid, b.pos, o.pos, b.qnameLen, o.qnameLen, (char*)(p.uncompressData.dataBuf + b.offset + b.QnameOffset)); +#endif + + // 开启排序并写入中间文件 + phase1Sort(&p); + phase1MergeCompress(&p); + + p.uncompressData.NextRound(); + p.allBams.Clear(); + } + + PROF_G_BEG(mem_copy); + 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); + // 更新全局bam数量和偏移 - 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(); - } + p.allBams.curIdx += lastThreadData.bamOffset + lastThreadData.bamArr.Size() + + lastThreadData.firstBam.Size(); + p.uncompressData.usedBufSize += newDataLen; + p.uncompressData.lastEndPos = + p.uncompressData.usedBufSize - (lastThreadData.blockBuf.curLen - lastThreadData.blockBuf.lastPos); + + // spdlog::info("usedBuf: {}, lastEndPos: {}", p.uncompressData.usedBufSize, p.uncompressData.lastEndPos); + + uncompressWrap.ResetBlockArr(); // 清空这一轮的线程相关数据 + PROF_G_END(mem_copy); } /* phase1Uncompress step-3 拷贝线程 */ @@ -410,6 +422,10 @@ void* phase1MemCopy(void* data) { doMemCopy(p); p.memCopyOrder += 1; } + + // 这里需要再检查一次缓冲区，有数据的话需要处理，只需要线程内排序，再线程间归并排序，不需要压缩写入中间文件了 + // todo + break; } doMemCopy(p); diff --git a/src/sort/phase_1_write.cpp b/src/sort/phase_1_write.cpp index e69de29..eed7cf4 100644 --- a/src/sort/phase_1_write.cpp +++ b/src/sort/phase_1_write.cpp @@ -0,0 +1,53 @@ +/* + Description: 第一阶段写入中间文件 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/02 +*/ + +#include "phase_1_write.h" + +#include +#include + +#include +#include + +#include "common_data.h" +#include "const_val.h" +#include "phase_1.h" +#include "phase_1_compress.h" +#include "phase_1_write.h" +#include "sort.h" +#include "util/profiling.h" + +static void doWrite(Phase1PipelineArg& p) { + DataBuffer& compressBuf = p.compressBuf[p.compressOrder % p.COMPRESS_BUF_NUM]; + fwrite(compressBuf.data, 1, compressBuf.curLen, p.midFilePtr); +} + +void* phase1Write(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, 0); + + if (p.compressFinish) { + while (p.writeOrder < p.compressOrder) { + doWrite(p); + p.writeOrder += 1; + } + break; + } + doWrite(p); + // update status + yarn::CONSUME_SIGNAL(p.compressSig); + p.writeOrder += 1; + } + + spdlog::info("End write order: {}", p.writeOrder); + return nullptr; +} \ No newline at end of file diff --git a/src/sort/phase_1_write.h b/src/sort/phase_1_write.h index e69de29..b5496c7 100644 --- a/src/sort/phase_1_write.h +++ b/src/sort/phase_1_write.h @@ -0,0 +1,11 @@ +/* + Description: 第一阶段写入中间文件 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/02 +*/ +#pragma once + +void* phase1Write(void* data); \ No newline at end of file diff --git a/src/sort/sam_io.cpp b/src/sort/sam_io.cpp index 08705fe..2efc87d 100644 --- a/src/sort/sam_io.cpp +++ b/src/sort/sam_io.cpp @@ -40,7 +40,9 @@ static inline void packInt32(uint8_t* buffer, uint32_t value) { buffer[3] = value >> 24; } +// dlen本身长度一定得够存放压缩后的数据，否则会返回错误 int bgzfCompress(void* _dst, size_t* dlen, const void* src, size_t slen, int level) { + // fprintf(stderr, "bgzfCompress, slen: %zu\n", slen); if (slen == 0) { // EOF block if (*dlen < 28) diff --git a/src/sort/sort.cpp b/src/sort/sort.cpp index 150c351..3c735e4 100644 --- a/src/sort/sort.cpp +++ b/src/sort/sort.cpp @@ -240,8 +240,8 @@ static void mtUncompressBlock(void *data, long idx, int tid) { while (nextBamStart + 4 <= blockItem.blockLen) { OneBam& bam = bamItemArr.add(); - bam.blockThread = &blockItemArr; - bam.blockIdx = blockItemArr.curIdx - 1; + //bam.blockThread = &blockItemArr; + //bam.blockIdx = blockItemArr.curIdx - 1; bam.offset = nextBamStart; uint8_t *curAddr = &blockItem.data[nextBamStart]; #if 0 @@ -251,7 +251,7 @@ static void mtUncompressBlock(void *data, long idx, int tid) { memcpy(&bamLen, curAddr, 4); curAddr += 4; if (nsgv::gIsBigEndian) ed_swap_4p(&bamLen); - bam.bamLen = bamLen; + //bam.bamLen = bamLen; bam.tid = le_to_u32(curAddr); bam.pos = le_to_i32(curAddr + 4); uint32_t x2 = le_to_u32(curAddr + 8); @@ -325,7 +325,7 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) { curAddr += 4; if (nsgv::gIsBigEndian) ed_swap_4p(&bamLen); - bam.bamLen = bamLen; + // bam.bamLen = bamLen; bam.tid = le_to_u32(curAddr); bam.pos = le_to_i32(curAddr + 4); uint32_t x2 = le_to_u32(curAddr + 8); @@ -371,8 +371,8 @@ static void mtInThreadSort(void* data, long idx, int tid) { } else if (nsgv::gSortArg.QUERY_NAME_TYPE == nsmd::QueryNameType::PICARD) { std::sort(arr.begin(), arr.begin() + bamItemArr.curIdx, [](const OneBam& b1, const OneBam& b2) { int cmp = 0; - strncmp((char*)(b1.blockThread->blockArr[b1.blockIdx].data + b1.offset + OneBam::QnameOffset), - (char*)(b2.blockThread->blockArr[b2.blockIdx].data + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen)); +// strncmp((char*)(b1.blockThread->blockArr[b1.blockIdx].data + b1.offset + OneBam::QnameOffset), +// (char*)(b2.blockThread->blockArr[b2.blockIdx].data + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen)); if (cmp == 0) return b1.qnameLen < b2.qnameLen; return cmp < 0; @@ -399,8 +399,8 @@ static void mtCompressBlock(void* data, long idx, int tid) { int curAddr = 0; for (int i = t.idx; i < t.idx + t.num; ++i) { // copy data - memccpy(&block[curAddr], barr[i]->blockThread->blockArr[barr[i]->blockIdx].data + barr[i]->offset, 1, barr[i]->bamLen); - curAddr += barr[i]->bamLen; + //memccpy(&block[curAddr], barr[i]->blockThread->blockArr[barr[i]->blockIdx].data + barr[i]->offset, 1, barr[i]->bamLen); + //curAddr += barr[i]->bamLen; } size_t dlen = BGZF_MAX_BLOCK_SIZE; diff --git a/src/sort/sort.h b/src/sort/sort.h index 12cb752..52d3da4 100644 --- a/src/sort/sort.h +++ b/src/sort/sort.h @@ -123,16 +123,12 @@ struct OneBam { // char *qnameAddr = 0; // qname的地址 // uint64_t pos = 0; // mapping 位置 // uint8_t *addr = 0; // 地址 - static constexpr int QnameOffset = 36; // 距离该sam记录开头地址的偏移量 - int blockIdx = 0; // 在哪个block里 - uint16_t bamLen = 0; - uint16_t qnameLen; // 序列名字长度 - uint32_t offset = 0; // 距离首地址的偏移量 + static constexpr int QnameOffset = 36; // 距离该sam记录开头地址的偏移量，包含了4字节bam长度，32字节的bam core信息，后面紧跟着qname字符串 + uint32_t qnameLen; // 序列名字长度 + uint32_t wholeBamLen = 0; // 包含4字节bam长度的所有bam内容长度，uint16只适合二代 int32_t tid = 0; // 比对到的染色体 int64_t pos = 0; // mapping 位置 - ThreadBlockArr* blockThread; - // for test - // bam1_t b; + uint64_t offset = 0; // 距离首地址的偏移量 }; diff --git a/src/sort/sort_impl.cpp b/src/sort/sort_impl.cpp index 44c8aaa..884cf8c 100644 --- a/src/sort/sort_impl.cpp +++ b/src/sort/sort_impl.cpp @@ -24,8 +24,8 @@ int ks_radixsort(size_t n, OneBam* buf, const sam_hdr_t* h, int tid) { // std::sort(buf, buf + n, [](const OneBam& b1, const OneBam& b2) { return b1.pos < b2.pos; }); std::sort(buf, buf + n, [](const OneBam& b1, const OneBam& b2) { int cmp = 0; - strncmp((char*)(b1.blockThread->blockArr[b1.blockIdx].data + b1.offset + OneBam::QnameOffset), - (char*)(b2.blockThread->blockArr[b2.blockIdx].data + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen)); +// strncmp((char*)(b1.blockThread->blockArr[b1.blockIdx].data + b1.offset + OneBam::QnameOffset), +// (char*)(b2.blockThread->blockArr[b2.blockIdx].data + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen)); if (cmp == 0) return b1.qnameLen < b2.qnameLen; return cmp < 0; diff --git a/src/util/yarn.h b/src/util/yarn.h index cac9274..094925c 100644 --- a/src/util/yarn.h +++ b/src/util/yarn.h @@ -166,4 +166,8 @@ void free_lock_(lock_t *, char const *, long); possess_(sig, __FILE__, __LINE__); \ twist_(sig, yarn::BY, -1, __FILE__, __LINE__); +#define INIT_SIG(sig) \ + possess_(sig, __FILE__, __LINE__); \ + twist_(sig, yarn::TO, 0, __FILE__, __LINE__); + }; // namespace yarn