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)

}