mlir的Linalg方言介绍

最新推荐文章于 2025-04-03 09:51:51 发布

早睡的叶子

最新推荐文章于 2025-04-03 09:51:51 发布

阅读量1.1k

点赞数 14

分类专栏： AI编译器文章标签： mlir

本文链接：https://blog.csdn.net/sexyluna/article/details/143925255

版权

AI编译器专栏收录该内容

19 篇文章

订阅专栏

Linalg Dialect

基本原理

编译器友好的自定义OP方言

在这里插入图片描述

Linalg 旨在解决 MLIR 中的高级层次优化（HHO 框），并在专家编译器混合环境（即CGSel框）中良好地互操作。

这项工作的灵感来自于该领域丰富的现有技术，并试图从中吸取重要的经验教训。此文档和内省工作也是在建议成立一个工作组来讨论高级张量计算基元方言和转换的开发的背景下进行的。我们希望现有技术的经验教训、本文档中概述的设计原则以及 Linalg 的架构能够帮助社区了解定义这些高级张量计算基元的路径。

“‘linalg’ Dialect” 页面主要介绍了 MLIR（多级中间表示）中的 “linalg” 方言，其旨在解决 MLIR 中的高级分层优化（HHO）问题，并在混合专家编译器环境中良好协作。以下是该页面的主要内容总结：

设计目的与理念

解决优化问题：用于处理 MLIR 中的高级分层优化，在混合专家编译器环境中实现良好交互，设计决策细节参考相关原理文档。
关键变换集合：一系列关键变换推动了 “linalg” 的设计，包括渐进式缓冲区分配、参数化切片、提升到快速内存中的临时缓冲区、带参数的切片与融合的生产者 - 消费者融合、映射到并行和归约循环及硬件、向量化、降低到循环、降低到库调用或特殊指令、部分降低到更细粒度 “linalg” 操作的迭代等。

“linalg” 操作概述

受启发来源与设计优势：从先前艺术中获取灵感，其设计允许定义具有通用属性的自定义操作，支持关键变换，操作数可为张量或缓冲区，输出张量有 “init 张量”（提供初始值，通过迭代更新结果，可能被优化为寄存器级 SSA 值，期望可重写为缓冲区的原位更新）和 “仅形状”（仅携带形状信息，未来可能被合适形状类型替代）两种类型。
承载有效负载操作的属性
- 属性 1：输入和输出操作数定义迭代空间：“linalg.generic” 操作从其操作数完全派生迭代空间规范，保证按类型迭代操作数时不会发生越界访问，操作语义限于结构化数据类型，虽有一定限制，但 MLIR 可混合不同抽象级别，该属性有助于实现扩展性。
- 属性 2：控制和数据结构之间的可逆映射：定义了迭代空间（循环）和数据之间的映射，此映射需可逆，用于实现切片、融合和提升到临时缓冲区等变换，当前使用仿射映射列表，长远来看可能动态评估。
- 属性 3：迭代器类型显式定义：“linalg.generic” 操作显式声明迭代器类型，信息用于变换，传统上这些类型源于复杂依赖分析，在 “linalg” 中提前声明可传递难以从低级信息派生的属性，当前主要使用并行和归约循环类型。
- 属性 4：计算有效负载使用区域指定：“linalg.generic” 操作通过区域指定计算有效负载，区域接受张量或缓冲区操作数的标量元素类型及可能的特殊值，当前对区域语义无额外限制，旨在探索区域和迭代器类型交叉处的设计权衡，降低到循环时会遇到类似需求，期望重用低级基础架构。
- 属性 5：可映射到外部库调用：通过指定 “SymbolAttr”，“linalg.generic” 操作可映射到外部库调用，需考虑操作语义保留和 ABI 集成，示例展示了添加属性后的转换过程，“linalg” 方言采用与 “BLAS” 类似的约定与外部库交互，当前存在关于扩展互操作性的讨论。
- 属性 6：对整个输出操作数的完美嵌套写入：“linalg.generic” 操作表示完美嵌套循环巢，写入整个内存区域，此结构便于关键循环变换，转换不完美嵌套代码为完美嵌套代码可通过循环分布和条件嵌入实现，虽可能自动完成，但当前 “linalg.generic” 强制此语义，完成变换后需撤销转换，依赖后期优化可能影响性能。

开放问题与设计替代方案

提出多个开放问题，如 “linalg.generic” 是否支持嵌套、区域接受类型、是否解决自动微分、六个属性是否必要、对声明性规范的依赖程度、是否满足社区需求及如何扩展等，强调应在 MLIR 整体环境中考虑这些问题，不同 IR 级别可协作解决问题。

操作列表

列举了众多 “linalg” 方言中的操作，如 “linalg.abs”（逐元素应用绝对值）、“linalg.add”（逐元素相加张量）、“linalg.batch_matmul”（执行批量矩阵乘法）等，每个操作介绍了功能、输入输出、特性、接口、属性（部分操作有）等信息，如 “linalg.broadcast”（执行静态广播操作，通过指定维度将输入广播为给定形状）等操作。

实际案例

逐元素加法（linalg.add）

#map = affine_map<(i, j) -> (i, j)>

#attrs = {
  indexing_maps = [#map, #map, #map],
  iterator_types = ["parallel", "parallel"]
}

func.func @elementwise_add(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
  linalg.generic #attrs
  ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
  outs(%C: memref<?x?xf32>) {
    ^bb0(%a: f32, %b: f32, %c: f32):
      %d = arith.addf %a, %b : f32
      linalg.yield %d : f32
  }
  return
}

这个例子展示了如何使用linalg.add（通过linalg.generic实现）对两个二维矩阵%A和%B进行逐元素相加，并将结果存储到%C中。迭代器类型为并行，通过仿射映射指定了输入和输出的对应关系。