Tutorial: Laziness and Dataflow

This tutorial is about lazy evaluation and dataflow programming in Owl. Different from Haskell, we know that OCaml eagerly evaluates the expressions in a program. However, laziness can be very helpful in numerical applications which often involve complicated computation graphs and large data chunks such as matrices. Laziness allows us to reduce the memory allocation, optimise computation graph, incrementally update results, and etc.

Owl is able to simulate the lazy evaluation atop of OCaml using it Lazy functor. The functor has been optimised for memory allocation. Moreover, Lazy functor also provides a dataflow programming model to allow us explicitly constructing a computation graph.

The examples used in this tutorial can be found in lazy_eval.ml.

Core APIs

The Lazy functor in Owl is specifically designed for numerical computing which uses ndarray as its core data structure. It does not aim to provide a general framework beyond aforementioned numerical scope. To generate a new module, you need to pass in Owl's ndarray module.

module M = Owl.Lazy.Make (Dense.Ndarray.D);;
module N = Owl.Lazy.Make (Dense.Ndarray.S);;

The code above generates two lazy modules supporting both single and double precision float numbers respectively. In the generated module, there are several core APIs I want you to know before you start playing with it.

val variable : unit -> t : this function creates a placeholder for a variable in the computation graph. This is where you can start constructing a computation graph. The variables are the inputs of the graph, and you can plug in actual values later.
val assign_arr : t -> A.arr -> unit : this function assigns a specific ndarray value to the variable previously created by variable function. Note that you can only assign values to those created by variable function, otherwise an exception will be thrown.
val assign_elt : t -> A.elt -> unit : similar to assign_arr, this function assigns a value of type A.elt to a variable. A.elt depends on the module you plug into Lazy functor, but it refers to a float number in our current context.
val eval : t -> unit : this function evaluates the associated computation graph for a given variable of type [t].
val to_arr : t -> A.arr : this function unpacks the value of an ndarray from a variable of type t. You can only use this function after you call eval.
val to_elt : t -> A.elt : Similarly, to_elt is used for unpacking the value of float number out of a variable.

What I want to emphasis here is: assign_arr and assign_elt can only be applied to those variables previously created by the variable function. There are other ways to plug in the inputs into a computation graph, e.g. the following two functions.

val of_arr : A.arr -> t : pack an ndarray and return it as a constant of type t.
val of_elt : A.elt -> t : pack an element and return it as a constant of type t.

The difference is that inputs created by of_arr and of_elt are treated as constants and their values cannot be changed in the evaluation of a computation graph. Therefore, you cannot use assign_arr and assign_elt to change the value of inputs created by these two functions.

A Simple Example

The following code snippet shows how to perform a sequence of computation using ndarray with eager evaluation. It is just like how you write normal application in Owl normally, you must be very familiar with this already.

let eager_eval x () =
  Arr.(add x x |> log |> sin |> neg |> cos |> abs |> sinh |> round |> atan)

But how can we perform the same computation using lazy evaluation? The following code demonstrates how Lazy functor can help us to achieve this.

let lazy_eval x () =
  let a = M.variable () in
  let z = M.(add a a |> log |> sin |> neg |> cos |> abs |> sinh |> round |> atan) in
  M.assign_arr a x;
  M.eval z

As you can see, we first define a placeholder for variable a. Then we use the predefined math operators in Lazy functor to construct the computation graph, and the code looks actually exactly the same as that in eager evaluation.

Comparing to eager evaluation, we need to do some extra steps: first, assign value x to variable a; second, evaluate the computation graph by calling eval function. z is the output of the computation graph, you can also think z represents the expression that produces the final value.

We did not really extract the final value from z, but it is trivial if you want to do so, just call M.to_arr z.

Visualise Computation Graph

Lazy functor provides two APIs that help you to visualise the computation graph corresponding to an expression.

val to_trace : t list -> string : the function returns a string which can be printed out on the terminal directly.
val to_dot : t list -> string : the function also returns a dot-formatted string to represent the graph structure. Because it is dot format, you can save it to a file and use another tool to produce the figure such as GraphViz.

For example, the following code prints out the trace using to_trace function.

let print_trace a =
  let x = M.variable () in
  let y = M.variable () in
  let z = M.(add x y |> sin |> abs |> log) in
  M.assign_arr x a;
  M.assign_arr y a;
  M.eval z;
  M.to_trace [z] |> print_endline

You can probably see similar output on your screen if you try it out in utop. "valid" and "invalid" states relate to the internal mechanism of Lazy functor and we can simply ignore them at the moment.

