ATB当前常用的多线程场景有多机多卡场景、单机多卡场景、host接口并行场景等,本部分主要介绍在单机多卡或多机多卡场景下的多线程使用。
多卡场景下的多线程部分的实现主要有以下两个要点:1、如何多线程调用ATB的Operation接口 2、Context的创建与分配。示例中使用基础的std::thread类来实现多线程。示例中根据卡的数量创建了与卡数相同的线程个数,每个线程对应一张卡。线程在刚开始时会调用aclrtSetDevice函数进行绑卡,绑定对应的NPU卡后,开始创建执行所需的硬件相关资源,如:ATB的context,device stream等。
ATB当前的功能接口都是通过Operation对外提供的,又由于同一个Operation对象的Setup接口与Execute接口存在依赖关系,当前推荐的多线程使用方式是多个线程/进程单独创建自己的Operation对象,这样可以保证同一Operation对象Setup接口与Execute接口的串行,保证功能的正确性。
由于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_); } |