figure 8. Code generation interface and skeleton of scala
compilation component.
trait Compile extends Base {
def compile[A,B](f: Rep[A] => Rep[B]): A=>B
}
trait CompileScala extends Compile with ScalaGenBase =>
def compile[A,B](f: Exp[A] => Exp[B]) = {
val x = fresh[A]
val y = reifyBlock {f(x)}
// emit header
emitBlock(y)
//
//
//
//
}
}
figure 9. scala code generation for selected expressions.
trait ScalaGenBase extends BaseExp {
def emitBlock[T](Block[T]) = ...
def emitNode[T](sym: Sym[T], node: Def[T]) =
throw new Exception("node " + node + " not supported")
}
trait ScalaGenArith extends ScalaGenBase with ArithExp {
override def emitNode(sym: Sym[T], node: Def[T]) =
node match {
case Plus(a,b) => println("val %s =
%a + %b".format(sym,a,b))
case Times(a,b) => println("val %s =
%a %b".format(sym,a,b))
case _ => super.emitNode(sym, rhs)
}
}
figure 10. extending the ff T component from figure 1 with explicit
compilation.
trait FFTC extends FFT { this: Arrays with Compile =>
def fftc(size: Int) = compile {
input: Rep[Array[Double]] =>
assert(<size is power of 2>) // staging time
val arg = Array.tabulate(size) {i =>
Complex(input( 2*i), input( 2*i+ 1))
}
val res = fft(arg)
updateArray(input, res.flatMap {
case Complex(re,im) => Array(re,im)
})
}
}
type Rep[Array[Double] ]. Figure 10 shows how we
can extend the trait FFT to FFTC to obtain compiled FFT
implementations that realize the necessary data interface
for a fixed input size.
We can then define a code that creates and uses compiled
FFT “codelets” by extending FFTC:
trait TestFFTC extends FFTC {
val fft4: Array[Double] => Array[Double] =
// embedded code using fft4, fft8, ...
}
Constructing an instance of this subtrait (mixed in with
the appropriate LMS traits) will execute the embedded
code:
val OP: TestFFC = new TestFFTC with CompileScala
with ArithExpOpt with ArithExpOptFFT with
ScalaGenArith
with TrigExpOpt with ScalaGenTrig
with ArraysExpOpt with ScalaGenArrays
We can also use the compiled methods from outside the
object:
OP.fft4(Array( 1.0,0.0, 1.0,0.0, 2.0,0.0, 2.0,0.0))
Array( 6.0, 0.0, − 1.0, 1.0, 0.0, 0.0, − 1.0, − 1.0)
Providing an explicit type in the definition val OP: TestFFC =
… ensures that the internal representation is not accessible
to the outside, but only the members defined by TestFFC
are accessible.
3. ADDinG MoRe feATuRes
Many features can be added in a way that is analogous
to what we have seen above, but some require a bit more
thought. In this section, we take a closer look at staged
functions. Basic support for staged function definitions
and function applications can be defined in terms of a
simple higher order abstract syntax (HOAS)
16 representation, similar to those of Carette et al.
2 and Hofer et al.
8 (see
Figure 11). The idea is to provide a lambda operation that
transforms present-stage functions over staged values
(type Rep[A] => Rep[B]) to staged function values (type
Rep[A=>B]). Note how this is similar to the signature of
compile. To give an example, the staged recursive factorial function will look like this:
def fac: Rep[Int => Int] = lambda {n =>
if (n == 0) 1
else n fac(n − 1)
}