如何告诉 LLVM 它可以优化离店？

Posted 2023-02-21

【中文标题】如何告诉 LLVM 它可以优化离店？

【英文标题】：How to tell LLVM that it can optimize away stores?

【发布时间】：2020-03-09 11:53:54

【问题描述】：

背景（可能有更好的方法来做到这一点）： 我正在开发一个手动管理内存的 Julia 库；我mmap 一个大块，然后主要把它当作一个堆栈：函数接收指针作为参数，如果它们分配一个对象，它们将返回一个递增的指针给被调用者。 该被调用者本身可能不会增加指针，并且只返回它收到的原始指针，如果它返回指针的话。

每当一个函数返回时，就我的库而言，指针当前位置之外的任何东西都是垃圾。我希望 LLVM 能够意识到这一点，以便它可以优化掉任何不必要的存储。

这是一个演示问题的测试用例：取两个长度为 16 的向量的点积。首先，一些初步加载（这些是我的库，在 GitHub 上：SIMDPirates、PaddedMatrices）：

using SIMDPirates, PaddedMatrices
using SIMDPirates: lifetime_start, lifetime_end
b = @Mutable rand(16);
c = @Mutable rand(16);
a = FixedSizeVector16,Float64(undef);
b' * c # dot product
# 3.9704768664758925

当然，如果我们手工编写点积，我们永远不会包含商店，但是当您尝试为任意模型生成代码时，这要困难得多。所以我们会写一个坏点积存储到一个指针中：

@inline function storedot!(ptr, b, c)
    ptrb = pointer(b)
    ptrc = pointer(c)
    ptra = ptr
    for _ ∈ 1:4
        vb = vload(Vec4,Float64, ptrb)
        vc = vload(Vec4,Float64, ptrc)
        vstore!(ptra, vmul(vb, vc))
        ptra += 32
        ptrb += 32
        ptrc += 32
    end
    ptra = ptr
    out = vload(Vec4,Float64, ptra)
    for _ ∈ 1:3
        ptra += 32
        out = vadd(out, vload(Vec4,Float64, ptra))
    end
    vsum(out)
end

我们不是循环一次并用fma指令累加点积，而是循环两次，首先计算和存储产品，然后求和。 我想要的是让编译器找出正确的东西。

下面是两个版本。第一个使用 llvm lifetime 内部函数尝试将指针内容声明为垃圾：

function test_lifetime!(a, b, c)
    ptra = pointer(a)
    lifetime_start(Val(128), ptra)
    d = storedot!(ptra, b, c)
    lifetime_end(Val(128), ptra)
    d
end

第二个不是使用预分配的指针，而是使用alloca创建一个指针

function test_alloca(b, c)
    ptra = SIMDPirates.alloca(Val(16), Float64)
    storedot!(ptra, b, c)
end

当然都得到正确答案

test_lifetime!(a, b, c)
# 3.9704768664758925
test_alloca(b, c)
# 3.9704768664758925

但只有 alloca 版本被正确优化。 alloca 的程序集（AT&T 语法）：

# julia> @code_native debuginfo=:none test_alloca(b, c)
        .text
        vmovupd (%rsi), %ymm0
        vmovupd 32(%rsi), %ymm1
        vmovupd 64(%rsi), %ymm2
        vmovupd 96(%rsi), %ymm3
        vmulpd  (%rdi), %ymm0, %ymm0
        vfmadd231pd     32(%rdi), %ymm1, %ymm0 # ymm0 = (ymm1 * mem) + ymm0
        vfmadd231pd     64(%rdi), %ymm2, %ymm0 # ymm0 = (ymm2 * mem) + ymm0
        vfmadd231pd     96(%rdi), %ymm3, %ymm0 # ymm0 = (ymm3 * mem) + ymm0
        vextractf128    $1, %ymm0, %xmm1
        vaddpd  %xmm1, %xmm0, %xmm0
        vpermilpd       $1, %xmm0, %xmm1 # xmm1 = xmm0[1,0]
        vaddsd  %xmm1, %xmm0, %xmm0
        vzeroupper
        retq
        nopw    %cs:(%rax,%rax)
        nopl    (%rax,%rax)

