We investigate how to build sophisticated program generators like Spiral using modern language features.

What is it?

We reimplemented a subset of the original Spiral system using modern program language features. In the process we make heavy use of embedded DSLs in Scala and staging utilizing the Lightweight Modular Staging framework to achieve to following goals:

  • Describing problem and algorithmic knowledge through multiple levels of DSLs.
  • Specifying optimizations and algorithmic choices as rewrite rules on DSL programs.
  • Build a infrastructure to easily abstract over recurring low-level transformations such as unrolling and scalarization; this is done by abstracting over staging decisions.
  • Designing high-level data structures that can be parametrized to generate multiple low-level representations including different data layouts.

Technical Highlights

Multi-stage programming inherently fits nicely in the context of program generators. In this work we go one step further and explore the possibilities that staging done through types, as implemented in Lightweight Modular Staging, combined with modern program language constructs such as Type Parameterization and Type Classes offers. We give a teaser on some of the powerful constructs that can be build through this combination in the following sections. More details are in [1,2].

Selective Staging

In the case of program generators its beneficial to selectively decide if code should be executed at compile time, or if it should be part of the final program. To avoid code duplication we can encode both behaviors by abstraction over the staging decision. Since LMS uses types to annotate staging, we can utilize parametric polymorphism to achieve this as depicted below:

def somefunction(x0: T) : T = {
  val x1: T = x0 * 2.0
  x1
}

Where type T can now be either Double or Rep[Double] (the staged type). This allows us to write code that is universal for precomputation and as part of the final program.

Different Code Styles

We observe that moving the staging annotation in the code changes the code style:

Meta Program (Scala) Resulting Intermediate Representation Unparsed Result (C) 
type T = Array[Rep[Double]]
def f3(x: T, y: T) = {
  for (i <- 0 until 2) {
    y(2*i) = x(i*2) + y(i*2+1)
    y(2*i+1) = x(i*2) - y(i*2+1)
  }
}
 scalar code 
t0 = s0 + s1;
t1 = s0 - s1;
t2 = s2 + s3;
t2 = s2 - s3;
 
type T = Rep[Array[Double]]
def f3(x: T, y: T) = {
  for (i <- 0 until 2) {
    y(2*i)   = x(i*2) + y(i*2+1)
    y(2*i+1) = x(i*2) - y(i* 2+1)
  }
}
 scalar code 
t0 = x[0];
t1 = x[1];
t2 = t0 + t1;
y[0] = t2;
t3 = t0 - t1;
y[1] = t3;
t4 = x[0];
t5 = x[1];
t6 = t4 + x5;
y[0] = t6;
t7 = t4 - x5;
y[3] = t7;
 
type T = Rep[Array[Double]]
def f3(x: T, y: T) = {
  for (i <- (0 until 2)): Rep[Range] {
    y(2*i)   = x(i*2) + y(i*2+1)
    y(2*i+1) = x(i*2) - y(i*2+1)
  }
}
 scalar code 
int i;
for (i=0;i<2;i++){
  t0 = x[i];
  t1 = x[i+1];
  t2 = t0 + t1;
  y[i] = t2;
  t3 = t0 - t1;
  y[i+1] = t3;
}

Abstraction over Code Style

We then extend the idea of selective staging to also abstract over all these code styles.

We modify the function from above to the form:

type NoRep[T] = T
def f3[L[_],A[_],T](looptype: L, x: A[Array[T]], y: A[Array[T]])= {
  for ( i <- (0 until 2)): L[Range] {
    y(2 * i) = x(i * 2) + y(i * 2 + 1)
    y(2 * i + 1) = x(i * 2) - y(i * 2 + 1)
}

This allows us to instantiate all of the above code styles from a single function.

Abstraction over Data Layout

We hide the type information into an abstract base class. We extend the base class with different types of instantiations of the typing such that we realize code such as:

val x: Looped_SplitComplex<
val y: Scalar_AVXUnrolled
y = x // unrolling with scalarization, data format change and vectorization

The actual implementation is quite involved and is shown in full detail in [1,2].

Source Code

The current version of SpiralS can be found at GeorgOfenbeck / SpiralS.

Main References

  1. Georg Ofenbeck, Tiark Rompf and Markus Püschel
    Staging for Generic Programming in Space and Time
    Proc. International Conference on Generative Programming: Concepts & Experiences (GPCE), pp. 15-28, 2017

  2. Georg Ofenbeck, Tiark Rompf, Alen Stojanov, Martin Odersky and Markus Püschel
    Spiral in Scala: Towards the Systematic Construction of Generators for Performance Libraries
    Proc. International Conference on Generative Programming: Concepts Experiences (GPCE), pp. 125-134, 2013

  3. Georg Ofenbeck, Tiark Rompf, Markus Püschel
    RandIR: Differential Testing for Embedded Compilers
    Proc. Scala Symposium, pp. 21-30, 2016

Auxiliary References