📅 创建时间：2026-06-03 🏷️ 标签：#TorchMLIR #ATen #Functionalization #IREE #PyTorch导出 #MLIR #算子转换 📚 前置知识：[[/04-ai/01-llm-engineering/07-llm-evolution]]（LLM 发展脉络） 📚 相关知识：[[04-mlir-architecture]]（MLIR 架构） [[24-torch-compile]]（torch.compile）

Torch-MLIR：从 PyTorch 算子到 MLIR 方言 / Torch-MLIR from PyTorch Operators to MLIR Dialects

┌─────────────────────────────────────────────────────────────────────────────┐
│  场景：PyTorch 到 MLIR 的导出迷雾                                            │
├─────────────────────────────────────────────────────────────────────────────┤
│  你想用 MLIR 的 IREE 编译器部署 PyTorch 模型到边缘设备（树莓派）。           │
│  查了资料发现 Torch-MLIR 能把 PyTorch 模型 convert 到 MLIR：                 │
│                                                                             │
│  import torch_mlir                                                          │
│  module = torch_mlir.compile(torch_model, example_input,                    │
│                              output_type="linalg-on-tensors")               │
│                                                                             │
│  但"linalg-on-tensors""是什么意思？                                        │
│  为什么输出类型有十几种选择（TOSA, Linalg, MHLO, ...）？                     │
│  我应该选哪个？                                                              │
└─────────────────────────────────────────────────────────────────────────────┘

第1节 PyTorch 算子系统

1.1 ATen：PyTorch 的核心算子库

ATen（A Tensor Library） 是 PyTorch 的底层张量运算库，提供 CPU/GPU 统一的算子接口：

python

# PyTorch 算子调用链
import torch

# 用户代码
x = torch.randn(64, 128)
w = torch.randn(128, 256, requires_grad=True)
y = torch.matmul(x, w)  # 内部调用 ATen
loss = y.sum()
loss.backward()

# ATen 的内部层次
"""
┌─────────────────────────────────────────────────────────────────┐
│  Python API (torch.matmul)                                       │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│  Dispatch (device dispatch, schema validation)                  │
│  torch::dispatch(... )                                           │
└─────────────────────────────────────────────────────────────────┘
                              │
              ┌───────────────┼───────────────┐
              ▼               ▼               ▼
┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐
│     CPU          │ │     CUDA        │ │     HIP         │
│   (ATen CPU)     │ │   (ATen CUDA)   │ │   (ATen HIP)    │
│                  │ │                  │ │                  │
│ · native/        │ │ · native/cuda/  │ │ · native/hip/   │
│   LegacyATen     │ │ · cudnn/        │ │ · hipblas/      │
└──────────────────┘ └──────────────────┘ └──────────────────┘
              │               │               │
              ▼               ▼               ▼
┌──────────────────┐ ┌──────────────────┐ ┌──────────────────┐
│  Eigen (CPU)    │ │  cuBLAS/cuDNN    │ │  rocBLAS/MIOpen  │
│  MKL/OpenBLAS   │ │  (NVIDIA libs)  │ │  (AMD libs)     │
└──────────────────┘ └──────────────────┘ └──────────────────┘
"""

1.2 PyTorch IR 的特殊性

PyTorch 的算子系统与其他框架不同，有几个关键特点：

python

# 特点1：PyTorch 支持动态图（Eager Mode）
import torch

# Shape 在运行时才知道
x = torch.randn(input_dim, 512)  # input_dim 可以是任何值

# 特点2：PyTorch 有副作用（in-place 操作）
a = torch.randn(10)
a.add_(1)  # 下划线表示 in-place，a 本身被修改

# 特点3：PyTorch 有 Python control flow
for i in range(n_layers):  # n_layers 是 Python 值
    x = layer(x)           # 循环次数在运行时决定

# 特点4：PyTorch 张量可以有别名（aliasing）
b = a[::2]  # b 是 a 的视图，共享底层数据
b.add_(1)   # 这也会影响 a

