论文精读(三)

TensorIR: An Abstraction for Automatic Tensorized Program Optimization

题目和作者

  • 关键概念
    • Tensor IR
    • Automatic Tensorized Optimization
  • 冯思远:上交本博
  • 侯博涵:上交->CMU

摘要

  • 需求:各种后端设备、加速原语和机器学习模型带来了巨大的工程挑战。
  • 贡献:提出 TensorIR,一种编译器抽象,用于张量计算优化
    • 推广了现有机器学习编译器中的循环表示
    • 把张量计算作为一等公民
    • 构建了一个端到端的系统
  • 实验:效果拔群
    • 自动优化
    • 高效
    • 跨平台

简介

  • 需求:张量化程序自动优化
  • 挑战
    • 张量化程序的抽象:合理表达不同硬件后端,包括内存访问,线程层次和张量化计算原语,各种操作符
    • 巨大的张量化程序优化空间
  • 贡献
    • 新的抽象:引入新结构,“块”,隔离张量化程序和外层嵌套循环
    • 构建了对应的转换原语,生成张量化程序优化的丰富搜索空间
    • 新的张量感知调度器
  • 最高加速 7 倍

概述

  • 核心 Idea:把张量化程序和外层嵌套循环分开,原因:张量化程序通常和硬件原语相关

TensorIR 抽象

  • 三个核心抽象
    • 多维缓冲区
    • 嵌套循环,GPU 中存在线程绑定
    • 块:块内能包含嵌套循环

  • 包含三个部分

    • 块签名
    • 初始化语句
    • 块体
  • 块签名

    • 块签名:包含影响块外部变换的信息,是保证外部变换安全的约束
    • 原理:隔离张量化计算,在不查看块主体的情况下转换块外循环嵌套,所以引入块签名来提取转换需要的依赖信息
    • 块迭代域:将块内部迭代器的信息绑定到外部迭代器;可以看作表示整型域集和迭代器之间关系的一种特定方式,简单高效
    • 访问区域和依赖项
    • 规约块和初始化:对于规约块往往包含初始化步骤,可以把它们分成两个块来处理,但是把他们合并处理通常是有帮助的,所以提供初始化块签名来辅助代码变换 #### 变换调度
  • 如果一个块只包含以子块为叶的嵌套循环,则它是可调度的
    • 可调度意味着,可以通过分析子块签名来变换可调度块的嵌套循环和子块计算位置
    • 其中子块是否可调度无关紧要,块之间的调度是互相隔离的
  • 循环变换
    • 循环平铺,变换计算位置等,还支持把循环绑定到 GPU 线程,标记向量化和循环展开
    • 与以往的张量编译器不同,只需要检查块签名来计算依赖,而不是提取每个张量计算语句的依赖
    • 块隔离不会阻止重要的跨块协作优化
  • 块化:方便缓存,规约块初始化等操作
  • 分离调度和 TensorIR
    • 与以往的张量编译器不同,调度不需要添加新的数据结构和修改相应的 lower 规则
    • 本文的调度实现为一个 TensorIR 程序到另一个 TensorIR 程序的独立转换感觉完全是 Pass 了
    • 以往的调度主要基于 TE (张量表达式)

验证

  • 验证
    • 确保 TensorIR 程序转换前后等价
    • 给用户报 Warning 或 Error
  • 循环嵌套验证
    • 域约束,域大小,迭代器独立验证
    • 生产者-消费者关系验证
  • 线程验证
    • 线程绑定验证
    • 协同内存访问验证
    • 执行作用域验证

编程工作

  • 实现了一种 DSA实际上指的是 TVMScript来允许用户直接编写 TensorIR

自动调度张量程序

  • 需要一种自动解决方案来优化一组调度原语:张量感知自动调度器

张量原语抽象

  • 每种张量原语包含两个部分
    • TensorIR 计算语义
    • 张量化计算的底层实现
  • 需要注意的是张量原语通常和内存指令一起应用

生成张量化候选

  • 给定一个输入程序和一个后端目标,将程序与可能的张量原语计算语义匹配,生成张量化候选
    • 将输入程序重索引:把计算式写成等效的形式,形成映射
    • 匹配:匹配重索引后的程序和张量原语中各个索引所在的项
    • 匹配成功,则可以替换索引
    • 循环检查(重组块实例空间的形状需要被张量原语的子问题大小整除),循环平铺,块化
  • 相关数学证明

生成张量化程序草图

  • 程序草图此时已经固定了部分程序结构,同时也为其它程序优化留下了搜索空间
  • 迭代地应用预定义的草图生成规则来生成草图
  • 数据移动作为一等公民
    • 在生成草图时插入“自动复制块”
    • 重视数据移动,将其与计算调度解耦

进化搜索

  • 此时获得了数十亿个张量化程序草图,利用进化搜索来找到一个优化的张量化程序
  • 特征提取:块签名、块体
  • cost model:张量化程序 cost model 的推广,是一个有前途的领域
  • 验证以拒绝无效程序

评估

  • 在 TVM 上实现
  • 核心问题
    • TensorIR 能不能优化通用机器学习算子
    • TensorIR 能不能改善端到端的网络推理
    • TensorIR 能不能支持不同硬件平台的张量原语

算子

  • Benchmark
    • 各种卷积
    • GMM
  • 环境:RTX3080
  • 比较
    • Ansor 版 TVM
    • AMOS
    • CUTLASS
    • TensorRT
  • 和机器学习编译器比:效果拔群
  • 和手工算子库比:有输有赢

端到端

  • Benchmark
    • ResNet-50
    • MobileNet
    • V2 BERT_large
    • ViT
  • 环境:RTX3080
  • 比较
    • AMOS
    • TVM
    • PyTorch
    • TensorRT
  • 有输有赢,支持更好
  • 调优时间:比 TVM 短

ARM CPU

  • 环境:AWS Graviton2 CPU
  • 算子:和 TVM、ArmComputeLib 比
    • 有输有赢
  • 端到端:和 PyTorch、TVM 比
    • 效果拔群

相关工作

  • 各种算子库
  • 各种机器学习编译器
  • 各种机器学习编译器自动优化
  • 各种自动张量化

结论

  • 提出了一种自动张量化程序优化的 IR

评论


论文精读(三)
http://example.com/2023/04/19/论文精读(三)/
作者
zty
发布于
2023年4月19日
许可协议