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