XLA（只读源码）

主线位置：traced 主线的另一个终点。GraphCompiler::compile 得到 CompiledGraph 之后，通常交给 Runtime 去 prepare / run（见执行与缓存）；也可以交给 tenferro-xla，翻成 StableHLO，由外部的 PJRT 插件编译执行。本机没有 PJRT 插件，页上对 PJRT 执行的描述来自源码和上游文档，没有运行过。实际跑过的是两件事：StableHLO lowering（纯字符串生成，不需要插件），以及没有插件时 PJRT 入口返回什么。

flowchart TB
  tt["TracedTensor"] --> gc["GraphCompiler"]
  gc --> cg["CompiledGraph"]
  cg --> rt["Runtime 引擎"]
  cg --> low["lower_to_stablehlo"]
  low --> mlir["StableHLO 文本"]
  mlir --> exe["XlaExecutor"]
  exe --> plug["PJRT 插件"]
  plug --> dev["Host 或 CUDA"]

实验 crate 是 exp-tenferro8/gpu-xla：crates.io 上的 tenferro-xla 0.8.0 开 pjrt，加 tenferro-runtime 0.8.0，release，M2 Max、macOS 27、rustc 1.99.0。

它是什么

tenferro-xla 的 src/ 有 3526 行（不算测试），包描述是「Experimental StableHLO lowering and runtime PJRT plugin loading for tenferro static graph programs」。lib.rs 开头说清了它的位置：「an optional peer executor over tenferro-runtime SemanticProgram values. It does not implement Engine」。它不注册进 Runtime，不经过运行时的准备、调度和缓存，直接吃 CompiledGraph 里的语义程序。

从 v0.7.1 到 v0.8.0，这个 crate 的源码只改了两处：文档里的 TensorBackend 换成 Engine，以及新加的 erf 原语被明确拒绝。同一句文档还写着「it does not change native CPU, CUDA, or WebGPU execution」，WebGPU 在 v0.8.0 已经不存在了。

lowering：拼一段 MLIR 文本

lower_compiled_to_stablehlo 取出 CompiledGraph 里的语义程序，交给 lower_semantic_program：先给每个输入一个 %argN，再按顺序对每个运算调 lower_semantic_operation，用字符串拼出一行 stablehlo.* 指令。没有 MLIR 库，也没有 XLA 的任何依赖。gpu-xla 编一个 exp(a @ b)：

module {
  func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<2x2xf32> {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<3x2xf32>) -> tensor<2x2xf32>
    %1 = stablehlo.exponential %0 : tensor<2x2xf32>
    return %1 : tensor<2x2xf32>
  }
}

fingerprint 1aa86a0f09d14c92

a、b 是用 TracedTensor::from_vec_col_major 建的，在语义程序里是带默认值的输入，所以成了 @main 的两个参数，不是常量。StableHloModule 给文本算一个 SHA-256 指纹（上面打印了前 16 个十六进制字符）。

StableHLO 的类型只说逻辑形状。tenferro 的主机张量是列主序（见列主序与视图），布局由 PJRT 那一侧处理。上游设计文档 xla-backend.md 的「Shape and Layout」一节写道：上传时把列主序的 byte_strides 交给 PJRT_Client_BufferFromHostBuffer，下载时要列主序的 host_layout。dot_general 还有一处顺序差异：StableHLO 的批量 dot_general 输出是「[batch..., lhs_free..., rhs_free...]」，tenferro 的约定是「[lhs_free..., rhs_free..., batch...]」，lowering 在后面补一个 transpose。

能翻什么

dtype 只有 F32、F64（stablehlo_dtype），整数、Bool、复数一律拒绝。核心算子 23 个：Constant、Add、Mul、Neg、Div、Abs、Exp、Log、Sin、Cos、Tanh、Sqrt、Rsqrt、Pow、Expm1、Log1p、Convert、Reshape、BroadcastInDim、Transpose、ReduceSum、ReduceSumSquares、DotGeneral。扩展算子只有提供了「标准算子 lowering」的才行（einsum 在固定形状下展开成 dot_general，见扩展算子）。形状必须精确：compile_with_input_specs 给出的形状是精确的，dynamic_truncate 这种只有上界的输出不行。实测：

== upper-bound shape (dynamic_truncate)
Err: XLA lowering supports only exact static shapes; ExecOp::DynamicTruncate output 0 axis 0 is an upper bound
  kind: Validation(ShapeMismatch)
== ops and dtypes
sub: Err: XLA lowering does not support ExecOp::Sub: operation is outside the initial StableHLO lowering subset
maximum: Err: XLA lowering does not support ExecOp::Maximum: operation is outside the initial StableHLO lowering subset
erf: Err: XLA lowering does not support ExecOp::Erf: erf is a CHLO operation (chlo.erf); the StableHLO lowering does not emit CHLO
...
i64 add: Err: XLA lowering does not support dtype I64 in program input

