julia - 使文字常量的类型依赖于其他变量

Question

我在 Julia 中有以下代码，其中文字常量 2. 对数组元素进行乘法运算。我现在将文字常量设置为单精度 (2.f0)，但我想让类型取决于其他变量(这些都是 Float64 或都是 Float32)。我该如何优雅地做到这一点？

function diff!(
        at, a,
        visc, dxidxi, dyidyi, dzidzi,
        itot, jtot, ktot)
​
    @tturbo for k in 2:ktot-1
        for j in 2:jtot-1
            for i in 2:itot-1
                at[i, j, k] += visc * (
                    (a[i-1, j  , k  ] - 2.f0 * a[i, j, k] + a[i+1, j  , k  ]) * dxidxi +
                    (a[i  , j-1, k  ] - 2.f0 * a[i, j, k] + a[i  , j+1, k  ]) * dyidyi +
                    (a[i  , j  , k-1] - 2.f0 * a[i, j, k] + a[i  , j  , k+1]) * dzidzi )
            end
        end
    end
end

Answer 1

通常，如果您有一个标量 x 或一个数组 A，您可以使用 T = typeof(x) 获取类型或 T = eltype(A)，然后使用它来将文字转换为等效类型，例如

julia> A = [1.0]
1-element Vector{Float64}:
 1.0

julia> T = eltype(A)
Float64

julia> T(2)
2.0

因此原则上您可以在函数内使用它，如果一切都是类型稳定的，那么这应该实际上是无开销的:

julia> @code_native 2 * 1.0f0
    .section    __TEXT,__text,regular,pure_instructions
; ┌ @ promotion.jl:322 within `*'
; │┌ @ promotion.jl:292 within `promote'
; ││┌ @ promotion.jl:269 within `_promote'
; │││┌ @ number.jl:7 within `convert'
; ││││┌ @ float.jl:94 within `Float32'
    vcvtsi2ss   %rdi, %xmm1, %xmm1
; │└└└└
; │ @ promotion.jl:322 within `*' @ float.jl:331
    vmulss  %xmm0, %xmm1, %xmm0
; │ @ promotion.jl:322 within `*'
    retq
    nopw    (%rax,%rax)
; └

julia> @code_native 2.0f0 * 1.0f0
    .section    __TEXT,__text,regular,pure_instructions
; ┌ @ float.jl:331 within `*'
    vmulss  %xmm1, %xmm0, %xmm0
    retq
    nopw    %cs:(%rax,%rax)
; └

julia> @code_native Float32(2) * 1.0f0
    .section    __TEXT,__text,regular,pure_instructions
; ┌ @ float.jl:331 within `*'
    vmulss  %xmm1, %xmm0, %xmm0
    retq
    nopw    %cs:(%rax,%rax)
; └

然而，碰巧的是，在 Julia 中有一个更优雅的模式来编写一个函数签名，这样它将参数化地专注于您传递给该函数的数组的元素类型，然后您应该能够使用它没有开销来确保您的文字是适当的类型，如下所示:

function diff!(at::AbstractArray{T}, a::AbstractArray{T},
        visc, dxidxi, dyidyi, dzidzi,
        itot, jtot, ktot) where T <: Number

    @tturbo for k in 2:ktot-1
        for j in 2:jtot-1
            for i in 2:itot-1
                at[i, j, k] += visc * (
                    (a[i-1, j  , k  ] - T(2) * a[i, j, k] + a[i+1, j  , k  ]) * dxidxi +
                    (a[i  , j-1, k  ] - T(2) * a[i, j, k] + a[i  , j+1, k  ]) * dyidyi +
                    (a[i  , j  , k-1] - T(2) * a[i, j, k] + a[i  , j  , k+1]) * dzidzi )
            end
        end
    end
end

关于 parametric methods 的文档在某种程度上讨论了这种方法在 Julia