爱笑的小姐姐 · 2021年11月10日

Tengine推理引擎 | 异构计算 源码解读

异构计算主要是指使用不同类型指令集和体系架构的计算单元组成系统的计算方式

本人理解就是不同加速设备,如何执行同一套代码。

之前已经通过四篇文章对Tengine推理引擎的apigraph调度策略切图策略做了讲解,本文继续对Tengine的异构计算的NPU部分的核心代码进行简单介绍。

int timvx_describe(struct device* device, struct vector* allowed_ops, struct vector* blocked_ops, struct vector* precision)
{
    (void)device;
    // 将npu支持的算子添加到allowed_ops列表中
    for (int op_type : timvx_supported_ops)
    {
        push_vector_data(allowed_ops, &op_type);
    }
    // 遍历所有的OP算子,把NPU不支持的算子添加到blocked_ops列表中
    for (int i = 0; i < OP_BUILTIN_LAST; i++)
    {
        bool in_list = false;
​
        for (const auto& type : timvx_supported_ops)
        {
            if (type == i)
            {
                in_list = true;
                break;
            }
        }
​
        if (!in_list)
        {
            push_vector_data(blocked_ops, &i);
        }
    }
    // 把支持的数据类型添加到prcision列表中
    int precision_var = TENGINE_DT_UINT8;
    push_vector_data(precision, &precision_var);
    precision_var = TENGINE_DT_FP16;
    push_vector_data(precision, &precision_var);
​
    return 0;
}
​
​
​
int timvx_allocate(struct device* device, struct subgraph* sub_graph)
{
    if (nullptr == device)
    {
        return -1;
    }
    // 初始化子图的输入等待数为0
    sub_graph->input_wait_count = 0;
    // 遍历子图的所有input tensor,如果tensor的张量类型是VAR,则输入等待数要自增
    for (int i = 0; i < sub_graph->input_num; i++)
    {
        struct tensor* tensor = get_ir_graph_tensor(sub_graph->graph, sub_graph->input_tensor_list[i]);
​
        if (tensor->tensor_type == TENSOR_TYPE_VAR)
            sub_graph->input_wait_count++;
    }
    return 0;
}
​
​
int timvx_split_graph(struct graph* ir_graph)
{
    // 获取graph的当前执行设备
    struct device* cur_dev = ir_graph->attribute->context->device;
    // 如果设备不是NPU名字,那就直接返回
    if (0 != strcmp(TIMVX_DEV_NAME, cur_dev->name))
    {
        return -1;
    }
    // 开辟算子列表、支持精度的内存空间
    struct vector* allowed_ops = create_vector(sizeof(int), nullptr);
    struct vector* blocked_ops = create_vector(sizeof(int), nullptr);
    struct vector* precision = create_vector(sizeof(int), nullptr);
    // 描述设备信息,把NPU支持的算子和精度添加到描述信息中
    cur_dev->allocator->describe(cur_dev, allowed_ops, blocked_ops, precision);
    // 切图
    split_graph_node_to_sub_graph(ir_graph, allowed_ops, blocked_ops, precision);
    // 释放内存空间
    release_vector(allowed_ops);
    release_vector(blocked_ops);
    release_vector(precision);
​
    // 为graph的每一个subgraph中的input tensor列表和output tensor list设置输入输出tensor id
    generate_sub_graph_io(ir_graph);
    // 子图添加到graph中
    add_sub_graph_to_ir_graph(ir_graph);
​
    // add node sub graph id
    // 遍历graph的子图列表,设置子图索引
    for (int i = 0; i < (uint16_t)get_vector_num(ir_graph->subgraph_list); i++)
    {
        struct subgraph* sub_graph = *(struct subgraph**)get_vector_data(ir_graph->subgraph_list, i);
        sub_graph->index = i;
        // 遍历子图的所有节点,设置节点对应的子图id为当前子图的索引
        for (uint16_t j = 0; j < sub_graph->node_num; j++)
        {
            uint16_t node_id = sub_graph->node_list[j];
            struct node* ir_node = get_ir_graph_node(ir_graph, node_id);
            ir_node->subgraph_idx = sub_graph->index;
        }
    }
​
    return 0;
}
}
原文链接:https://zhuanlan.zhihu.com/p/423317500
作者:闪电侠的右手

推荐阅读

推荐阅读
关注数
3391
内容数
68
Tengine是一款轻量级模块化高性能的神经网络推理引擎 ;欢迎体验Tengine,[链接] 《Tengine开发者入门资料包》[链接]
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息