Tracking Effects

Previous: Home

---

The linear algebra DSL presented in the previous part had no operation performing side-effects.

Tracking effects requires some additional machinery, this page explains why and how.

Defining a Logging DSL

Let’s define a logging DSL with just one operation that logs a message:

// Interface: just a “log” operation
trait Log extends Base {
  def log(msg: Rep[String]): Rep[Unit]
}

// Implementation reifying the operations into an intermediate representation
trait LogExp extends Log with BaseExp {
  case class Log(msg: Exp[String]) extends Def[Unit]
  def log(msg: Exp[String]) = Log(msg)
}

// Code generator
trait ScalaGenLog extends ScalaGenBase {
  val IR: LogExp
  import IR._

  override def emitNode(sym: Sym[Any], rhs: Def[Any]) = rhs match {
    case Log(msg) =>
      stream.println(s"println(${quote(msg)})")
    case _ => super.emitNode(sym, rhs)
  }
}

We just followed the pattern previously explained: the Log trait defines the DSL interface, the LogExp trait implements the DSL by returning an abstract representation of its operations and finally the ScalaGenLog trait generates scala code for the DSL.

We can try to use it as follows:

trait LogProg extends LinearAlgebra with Log {
  def double(v: Rep[Vector]): Rep[vector] = {
    log(unit("Hello!"))
    v * unit(2.0)
  }
}

object Main extends App {
  val prog = new LogProg with EffectExp with LinearAlgebraExp with LogExp
  val codegen = new ScalaGenEffect with ScalaGenLinearAlgebra with ScalaGenLog { val IR: prog.type = prog }
  codegen.emitSource(prog.double, "Double", new java.io.PrintWriter(System.out))
}

Surprisingly, we get the following implementation for the double function:

def apply(x0:scala.collection.Seq[Double]): scala.collection.Seq[Double] = {
  val x2 = x0.map(x => x * 2.0)
  x2
}

There is no trace of the log statement. What happened?

The emitSource statement produces the code by calling the double function with an arbitrary symbol as parameter. The double function evaluates the log expression, which returns a Log node, and then evaluates the v * unit(2.0) expression, which returns a VectorScale node, and returns it. The Log node is discarded because it is not returned by the double function. The intermediate representation returned by double looks like the following:

Instead, we would like to return a block containing the Log statement and the VectorScale expression:

However, the Scala evaluation process of blocks consists in evaluating the sequence of expressions in the block and returning the value of the last expression only. we need to reify the concept of block: opening a block consists in opening a new context in which side-effectful statements are captured.

The EffectExp trait of LMS helps us to do that. It defines a reifyEffects function that opens a context (a block). Then, side-effectful statements can be captured in this context using the reflectEffect function:

trait LogExp extends Log with EffectExp {
  case class Log(msg: Exp[String]) extends Def[Unit]
  def log(msg: Exp[String]) = reflectEffect(Log(msg))
}


object Main extends App {
  val prog = new LogProg with LinearAlgebraExp with LogExp
  val codegen = new ScalaGenLinearAlgebra with ScalaGenLog { val IR: prog.type = prog }
  codegen.emitSource(prog.double, "Double", new java.io.PrintWriter(System.out))
}

We don’t need to call reifyEffects because the code generation process does it for us, before calling the double function.