Macro to serialize a linear solver code

New to Julia
,
1

Hi,

I am converting a Matlab code to Julia, learning metaprogramming and macros.

The code converts a linear solver algorithm into a sequence of operations

  • mapping all variables to a single vector (target language limitation)
  • avoiding operations involving zero (sparse matrix and RHS, Markowitz reordering)
A[1, 1] = A[1, 1] + A[1, 1] * b[1] --> Ab[1] = Ab[1] + Ab[1] * Ab[2]
A[1, 2] = A[1, 2] + A[2, 1] * b[1] --> Ab[3] = Ab[3] + Ab[4] * Ab[2]
A[2, 2] = A[2, 2] + A[2, 2] * b[2] --> Ab[5] = Ab[5] + Ab[5] * Ab[6]

This is my first code using Expr type.
For the macro I see several obstacles and would be happy for your guidance:

  • The dictionary is globally defined. How to avoid passing it through all hierarchies?
  • I needed to convert the variable expression into a string in order to use it as dictionary key.
    vardict = Dict{Expr, Int}() and key = e.args did not add new keys.
  • The formulation of ex = :(A[$i,$j] = A[$i,$j] + A[$j,$i]*b[$i]) is cumbersome, hopefully can be expressed by a macro like @expr2vec A[i,j] = A[i,j] + A[j,i]*b[i]
  • The original and serialized codes are to be executed side-by side to see errors or numerical issues.
vardict = Dict{String, Int}()

function var2idx(e)
  for ii in 2:length(e.args)
    e.args[ii] = eval(e.args[ii])
  end
  key = string(e.args) # alternative?
  if !haskey(vardict, key)
    vardict[key] = length(vardict) + 1; 
  end
  vardict[key]
end

function map2vec!(e)
  e.args[2] = var2idx(e)
  if e.args[1] == :A
    pop!(e.args)
  end
  e.args[1] = :Ab
end

function expr2vec(ex)
  for e in filter(e->isa(e, Expr), ex.args)
    if e.head == :ref
      map2vec!(e)
    elseif e.head == :call
      expr2vec(e)
    else
      error("unknown e.head: $(e.head)")
    end
  end
  ex
end

function solver()
  n = 2
  A = randn(n,n)
  b = randn(n)
  for i = 1:n, j = i:n
    
    #  @expr2vec A[i,j] = A[i,j] + A[j,i]*b[i]

    ex = :(A[$i,$j] = A[$i,$j] + A[$j,$i]*b[$i])
    print(ex)
    vex = expr2vec(ex)
    println(" --> ", vex)
  end
end

solver()
2

I wouldn’t do this with a macro. If you use StaticArrays.jl you’ll get the same codegen, and if you use Symbolics.jl you can easily trace it to symbolic form and call build_function for codegen.

using Symbolics, StaticArrays
@variables A[1:4,1:4] B[1:4,1:4]
A,B = collect(A), collect(B)
A*B

4×4 Matrix{Num}:
 A[1, 1]*B[1, 1] + A[1, 2]*B[2, 1] + A[1, 3]*B[3, 1] + A[1, 4]*B[4, 1]  …  A[1, 1]*B[1, 4] + A[1, 2]*B[2, 4] + A[1, 3]*B[3, 4] + A[1, 4]*B[4, 4]
 A[2, 1]*B[1, 1] + A[2, 2]*B[2, 1] + A[2, 3]*B[3, 1] + A[2, 4]*B[4, 1]     A[2, 1]*B[1, 4] + A[2, 2]*B[2, 4] + A[2, 3]*B[3, 4] + A[2, 4]*B[4, 4]
 A[3, 1]*B[1, 1] + A[3, 2]*B[2, 1] + A[3, 3]*B[3, 1] + A[3, 4]*B[4, 1]     A[3, 1]*B[1, 4] + A[3, 2]*B[2, 4] + A[3, 3]*B[3, 4] + A[3, 4]*B[4, 4]
 A[4, 1]*B[1, 1] + A[4, 2]*B[2, 1] + A[4, 3]*B[3, 1] + A[4, 4]*B[4, 1]     A[4, 1]*B[1, 4] + A[4, 2]*B[2, 4] + A[4, 3]*B[3, 4] + A[4, 4]*B[4, 4]

