2. Building a JIT: Adding Optimizations – An introduction to ORC Layers¶

Chapter 2 Introduction
Optimizing Modules using the IRTransformLayer
Full Code Listing

This tutorial is under active development. It is incomplete and details may change frequently. Nonetheless we invite you to try it out as it stands, and we welcome any feedback.

2.1. Chapter 2 Introduction ¶

Warning: This text is currently out of date due to ORC API updates.

The example code has been updated and can be used. The text will be updated once the API churn dies down.

Welcome to Chapter 2 of the “Building an ORC-based JIT in LLVM” tutorial. In Chapter 1 of this series we examined a basic JIT class, KaleidoscopeJIT, that could take LLVM IR modules as input and produce executable code in memory. KaleidoscopeJIT was able to do this with relatively little code by composing two off-the-shelf ORC layers: IRCompileLayer and ObjectLinkingLayer, to do much of the heavy lifting.

In this layer we’ll learn more about the ORC layer concept by using a new layer, IRTransformLayer, to add IR optimization support to KaleidoscopeJIT.

2.2. Optimizing Modules using the IRTransformLayer ¶

In Chapter 4 of the “Implementing a language with LLVM” tutorial series the llvm FunctionPassManager is introduced as a means for optimizing LLVM IR. Interested readers may read that chapter for details, but in short: to optimize a Module we create an llvm::FunctionPassManager instance, configure it with a set of optimizations, then run the PassManager on a Module to mutate it into a (hopefully) more optimized but semantically equivalent form. In the original tutorial series the FunctionPassManager was created outside the KaleidoscopeJIT and modules were optimized before being added to it. In this Chapter we will make optimization a phase of our JIT instead. For now this will provide us a motivation to learn more about ORC layers, but in the long term making optimization part of our JIT will yield an important benefit: When we begin lazily compiling code (i.e. deferring compilation of each function until the first time it’s run), having optimization managed by our JIT will allow us to optimize lazily too, rather than having to do all our optimization up-front.

To add optimization support to our JIT we will take the KaleidoscopeJIT from Chapter 1 and compose an ORC IRTransformLayer on top. We will look at how the IRTransformLayer works in more detail below, but the interface is simple: the constructor for this layer takes a reference to the layer below (as all layers do) plus an IR optimization function that it will apply to each Module that is added via addModule:

class KaleidoscopeJIT {
private:
  std::unique_ptr<TargetMachine> TM;
  const DataLayout DL;
  RTDyldObjectLinkingLayer<> ObjectLayer;
  IRCompileLayer<decltype(ObjectLayer)> CompileLayer;

  using OptimizeFunction =
      std::function<std::shared_ptr<Module>(std::shared_ptr<Module>)>;

  IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;

public:
  using ModuleHandle = decltype(OptimizeLayer)::ModuleHandleT;

  KaleidoscopeJIT()
      : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
        ObjectLayer([]() { return std::make_shared<SectionMemoryManager>(); }),
        CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
        OptimizeLayer(CompileLayer,
                      [this](std::unique_ptr<Module> M) {
                        return optimizeModule(std::move(M));
                      }) {
    llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
  }

Our extended KaleidoscopeJIT class starts out the same as it did in Chapter 1, but after the CompileLayer we introduce a typedef for our optimization function. In this case we use a std::function (a handy wrapper for “function-like” things) from a single unique_ptr<Module> input to a std::unique_ptr<Module> output. With our optimization function typedef in place we can declare our OptimizeLayer, which sits on top of our CompileLayer.

To initialize our OptimizeLayer we pass it a reference to the CompileLayer below (standard practice for layers), and we initialize the OptimizeFunction using a lambda that calls out to an “optimizeModule” function that we will define below.

// ...
auto Resolver = createLambdaResolver(
    [&](const std::string &Name) {
      if (auto Sym = OptimizeLayer.findSymbol(Name, false))
        return Sym;
      return JITSymbol(nullptr);
    },
// ...

// ...
return cantFail(OptimizeLayer.addModule(std::move(M),
                                        std::move(Resolver)));
// ...

// ...
return OptimizeLayer.findSymbol(MangledNameStream.str(), true);
// ...

// ...
cantFail(OptimizeLayer.removeModule(H));
// ...

Next we need to replace references to ‘CompileLayer’ with references to OptimizeLayer in our key methods: addModule, findSymbol, and removeModule. In addModule we need to be careful to replace both references: the findSymbol call inside our resolver, and the call through to addModule.

std::shared_ptr<Module> optimizeModule(std::shared_ptr<Module> M) {
  // Create a function pass manager.
  auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get());

  // Add some optimizations.
  FPM->add(createInstructionCombiningPass());
  FPM->add(createReassociatePass());
  FPM->add(createGVNPass());
  FPM->add(createCFGSimplificationPass());
  FPM->doInitialization();

  // Run the optimizations over all functions in the module being added to
  // the JIT.
  for (auto &F : *M)
    FPM->run(F);

  return M;
}

At the bottom of our JIT we add a private method to do the actual optimization: optimizeModule. This function sets up a FunctionPassManager, adds some passes to it, runs it over every function in the module, and then returns the mutated module. The specific optimizations are the same ones used in Chapter 4 of the “Implementing a language with LLVM” tutorial series. Readers may visit that chapter for a more in-depth discussion of these, and of IR optimization in general.

And that’s it in terms of changes to KaleidoscopeJIT: When a module is added via addModule the OptimizeLayer will call our optimizeModule function before passing the transformed module on to the CompileLayer below. Of course, we could have called optimizeModule directly in our addModule function and not gone to the bother of using the IRTransformLayer, but doing so gives us another opportunity to see how layers compose. It also provides a neat entry point to the layer concept itself, because IRTransformLayer turns out to be one of the simplest implementations of the layer concept that can be devised:

template <typename BaseLayerT, typename TransformFtor>
class IRTransformLayer {
public:
  using ModuleHandleT = typename BaseLayerT::ModuleHandleT;

