多线程部分介绍

ATB当前常用的多线程场景有多机多卡场景、单机多卡场景、host接口并行场景等,本部分主要介绍在单机多卡或多机多卡场景下的多线程使用。

概述

多卡场景下的多线程部分的实现主要有以下两个要点:1、如何多线程调用ATB的Operation接口 2、Context的创建与分配。示例中使用基础的std::thread类来实现多线程。示例中根据卡的数量创建了与卡数相同的线程个数,每个线程对应一张卡。线程在刚开始时会调用aclrtSetDevice函数进行绑卡,绑定对应的NPU卡后,开始创建执行所需的硬件相关资源,如:ATB的context,device stream等。

多线程调用ATB的Operation接口

ATB当前的功能接口都是通过Operation对外提供的,又由于同一个Operation对象的Setup接口与Execute接口存在依赖关系,当前推荐的多线程使用方式是多个线程/进程单独创建自己的Operation对象,这样可以保证同一Operation对象Setup接口与Execute接口的串行,保证功能的正确性。

Context的创建于分配

由于1个Context中包含32*3Mb大小的显存空间及Event等与device强相关的资源,且多线程/进程场景下一个线程/进程通常对应一个device设备,因此Context推荐在线程/进程内绑定好device设备后进行创建。同时Context是与device绑定的,使用device设备1的资源创建的Context无法在device设备2上重复使用。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
int main()
{
	...
	
    /* 线程创建部分 */

    std::vector<std::thread> threadArray(THREAD_SIZE);
    for (size_t i = 0; i < THREAD_SIZE; i++) {
        Model &model = modelArray.at(i);
        threadArray.at(i) = std::thread([i, &model]{ModelExecute(i, model);}); // 线程创建及函数绑定
    }
    for (size_t i = 0; i < THREAD_SIZE; i++) {
        threadArray.at(i).join(); // 等待子线程结束
    }

	...
}

/* 线程启动后立马进行device资源创建 */
void ModelExecute(uint32_t deviceId, Model &model)
{
    // 初始化模型,创建需要的context,stream
    model.InitResource(deviceId);

    ...
}

void Model::InitResource(uint32_t deviceId)
{
    /* 先绑定对应的device设备,再进行device资源的分配 */
    // 配置deviceId
    deviceId_ = deviceId;
    auto ret = aclrtSetDevice(deviceId_);
    CHECK_RET(ret, "aclrtSetDevice failed. ret: " + std::to_string(ret));

    // 创建context
    ret = atb::CreateContext(&mode_context_);
    CHECK_RET(ret, "ATB CreateContext failed. ret: " + std::to_string(ret));
	
    // 创建Stream
    ret = aclrtCreateStream(&model_stream_);
    CHECK_RET(ret, "aclrtCreateStream failed. ret: " + std::to_string(ret));

    // 配置Stream
    mode_context_->SetExecuteStream(model_stream_);
}