# 这些特性使得 PyTorch → MLIR 的转换具有挑战性

1.3 ATen 算子到 MLIR 的映射

python

# ATen 算子和 MLIR dialect 的对应关系
"""
┌────────────────────────────────────────────────────────────────────┐
│  ATen 算子                                                        │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐        │
│  │ matmul   │  │  relu    │  │  conv2d   │  │  layer_norm│       │
│  └────┬─────┘  └────┬─────┘  └────┬─────┘  └────┬─────┘        │
└───────┼─────────────┼─────────────┼─────────────┼────────────────┘
        │             │             │             │
        ▼             ▼             ▼             ▼
┌────────────────────────────────────────────────────────────────────┐
│  Torch IR (Torch-MLIR 的中间层)                                     │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐        │
│  │ aten::   │  │ aten::   │  │ aten::   │  │ aten::   │        │
│  │ mm/mv    │  │ relu     │  │ conv2d   │  │ layer_norm│       │
│  └──────────┘  └──────────┘  └──────────┘  └──────────┘        │
└────────────────────────────────────────────────────────────────────┘
        │             │             │             │
        ▼             ▼             ▼             ▼
┌────────────────────────────────────────────────────────────────────┐
│  导出目标 dialects (用户选择)                                        │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐        │
│  │  TOSA    │  │  Linalg  │  │   MHLO   │  │  Stablehlo│       │
│  │  (edge)  │  │ (tensor) │  │  (XLA)   │  │ (HLO v2) │        │
│  └──────────┘  └──────────┘  └──────────┘  └──────────┘        │
└────────────────────────────────────────────────────────────────────┘
"""

第2节 Torch-MLIR 架构详解

2.1 整体架构

┌─────────────────────────────────────────────────────────────────────────────┐
│                        Torch-MLIR Architecture                               │
└─────────────────────────────────────────────────────────────────────────────┘

    ┌──────────────────────────────────────────────────────────────────────┐
    │  PyTorch Model (Python)                                              │
    │  import torch                                                        │
    │  model = torch.nn.Transformer(...)                                    │
    │  model.eval()                                                        │
    └──────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
    ┌──────────────────────────────────────────────────────────────────────┐
    │  Torch-MLIR: torch_mlir.compile()                                   │
    │  ┌────────────────────────────────────────────────────────────────┐  │
    │  │  Step 1: FX Graph Extraction (via TorchDynamo)                  │  │
    │  │  torch.fx.Graph → torch_mlir.fb.GraphModule                    │  │
    │  └────────────────────────────────────────────────────────────────┘  │
    │                                    │                                  │
    │                                    ▼                                  │
    │  ┌────────────────────────────────────────────────────────────────┐  │
    │  │  Step 2: Conversion to Torch IR                                 │  │
    │  │  FX Graph → torch::lazy (LazyTensor) → torch dialect           │  │
    │  │                                                                 │  │
    │  │  torch Dialect (high-level, PyTorch semantics)                 │  │
    │  │  - torch.operator: 对应 ATen 算子                               │  │
    │  │  - torch.tensor: 动态 shape tensor                              │  │
    │  │  - torch.nn_module: 模块结构                                   │  │
    │  └────────────────────────────────────────────────────────────────┘  │
    │                                    │                                  │
    │                                    ▼                                  │
    │  ┌────────────────────────────────────────────────────────────────┐  │
    │  │  Step 3: Functionalization (Pure Functions)                   │  │
    │  │  - 消除 in-place mutation                                      │  │
    │  │  - 消除 tensor aliasing                                        │  │
    │  │  - 转换为纯函数式计算                                           │  │
    │  └────────────────────────────────────────────────────────────────┘  │
    │                                    │                                  │
    └────────────────────────────────────┼────────────────────────────────┘
                                          │
                    ┌─────────────────────┼─────────────────────┐
                    │                     │                     │
                    ▼                     ▼                     ▼
    ┌──────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
    │  TOSA Dialect         │ │  Linalg on Tensors   │ │  MHLO Dialect         │
    │  (Tensor Operator     │ │  (Optimizable        │ │  (XLA-style           │
    │   Set Architecture)   │ │   tensor ops)       │ │   representation)     │
    └──────────────────────┘ └──────────────────────┘ └──────────────────────┘
                    │                     │                     │
                    └─────────────────────┼─────────────────────┘
                                          ▼
    ┌──────────────────────────────────────────────────────────────────────┐
    │  Backend Integration                                                 │
    │  ┌────────────┐  ┌────────────┐  ┌────────────┐  ┌────────────┐   │
    │  │  IREE      │  │  LLVM      │  │  SPIR-V    │  │  VM        │   │
    │  │  (Runtime) │  │  (CPU/GPU) │  │  (Vulkan)   │  │  (Embed)   │   │
    │  └────────────┘  └────────────┘  └────────────┘  └────────────┘   │
    └──────────────────────────────────────────────────────────────────────┘
                                          │
                                          ▼
    ┌──────────────────────────────────────────────────────────────────────┐
    │  Target Code                                                         │
    │  x86asm | ARMasm | WASM | C++ | ...                                 │
    └──────────────────────────────────────────────────────────────────────┘