如您所见，内存中没有移动，我们有一个 vmul 和三个 vfmadds 来计算点积（在进行向量缩减之前）。

不幸的是，这不是我们从尝试使用生命周期的版本中得到的：

 # julia> @code_native debuginfo=:none test_lifetime!(a, b, c)
        .text
        vmovupd (%rdx), %ymm0
        vmulpd  (%rsi), %ymm0, %ymm0
        vmovupd %ymm0, (%rdi)
        vmovupd 32(%rdx), %ymm1
        vmulpd  32(%rsi), %ymm1, %ymm1
        vmovupd %ymm1, 32(%rdi)
        vmovupd 64(%rdx), %ymm2
        vmulpd  64(%rsi), %ymm2, %ymm2
        vmovupd %ymm2, 64(%rdi)
        vmovupd 96(%rdx), %ymm3
        vaddpd  %ymm0, %ymm1, %ymm0
        vaddpd  %ymm0, %ymm2, %ymm0
        vfmadd231pd     96(%rsi), %ymm3, %ymm0 # ymm0 = (ymm3 * mem) + ymm0
        vextractf128    $1, %ymm0, %xmm1
        vaddpd  %xmm1, %xmm0, %xmm0
        vpermilpd       $1, %xmm0, %xmm1 # xmm1 = xmm0[1,0]
        vaddsd  %xmm1, %xmm0, %xmm0
        vzeroupper
        retq
        nopw    %cs:(%rax,%rax)
        nop

在这里，我们只得到写好的循环：vmul，存储到内存中，然后是vadd。然而，这 4 个中的一个已被替换为 fmadd。

此外，它不会从任何商店读取，所以我认为死店消除通行证应该没有问题。

相关的llvm：

;; julia> @code_llvm debuginfo=:none test_alloca(b, c)

define double @julia_test_alloca_17840(%jl_value_t addrspace(10)* nonnull align 8 dereferenceable(128), %jl_value_t addrspace(10)* nonnull align 8 dereferenceable(128)) 
top:
  %2 = addrspacecast %jl_value_t addrspace(10)* %0 to %jl_value_t addrspace(11)*
  %3 = addrspacecast %jl_value_t addrspace(11)* %2 to %jl_value_t*
  %4 = addrspacecast %jl_value_t addrspace(10)* %1 to %jl_value_t addrspace(11)*
  %5 = addrspacecast %jl_value_t addrspace(11)* %4 to %jl_value_t*
  %ptr.i20 = bitcast %jl_value_t* %3 to <4 x double>*
  %res.i21 = load <4 x double>, <4 x double>* %ptr.i20, align 8
  %ptr.i18 = bitcast %jl_value_t* %5 to <4 x double>*
  %res.i19 = load <4 x double>, <4 x double>* %ptr.i18, align 8
  %res.i17 = fmul fast <4 x double> %res.i19, %res.i21
  %6 = bitcast %jl_value_t* %3 to i8*
  %7 = getelementptr i8, i8* %6, i64 32
  %8 = bitcast %jl_value_t* %5 to i8*
  %9 = getelementptr i8, i8* %8, i64 32
  %ptr.i20.1 = bitcast i8* %7 to <4 x double>*
  %res.i21.1 = load <4 x double>, <4 x double>* %ptr.i20.1, align 8
  %ptr.i18.1 = bitcast i8* %9 to <4 x double>*
  %res.i19.1 = load <4 x double>, <4 x double>* %ptr.i18.1, align 8
  %res.i17.1 = fmul fast <4 x double> %res.i19.1, %res.i21.1
  %10 = getelementptr i8, i8* %6, i64 64
  %11 = getelementptr i8, i8* %8, i64 64
  %ptr.i20.2 = bitcast i8* %10 to <4 x double>*
  %res.i21.2 = load <4 x double>, <4 x double>* %ptr.i20.2, align 8
  %ptr.i18.2 = bitcast i8* %11 to <4 x double>*
  %res.i19.2 = load <4 x double>, <4 x double>* %ptr.i18.2, align 8
  %res.i17.2 = fmul fast <4 x double> %res.i19.2, %res.i21.2
  %12 = getelementptr i8, i8* %6, i64 96
  %13 = getelementptr i8, i8* %8, i64 96
  %ptr.i20.3 = bitcast i8* %12 to <4 x double>*
  %res.i21.3 = load <4 x double>, <4 x double>* %ptr.i20.3, align 8
  %ptr.i18.3 = bitcast i8* %13 to <4 x double>*
  %res.i19.3 = load <4 x double>, <4 x double>* %ptr.i18.3, align 8
  %res.i17.3 = fmul fast <4 x double> %res.i19.3, %res.i21.3
  %res.i12 = fadd fast <4 x double> %res.i17.1, %res.i17
  %res.i12.1 = fadd fast <4 x double> %res.i17.2, %res.i12
  %res.i12.2 = fadd fast <4 x double> %res.i17.3, %res.i12.1
  %vec_2_1.i = shufflevector <4 x double> %res.i12.2, <4 x double> undef, <2 x i32> <i32 0, i32 1>
  %vec_2_2.i = shufflevector <4 x double> %res.i12.2, <4 x double> undef, <2 x i32> <i32 2, i32 3>
  %vec_2.i = fadd <2 x double> %vec_2_1.i, %vec_2_2.i
  %vec_1_1.i = shufflevector <2 x double> %vec_2.i, <2 x double> undef, <1 x i32> zeroinitializer
  %vec_1_2.i = shufflevector <2 x double> %vec_2.i, <2 x double> undef, <1 x i32> <i32 1>
  %vec_1.i = fadd <1 x double> %vec_1_1.i, %vec_1_2.i
  %res.i = extractelement <1 x double> %vec_1.i, i32 0
  ret double %res.i

