Add Graal JIT Compilation to Your JVM Language in 5 Easy Steps, Step 3

Step 3: Interpreting a Simple Fibonacci Function with Golo+Truffle

First, a word of warning. This post is much longer than I had hoped. The goal for today is to be able to execute a Fibonacci function with Golo+Truffle. However, this means, we need to get acquainted with most of the basic Truffle concepts. And in the Golo implementation that corresponds to a patch that is about 1500 LOC large. So please bear with me, this gets a little long-winded.

In case you missed the previous posts, this series starts with a brief intro and motivation. Today's post discusses the basics of self-optimizing interpreters using the Truffle framework. Some of the discussions are also already including details on how the compilation with Graal works. However, the main focus is on the basic interpreter functionality.

This post discusses how to create a Truffle abstract syntax tree (AST) from the Golo intermediate representation (IR), explains the basic mechanism of the Truffle framework, e.g., execute() methods, the TruffleDSL that generates much of the boilerplate code, self-optimization and node specialization, basic interpreter functionality such as function arguments, function invocation, and control flow constructs. So, just enough functionality to get a basic Fibonacci function working. At the end of the post, we also provide a brief overview of the most important bits.

Adding Basic Truffle Support

The first step is to configure the Gradle build system to include the Truffle dependencies. As of this writing, the latest released of Truffle and Graal is version 0.9 (binary builds).

We need the basic Truffle libraries as well as the TruffleDSL, which makes it easier to implement the basic operations for our interpreter. The build dependencies are configured in the build.gradle file. We need to add the Truffle jar as a compile dependency. Furthermore, we need a build plugin to enable annotation processing support for the TruffleDSL. The DSL uses annotations to generate implementation classes for us. Since the DSL itself also depends on Truffle, the Truffle jar needs to be registered as an apt dependency as well. Note, at the moment we also need to use a custom Maven repository with the latest Truffle jars:

repositories {
  maven {
    url "http://lafo.ssw.uni-linz.ac.at/nexus/content/repositories/releases/"
  }
  jcenter()
}

plugins {
  // ...
  id "net.ltgt.apt" version "0.3"
}

dependencies {
  // ...
  compile 'com.oracle.truffle:truffle-api:0.9'
  apt     'com.oracle.truffle:truffle-api:0.9'
  apt     'com.oracle.truffle:truffle-dsl-processor:0.9'

  // ...
}

Working in the Debugger

The next step might seem superfluous, but it is going to be extremely useful. The first class we add to the Golo codebase is a NotYetImplemented exception:

class NotYetImplemented extends RuntimeException {}

With this class, we can mark all open issues/todo items and make the compilers happy, which allows us to do all our changes in the debugger. Why in the debugger? Because, I don't want to guess what Golo is doing. I don't know Golo. So, I just put a breakpoint where I assume I'll need to change something, and then start hacking there. The debugger then can just tell me what arguments I got, their values, etc. Trivial, but very convenient in an unknown any codebase.

Goal: Convert Golo's IR to a Truffle AST for to run Fibonacci

The next step is to add an option to Golo's start mechanism to convert the Golo IR not into bytecodes but to give us a Truffle AST that can be executed. For that, we add a --truffle option to the golo command. If we set it on the command line, Golo will parse the files as normal and then use a visitor to convert the resulting GoloModule into a set of Truffle-based functions/ASTs (cf. change set).

The resulting GoloModules represent Golo code in an IR that is already pretty high-level, and it would require only small changes to be compatible with Truffle. However, for clarity and to avoid artifacts that are not idiomatic in Truffle, I chose to keep the IR and the Truffle AST clearly separated. In a practical system, it might not be necessary.

As a result of separating the Truffle trees from the IR, we only need to add accept(.) methods to the IR. They perform a double dispatch to the corresponding method in the TruffleGenerationGoloIrVisitor. In a first step, the visitor could be left almost empty and just needs visitor methods for all IR elements, which in turn can raise the NotYetImplemented exception. The full visitor implementation was added in my branch with this change.

With this basic infrastructure in place, we can run Golo from the debugger executing the following program: golo --truffle --files samples/fib.golo

module samples.Fibonacci

function fib = |n| {
  if n < 2 {
    return n
  } else {
    return fib(n - 1) + fib(n - 2)
  }
}

function main = |args| {
  println(fib(10))
}

