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.
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 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 |
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"
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:
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 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.
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 %}
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 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:
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.