2.2 Torch Dialect 详解

Torch Dialect 是 Torch-MLIR 定义的 MLIR Dialect，用于表示 PyTorch 语义：

mlir

// Torch Dialect 示例：简单的 Linear 层
module {
  // torch.nn.Linear 的 Torch Dialect 表示
  func.func @forward(%input: !torch.tensor<[?, 784], f32>,
                     %weight: !torch.tensor<[10, 784], f32>,
                     %bias: !torch.tensor<[10], f32>) 
      -> !torch.tensor<[?, 10], f32> {
    
    // torch.mm: 矩阵乘法
    %0 = torch.aten.mm %input, %weight 
        : !torch.tensor<[?, 784], f32>, 
          !torch.tensor<[784, 10], f32> -> 
          !torch.tensor<[?, 10], f32>
    
    // torch.add: 加偏置 (tensor + scalar broadcast)
    %1 = torch.aten.add.Tensor %0, %bias, %cst_1 
        : !torch.tensor<[?, 10], f32>, 
          !torch.tensor<[10], f32>, 
          !torch.int -> 
          !torch.tensor<[?, 10], f32>
    
    return %1 : !torch.tensor<[?, 10], f32>
  }
  
  // 常量定义
  %cst_1 = torch.constant.int 1
}

第3节 Functionalization（函数化）

3.1 为什么需要 Functionalization

PyTorch 的 Eager Mode 有副作用（in-place mutation），而 MLIR 是函数式的 SSA（Static Single Assignment）。Functionalization 把有副作用的 PyTorch 代码转换为纯函数式：

python

