【AI系统】常量折叠原理

简介: 常量折叠是一种编译器优化技术,通过在编译阶段计算常量表达式,用结果替换原表达式,减少运行时计算。传统编译器在编译期间识别并计算常量表达式,如 Python 中 `day_sec = 24*60*60` 被优化为 `day_sec = 86400`。AI 编译器则在计算图中进行类似优化,如 TensorFlow 通过分析计算图节点,提前计算确定结果的节点,提高执行效率。

常量折叠(Constant Folding)是编译器的一种优化技术,它通过在编译期间对常量表达式进行计算,将其结果替换为常量值,从而减少程序运行时的计算和开销。

传统编译器的常量折叠

传统编译器在编译期间,编译器会设法识别出常量表达式,对其进行求值,然后用求值的结果来替换表达式,从而使得运行时更精简。

day_ sec = 24*60*60

当编译器遇到这样的一个常量表达式时,表达式会被计算值所替换。因此上述表达式可以等效地被执行为 :

day_sec = 86400

以 python 为例,在 python 中,使用反汇编模块 Disassembler 获取 day_ sec = 24×60×60 的 CPython 的字节码。

import dis

dis.dis("day_sec=24*60*60")
  1           0 LOAD_CONST               0 (86400)
              2 STORE_NAME               0 (day_sec)
              4 LOAD_CONST               1 (None)
              6 RETURN_VALUE

上述的 CPython 的字节码表明,python 在对 day_sec 赋值是直接加载计算结果 86400,相比于 3 次载入数据和两次乘法,更加地高效。

表达式可进行常量折叠,当且仅当表达式的所有子表达式都是常量。而子表达式被判断为常量通常需要常量传播的帮助。

举个例子:

int x = 14;
int y = 7 - x / 2;
return y * (28 / x + 2);

编译器首先会进行一次常量传播,将所有对 x 的引用替换成 14:

int x = 14;
int y = 7 - 14 / 2;
return y * (28 / 14 + 2);

对于 $7 - 14 /2 \quad\quad\quad$ 和 $28/14+2\quad\quad\quad\quad$ 这两个表达式,由于其所有的子表达式都是常量(字面值),所以这两个表达式可以进常量折叠优化,优化后得到:

int x = 14;
int y = 0;
return y * 4;

编译器再次对 y 进行常量传播,将所有对 y 的可达引用都替换 0 得到:

int x = 14;
int y = 0;
return 0 * 4;

对于表达式 $0 * 4 \quad\quad$,由于其所有的子表达式都是常量(字面值),所以这个表达式可以进行常量折叠优化,优化后得到:

int x = 14;
int y = 0;
return 0;

由例子可见,常量传播对于常量折叠的重要性。在传统编译器中,常量传播主要是通过对控制流图(CFG)进行可达性分析,为每个基本块维护一个可达集合,记为 $Reaches(n)\quad\quad\quad\quad\quad$。其含义为若定义 $d\in Reaches(n)\quad\quad\quad\quad\quad\quad$,则意味着存在一条从入口基本块到基本块 $b_n$ 的路径,d 没有被重新定义。计算公式如下:

$$ Reaches(n) = \bigcup_{m\in preds(n)}(DEDef(m)\ \cup (Reaches(m)\ \cap\ \overline{DefKill(m)})) \\ $$

方程的初始条件为:$Reaches(n) = \emptyset , \forall n$

其中:

  • $preds(n)\quad\quad\quad\quad$ 表示 n 的前趋结点集。
  • $DEDef(m)\quad\quad\quad\quad\quad$ 表示基本块 $b_m$ 中向下展示的定义,其含义为若定义 $d\in DEDef(m)\quad\quad\quad\quad\quad\quad$,则意味着从 d 定义处到 $b_m$ 的出口处都没有被重新定义。
  • $DefKill(m)\quad\quad\quad\quad\quad$ 表示在基本块 $b_m$ 中被杀死的定义。其含义若定义 $d\in DefKill(m)\quad\quad\quad\quad\quad\quad$,则意味着从 d 定义处到 $b_m$ 的出口处被重新定义。因此 $\overline{DefKill(m)}$ 包含了 m 中可见的所有定义位置。

从公式上看,如果定义 d 在基本块的出口处是可达的,当且仅当定义 d 是基本块中向下展示的定义,或者定义 d 在基本块的入口处是可定义的,并且在基本块内没有被杀死。根据入口可达集合的定义,存在一条路径即可,所以定义 d 在基本块的入口处是可达的,只需要在其任意前趋结点的出口处是可达的即可。

