Lab 06: Code introspection and metaprogramming

In this lab we are first going to inspect some tooling to help you understand what Julia does under the hood such as:

• looking at the code at different levels

• understanding what method is being called

• showing different levels of code optimization

Secondly we will start playing with the metaprogramming side of Julia, mainly covering:

• how to view abstract syntax tree (AST) of Julia code

• how to manipulate AST

These topics will be extended in the next lecture/lab, where we are going use metaprogramming to manipulate code with macros.

We will be again a little getting ahead of ourselves as we are going to use quite a few macros, which will be properly explained in the next lecture as well, however for now the important thing to know is that a macro is just a special function, that accepts as an argument Julia code, which it can modify.

Quick reminder of introspection tooling

Let's start with the topic of code inspection, e.g. we may ask the following: What happens when Julia evaluates [i for i in 1:10]?

parsing

julia> :([i for i in 1:10]) |> dump
Expr
  head: Symbol comprehension
  args: Array{Any}((1,))
    1: Expr
      head: Symbol generator
      args: Array{Any}((2,))
        1: Symbol i
        2: Expr
          head: Symbol =
          args: Array{Any}((2,))
            1: Symbol i
            2: Expr
              head: Symbol call
              args: Array{Any}((3,))
                1: Symbol :
                2: Int64 1
                3: Int64 10

lowering

Meta.@lower debuginfo=:none [i for i in 1:10]

typing

f() = [i for i in 1:10]
@code_typed debuginfo=:none f()

LLVM code generation

@code_llvm debuginfo=:none f()

native code generation

@code_native debuginfo=:none f()

Let's see how these tools can help us understand some of Julia's internals on examples from previous labs and lectures.

Understanding runtime dispatch and type instabilities

We will start with a question: Can we spot internally some difference between type stable/unstable code?

Exercise

Inspect the following two functions using @code_lowered, @code_typed, @code_llvm and @code_native.

x = rand(10^5)
function explicit_len(x)
    length(x)
end

function implicit_len()
    length(x)
end

For now do not try to understand the details, but focus on the overall differences such as length of the code.

Redirecting stdout

If the output of the method introspection tools is too long you can use a general way of redirecting standard output stdout to a file

open("./llvm_fun.ll", "w") do file
    original_stdout = stdout
    redirect_stdout(file)
    @code_llvm debuginfo=:none fun()
    redirect_stdout(original_stdout)
end

In case of @code_llvm and @code_native there are special options, that allow this out of the box, see help ? for underlying code_llvm and code_native. If you don't mind adding dependencies there is also the @capture_out from Suppressor.jl

:::

Details

@code_warntype explicit_len(x)
@code_warntype implicit_len()

@code_typed debuginfo=:none explicit_len(x)
@code_typed debuginfo=:none implicit_len()

@code_llvm debuginfo=:none explicit_len(x)
@code_llvm debuginfo=:none implicit_len()

@code_native debuginfo=:none explicit_len(x)
@code_native debuginfo=:none implicit_len()

In this case we see that the generated code for such a simple operation is much longer in the type unstable case resulting in longer run times. However in the next example we will see that having longer code is not always a bad thing.

Loop unrolling

In some cases the compiler uses loop unrolling^[1] optimization to speed up loops at the expense of binary size. The result of such optimization is removal of the loop control instructions and rewriting the loop into a repeated sequence of independent statements.

Exercise

Inspect under what conditions does the compiler unroll the for loop in the polynomial function from the last lab.

function polynomial(a, x)
    accumulator = a[end] * one(x)
    for i in length(a)-1:-1:1
        accumulator = accumulator * x + a[i]
    end
    accumulator
end

Compare the speed of execution with and without loop unrolling.

HINTS:

• these kind of optimization are lower level than intermediate language

• loop unrolling is possible when compiler knows the length of the input

Details

using BenchmarkTools
a = Tuple(ones(20)) # tuple has known size
ac = collect(a)
x = 2.0

@code_lowered polynomial(a,x)       # cannot be seen here as optimizations are not applied
@code_typed debuginfo=:none polynomial(a,x)         # loop unrolling is not part of type inference optimization

More than 2x speedup

julia> @btime polynomial($a,$x)
  7.769 ns (0 allocations: 0 bytes)
1.048575e6

julia> @btime polynomial($ac,$x)
  13.005 ns (0 allocations: 0 bytes)
1.048575e6

@code_llvm debuginfo=:none polynomial(a,x)
@code_llvm debuginfo=:none polynomial(ac,x)

Recursion inlining depth

Inlining^[2] is another compiler optimization that allows us to speed up the code by avoiding function calls. Where applicable compiler can replace f(args) directly with the function body of f, thus removing the need to modify stack to transfer the control flow to a different place. This is yet another optimization that may improve speed at the expense of binary size.