This is a simple recursive version of Fibonacci. To run it with Truffle, we need to implement function calls, function arguments, control flow constructs, and basic arithmetic operations. Sounds simple, but it will keep us busy for the rest of this post.

By now, the debugger should have hit the first NotYetImplemented exception in the visitor method for modules. Ignoring some Golo details, it is sufficient to iterate over all GoloFunctions and visit them one by one. The corresponding visitFunction(.) eventually needs to assemble a Truffle-compatible function. But, following the relevant execution flow in the debugger, we first transform the fib function's block, i.e., its body to a Truffle AST and worry about the details of creating a function object later.

Truffle Basics: Nodes, execute*() Methods, TypeSystems, and Frames

The first element we need is a common root node for all possible AST nodes. In Golo's IR, the corresponding class is essentially GoloStatement. For our purposes we add ExpressionNode,11 We use the term expression here, because most nodes will need to provide a return value for the execute*() methods. which needs to extend Truffle's Node class. This ExpressionNode needs to provide a set of execute*(frame) methods. The execute*() methods implement the semantics of an AST node. To be able to specialize the behavior of a node on the kind of value it should return, the ExpressionNode provides a set of basic implementations for the types we want to be able to specialize on as return types for expressions. These types are communicated explicitly to the TruffleDSL to generate some type-checking helper methods. We do this with the following Types class:

@TypeSystem({
  int.class,
  boolean.class,
  Object[].class
})
class Types { }

With the @TypeSystem annotation, we list int, boolean, and Object[] as the three types of interest. Note, the object array is going to be used only for arguments to function calls. Because, our Fibonacci program only needs integers and booleans as supported data types.

For the ExpressionNode class, we need the following code:

public abstract class ExpressionNode extends Node {
  public abstract Object executeGeneric(VirtualFrame frame);

  public int executeInteger(VirtualFrame frame) throws UnexpectedResultException {
    return TypesGen.expectInteger(executeGeneric(frame));
  }

  public boolean executeBoolean(VirtualFrame frame) throws UnexpectedResultException {
    return TypesGen.expectBoolean(executeGeneric(frame));
  }

  public Object[] executeObjectArray(VirtualFrame frame) throws UnexpectedResultException {
    return TypesGen.expectObjectArray(executeGeneric(frame));
  }
}

The abstract executeGeneric(frame) method on the second line is the generic execute method each AST node needs to implement. The other three type-specific methods can be implemented if it makes sense for a node. Otherwise, the fallback implementation provided here will call executeGeneric(.), and then check that the result is of the expected type. For the check, we use the expect*() methods, which are generated from the Types class. Technically it is optional, but avoids some boilerplate code.

Generally, these fallback methods are useful for optimistic, i.e., speculative optimizations, which might not hold for all future executions of a node. We'll go into more details on this later.

One might also notice the frame argument. The frame is the activation record of a Truffle function and is used to access the actual arguments to a function call, local variables, and other state a language implementation might need in the scope of a function. Later we will see the argument access. For the other aspects, please see a later post in this series.

Sequences

With the ExpressionNode in place, we can now transform the body of the GoloFunction, which is called a block and corresponds to a distinct lexical scope in Golo. To represent the content of such blocks, which can be multiple expressions, we use the following SequenceNode that holds the expressions as child nodes:

class SequenceNode extends ExpressionNode {
  @Children private final ExpressionNode[] expressions;

  public SequenceNode(ExpressionNode[] expressions) {
    this.expressions = expressions;
  }

  @ExplodeLoop
  public Object executeGeneric(VirtualFrame frame) {
    for (int i = 0; i < expressions.length - 1; i++) {
      expressions[i].executeGeneric(frame);
    }
    return expressions[expressions.length - 1].executeGeneric(frame);
  }
}

For the Truffle framework, it is important that the expressions field of the class is annotated with @Children so that specializations are handled correctly and so that the Graal compiler knows these expressions can be considered constant. The final keyword is here not strictly necessary, but in general it is good practice to make fields final whenever possible to facilitate constant folding during compilation.22 Fields annotated with @Child cannot be final, because they need to be able to specialize themselves, and the compiler treats them specially.

The executeGeneric(.) method iterates over these child expressions and recursively calls executeGeneric(.). The last expression is different, because we need to return its result value.

