GitHub

Tape anatomy

Operations

The very core of every tape is a list of operations. Let's take a look at one particular tape:

using Umlaut

foo(x) = 2x + 1
_, tape = trace(foo, 2.0)
print(tape)
Tape{Umlaut.BaseCtx}
  inp %1::typeof(Main.foo)
  inp %2::Float64
  %3 = *(2, %2)::Float64
  %4 = +(%3, 1)::Float64

Each indented line in this output represents an operation. The first 2 designate the tape inputs and have type Input. Note that the traced function itself is also recorded as an input and can be referenced from other operations on the tape, which is a typical case in closures and other callable objects. We can set new inputs to the tape as inputs!(tape, foo, 3.0).

Operations 3 and 4 represent function calls and have type Call. For example, the notation %4 = +(%3, 1) means that variable %4 is equal to the addition of variable %3 and a constant 1 (we will talk about variables in a minute). The easiest way to construct this operation is by using mkcall.

Although constants can be used directly inside Calls, sometimes we need them as separate objects on the tape. Constant operation serves exactly this role.

Finally, there's an experimental Loop operation which presents whole loops in a computational graphs and contain their own subtapes.

Variables

Variable (also aliased as just V) is a reference to an operation on tape. Variables can be bound or unbound.

Unbound variables are constructed as V(id) and point to an operation by its position on a tape. Their primary use is for indexing and short-living handling, e.g.:

import Umlaut.V

op = tape[V(4)]
%4 = +(%3, 1)::Float64

On the contrary, bound variables (created as V(op)) point to a specific operation on the tape. Even if the tape is modified, the reference is preserved. Here's an illustrative example:

vu = V(4)         # unbound
vb = V(tape[vu])  # bound, can also be created as `bound(tape, vu)`

# insert a dummy operation
insert!(tape, 3, Constant(42))
println(tape)
println("Unbound variable is still $vu")
println("Bound variable is now $vb")
Tape{Umlaut.BaseCtx}
  inp %1::typeof(Main.foo)
  inp %2::Float64
  const %3 = 42::Int64
  %4 = *(2, %2)::Float64
  %5 = +(%4, 1)::Float64

Unbound variable is still %4
Bound variable is now %5

Most functions in Umlaut create bound variables to make them resistant to transformations. Note, for example, how in the tape above the last operation automatically updated itself from +(%3, 1) to +(%4, 1). Yet sometimes explicit rebinding is neccessary, in which case rebind! can be used. Note that for rebind! to work properly with a user-defined tape context (see below), one must also implement rebind_context!

Transformations

Tapes can be modified in a variaty of ways. For this set of examples, we won't trace any function, but instead construct a tape manually:

using Umlaut
import Umlaut: Tape, V, inputs!, mkcall

tape = Tape()
# record inputs, using nothing instead of a function argument
v1, v2, v3 = inputs!(tape, nothing, 1.0, 2.0)
3-element Vector{Variable}:
 %1
 %2
 %3

push! is the standard way to add new operations to the tape, e.g.:

v4 = push!(tape, mkcall(*, v2, v3))
println(tape)
Tape{Dict{Any, Any}}
  inp %1::Nothing
  inp %2::Float64
  inp %3::Float64
  %4 = *(%2, %3)::Float64

insert! is similar to push!, but adds operation to the specified position:

v5 = insert!(tape, 4, mkcall(-, v2, 1))  # inserted before v4
println(tape)
Tape{Dict{Any, Any}}
  inp %1::Nothing
  inp %2::Float64
  inp %3::Float64
  %4 = -(%2, 1)::Float64
  %5 = *(%2, %3)::Float64

replace! is useful when you need to replace an operation with one or more other operations.

new_op1 = mkcall(/, V(2), 2)
new_op2 = mkcall(+, V(new_op1), 1)
replace!(tape, 4 => [new_op1, new_op2]; rebind_to=2)
println(tape)
Tape{Dict{Any, Any}}
  inp %1::Nothing
  inp %2::Float64
  inp %3::Float64
  %4 = /(%2, 2)::Umlaut.UncalculatedValue
  %5 = +(%4, 1)::Umlaut.UncalculatedValue
  %6 = *(%2, %3)::Float64

deleteat! is used to remove entries from the tape.

deleteat!(tape, 5; rebind_to = 1)
println(tape)
Tape{Dict{Any, Any}}
  inp %1::Nothing
  inp %2::Float64
  inp %3::Float64
  %4 = *(%2, 2)::Float64
  %5 = +(%1, %3)::Float64
  %6 = +(%1, %2)::Float64

Although trace creates a tape consisting only of primitives, tape itself can hold any function calls. It's possible to decompose all non-primitive calls on the tape to lists of corresponding primitives using primitivize!.

import Umlaut: primitivize!

f(x) = 2x - 1
g(x) = f(x) + 5

tape = Tape()
_, x = inputs!(tape, g, 3.0)
y = push!(tape, mkcall(f, x))
z = push!(tape, mkcall(+, y, 5))
tape.result = z

primitivize!(tape)

Tape execution & compilation

There are 2 ways to execute a tape. For debug purposes it's easiest to run play!:

using Umlaut
import Umlaut: play!

foo(x) = 2x + 1
_, tape = trace(foo, 2.0)

play!(tape, foo, 3.0)
7.0

compile turns the tape into a normal Julia function (subject to the World Age restriction):

using Umlaut
import Umlaut: compile

foo(x) = 2x + 1
_, tape = trace(foo, 2.0)

foo2 = compile(tape)
foo2(foo, 3.0)   # note: providing the original `foo` as the 1st argument
7.0

It's possible to see what exactly is being compiled using to_expr function.

Context (again)

We have already discussed contexts as a way to customize tracing in the Linearized traces section, but here we need to emphasize context's role as a storage. Context is attached to a Tape and can be accessed throughout its lifetime. For instance, imagine that you are working on a DSL engine which traces function execution and enriches the resulting tape with domain-specific operations. You also want to keep track of all added operations, but don't want to pass around an additional object holding them. You can attach a custom context to the tape and reference it as tape.c:

using Umlaut
import Umlaut: Variable

dsl_function(x) = ...


mutable struct DSLContext
    added_variables::Vector{Variable}
end

_, tape = trace(dsl_function, 2.0; ctx=DSLContext([]))


function add_operations(tape::Tape{DSLContext})
    v = push!(tape, ...)
    push!(tape.c.added_variables, v)
    ...
end

function process_dsl_tape(tape::Tape{DSLContext})
    vars = tape.c.added_variables
    ...
end

Just to remind you, if your context contains variables and you plan to use rebind!, you must also implement rebind_context! for your specific context type.