Exercise

Rewrite the polynomial function from the last lab using recursion and find the length of the coefficients, at which inlining of the recursive calls stops occurring.

function polynomial(a, x)
    accumulator = a[end] * one(x)
    for i in length(a)-1:-1:1
        accumulator = accumulator * x + a[i]
    end
    accumulator  
end

HINTS:

• define two methods _polynomial!(ac, x, a...) and _polynomial!(ac, x, a) for the case of ≥2 coefficients and the last coefficient

• use splatting together with range indexing a[1:end-1]...

• the correctness can be checked using the built-in evalpoly

• recall that these kind of optimization are possible just around the type inference stage

• use container of known length to store the coefficients

Splatting/slurping operator ...

The operator ... serves two purposes inside function calls ^[3][4]:

• combines multiple arguments into one

julia> function printargs(args...)
           println(typeof(args))
           for (i, arg) in enumerate(args)
               println("Arg #$i = $arg")
           end
       end
printargs (generic function with 1 method)

julia> printargs(1, 2, 3)
Tuple{Int64, Int64, Int64}
Arg #1 = 1
Arg #2 = 2
Arg #3 = 3

• splits one argument into many different arguments

julia> function threeargs(a, b, c)
           println("a = $a::$(typeof(a))")
           println("b = $b::$(typeof(b))")
           println("c = $c::$(typeof(c))")
       end
threeargs (generic function with 1 method)

julia> threeargs([1,2,3]...) # or with a variable threeargs(x...)
a = 1::Int64
b = 2::Int64
c = 3::Int64

:::

Details

_polynomial!(ac, x, a...) = _polynomial!(x * ac + a[end], x, a[1:end-1]...)
_polynomial!(ac, x, a) = x * ac + a
polynomial(a, x) = _polynomial!(a[end] * one(x), x, a[1:end-1]...)

# the coefficients have to be a tuple
a = Tuple(ones(Int, 21)) # everything less than 22 gets inlined
x = 2
polynomial(a,x) == evalpoly(x,a) # compare with built-in function

# @code_llvm debuginfo=:none polynomial(a,x)    # seen here too, but code_typed is a better option
@code_lowered polynomial(a,x) # cannot be seen here as optimizations are not applied

@code_typed debuginfo=:none polynomial(a,x)

AST manipulation: The first steps to metaprogramming

Julia is so called homoiconic language, as it allows the language to reason about its code. This capability is inspired by years of development in other languages such as Lisp, Clojure or Prolog.

There are two easy ways to extract/construct the code structure ^[5]

• parsing code stored in string with internal Meta.parse

code_parse = Meta.parse("x = 2")    # for single line expressions (additional spaces are ignored)
code_parse_block = Meta.parse("""
begin
    x = 2
    y = 3
    x + y
end
""") # for multiline expressions

• constructing an expression using quote ... end or simple :() syntax

code_expr = :(x = 2)    # for single line expressions (additional spaces are ignored)
code_expr_block = quote
    x = 2
    y = 3
    x + y   
end # for multiline expressions

Results can be stored into some variables, which we can inspect further.

julia> typeof(code_parse)
ERROR: UndefVarError: `code_parse` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

julia> dump(code_parse)
ERROR: UndefVarError: `code_parse` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

julia> typeof(code_parse_block)
ERROR: UndefVarError: `code_parse_block` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

julia> dump(code_parse_block)
ERROR: UndefVarError: `code_parse_block` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

The type of both multiline and single line expression is Expr with fields head and args. Notice that Expr type is recursive in the args, which can store other expressions resulting in a tree structure - abstract syntax tree (AST) - that can be visualized for example with the combination of GraphRecipes and Plots packages.

using GraphRecipes #hide
using Plots #hide
plot(code_expr_block, fontsize=12, shorten=0.01, axis_buffer=0.15, nodeshape=:rect)

This recursive structure has some major performance drawbacks, because the args field is of type Any and therefore modifications of this expression level AST won't be type stable. Building blocks of expressions are Symbols and literal values (numbers).

A possible nuisance of working with multiline expressions is the presence of LineNumber nodes, which can be removed with Base.remove_linenums! function.

julia> Base.remove_linenums!(code_parse_block)
ERROR: UndefVarError: `code_parse_block` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

Parsed expressions can be evaluate using eval function.

julia> eval(code_parse)    # evaluation of :(x = 2)
ERROR: UndefVarError: `code_parse` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

julia> x                   # should be defined
ERROR: UndefVarError: `x` not defined in `Main`
Suggestion: check for spelling errors or missing imports.

Exercise