To enable the Graal compiler later on to optimize this code as much as possible, it is necessary to mark this loop with the @ExplodeLoop annotation to ensure it is unrolled. This makes sure that the call to executeGeneric(.) is not treated as polymorphic. By unrolling the loop, the compiler can inline the specific executeGeneric(.) methods for each expression[i]. This enables many common compiler optimizations. The unrolling is safe here, because the number of statements in a block is lexically defined, and thus constant. Depending on language semantics, this might for instance also apply to the processing of function arguments, as we will see later.

Note, while stepping through the debugger, the block we encountered should be the one that contains the if ... {} else {} construct of the fib function, which we dive into now.

Argument Reads

Continuing from the visitBlock(.) function in the debugger, transforming our fib function step-by-step, the next element is going to be the conditional branch. It is transformed to Truffle nodes in the visitConditionalBranching(.) method. As a first step, we visit the condition n < 2. When stepping into it, the left operand is a argument read of n.

As mentioned earlier, function arguments are passed as Object[] arrays in the frame, i.e., activation record of the function. To access them, we add the LocalArgumentReadNode class, which is sketched below. Such a node keeps the index of the argument it corresponds to, and on execution, simply accesses the array from the frame object and returns the value:

class LocalArgumentReadNode extends ExpressionNode {
  protected final int index;

  public LocalArgumentReadNode(int index) {
    this.index = index;
  }

  public Object executeGeneric(VirtualFrame frame) {
    return frame.getArguments()[index];
  }
}

Literals

The right operand is a literal integer 2. Literal nodes are the simplest nodes we have, because it simply returns the stored value. For integer values, we add the IntegerLiteralNode below. Note, for the compiler later on, it is important that the val field, in which the value is stored, is final. This ensures that the compiler can treat the value as constant to enable optimizations.

class IntegerLiteralNode extends LiteralNode {
  private final int val;

  public IntegerLiteralNode(int val) {
    this.val = val;
  }

  public int executeInteger(VirtualFrame frame) {
    return val;
  }

  public Object executeGeneric(VirtualFrame frame) {
    return val;
  }
}

TruffleDSL, Basic Operators, and Returning from a Function

So far, we needed only very basic mechanisms of Truffle and mostly followed the framework by implementing abstract methods. Now, we are moving on to use the TruffleDSL to provide functionality that is, for instance, specific to certain types that we want to optimize for. The TruffleDSL facilitates this by providing annotations that are used to generate much of the necessary boilerplate code.

The DSL provides the notions of nodes and their children, i.e., subexpressions. Furthermore, it enables us to implement specializations based on types as well as other runtime information. The previously discussed execute*() methods are essentially the interface to subnodes. Using this interface, the DSL generates the self-specializing code to instantiate node objects, provide the specialization checks, evaluate subnodes, do node rewriting, and handle potential polymorphism, i.e., the case that more than one specialization is required for a specific point in the program (AST node).

In the following, we discuss the basic concepts to implement binary operations in the Truffle style.

Node Children and Specializations

Back to the Fibonacci function, the comparison operation for our n < 2 expression is implemented as LessThanNode. For convenience, we extend the BinaryNode class below:

@NodeChildren({
  @NodeChild(value = "left",  type = ExpressionNode.class),
  @NodeChild(value = "right", type = ExpressionNode.class)
})
public abstract class BinaryNode extends ExpressionNode { }

The BinaryNode is an abstract class without any methods. Its main use is to provide the left and right subnodes, i.e., node children by using the TruffleDSL. As the code above shows, we use the @NodeChildren annotation with two @NodeChild annotations that define left and right to be of type ExpressionNode. The main benefit of using the DSL is that it generates the boilerplate code for specializations and handling many of the common issues for optimistic optimizations. In our specific case here, it enables use to implement the LessThanNode with the following code:

public abstract class LessThanNode extends BinaryNode {

  @Specialization
  public boolean doIntegers(int left, int right) {
    return left < right;
  }
}

Note that LessThanNode is again an abstract class. However, it provides the doIntegers(.,.) method that is annotated with @Specialization. This means, the DSL will generate the code to evaluate the left and right subexpressions of the < operation, which in our case are an argument read and an integer literal. If the evaluation returns the expected integers, it will make sure that the AST is rewritten to use the doIntegers method. Otherwise, it would throw an UnsupportedSpecializationException at runtime. For our Fibonacci support, this behavior is good enough. For a complete Golo implementation, we would of course need to provide the other necessary specializations as well.