# PyTorch 有副作用的代码
class ModelWithInPlace(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.running_mean = None  # buffer，in-place 更新
    
    def forward(self, x):
        # In-place 操作：更新 running_mean
        self.running_mean = 0.9 * self.running_mean + 0.1 * x.mean()
        return self.layer(x)

# Functionalization 后变成纯函数
def functional_forward(params, buffers, x):
    # 原来的 in-place 更新变成返回值
    new_running_mean = 0.9 * buffers['running_mean'] + 0.1 * x.mean()
    output = layer_forward(params, x)
    return output, new_running_mean  # 返回新值

# 外部调用者负责更新 buffer
output, new_running_mean = functional_forward(params, buffers, x)
buffers['running_mean'] = new_running_mean  # 显式更新

3.2 Functionalization 的具体转换

python

# Functionalization 转换规则
"""
转换前（In-place PyTorch）:
  a = torch.tensor([1, 2, 3])
  b = a.add_(1)     # add_ 是 in-place，a 和 b 指向同一内存
  c = a * 2         # c 基于被修改后的 a
  
转换后（Functional）:
  a_orig = torch.tensor([1, 2, 3])
  a_modified = torch.add(a_orig, 1)  # 返回新 tensor
  c = torch.mul(a_modified, 2)       # 基于新 tensor
  
关键转换:
  tensor.add_(other)  →  (new_tensor, tensor)
  tensor.copy_(other)  →  new_tensor
  tensor.mul_(scalar)  →  (new_tensor, tensor)
"""

# 实际的 MLIR 转换示例
"""
Functionalization 前 (torch dialect):
  %1 = torch.aten.add_.Tensor %a, %b
  %2 = torch.aten.mul.Tensor %a, %c  # %a 已被修改

Functionalization 后:
  %1, %updated_a = torch.aten.add.Tensor %a, %b  # 返回新值和更新后的张量
  %2 = torch.aten.mul.Tensor %updated_a, %c
"""

Functionalization 后，可能还需要 Refinement 把纯函数优化回 PyTorch 兼容形式：

python

# Refinement 过程
"""
阶段1：Functionalization
  In-place ops → Pure functional form

阶段2：Optimization
  在纯函数形式上进行算子融合、常数折叠等优化

阶段3：Refinement
  把优化后的纯函数转换回 PyTorch 可执行的形式
  - 消除不必要的中间张量分配
  - 恢复部分 in-place 操作以节省内存
  - 保留别名关系用于原地更新
"""

第4节输出类型对比

4.1 TOSA（Tensor Operator Set Architecture）

TOSA 是面向嵌入式和边缘计算的标准化算子集：

python

# TOSA 的特点
tosa_features = {
    "目标": "嵌入式/边缘设备部署",
    "算子覆盖": "标准化 50+ 基础算子",
    "精度支持": "INT4/INT8/FP16/FP32",
    "硬件支持": "CPU, GPU, DSP, NPU",
    "优点": "硬件友好，功耗低",
    "缺点": "算子种类有限，灵活性低"
}

# TOSA 算子示例（MLIR）
"""
// TOSA matmul（输入需要量化）
func.func @matmul_tosa(%a: tensor<1x64x128xui8, #tosa.scales<...>>,
                        %b: tensor<1x64x256xui8, #tosa.scales<...>>) 
    -> tensor<1x64x256xi32, #tosa.scales<...>> {
  %0 = tosa.matmul %a, %b : (tensor<1x64x128xui8>, tensor<1x64x256xui8>)
                              -> tensor<1x64x256xi32>
  return %0 : tensor<1x64x256xi32>
}
"""

4.2 Linalg on Tensors

Linalg on Tensors 是可优化的 MLIR 张量操作表示：

python

# Linalg 的特点
linalg_features = {
    "目标": "通用优化编译器",
    "算子覆盖": "可组合的通用算子",
    "优化能力": "算子融合、tiling、vectorization",
    "层次结构": "Linalg → Affine → SCF → LLVM",
    "优点": "强大的优化能力，层次化 lowering",
    "缺点": "中间表示较复杂"
}

# Linalg 算子示例
"""
// Linalg matmul（Linalg Dialect）
func.func @matmul_linalg(%A: tensor<128x256xf32>,
                          %B: tensor<256x512xf32>,
                          %C: tensor<128x512xf32>) 
    -> tensor<128x512xf32> {
  
  %D = linalg.matmul 
       ins(%A, %B: tensor<128x256xf32>, tensor<256x512xf32>)
       outs(%C: tensor<128x512xf32>)
    -> tensor<128x512xf32>
  
  return %D : tensor<128x512xf32>
}

// 后续优化会将其 lower 到 Affine Loop Fusion
"""

4.3 MHLO（Meta HLO）

MHLO 是 XLA 风格的表示，用于与 JAX/TensorFlow 生态集成：

python

# MHLO 的特点
mhlo_features = {
    "目标": "XLA/JAX 生态兼容",
    "算子覆盖": "与 HLO 1:1 对应",
    "特性": "动态 shape 支持",
    "应用": "需要 JAX 互操作时",
    "优点": "与 XLA 生态无缝集成",
    "缺点": "优化 passes 相对较少"
}

# MHLO 算子示例
"""
// MHLO matmul（与 XLA HLO 对应）
func.func @matmul_mhlo(%A: tensor<128x256xf32>,
                        %B: tensor<256x512xf32>) 
    -> tensor<128x512xf32> {
  
  %D = "mhlo.dot"(%A, %B) 
       : (tensor<128x256xf32>, tensor<256x512xf32>) 
         -> tensor<128x512xf32>
  
  return %D : tensor<128x512xf32>
}
"""

4.4 三种输出类型对比表

特性	TOSA	Linalg on Tensors	MHLO
设计目标	边缘/嵌入式部署	通用编译器优化	XLA/JAX 生态兼容
算子数量	50+ 标准化算子	可组合通用算子	与 HLO 1:1 对应
动态 Shape	❌ 固定 shape	⚠️ 部分支持	✅ 完整支持
量化支持	✅ 内置	⚠️ 需要额外处理	⚠️ 需要额外处理
优化能力	中等	强（多层 lowering）	中等
硬件支持	CPU/GPU/DSP/NPU	通用	TPU/GPU/CPU
代码生成	直接编译	Linalg → Affine → SCF → LLVM	MHLO → XLA Backend
适用场景	边缘推理、功耗敏感	服务器端优化编译	JAX 互操作
与 IREE 集成	✅ 完整支持	✅ 完整支持	✅ 完整支持

4.5 如何选择输出类型

python

# 选择决策树
def choose_output_type(model, target_device):
    if target_device == "edge/raspberry_pi":
        # 边缘设备：选择 TOSA
        return "tosa"
    elif target_device == "server_cpu" or target_device == "gpu":
        # 服务器端优化：选择 Linalg
        return "linalg-on-tensors"
    elif need_jax_interop():
        # 需要与 JAX 互操作：选择 MHLO
        return "mhlo"
    else:
        # 默认：Linalg（有最强优化能力）
        return "linalg-on-tensors"

# 示例：部署到树莓派
import torch_mlir

# 树莓派是边缘设备，选择 TOSA
module = torch_mlir.compile(
    torch_model,
    example_input,
    output_type="tosa"  # 边缘部署用 TOSA
)

# 示例：服务器端优化
# 需要最强优化能力，选择 Linalg
module = torch_mlir.compile(
    torch_model,
    example_input,
    output_type="linalg-on-tensors"  # 服务器端用 Linalg
)

第5节 IREE 集成

5.1 IREE 是什么

IREE（Intermediate Representation Execution Environment） 是一个 MLIR 端到端编译器，用于部署 ML 模型到各种硬件：

python

# IREE 的编译流程
"""
PyTorch Model
     │
     ▼
Torch-MLIR (→ TOSA/Linalg/MHLO)
     │
     ▼
IREE Input (HAL dialect + Flow dialect)
     │
     ├──▶ Flow → Stream (运行时分配)
     ├──▶ Stream → HAL (硬件抽象)
     └──▶ HAL → Device (CPU/GPU/Vulkan)
     
最终产物：
  - .vmfb: IREE 虚拟机格式（跨平台）
  - 或者直接编译到 .o (LLVM)
"""

5.2 完整的 Torch-MLIR → IREE 工作流

python

# 完整示例：PyTorch 模型 → IREE 部署
import torch
import torch_mlir
import iree.compiler
import iree.runtime

# Step 1: 定义 PyTorch 模型
class SimpleTransformer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = torch.nn.Embedding(1000, 256)
        self.transformer = torch.nn.TransformerEncoderLayer(
            d_model=256, nhead=8, batch_first=True
        )
        self.output = torch.nn.Linear(256, 10)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer(x)
        return self.output(x[:, 0])  # 取第一个 token

model = SimpleTransformer().eval()

# Step 2: 创建示例输入
example_input = torch.randint(0, 1000, (1, 32))

# Step 3: 编译为 Torch-MLIR
print("Step 1: Compiling to Torch-MLIR...")
module = torch_mlir.compile(
    model,
    example_input,
    output_type="linalg-on-tensors"  # 选择 Linalg 作为优化目标
)

# Step 4: 保存为 .mlir 文件
with open("/tmp/module.mlir", "w") as f:
    f.write(str(module))

# Step 5: 用 IREE 进一步编译
print("Step 2: Compiling to IREE...")

# 编译为 CPU 目标
iree.compiler.compile_file(
    input_file="/tmp/module.mlir",
    output_file="/tmp/module_cpu.vmfb",
    target_backends=["llvm-cpu"],
    extra_args=["-iree-hal-target-device=local-task"]
)

# 编译为 Vulkan/SPIR-V 目标（GPU）
iree.compiler.compile_file(
    input_file="/tmp/module.mlir",
    output_file="/tmp/module_vulkan.vmfb",
    target_backends=["vulkan-spirv"],
    extra_args=[
        "-iree-vulkan-target-triple=radeon-rx6700xt"
    ]
)

# Step 6: 在 IREE Runtime 上运行
print("Step 3: Running on IREE Runtime...")

# 加载编译产物
config = iree.runtime.Config("local-task")
vm_module = iree.runtime.load_module("/tmp/module_cpu.vmfb", config)

# 准备输入
import numpy as np
input_np = np.random.randint(0, 1000, (1, 32)).astype(np.int64)

# 调用
results = vm_module.main_batch_function(input_np)
print(f"Output shape: {results.shape}")
print(f"Output dtype: {results.dtype}")

第6节 TorchDynamo 导出路径

6.1 torch.compile → FX Graph → MLIR

PyTorch 2.0 引入了 torch.compile，为 Torch-MLIR 提供了新的导出路径：

python

# 传统路径 vs torch.compile 路径
"""
传统路径 (TorchScript):
  PyTorch Model → TorchScript (tracing/scripting) → torch-mlir
  
新路径 (torch.compile):
  PyTorch Model → TorchDynamo (graph capture) → FX Graph → torch-mlir
"""

# 示例：torch.compile 导出
import torch
import torch_mlir

model = MyModel().eval()

# 使用 torch.compile 捕获计算图
compiled_model = torch.compile(model, backend="eager")  # 先试试 eager

# 用 MLIR 特定的 backend
# 注意：torch-mlir 提供自己的 torch.compile backend

6.2 TorchDynamo Graph Capture

python

import torch
import torch.fx
from torch._dynamo import optimize

# TorchDynamo 的 graph capture 过程
@optimize(nopython=True)
def compiled_forward(x):
    return model(x)

# 捕获的 FX Graph
"""
原始代码:
  def forward(x):
      y = layer1(x)
      z = layer2(y)
      return layer3(z)

FX Graph capture:
  Graph:
    %0 : torch.Tensor = l__self___layer1
    %1 : torch.Tensor = call_module %0, (%x)
    %2 : torch.Tensor = l__self___layer2
    %3 : torch.Tensor = call_module %2, (%1)
    %4 : torch.Tensor = l__self___layer3
    %5 : torch.Tensor = call_module %4, (%3)
    return %5
"""

升华

┌─────────────────────────────────────────────────────────────────────────────┐
│                      Torch-MLIR 核心原则                                    │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  1. 层次化转换                                                              │
│     PyTorch → Torch IR → 目标 Dialect → 硬件代码                           │
│                                                                             │
│  2. Functionalization 是桥梁                                               │
│     把有副作用的 PyTorch 转换为纯函数式，是 MLIR 兼容的关键                 │
│                                                                             │
│  3. 选择正确的输出类型                                                      │
│     TOSA（边缘）、Linalg（优化）、MHLO（XLA 生态）                          │
│                                                                             │
│  4. IREE 提供端到端部署                                                     │
│     从 PyTorch 到可执行文件，跨越多种硬件                                    │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

"AI 可查 vs 必须理解"清单

必须理解（不理解就等于不会）：

🔴 为什么需要 Functionalization：PyTorch 有副作用（in-place），MLIR 是纯函数式
🔴 三种输出类型的区别：TOSA（边缘）、Linalg（优化）、MHLO（XLA 兼容）
🔴 Torch Dialect 的作用：Torch-MLIR 定义的中间层，表示 PyTorch 语义
🔴 IREE 的定位：端到端 ML 编译器，把 MLIR 进一步编译到目标硬件
🔴 ATen 是什么：PyTorch 的底层张量库，CPU/GPU 统一的算子接口

AI 可查（知道去哪查就行）：

✅ TOSA 具体算子列表：TOSA 规范文档列出 50+ 标准化算子
✅ Linalg 到 Affine 的具体 lower 规则：MLIR 官方文档
✅ IREE HAL 抽象细节：IREE 架构文档
✅ Torch Dialect 完整的 op 定义：Torch-MLIR 源码 dialect definition
✅ 具体硬件的 IREE 支持情况：IREE 官方 supported targets 页面

学习状态：🟡 开始学习

Torch-MLIR：从 PyTorch 算子到 MLIR 方言 / Torch-MLIR from PyTorch Operators to MLIR Dialects ​

第1节 PyTorch 算子系统 ​

1.1 ATen：PyTorch 的核心算子库 ​

1.2 PyTorch IR 的特殊性 ​

1.3 ATen 算子到 MLIR 的映射 ​

第2节 Torch-MLIR 架构详解 ​

2.1 整体架构 ​

2.2 Torch Dialect 详解 ​

第3节 Functionalization（函数化） ​

3.1 为什么需要 Functionalization ​

3.2 Functionalization 的具体转换 ​

3.3 Refinement（精化） ​

第4节 输出类型对比 ​

4.1 TOSA（Tensor Operator Set Architecture） ​

4.2 Linalg on Tensors ​

4.3 MHLO（Meta HLO） ​

4.4 三种输出类型对比表 ​

4.5 如何选择输出类型 ​

第5节 IREE 集成 ​

5.1 IREE 是什么 ​

5.2 完整的 Torch-MLIR → IREE 工作流 ​

第6节 TorchDynamo 导出路径 ​

6.1 torch.compile → FX Graph → MLIR ​

6.2 TorchDynamo Graph Capture ​

升华 ​

"AI 可查 vs 必须理解"清单 ​

Torch-MLIR：从 PyTorch 算子到 MLIR 方言 / Torch-MLIR from PyTorch Operators to MLIR Dialects

第1节 PyTorch 算子系统

1.1 ATen：PyTorch 的核心算子库

1.2 PyTorch IR 的特殊性

1.3 ATen 算子到 MLIR 的映射

第2节 Torch-MLIR 架构详解

2.1 整体架构

2.2 Torch Dialect 详解

第3节 Functionalization（函数化）

3.1 为什么需要 Functionalization

3.2 Functionalization 的具体转换

3.3 Refinement（精化）

第4节输出类型对比

4.1 TOSA（Tensor Operator Set Architecture）

4.2 Linalg on Tensors

4.3 MHLO（Meta HLO）

4.4 三种输出类型对比表

4.5 如何选择输出类型

第5节 IREE 集成

5.1 IREE 是什么

5.2 完整的 Torch-MLIR → IREE 工作流

第6节 TorchDynamo 导出路径

6.1 torch.compile → FX Graph → MLIR

6.2 TorchDynamo Graph Capture

升华

"AI 可查 vs 必须理解"清单