当已知基本块入口处的可达定义集合,对于基本块中某个定义引用,若从引用处到基本块的入口都没有重定义,且该定义引用在可达定义集合中,则可以用可达定义集合中的值替换该定义引用。如果有重定义,则用重定义的值替换该定义引用,从而达到传播的目的。

如果传播的定义为常量定义,则称常量传播。但不同的编译器对于常量的定义不同,有些编译器需要特殊的关键字标识才会被认定为常量。这将会影响一个表达式能否进行常量折叠的判断。比如下面 java 进行字符串计算的例子 :

String a = "a";
String bc = "bc";
String s1 = a + bc;

在 java 中,只有被 final 修饰的变量,且该变量初始化的值是编译期常量,才是编译期常量。上述例子中,s1 不会被常量折叠成“abc”,而是根据 String 类特有的 + 运算符重载,变成类似下面这样的代码,不会进行常量折叠:

String s2 = new StringBuilder(a).append(b).toString();

所以了解编译器对于编译期常量的定义是一件很重要的事情,这将决定你的代码能否进行常量折叠优化。除此之外,还有其他的一些情况可能会影响常量折叠,比如在 python 中,下面两种情况就不会被常量折叠:

  • 当常量表达式计算的结果溢出时,编译器不会进行常量折叠。比如 2 的 64 次幂会被折叠,而 4 的 64 次幂不会被折叠:

    >>> dis.dis("day_sec=2**64")
    
    1           0 LOAD_CONST               0 (18446744073709551616)
                2 STORE_NAME               0 (day_sec)
                4 LOAD_CONST               1 (None)
                6 RETURN_VALUE
    
    >>> dis.dis("day_sec=4**64")
    
    1           0 LOAD_CONST               0 (4)
                2 LOAD_CONST               1 (64)
                4 BINARY_POWER
                6 STORE_NAME               0 (day_sec)
                8 LOAD_CONST               2 (None)
                10 RETURN_VALUE
    

上面两图分别是 2 的 64 次幂和 4 的 64 次幂的 CPython 字节码,可以看出 2 的 64 次幂会被折叠成一个具体的数,而 4 的 64 次幂不会。

  • 当进行字符串运算的时候,比如两个字符串相加,当且仅当字符串相加的结果大小小于等于 4096 时,该常量表达式才会被折叠,否者不会进行折叠:
>>> dis.dis("day_sec='-'*4097")
1           0 LOAD_CONST               0 ('-')
            2 LOAD_CONST               1 (4097)
            4 BINARY_MULTIPLY
            6 STORE_NAME               0 (day_sec)
            8 LOAD_CONST               2 (None)
            10 RETURN_VALUE

上面给出了结果大小为 4097 的字节码,可以看出并不会被折叠。而大小为 4096 会被折叠,由于会被折叠,得到的字符串是非常长的,这里就不给出例子了,读者可以自行尝试。

AI 编译器的常量折叠

常量折叠作为传统编译器的一种优化技术,其迁移到 AI 编译器依然适用。传统编译器通常是对抽象语法树进行常量折叠优化,而 AI 编译器是对计算图进行常量折叠优化。AI 编译器会对计算图中的每个操作节点进行分析,判断其是否可进行常量折叠。如果可以,则通过计算得到结果替换该节点。

以下是 AI 编译器常量折叠的几个类型:

  • 当计算图中某个节点的数据输入节点均为编译期常量(不同 AI 编译器的定义可能完全不相同)的节点,则可以提前计算出该节点的值来完全替换该节点。

    first

    以 AddN 为例,对于两个形状大小为(N,C,H,W)四维常量 Tensor,AddN 的计算结果在编译期间是可以确定的,可以生成一个与该结果等值的数据节点替换掉 AddN。这样就不需要给 Add 节点分配额外的存储资源,在计算图执行的过程中,也不需要反复计算 AddN 这个操作,可直接进行访问。

  • 输入形状确定的 Shape 类型的操作,比如 Size、Shape 等操作。这些操作都只与输入的形状有关,与输入的具体值无关,而相比于具体输入值来说,当输入的形状在编译期间是一个可以确定的值时,AI 编译器就可以直接计算出 Shape 类型操作的值。

    second

    以 Size 为例,对于形状大小为(1,2,3,4)的四维变量 Tensor,Size 的计算在编译期间是可以确定,AI 编译器会生成一个值为 24 的常量数据节点来替换 Size 节点。

  • 数据输入节点不全为编译常量节点,但是这部分节点是可常量折叠的节点,即经过一系列的常量折叠后,该节点会被替换成编译常量节点。

    third

    如上图,依旧以 AddN 为例,可以看出 AddN2 并不满足第一类的折叠规则,但是 AI 编译器发现 AddN1 是个可常量折叠的节点,AI 编译器生成一个常量数据节点替换掉 AddN1 后,发现 AddN2 也满足第一类的折叠规则,所以 AI 编译器会再生成一个常量数据节点替换掉 AddN2。

