docs/developers/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) albeit while preserving higher-level program constructs and semantics.
In Codon's optimization framework, we often found the need to reason about higher-level Python constructs, which wasn't feasible through LLVM IR. For instance, one optimization Codon performs pertains to removing redundant dictionary accesses, e.g. as found in this common code pattern:
d[k] = d.get(k, 0) + 1 # increment count of key 'k' in dict 'd'
The original code performs two lookups of key k, but only one lookup is required. To
address this, we want to find patterns of the form d[k] = d.get(k, i) + j and replace them
with an implementation that performs just one lookup of k. It turns out that this is nearly
impossible to do at the LLVM IR level as dict.__setitem__() (which corresponds to the
assignment d[k] = ...) and dict.get() both compile to large, complex LLVM IR that we'd
have no hope of recognizing and making sense of after the fact — in other words, the semantic
meaning of the original code (i.e. "increment the count of k in d") is effectively lost.
To solve this problem, Codon employs its own IR, which sits between the abstract syntax tree and
LLVM IR in the compilation process. Codon IR is much simpler than the AST, but still retains
enough information to reason about semantics and higher-level constructs. Identifying the
dictionary pattern above, for example, amounts to simply identifying a short sequence of IR
that resembles dict.__setitem__(d, k, dict.get(d, k, i) + j).
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:
+, -, etc. and instead expresses these via their corresponding
magic methods (__add__, __sub__, etc.).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 the node hierarchy:
graph TD
N(Node) --> M(Module);
N --> T(Type);
N --> VR(Var);
N --> VL(Value);
VR --> FN(Func);
VL --> C(Const);
VL --> I(Instr);
VL --> FL(Flow);
Here is a breakdown of the different types of nodes:
x = y, but they can also be produced from control-flow statements like
for i in range(10), which creates variable i.x if y else z) and so on.if-statements, while-loops,
try-except etc. These are first-class values in Codon IR.42, "hello" or 3.14.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:
We're regularly adding new standard passes, so this list is always growing.
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"
#include "codon/cir/util/irtools.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) {
long 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:
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.MyAddFolder::handle(CallInstr *) will be called on every call instruction
in the module.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.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.KEY field to uniquely identify them.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 mangled type name
/// @param generics the generics
/// @return the function or nullptr
types::Type *getOrRealizeType(const std::string &typeName,
std::vector<types::Generic> generics = {});
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"
#include "codon/cir/util/irtools.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) override {
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()});
if (!validate)
return;
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.
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
(note that Codon already has a native 32-bit float type; this is just for demonstration purposes):
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.
!!! tip
When subclassing nodes other than types (e.g. instructions, flows, etc.), be sure to use the `AcceptorExtend`
[CRTP](https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern) class, as in
`class MyNewInstr : public AcceptorExtend<MyNewInstr, dsl::CustomInstr>`.
The codon/cir/util/ directory has a number of utility and generally helpful functions, for things like
cloning IR, inlining/outlining, matching and more. codon/cir/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!
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, various I/O patterns, list operations and so on.
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:
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.
Library-specific optimizations, such as operator fusion for NumPy.
Codon plugins can inject their own passes into the pipeline as well.