今年我们正处在一个十字路口，预测未来走向并非易事。本文所勾勒的路径，很可能并非 agentic 时代 ML 编译器的唯一前进方向。随着 ML 编译器迈入这一新阶段，我们必须谨慎避免那些会限制 agent 能力的僵化抽象，同时也要认识到，完全缺乏分层会导致不必要的 token 开销和草率的结果。

我们的观点是，只要避免限制 agent 能力的僵化抽象，可组合的分层就能成功。一个 ML 编译器基础可以在支持 ML 编译器和 agent 的研发方面发挥重要作用。我们希望 TIRx 能提供这样一个基础，并与更广泛的 ML 编译器和 agentic 开发生态系统协同工作。

上一节讨论了整体设计动机。本节将通过代码示例介绍一些技术设计要素。为便于阅读，技术细节折叠在下方。

TIRx 提供了完整的接口，用于在 Python 中构造典型的张量程序：

from tvm_tirx.script import tirx as T n = T.dynamic("n") def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)): for i in range(n): B[i] = A[i] + 1

上面的示例构造了一个函数对象，表示在 CPU 上执行加一操作。TIRx 提供了内置 API，可将该函数编译为能与 Python 数组 API 生态中的库（如 NumPy 和 PyTorch）互操作的代码。

import torch import tvm_tirx as tirx add_one_compiled = tirx.compile(add_one, target="llvm") a = torch.arange(10, dtype=torch.float32) b = torch.empty(a.shape, dtype=a.dtype) add_one_compiled(a, b)

我们的设计力求最小化、忠实且可预测。调用 tirx.compile 不会施加额外的启发式优化；程序会直接传递给 LLVM 编译器工具链以生成 CPU 代码。

到目前为止的示例展示了如何将 TIRx 用作典型的黑盒领域特定编译器。如前所述，我们的目标不仅是暴露编程接口，还要向开发者和 agent 暴露编译器 API 以及每一个中间 AST 对象。我们将每个 AST 节点存储在一个由 TVM FFI 支撑的跨语言对象中，使其可以直接从 Python、C++ 和 Rust 等语言访问。以下代码展示了如何检查程序结构，例如类型和形状值：

查询参数 A 的类型

ty_A = add_one.params[0].ty

查询参数 B 的长度

n = add_one.params[1].ty.shape[0]

遍历、映射与分析的通用工具

直接访问编译器内部对象使我们能够进一步检查和变换程序。我们使用一套内置于 TIRx 和 TVM FFI 的轻量级通用工具来遍历、比较和变换程序结构。例如，以下代码片段将动态维度 n 特化为 128：

import tvm_ffi import tvm_tirx add_one_specialized = tvm_ffi.structural_map( add_one, (tvm_tirx.ir.Var, lambda v: 128 if v.same_as(n) else v), ) add_one_specialized.show() def add_one(A: T.Tensor((128,), T.float32), B: T.Tensor((128,), T.float32)): for i in range(128): B[i] = A[i] + 1

符号形状和索引在 kernel 和更高层程序中都很常见。TIRx 的表示允许 TensorType 这样的类型依赖于表达式值，并且我们提供了通用工具，在程序分析过程中检查符号关系。得益于 TileLang 社区的反馈，我们还增加了带有符号约束上下文和更强 Z3 证明的结构遍历变体。这些工具被设计为可处理动态 TVM FFI 对象，因此随着 AST 节点集合的扩展（无论是在主项目中还是在下游项目中），它们都能继续正常工作。

在编译的每个阶段，程序都可以用 Python 脚本语法打印出来，如上所示。这使得 agent 和开发者能够检查中间程序、试验变换并构建新的程序。

虽然编译器基础提供了最小化的骨干，但 TIRx 的设计允许在树内和下游进行新的定制。扩展 TIRx 基础的主要方式是引入新的类型和操作。下游项目可以将新类型定义为 TVM-FFI 对象，这些对象继承原始类型节点并携带额外的类型和表达式信息。例如，我们引入了 tirx.TensorMapType 来支持 CUDA 程序中的 TMA tensor map。

由于新类型可以携带依赖于其他类型和表达式的信息，它们能够在编译的中间阶段之间保留这些信息。一种典型的方法是将这些中间类型降级为现有类型，或为新类型添加专用的后端代码生成支持。为保持简单性和一致性，我们允许在脚本语法中直接使用 Python 类型构造器。例如，T.Tensor((n,), T.float32) 只是用元组和 T.float32 调用类型构造器，以获得具体类型的实例。

大多数扩展会引入新的操作（op）。一个 op 可以用于 Call AST 节点（类似于内建函数调用），也可以用于 RegionStmt，定义带有额外语义（如线程启动或 kernel 区域）的区域。我们通过引入 device_entry 区域 op，并将对 T.ptx.add.rn.f32 的调用表示为对应 PTX op 的 Call 节点，构建了最小化的 CUDA 和 PTX 支持。

from tvm_tirx.script import tirx as T n = T.dynamic("n") def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)): T.device_entry( launch=T.cuda.LaunchConfig(grid=((n + 255) // 256,), block=(256,)) ) i = T.cuda.block_idx("x") * 256 + T.cuda.thread_idx("x") if i < n: T.ptx.add.rn.f32(B[i], A[i], T.float32(1))

作为这些动态扩展能力的一个例子，我们直接通过 Python 注册了 T.ptx 命名空间中的大部分 op，使我们能够快速更新编译器支持而无需修改核心库。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）