figure 3. interface traits defining staged operations. for simplicity, operations are defined for Double only.

trait Base {
type Rep[+T]
}
trait Arith extends Base {
implicit def unit(x: Double): Rep[Double]
def infix_+(x: Rep[Double], y: Rep[Double]):
Rep[Double]
def infix_*(x: Rep[Double], y: Rep[Double]):
Rep[Double]
...
}
trait Trig extends Base {
def cos(x: Rep[Double]): Rep[Double]
def sin(x: Rep[Double]): Rep[Double]
}

implicit lifting of Doubles to staged values and the usual arithmetic operations on staged expressions of type Rep[Double]. The restriction to Doubles is just to keep the presentation concise. Any suitable means to abstract over numeric types, such as the “type class” Numeric from the Scala standard library, could be used to define Arith in a generic way for a range of numeric types. Analogous to Double, we could define arithmetic on matrices and vectors and implement optimizations on those operations in exactly the same way. Trait Trig is similar to Arith but defines trigonometric operations.

One way to look at Base, Arith, and Trig is as the definition of a typed embedded language (or its syntax). The embedding is tagless (i.e., method resolution happens at compile time without additional runtime dispatch overhead) 2 and polymorphic, 8 in the sense that we are free to pick any suitable concrete implementation that fulfills the given interface.

From a safety point of view, keeping the actual representation inaccessible to the program generator is very important. Otherwise, the program generator could execute a different code depending on the exact structure of a staged expression. Optimizations that replace staged code with simpler but semantically equivalent expressions would risk changing the meaning of the generated program.

2. 1. Representing staged code

With the aim of generating code, we could represent staged expressions directly as strings. But for optimization purposes we would rather have a structured intermediate representation that we can analyze in various ways.

We choose a representation on the basis of expression trees or, more exactly, a “sea of nodes” representation that is in fact a directed (and for the moment, acyclic) graph but can be accessed through a tree-like interface. The necessary infrastructure is defined in trait Expressions, shown in Figure 4.

figure 4. expression representation (method implementations have been omitted).

 

trait Expressions { // expressions (atomic) abstract class Exp[+T] case class Const[T](x: T) extends Exp[T] case class Sym[T](n: Int) extends Exp[T]

def fresh[T]: Sym[T] // block scopes abstract class Block[T]

def reifyBlock[T](x: => Exp[T]): Block[T] // definitions (composite, subclasses provided // by other traits) abstract class Def[T]

def findDefinition[T](s: Sym[T]): Option[Def[T]]
def findDefinition[T](d: Def[T]): Option[Sym[T]]
def findOrCreateDefinition[T](d: Def[T]): Sym[T]
// bind definitions to symbols automatically
implicit def toAtom[T](d: Def[T]): Exp[T] =
// pattern match on definition of a given symbol
object Def {
}
}

findOrCreateDefinition(d)

def unapply[T](s: Sym[T]): Option[Def[T]] = findDefinition(s)

There are three categories of objects involved: expressions, which are atomic (subclasses of Exp: constants and symbols; with a “gensym” operator fresh to create fresh symbols), definitions, which represent composite operations (subclasses of Def, to be provided by other components), and blocks, which model nested scopes.

The guiding principle is that each definition has an associated symbol and refers to other definitions only via their symbols. In effect, this means that every composite value will be named, similar to administrative normal form (ANF). The trait Expressions provides methods to find a definition given a symbol, or vice versa. The extractor object Def allows one to pattern-match on the definition of a given symbol, a facility that is used for implementing rewritings (see below).

The framework ensures that the code that contains staging operations will always be executed within the dynamic scope of at least one invocation of reifyBlock, which returns a block object and takes as call-by-name argument the present-stage expression that will compute the staged block result. Block objects can be part of definitions, for example, for loops or conditionals.

The implicit conversion method toAtom converts a definition to an atomic expression and links it to the scope being built up by the innermost enclosing reifyBlock