At this point, note that the signature of doIntegers(.,.) as well as the return type of executeInteger(.) uses the primitive int and thus avoid boxing of integers.

Control Flow: The return Keyword

After completing the condition, we are back in the visitConditionalBranching(.) method with the debugger. The next element, for which we need to provide an implementation is the then-branch of the condition. Stepping into it, we reach the visitReturnStatement(.) method. In our fib function, we know this is the return n branch. So, there will be an argument read again, which we already covered. Thus, we can focus on the return itself.

Imagine for a moment what happens when we execute the kind of AST that we are just building up. The Java stack trace might look something like this:

ReturnNode.executeGeneric()
Sequence.executeGeneric()
IfNode.executeGeneric()
Sequence.executeGeneric()
Function.execute()

At the bottom of the stack, we see the function call itself, and then we see the different AST nodes with their executeGeneric() methods executing. At the top, we see the node that corresponds to the return operation we want to implement. So, how can we return from this node back to the function and at the same time also pass the desired return value? In Truffle, the solution is rather straightforward and uses ControlFlowExceptions. The class ControlFlowException minimizes the cost of an exception at runtime, for instance, by not capturing a stack trace. To use it, we subclass it with the following ReturnException:

class ReturnException extends ControlFlowException {
  private final Object result;

  public ReturnException(Object result) {
    this.result = result;
  }

  public Object getResult() {
    return result;
  }
}

This ReturnException takes the result we want to return from the function, and when it is thrown, it will unwind the Java stack for us. Thus, we only need to make sure that we catch it later on in the Function class that we still need to implement. More on that later.

In addition to the ReturnException, we also need the corresponding AST node:

class ReturnNode extends ExpressionNode {
  @Child protected ExpressionNode expr;

  public ReturnNode(ExpressionNode expr) {
    this.expr = expr;
  }

  public Object executeGeneric(VirtualFrame frame) {
    throw new ReturnException(expr.executeGeneric(frame));
  }
}

ReturnNode evaluates its child node expr, which is marked with @Child. It uses the resulting value to create a new ReturnException, which is thrown to unwind the stack and return from the function. So, in case of the return n expression in the fib function, this means, we read the argument, wrap the result into an exception object that is than thrown. And that's basically it. To see how the exception is caught, see section 6.5 on implementing the Function class. Note, we did not mark the expr field as final. Since the expression node is supposed to specialize itself during execution, it needs to mutate this field. For the compilation, the @Child annotation is sufficient to consider the field as a constant for optimization.

Addition and Subtraction Operators

Moving on in the debugger should get us back to the visitConditionalBranching(.) method, which we continue to transform by processing the else branch: return fib(n - 1) + fib(n - 2). Handling the return was just covered. So, the next elements of interest are the binary operators + and -. They are implemented essentially in the same way as the less-than operator <. We use the TruffleDSL and its support for specializations:

public abstract class PlusNode extends BinaryNode {
  @Specialization
  public int doIntegers(int left, int right) {
    return left + right;
  }
}

public abstract class MinusNode extends BinaryNode {
  @Specialization
  public int doIntegers(int left, int right) {
    return left - right;
  }
}

Both nodes are implemented as subclasses of BinaryNode and provide the specialized behavior for adding or subtracting integers. With this minimalism, the same caveats apply as for the LessThanNode, i.e., for a full Golo implementation, we would need to provide support for other possible types as well.

Functions and Conditional Control Flow

So far, we merely dabbled with the very basics of language implementation. In this section, we finally make the step to add function calls and conditional control flow. Note, this time, we will not rely on the DSL. Instead, we are going to do some manual node specialization. I hope this provides some additional insights on how these types of self-optimizing interpreters work.

Function Arguments

The generation of the fib(.) function invocations are handled by visitFunctionInvocation(.). As the first step, it needs to transform the argument expression of the function call. Ignoring Golo's named arguments, we can transform the argument expressions into a plain ExpressionNode[]. To avoid having different function invocation nodes for different numbers of arguments, we introduce an EvalArgumentsNode that will create an array of objects that is passed to the function call. That's also the reason why we needed the Object[] type in our Types class earlier.

class EvalArgumentsNode extends ExpressionNode {
  @Children protected final ExpressionNode[] argumentNodes;

  public EvalArgumentsNode(ExpressionNode[] argumentNodes) {
    this.argumentNodes = argumentNodes;
  }

