Julia: @inbounds 和 @propagate_inbounds

探究 Julia 的 @inbounds 和 @propagate_inbounds。

背景

近期在开发一个软件包，需要创建一些自定义矩阵，在设计 getindex 方法时，需要加入检查索引边界的代码，并且提供跳过边界检查的方式。

官方文档 详细讨论了这个主题，包括 @inbounds 和 @propagate_inbounds 的使用方式，但是较为简略，且没有提供示例代码，下文将演示这两个宏的使用方式，以及它们的性能差异。

@inbounds

当我们尝试访问不存在的索引时，Julia 会抛出 BoundsError 异常，我们可以通过 @inbounds 来跳过边界检查。

@inline function g(A, i)
    @boundscheck checkbounds(A, i)
    return "accessing ($A)[$i]"
end

function test()
    g(1:2, -1)
end

test() # ERROR: BoundsError: attempt to access 2-element UnitRange{Int64} at index [-1]

只需要在调用 g 时加入 @inbounds 即可跳过边界检查。

function test()
    @inbounds g(1:2, -1)
end

test() # "accessing (1:2)[-1]"

多层检查：@propagate_inbounds

问题

如果函数有多层的边界检查，最外层的 @inbounds 只能跳过最外层的边界检查，内层的边界检查仍然会被执行。

@inline function f(A, i)
    @boundscheck checkbounds(A, i)
    return g(1:2, -1)
end

function test()
    @inbounds f(1:2, -1)
end

test() # ERROR: BoundsError: attempt to access 2-element UnitRange{Int64} at index [-1]

经典解决方案

修改 f 函数，在每一层函数调用，加入 @inbounds，可以跳过边界检查。

@inline function f(A, i)
    @boundscheck checkbounds(A, i)
    return @inbounds g(1:2, -1)
end

test() # "accessing (1:2)[-1]"

在 Julia 官方库 LinearAlgebra.jl 中，广泛使用了 @inbounds，例如 Tridiagonal 矩阵：

@propagate_inbounds

@propagate_inbounds 可以帮我们解决在多层调用时，需要为每一层函数调用加入 @inbounds 的问题。它可以用来传播边界检查的上下文，允许我们在多层函数调用时跳过边界检查。

@inline function g(A, i)
    @boundscheck checkbounds(A, i)
    return "accessing ($A)[$i]"
end

@inline Base.@propagate_inbounds function f(A, i)
    @boundscheck checkbounds(A, i)
    return g(1:2, -1)
end

function test()
    @inbounds f(1:2, -1)
end

test() # "accessing (1:2)[-1]"

这样， f 就可以把边界检查的忽略信息传递给 g，只要调用 f 时加上 @inbounds，就可以跳过内部的边界检查。

需要注意 @propagate_inbounds 并不会无限传递下去，如果 g 函数内部调用了其他有边界检查的函数，需要使用 @propagate_inbounds 或 @inbounds。

性能差异

下面我们来测试 @inbounds 和 @propagate_inbounds 的性能差异，测试代码如下，f1 使用 @inbounds，f2 使用 @propagate_inbounds。

using BenchmarkTools

@inline function f1(A, i)
    @boundscheck checkbounds(A, i)
    return @inbounds g(1:2, -1)
end

@inline Base.@propagate_inbounds function f2(A, i)
    @boundscheck checkbounds(A, i)
    return g(1:2, -1)
end

@inline function g(A, i)
    @boundscheck checkbounds(A, i)
    return "g: accessing ($A)[$i]"
end

function test_f1()
    for i = 1:100000
        @inbounds f1(1:2, -1)
    end
end

function test_f2()
    for i = 1:100000
        @inbounds f2(1:2, -1)
    end
end

function main()
    @inbounds f1(1:2, -1)
    @inbounds f2(1:2, -1)

    t1 = @benchmark test_f1()
    t2 = @benchmark test_f2()
    display(t1)
    display(t2)
end

可以认为两者没有性能差异。

结论

@inbounds 可以跳过单层函数调用的边界检查，在最外调用函数的时候可以使用，有助于提高性能，但也要注意访问的安全性。

@propagate_inbounds 可以跳过多层函数调用的边界检查，适用于多层函数调用的场景，例如开发软件包。这个宏能够提供更好的代码整洁性，避免了在每一层函数调用都加入 @inbounds 的麻烦。

综合来看，两者并没有性能差异，在开发软件包时使用 @propagate_inbounds 是更好的选择。