From 86b26c95127ee520a1afa07e9fca82e43f4cea3e Mon Sep 17 00:00:00 2001 From: zzh Date: Mon, 1 Jun 2026 16:51:12 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3=E4=BA=86bam=E5=92=8Cblock?= =?UTF-8?q?=E4=B8=8D=E5=AF=B9=E9=BD=90=E6=83=85=E5=86=B5=E7=9A=84=E5=A4=84?= =?UTF-8?q?=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Copilot --- .gitignore | 1 + convert.sh | 1 + src/sort/phase_1.cpp | 4 +- src/sort/phase_1.h | 84 +++++--- src/sort/phase_1_uncompress.cpp | 342 +++++++++++++++++++++++++------- src/sort/sam_io.cpp | 4 +- src/sort/sam_io.h | 18 +- src/sort/sort.cpp | 22 +- src/sort/sort.h | 77 +++++-- src/util/profiling.cpp | 2 + src/util/profiling.h | 3 + 11 files changed, 429 insertions(+), 129 deletions(-) create mode 100755 convert.sh diff --git a/.gitignore b/.gitignore index 51c06a1..7faad8d 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ /build build.sh run.sh +/output # Compiled Object files *.slo diff --git a/convert.sh b/convert.sh new file mode 100755 index 0000000..9cb2d0a --- /dev/null +++ b/convert.sh @@ -0,0 +1 @@ +time samtools view -h -@ 32 ~/mini.bam | samtools view -@ 32 -h -b -o ~/mini-samtools.bam \ No newline at end of file diff --git a/src/sort/phase_1.cpp b/src/sort/phase_1.cpp index d4c155e..3710f9f 100644 --- a/src/sort/phase_1.cpp +++ b/src/sort/phase_1.cpp @@ -23,10 +23,10 @@ void phase1Pipeline() { /* set up*/ Phase1PipelineArg phase1Arg; phase1Arg.numThread = nsgv::gSortArg.NUM_THREADS; - const size_t kReadBufSize = 1L * 1024 * 1024 * phase1Arg.numThread; // 平均每线程1M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略） + const size_t kReadBufSize = 4L * 1024 * 1024 * phase1Arg.numThread; // 平均每线程4M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略） phase1Arg.uncompressBufBytes = nsgv::gSortArg.MAX_MEM * 0.9; // 比最大内存参数小点 - phase1Arg.threadBlocksWrap.Resize(phase1Arg.numThread); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容 + phase1Arg.threadUncompressWrap.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) { diff --git a/src/sort/phase_1.h b/src/sort/phase_1.h index 764eb12..3d5f2a4 100644 --- a/src/sort/phase_1.h +++ b/src/sort/phase_1.h @@ -24,47 +24,74 @@ #define START_IDX(i, nt, nele) ((i) * (nele) / (nt)) #define STOP_IDX(i, nt, nele) (((i) + 1) * (nele) / (nt)) + +// 把线程解压需要的数据放到一个结构体里 +struct ThreadUncompressData { + DataBuffer blockBuf; // 解压的block放在这里 + size_t blockNum = 0; // 解压的block数量 + BamArr bamArr; // 解析后的bam数据放在这里 + BamArr firstBam; // 连接上一个不完整的bam数据，当作当前block的第一个bam + DataBuffer lastBamBuf; // 最后那个不完整的bam解析时用到的缓冲区 + size_t memOffset = 0; // 这个线程解压的数据在全局解压数据中的偏移位置 + size_t bamOffset = 0; // 这个线程解析的bam数据在全局解压数据中的偏移位置 + + + void Resize(int vecInitSize) { + blockBuf.AllocMem(vecInitSize * SINGLE_BLOCK_SIZE); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容 + bamArr.arr.reserve(vecInitSize); + firstBam.arr.reserve(1); + lastBamBuf.AllocMem(SINGLE_BLOCK_SIZE); + } + + void Reset() { + blockNum = 0; + blockBuf.Clear(); + bamArr.Clear(); + firstBam.Clear(); + lastBamBuf.Clear(); + memOffset = 0; + bamOffset = 0; + } + + size_t GetBlockNum() const { return blockNum; } + size_t GetBamNum() const { return bamArr.Size() + firstBam.Size(); } +}; + // 第一阶段用到的数据结构 // 第一阶段的解压、排序、归并、压缩 -struct UncompressBlocksWrap { - vector threadBlocks; // 每个thread一个，用来保存解压后的block数据 +struct ThreadUncompressWrap { + vector threadUncompressDataArr; // 每个线程一个，用来保存解压数据和解析bam数据 - vector threadBlockBuf; // 每个线程一个，用来保存解压后的block数据，和上边的threadBlocks重复了，暂时保留，后续优化 - - UncompressBlocksWrap() {} - UncompressBlocksWrap(int numThread) { Resize(numThread); } - UncompressBlocksWrap(int numThread, int vecInitSize) { Resize(numThread, vecInitSize); } + ThreadUncompressWrap() {} + ThreadUncompressWrap(int numThread) { Resize(numThread); } + ThreadUncompressWrap(int numThread, int vecInitSize) { Resize(numThread, vecInitSize); } void Resize(int numThread) { Resize(numThread, 128); } void Resize(int numThread, int vecInitSize) { - threadBlocks.resize(numThread); - threadBlockBuf.resize(numThread); - for (int i = 0; i < numThread; ++i) { threadBlocks[i].blockArr.reserve(vecInitSize); } + threadUncompressDataArr.resize(numThread); for (int i = 0; i < numThread; ++i) { - threadBlocks[i].blockArr.reserve(vecInitSize); - threadBlockBuf[i].allocMem(vecInitSize * SINGLE_BLOCK_SIZE); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容 + threadUncompressDataArr[i].Resize(vecInitSize); } } void ResetBlockArr() { - for (int i = 0; i < threadBlocks.size(); ++i) { - threadBlocks[i].clear(); - threadBlockBuf[i].clear(); + for (int i = 0; i < threadUncompressDataArr.size(); ++i) { + threadUncompressDataArr[i].Reset(); } } uint64_t GetTotalBlockNum() { uint64_t totalBlockNum = 0; - for (int i = 0; i < threadBlocks.size(); ++i) { - totalBlockNum += threadBlocks[i].curIdx; + for (int i = 0; i < threadUncompressDataArr.size(); ++i) { + totalBlockNum += threadUncompressDataArr[i].GetBlockNum(); } return totalBlockNum; } - uint64_t GetHeapBlockNum() { - uint64_t heapBlockNum = 0; - for (int i = 0; i < threadBlocks.size(); ++i) { - heapBlockNum += threadBlocks[i].blockHeap.size(); + uint64_t GetTotalBamNum() { + uint64_t totalBamNum = 0; + for (int i = 0; i < threadUncompressDataArr.size(); ++i) { + totalBamNum += threadUncompressDataArr[i].GetBamNum(); } - return heapBlockNum; + return totalBamNum; } }; @@ -77,7 +104,6 @@ struct Phase1PipelineArg { // common parameters int numThread = 0; // 线程数 - uint64_t bamNum = 0; // 解压后的bam数量 uint64_t singleThreadMemBytes = 0; // 单线程开辟的内存字节上限 uint64_t uncompressBufBytes = 0; // 总的解压缓冲区大小 uint64_t startBlockId = 0; // 当前轮次起始block id @@ -85,14 +111,24 @@ struct Phase1PipelineArg { // for read-uncompress-parse uint64_t readOrder = 0; // 读取文件轮次编号，与下边的uncompressOrder对应 uint64_t uncompressOrder = 0; // 并行解压gz block, 包含排序（缓冲区满之后排序），以及合并之后的压缩 + uint64_t memCopyOrder = 0; // 串行拷贝解压数据到uncompressData的轮次编号，与上边的uncompressOrder对应 volatile int readFinish = 0; yarn::lock_t* readSig; yarn::lock_t* uncompressSig; ReadBuffer readData[READ_BUF_NUM]; // 用来读如数据，双缓冲 - UncompressBlocksWrap threadBlocksWrap; // 每个thread一个，用来保存解压后的block数据 + ThreadUncompressWrap threadUncompressWrap; // 每个thread一个，用来保存解压后的block数据 UncompressBlockBuffer uncompressData; // 所有线程共用一个，串行往这里添加解压后的block数据 + BamArr allBams; // 所有线程共用一个，串行往这里添加解析后的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 merge-compress-write uint64_t compressOrder = 0; // 排序后压缩，这个和下边的writeOder对应，跟上边的order不相关 diff --git a/src/sort/phase_1_uncompress.cpp b/src/sort/phase_1_uncompress.cpp index a26f7d8..a79a3af 100644 --- a/src/sort/phase_1_uncompress.cpp +++ b/src/sort/phase_1_uncompress.cpp @@ -22,47 +22,114 @@ #include "util/profiling.h" #include "util/yarn.h" -/* 多线程解压 */ -static void mtUncompressBlock(void* data, long idx, int tid) { - PROF_T_BEG(mem_copy); +int GetBamLen(uint8_t* dataAddr) { + uint32_t bamLen = 0; + memcpy(&bamLen, dataAddr, 4); + if (nsgv::gIsBigEndian) + ed_swap_4p(&bamLen); + return bamLen; +} - Phase1PipelineArg& p = *(Phase1PipelineArg*)data; - ReadBuffer & readData = p.readData[p.uncompressOrder % p.READ_BUF_NUM]; +void ParseBam(uint8_t* dataAddr, OneBam& bam) { + uint32_t bams = 0; + bam.bamLen = GetBamLen(dataAddr); + dataAddr += 4; + bam.tid = le_to_u32(dataAddr); + bam.pos = le_to_i32(dataAddr + 4); + uint32_t x2 = le_to_u32(dataAddr + 8); + bam.qnameLen = x2 & 0xff; +} - auto& blockArr = p.threadBlocksWrap.threadBlocks[tid]; - auto& blockItem = blockArr.add(); - uint8_t* block = readData.startAddrArr[idx]; +// 解析一个bam并放入bamArr +void ParseAddBam(uint8_t* dataAddr, BamArr& bamArr) { + OneBam& bam = bamArr.Add(); + ParseBam(dataAddr, bam); +} - size_t dlen = SINGLE_BLOCK_SIZE; // 65535 - int block_length = unpackInt16(&block[16]) + 1; - uint32_t crc = le_to_u32(block + block_length - 8); - int ret = bgzfUncompress(blockItem.data, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, block_length - BLOCK_HEADER_LENGTH, crc); - if (ret != 0) { - spdlog::error("uncompress error, block id: {}, len: {}, ret: {}", idx, block_length, ret); - exit(0); +// 返回解析bam的个数 +size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, BamArr& bamArr, size_t* nextBamStartPtr = nullptr, + size_t* lastPosPtr = nullptr) { + size_t nextBamStart = startOffset; + size_t lastPos = 0; + uint32_t bamLen = 0; + uint32_t bams = 0; + + while (nextBamStart + 4 <= endOffset) { + uint8_t* curAddr = dataAddr + nextBamStart; + memcpy(&bamLen, curAddr, 4); + if (nsgv::gIsBigEndian) + ed_swap_4p(&bamLen); + nextBamStart += 4 + bamLen; + if (nextBamStart == endOffset) { // 刚好解析到最后，说明这个block的内容都是完整bam + lastPos = endOffset; // 继续解析当前的block + } else if (nextBamStart > endOffset) { // 当前bam不完整，不能解析 + 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; + ++bams; } - blockItem.blockId = idx + p.startBlockId; - blockItem.blockLen = dlen; - blockArr.blockHeap.push({blockItem.blockId, blockArr.curIdx - 1}); // 解压完成后，将block的id和在block数组里的索引加入堆中，方便后续排序和合并 - -#if 0 - // 放入全局缓冲区 - // spdlog::info("top id: {}, block id: {}", blockArr.blockHeap.top().blockId, p.uncompressData.nextBlockId); - while (blockArr.blockHeap.top().blockId == p.uncompressData.nextBlockId) { - auto& top = blockArr.blockHeap.top(); - // auto& topBlock = blockArr.blockArr[top.blockArrIdx]; - // memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, topBlock.data, topBlock.blockLen); - // p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize); - // p.uncompressData.usedBufSize += topBlock.blockLen; - // p.bamNum += topBlock.bamNum; - - blockArr.blockHeap.pop(); - p.uncompressData.nextBlockId += 1; - p.uncompressData.blockNum += 1; + if (nextBamStart < endOffset) { + lastPos = nextBamStart; // 最后一个不完整的bam，连长度都不够数据解析 } -#endif + if (nextBamStartPtr) *nextBamStartPtr = nextBamStart; + if (lastPosPtr) *lastPosPtr = lastPos; + return bams; +} - PROF_T_END(tid, mem_copy); +// 检查当前内存对应的bam是否能正确解析 +static bool isValidBam(uint8_t* dataAddr, uint32_t &bamLen, int maxBamLen, int maxSeqLen) { + memcpy(&bamLen, dataAddr, 4); + if (nsgv::gIsBigEndian) + ed_swap_4p(&bamLen); + int32_t tid = 0; + int64_t pos = 0; + int qnameLen = 0; + uint16_t flag = 0; + uint8_t* x = dataAddr + 4; + uint32_t x2 = le_to_u32(x + 8); + qnameLen = x2 & 0xff; + if (qnameLen >= bamLen) { + return false; + } + int32_t seqLen = le_to_u32(x + 16); + if (seqLen >= bamLen) { + return false; + } + tid = le_to_u32(x); + pos = le_to_i32(x + 4); + uint32_t x3 = le_to_u32(x + 12); + flag = x3 >> 16; + + if ((flag & BAM_FUNMAP) && tid == -1 && pos == -1) { + return true; + } + int nref = sam_hdr_nref(nsgv::gInHdr.header); + if (tid < 0 || tid >= nref) { + // 非法 tid + return false; + } + hts_pos_t ref_len = sam_hdr_tid2len(nsgv::gInHdr.header, tid); + if (pos < 0 || pos >= ref_len) { + // 非法 pos：0 ≤ pos < ref_len + return false; + } + // uint32_t n_cigar = x3 & 0xffff; + if (seqLen * 10 < maxSeqLen || maxSeqLen * 10 < seqLen) { // 猜测条件可以仔细考虑下 + //spdlog::info("invalid bam(seqlen), bamlen: {}, seqLen: {}, maxSeqLen: {}, maxSeqLen: {}", bamLen, seqLen, maxSeqLen, maxSeqLen); + //return false; + } + if (bamLen * 10 < maxBamLen || maxBamLen * 10 < bamLen) { // 猜测条件可以仔细考虑下 + //spdlog::info("invalid bam(bamlen), bamlen: {}, maxBamLen: {}, seqLen: {}, qnameLen: {}, n_cigar: {}", bamLen, maxBamLen, seqLen, qnameLen, n_cigar); + // return false; + } + return true; } // 多线程解压，静态分配任务，此时用idx代替tid，multi-thread uncompress bam blocks @@ -76,12 +143,15 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) { int startIdx = START_IDX(idx, p.numThread, readData.startAddrArr.size()); int stopIdx = STOP_IDX(idx, p.numThread, readData.startAddrArr.size()); - auto &blockBuf = p.threadBlocksWrap.threadBlockBuf[tid]; + auto &blockBuf = p.threadUncompressWrap.threadUncompressDataArr[tid].blockBuf; + auto &bamArr = p.threadUncompressWrap.threadUncompressDataArr[tid].bamArr; + // 开辟足够的内存 if (stopIdx - startIdx > blockBuf.maxLen / SINGLE_BLOCK_SIZE) { - blockBuf.reAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE); + blockBuf.ReAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE); } + // 解压block for (int i = startIdx; i < stopIdx; ++i) { uint8_t* block = readData.startAddrArr[i]; size_t dlen = SINGLE_BLOCK_SIZE; // 65535 @@ -94,71 +164,161 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) { } blockBuf.curLen += dlen; } + + // 用readPos来表示第一个bam起始位置，默认是ngs的bam，如果是三代bam，应该不需要计算这个了，因为三代bam很长 + blockBuf.readPos = 0; + for (int i = 0; i < SINGLE_BLOCK_SIZE; ++i) { + uint32_t nextBamStart = i; // 这个要注意，每次应该要计算一下 + uint32_t bamLen = 0; + uint8_t nextPos = 0; // 是否检查下一个可能的位置 + while (nextBamStart + 4 <= SINGLE_BLOCK_SIZE) { + if (isValidBam(blockBuf.data + nextBamStart, bamLen, p.maxBamLen, p.maxSeqLen)) { + nextBamStart += 4 + bamLen; + } else { + nextPos = 1; + break; + } + } + if (!nextPos) { + blockBuf.readPos = i; + break; + } + } + // 解析bam + ParseAddAllBams(blockBuf.data, blockBuf.readPos, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos); PROF_T_END(tid, mem_copy); } +// 处理相邻线程的block数据，可能有bam跨越这两个线程的block（GATK的bam） +static void handleAdjacentThreadBlock(Phase1PipelineArg& p) { + auto& uncompressData = p.uncompressData; + auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr; + size_t offset = 0; // 当前线程对应的全局数据的起始偏移量 + size_t bamOffset = p.allBams.Size(); // 当前线程解析的bam在全局数据中的偏移量 + + for (int tid = 0; tid < p.numThread; ++tid) { + threadUncompressDataArr[tid].memOffset = offset; + threadUncompressDataArr[tid].bamOffset = bamOffset; + auto& blockBuf = threadUncompressDataArr[tid].blockBuf; + auto& bamArr = threadUncompressDataArr[tid].bamArr; + auto& firstBam = threadUncompressDataArr[tid].firstBam; + auto& lastBamBuf = threadUncompressDataArr[tid].lastBamBuf; + + bool hasLastData = false; // 上一个block里有不完整bam数据 + int lastDataLen = 0; // 上一个block里不完整bam数据的长度 + int leftDataLen = blockBuf.readPos; // 本轮剩余的不完整的bam数据 + + if (tid == 0) { // 第一个线程 + // 检查一下bam的定位是否正确 + hasLastData = uncompressData.usedBufSize != uncompressData.lastEndPos; + if (hasLastData || blockBuf.readPos > 0) { + lastDataLen = uncompressData.usedBufSize - uncompressData.lastEndPos; + leftDataLen = blockBuf.readPos; // 本轮剩余的不完整的bam数据 + lastBamBuf.MemCopy(uncompressData.dataBuf + uncompressData.lastEndPos, lastDataLen); + lastBamBuf.MemCopy(blockBuf.data, leftDataLen); + } + } else { + hasLastData = threadUncompressDataArr[tid - 1].blockBuf.curLen != threadUncompressDataArr[tid - 1].blockBuf.lastPos; + if (hasLastData || blockBuf.readPos > 0) { // 上一轮有遗留数据 + lastDataLen = threadUncompressDataArr[tid - 1].blockBuf.curLen - threadUncompressDataArr[tid - 1].blockBuf.lastPos; + leftDataLen = blockBuf.readPos; // 本轮剩余的不完整的bam数据 + lastBamBuf.MemCopy(threadUncompressDataArr[tid - 1].blockBuf.data + threadUncompressDataArr[tid - 1].blockBuf.lastPos, lastDataLen); + lastBamBuf.MemCopy(blockBuf.data, leftDataLen); + } + } + if (hasLastData || blockBuf.readPos > 0) { + int bamLen = GetBamLen(lastBamBuf.data); + int claculatedBamLen = lastDataLen + leftDataLen - 4; + int additionDataLen = 0; + if (bamLen != claculatedBamLen) { // 猜测错了 + if (lastDataLen + leftDataLen < 4) { // 不够解析bam长度 + lastBamBuf.MemCopy(blockBuf.data + leftDataLen, 4); + leftDataLen += 4; + bamLen = GetBamLen(lastBamBuf.data); // 真正的长度 + } + additionDataLen = bamLen + 4 - (lastDataLen + leftDataLen); // 还缺多少数据能解析出完整的bam + if (additionDataLen > 0) { + lastBamBuf.MemCopy(blockBuf.data + leftDataLen, additionDataLen); + leftDataLen += additionDataLen; + } + // 重新解析 + bamArr.Clear(); + ParseAddAllBams(blockBuf.data, leftDataLen, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos); + spdlog::error("bam len mismatch {}: {}, {}", tid, bamLen, claculatedBamLen); + } + ParseAddBam(lastBamBuf.data, firstBam); + } + offset += threadUncompressDataArr[tid].blockBuf.curLen; + bamOffset += threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size(); + } +} + static void mtMemCopy(void* data, long idx, int tid) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; tid = idx; // 静态分配任务，此时用idx代替tid - uint64_t offset = 0; - for (int i = 0; i < tid; ++i) { - offset += p.threadBlocksWrap.threadBlockBuf[i].curLen; + + 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(); } - memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize + offset, p.threadBlocksWrap.threadBlockBuf[tid].data, - p.threadBlocksWrap.threadBlockBuf[tid].curLen); } /* 将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(); - // kt_for(p.numThread, mtUncompressBlock, &p, blockNum); + p.blockNum += blockNum; + kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread); - // 串行拷贝所有blocks + PROF_G_END(uncompress); + + // 并行拷贝所有blocks PROF_G_BEG(mem_copy); -#if 1 + handleAdjacentThreadBlock(p); + p.allBams.Add(p.threadUncompressWrap.GetTotalBamNum()); + p.bamNum += p.threadUncompressWrap.GetTotalBamNum(); kt_for(p.numThread, mtMemCopy, &p, p.numThread); -#else - for (int i = 0; i < p.numThread; ++i) { - memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, p.threadBlocksWrap.threadBlockBuf[i].data, p.threadBlocksWrap.threadBlockBuf[i].curLen); - // p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize); - p.uncompressData.usedBufSize += p.threadBlocksWrap.threadBlockBuf[i].curLen; - } -#endif - -#if 0 - for (int i = 0; i < 1; ++i) { - auto& blockArr = p.threadBlocksWrap.threadBlocks[i]; - for (int j = 0; j < blockArr.curIdx; ++j) { - auto& blockItem = blockArr.blockArr[j]; - memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, blockItem.data, blockItem.blockLen); - p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize); - p.uncompressData.usedBufSize += blockItem.blockLen; - p.uncompressData.blockNum += 1; - } - } -#endif - PROF_G_END(mem_copy); - p.startBlockId += blockNum; + PROF_G_BEG(parse_block); + PROF_G_END(parse_block); if (true) { // 缓冲区满了 - spdlog::info("blocks num: {}, left: {}, uncompressed: {}", p.threadBlocksWrap.GetTotalBlockNum(), p.threadBlocksWrap.GetHeapBlockNum(), - p.uncompressData.blockNum); - p.uncompressData.Clear(); - p.threadBlocksWrap.ResetBlockArr(); - p.uncompressData.nextBlockId = p.startBlockId; - p.uncompressData.blockNum = 0; + 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 解压线程 */ void* phase1Uncompress(void* data) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; - + int parseFirstBlock = 1; /* 2. do the work */ while (true) { // previous dependency @@ -171,7 +331,39 @@ void* phase1Uncompress(void* data) { } break; } + if (parseFirstBlock) { + parseFirstBlock = 0; + // 计算bam的平均长度，以及第一个block里的bam个数，用来指导后续的解压和排序 + uint8_t* block = p.readData[p.uncompressOrder % p.READ_BUF_NUM].startAddrArr[0]; + size_t dlen = SINGLE_BLOCK_SIZE; // 65535 + int block_length = unpackInt16(&block[16]) + 1; + uint32_t crc = le_to_u32(block + block_length - 8); + uint8_t oneBlock[SINGLE_BLOCK_SIZE]; + int ret = bgzfUncompress(oneBlock, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, block_length - BLOCK_HEADER_LENGTH, crc); + if (ret != 0) { + spdlog::error("First block uncompress error, len: {}, ret: {}", block_length, ret); + exit(0); + } + uint64_t nextBamStart = 0; // 第一个block的起始bam位置 + uint32_t bamLen = 0; + uint64_t allBamLen = 0; + uint64_t bamNum = 0; + /* 解析每个bam */ + while (nextBamStart + 4 <= dlen) { + bamLen = GetBamLen(oneBlock + nextBamStart); + p.maxBamLen = p.maxBamLen < bamLen ? bamLen : p.maxBamLen; + uint8_t* x = oneBlock + nextBamStart + 4; + int32_t seqLen = le_to_u32(x + 16); + p.maxSeqLen = p.maxSeqLen < seqLen ? seqLen : p.maxSeqLen; + nextBamStart += 4 + bamLen; + allBamLen += bamLen; + ++bamNum; + } + p.uncompressData.avgBamSize = bamNum == 0 ? 0 : allBamLen / bamNum; + 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); + } doPhase1Uncompress(p); // update status diff --git a/src/sort/sam_io.cpp b/src/sort/sam_io.cpp index 1795256..08705fe 100644 --- a/src/sort/sam_io.cpp +++ b/src/sort/sam_io.cpp @@ -137,7 +137,7 @@ size_t readUncompressOneBlock(FILE *fpr, uint8_t *fBuf, DataBuffer *uDataPtr) { uint32_t crc = le_to_u32(fBuf + blockLen - 8); size_t newDataSize = uData.maxLen; while (uData.curLen + SINGLE_BLOCK_SIZE > newDataSize) newDataSize *= 2; - uData.reAllocMem(newDataSize); // 需要重新开辟空间 + uData.ReAllocMem(newDataSize); // 需要重新开辟空间 int ret = bgzfUncompress(&uData.data[uData.curLen], &dlen, (Bytef *)fBuf + BLOCK_HEADER_LENGTH, blockLen - BLOCK_HEADER_LENGTH, crc); if (ret < 0) { @@ -159,7 +159,7 @@ void parseSamHeader(FILE *fpr, HeaderBuf &hdrBuf) { int32_t i, nameLen, numNames = 0; header = sam_hdr_init(); // 初始化header - uData.allocMem(kMaxBlockSize); // 初始化解压数据的buffer + uData.AllocMem(kMaxBlockSize); // 初始化解压数据的buffer readUncompressOneBlock(fpr, fBuf, &uData); // 读取第一个gz block // 解析header diff --git a/src/sort/sam_io.h b/src/sort/sam_io.h index 1063a12..4903e26 100644 --- a/src/sort/sam_io.h +++ b/src/sort/sam_io.h @@ -8,6 +8,7 @@ struct DataBuffer { uint8_t *data; size_t readPos = 0; // 当前读取的位置 + size_t lastPos = 0; // 最后一个bam的起始位置 size_t curLen = 0; // 当前使用的空间 size_t maxLen = 0; // 最大空间 @@ -25,19 +26,30 @@ struct DataBuffer { if (data) free(data); } - void allocMem(size_t memSize) { + void AllocMem(size_t memSize) { curLen = 0; maxLen = memSize; data = (uint8_t *)realloc(data, maxLen); } - void reAllocMem(size_t memSize) { + void ReAllocMem(size_t memSize) { if (memSize > maxLen) { maxLen = memSize; data = (uint8_t *)realloc(data, maxLen); } } - void clear() { curLen = 0; readPos = 0; } + + void MemCopy(uint8_t *src, size_t len) { + ReAllocMem(curLen + len); + memcpy(&data[curLen], src, len); + curLen += len; + } + + void Clear() { + curLen = 0; + readPos = 0; + lastPos = 0; + } }; struct HeaderBuf { diff --git a/src/sort/sort.cpp b/src/sort/sort.cpp index 3f5c7ec..53403eb 100644 --- a/src/sort/sort.cpp +++ b/src/sort/sort.cpp @@ -661,10 +661,17 @@ static void samSortFirstPipe() { // 排序的入口函数，entry function int doSort() { +#if 1 + gfp[0] = fopen("f0.txt", "w"); + gfp[1] = fopen("f1.txt", "w"); + gfp[2] = fopen("f2.txt", "w"); + gfp[3] = fopen("f3.txt", "w"); +#endif + #if 1 nsgv::gIsBigEndian = ed_is_big(); // 第一轮排序，外排到多个中间文件 - // bamSortFirstPipe(); + //bamSortFirstPipe(); phase1Pipeline(); @@ -701,6 +708,11 @@ int doSort() { if (bamp->l_data > 1000) { spdlog::info("large record len: {}", bamp->l_data); } + if (bam_num % 10000000 == 0) { + spdlog::info("bam num: {}, max bam len: {}", bam_num, max_bam_len); + } + + // fprintf(gfp[0], "%d-%ld\n", bamp->core.tid, bamp->core.pos); } sam_close(inBamFp); spdlog::info("max record len: {}", max_bam_len); @@ -708,6 +720,12 @@ int doSort() { #endif +#if 1 + fclose(gfp[0]); + fclose(gfp[1]); + fclose(gfp[2]); + fclose(gfp[3]); +#endif + return 0; } - diff --git a/src/sort/sort.h b/src/sort/sort.h index 69a822b..12cb752 100644 --- a/src/sort/sort.h +++ b/src/sort/sort.h @@ -19,10 +19,10 @@ using std::vector; struct ReadBuffer { uint8_t *dataBuf = nullptr; uint8_t *blockBuf = nullptr; // 用来保存上一轮没能完整读取的block - int readBufSize = 0; // 读入的buf大小 + size_t readBufSize = 0; // 读入的buf大小 vector startAddrArr; // 存放每个block的起始地址 ReadBuffer() { } - ReadBuffer(int readBufSize_) { + ReadBuffer(size_t readBufSize_) { readBufSize = readBufSize_; dataBuf = (uint8_t *)malloc(readBufSize); blockBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE); @@ -31,7 +31,7 @@ struct ReadBuffer { if (dataBuf) free(dataBuf); if (blockBuf) free(blockBuf); } - void Resize(int readBufSize_) { + void Resize(size_t readBufSize_) { if (dataBuf) free(dataBuf); if (blockBuf) free(blockBuf); readBufSize = readBufSize_; @@ -43,38 +43,72 @@ struct ReadBuffer { /* 用来串行保存解压后的block数据, 跟ReadBuffer差不多*/ struct UncompressBlockBuffer { uint8_t *dataBuf = nullptr; // 用来保存解压后的block数据，串行往这里添加解压后的block数据 - uint8_t *unCompleteBuf = nullptr; // 用来保存上一轮没能完整解析的bam数据（前一部分，需要后续拷贝到dataBuf里） uint64_t dataBufSize = 0; // 存放的解压之后的block的buf大小 uint64_t usedBufSize = 0; // 已经使用的buf大小 - uint64_t usedUnCompleteBufSize = 0; // 已经使用的unCompleteBuf大小 - int unCompleteBufSize = 0; // unCompleteBuf里数据的字节数，即上一轮不完整的bam的前半部分大小 - vector startAddrArr; // 存放每个bam的起始地址 - uint64_t nextBlockId = 0; // 下一个需要放入的block id，按照顺序放入dataBuf里 + uint64_t lastEndPos = 0; // 上一轮计算完整bam结束位置，也是下一轮第一个bam（不完整）的起始数据 - uint64_t blockNum = 0; // 解压后的block数量 + int avgBamSize = 0; // bam的平均长度，用来指导后续的解压和排序 + int avgBamNumPerBlock = 0; // 每个block里bam的平均数量，用来指导后续的解压和排序 UncompressBlockBuffer() { } UncompressBlockBuffer(uint64_t dataBufSize_) { dataBufSize = dataBufSize_; dataBuf = (uint8_t *)malloc(dataBufSize); - unCompleteBufSize = SINGLE_BLOCK_SIZE; - unCompleteBuf = (uint8_t*)malloc(SINGLE_BLOCK_SIZE); } ~UncompressBlockBuffer() { if (dataBuf) free(dataBuf); - if (unCompleteBuf) free(unCompleteBuf); } void Resize(uint64_t dataBufSize_) { if (dataBuf) free(dataBuf); - if (unCompleteBuf) free(unCompleteBuf); dataBufSize = dataBufSize_; - unCompleteBufSize = SINGLE_BLOCK_SIZE; dataBuf = (uint8_t *)malloc(dataBufSize); - unCompleteBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE); } - void Clear() { usedBufSize = 0; usedUnCompleteBufSize = 0; startAddrArr.clear(); } + void Clear() { + usedBufSize = 0; + lastEndPos = 0; + } + + void NextRound() { + usedBufSize -= lastEndPos; + memcpy(dataBuf, dataBuf + lastEndPos, usedBufSize); + lastEndPos = 0; + } }; +/* 对vector的一个包装，避免频繁内存分配和释放 */ +template +struct FastVector { + vector arr; + size_t curIdx = 0; // 当前已经添加的元素数量，也是下一个要添加的元素的索引 + T& Add() { + if (curIdx < arr.size()) + return arr[curIdx++]; + else { +#if 0 + arr.resize((arr.size() + 1) << 1); + return arr[curIdx++]; + +#else + arr.push_back(T()); + curIdx++; + return arr.back(); +#endif + } + } + void Add(size_t num) { + if (curIdx + num > arr.size()) { + arr.resize(curIdx + num); + } + } + void ReAllocate(size_t num) { + if (num > arr.size()) { + arr.resize(num); + } + } + size_t Size() const { return curIdx; } + size_t Capacity() const { return arr.size(); } + void Clear() { curIdx = 0; } +}; /* */ @@ -94,8 +128,8 @@ struct OneBam { uint16_t bamLen = 0; uint16_t qnameLen; // 序列名字长度 uint32_t offset = 0; // 距离首地址的偏移量 - uint32_t tid = 0; // 比对到的染色体 - uint64_t pos = 0; // mapping 位置 + int32_t tid = 0; // 比对到的染色体 + int64_t pos = 0; // mapping 位置 ThreadBlockArr* blockThread; // for test // bam1_t b; @@ -145,8 +179,6 @@ struct ThreadBlockArr { vector blockArr; // 解压后的数据 int curIdx = 0; // 当前解压数据对应的vector的索引 uint64_t bamNum = 0; // 解压后的bam数量 - // 最小堆 - std::priority_queue, std::greater> blockHeap; OneBlock& add() { if (curIdx < blockArr.size()) @@ -172,7 +204,6 @@ struct ThreadBlockArr { void clear() { curIdx = 0; bamNum = 0; - blockHeap = {}; } }; @@ -355,6 +386,10 @@ struct MergeSortData { } }; +typedef FastVector BamArr; +typedef FastVector BlockArr; + + /* 第一阶段的多线程流水线参数 */ struct FirstPipeArg { static const int READ_BUF_NUM = 2; // 读入的buf数量 diff --git a/src/util/profiling.cpp b/src/util/profiling.cpp index 76a14fa..08f18b7 100644 --- a/src/util/profiling.cpp +++ b/src/util/profiling.cpp @@ -9,6 +9,8 @@ uint64_t tprof[LIM_THREAD_PROF_TYPE][LIM_THREAD] = {0}; uint64_t proc_freq = 1000; uint64_t gprof[LIM_GLOBAL_PROF_TYPE] = {0}; + +FILE* gfp[4] = {NULL, NULL, NULL, NULL}; #endif uint64_t realtimeMsec(void) { diff --git a/src/util/profiling.h b/src/util/profiling.h index a0fc35c..68f3511 100644 --- a/src/util/profiling.h +++ b/src/util/profiling.h @@ -2,6 +2,7 @@ #include #include #include +#include // #define SHOW_PERF @@ -19,6 +20,8 @@ extern "C" { extern uint64_t proc_freq; extern uint64_t tprof[LIM_THREAD_PROF_TYPE][LIM_THREAD]; extern uint64_t gprof[LIM_GLOBAL_PROF_TYPE]; + +extern FILE* gfp[4]; #endif #ifdef SHOW_PERF