codon/docs/advanced/ir.md

After type checking but before native code generation, the Codon compiler makes use of a new intermediate representation called CIR, where a number of higher-level optimizations, transformations and analyses take place. CIR offers a comprehensive framework for writing new optimizations or analyses without having to deal with cumbersome abstract syntax trees (ASTs). In this section we'll give an overview of CIR, discuss the types of things you might want to use it for, and give a few examples.

At a glance

Here is a small (simplified) example showcasing CIR in action. Consider the code:

def fib(n):
    if n < 2:
        return 1
    else:
        return fib(n - 1) + fib(n - 2)

When instantiated with an int argument, the following IR gets produced (the names have been cleaned up for simplicity):

(bodied_func
  '"fib[int]"
  (type '"fib[int]")
  (args (var '"n" (type '"int") (global false)))
  (vars)
  (series
    (if (call '"int.__lt__[int,int]" '"n" 2)
      (series (return 1))
      (series
        (return
          (call
            '"int.__add__[int,int]"
            (call
              '"fib[int]"
              (call '"int.__sub__[int,int]" '"n" 1))
            (call
              '"fib[int]"
              (call '"int.__sub__[int,int]" '"n" 2))))))))

A few interesting points to consider:

• CIR is hierarchical like ASTS, but unlike ASTs it uses a vastly reduced set of nodes, making it much easier to work with and reason about.
• Operators are expressed as function calls. In fact, CIR has no explicit concept of +, -, etc. and instead expresses these via their corresponding magic methods (__add__, __sub__, etc.).
• CIR has no concept of generic types. By the time CIR is generated, all types need to have been resolved.

Structure

CIR is comprised of a set of nodes, each with a specific semantic meaning. There are nodes for representing constants (e.g. 42), instructions (e.g. call) control flow (e.g. if), types (e.g. int) and so on.

Here is a table showing the different types of nodes, LLVM IR equivalents, and some examples:

Node	LLVM equivalent	Examples
Node	n/a	all of the below
Module	Module	n/a
Type	Type	IntType, FuncType, RefType
Var	AllocaInst	Var, Func
Func	Function	BodiedFunc, ExternalFunc, LLVMFunc
Value	Value	all of the below
Const	Constant	IntConst, FloatConst, StringConst
Instr	Instruction	CallInstr, TernaryInstr, ThrowInstr
Flow	n/a	IfFlow, WhileFlow, ForFlow

Uses

CIR provides a framework for doing program optimizations, analyses and transformations. These operations are collectively known as IR passes.

A number of built-in passes and other functionalities are provided by CIR. These can be used as building blocks to create new passes. Examples include:

• Control-flow graph creation
• Reaching definitions
• Dominator analysis
• Side effect analysis
• Constant propagation and folding
• Canonicalization
• Inlining and outlining
• Python-specific optimizations targeting several common Python idioms

We're regularly adding new standard passes, so this list is always growing.

An example

Let's look at a real example. Imagine we want to write a pass that transforms expressions of the form <int const> + <int const> into a single <int const> denoting the result. In other words, a simple form of constant folding that only looks at addition on integers. The resulting pass would like this:

#include "codon/cir/transform/pass.h"

using namespace codon::ir;

class MyAddFolder : public transform::OperatorPass {
public:
  static const std::string KEY;
  std::string getKey() const override { return KEY; }

  void handle(CallInstr *v) override {
    auto *f = util::getFunc(v->getCallee());
    if (!f || f->getUnmangledName() != "__add__" || v->numArgs() != 2)
        return;

    auto *lhs = cast<IntConst>(v->front());
    auto *rhs = cast<IntConst>(v->back());

    if (lhs && rhs) {
      auto sum = lhs->getVal() + rhs->getVal();
      v->replaceAll(v->getModule()->getInt(sum));
    }
  }
};

const std::string MyAddFolder::KEY = "my-add-folder";

So how does this actually work, and what do the different components mean? Here are some notable points:

• Most passes can inherit from transform::OperatorPass. OperatorPass is a combination of an Operator and a Pass. An Operator is a utility visitor that provides hooks for handling all the different node types (i.e. through the handle() methods). Pass is the base class representing a generic pass, which simply provides a run() method that takes a module.
• Because of this, MyAddFolder::handle(CallInstr *) will be called on every call instruction in the module.
• Within our handle(), we first check to see if the function being called is __add__, indicating addition (in practice there would be a more specific check to make sure this is the __add__), and if so we extract the first and second arguments.
• We cast these arguments to IntConst. If the results are non-null, then both arguments were in fact integer constants, meaning we can replace the original call instruction with a new constant that represents the result of the addition. In CIR, all nodes are "replaceable" via a replaceAll() method.
• Lastly, notice that all passes have a KEY field to uniquely identify them.

Bidirectionality

An important and often very useful feature of CIR is that it is bidirectional, meaning it's possible to return to the type checking stage to generate new IR nodes that were not initially present in the module. For example, imagine that your pass needs to use a List with some new element type; that list's methods need to be instantiated by the type checker for use in CIR. In practice this bidirectionality often lets you write large parts of your optimization or transformation in Codon, and pull out the necessary functions or types as needed in the pass.