  @ExplodeLoop
  public Object[] executeObjectArray(VirtualFrame frame) {
    Object[] arguments = new Object[argumentNodes.length];
    for (int i = 0; i < argumentNodes.length; i++) {
      arguments[i] = argumentNodes[i].executeGeneric(frame);
    }
    return arguments;
  }

  public Object executeGeneric(VirtualFrame frame) {
    return executeObjectArray(frame);
  }
}

The implementation of EvalArgumentsNode is very similar to the one of SequenceNode. The main difference is that we create an Object[] array that captures the actual argument values for the call. Note again that the argumentNodes field is annotated with @Children and the executeObjectArray(frame) method is annotated with @ExplodeLoop to help the optimizers.

Function Invocation in Self-Optimizing Interpreters

Once we got the EvalArgumentsNode object, we can create the actual FunctionInvocationNode. This one is quite a bit more complex than any of the nodes we had before. Before we start looking at it, note that we care for the moment only about simple function invocation to targets that can be resolved statically. Thus, we do not support any form of polymorphism yet. This also means that we are going to implement the node manually. The TruffleDSL provides with the @Cached mechanism a nice way of implementing custom polymorphic inline caches. But, here this approach would be more complicated without providing a benefit.

Before we go into the details, let's back up even more. One of the main assumptions for self-optimizing interpreters in the Truffle style is that self-optimizations are designed so that they always have converging, i.e., stabilizing behavior. To be able to generate native code, we need to avoid getting into an endless loop of back and forth switching between different optimizations. Thus, one needs to design node specializations as a state machine that will reach a final state within a finite and small number of steps.

With this in mind, let's have a look at what we need to do for the function invocation. As already mentioned, we do not support polymorphic functions. And for our Fibonacci example, we have two cases. The first one is the call to the fib(.) function itself, which we just hit in the debugger. And the second one is the call to the println(.) function, which we need later to generate output on the terminal. The fib(.) function is going to be resolved statically, and println(.) is a method on Golo's Predefined class, which provides such convenience methods. We will handle println(.) specially, just to show a little more what we can do with Truffle.33 In a full Golo+Truffle backend, one probably wants a more generic mechanism. But, we are going to look at that only in the next part of this series of blog posts.

So, overall, we got three cases to cover. The initial function invocation node that is uninitialized, a normal static function call, and well-know builtins that we want to treat specially. This means, we can design a node that goes from an uninitialized case to either of the two specializations depending on the result of a lookup function.

We'll start with the abstract FunctionInvocationNode:

abstract class FunctionInvocationNode extends ExpressionNode implements PreEvaluated {
  protected final String name;
  protected final GoloModule module;

  @Child protected EvalArgumentsNode argumentsNode;

  FunctionInvocationNode(String name, GoloModule module, EvalArgumentsNode argumentsNode) {
    this.name          = name;
    this.module        = module;
    this.argumentsNode = argumentsNode;
  }

  public abstract Object executeEvaluated(VirtualFrame frame, Object[] args);

  public Object doEvaluated(VirtualFrame frame, Object[] args) {
    return executeEvaluated(frame, args);
  }

  public Object executeGeneric(VirtualFrame frame) {
    Object[] args = argumentsNode.executeObjectArray(frame);
    return executeEvaluated(frame, args);
  }
}

It is a subclass of ExpressionNode and as such we implement executeGeneric(frame). Furthermore, we implement argument evaluation here, since it is always the same independently of how an invocation is done. Thus, we take the argumentsNode and ask it to evaluate to an Object[], which it will by construction (see the just implemented EvalArgumentsNode).

With the result, we call executeEvaluated(frame, args). This method is a convention that is also used by the DSL. The idea is that a node will first evaluate all child nodes, and then at some point execute the node's own logic. In this case, we only provide an abstract method that is implemented in the concrete subclasses. Note also that we implemented an interface PreEvaluated, which gives us a way to always execute only the node-specific behavior without evaluating child nodes. This is useful for instance during specializing a node. First we need to evaluate the child nodes to be able to determine the applicable specialization, and then we want to execute only the node-specific behavior. Since Golo has side-effects, just repeating the evaluation of the child nodes could be incorrect.

For the specialization of function invocations, the code looks as follows:

class UninitializedInvocationNode extends FunctionInvocationNode {

  UninitializedInvocationNode(String name, GoloModule module,
      EvalArgumentsNode argumentsNode) {
    super(name, module, argumentsNode);
  }