build_function(A*B,A,B)

:(function (ˍ₋out, ˍ₋arg1, ˍ₋arg2)
      #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:303 =#
      #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:304 =#
      let var"A[1, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[1]), var"A[2, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[2]), var"A[3, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[3]), var"A[4, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[4]), var"A[1, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[5]), var"A[2, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[6]), var"A[3, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[7]), var"A[4, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[8]), var"A[1, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[9]), var"A[2, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[10]), var"A[3, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[11]), var"A[4, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[12]), var"A[1, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[13]), var"A[2, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[14]), var"A[3, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[15]), var"A[4, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg1[16]), var"B[1, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[1]), var"B[2, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[2]), var"B[3, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[3]), var"B[4, 1]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[4]), var"B[1, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[5]), var"B[2, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[6]), var"B[3, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[7]), var"B[4, 2]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[8]), var"B[1, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[9]), var"B[2, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[10]), var"B[3, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[11]), var"B[4, 3]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[12]), var"B[1, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[13]), var"B[2, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[14]), var"B[3, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[15]), var"B[4, 4]" = #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:190 =# @inbounds(ˍ₋arg2[16])
          #= C:\Users\accou\.julia\packages\Symbolics\LfLYY\src\build_function.jl:374 =#
          #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:350 =# @inbounds begin
                  #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:346 =#
                  ˍ₋out[1] = (+)((+)((+)((*)(var"A[1, 1]", var"B[1, 1]"), (*)(var"A[1, 2]", var"B[2, 1]")), (*)(var"A[1, 3]", var"B[3, 1]")), (*)(var"A[1, 4]", var"B[4, 1]"))
                  ˍ₋out[2] = (+)((+)((+)((*)(var"A[2, 1]", var"B[1, 1]"), (*)(var"A[2, 2]", var"B[2, 1]")), (*)(var"A[2, 3]", var"B[3, 1]")), (*)(var"A[2, 4]", var"B[4, 1]"))
                  ˍ₋out[3] = (+)((+)((+)((*)(var"A[3, 1]", var"B[1, 1]"), (*)(var"A[3, 2]", var"B[2, 1]")), (*)(var"A[3, 3]", var"B[3, 1]")), (*)(var"A[3, 4]", var"B[4, 1]"))
                  ˍ₋out[4] = (+)((+)((+)((*)(var"A[4, 1]", var"B[1, 1]"), (*)(var"A[4, 2]", var"B[2, 1]")), (*)(var"A[4, 3]", var"B[3, 1]")), (*)(var"A[4, 4]", var"B[4, 1]"))
                  ˍ₋out[5] = (+)((+)((+)((*)(var"A[1, 1]", var"B[1, 2]"), (*)(var"A[1, 2]", var"B[2, 2]")), (*)(var"A[1, 3]", var"B[3, 2]")), (*)(var"A[1, 4]", var"B[4, 2]"))
                  ˍ₋out[6] = (+)((+)((+)((*)(var"A[2, 1]", var"B[1, 2]"), (*)(var"A[2, 2]", var"B[2, 2]")), (*)(var"A[2, 3]", var"B[3, 2]")), (*)(var"A[2, 4]", var"B[4, 2]"))
                  ˍ₋out[7] = (+)((+)((+)((*)(var"A[3, 1]", var"B[1, 2]"), (*)(var"A[3, 2]", var"B[2, 2]")), (*)(var"A[3, 3]", var"B[3, 2]")), (*)(var"A[3, 4]", var"B[4, 2]"))
                  ˍ₋out[8] = (+)((+)((+)((*)(var"A[4, 1]", var"B[1, 2]"), (*)(var"A[4, 2]", var"B[2, 2]")), (*)(var"A[4, 3]", var"B[3, 2]")), (*)(var"A[4, 4]", var"B[4, 2]"))
                  ˍ₋out[9] = (+)((+)((+)((*)(var"A[1, 1]", var"B[1, 3]"), (*)(var"A[1, 2]", var"B[2, 3]")), (*)(var"A[1, 3]", var"B[3, 3]")), (*)(var"A[1, 4]", var"B[4, 3]"))
                  ˍ₋out[10] = (+)((+)((+)((*)(var"A[2, 1]", var"B[1, 3]"), (*)(var"A[2, 2]", var"B[2, 3]")), (*)(var"A[2, 3]", var"B[3, 3]")), (*)(var"A[2, 4]", var"B[4, 3]"))
                  ˍ₋out[11] = (+)((+)((+)((*)(var"A[3, 1]", var"B[1, 3]"), (*)(var"A[3, 2]", var"B[2, 3]")), (*)(var"A[3, 3]", var"B[3, 3]")), (*)(var"A[3, 4]", var"B[4, 3]"))
                  ˍ₋out[12] = (+)((+)((+)((*)(var"A[4, 1]", var"B[1, 3]"), (*)(var"A[4, 2]", var"B[2, 3]")), (*)(var"A[4, 3]", var"B[3, 3]")), (*)(var"A[4, 4]", var"B[4, 3]"))
                  ˍ₋out[13] = (+)((+)((+)((*)(var"A[1, 1]", var"B[1, 4]"), (*)(var"A[1, 2]", var"B[2, 4]")), (*)(var"A[1, 3]", var"B[3, 4]")), (*)(var"A[1, 4]", var"B[4, 4]"))
                  ˍ₋out[14] = (+)((+)((+)((*)(var"A[2, 1]", var"B[1, 4]"), (*)(var"A[2, 2]", var"B[2, 4]")), (*)(var"A[2, 3]", var"B[3, 4]")), (*)(var"A[2, 4]", var"B[4, 4]"))
                  ˍ₋out[15] = (+)((+)((+)((*)(var"A[3, 1]", var"B[1, 4]"), (*)(var"A[3, 2]", var"B[2, 4]")), (*)(var"A[3, 3]", var"B[3, 4]")), (*)(var"A[3, 4]", var"B[4, 4]"))
                  ˍ₋out[16] = (+)((+)((+)((*)(var"A[4, 1]", var"B[1, 4]"), (*)(var"A[4, 2]", var"B[2, 4]")), (*)(var"A[4, 3]", var"B[3, 4]")), (*)(var"A[4, 4]", var"B[4, 4]"))
                  #= C:\Users\accou\.julia\packages\SymbolicUtils\0KTj4\src\code.jl:348 =#
                  nothing
              end
      end
  end)
4 Likes
3

Thanks Chris for pointing out Symbolics.jl as alternative!

However, I am not sure if this approach can lead to the wanted output format:

As said, in my code, the “intelligence” (iszero testing) is in the logic around the for-loops using Markowitz. I assume Symbolics.jl uses a mix of CSE and Dulmage–Mendelsohn.
It would be interesting to create a benchmark.

4

@shashi do we have a pass that does the same thing as GitHub - SciML/MuladdMacro.jl: This package contains a macro for converting expressions to use muladd calls and fused-multiply-add (FMA) operations for high-performance in the SciML scientific machine learning ecosystem ?

The C-code in that tutorial is the C code that is automatically generated.

It would be. Though if you’re looking for very efficient looping stuff, you should take a look at GitHub - JuliaSIMD/LoopVectorization.jl: Macro(s) for vectorizing loops. which has been used to write BLAS functions that outperform MKL on some systems.