x - y 翻不过去：Sub 原语比 XLA lowering 加得晚，之后一直没进表；上游用户指南 xla.md 列的「Phase 1 operation subset」里也没有它（那张单子同样漏了代码里支持的 ReduceSumSquares）。erf 是 0.8.0 新加的原语，lowering 里专门有一支拒绝它，理由写在错误信息里：StableHLO 没有 erf，它在 CHLO 方言里，而这个生成器只产出 StableHLO。

0.8.0 新加的激活、softmax、归一化不是原语，是在 composite.rs 里用核心算子拼出来的（见 算子级 AD），traced 程序里只有拼出来的那些算子。所以一个组合运算能不能交给 XLA，取决于它的分解用了哪些算子。

执行：每次都从头来

执行要开 pjrt feature，它只多一个 libloading 依赖。PjrtPlugin::load_path 打开插件文件、取 GetPjrtApi 符号；XlaExecutor::from_env 读环境变量 TENFERRO_PJRT_PLUGIN。没有插件时的三种情况：

== PJRT without a plugin
from_env: Err: environment variable TENFERRO_PJRT_PLUGIN is not set; set it to a PJRT plugin .so path  kind RuntimeState
load_path: Err: failed to load PJRT plugin from /nonexistent/pjrt_c_api_cpu_plugin.so: dlopen(/nonexistent/pjrt_c_api_cpu_plugin.so, 0x0005): tried: ... '/nonexistent/pjrt_c_api_cpu_plugin.so' (no such file), ...  kind Io
run without plugin: Err: PJRT execution requires an executor created from a loaded plugin  kind RuntimeState

都是类型化错误，不 panic，kind() 分得清「没配置」（RuntimeState）和「文件打不开」（Io）。没开 pjrt 时，from_env 返回 PjrtFeatureDisabled。

有插件时，run_many_with_inputs 每次调用都做全套：

    validate_inputs(program, inputs)?;
    let output_specs = output_specs(program)?;
    let module = lower_to_stablehlo(program)?;
    // ...
    let mut client = PjrtClient::create(api)?;
    let device = client.first_addressable_device()?;
    let mut executable = client.compile(module.as_str(), &output_specs)?;
    // ...
    let mut input_buffers = inputs
        .iter()
        .map(|input| client.buffer_from_host_tensor(input, device))
        .collect::<Result<Vec<_>>>()?;
    let mut output_buffers = executable.execute(device, &mut input_buffers, num_outputs)?;

校验输入、lowering、建 PJRT 客户端、编译、上传、执行、下载，函数返回时客户端和可执行文件都 drop。模块的指纹算了，但没有地方拿它做缓存键；xla-backend.md 把「Add executable cache entries keyed by StableHLO fingerprint, plugin identity, platform, compiler options, and layout policy」列在「Future Work」里，并写明这一阶段的包装「do not add a runtime-owned SubgraphCompiler implementation or a PJRT executable cache」。所以同一个程序跑两次就编译两次。

动脑筋答案

sigmoid: Err: XLA lowering does not support dtype Bool in instruction output
silu: Err: XLA lowering does not support dtype Bool in instruction output
gelu: Err: XLA lowering does not support ExecOp::Erf: erf is a CHLO operation (chlo.erf); the StableHLO lowering does not emit CHLO
gelu_tanh: Ok (22 lines)
softmax: Err: XLA lowering does not support ExecOp::ReduceMax: operation is outside the initial StableHLO lowering subset

• 只有 gelu_tanh 能翻：它的分解是乘、加、乘常数、tanh，全在 23 个算子里，翻出来 22 行。
• sigmoid 卡在 Bool 上。sigmoid 为了不溢出，用 e = exp(-|x|)，再按 x > 0 在 1 和 e 之间 select 分子；那个比较的结果是 Bool，lowering 不支持 Bool dtype（Compare、Select 本身也不在表里，dtype 检查先报了错）。silu 是 x * sigmoid(x)，原因相同。
• gelu 是精确版 x/2 * (1 + erf(x / sqrt(2)))，卡在 erf。
• softmax 先减去沿轴的最大值，ReduceMax 不在表里。
• 实验里还有一个前置条件：这几个输入用 TracedTensor::input_concrete_shape(DType::F32, &[2]) 建。用 input_symbolic_shape 建的输入直接调 sigmoid，在建图时就失败：「TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 7」，因为分解要按输入形状造常数 1。

小结

• tenferro-xla 是和 Runtime 平级的执行器，不是 engine：把 CompiledGraph 的语义程序拼成 StableHLO 文本，交给运行期加载的 PJRT 插件；从 0.7.1 到 0.8.0 几乎没改。
• 只收精确的静态形状、F32 / F64 和 23 个核心算子，Sub、Maximum、ReduceMax、Erf 和一切 Bool 都不在内；0.8.0 的组合运算里只有 gelu_tanh 翻得过去。
• 没有插件时三个入口都返回类型化错误；有插件时每次执行都重新建客户端、重新编译，没有可执行文件缓存。