zzh
/
FastSort
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

解决了bam和block不对齐情况的处理 

Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
zzh 2026-06-01 16:51:12 +08:00
parent 00286d9899
commit 86b26c9512
11 changed files with 429 additions and 129 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -8,6 +8,7 @@
/build
build.sh
run.sh
/output
# Compiled Object files
*.slo

1
convert.sh 100755
View File

@ -0,0 +1 @@
time samtools view -h -@ 32 ~/mini.bam | samtools view -@ 32 -h -b -o ~/mini-samtools.bam

4
src/sort/phase_1.cpp
View File

@ -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) {

84
src/sort/phase_1.h
View File

@ -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<ThreadBlockArr> threadBlocks; // 每个thread一个，用来保存解压后的block数据
struct ThreadUncompressWrap {
vector<ThreadUncompressData> threadUncompressDataArr; // 每个线程一个，用来保存解压数据和解析bam数据
vector<DataBuffer> 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不相关

342
src/sort/phase_1_uncompress.cpp
View File

@ -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;
}
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;
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;
}
#endif
if (nextBamStart < endOffset) {
lastPos = nextBamStart; // 最后一个不完整的bam，连长度都不够数据解析
}
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

4
src/sort/sam_io.cpp
View File

@ -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

18
src/sort/sam_io.h
View File

@ -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 {

22
src/sort/sort.cpp
View File

@ -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;
}

77
src/sort/sort.h
View File

@ -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<uint8_t *> 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<uint8_t *> 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 <typename T>
struct FastVector {
vector<T> 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<OneBlock> blockArr; // 解压后的数据
int curIdx = 0; // 当前解压数据对应的vector的索引
uint64_t bamNum = 0; // 解压后的bam数量
// 最小堆
std::priority_queue<BlockIdIdx, std::vector<BlockIdIdx>, std::greater<BlockIdIdx>> 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<OneBam> BamArr;
typedef FastVector<OneBlock> BlockArr;
/* 第一阶段的多线程流水线参数 */
struct FirstPipeArg {
static const int READ_BUF_NUM = 2; // 读入的buf数量

2
src/util/profiling.cpp
View File

@ -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) {

3
src/util/profiling.h
View File

@ -2,6 +2,7 @@
#include <stdint.h>
#include <stdlib.h>
#include <sys/time.h>
#include <stdio.h>
// #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