Before doing anything more fancy let's start with some simple manipulation of ASTs.

• Define a variable code to be as the result of parsing the string "j = i^2".

• Copy code into a variable code2. Modify this to replace the power 2 with a power 3. Make sure that the original code variable is not also modified.

• Copy code2 to a variable code3. Replace i with i + 1 in code3.

• Define a variable i with the value 4. Evaluate the different code expressions using the eval function and check the value of the variable j.

Details

julia> code = Meta.parse("j = i^2")
:(j = i ^ 2)

julia> code2 = copy(code)
:(j = i ^ 2)

julia> code2.args[2].args[3] = 3
3

julia> code3 = copy(code2)
:(j = i ^ 3)

julia> code3.args[2].args[2] = :(i + 1)
:(i + 1)

julia> i = 4
4

julia> eval(code), eval(code2), eval(code3)
(16, 64, 125)

Following up on the more general substitution of variables in an expression from the lecture, let's see how the situation becomes more complicated, when we are dealing with strings instead of a parsed AST.

Exercise

replace_i(s::Symbol) = s == :i ? :k : s
replace_i(e::Expr) = Expr(e.head, map(replace_i, e.args)...)
replace_i(u) = u

Given a function replace_i, which replaces variables i for k in an expression like the following

julia> ex = :(i + i*i + y*i - sin(z))
:((i + i * i + y * i) - sin(z))

julia> @test replace_i(ex) == :(k + k*k + y*k - sin(z))
Test Passed

write a different function sreplace_i(s), which does the same thing but instead of a parsed expression (AST) it manipulates a string, such as

julia> s = string(ex)
"(i + i * i + y * i) - sin(z)"

HINTS:

• Use Meta.parse in combination with replace_i ONLY for checking of correctness.

• You can use the replace function in combination with regular expressions.

• Think of some corner cases, that the method may not handle properly.

Details

The naive solution

julia> sreplace_i(s) = replace(s, 'i' => 'k')
sreplace_i (generic function with 1 method)

@test Meta.parse(sreplace_i(s)) == replace_i(Meta.parse(s))

does not work in this simple case, because it will replace "i" inside the sin(z) expression. We can play with regular expressions to obtain something, that is more robust

sreplace_i(s) = replace(s, r"([^\w]|\b)i(?=[^\w]|\z)" => s"\1k")
@test Meta.parse(sreplace_i(s)) == replace_i(Meta.parse(s))

however the code may now be harder to read. Thus it is preferable to use the parsed AST when manipulating Julia's code.

If the exercises so far did not feel very useful let's focus on one, that is similar to a part of the IntervalArithmetics.jl pkg.

Exercise

Write function wrap!(ex::Expr) which wraps literal values (numbers) with a call to f(). You can test it on the following example

f = x -> convert(Float64, x)
ex = :(x*x + 2*y*x + y*y)     # original expression
rex = :(x*x + f(2)*y*x + y*y) # result expression

HINTS:

• use recursion and multiple dispatch

• dispatch on ::Number to detect numbers in an expression

• for testing purposes, create a copy of ex before mutating

Details

julia> function wrap!(ex::Expr)
           args = ex.args
       
           for i in 1:length(args)
               args[i] = wrap!(args[i])
           end
       
           return ex
       end
wrap! (generic function with 1 method)

julia> wrap!(ex::Number) = Expr(:call, :f, ex)
wrap! (generic function with 2 methods)

julia> wrap!(ex) = ex
wrap! (generic function with 3 methods)

julia> ext, x, y = copy(ex), 2, 3
(:(x * x + 2 * y * x + y * y), 2, 3)

julia> @test wrap!(ex) == :(x*x + f(2)*y*x + y*y)
Test Passed

julia> eval(ext)
25

julia> eval(ex)
25.0

This kind of manipulation is at the core of some pkgs, such as aforementioned IntervalArithmetics.jl where every number is replaced with a narrow interval in order to find some bounds on the result of a computation.

---

Resources

• Julia's manual on metaprogramming

• David P. Sanders' workshop @ JuliaCon 2021

• Steven Johnson's keynote talk @ JuliaCon 2019

• Andy Ferris's workshop @ JuliaCon 2018

• From Macros to DSL by John Myles White

• Notes on JuliaCompilerPlugin

---

1. https://en.wikipedia.org/wiki/Loop_unrolling ↩︎

2. https://en.wikipedia.org/wiki/Inline_expansion ↩︎

3. https://docs.julialang.org/en/v1/manual/faq/#What-does-the-...-operator-do? ↩︎

4. https://docs.julialang.org/en/v1/manual/functions/#Varargs-Functions ↩︎

5. Once you understand the recursive structure of expressions, the AST can be constructed manually like any other type. ↩︎