在Julia NFFT中，用单个数组方法替换向量和矩阵方法

Question

在NFFT包中有用于向量和矩阵的快速专用函数和用于一般数组的慢函数。我想让通用函数更快，而专用函数过时，但我遇到了麻烦。

一组专门化函数本质上是这样的(offset是在实际代码中计算的，但它的值对于这个问题并不重要):

myfunc!(f::Vector, g::Vector)
    offset = 5
    n = length(g)
    N = length(f)
    for l = 1:N
        g[ ((l+offset)%n)+1 ] = f[l]
    end
end

和

myfunc!(f::Matrix, g::Matrix)
    offsetx = 5
    offsety = 5
    n1, n2 = size(g)
    N1, N2 = size(f)
    for ly = 1:N2
        for lx = 1:N1
            g[ ((lx+offsetx)%n1)+1, ((ly+offsety)%n2)+1 ] = f[lx, ly]
        end
    end
end

一般函数可以这样写

myfunc!{D}(f::Array{D}, g::Array{D})
    offset = ntuple(d -> 5, D)
    n = size(g)
    for l in CartesianRange(size(f))
        idx = CartesianIndex{D}( ntuple(d -> ((l[d]+offset[d])%n[d])+1, D) )
        g[ idx ] = f[l]
    end
end

不幸的是，这非常慢。大部分时间花在循环中的ntuple上。

idx的其他可能性包括让idx = Array{Int}(D)并使内环看起来像

for d = 1:D
    idx[d] = ((l[d]+offset[d])%n[d])+1
end
g[idx...] = f[l]

这个也很慢。

我认为，由于维度D是一个类型参数，@generated函数可以用于计算idx，但我不知道如何做到这一点(或者如果有更好的方法)。

我使用的是Julia v0.4.5

Answer 1

如何使用生成的函数做到这一点的答案是使用基地。笛卡尔辅助宏。

using Base.Cartesian 
@generated function myfunc!{T,D}(f::Array{T,D}, g::Array{T,D})
    quote
        @nexprs $D d->offset_d=5 #Declase offset_1=5, offset_2=5 etc
        @nloops $D l f begin
            (@nref $D g d->(l_d+offset_d)%size(g,d)+1) = @nref $D f l 
        end
    end
end

我已经确认这是正确的，至少对于2D。我把它留给读者来分析。