它省略了alloca 和stores。 但是，尝试使用生命周期：

;; julia> @code_llvm debuginfo=:none test_lifetime!(a, b, c)

define double @"julia_test_lifetime!_17839"(%jl_value_t addrspace(10)* nonnull align 8 dereferenceable(128), %jl_value_t addrspace(10)* nonnull align 8 dereferenceable(128), %jl_value_t addrspace(10)* nonnull align 8 dereferenceable(128)) 
  980 top:
  %3 = addrspacecast %jl_value_t addrspace(10)* %0 to %jl_value_t addrspace(11)*
  %4 = addrspacecast %jl_value_t addrspace(11)* %3 to %jl_value_t*
  %.ptr = bitcast %jl_value_t* %4 to i8*
  call void @llvm.lifetime.start.p0i8(i64 256, i8* %.ptr)
  %5 = addrspacecast %jl_value_t addrspace(10)* %1 to %jl_value_t addrspace(11)*
  %6 = addrspacecast %jl_value_t addrspace(11)* %5 to %jl_value_t*
  %7 = addrspacecast %jl_value_t addrspace(10)* %2 to %jl_value_t addrspace(11)*
  %8 = addrspacecast %jl_value_t addrspace(11)* %7 to %jl_value_t*
  %ptr.i22 = bitcast %jl_value_t* %6 to <4 x double>*
  %res.i23 = load <4 x double>, <4 x double>* %ptr.i22, align 8
  %ptr.i20 = bitcast %jl_value_t* %8 to <4 x double>*
  %res.i21 = load <4 x double>, <4 x double>* %ptr.i20, align 8
  %res.i19 = fmul fast <4 x double> %res.i21, %res.i23
  %ptr.i18 = bitcast %jl_value_t* %4 to <4 x double>*
  store <4 x double> %res.i19, <4 x double>* %ptr.i18, align 8
  %9 = getelementptr i8, i8* %.ptr, i64 32
  %10 = bitcast %jl_value_t* %6 to i8*
  %11 = getelementptr i8, i8* %10, i64 32
  %12 = bitcast %jl_value_t* %8 to i8*
  %13 = getelementptr i8, i8* %12, i64 32
  %ptr.i22.1 = bitcast i8* %11 to <4 x double>*
  %res.i23.1 = load <4 x double>, <4 x double>* %ptr.i22.1, align 8
  %ptr.i20.1 = bitcast i8* %13 to <4 x double>*
  %res.i21.1 = load <4 x double>, <4 x double>* %ptr.i20.1, align 8
  %res.i19.1 = fmul fast <4 x double> %res.i21.1, %res.i23.1
  %ptr.i18.1 = bitcast i8* %9 to <4 x double>*
  store <4 x double> %res.i19.1, <4 x double>* %ptr.i18.1, align 8
  %14 = getelementptr i8, i8* %.ptr, i64 64
  %15 = getelementptr i8, i8* %10, i64 64
  %16 = getelementptr i8, i8* %12, i64 64
  %ptr.i22.2 = bitcast i8* %15 to <4 x double>*
  %res.i23.2 = load <4 x double>, <4 x double>* %ptr.i22.2, align 8
  %ptr.i20.2 = bitcast i8* %16 to <4 x double>*
  %res.i21.2 = load <4 x double>, <4 x double>* %ptr.i20.2, align 8
  %res.i19.2 = fmul fast <4 x double> %res.i21.2, %res.i23.2
  %ptr.i18.2 = bitcast i8* %14 to <4 x double>*
  store <4 x double> %res.i19.2, <4 x double>* %ptr.i18.2, align 8
  %17 = getelementptr i8, i8* %10, i64 96
  %18 = getelementptr i8, i8* %12, i64 96
  %ptr.i22.3 = bitcast i8* %17 to <4 x double>*
  %res.i23.3 = load <4 x double>, <4 x double>* %ptr.i22.3, align 8
  %ptr.i20.3 = bitcast i8* %18 to <4 x double>*
  %res.i21.3 = load <4 x double>, <4 x double>* %ptr.i20.3, align 8
  %res.i19.3 = fmul fast <4 x double> %res.i21.3, %res.i23.3
  %res.i13 = fadd fast <4 x double> %res.i19.1, %res.i19
  %res.i13.1 = fadd fast <4 x double> %res.i19.2, %res.i13
  %res.i13.2 = fadd fast <4 x double> %res.i19.3, %res.i13.1
  %vec_2_1.i = shufflevector <4 x double> %res.i13.2, <4 x double> undef, <2 x i32> <i32 0, i32 1>
  %vec_2_2.i = shufflevector <4 x double> %res.i13.2, <4 x double> undef, <2 x i32> <i32 2, i32 3>
  %vec_2.i = fadd <2 x double> %vec_2_1.i, %vec_2_2.i
  %vec_1_1.i = shufflevector <2 x double> %vec_2.i, <2 x double> undef, <1 x i32> zeroinitializer
  %vec_1_2.i = shufflevector <2 x double> %vec_2.i, <2 x double> undef, <1 x i32> <i32 1>
  %vec_1.i = fadd <1 x double> %vec_1_1.i, %vec_1_2.i
  %res.i = extractelement <1 x double> %vec_1.i, i32 0
  call void @llvm.lifetime.end.p0i8(i64 256, i8* %.ptr)
  ret double %res.i

