2
votes

I am trying to write a functionality, (using macro, generated function or something), that effectively vectorizes Julia function calls to functions that I've written. Basically, I'm trying to write my own version of the @. macro, but instead, I'd like it to accept functions instead of a for loop--- if I understand this correctly. Here are some documents that I've read on the subject:

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

Here is my preliminary toy example that I'm working with to achieve such a functionality:

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")

The problem I'm having so far is that I can't even get the de-vectorized version running using macros. This example results in the LoadError: UndefVarError: args not defined. I would definitely appreciate any help in getting this script working as expected (input is [1, 2, 3], and output should be [5, 6, 7]).

Any help/suggestions are greatly appreciated.

Update

More concretely, given the following defined functions:

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

I would like to be able to use a macro to convert the following lines of code:

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)
)

To generate code that is compiled like this:

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

My goal is to write code that requires a single memory traversal.

Even better, if I could generate code that looks like this at compile time:

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

But I'm not sure if this is as feasable for more complex cases that I'm considering compared to this toy example.

** Update 2**

Here is a MWE of @Bogumił Kamiński's solution---except that I've moved the macro call into a function, so now it doesn't work because it complains that v_test is not defined.

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
I am not 100% clear what you want to achieve. Could you please give a description of an input and expected processing workflow? - Bogumił Kamiński
@BogumiłKamiński, I added a more concrete example of the functionality that I'm looking to achieve. - Charles
EDIT and EDIT 2 in my answer work with your example (I leave both of them as I think it is helpful to analyze the difference between them). - Bogumił Kamiński

1 Answers

1
votes

Is this what you want?

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

and now

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))

Note that I have changed a and b definitions as your add! required Float64 as a second argument.

EDIT: If you want to use this macro inside a function the simplest thing to do is to esc its whole return value:

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

Then you can define e.g.:

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

and run f() to get the same result as above in global scope.

EDIT 2: I just realized that actually I have to sanitize j as otherwise the following code will fail:

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

Here is how you should do it:

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

As you can see developing macros is a non-obvious task (hopefully the final recipe is bug-free :)).

EDIT 3: Here is the code that correctly handles length. Also now in each expression actually you can pass a different value as a first argument (so you can independently process different vectors). If you do want to process the same vector check is a.args[2] is always the same symbol:

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