  public Object executeEvaluated(VirtualFrame frame, Object[] args) {
    return specialize(args).
        doEvaluated(frame, args);
  }

  private PreEvaluated specialize(Object[] args) {
    Object lookupResult = lookup(args);

   if (lookupResult instanceof Function) {
      return replace(new DirectFunctionInvokeNode(this, (Function) lookupResult));
    } else if (lookupResult instanceof PreEvaluated) {
      return (PreEvaluated) replace((ExpressionNode) lookupResult);
    }
    throw new NotYetImplemented();
  }

  private Object lookup(Object[] args) {
    // ...
  }
}

Our UninitializedInvocationNode class implements the abstract executeEvaluated method. As discussed, it already gets the evaluated arguments, passes them to the specialize(.) function, and then delegates the evaluation to the resulting node by calling doEvaluated(.,.). The specialize(.) function performs the lookup, and then based on the lookup result, it chooses an appropriate implementation for the invocation. At this point, we only need the two previously discussed ways for invocation. The first one is the DirectFunctionInvokeNode. In case the lookup result was a Golo+Truffle function, we instantiate the invocation node and replace(.) the UninitializedInvocationNode node, i.e., the current this with it.

Just as a side note, the direct use of replace(.) is only advisable for advanced use cases and requires extra care. Normally, the DSL shields us from this operation. replace(.) is going to adapt the AST and makes sure that the parent pointer in the newly inserted node is set correctly. It also performs a few checks to make sure the replacement is legal. If the parent node's field type is not compatible with the new node, we will get an error. While there are such safeguards, we have to manually make sure that a node is only used once in the AST. This is required because the nodes can only have a single parent, which is required to make the specialization logic work.

In case the lookup result is some node that implements the PreEvaluated interface, we use it immediately to replace the current node.

Invoking Truffle Functions

The DirectFunctionInvokeNode only has to perform the actual function invocation. As discussed earlier, here we only handle cases that are statically resolved. Thus, there is no further dynamicity that needs to be handled. The only special thing we have to do is to use Truffle's mechanism to function invocation so that the Truffle framework can apply optimizations such as splitting and inlining. To this end, the following implementation uses a DirectCallNode:

class DirectFunctionInvokeNode extends FunctionInvocationNode {
  @Child protected DirectCallNode call;

  protected DirectFunctionInvokeNode(
      FunctionInvocationNode uninit, Function function) {
    super(uninit.name, uninit.module, uninit.argumentsNode);
    call = Truffle.getRuntime().createDirectCallNode(
                function.getCallTarget());
  }

  public Object executeEvaluated(VirtualFrame frame, Object[] args) {
    return call.call(frame, args);
  }
}

The DirectCallNode caches the so-called CallTarget. CallTarget's are Truffle's notion of invokable or callable entities. Thus, they are the executable elements, and define the calling convention. Without going into the details of the Function implementation already, our Truffle Functions will cache such a CallTarget that is used in DirectCallNodes. These nodes are created, as shown here, by the Truffle runtime. They represent statically bound calls, and communicate this knowledge to the compiler.44 They can also be used, for instance in combination with the @Cached annotation to build polymorphic inline caches. When a call-site is highly variable (megamorphic), i.e., corresponds to many different functions, we could use the IndirectCallNode instead. In that case, the optimizer would not perform inlining, because it would most likely slow down execution.

With this support for function invocation, we can transform the remainder of the fib() function. Stepping through the debugger should eventually lead us back to visitReturnStatement() where the ReturnNode for the else branch is created. After this is done, we are finally back at visitBlock(.), where the SequenceNode for the block of the else branch is created.

Control Flow: Conditionals

With the SequenceNode for the else branch of the conditional, we finally got all elements needed to assemble the conditional itself. It is implemented as the following IfNode class:

class IfNode extends ExpressionNode {
  @Child ExpressionNode condition;
  @Child ExpressionNode thenNode;
  @Child ExpressionNode elseNode;

  public IfNode(ExpressionNode condition,
      ExpressionNode thenNode, ExpressionNode elseNode) {
    this.condition = condition;
    this.thenNode  = thenNode;
    this.elseNode  = elseNode;
  }

  @Override
  public Object executeGeneric(VirtualFrame frame) {
    if (doCondition(frame)) {
      thenNode.executeGeneric(frame);
    } else if (elseNode != null) {
      elseNode.executeGeneric(frame);
    }
    return null;
  }