生命周期开始和生命周期结束都在那里，但四家商店中的三家也是如此。 我可以确认第四家店已经不见​​了：

julia> fill!(a, 0.0)'
1×16 LinearAlgebra.AdjointFloat64,FixedSizeArrayTuple16,Float64,1,Tuple1,16:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

julia> test_lifetime!(a, b, c)
3.9704768664758925

julia> a'
1×16 LinearAlgebra.AdjointFloat64,FixedSizeArrayTuple16,Float64,1,Tuple1,16:
 0.157677  0.152386  0.507693  0.00696963  0.0651712  0.241523  0.129705  0.175321  0.236032  0.0314141  0.199595  0.404153  0.0  0.0  0.0  0.0

虽然没有指定生命周期，但所有四个当然都必须发生：

julia> function teststore!(a, b, c)
       storedot!(pointer(a), b, c)
       end
test_store! (generic function with 1 method)

julia> fill!(a, 0.0); test_store!(a, b, c)
3.9704768664758925

julia> a'
1×16 LinearAlgebra.AdjointFloat64,FixedSizeArrayTuple16,Float64,1,Tuple1,16:
 0.157677  0.152386  0.507693  0.00696963  0.0651712  0.241523  0.129705  0.175321  0.236032  0.0314141  0.199595  0.404153  0.256597  0.0376403  0.889331  0.479269

