向量化函数调用组

时间:2018-11-14 19:31:03

标签: function macros julia vectorization metaprogramming

我正在尝试编写一个功能(使用宏,生成的函数等),以有效地将Julia函数调用向量化为我编写的函数。基本上,我正在尝试编写自己的@版本。宏,但我希望它接受函数而不是for循环-如果我正确理解这一点。这是我已阅读的一些文档:

https://docs.julialang.org/en/v1/manual/functions/#man-vectorized-1

https://github.com/JuliaLang/julia/blob/master/base/broadcast.jl

https://julialang.org/blog/2017/01/moredots

https://docs.julialang.org/en/v1/manual/metaprogramming/index.html#Code-Generation-1

这是我正在使用的初步玩具示例,用于实现这种功能:

function add!(v_add::Vector{Float64}, a_add::Float64, j::Int64)
  v_add[j] = v_add[j]+a_add
end

function add!(v_add::Vector{Float64}, a_add::Float64)
  for j in 1:length(v_add)
    v_add[j] = v_add[j]+a_add
  end
end

macro vectorize(args)
  print("\n****************** args\n")
  print(args)
  print("\n******************\n")
  e = :(:call,
    $args[1],
    $args[2],
    $args[3])
  print("\n****************** expression\n")
  show(e)
  print(e)
  print("\n******************\n")
  return e
end

function test!(v_test, a_test)
  # # Traverse vector twice
  # add!(v_test, a_test)
  # add!(v_test, a_test)
  # Traverse vector once
  args = [
  add!, v_test, a_test,
  add!, v_test, a_test
  ]
  e = @vectorize(args)
  # eval(e) # Next step
end

v_main = Vector([Float64(i) for i in 1:3])
a_main = Float64(2.0)
print("\n",v_main, "\n")
Main.test!(v_main, a_main)
print("\n",v_main, "\n")

到目前为止,我遇到的问题是我什至无法使用宏运行去矢量化的版本。此示例导致LoadError:UndefVarError:args未定义。对于使此脚本按预期运行(如果输入为[1、2、3],输出应该为[5、6、7]),我将不胜感激。

任何帮助/建议都将不胜感激。

更新

更具体地讲,给出以下定义的功能:

function add!(v::Vector{Float64}, a::Float64)
  for j in 1:length(v)
    v[j]+= a
  end
end
function add!(v::Vector{Float64}, a::Float64, j::Int64)
  v[j]+= a
end

我希望能够使用宏来转换以下代码行:

v = [Float64(j) for j in 1:10]
a = 1
b = 2
@vectorize_I_would_like_to_define(
# I don't know the exact form that the args to this macro should take.
add!(v, a),
add!(v, b)
)

要生成像这样编译的代码:

v = [Float64(j) for j in 1:10]
a = 1
b = 2
for j in 1:length(v)
  add!(v, a, j)
  add!(v, b, j)
end

我的目标是编写只需要遍历一次内存的代码。

更好的是,如果我可以在编译时生成如下代码:

v = [Float64(j) for j in 1:10]
a = 1
b = 2
for j in 1:length(v)
  v[j]+= a # taken from add!(v::Vector{Float64}, a::Float64, j::Int64)
  v[j]+= b # taken from add!(v::Vector{Float64}, a::Float64, j::Int64)
end

但是我不确定与这个玩具示例相比,对于我正在考虑的更复杂的情况,这是否可行?

**更新2 **

这里是@BogumiłKamiński的解决方案的MWE,除了我已将宏调用移到一个函数中之外,现在它不起作用了,因为它抱怨未定义v_test

macro vectorize(args...)
    expr = :()
    for arg in args
        a = deepcopy(arg) # for safety in case arg is also used somewhere else
        push!(a.args, :j)
        expr = :($expr; $a)
    end
    quote
        for j in 1:length(v)
            $expr
        end
    end
end

function add!(v::Vector{Float64}, a::Float64)
  for j in 1:length(v)
    v[j]+= a
  end
end

function add!(v::Vector{Float64}, a::Float64, j::Int64)
  v[j]+= a
end

v = [Float64(j) for j in 1:10]
a = 1.0
b = 2.0

function test!(v_test, a_test, b_test)
  @vectorize(
  add!(v_test, a_test),
  add!(v_test, b_test)
  )
end

test!(v, a, b)

1 个答案:

答案 0 :(得分:1)

这是您想要的吗?

macro vectorize(args...)
    expr = :()
    for arg in args
        a = deepcopy(arg) # for safety in case arg is also used somewhere else
        push!(a.args, :j)
        expr = :($expr; $a)
    end
    quote
        for j in 1:length(v)
            $expr
        end
    end
end

现在

function add!(v::Vector{Float64}, a::Float64)
  for j in 1:length(v)
    v[j]+= a
  end
end

function add!(v::Vector{Float64}, a::Float64, j::Int64)
  v[j]+= a
end

v = [Float64(j) for j in 1:10]
a = 1.0
b = 2.0

@vectorize(add!(v, a), add!(v, b))

请注意,我已经更改了ab的定义,因为您的add!需要Float64作为第二个参数。

编辑:如果要在函数内部使用此宏,最简单的操作是esc其整个返回值:

macro vectorize(args...)
    expr = :()
    for arg in args
        a = deepcopy(arg) # for safety in case arg is also used somewhere else
        push!(a.args, :j)
        expr = :($expr; $a)
    end
    esc(quote
        for j in 1:length(v)
            $expr
        end
    end)
end

然后您可以定义例如:

function f()
    v = [Float64(j) for j in 1:10]
    a = 1.0
    b = 2.0
    @vectorize(add!(v, a), add!(v, b))
    v
end

并运行f()在全局范围内获得与上述相同的结果。

编辑2:我刚刚意识到实际上我必须清理j,否则以下代码将失败:

test!(v_test, j, b_test) =
    @vectorize(add!(v_test, j), add!(v_test, b_test))

这是您应该怎么做:

macro vectorize(args...)
    expr = :()
    j = gensym()
    for arg in args
        a = deepcopy(arg) # for safety in case arg is also used somewhere else
        push!(a.args, j)
        expr = :($expr; $a)
    end
    esc(quote
        for $j in 1:length(v)
            $expr
        end
    end)
end

如您所见,开发宏不是一项显而易见的任务(希望最终的配方没有错误:))。

编辑3:这是正确处理length的代码。同样,现在在每个表达式中,实际上您可以传递一个不同的值作为第一个参数(因此您可以独立处理不同的向量)。如果确实要处理相同的向量,则检查a.args[2]始终是相同的符号:

macro vectorize(args...)
    expr = :()
    j = gensym()
    for arg in args
        a = deepcopy(arg) # for safety in case arg is also used somewhere else
        var = a.args[2]
        push!(a.args, j)
        q = quote
            for $j in 1:length($var)
                $a
            end
        end
        expr = :($expr; $q)
    end
    esc(expr)
end