  private boolean doCondition(VirtualFrame frame) {
    try {
      return condition.executeBoolean(frame);
    } catch (UnexpectedResultException e) {
      CompilerDirectives.transferToInterpreter();
      throw new UnsupportedSpecializationException(
          this, new Node[]{condition}, e.getResult());
    }
  }
}

As one might expect, the IfNode got three @Child nodes. The condition, the thenNode, and the elseNode. In executeGeneric(.), we first evaluate the condition, and depending on the result, either the thenNode or elseNode. Note, the elseNode is optional. The check here is also not an issue for performance, because the compiler knows from the @Child annotation that it can consider the value of elseNode as a constant. Thus, the null check will be removed during compilation.

Evaluating the condition is actually a little more complex.55 I refrained from putting in a ConditionProfile to avoid additional complexity. But usually it is a good idea, e.g., to enable elimination of unused branches. The doCondition(.) method needs to handle the case that the condition does not actually return a boolean value. With Golo being a dynamic language, this could of course happen at runtime. For our purposes, we do not provide a fallback implementation, but instead just raise an error. Hence, if executeBoolean(.) fails, we catch an UnexpectedResultException and simply throw an UnsupportedSpecializationException.

Note, the first operation in the catch clause is CompilerDirectives.transferToInterpreter(). This operation makes sure that the compiler does not include the catch clause in the compilation. We assume, well behaved programs won't use non-boolean conditions and thereby avoid the complexity of handling this special case during compilation. Instead, execution will deoptimize. Thus, it leaves the native code Graal produced and returns to this program point, i.e., to the interpreter, which is slower, but does not prevent many essential optimizations to happen in the compiler.

Truffle Functions aka RootNodes

With the IfNode constructed, we finally got all elements for our fib(.) function transformed to a Truffle representation. Now, we can create the actual Function object:

class Function extends RootNode {
  @Child protected ExpressionNode expr;

  public Function(ExpressionNode expr) {
    this.expr = expr;
    Truffle.getRuntime().createCallTarget(this);
  }

  @Override
  public Object execute(VirtualFrame frame) {
    try {
      return expr.executeGeneric(frame);
    } catch (ReturnException ex) {
      return ex.getResult();
    }
  }
}

For the moment, we need to cover three aspects here. The first is that we need to subclass Truffle's RootNode class, implement its execute(.) method, and handle the ReturnException from earlier. RootNode is the framework's class to represent the main node of a Truffle function. Its execute(.) method takes only a frame object as argument and is the main entry point for evaluating Truffle ASTs. In the case of Golo, a Function is simply referring to some ExpressionNode that we ask to execute. As you can see in the execute(.), this is also the right place to handle the ReturnException. The Function is the inner part of the call boundary. Hence, we catch the exception, ask it for the result and can then return it. In case there was no exception, we simply return the value to which the expr node evaluated.

Previously, we also discussed the CallTarget as Truffle's mechanism to optimize function calls. Since we need a CallTarget for each of our Truffle functions, we create it in the constructor and it is automatically set on this, so that it can be used later.

An Intrinsic Node for Printing

With that, we are almost set for executing the Fibonacci program. It remains only the main() function itself. When discussing function invocation, I mentioned that we are going to treat println(.) special to show a little more of how Truffle works.

As a basic implementation of the printing, we use the following node class:

@NodeChild(value = "expr", type = ExpressionNode.class)
abstract class UnaryNode extends ExpressionNode { }

abstract class PrintlnNode extends UnaryNode implements PreEvaluated {
  public abstract Object executeEvaluated(VirtualFrame frame, Object value);

  @Override
  public Object doEvaluated(VirtualFrame frame, Object[] args) {
    return executeEvaluated(frame, args[0]);
  }

  @Specialization
  public Object printObject(Object obj) {
    System.out.println(obj);
    return null;
  }
}

First, we define UnaryNode, very similar to the previously defined BinaryNode. It has one node child expr and otherwise is an empty abstract class. PrintlnNode extends UnaryNode and implements again the PreEvaluated interface so that we can use it in the specialize(.) method of UninitializedInvocationNode. We also have the abstract executeEvaluated(.), which the DSL is going to generate for us. The actual implementation is done in the printObject(obj) method, which is annotated with @Specialization. Based on this specialization, the TruffleDSL generates the missing implementation methods, so that we can use the node in the AST.

