LinalgExt Dialect 详解

Dialect 名称: iree_linalg_ext C++ 命名空间: ::mlir::iree_compiler::IREE::LinalgExt 主要用途: Linalg 非结构化/实验性扩展

目录

1. 概述
2. 核心概念
3. 类型和属性
4. 核心 Op 列表
5. 关键 Interface
6. 在编译 Pipeline 中的位置

1. 概述

LinalgExt (Linalg Extensions) Dialect 用于试验无法由 Linalg Dialect 高效/直接表示的非结构化操作。

设计目标

1. 实验平台 - 试验新的线性代数操作表示
2. 填补空白 - 提供 Linalg 无法有效表示的操作
3. 向上游反馈 - 成熟后可能合并到上游 MLIR Linalg

与 Linalg 的关系

LinalgExt 扩展而不是替代 Linalg：

• 继承 Linalg 的结构化操作原则
• 提供 Linalg 中缺失的特殊操作
• 可以与 Linalg ops 混合使用
• 最终可以降低为 Linalg 或循环

2. 核心概念

2.1 迭代器类型

LinalgExt 继承并扩展了 Linalg 的迭代器类型概念：

• parallel - 可并行迭代
• reduction - 归约迭代
• window - 窗口/滑动迭代（扩展）

2.2 拆分归约

SplitReduction 是 LinalgExt 的关键特性，支持：

• 将大型归约拆分为多个部分
• 每个部分可以并行计算
• 最后合并所有部分的结果

策略:

• PartialReductionOuterParallel - 外部并行部分归约

2.3 设备映射

LinalgExt ops 支持 DeviceMappingAttrInterface，允许：

• 将循环维度映射到处理元素
• 支持 GPU thread/workgroup 映射
• 支持分布式计算

3. 类型和属性

3.1 类型约束

3.2 属性

3.3 SplitReductionMappingAttr

结构:

SplitReductionMappingAttr {
  int64_t dimension;  // 要拆分和分发的维度
}

用途:

• 指示 scf.forall 表示通过 PartialReductionOuterParallel 策略生成的并行部分归约
• 用于归约 tiling 和分发

4. 核心 Op 列表

4.1 Attention 操作

4.2 卷积操作

4.3 排序和 TopK

4.4 填充和打包

4.5 数据移动

4.6 张量形状操作

4.7 量化操作

4.8 聚合操作

4.9 测试和调试

5. 关键 Interface

5.1 DeviceMappingAttrInterface

设备映射属性接口，允许：

• 将循环维度映射到处理元素
• 支持 GPU thread/workgroup 层级
• 支持分布式计算设备
• 由 SplitReductionMappingAttr 实现

6. 在编译 Pipeline 中的位置

LinalgExt Dialect 主要在全局优化阶段和代码生成阶段使用：

Input → GlobalOptimization → DispatchCreation → Flow → ...
                           ↑
                    LinalgExt ops

典型 Pass 顺序

降低策略

LinalgExt ops 可以通过多种方式降低：

总结

LinalgExt Dialect 是 IREE 的线性代数实验平台：

• ✅ 扩展 Linalg 提供非结构化操作
• ✅ 支持 Attention、Winograd 卷积、TopK 等高级操作
• ✅ 提供拆分归约和设备映射功能
• ✅ 作为新操作的试验场
• ✅ 可以降低为 Linalg、循环或张量操作
• ✅ 位于 GlobalOptimization 和 DispatchCreation 阶段

LinalgExt 使得 IREE 能够：

1. 试验前沿的线性代数操作
2. 优化无法用标准 Linalg 有效表示的操作
3. 提供更丰富的操作集支持现代 AI 模型
4. 向上游 MLIR 反馈成功的设计