在索引处增加StaticVector

Question

我想A在location iby 处增加一个静态数组x.如果它是一个可变数组,我会这样做A[i] += x.但由于它是StaticArray,我需要创建一个新的.但是,如果我新的大小,A那么我会做类似的事情

A = A + @SVector [0,0,x]

并为每个人分支i.但在这种情况下,这SVector是一个用户输入,所以我只提前知道使用类型信息.我宁愿不让我的核心逻辑都只是一个生成函数来处理这个,所以我希望有一个简单的解决方案,或者这可能需要一个@generated辅助函数.

请注意,此问题等同于创建在位置@SVector具有值但在其他x位置i为零的问题.如果有一个简单的方法,那么我的问题也解决了.

Answer 1

使用数组理解的天真方法将是

julia> k = 4 4

julia> @SVector [i == k? 1.0 : 0 for i in 1:5] 5-element SVector{10,Float64}: 0.0 0.0 0.0 1.0 0.0

这是阅读README 的StaticArrays.jl "快速入门"部分时可以采取的良好的第一步.

但是,我们在Julia中关注类型稳定性和通用代码,因为:

1. 函数中的类型稳定性让编译器优化 - >速度
2. 利用多个调度,可以以强大的方式重用和扩展代码.

所以更多的朱利安方式将是使用

julia> function increment_value(A::SVector{L,T},x,k) where {L,T} _A = [i == k ? x : zero(x) for i in 1:L] A+_A end

julia> A = @SVector [0, 0, 0, 0, 10] 5-element SVector{5,Int64}: 0 0 0 0 10

julia> increment_value(A,5,2) 5-element SVector{5,Int64}: 0 5 0 0 10

但是,我们的最终答案应该包括一种避免额外变量分配的方法,并利用一些方便的函数利用一些编译器流水线ifelse:

`julia>使用StaticArrays,BenchmarkTools

julia> function increment_value(A :: SVector{L,T}, x,k) where {T,L} SVector(ntuple(i->ifelse(i == k, A[i]+x, A[i]), Val{L})) end increment_value (generic function with 1 method)

julia> a = @SVector [ 1, 2, 3, 4, 5] 5-element SVector{5,Int64}: 1 2 3 4 5

julia> @benchmark increment_value($a,$3,$5) BenchmarkTools.Trial: memory estimate: 0 bytes allocs estimate: 0

minimum time: 3.178 ns (0.00% GC) median time: 3.285 ns (0.00% GC) mean time: 3.293 ns (0.00% GC) maximum time: 13.620 ns (0.00% GC) samples: 10000 evals/sample: 1000

Answer 2

在编译时获取此问题的所有必要值有点棘手。我现在得到的是：

@generated updateindex(s::SVector{L,T},j::Type{Val{I}},v) where {L,T,I} = 
  Expr(:call, :(SVector{L,T}), (ifelse(i==I, :(s[$i]+v), :(s[$i])) for i=1:L)...)

或者只是设置一个坐标：

@generated setindex(s::SVector{L,T},j::Type{Val{I}},v) where {L,T,I} = 
  Expr(:call, :(SVector{L,T}), (ifelse(i==I, :v, :(s[$i])) for i=1:L)...)

这可以用作：

julia> Z = @SVector [1,1,1,1,1];

julia> updateindex(Z,Val{3},4)
5-element SVector{5,Int64}:
 1
 1
 5
 1
 1

并进行基准测试：

julia> using BenchmarkTools

julia> @btime updateindex($Z,Val{3},4);
  2.032 ns (0 allocations: 0 bytes)

代码是最少的：

julia> @code_native updateindex(Z,Val{3},4)
    .text
Filename: REPL[13]
    pushq   %rbp
    movq    %rsp, %rbp
Source line: 1
    vmovups (%rsi), %xmm0
    addq    16(%rsi), %rcx
    movq    24(%rsi), %rax
    movq    32(%rsi), %rdx
    vmovups %xmm0, (%rdi)
    movq    %rcx, 16(%rdi)
    movq    %rax, 24(%rdi)
    movq    %rdx, 32(%rdi)
    movq    %rdi, %rax
    popq    %rbp
    retq
    nopl    (%rax)

这能解决难题吗？

顺便说一句，如果有办法将其重写为更易读的形式，我将很高兴在评论中看到（并将相应地更新答案）。

更新

j克里斯的评论正确地指出可以制作具有非值类型的版本：

@generated setindex(s::SVector{L,T},j,v) where {L,T} =
  Expr(:call, :(SVector{L,T}), (:(ifelse($i==j, v, s[$i])) for i=1:L)...)

Demo和底层代码（可见不知道更新哪个索引是要付出一点性能代价的）：

julia> setindex(Z,4,3)
5-element SVector{5,Int64}:
 1
 1
 1
 3
 1

julia> @code_native setindex(Z,4,3)
    .text
Filename: REPL[15]
    pushq   %rbp
    movq    %rsp, %rbp
Source line: 1
    cmpq    $1, %rdx
    movq    (%rsi), %r8
    cmoveq  %rcx, %r8
    cmpq    $2, %rdx
    movq    8(%rsi), %r9
    cmoveq  %rcx, %r9
    cmpq    $3, %rdx
    movq    16(%rsi), %r10
    cmoveq  %rcx, %r10
    cmpq    $4, %rdx
    movq    24(%rsi), %rax
    cmoveq  %rcx, %rax
    cmpq    $5, %rdx
    cmovneq 32(%rsi), %rcx
    movq    %r8, (%rdi)
    movq    %r9, 8(%rdi)
    movq    %r10, 16(%rdi)
    movq    %rax, 24(%rdi)
    movq    %rcx, 32(%rdi)
    movq    %rdi, %rax
    popq    %rbp
    retq
    nopw    %cs:(%rax,%rax)