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

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

Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
zzh 2026-06-01 23:29:14 +08:00
parent 86b26c9512
commit c723476a0f
6 changed files with 135 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/sort/phase_1.cpp
View File

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

11
src/sort/phase_1.h
View File

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

169
src/sort/phase_1_uncompress.cpp
View File

@ -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 &bamArr = p.threadUncompressWrap.threadUncompressDataArr[tid].bamArr;
auto& blockBuf = uncompressWrap.threadUncompressDataArr[tid].blockBuf;
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();
}
}
static void mtMemCopy(void* data, long idx, int tid) {
Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
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();
if (tid == p.numThread - 1) { // 最后一个线程
lastRoundBuf.Clear();
int lastBlockLeftDataLen = blockBuf.curLen - blockBuf.lastPos;
if (lastBlockLeftDataLen > 0) {
lastRoundBuf.MemCopy(blockBuf.data + blockBuf.lastPos, lastBlockLeftDataLen);
}
}
}
}
/* 将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
PROF_G_BEG(mem_copy);
// 处理相邻线程的block数据，可能有bam跨越这两个线程的block（GATK的bam）
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;
}

2
src/sort/phase_1_uncompress.h
View File

@ -11,3 +11,5 @@
/* phase1Uncompress step-2 解压线程 */
void* phase1Uncompress(void* data);
void* phase1MemCopy(void* data);

8
src/sort/sam_io.h
View File

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

8
src/sort/sort.cpp
View File

@ -662,10 +662,10 @@ static void samSortFirstPipe() {
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");
gfp[0] = fopen("output/f0.txt", "w");
gfp[1] = fopen("output/f1.txt", "w");
gfp[2] = fopen("output/f2.txt", "w");
gfp[3] = fopen("output/f3.txt", "w");
#endif
#if 1