与传统编译器相同,AI 编译器在进行常量折叠的时候也会被诸多因素所影响,比如 如果常量的大小(以字节为单位)太大,则不替换它。这可以防止图的大小变得过大。

实现案例

传统编译器实现案例

不同编译器的常量折叠的实现细节不尽相同,下面以 python 为例来描述传统编译器的常量折叠的一种实现。

在 python 中,CPython 会调用 astfold_expr 来对表达式进行常量折叠。astfold_expr 以递归的方式遍历 AST(抽象语法树),并尝试折叠所有的表达式。比如二值运算操作,astfold_expr 会先递归处理该二值操作的左操作数和右操作数,然后将此折叠操作代理给特定的折叠函数 fold_binop。在 fold_binop 中,首先会判断左操作数和右操作数是否都是常量,如果为常量,则判断该二值操作的具体操作类型,然后调用对应基本运算操作,比如 ADD 运算,会调用 PyNumber_Add。最后将计算出来的结果更新到 AST 对应的节点中。

AI 编译器实现案例

不同的 AI 编译器对于常量折叠的细节不尽相同,这里以 tensorflow 为例,描述其常量折叠的实现细节:

  1. 获得逆后续节点集。使用反向深度优先搜索对计算图进行遍历,获取逆后续节点集。这样可以保证在处理当前节点时,其所有输入的节点都已经处理完毕了。

  2. 构建常量集和 Shape 操作节点集。常量集存储编译期常量,可常量折叠节点(包括 Shape 操作节点)。Shape 操作节点集用于存储可常量折叠 Shape 操作节点,由于可常量折叠的 Shape 的特殊性——可以直接计算出结果,所以需要单独存储可常量折叠的 Shape 操作节点及其对应的计算结果,以区分于其他的可常量折叠节点——这些节点有可能是第一类或第三类可常量折叠节点,无法立刻确定其计算结果。

  3. 构建常量图。遍历常量集,对于常量集中的每个节点,如果其在 Shape 操作节点集中,则使用 Shape 操作节点集中提前存储的计算结果生成一个常量数据节点,加入常量图中;如果其不在 Shape 操作节点集中,记该节点为 n,先将 n 的副本加入常量图中,然后遍历 n 在原图中的所有数据入边,重新在常量图中构建这些边,由于 n 的所有数据输入均是常量且由于逆后续节点集的关系,n 的所有输入节点都已经处理过,并已处于常量图中。

  4. 获取需要可常量折叠节点集。遍历常量图中的所有节点,对于某个节点 n,如果其存在一个输出节点不在常量图中,则将其加入需要可常量折叠节点集中。

  5. 常量折叠。判断需要可常量折叠节点集是否为空,如果为空,则说明常量图中的所有常量节点无论替换与否都不会影响原图的计算,这些点通常可能是死代码或者冗余节点,会被后续的优化处理给删除掉,这种情况则不需要进行常量折叠;如果不为空,那么会将生成一个值与该节点的计算结果相同的常量数据节点来替换该节点。

    third test

    如上图中,AddN2 的输出为 Conv(不在常量图中),tensorflow 会计算 AddN2 的输出结果,由于可以肯定其输入要么是编译期常量要么是可折叠常量节点,递归处理其所有输入节点即可。

如果您想了解更多AI知识,与AI专业人士交流,请立即访问昇腾社区官方网站https://www.hiascend.com/ 或者深入研读《AI系统:原理与架构》一书,这里汇聚了海量的AI学习资源和实践课程,为您的AI技术成长提供强劲动力。不仅如此,您还有机会投身于全国昇腾AI创新大赛和昇腾AI开发者创享日等盛事,发现AI世界的无限奥秘~
d