  IRTransformLayer(BaseLayerT &BaseLayer,
                   TransformFtor Transform = TransformFtor())
    : BaseLayer(BaseLayer), Transform(std::move(Transform)) {}

  Expected<ModuleHandleT>
  addModule(std::shared_ptr<Module> M,
            std::shared_ptr<JITSymbolResolver> Resolver) {
    return BaseLayer.addModule(Transform(std::move(M)), std::move(Resolver));
  }

  void removeModule(ModuleHandleT H) { BaseLayer.removeModule(H); }

  JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
    return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
  }

  JITSymbol findSymbolIn(ModuleHandleT H, const std::string &Name,
                         bool ExportedSymbolsOnly) {
    return BaseLayer.findSymbolIn(H, Name, ExportedSymbolsOnly);
  }

  void emitAndFinalize(ModuleHandleT H) {
    BaseLayer.emitAndFinalize(H);
  }

  TransformFtor& getTransform() { return Transform; }

  const TransformFtor& getTransform() const { return Transform; }

private:
  BaseLayerT &BaseLayer;
  TransformFtor Transform;
};

This is the whole definition of IRTransformLayer, from llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h, stripped of its comments. It is a template class with two template arguments: BaesLayerT and TransformFtor that provide the type of the base layer and the type of the “transform functor” (in our case a std::function) respectively. This class is concerned with two very simple jobs: (1) Running every IR Module that is added with addModule through the transform functor, and (2) conforming to the ORC layer interface. The interface consists of one typedef and five methods:

Interface	Description
ModuleHandleT	Provides a handle that can be used to identify a module set when calling findSymbolIn, removeModule, or emitAndFinalize.
addModule	Takes a given set of Modules and makes them “available for execution”. This means that symbols in those modules should be searchable via findSymbol and findSymbolIn, and the address of the symbols should be read/writable (for data symbols), or executable (for function symbols) after JITSymbol::getAddress() is called. Note: This means that addModule doesn’t have to compile (or do any other work) up-front. It can, like IRCompileLayer, act eagerly, but it can also simply record the module and take no further action until somebody calls JITSymbol::getAddress(). In IRTransformLayer’s case addModule eagerly applies the transform functor to each module in the set, then passes the resulting set of mutated modules down to the layer below.
removeModule	Removes a set of modules from the JIT. Code or data defined in these modules will no longer be available, and the memory holding the JIT’d definitions will be freed.
findSymbol	Searches for the named symbol in all modules that have previously been added via addModule (and not yet removed by a call to removeModule). In IRTransformLayer we just pass the query on to the layer below. In our REPL this is our default way to search for function definitions.
findSymbolIn	Searches for the named symbol in the module set indicated by the given ModuleHandleT. This is just an optimized search, better for lookup-speed when you know exactly a symbol definition should be found. In IRTransformLayer we just pass this query on to the layer below. In our REPL we use this method to search for functions representing top-level expressions, since we know exactly where we’ll find them: in the top-level expression module we just added.
emitAndFinalize	Forces all of the actions required to make the code and data in a module set (represented by a ModuleHandleT) accessible. Behaves as if some symbol in the set had been searched for and JITSymbol::getSymbolAddress called. This is rarely needed, but can be useful when dealing with layers that usually behave lazily if the user wants to trigger early compilation (for example, to use idle CPU time to eagerly compile code in the background).

This interface attempts to capture the natural operations of a JIT (with some wrinkles like emitAndFinalize for performance), similar to the basic JIT API operations we identified in Chapter 1. Conforming to the layer concept allows classes to compose neatly by implementing their behaviors in terms of the these same operations, carried out on the layer below. For example, an eager layer (like IRTransformLayer) can implement addModule by running each module in the set through its transform up-front and immediately passing the result to the layer below. A lazy layer, by contrast, could implement addModule by squirreling away the modules doing no other up-front work, but applying the transform (and calling addModule on the layer below) when the client calls findSymbol instead. The JIT’d program behavior will be the same either way, but these choices will have different performance characteristics: Doing work eagerly means the JIT takes longer up-front, but proceeds smoothly once this is done. Deferring work allows the JIT to get up-and-running quickly, but will force the JIT to pause and wait whenever some code or data is needed that hasn’t already been processed.

