跳到主要内容

Triton + RISC-V

如何深入学习 Triton 的编译器技术栈？

为此，我设定了一个 Triton + RISC-V 的学习目标：尝试将一个简单的向量加法算子，通过 Triton 编译并运行到 RISC-V 仿真环境上。

结果，发现将 Triton 引入 RISC-V 生态是当前 AI 编译器与底层硬件领域非常活跃且前沿的方向。

目前，随着 RISC-V 向量扩展（RVV）的成熟以及 AI 算力泛化的需求，社区（包括 PyTorch 团队、阿里平头哥玄铁、Terapines、SpacemiT 等）正在大力推进 Triton 在 RISC-V 上的落地。

因此，我想与大家分享一下目前 Triton + RISC-V 的技术路径、演进情况、以及实践方法。

1. 技术路径​

Triton 在 RISC-V 上的技术路径经历了从“能跑”到“高性能”的演进，核心挑战是如何将 Triton 基于 Tile（分块） 的编程模型高效映射到 RISC-V 的 SIMD（向量） 架构上。

1.1 编译管线​

目前主流的技术路径是基于 MLIR 构建的：

  • Triton AST -> TTIR：Python 代码被解析为 Triton 特有的中间表示（TTIR）。
  • TTIR -> MLIR Dialects：由于目标是 RISC-V CPU 或类似 CPU 的 AI 加速器，目前有两种主要降级路径：
    • 直接降级到 Vector Dialect（主导方向）：直接映射为底层向量操作，性能较好。
    • 降级到 Linalg/Tensor Dialect (TritonShared 路径)：通用性更强，但目前针对特定硬件的极致性能尚未完全释放。
  • MLIR -> LLVM IR：将上述表示转换为泛化的 LLVM IR。
  • LLVM IR -> RISC-V Machine Code：LLVM 后端通过指令选择，将 IR 映射为 RISC-V 的 RVV 指令集。

1.2 核心技术依赖：RVV​

  • Triton 代码中大量的张量并行计算（如矩阵乘法），在 RISC-V 上完全依靠 RVV (RISC-V Vector Extension) 来加速。
  • 固定长度 vs 可变长度：Triton 编译器天然倾向于生成固定向量长度（Fixed Vector Length）的掩码操作，而 RVV 是基于 VLEN（可变向量长度）设计的。目前编译路径中会强制指定 VLEN（如 256 位）进行对齐匹配。

1.3 性能表现与瓶颈​

根据最近（2024年底至2025年）在 SpacemiT K1 板等硬件上的测试反馈：

  • 单线程性能：Triton 编译出的内核（通过 LLVM/ZCC）在执行类似 Softmax 或 Matmul 算子时，性能已经基本追平甚至在某些场景下略超（2%~4%）基于 C 语言手写的 GCC/LLVM 自动向量化代码。
  • 多线程性能：通过 OpenMP 包装的 Triton 内核在多核扩展上表现优异，但在超多线程时（如 8 线程），可能会因为寄存器溢出 (Register Spill) 和上下文切换存储开销而略有性能损耗。
  • 瓶颈点：离散内存访问（Gather/Scatter）的向量化效率、寄存器分配压力、以及 Mask 掩码生成的冗余指令是当前重点优化的方向。

2. 演进情况​

根据 PyTorch 社区与玄铁等团队在 2026 年最新制定的 PyTorch + RISC-V Roadmap (#171659)，Triton 位于该生态发展的第四阶段（Phase 4），也是终极目标：

  1. 原生 Triton Backend for RISC-V： 社区正在开发原生 RISC-V 代码生成后端，意味着未来不需要强依赖于 Triton-CPU 的妥协路径，而是直接输出对 RISC-V 友好的汇编。
  2. 拥抱 RVM (Matrix Extension)： 除了 RVV（向量扩展），RISC-V 最新的矩阵扩展（RVM）将直接对接 Triton 的 Block 操作，极大提升 GEMM（矩阵乘法）算子的性能。
  3. 与 torch.compile 无缝集成： 最终目标是让 RISC-V 像 Nvidia GPU 一样，当你在 PyTorch 里调用 torch.compile 时，底层 Inductor 能够直接调用 Triton 编译出高性能的 RISC-V 算子，无需开发者手动写 C++ / 汇编内核。
  4. TileLang 的融合： 除了 Triton，社区还在探索 TileLang 对 RISC-V AI 专用加速器（NPU）的支持，为基于 RISC-V 指令集魔改的 AI 芯片提供统一的 Tile 级编程范式。

3. 实践方法​

由于 Triton 最初是为 Nvidia GPU (PTX) 设计的，想要在 RISC-V 跑起来，目前主要依赖 Triton-CPU 项目的底座并结合 RISC-V 交叉编译工具链。

3.1 硬件准备​

  • 支持 RVV 1.0 的开发板：目前最佳的验证平台是带有 RISC-V Vector 扩展（RVV 1.0）的开发板。例如搭载 SpacemiT K1 芯片的开发板（如进迭时空的生态板，支持 256-bit RVV）或阿里平头哥玄铁系列（如 C920 等支持高性能向量计算的内核）。
  • 模拟器：如果没有物理板，可以使用 QEMU RISC-V 模拟器跑通功能，但无法测试真实的算力性能。可参考『RISC-V 软件环境』。

3.2 软件与工具链配置​

  • 编译器前端：使用开源的 OpenAI Triton（基于 Triton-CPU 分支，该分支最初由 Intel/Meta 发起，现已用于扩展泛化 CPU/非 GPU 架构）。
  • 后端交叉编译器：需要支持 RVV 的编译器。目前常用的有最新的 RISC-V GCC (如 14.0/15.0)，或者基于 LLVM 的高度优化编译器（如 Terapines 的 ZCC）。

3.3 运行与开发工作流​

  • 编写 Triton Kernel：像往常一样用 Python DSL 编写算子（例如 gemv, softmax 或 matmul）。
  • 生成 LLVM IR：调用 Triton-CPU 的内置编译器，将 Triton 代码 JIT/AOT 编译。它会将 Triton 代码降级（Lower）为包含向量操作的 LLVM IR。
  • 交叉编译生成汇编：利用 x86 平台上的 RISC-V 交叉编译器，将上述 LLVM IR 编译成带有 RVV 向量指令（如 vle32.v, vfmac.vv）的 RISC-V 二进制目标文件或动态库。
  • 板端执行：将编译好的动态库或二进制文件推送到 RISC-V 开发板上运行验证，Triton 社区还会通过生成 OpenMP 包装代码来支持多线程执行。

4. 总结建议​

如果你现在想“玩起来”，建议基于 x86 主机搭建 Triton-CPU + LLVM RISC-V Toolchain，买一块平价的 RVV 1.0 单板计算机，尝试用 Triton 写一个简单的向量加法或 GEMV，通过交叉编译并在板子上运行 objdump 观察其生成的 vsetvli 和 vfmacc 等 RVV 汇编指令。这是目前切入 RISC-V AI 编译器底层最直观的练手方式。