目录
相关文章
|
7天前
|
人工智能 自动驾驶 大数据
预告 | 阿里云邀您参加2024中国生成式AI大会上海站,马上报名
大会以“智能跃进 创造无限”为主题,设置主会场峰会、分会场研讨会及展览区,聚焦大模型、AI Infra等热点议题。阿里云智算集群产品解决方案负责人丛培岩将出席并发表《高性能智算集群设计思考与实践》主题演讲。观众报名现已开放。
|
23天前
|
存储 人工智能 弹性计算
阿里云弹性计算_加速计算专场精华概览 | 2024云栖大会回顾
2024年9月19-21日,2024云栖大会在杭州云栖小镇举行,阿里云智能集团资深技术专家、异构计算产品技术负责人王超等多位产品、技术专家,共同带来了题为《AI Infra的前沿技术与应用实践》的专场session。本次专场重点介绍了阿里云AI Infra 产品架构与技术能力,及用户如何使用阿里云灵骏产品进行AI大模型开发、训练和应用。围绕当下大模型训练和推理的技术难点,专家们分享了如何在阿里云上实现稳定、高效、经济的大模型训练,并通过多个客户案例展示了云上大模型训练的显著优势。
|
27天前
|
存储 人工智能 调度
阿里云吴结生:高性能计算持续创新,响应数据+AI时代的多元化负载需求
在数字化转型的大潮中,每家公司都在积极探索如何利用数据驱动业务增长,而AI技术的快速发展更是加速了这一进程。
|
18天前
|
并行计算 前端开发 物联网
全网首发!真·从0到1!万字长文带你入门Qwen2.5-Coder——介绍、体验、本地部署及简单微调
2024年11月12日,阿里云通义大模型团队正式开源通义千问代码模型全系列,包括6款Qwen2.5-Coder模型,每个规模包含Base和Instruct两个版本。其中32B尺寸的旗舰代码模型在多项基准评测中取得开源最佳成绩,成为全球最强开源代码模型,多项关键能力超越GPT-4o。Qwen2.5-Coder具备强大、多样和实用等优点,通过持续训练,结合源代码、文本代码混合数据及合成数据,显著提升了代码生成、推理和修复等核心任务的性能。此外,该模型还支持多种编程语言,并在人类偏好对齐方面表现出色。本文为周周的奇妙编程原创,阿里云社区首发,未经同意不得转载。
11735 12
|
12天前
|
人工智能 自然语言处理 前端开发
100个降噪蓝牙耳机免费领,用通义灵码从 0 开始打造一个完整APP
打开手机,录制下你完成的代码效果,发布到你的社交媒体,前 100 个@玺哥超Carry、@通义灵码的粉丝,可以免费获得一个降噪蓝牙耳机。
5399 14
|
19天前
|
人工智能 自然语言处理 前端开发
用通义灵码,从 0 开始打造一个完整APP,无需编程经验就可以完成
通义灵码携手科技博主@玺哥超carry 打造全网第一个完整的、面向普通人的自然语言编程教程。完全使用 AI,再配合简单易懂的方法,只要你会打字,就能真正做出一个完整的应用。本教程完全免费,而且为大家准备了 100 个降噪蓝牙耳机,送给前 100 个完成的粉丝。获奖的方式非常简单,只要你跟着教程完成第一课的内容就能获得。
9612 15
|
1月前
|
缓存 监控 Linux
Python 实时获取Linux服务器信息
Python 实时获取Linux服务器信息
|
17天前
|
人工智能 自然语言处理 前端开发
什么?!通义千问也可以在线开发应用了?!
阿里巴巴推出的通义千问,是一个超大规模语言模型,旨在高效处理信息和生成创意内容。它不仅能在创意文案、办公助理、学习助手等领域提供丰富交互体验,还支持定制化解决方案。近日,通义千问推出代码模式,基于Qwen2.5-Coder模型,用户即使不懂编程也能用自然语言生成应用,如个人简历、2048小游戏等。该模式通过预置模板和灵活的自定义选项,极大简化了应用开发过程,助力用户快速实现创意。
|
5天前
|
机器学习/深度学习 人工智能 安全
通义千问开源的QwQ模型,一个会思考的AI,百炼邀您第一时间体验
Qwen团队推出新成员QwQ-32B-Preview,专注于增强AI推理能力。通过深入探索和试验,该模型在数学和编程领域展现了卓越的理解力,但仍在学习和完善中。目前,QwQ-32B-Preview已上线阿里云百炼平台,提供免费体验。
|
13天前
|
人工智能 C++ iOS开发
ollama + qwen2.5-coder + VS Code + Continue 实现本地AI 辅助写代码
本文介绍在Apple M4 MacOS环境下搭建Ollama和qwen2.5-coder模型的过程。首先通过官网或Brew安装Ollama,然后下载qwen2.5-coder模型,可通过终端命令`ollama run qwen2.5-coder`启动模型进行测试。最后,在VS Code中安装Continue插件,并配置qwen2.5-coder模型用于代码开发辅助。
917 5