Phobos：支持tile-DAG的微型无尺度内核语言

背景

作者从小对神经网络和AI充满热情，曾在卧室用劣质摄像头和1.8 GHz CPU训练目标识别软件。大学就读于比勒费尔德大学认知信息学专业，但仅四天后就因巴黎一家初创公司的工作机会而辍学。此后多年，他专注于优化、编译器以及机器学习相关领域，虽持续关注前沿研究，但自认是AI前沿的局外人。现在，他带着一张老旧的RTX 2080 SUPER显卡重新出发，渴望深入学习底层GPU概念。

作者坚信要真正理解某件事就必须亲自动手。凭借编译器、优化和LLVM的经验，他认为编写一个面向GPU的小型语言和编译器是可行的。他也有足够的软件工程和编译器经验，知道必须谨慎行事，避免项目失控占用太多时间。编译器尤其棘手：如果不能从头到尾编译一个程序，就什么都看不到。因此他决定走些捷径，不处理链接或阶段模型。

核心内容

Phobos 是一个受 Triton 启发的小型内核语言，直接编译到 PTX，能在 NVIDIA GPU 上运行。在 RTX 2080 SUPER 上，其 SGEMM 性能达到 cuBLAS 的 76%（或理论峰值 GFLOP/s 的 74%）。Phobos 天然支持分布式 tile-DAG（瓦片有向无环图）设计，但作者仅在一台机器上验证了集群原型，未提供多节点基准测试。该项目是个人研究，旨在深入理解底层GPU概念，实现局部 tile 优化时才意识到其分布式计算潜力。

语言设计围绕 Tile（张量块）展开，支持一等公民的 tensor/tile 类型。以下是两个示例程序：

向量加法：c = a + b

@autotune(BLOCK in [1024])
kernel add(a: tensor<f32>[N], b: tensor<f32>[N], c: tensor<f32>[N]) {
    let base = program_id(0) * BLOCK
    c[base :+ BLOCK] = a[base :+ BLOCK] + b[base :+ BLOCK]
}

矩阵乘法：C = α A B + β C

@autotune(TILE_M in [32, 128], TILE_N in [64, 256], TILE_K in [16, 64])
kernel sgemm(A: tensor<f32>[M, K],
             B: tensor<f32>[K, N],
             C: tensor<f32>[M, N],
             alpha: f32,
             beta: f32) {
    let pm = program_id(0)
    let pn = program_id(1)
    var acc: tile<f32>[TILE_M, TILE_N] = 0.0
    for kt in range(0, K, TILE_K) {
        var a = A[pm * TILE_M :+ TILE_M, kt :+ TILE_K]
        var b = B[kt :+ TILE_K, pn * TILE_N :+ TILE_N]
        acc += dot(a, b)
    }
    let c_old = C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N]
    C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] = alpha * acc + beta * c_old
}

作者自嘲这“基本上就是迷你版Triton”。

为了实现编译链，作者探索了CUDA底层技术。他关注的是PTX（并行线程执行）而非高层CUDA C++代码。PTX是NVIDIA定义的稳定编程模型和指令集，高层编译器（如CUDA C/C++）生成PTX指令，再被优化和翻译为原生目标指令。LLVM已有NVPTX后端可生成PTX。此外还有MLIR（多级中间表示）这条路线。作者最终选用了MLIR路径，但配置过程极其痛苦：使用Rust的inkwell（LLVM绑定）和melior（MLIR绑定），花费大量时间修修补补（yak-shaving），包括修补melior、mlir-sys和tblgen-rs等crate，直到能编译简单程序。

由于可能遇到许多问题，作者采用“从右向左”的开发策略：先从最隔离的任务开始——在GPU上运行PTX代码。他手动编写了向量加法的PTX代码，并用cust crate运行。首先通过nvidia-smi查询GPU计算能力（7.5，Turing架构，缺float32 tensor core但有16位精度）。然后手动编写PTX代码，用ptxas add.ptx直接汇编验证正确性，隔离Rust和CUDA crate的问题。也可以通过nvcc -ptx将CUDA内核编译为PTX来辅助开发。

示例PTX代码（向量加法）展示了从ld.param加载参数、计算线程索引、条件分支、从全局内存加载数据、浮点加法、存储结果等操作。作者指出这是唯一一次手动编写PTX，后续将交由编译器生成。

关键要点

• 项目定位：一个个人研究项目，旨在通过实践学习底层GPU概念，非生产级产品。
• 语言特点：面向Tile（张量块）设计，支持@autotune自动调优，语法简洁，类似微型Triton。
• 编译链：从Phobos源码 → Lexer → Parser → Codegen → MLIR → PTX → GPU。作者选择了MLIR路径而非直接LLVM IR，尽管配置MLIR耗时两天，而配置inkwell只需几小时。
• 性能表现：在RTX 2080 SUPER上，SGEMM性能达cuBLAS的76%或理论峰值的74%，属于可接受水平。
• 分布式潜力：Phobos自然映射到分布式tile-DAG设计，但作者仅单机验证，无多节点基准。
• 开发策略：采用“从右向左”方法，先手动编写并运行PTX代码，确保底层路径可行，再逐步构建前端。
• 技术栈：使用Rust语言，inkwell（LLVM绑定）、melior（MLIR绑定）及其底层crate（mlir-sys, tblgen-rs）。
• 已知限制：缺乏链接和阶段模型，不处理多节点，仅支持float32（但GPU缺少float32 tensor core），仅在一台机器上测试。

意义与影响

• 对个人成长：作者通过亲手构建一个完整的GPU编译工具链，深入理解了PTX、MLIR、LLVM以及GPU的执行模型。这种“脏手”实践极有助于掌握底层概念。
• 对开源社区：Phobos提供了一个轻量级的、面向Tile的GPU内核语言原型，可作为教育和实验平台。它展示了如何用Rust结合MLIR快速构建自定义GPU编译前端。
• 对分布式计算：虽然当前未验证多节点，但Phobos的tile-DAG映射天然适合分布式执行。这为未来探索异构集群上的张量计算提供了基础。
• 与Triton的对比：Phobos与Triton理念相似但更精简，更侧重于个人学习而非工程化。它的代码量和复杂度可控，适合作为理解Triton/MLIR等技术的踏脚石。
• 对编译器设计：作者强调“从右向左”的开发策略，先保证后端可执行，再逐步扩展前端，这种自底向上的方法在复杂系统开发中极具参考价值。