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

In Julia NFFT, replace vector and matrix methods with a single array method

在 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。

关于如何使用生成的函数执行此操作的答案是 使用 Base.Cartesian 辅助宏。

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。 我把它留给 reader 来分析它。 它或多或少不会分配。