本文链接

https://www.cnblogs.com/wanger-sjtu/p/16818492.html

调用链

tvm搜索算子在需要多线程运行的算子,是在codegen阶段时插入TVMBackendParallelLaunch的调用。

TVMBackendParallelLaunch 是tvm的线程池并行化入口,具体如下

/*!
* \brief The callback function to execute a parallel lambda
* \param task_id the task id of the function. //这里实际就是线程池线程编码,对应第几个线程
* \param penv The parallel environment backs the execution. // num_task, sync
* \param cdata The supporting closure data.
*/
typedef int (*FTVMParallelLambda)(int task_id, TVMParallelGroupEnv* penv, void* cdata); /*!
* \brief Backend function for running parallel jobs.
*
* \param flambda The parallel function to be launched.
* \param cdata The closure data. // 可以认为时循环的变量 codegen时生成
* \param num_task Number of tasks to launch, can be 0, means launch
* with all available threads. // codegen 时写入的是0,运行时根据配置写入
*
* \return 0 when no error is thrown, -1 when failure happens
*/
int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int num_task);

flambda的调用在单线程和多线程下略有区别。

单线程运行时

if (num_workers == 1) {
std::atomic<int32_t> sync_counter{0};
TVMParallelGroupEnv env;
env.num_task = 1;
env.sync_handle = &sync_counter;
(*flambda)(0, &env, cdata);
return 0;
}

多线程运行时

// launcher->Init(flambda, cdata, num_task, need_sync != 0);
this->cdata = cdata;
this->flambda = flambda;
this->env.num_task = num_task; while (queue->Pop(&task, spin_count)) {
ICHECK(task.launcher != nullptr);
TVMParallelGroupEnv* penv = &(task.launcher->env);
void* cdata = task.launcher->cdata;
if ((*task.launcher->flambda)(task.task_id, penv, cdata) == 0) {
task.launcher->SignalJobFinish();
} else {
task.launcher->SignalJobError(task.task_id);
}
}

可以看到 待并行函数中 TVMParallelGroupEnv* penv 包含了实际的运行时线程,运行时可以根据这个确定每个线程的工作区间和步长。

cdata则是线程运行时需要变量信息,闭包变量。

总结

对要并行的函数,实际上是按照lambda表达式的方式生成的。FTVMParallelLambda 的输入参数前两个是运行时确定的,第三个是捕获的外部变量。

codegen 过程

下面验证一下上述的猜测。

codegen过程中,实际上是在遍历tir Stmt的AST,因为生成的循环都是基于For的,调用过程也比较简单了。

void CodeGenCPU::VisitStmt_(const ForNode* op)  // ->
CreateParallelLaunch(For(op->loop_var, op->min, op->extent, op->kind, op->body,
op->thread_binding, op->annotations),
0, std::string("loop_parallel_") + op->loop_var->name_hint.c_str()); // ->
CodeGenCPU::VisitStmt_(const ForNode* op);

当遍历到For节点时, 根据属性判断是否并行加速。这里只分析加速场景。此时parallel_env_.penv == nullptr 创建多线程调用函数,进入CreateParallelLaunch函数。

