论文精读（三）

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

评论

#笔记 #TVM #机器学习编译 #论文 #TensorIR
论文精读（三）
http://example.com/2023/04/19/论文精读（三）/
作者
zty
发布于
2023年4月19日
许可协议