[ #5 name:abs state:invalid ] -> [ #6 name:log state:valid ]
[ #4 name:sin state:invalid ] -> [ #5 name:abs state:invalid ]
[ #3 name:add state:invalid ] -> [ #4 name:sin state:invalid ]
[ #1 name:variable state:valid ] -> [ #3 name:add state:invalid ]
[ #2 name:variable state:valid ] -> [ #3 name:add state:invalid ]

What if you want to actually visualise the graph in to a picture, then you need to use to_trace function as the following code snippet does.

let print_graph () =
  let x = M.variable () in
  let y = M.variable () in
  let z = M.(add x y |> dot x |> sin |> abs |> sum' |> add_scalar x |> log |> atan2 y |> neg |> relu) in
  print_endline "visualise computation graph to dot file ...";
  M.to_dot [z] |> Utils.write_file "plot_lazy.dot";
  Sys.command "dot -Tpdf plot_lazy.dot -o plot_lazy.pdf"

With the code above, you can generate a nice figure as below. Note you need to install GraphViz on you local computer because Sys.command calls dot command which is a tool in GraphViz.

Incremental Computation

What is incremental computation? Many realistic numerical applications involve calculations on many variables which further creates complicated computation graph. Quite often, we only change a couple of variables and we need to re-evaluate the graph. However, re-evaluating the whole computation graph is very expensive and is not necessary in many cases.

Incremental computation refers to the situation we only re-evaluate the subgraph depending on the updated variables. This can also be referred to Self-Adjusting Computing (SAC), or dataflow programming. It has many names but the general idea is the same.

The flowing code shows how to perform incremental computation with Lazy functor.

let incremental x =
  let a = M.variable () in
  let b = M.variable () in
  let c = M.(a |> sin |> cos |> abs |> log |> sum' |> mul_scalar a |> scalar_add b) in
  M.assign_arr a x;
  M.assign_elt b 5.;
  Printf.printf "Incremental#0:\t%.3f ms\n" (Utils.time (fun () -> M.eval c));
  M.assign_elt b 6.;
  Printf.printf "Incremental#1:\t%.3f ms\n" (Utils.time (fun () -> M.eval c));
  M.assign_elt b 7.;
  Printf.printf "Incremental#2:\t%.3f ms\n" (Utils.time (fun () -> M.eval c))

In the code, a and b are variables and c represents the expression we want to evaluate. After the first evaluation, we changed the assignment of b twice, then re-evaluate the expression.

Because of the capability of incremental computation, Owl does not have to re-calculate everything in the latter evaluations. In fact, as we can see in the following figure, only the part of the computation graph in the red-dotted rectangular will be re-evaluated if we update variable b. This mechanism can significantly reduce the computation time. On my computer, the first full evaluation took about 275ms whereas the latter two incremental ones took only 30ms each, a big win!

Behind the Scene

In this section, I want to show you something under the hood. It is not critical if you just want to use Lazy functor, but knowing these details help you understand better about Owl's design.

Lazy functor is optimised for memory allocation in numerical computing, because keep allocating large chunk of memory is really expensive. To reuse the allocated one, you can use in-place modification function in Owl. E.g., Arr.sin is a pure function and always returns a new ndarray whilst Arr.sin_ performs in-place modification and overwrites the original ndarray. However, this is not safe and you should avoid using in-place modification unless you really know what you are doing.

Lazy functor is an ideal solution in this case. It automatically tracks the reference number of each node in the computation graph and tries to reuse the allocated memory as much as possible. By so doing, Lazy module often brings us a huge benefit in both performance and memory usage when large ndarrays and matrices are involved in computation. You can safely construct very complicated computation graphs without worrying about blowing up your memory. For example, you can keep transforming a big ndarray using sin function one million times and only constant memory will be used (plus some additional overhead in maintaining the graph structure).

One thing worth noting is that Lazy functor will never overwrite the value you have assigned on a variable, so you can always safely reuse them later, the functor only tries to re-use the memory of internal nodes which are opaque to the outside. In short, you do not need to worry about the issues that something might get overwritten then later cause problems in your application, Owl will take care of it.

Regarding incremental computation, every time you re-assign a new value on a variable, Lazy functor will invalidate all the nodes in the subgraph that depends on this variable, i.e. its descendant nodes. When you evaluate an expression, only the nodes in the subgraph this expression depends on will be evaluated, i.e. its ancestor nodes. As you see, the directions are different.

All right, I think you have learnt enough today in order to be Lazy in Owl :)

Enjoy Owl! Happy hacking!

Provide feedback

Saved searches