然后 再生成 For的遍历逻辑。this->VisitStmt(body); 这里的body其实还是For ,这时候就进入

} else {
// already in parallel env.

前文的猜测也在这里得到验证。


void CodeGenCPU::VisitStmt_(const ForNode* op) {
ICHECK(is_zero(op->min));
if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) {
CodeGenLLVM::VisitStmt_(op);
} else if (op->kind == ForKind::kParallel) {
if (parallel_env_.penv == nullptr) {
CreateParallelLaunch(For(op->loop_var, op->min, op->extent, op->kind, op->body,
op->thread_binding, op->annotations),
0, std::string("loop_parallel_") + op->loop_var->name_hint.c_str());
} else {
// already in parallel env.
ICHECK(parallel_env_.task_id.defined());
ICHECK(parallel_env_.num_task.defined());
ICHECK(parallel_env_.penv != nullptr);
DataType t = op->extent.dtype();
PrimExpr num_task = cast(t, parallel_env_.num_task);
PrimExpr task_id = cast(t, parallel_env_.task_id);
ICHECK(!parallel_env_.in_parallel_loop)
<< "Nested parallel loop is not supported by threadpool, try fuse them instead";
parallel_env_.in_parallel_loop = true;
if (parallel_env_.stride_pattern) {
CreateSerialFor(MakeValue(task_id), MakeValue(op->extent), MakeValue(num_task),
op->loop_var, op->body);
} else {
PrimExpr step = (op->extent + num_task - make_const(t, 1)) / num_task;
PrimExpr begin = min(task_id * step, op->extent);
PrimExpr end = min((task_id + make_const(t, 1)) * step, op->extent);
CreateSerialFor(MakeValue(begin), MakeValue(end),
llvm::ConstantInt::getSigned(GetLLVMType(end), 1), op->loop_var, op->body);
}
parallel_env_.in_parallel_loop = false;
++parallel_env_.parallel_loop_count;
}
} else {
LOG(FATAL) << "cannot handle for type " << op->kind;
}
} /*
const Stmt& body For 循环的statement
int num_task, 这里设置的是0,根据运行时参数确定使用线程
std::string name
*/
void CodeGenCPU::CreateParallelLaunch(const Stmt& body, int num_task, std::string name) {
// closure data
llvm::Function* f =
llvm::Function::Create(ftype_tvm_parallel_lambda_, llvm::Function::PrivateLinkage,
"__tvm_parallel_lambda", module_.get());
SetTargetAttributes(f); // allocate and setup the closure, call the closure. //For 循环内部变量。这里需要声明一下
Array<Var> vfields = tir::UndefinedVars(body, {});
uint64_t nbytes;
TypedPointer cdata = PackClosureData(vfields, &nbytes, "closure_" + name); // 可以认为时循环的变量
#if TVM_LLVM_VERSION >= 90
auto launch_callee = llvm::FunctionCallee(ftype_tvm_parallel_launch_, RuntimeTVMParallelLaunch());
#else
auto launch_callee = RuntimeTVMParallelLaunch();
#endif
llvm::BasicBlock* par_launch_end = CheckCallSuccess(builder_->CreateCall(
launch_callee,
{f, builder_->CreatePointerCast(cdata.addr, t_void_p_), ConstInt32(num_task)}));
// Setup the closure function.
auto* lambda_entry =
llvm::BasicBlock::Create(*llvm_target_->GetContext(), "parallel_closure_entry", f);
builder_->SetInsertPoint(lambda_entry);
auto it = f->arg_begin();
llvm::Value* task_id = &(*it++);
task_id->setName("task_id");
llvm::Value* penv = &(*it++);
cdata.addr = builder_->CreatePointerCast(&(*it++), cdata.addr->getType());
// setup new variable map, swap it with current var context.
std::unordered_map<const VarNode*, llvm::Value*> new_vmap;
UnpackClosureData(cdata, vfields, &new_vmap);
// setup parallel env
ParallelEnv par_env;
par_env.task_id = Var("task_id", DataType::Int(32));
par_env.num_task = Var("num_task", DataType::Int(32));
new_vmap[par_env.task_id.get()] = task_id;
new_vmap[par_env.num_task.get()] = builder_->CreateLoad(
t_int32_,
builder_->CreateInBoundsGEP(t_tvm_parallel_group_env_, penv, {ConstInt32(0), ConstInt32(1)}),
"num_task");
par_env.penv = penv;
auto new_analyzer = std::make_unique<arith::Analyzer>();
std::swap(function_, f);
std::swap(parallel_env_, par_env);
std::swap(analyzer_, new_analyzer);
std::swap(var_map_, new_vmap);
this->VisitStmt(body);
builder_->CreateRet(ConstInt32(0));
// swap the var map back, now we are back on track.
std::swap(var_map_, new_vmap);
std::swap(analyzer_, new_analyzer);
std::swap(parallel_env_, par_env);
std::swap(function_, f);
ICHECK_NE(par_env.parallel_loop_count, 0) << "Cannot find parallel loop within parallel launch";
builder_->SetInsertPoint(par_launch_end);
}