Our current REPL is eager: Each function definition is optimized and compiled as soon as it’s typed in. If we were to make the transform layer lazy (but not change things otherwise) we could defer optimization until the first time we reference a function in a top-level expression (see if you can figure out why, then check out the answer below 1). In the next chapter, however we’ll introduce fully lazy compilation, in which function’s aren’t compiled until they’re first called at run-time. At this point the trade-offs get much more interesting: the lazier we are, the quicker we can start executing the first function, but the more often we’ll have to pause to compile newly encountered functions. If we only code-gen lazily, but optimize eagerly, we’ll have a slow startup (which everything is optimized) but relatively short pauses as each function just passes through code-gen. If we both optimize and code-gen lazily we can start executing the first function more quickly, but we’ll have longer pauses as each function has to be both optimized and code-gen’d when it’s first executed. Things become even more interesting if we consider interproceedural optimizations like inlining, which must be performed eagerly. These are complex trade-offs, and there is no one-size-fits all solution to them, but by providing composable layers we leave the decisions to the person implementing the JIT, and make it easy for them to experiment with different configurations.

Next: Adding Per-function Lazy Compilation

2.3. Full Code Listing ¶

Here is the complete code listing for our running example with an IRTransformLayer added to enable optimization. To build this example, use:

# Compile
clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -O3 -o toy
# Run
./toy

Here is the code:

//===- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope --------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// Contains a simple JIT definition for use in the kaleidoscope tutorials.
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H

#include "llvm/ADT/STLExtras.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/JITSymbol.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Mangler.h"
#include "llvm/Support/DynamicLibrary.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Transforms/InstCombine/InstCombine.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Scalar/GVN.h"
#include <algorithm>
#include <memory>
#include <string>
#include <vector>

namespace llvm {
namespace orc {

class KaleidoscopeJIT {
private:
  ExecutionSession ES;
  std::shared_ptr<SymbolResolver> Resolver;
  std::unique_ptr<TargetMachine> TM;
  const DataLayout DL;
  RTDyldObjectLinkingLayer ObjectLayer;
  IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;

  using OptimizeFunction =
      std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>;

  IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;

public:
  KaleidoscopeJIT()
      : Resolver(createLegacyLookupResolver(
            ES,
            [this](const std::string &Name) -> JITSymbol {
              if (auto Sym = OptimizeLayer.findSymbol(Name, false))
                return Sym;
              else if (auto Err = Sym.takeError())
                return std::move(Err);
              if (auto SymAddr =
                      RTDyldMemoryManager::getSymbolAddressInProcess(Name))
                return JITSymbol(SymAddr, JITSymbolFlags::Exported);
              return nullptr;
            },
            [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); })),
        TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
        ObjectLayer(ES,
                    [this](VModuleKey) {
                      return RTDyldObjectLinkingLayer::Resources{
                          std::make_shared<SectionMemoryManager>(), Resolver};
                    }),
        CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
        OptimizeLayer(CompileLayer, [this](std::unique_ptr<Module> M) {
          return optimizeModule(std::move(M));
        }) {
    llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
  }

  TargetMachine &getTargetMachine() { return *TM; }

  VModuleKey addModule(std::unique_ptr<Module> M) {
    // Add the module to the JIT with a new VModuleKey.
    auto K = ES.allocateVModule();
    cantFail(OptimizeLayer.addModule(K, std::move(M)));
    return K;
  }

  JITSymbol findSymbol(const std::string Name) {
    std::string MangledName;
    raw_string_ostream MangledNameStream(MangledName);
    Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
    return OptimizeLayer.findSymbol(MangledNameStream.str(), true);
  }

  void removeModule(VModuleKey K) {
    cantFail(OptimizeLayer.removeModule(K));
  }

private:
  std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) {
    // Create a function pass manager.
    auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get());

    // Add some optimizations.
    FPM->add(createInstructionCombiningPass());
    FPM->add(createReassociatePass());
    FPM->add(createGVNPass());
    FPM->add(createCFGSimplificationPass());
    FPM->doInitialization();

    // Run the optimizations over all functions in the module being added to
    // the JIT.
    for (auto &F : *M)
      FPM->run(F);

    return M;
  }
};

} // end namespace orc
} // end namespace llvm

#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H

1: When we add our top-level expression to the JIT, any calls to functions that we defined earlier will appear to the RTDyldObjectLinkingLayer as external symbols. The RTDyldObjectLinkingLayer will call the SymbolResolver that we defined in addModule, which in turn calls findSymbol on the OptimizeLayer, at which point even a lazy transform layer will have to do its work.

2. Building a JIT: Adding Optimizations – An introduction to ORC Layers¶

2.1. Chapter 2 Introduction¶

2.2. Optimizing Modules using the IRTransformLayer¶

2.3. Full Code Listing¶

2.1. Chapter 2 Introduction ¶

2.2. Optimizing Modules using the IRTransformLayer ¶

2.3. Full Code Listing ¶