课程 | AI 编译器(二)

项目简介

InfiniTensor 与 RefactorGraph.

关于 ONNX

正确打开方式 - 对照 ONNX Operator Docs 与 netron.app 食用。

AI 编译器对 AI 程序的抽象

这部分主要翻笔记和讲义吧。只记录一些对 RefactorGraph 代码的说明。

graph_topo 相关定义

struct Node {
    count_t
        _localEdgesCount,
        _inputsCount,
        _outputsCount;
};
using OutputEdge = count_t;

count_t _lenIn, _lenOut;
std::vector<OutputEdge> _connections;
std::vector<Node> _nodes;

这里只存了节点表,边表是虚拟的,通过计算得到。关注以下代码

auto GraphTopo::globalInputsCount() const noexcept -> size_t { return _lenIn; }
auto GraphTopo::globalOutputsCount() const noexcept -> size_t { return _lenOut; }
auto GraphTopo::nodeCount() const noexcept -> size_t { return static_cast<count_t>(_nodes.size()); }
auto GraphTopo::edgeCount() const noexcept -> size_t {
    return std::accumulate(_nodes.begin(), _nodes.end(), _lenIn,
                            [](auto const acc, auto const &n) { return acc + n._localEdgesCount + n._outputsCount; });
}

其中,edgeCount 的计算为累加全局输入 / 各节点张量常量输入 / 各节点输出。这样可以吗?答案是肯定的,节点间的边不需要重复计算,而输出和输入只有一侧绑定到对 nodes 的迭代上即可。最终,得到边表。

关于 connections 的说明

关注如下代码

auto GraphTopo::globalInputs() const noexcept -> range_t<count_t> {
    return range0_(_lenIn);
}
auto GraphTopo::globalOutputs() const noexcept -> slice_t<count_t> {
    return slice(_connections.data(), static_cast<size_t>(_lenOut));
}
auto GraphTopo::connections() const noexcept -> slice_t<count_t> {
    return slice(_connections.data(), _connections.size());
}

可以看到,获取全局输入,返回 0..(lenIn - 1),而输出则为一个切片操作,从 connections.data() 中切出前 lenOut 个元素。

总结

实际上整个图拓扑本质上是由 3 个列表组成,节点表、边表和连接关系表,三个表的元素都是严格有序的。但是在数据结构定义中只能看到 2 个表,是因为边表实际上不需要存储任何东西,所以省略掉了,或者也可以理解成这里还有一个由边数个大小为零的元素组成的虚拟的数组:

Node       nodes      [node_count];
Edge       edges      [edge_count];       // size = 0
Connection connections[connection_count];
posted @ 2025-03-08 22:21  Miya_Official  阅读(27)  评论(0)    收藏  举报