下面给出我的问题的最小可运行示例（MWE）。我特别想知道jit编译器在这个流程中会怎么处理：

def blackbox_function(x, y): # Blackbox function I'm given. Can't break it up.
    x = x + 1 # *Guaranteed* that x is updated in the same way, independent of y
    y = x * y
    return x, y

@jax.jit
def my_mapped_function(x, y):
    x, y = jax.vmap(blackbox_function, in_axes=(None,0))(x, y) # x fixed during parallelization and output will be the same for all elements of y
    return x[0], y # Return only the first element of x

x = 1
y = jnp.arange(10)
print(my_mapped_function.lower(x, y).compile().as_text()) # Inspect what the compiler is doing

这个流程有几个特别的性质：

• 输入 x 总是以相同的方式更新，与 y 无关。
• 我只输出 x 的第一个映射值（与其他值相同）。

因此，最优的计算应该只计算 x = x+1 一次，并把它输出到第一个输出参数中。为了验证这一点，我进行了如下操作：

print(my_mapped_function.lower(x, y).compile().as_text()) # Inspect what the compiler is doing

输出结果是：

HloModule jit_my_mapped_function, is_scheduled=true, entry_computation_layout={(s64[], s64[10]{0})->(s64[], s64[10]{0})}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true,true}

%fused_computation (param_0.1: s64[10], param_1.2: s64[]) -> s64[10] {
  %param_1.2 = s64[] parameter(1)
  %constant.0 = s64[] constant(1)
  %add.0 = s64[] add(%param_1.2, %constant.0), metadata={op_name="jit(my_mapped_function)/jit(main)/add" source_file="/var/folders/x0/28x522xx1vb2xl75tn781lqr0000gn/T/ipykernel_33608/2232785945.py" source_line=2}
  %broadcast.3 = s64[10]{0} broadcast(%add.0), dimensions={}, metadata={op_name="jit(my_mapped_function)/jit(main)/mul" source_file="/var/folders/x0/28x522xx1vb2xl75tn781lqr0000gn/T/ipykernel_33608/2232785945.py" source_line=3}
  %param_0.1 = s64[10]{0} parameter(0)
  ROOT %multiply.0 = s64[10]{0} multiply(%broadcast.3, %param_0.1), metadata={op_name="jit(my_mapped_function)/jit(main)/mul" source_file="/var/folders/x0/28x522xx1vb2xl75tn781lqr0000gn/T/ipykernel_33608/2232785945.py" source_line=3}
}

ENTRY %main.11 (Arg_0.1: s64[], Arg_1.2: s64[10]) -> (s64[], s64[10]) {
  %Arg_0.1 = s64[] parameter(0), metadata={op_name="x"}
  %Arg_1.2 = s64[10]{0} parameter(1), metadata={op_name="y"}
  %constant.3 = s64[] constant(1)
  %broadcast_multiply_fusion = s64[10]{0} fusion(%Arg_1.2, %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(my_mapped_function)/jit(main)/mul" source_file="/var/folders/x0/28x522xx1vb2xl75tn781lqr0000gn/T/ipykernel_33608/2232785945.py" source_line=3}
  %add.4 = s64[] add(%Arg_0.1, %constant.3), metadata={op_name="jit(my_mapped_function)/jit(main)/add" source_file="/var/folders/x0/28x522xx1vb2xl75tn781lqr0000gn/T/ipykernel_33608/2232785945.py" source_line=2}
  ROOT %tuple.10 = (s64[], s64[10]{0}) tuple(%add.4, %broadcast_multiply_fusion)
}

我在逐行理解这个过程。对于融合计算，我大致得到如下结论：

%fused_computation (param_0.1: s64[10], param_1.2: s64[]) -> s64[10] {
  %param_1.2 = s64[] parameter(1) # Getting the initial `x`
  %constant.0 = s64[] constant(1) # Declaring the constant 1 in `x=x+1`
  %add.0 = s64[] add(%param_1.2, %constant.0) ... # Updating `x=x+1`
  %broadcast.3 = s64[10]{0} broadcast(%add.0) ... # Broadcasting the new `x` to the shape of `y`
  %param_0.1 = s64[10]{0} parameter(0) # Getting the initial `y`
  %ROOT %multiply.0 = s64[10]{0} multiply(%broadcast.3, %param_0.1) ... # Multiply the new `x` with all the `y`, in parallel
}

然后是在主程序中：

ENTRY %main.11 (Arg_0.1: s64[], Arg_1.2: s64[10]) -> (s64[], s64[10]) {
  %Arg_0.1 = s64[] parameter(0), metadata={op_name="x"} # Get the initial `x`
  %Arg_1.2 = s64[10]{0} parameter(1), metadata={op_name="y"} # Get the initial `y`
  %constant.3 = s64[] constant(1) # Declare the constant 1 to be added
  %broadcast_multiply_fusion = s64[10]{0} fusion(%Arg_1.2, %Arg_0.1), kind=kLoop, calls=%fused_computation ... # Call the fused computation above
  %add.4 = s64[] add(%Arg_0.1, %constant.3) ... # Calculate `x=x+1`
  ROOT %tuple.10 = (s64[], s64[10]{0}) tuple(%add.4, %broadcast_multiply_fusion) # Output result
}

让我有点困惑的两行是以下两行：

 %add.0 = s64[] add(%param_1.2, %constant.0) ... # Updating `x=x+1` in fused computation
...
 %add.4 = s64[] add(%Arg_0.1, %constant.3) ... # Calculate `x=x+1`. Again?

是在 %fused_computation 里面执行了一次加法，然后在 %main 里面再执行一次吗？如果是，为什么融合计算不直接返回结果呢，因为它已经知道要只计算一次？

还是说我理解错了？

解决方案

从HLO来看，编译器似乎选择将 (x + 1) * y 融合成一个单一的融合操作。大致来说，可以把融合操作理解为可以在一次遍历中直接执行，而不需要对中间数组进行物化。

你的函数也返回 x + 1，必须单独计算，因为在用于计算 y 的融合操作中不会对该中间值进行物化。

编译器在这里会不会错过一个优化？大概不会：把 1 加到一个标量上本来就是一项非常便宜的运算，等该运算完成再分派广播乘法的执行，可能不如把这个廉价的运算执行两次以避免那种运行时依赖来得高效。