Julia使用生成器表达式引发函数类型不稳定问题咨询

在尝试借助生成器（Generator）表达式优化似然计算函数、降低内存分配量时，会遇到两个功能完全一致的函数存在明显性能差异的现象，两个函数实现如下：

function testMax!(x,X,β)
    xmax = 0.0
    @inbounds for i ∈ eachindex(x)
        x[i] = X[i,2] * β[1] + X[i,2] * β[2]
        if x[i] > xmax 
            xmax = x[i]
        end 
    end 
    y = 0.0 
    for i ∈ eachindex(x)
        y += exp(x[i]-xmax)
    end
    return xmax, y
end

function testMaxWeird!(x,X,β)
    xmax = 0.0
    @inbounds for i ∈ eachindex(x)
        x[i] = X[i,2] * β[1] + X[i,2] * β[2]
        if x[i] > xmax 
            xmax = x[i]
        end 
    end 
    y = sum(exp(x[j]-xmax) for j ∈ eachindex(x))
    return xmax, y
end

两个函数的输出结果完全一致，测试代码如下：

using Random
Random.seed!(1234)
H = 10000;
X = rand(H,2);
β = rand(2);

x = zeros(H);
testMax!(x,X,β)
x = zeros(H);
testMaxWeird!(x,X,β)

运行后两个函数均返回(1.0772897308017204, 6101.682959406999)，但testMax!是类型稳定（type stable）的，testMaxWeird!存在类型不稳定问题，运行速度慢得多。
通过@code_warntype查看类型推断结果可见核心差异：

• testMax!中变量类型完全确定：

y::Float64
xmax::Float64

• testMaxWeird!中存在类型不确定的变量：

y::Any
xmax@_9::Core.Box

核心疑问为：该问题的具体成因是循环中多次对xmax赋值的写法不规范，还是生成器表达式的使用方式存在问题？

补充追问 #

现有参考资料和解决方案可以解决表层问题，但对问题触发条件仍存疑：该问题是否源于xmax在for循环中被更新的写法？它的作用域是否和函数内定义的其他局部变量存在差异？为什么如下这种效率更低的最大值计算写法，不会触发同类闭包问题？

function testMax2!(x,X,β)
    @inbounds for i ∈ eachindex(x)
        x[i] = X[i,2] * β[1] + X[i,2] * β[2]
    end 
    xmax = maximum(x)
    y = sum(exp(x[j]-xmax) for j ∈ eachindex(x))            
    return xmax, y
end

二次补充：在Julia官方性能建议文档中可找到相关说明："The parser, when translating it into lower-level instructions, substantially reorganizes the above code by extracting the inner function to a separate code block." 可推测问题成因是变量被多次赋值后，代码重排过程中出现了类型推断障碍。

问题根因 #

这是Julia闭包捕获变量的典型行为，和循环写法本身是否规范无关：

• 生成器表达式本质是一个迭代器对象，编译器在底层翻译阶段会把它提取为独立的内部函数，这个内部函数会捕获外层作用域的xmax变量，形成闭包。
• 在testMaxWeird!中，xmax在前置for循环中被多次修改赋值，编译器在提取闭包时无法静态确定该变量的最终类型，只能将其包装为Core.Box做动态类型检查，最终导致求和结果y的类型无法被推断为Float64，退化为Any，产生大量额外运行开销。
• 在testMax2!中，xmax在for循环执行完毕后才被赋值，整个变量生命周期内仅赋值一次，编译器可以明确推断出它的类型为Float64，闭包捕获时不需要做装箱处理，自然不会出现类型不稳定。

修复方式非常简单，只需要在生成器中给捕获的xmax加显式类型标注，比如把求和行写成y = sum(exp(x[j] - xmax::Float64) for j ∈ eachindex(x))，就能让编译器明确变量类型，完全消除装箱带来的性能损失。

内容的提问来源于stack exchange，提问作者ekboehm