然而，与 alloca 不同的是，它无法忽略所有 4 家商店。

作为参考，我使用 LLVM 8.0.1 构建了 Julia。

我没有使用alloca 代替我的堆栈指针，原因有两个： a) 使用alloca-created 指针调用非内联函数时出现错误。用其他指针替换这些指针会使错误消失，内联函数也是如此。如果有办法解决这个问题，我至少可以在更多地方使用alloca。 b) 我不知道如何让 Julia 每个线程有超过 4MB 的堆栈可供 alloca 使用。我认为 4MB 对于我的许多用例来说已经足够了，但不是全部。如果我的目标是编写相当通用的软件，这样的限制并不是很好。

我的问题：

有什么方法可以让 LLVM 复制它在 alloca 中显示的行为？ 我做的事情是否正确，并允许 LLVM 显示所需的行为，但优化器由于某种原因比 alloca 更受限制？ 因此有望在未来的版本中得到改进。 有关如何处理此问题、更好地启用优化器或我一般缺少的东西的任何建议？ 鉴于只有最后一个被省略，问题是假设它们可能有别名吗？

【参考方案1】：

在最初发布问题后，我在以下项目符号中进行了编辑：

鉴于只有最后一个被省略，问题是假设它们可能有别名吗？

原来正是这个问题。如果ptra 做了别名b 或c，则省略商店将无效。

改写：

a = @Mutable rand(48);
a[Static(1:16)]' * a[Static(17:32)]
# 2.5295415040590425

function test_lifetime!(a)
    ptra = pointer(a)
    b = PtrVector16,Float64,16(ptra)
    c = PtrVector16,Float64,16(ptra + 128)
    ptra += 256
    lifetime_start(Val(128), ptra)
    d = storedot!(ptra, b, c)
    lifetime_end(Val(128), ptra)
    d
end

test_lifetime!(a)
# 2.5295415040590425

实际上忽略了所有商店：

# julia> @code_native debuginfo=:none test_lifetime!(a)
        .text
        vmovupd 128(%rdi), %ymm0
        vmovupd 160(%rdi), %ymm1
        vmovupd 192(%rdi), %ymm2
        vmovupd 224(%rdi), %ymm3
        vmulpd  (%rdi), %ymm0, %ymm0
        vfmadd231pd     32(%rdi), %ymm1, %ymm0 # ymm0 = (ymm1 * mem) + ymm0
        vfmadd231pd     64(%rdi), %ymm2, %ymm0 # ymm0 = (ymm2 * mem) + ymm0
        vfmadd231pd     96(%rdi), %ymm3, %ymm0 # ymm0 = (ymm3 * mem) + ymm0
        vextractf128    $1, %ymm0, %xmm1
        vaddpd  %xmm1, %xmm0, %xmm0
        vpermilpd       $1, %xmm0, %xmm1 # xmm1 = xmm0[1,0]
        vaddsd  %xmm1, %xmm0, %xmm0
        vzeroupper
        retq
        nop

所以答案是：LLVM 知道 alloca 指针不能为输入之一设置别名，因此不存储是安全的。 我在我的问题中想要的行为（没有别名检查）将是不安全的/可能会得到不正确的结果：ptra 中的一个商店可能会更改 b 或 c 的内容。因此，实际上必须执行除最后一个存储之外的所有存储。

在最后的测试中，我在同一个指针的不同偏移量处定义了 a、b 和 c 中的每一个，因此可以保证存储到 a 中的内容不会更改 b 或 @987654334 @，让 LLVM 实际上消除了商店。完美！