tvm-多线程代码生成和运行的更多相关文章

  1. ExecutorService java多线程分割list运行

    调用方法 int threadNum = 7; while(true) { List<FaceAnalyseImage> list = faceAnalyseImageMapper.sel ...

  2. Microsoft SDK 中Sample案例之Amcap項目 的运行方法(转)

    http://blog.csdn.net/erick08/article/details/7194575 Microsoft  SDK 中Sample之Amcap 的运行方法      写这篇文章的由 ...

  3. [转帖]运行时库(runtime library)

    运行时库(runtime library) https://blog.csdn.net/xitie8523/article/details/82712105 没学过这些东西 或者当时上课没听 又或者 ...

  4. TVM如何训练TinyML

    TVM如何训练TinyML 机器学习研究人员和从业人员对"裸机"(低功耗,通常没有操作系统)设备产生了广泛的兴趣.尽管专家已经有可能在某些裸机设备上运行某些模型,但是为各种设备优化 ...

  5. .NET基础拾遗(5)多线程开发基础

    Index : (1)类型语法.内存管理和垃圾回收基础 (2)面向对象的实现和异常的处理基础 (3)字符串.集合与流 (4)委托.事件.反射与特性 (5)多线程开发基础 (6)ADO.NET与数据库开 ...

  6. java之多线程 二

    线程的生命周期: 当线程被创建并被启动时,它既不是一启动就进入了执行状态,在线程的生命周期中,它要经过new(新建),就绪(Runnable),运行(Running),阻塞(Blocked),dead ...

  7. Python 多线程教程:并发与并行

    转载于: https://my.oschina.net/leejun2005/blog/398826 在批评Python的讨论中,常常说起Python多线程是多么的难用.还有人对 global int ...

  8. WebDriver多线程并发

    要想多线程并发的运行WebDriver,必须同时满足2个条件,首先你的测试程序是多线程,其次需要用到Selenium Server.下载位置如下图: 下载下来后是一个jar包,需要在命令行中运行.里面 ...

  9. JavaEE基础(二十四)/多线程

    1.多线程(多线程的引入) 1.什么是线程 线程是程序执行的一条路径, 一个进程中可以包含多条线程 多线程并发执行可以提高程序的效率, 可以同时完成多项工作 2.多线程的应用场景 红蜘蛛同时共享屏幕给 ...

  10. php模拟多线程

    一:应该知道的: php本身是不支持多线, 但是php的好搭档,apache和linux是支持的,故lamp才是最佳组合,还在使用win服务器的现在知道为什么要用linux吧.既然是模拟的, 就不是真 ...

随机推荐

  1. 本地Navicat无法连接服务器mysql8.0

    本地Navicat无法连接服务器mysql8.0 原因: mysql未开启远程连接权限 navivat与mysql密码加密不一致,需一致加密规则 允许远程连接  use msyql; // 1.先查询 ...

  2. C# Kafka重置到最新的偏移量,即从指定的Partition订阅消息使用Assign方法

    在使用Kafka的过程中,消费者断掉之后,再次开始消费时,消费者会从断掉时的位置重新开始消费. 场景再现:比如昨天消费者晚上断掉了,今天上午我们会发现kafka消费的数据不是最新的,而是昨天晚上的数据 ...

  3. DG:Oracle查看是否搭建DataGuard

    Oracle查看是否是DataGuard 1.查看归档路径 show parameter log_archive_dest LOG_ARCHIVE_DEST_n, 归档文件的生成路径, LOCATIO ...

  4. PHP利用 JSON 将XML转换为数组

    在很多开发项目中,我们都会遇到将XML文件转换为数组使用,因此在本篇 PHP教程 中,UncleToo和大家一起学习 如何转换XML为数组 . 现在有一个uncletoo.xml的配置文件,格式如下: ...

  5. PWN 学习日志(1): pwntools简单使用与栈溢出实践

    常用的模块 模块 功能 asm 汇编与反汇编 dynelf 远程符号泄漏 elf 对elf文件进行操作 memleak 用于内存泄漏 shellcraft shellcode生成器 gdb 配合gdb ...

  6. cocos2dx返回Android游戏黑屏解决办法

    用来解决返回Android游戏加载资源时黑屏的问题.帖子过些日子估计就沉了,所以转出来,以供后面查询. 需要修改三个文件: 1) cocos2dx/platform/CCPlatformMacros. ...

  7. scrapy框架简介

    一.安装scrapy环境 -mac或linux:pip install scrapy -windows: 1.pip install wheel 2.pip install twinsted 3.pi ...

  8. ☆常用的Sql语句汇总(DDL/DML)

    常用的sql语句汇总 1.获取所有表名.表信息 里面有表注释 数据库种类 sql 备注 mysql -- 获取所有表名.视图名show tables-- 获取 dev_test_data数据库 所有表 ...

  9. C# POST提交以及 解析 JSON 实例

    一.解析的JSON字符串如下 {"tinyurl":"http:\/\/dwz.cn\/v9BxE","status":0,"lo ...

  10. #AI 1分钟学会,利用AI制作思维导图 (NewBing&X-Mind )

    思维导图是一种有效的思考和学习工具,它可以帮助你整理和呈现信息,激发你的创造力和记忆力.但是,传统的思维导图软件往往需要你花费大量的时间和精力来设计和绘制,而且难以修改和分享.有没有一种更简单和智能的 ...