The question is now, how do we actually create the node. Below we see a snippet from the corresponding Golo lookup routine:

if (importClass == Predefined.class) {
  switch (lookup) {
    case "println":
      return PrintlnNodeGen.create(argumentsNode.getArgumentNodes()[0]);
  }
}

Since the Predefined class is completely under our control, we can provide this node implementation for println(.). In the lookup routine, we test whether the class on which the lookup is done is the Predefined class, and in case it is, and the desired function is "println", we instantiate a PrintlnNode. Note here that the TruffleDSL generated the implementation with the name PrintlnNodeGen. This class has a factory method create(expr) that takes the argument expression. In our case, we execute this code as part of the UninitializedInvocationNode that has the argument expressions stored in argumentsNode. So, we can take the corresponding expression from there. Once the lookup is done, the created node object is passed to the specialize(.), which will eventually replace the invocation node with our new PrintlnNode. This means, the print operation is directly inlined in the AST.

Wrap Up and Conclusions

In this post, many basic Truffle concepts have been introduced. From the basic AST nodes, over the design of an AST-based interpreter, arithmetic operations, control flow constructs, until function calls, we covered everything we needed to execute a simple Fibonacci function. For more on the basic mechanisms, the paper on Self-Optimizing AST Interpreters is certainly worth a read.

The code for this post is available on the truffle/fib branch. And with the following steps, it should be possible to execute the fib.golo program:

./gradlew installDist
build/install/golo/bin/golo golo --truffle --files samples/fib.golo

If everything works out, it should print 55 and exit cleanly.

In the next post, we will see the remaining features that are needed to get our Mandelbrot program running. But before we close for today, I wanted to briefly point at the differences between the bytecode generation in the last post, and how Truffle interpreters work. From my perspective, the two approaches require quite a bit of a different mindset. When generating bytecodes, one has to be very aware of how the Java Virtual Machines works. Specifically, we have to constantly track in our mind how it interacts with the stack, what the effect is of the bytecode that is eventually executing, and so on. To truly understand it, we need to be in a mindset for compilation. Thus, we need to imagine how the JVM executes and then generate instructions for it so that it does what we want.

With a Truffle-style interpreter on the other hand, we do only need to imagine what we want to do, and then can codify it in normal Java. We do not have to imagine an external machine that is doing something on our behalves, instead we just need to imagine what needs to be done for a specific piece of input program. Of course, this only goes so far. With the Graal compiler in the mix, we actually need a pretty good understanding of how the optimizers work, but at least we do not need to mess with a stack machine.

Truffle Concepts, Conventions, and Other Important Bits

  • @TypeSystem (sec. 4), optional, generates helper methods for type tests
  • execute*() methods (sec. 4) implement semantics for an expected return type
  • executeEvaluated(...) methods (sec. 6.2) are a convention used by Truffle to implement a node's semantics taking all subexpressions already fully evaluated.
  • @Specialization, (sec. 5.1) annotation of the DSL to generate the node for this functionality, i.e., optimized variant of a generic operation
  • @Children and @Child (sec. 4.1) mark subnodes to allow framework to do correct node replacement in AST. They are implicitly compilation final.
  • @NodeChildren, @NodeChild (sec. 5.1), annotation of the DSL to add subnodes in generate code
  • abstract get*() methods, for instance getLeft() and getRight() DSL generates implementations to be able to reference a subnode that is managed by DSL
  • @ExplodeLoop (sec. 4.1), annotation to ask for unrolling of loop during compilation
  • @CompilationFinal, annotation to mark a field as final for the Graal compiler even so it cannot be marked with the keyword final, for instance because of lazy initialization
  • ControlFlowException (sec. 5.2), the root class for all exceptions that are expressing control flow in a Truffle AST, optimized to avoid runtime overhead
  • nodes can be only part of an AST once, because they only have one parent

Acknowledgements

I'd like to thank Julien Ponge for his comments on drafts of this tutorial and his help with the technical details of Golo. Furthermore, I'd like to thank Andreas Wöß for his comments and clarifications of Truffle-related details.

Creative Commons License
“Add Graal JIT Compilation to Your JVM Language in 5 Easy Steps” by Stefan Marr is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.
Permissions beyond the scope of this license are available on request.

Tags: , , , , , ,

Comments are closed.