CIR's Module class has three methods to enable this feature:

  /// Gets or realizes a function.
  /// @param funcName the function name
  /// @param args the argument types
  /// @param generics the generics
  /// @param module the module of the function
  /// @return the function or nullptr
  Func *getOrRealizeFunc(const std::string &funcName, std::vector<types::Type *> args,
                         std::vector<types::Generic> generics = {},
                         const std::string &module = "");

  /// Gets or realizes a method.
  /// @param parent the parent class
  /// @param methodName the method name
  /// @param rType the return type
  /// @param args the argument types
  /// @param generics the generics
  /// @return the method or nullptr
  Func *getOrRealizeMethod(types::Type *parent, const std::string &methodName,
                           std::vector<types::Type *> args,
                           std::vector<types::Generic> generics = {});

  /// Gets or realizes a type.
  /// @param typeName the type name
  /// @param generics the generics
  /// @param module the module of the type
  /// @return the function or nullptr
  types::Type *getOrRealizeType(const std::string &typeName,
                                std::vector<types::Generic> generics = {},
                                const std::string &module = "");

Let's see bidirectionality in action. Consider the following Codon code:

def foo(x):
    return x*3 + x

def validate(x, y):
    assert y == x*4

a = foo(10)
b = foo(1.5)
c = foo('a')

Assume we want our pass to insert a call to validate() after each assignment that takes the assigned variable and the argument passed to foo(). We would do something like the following:

#include "codon/cir/transform/pass.h"

using namespace codon::ir;

class ValidateFoo : public transform::OperatorPass {
public:
  static const std::string KEY;
  std::string getKey() const override { return KEY; }

  void handle(AssignInstr *v) {
    auto *M = v->getModule();
    auto *var = v->getLhs();
    auto *call = cast<CallInstr>(v->getRhs());
    if (!call)
      return;

    auto *foo = util::getFunc(call->getCallee());
    if (!foo || foo->getUnmangledName() != "foo")
      return;

    auto *arg1 = call->front();         // argument of 'foo' call
    auto *arg2 = M->Nr<VarValue>(var);  // result of 'foo' call
    auto *validate =
      M->getOrRealizeFunc("validate", {arg1->getType(), arg2->getType()});
    auto *validateCall = util::call(validate, {arg1, arg2});

    insertAfter(validateCall);  // call 'validate' after 'foo'
  }
};

const std::string ValidateFoo::KEY = "validate-foo";

Note that insertAfter is a convenience method of Operator that inserts the given node "after" the node being visited (along with insertBefore which inserts before the node being visited).

Running this pass on the snippet above, we would get:

a = foo(10)
validate(10, a)

b = foo(1.5)
validate(1.5, b)

c = foo('a')
validate('a', c)

Notice that we used getOrRealizeFunc to create three different instances of validate: one for int arguments, one for float arguments and finally one for str arguments.

Extending the IR

CIR is extensible, and it is possible to add new constants, instructions, flows and types. This can be done by subclassing the corresponding custom base class; to create a custom type, for example, you would subclass CustomType. Let's look at an example where we extend CIR to add a 32-bit float type:

using namespace codon::ir;

#include "codon/cir/dsl/nodes.h"
#include "codon/cir/llvm/llvisitor.h"

class Builder : public dsl::codegen::TypeBuilder {
public:
  llvm::Type *buildType(LLVMVisitor *v) override {
    return v->getBuilder()->getFloatTy();
  }

  llvm::DIType *buildDebugType(LLVMVisitor *v) override {
    auto *module = v->getModule();
    auto &layout = module->getDataLayout();
    auto &db = v->getDebugInfo();
    auto *t = buildType(v);
    return db.builder->createBasicType(
           "float_32",
           layout.getTypeAllocSizeInBits(t),
           llvm::dwarf::DW_ATE_float);
  }
};

class Float32 : public dsl::CustomType {
public:
  std::unique_ptr<TypeBuilder> getBuilder() const override {
    return std::make_unique<Builder>();
  }
};

Notice that, in order to specify how to generate code for our Float32 type, we create a TypeBuilder subclass with methods for building the corresponding LLVM IR type. There is also a ValueBuilder for new constants and converting them to LLVM IR, as well as a CFBuilder for new instructions and creating control-flow graphs out of them.

{% hint style="info" %} When subclassing nodes other than types (e.g. instructions, flows, etc.), be sure to use the AcceptorExtend CRTP class, as in class MyNewInstr : public AcceptorExtend<MyNewInstr, dsl::CustomInstr>. {% endhint %}

Utilities

The codon/ir/util/ directory has a number of utility and generally helpful functions, for things like cloning IR, inlining/outlining, matching and more. codon/ir/util/irtools.h in particular has many helpful functions for performing various common tasks. If you're working with CIR, be sure to take a look at these functions to make your life easier!

Standard pass pipeline

These standard sets of passes are run in release-mode:

• Python-specific optimizations: a series of passes to optimize common Python patterns and idioms. Examples include dictionary updates of the form d[k] = d.get(k, x) <op> y, and optimizing them to do just one access into the dictionary, as well as optimizing repeated string concatenations or various I/O patterns.

• Imperative for-loop lowering: loops of the form for i in range(a, b, c) (with c constant) are lowered to a special IR node, since these loops are important for e.g. multithreading later.

• A series of program analyses whose results are available to later passes:

  • Control-flow analysis
  • Reaching definition analysis
  • Dominator analysis
  • Capture (or escape) analysis

• Parallel loop lowering for multithreading or GPU

• Constant propagation and folding. This also includes dead code elimination and (in non-JIT mode) global variable demotion.

Codon plugins can inject their own passes into the pipeline as well.