5. Kaleidoscope: Extending the Language: Control Flow

5.1 Chapter 5 Introduction

Welcome to Chapter 5 of the "Implementing a language with MLIR" tutorial. Parts 1-4 described the implementation of the simple Kaleidoscope language and included support for generating MLIR, followed by optimizations and a JIT compiler. Unfortunately, as presented, Kaleidoscope is mostly useless: it has no control flow other than call and return. This means that you can't have conditional branches in the code, significantly limiting its power. In this episode of "build that compiler", we'll extend Kaleidoscope to have an if/then/else expression plus a simple 'for' loop.

5.2 If/Then/Else

Extending Kaleidoscope to support if/then/else is quite straightforward. It basically requires adding support for this "new" concept to the lexer, parser, AST, and using an additional MLIR dialect. This example is nice, because it shows how easy it is to "grow" a language over time, incrementally extending it as new ideas are discovered.

Before we get going on "how" we add this extension, let's talk about "what" we want. The basic idea is that we want to be able to write this sort of thing:

def fib(x)
  if x < 3 then
    1
  else
    fib(x-1)+fib(x-2);

In Kaleidoscope, every construct is an expression: there are no statements. As such, the if/then/else expression needs to return a value like any other. Since we're using a mostly functional form, we'll have it evaluate its conditional, then return the 'then' or 'else' value based on how the condition was resolved. This is very similar to the C "?:" expression.

The semantics of the if/then/else expression is that it evaluates the condition to a boolean equality value: 0.0 is considered to be false and everything else is considered to be true. If the condition is true, the first subexpression is evaluated and returned, if the condition is false, the second subexpression is evaluated and returned. Since Kaleidoscope allows side-effects, this behavior is important to nail down.

Now that we know what we "want", let's break this down into its constituent pieces.

5.2.1 Lexer Extensions for If/Then/Else

The lexer extensions are straightforward. First we add new enum values for the relevant tokens:

// control
tok_if = -6,
tok_then = -7,
tok_else = -8,

Once we have that, we recognize the new keywords in the lexer. This is pretty simple stuff:

...
if (IdentifierStr == "def")
  return tok_def;
if (IdentifierStr == "extern")
  return tok_extern;
if (IdentifierStr == "if")
  return tok_if;
if (IdentifierStr == "then")
  return tok_then;
if (IdentifierStr == "else")
  return tok_else;
return tok_identifier;

5.2.2 AST Extensions for If/Then/Else

To represent the new expression we add a new AST node for it:

/// IfExprAST - Expression class for if/then/else.
class IfExprAST : public ExprAST {
  std::unique_ptr<ExprAST> Cond, Then, Else;

public:
  IfExprAST(std::unique_ptr<ExprAST> Cond, std::unique_ptr<ExprAST> Then,
            std::unique_ptr<ExprAST> Else)
    : Cond(std::move(Cond)), Then(std::move(Then)), Else(std::move(Else)) {}

  Value codegen() override;
};

The AST node just has pointers to the various subexpressions.

5.2.3 Parser Extensions for If/Then/Else

Now that we have the relevant tokens coming from the lexer and we have the AST node to build, our parsing logic is relatively straightforward. First we define a new parsing function:

/// ifexpr ::= 'if' expression 'then' expression 'else' expression
static std::unique_ptr<ExprAST> ParseIfExpr() {
  getNextToken();  // eat the if.

  // condition.
  auto Cond = ParseExpression();
  if (!Cond)
    return nullptr;

  if (CurTok != tok_then)
    return LogError("expected then");
  getNextToken();  // eat the then

  auto Then = ParseExpression();
  if (!Then)
    return nullptr;

  if (CurTok != tok_else)
    return LogError("expected else");

  getNextToken();

  auto Else = ParseExpression();
  if (!Else)
    return nullptr;

  return std::make_unique<IfExprAST>(std::move(Cond), std::move(Then),
                                      std::move(Else));
}

Next we hook it up as a primary expression:

static std::unique_ptr<ExprAST> ParsePrimary() {
  switch (CurTok) {
  default:
    return LogError("unknown token when expecting an expression");
  case tok_identifier:
    return ParseIdentifierExpr();
  case tok_number:
    return ParseNumberExpr();
  case '(':
    return ParseParenExpr();
  case tok_if:
    return ParseIfExpr();
  }
}

5.2.4 MLIR for If/Then/Else

Now that we have it parsing and building the AST, the final piece is adding MLIR code generation support. This is the most interesting part of the if/then/else example, because this is where it starts to introduce new concepts. All of the code above has been thoroughly described in previous chapters.

To motivate the code we want to produce, let's take a look at a simple example. Consider:

extern foo();
extern bar();
def baz(x) if x then foo() else bar();

Running the example produces the following MLIR, annotated for clarity:

$ build/toy --dump-mlir --dump-llvm-ir
ready> extern foo();
Read extern:
func.func private @foo() -> f64
ready> extern bar();
Read extern:
func.func private @bar() -> f64
ready> def baz(x) if x then foo() else bar();
Read function definition:

// SCF hierarchy
// function region
// └── block
//     └── scf.if
//         ├── then region → block → scf.yield
//         └── else region → block → scf.yield

// func.func is an operation that owns the function body region.
func.func @baz(%arg0: f64) -> f64 {
  // The function body region contains one implicit entry block.
  // └── entry block

  %cst = arith.constant 0.000000e+00 : f64
  %0 = arith.cmpf one, %arg0, %cst : f64

  // scf.if is an operation in the function's entry block.
  // It owns a then region and an else region.
  // The net construction is not very different from 
  // a regular if construction.
  %1 = scf.if %0 -> (f64) {
    // Then region
    // └── implicit entry block
    %2 = func.call @foo() : () -> f64
    // return value and save in %1 (%1 = ...)
    scf.yield %2 : f64
  } else {
    // Else region
    // └── implicit entry block
    %2 = func.call @bar() : () -> f64
    // return value and save in %1 (%1 = ...)
    scf.yield %2 : f64
  }
  
  // return the value produced by either `then` or `else`
  return %1 : f64
}

The MLIR is expressed using the SCF dialect, which represents structured control flow. An scf.if contains nested then and else regions. Each region uses scf.yield to return its value, and the selected value becomes the result %1 of the complete scf.if operation.

5.2.5 Lowering SCF to CF

Structured control flow is convenient for our frontend to generate and for high-level transformations to analyze. Before reaching LLVM, however, it must be lowered into an explicit control-flow graph. A control-flow graph has no if or else; it uses jumps between blocks to express the same thing. MLIR represents that form with the CF dialect.

Unlike SCF, the CF dialect has no if operation with nested regions. cf.cond_br chooses between named basic blocks, and cf.br transfers control from one block to another. When we run --convert-scf-to-cf, the baz function becomes:

// CF hierarchy
// function region
// ├── entry block
// ├── then block
// ├── else block
// ├── merge block(%3: f64)
// └── return block

func.func @baz(%arg0: f64) -> f64 {
  // Entry block: test the condition and select a successor.
  %cst = arith.constant 0.000000e+00 : f64
  %0 = arith.cmpf one, %arg0, %cst : f64
  cf.cond_br %0, ^bb1, ^bb2

// Then block: pass the value returned by foo to the merge block.
^bb1:
  %1 = call @foo() : () -> f64
  cf.br ^bb3(%1 : f64)

// Else block: pass the value returned by bar to the merge block.
^bb2:
  %2 = call @bar() : () -> f64
  cf.br ^bb3(%2 : f64)

// Merge block: receive the value selected by the predecessor as %3.
^bb3(%3: f64):
  cf.br ^bb4

// Return block.
^bb4:
  return %3 : f64
}

The two branches pass different values to the same destination block:

cf.br ^bb3(%1 : f64)
cf.br ^bb3(%2 : f64)

The destination receives whichever value was passed as its block argument %3:

^bb3(%3: f64):

A block argument is a value listed in a block's label. Every branch to that block supplies the corresponding value, much like arguments supplied in a function call. Here, both branches target ^bb3, so each must supply the f64 received as %3.

The block argument therefore contains the value returned by foo() when control arrives from ^bb1, and the value returned by bar() when control arrives from ^bb2. This is how the CF dialect represents an SSA value that can come from more than one predecessor.

5.2.6 Lowering CF to LLVM IR

The --dump-llvm-ir option used above prints the result after the remaining MLIR operations have been lowered and the LLVM dialect has been translated to LLVM IR. The output is shown below with descriptive names and comments added for clarity:

declare double @foo()
declare double @bar()

define double @baz(double %x) {
entry:
  %ifcond = fcmp one double %x, 0.000000e+00
  ;; Compare x with 0.0, producing an i1 condition.

  br i1 %ifcond, label %then, label %else
  ;; Branch to %then if the condition is true, or %else if it is false.

then:
  %calltmp = call double @foo()
  br label %ifcont

else:
  %calltmp1 = call double @bar()
  br label %ifcont

ifcont:
  %iftmp = phi double [ %calltmp1, %else ], [ %calltmp, %then ]
  br label %return

return:
  ret double %iftmp
}

The LLVM IR dump normally uses numbered names such as %0, %1, and %2. In the listing above, we've replaced those numbers with descriptive names to make the control flow easier to follow. We've also added comments that won't appear in the actual output.

The generated code is fairly simple: the entry block evaluates the conditional expression ("x" in our case here) and compares the result to 0.0 with the "fcmp one" instruction ('one' is "Ordered and Not Equal"). Based on the result of this expression, the code jumps to either the "then" or "else" blocks, which contain the expressions for the true/false cases.

Once the then/else blocks are finished executing, they both branch back to the 'ifcont' block to execute the code that happens after the if/then/else. In this case the only thing left to do is to return to the caller of the function. The question then becomes: how does the code know which expression to return?

The answer to this question involves an important SSA operation: the PHI node. If you're not familiar with SSA, the wikipedia article is a good introduction and there are various other introductions to it available on your favorite search engine. The short version is that "execution" of the PHI node requires "remembering" which block control came from. The PHI node takes on the value corresponding to the input control block. In this case, if control comes in from the "then" block, it gets the value of %calltmp. If control comes from the "else" block, it gets the value of %calltmp1.

The CF block argument performs the same SSA merge as this LLVM PHI node. MLIR attaches each incoming value to the branch that enters the block; LLVM instead lists the incoming value and predecessor together in the PHI node:

MLIR CF	LLVM IR
`cf.br ^bb3(%1 : f64)` from `^bb1`	`[ %calltmp, %then ]`
`cf.br ^bb3(%2 : f64)` from `^bb2`	`[ %calltmp1, %else ]`
`^bb3(%3: f64)` receives the selected value	`%iftmp = phi double ...` produces the selected value

The complete progression is therefore:

scf.if result
    -> cf block argument
    -> LLVM PHI node

For the rest of the tutorial, we can work with MLIR block arguments. Lowering will translate them into LLVM PHI nodes when LLVM IR is generated.

If we were generating LLVM IR directly, our frontend would need to construct these basic blocks and the PHI node. By starting with SCF, our frontend can describe the conditional directly and leave both lowering steps to MLIR.

5.2.7 Code Generation for If/Then/Else

In order to generate code for this, we implement the codegen method for IfExprAST.

The first part emits the condition:

Value IfExprAST::codegen() {
  Value CondV = Cond->codegen();
  if (!CondV)
    return {};

  // Convert the condition to a boolean by comparing it with 0.0.
  Value Zero = TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(0.0));
  CondV = TheBuilder->create<arith::CmpFOp>(
      getLocation(), arith::CmpFPredicate::ONE, CondV, Zero);

This code is straightforward and similar to what we saw before. We emit the expression for the condition, then compare that value to zero to get an i1 truth value.

With the condition emitted, we can create an scf.if operation. The scf dialect represents structured control flow, allowing us to describe the if expression directly instead of constructing its basic blocks ourselves.

The first callback passed to scf::IfOp builds the then region:

  bool CodegenFailed = false;
  auto IfOp = TheBuilder->create<scf::IfOp>(
      getLocation(), CondV,
      [&](OpBuilder &Builder, Location Loc) {
        Value ThenV = Then->codegen();
        if (!ThenV) {
          CodegenFailed = true;
          ThenV = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }
        Builder.create<scf::YieldOp>(Loc, ThenV);
      },

We recursively generate the value of the then expression and finish the region with scf.yield. The yielded value becomes the result of the scf.if operation when its condition is true.

Every region of an scf.if that produces a result must end with an scf.yield providing that result. If code generation fails, we record the failure and emit a temporary value so that the region remains structurally complete.

The second callback builds the else region in the same way:

      [&](OpBuilder &Builder, Location Loc) {
        Value ElseV = Else->codegen();
        if (!ElseV) {
          CodegenFailed = true;
          ElseV = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }
        Builder.create<scf::YieldOp>(Loc, ElseV);
      });

The then and else regions must yield values of the same type. In Kaleidoscope, both values are doubles, so the scf.if operation itself produces a single f64 result.

Finally, we check whether either region failed and return the result of the scf.if operation:

  if (CodegenFailed)
    return {};
  return IfOp.getResult(0);
}

This result is the value computed by the complete if/then/else expression. In our example, it is either the value returned by foo() or the value returned by bar().

Overall, we now have the ability to execute conditional code in Kaleidoscope. With this extension, Kaleidoscope is a fairly complete language that can calculate a wide variety of numeric functions. Next up we'll add another useful expression that is familiar from non-functional languages...

5.3 'for' Loop Expression

Now that we know how to add basic control flow constructs to the language, we have the tools to add more powerful things. Let's add something more aggressive, a 'for' expression:

extern putchard(char);
def printstar(n)
  for i = 1, i < n, 1.0 in
    putchard(42);  # ascii 42 = '*'

# print 100 '*' characters
printstar(100);

This expression defines a new variable ("i" in this case) which iterates from a starting value, while the condition ("i < n" in this case) is true, incrementing by an optional step value ("1.0" in this case). If the step value is omitted, it defaults to 1.0. While the loop is true, it executes its body expression. Because we don't have anything better to return, we'll just define the loop as always returning 0.0. In the future when we have mutable variables, it will get more useful.

As before, let's talk about the changes that we need to Kaleidoscope to support this.

5.3.1 Lexer Extensions for the 'for' Loop

The lexer extensions are the same sort of thing as for if/then/else:

... in enum Token ...
// control
tok_if = -6, tok_then = -7, tok_else = -8,
tok_for = -9, tok_in = -10

... in gettok ...
if (IdentifierStr == "def")
  return tok_def;
if (IdentifierStr == "extern")
  return tok_extern;
if (IdentifierStr == "if")
  return tok_if;
if (IdentifierStr == "then")
  return tok_then;
if (IdentifierStr == "else")
  return tok_else;
if (IdentifierStr == "for")
  return tok_for;
if (IdentifierStr == "in")
  return tok_in;
return tok_identifier;

5.3.2 AST Extensions for the 'for' Loop

The AST node is just as simple. It basically boils down to capturing the variable name and the constituent expressions in the node.

/// ForExprAST - Expression class for for/in.
class ForExprAST : public ExprAST {
  std::string VarName;
  std::unique_ptr<ExprAST> Start, End, Step, Body;

public:
  ForExprAST(const std::string &VarName, std::unique_ptr<ExprAST> Start,
             std::unique_ptr<ExprAST> End, std::unique_ptr<ExprAST> Step,
             std::unique_ptr<ExprAST> Body)
    : VarName(VarName), Start(std::move(Start)), End(std::move(End)),
      Step(std::move(Step)), Body(std::move(Body)) {}

  Value codegen() override;
};

5.3.3 Parser Extensions for the 'for' Loop

The parser code is also fairly standard. The only interesting thing here is handling of the optional step value. The parser code handles it by checking to see if the second comma is present. If not, it sets the step value to null in the AST node:

/// forexpr ::= 'for' identifier '=' expr ',' expr (',' expr)? 'in' expression
static std::unique_ptr<ExprAST> ParseForExpr() {
  getNextToken();  // eat the for.

  if (CurTok != tok_identifier)
    return LogError("expected identifier after for");

  std::string IdName = IdentifierStr;
  getNextToken();  // eat identifier.

  if (CurTok != '=')
    return LogError("expected '=' after for");
  getNextToken();  // eat '='.


  auto Start = ParseExpression();
  if (!Start)
    return nullptr;
  if (CurTok != ',')
    return LogError("expected ',' after for start value");
  getNextToken();

  auto End = ParseExpression();
  if (!End)
    return nullptr;

  // The step value is optional.
  std::unique_ptr<ExprAST> Step;
  if (CurTok == ',') {
    getNextToken();
    Step = ParseExpression();
    if (!Step)
      return nullptr;
  }

  if (CurTok != tok_in)
    return LogError("expected 'in' after for");
  getNextToken();  // eat 'in'.

  auto Body = ParseExpression();
  if (!Body)
    return nullptr;

  return std::make_unique<ForExprAST>(IdName, std::move(Start),
                                       std::move(End), std::move(Step),
                                       std::move(Body));
}

And again we hook it up as a primary expression:

static std::unique_ptr<ExprAST> ParsePrimary() {
  switch (CurTok) {
  default:
    return LogError("unknown token when expecting an expression");
  case tok_identifier:
    return ParseIdentifierExpr();
  case tok_number:
    return ParseNumberExpr();
  case '(':
    return ParseParenExpr();
  case tok_if:
    return ParseIfExpr();
  case tok_for:
    return ParseForExpr();
  }
}

5.3.4 MLIR and LLVM IR for the 'for' Loop

Now we get to the good part: the MLIR we want to generate for this construct. With the simple example above, we get:

func.func private @putchard(f64) -> f64

func.func @printstar(%arg0: f64) -> f64 {
  %cst = arith.constant 4.200000e+01 : f64
  %cst_0 = arith.constant 0.000000e+00 : f64
  %cst_1 = arith.constant 1.000000e+00 : f64
  %0 = scf.while (%arg1 = %cst_1) : (f64) -> f64 {
    %1 = arith.cmpf ult, %arg1, %arg0 : f64
    scf.condition(%1) %arg1 : f64
  } do {
  ^bb0(%arg1: f64):
    %1 = func.call @putchard(%cst) : (f64) -> f64
    %2 = arith.addf %arg1, %cst_1 : f64
    scf.yield %2 : f64
  }
  return %cst_0 : f64
}

The loop is represented by an scf.while operation. Its loop-carried value, %arg1, is the current value of the induction variable. It begins with %cst_1, which is 1.0 in this example.

The first region tests the end condition before each iteration. The scf.condition operation determines whether the loop continues and forwards the current induction value to the do region. Consequently, a false initial condition prevents the body from running at all. The do region emits the body and step, then uses scf.yield to carry the next induction value back to the condition.

Lowering the structured loop to the cf dialect replaces its two regions with explicit condition, body, and exit blocks. The loop-carried value is passed between those blocks as a block argument:

// CF hierarchy
// function region
// ├── entry block
// ├── condition block(%0: f64)
// ├── body block(%2: f64)
// └── exit block

func.func @printstar(%arg0: f64) -> f64 {
  %cst = arith.constant 4.200000e+01 : f64
  %cst_0 = arith.constant 0.000000e+00 : f64
  %cst_1 = arith.constant 1.000000e+00 : f64

  // Entry block: pass the initial value to the condition block.
  cf.br ^bb1(%cst_1 : f64)

// Condition block: %0 is the current induction value.
^bb1(%0: f64):
  %1 = arith.cmpf ult, %0, %arg0 : f64
  cf.cond_br %1, ^bb2(%0 : f64), ^bb3

// Body block: receive the current value, execute the body, and pass the next
// value back to the condition block.
^bb2(%2: f64):
  %3 = call @putchard(%cst) : (f64) -> f64
  %4 = arith.addf %2, %cst_1 : f64
  cf.br ^bb1(%4 : f64)

// Exit block.
^bb3:
  return %cst_0 : f64
}

The first branch supplies the initial value %cst_1 to the condition block. After each iteration, the body supplies %4, the next induction value, to the same block argument %0. If the condition remains true, %0 is passed onward to the body as its block argument %2.

Lowering these block arguments to LLVM IR produces PHI nodes. The automatically numbered names in the actual output have been replaced with descriptive names below for clarity:

declare double @putchard(double)

define double @printstar(double %n) {
entry:
  br label %loop

loop:                                             ; preds = %body, %entry
  %i = phi double [ %nextvar, %body ], [ 1.000000e+00, %entry ]
  %loopcond = fcmp ult double %i, %n
  br i1 %loopcond, label %body, label %afterloop

body:                                             ; preds = %loop
  %bodyi = phi double [ %i, %loop ]
  %calltmp = call double @putchard(double 4.200000e+01)
  %nextvar = fadd double %bodyi, 1.000000e+00
  br label %loop

afterloop:                                        ; preds = %loop
  ret double 0.000000e+00
}

This loop contains the same basic blocks and PHI nodes that we saw in the lowered if/then/else expression. The PHI node selects 1.0 when control first enters the loop from entry, and %nextvar when control returns along the loop backedge. The condition is tested before branching to body. The scf.while region arguments express these relationships in the structured form; after SCF-to-CF lowering, the block arguments express them in the control-flow graph.

5.3.5 Code Generation for the 'for' Loop

The first part of codegen is very simple: we emit the start expression before putting the loop variable in scope:

Value ForExprAST::codegen() {
  // Emit the start value before putting the loop variable in scope.
  Value StartVal = Start->codegen();
  if (!StartVal)
    return {};

Next, we save any existing symbol with the same name as the loop variable:

  auto OldValue = NamedValues.find(VarName);
  bool HadOldValue = OldValue != NamedValues.end();
  Value SavedValue = HadOldValue ? OldValue->second : Value();
  bool CodegenFailed = false;

MLIR regions define the scope of their SSA values, but they do not automatically manage the NamedValues map used by our frontend. We still need that map to resolve a source-level name such as i while walking the AST. Saving its previous entry allows a loop variable to shadow a function argument or an enclosing loop variable without making the outer value inaccessible after the loop.

We can now create the scf.while operation. Its initial loop-carried value is StartVal. The first region evaluates the condition before the body is entered:

  // The "before" region tests the loop condition. The "after" region emits
  // the body and step, then carries the next induction value back to be tested.
  TheBuilder->create<scf::WhileOp>(
      getLocation(), TypeRange{TheBuilder->getF64Type()},
      ValueRange{StartVal},
      [&](OpBuilder &Builder, Location Loc, ValueRange Args) {
        NamedValues[VarName] = Args.front();

        Value EndCond = End->codegen();
        if (!EndCond) {
          CodegenFailed = true;
          EndCond = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }

        Value Zero = Builder.create<arith::ConstantOp>(
            Loc, Builder.getF64FloatAttr(0.0));
        EndCond = Builder.create<arith::CmpFOp>(
            Loc, arith::CmpFPredicate::ONE, EndCond, Zero);
        Builder.create<scf::ConditionOp>(Loc, EndCond, Args.front());
      },

Args.front() is the current SSA value of the induction variable. We enter it in NamedValues so that the end expression can refer to the loop variable. The scf.condition operation enters the second region only when the condition is true, forwarding the current induction value to it.

The second region emits the body and calculates the next value of the induction variable by adding the step expression, or 1.0 when no step was specified:

      [&](OpBuilder &Builder, Location Loc, ValueRange Args) {
        NamedValues[VarName] = Args.front();

        if (!Body->codegen())
          CodegenFailed = true;

        Value StepVal;
        if (Step)
          StepVal = Step->codegen();
        else
          StepVal = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(1.0));
        if (!StepVal) {
          CodegenFailed = true;
          StepVal = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(1.0));
        }

        Value NextVar =
            Builder.create<arith::AddFOp>(Loc, Args.front(), StepVal);
        Builder.create<scf::YieldOp>(Loc, NextVar);
      });

The scf.yield operation carries NextVar back to the first region, where the condition is evaluated again. MLIR handles the blocks and their arguments, so we do not need to construct the loop's block arguments or backedge ourselves. Those block arguments become PHI nodes when we later lower to LLVM IR.

After constructing the loop, we restore the source-level symbol that was shadowed, or remove the loop variable if no previous definition existed:

  // Restore any variable shadowed by the loop induction variable.
  if (HadOldValue)
    NamedValues[VarName] = SavedValue;
  else
    NamedValues.erase(VarName);

  if (CodegenFailed)
    return {};

  // A for expression always returns 0.0.
  return TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(0.0));
}

The generated SSA value remains scoped to the scf.while regions, while restoring NamedValues keeps the frontend's view of source-level scope in sync. Finally, code generation of the for loop always returns 0.0.

5.4 Control Flow Graph Visualization Tools

5.4.1 MLIR

To visualize the control flow graph, you can use MLIR's --view-op-graph option. The output above shows one REPL interaction at a time, so it is not quite a standalone MLIR file. Save the complete module below as t.mlir:

// t.mlir
module {
  func.func private @foo() -> f64
  func.func private @bar() -> f64

  func.func @baz(%arg0: f64) -> f64 {
    %cst = arith.constant 0.000000e+00 : f64
    %0 = arith.cmpf one, %arg0, %cst : f64
    %1 = scf.if %0 -> (f64) {
      %2 = func.call @foo() : () -> f64
      scf.yield %2 : f64
    } else {
      %2 = func.call @bar() : () -> f64
      scf.yield %2 : f64
    }
    return %1 : f64
  }
}

Then run:

mlir-opt \
  t.mlir \
  --convert-scf-to-cf \
  '--view-op-graph=print-control-flow-edges' \
  -o /dev/null 2>&1 \
  | perl -pe 's/style = filled/style = solid/g; s/fillcolor = "[^"]+"/fillcolor = "transparent"/g' \
  > t-mlir.dot

dot -Tsvg -Gbgcolor=transparent \
  -Gcolor=black -Gfontcolor=black \
  -Ncolor=black -Nfontcolor=black \
  -Ecolor=black -Efontcolor=black \
  -o t-mlir.svg t-mlir.dot

dot -Tsvg -Gbgcolor=transparent \
  -Gcolor=gray -Gfontcolor=gray \
  -Ncolor=gray -Nfontcolor=gray \
  -Ecolor=gray -Efontcolor=gray \
  -o t-mlir-gray.svg t-mlir.dot

The --convert-scf-to-cf pass first lowers scf.if into basic blocks connected by cf.cond_br and cf.br operations. --view-op-graph=print-control-flow-edges then writes a Graphviz representation of the operations, blocks, data-flow edges, and control-flow edges. The remaining commands remove the default node colors and render two graphs with transparent backgrounds: a black version for light mode and a gray version that remains visible in dark mode.

You can then open the t-mlir.svg or t-mlir-gray.svg file in a viewer of your choice.

5.4.2 LLVM IR

LLVM's opt tool can display the corresponding LLVM control-flow graph. Save the complete module below as t.ll:

; t.ll
declare double @foo()
declare double @bar()

define double @baz(double %x) {
entry:
  %ifcond = fcmp one double %x, 0.000000e+00
  br i1 %ifcond, label %then, label %else

then:
  %calltmp = call double @foo()
  br label %ifcont

else:
  %calltmp1 = call double @bar()
  br label %ifcont

ifcont:
  %iftmp = phi double [ %calltmp1, %else ], [ %calltmp, %then ]
  ret double %iftmp
}

Then ask opt to write its control-flow graph as a DOT file. LLVM gives the nodes heat-map colors by default, so we remove those explicit colors before rendering neutral light- and dark-mode versions:

opt -passes=dot-cfg -disable-output t.ll

perl -pe \
  's/fillcolor="#[0-9a-fA-F]+",?\s*//g; s/color="#[0-9a-fA-F]+",?\s*//g; s/style=filled,?\s*//g' \
  .baz.dot > t-llvm.dot

dot -Tsvg -Gbgcolor=transparent \
  -Gcolor=black -Gfontcolor=black \
  -Ncolor=black -Nfontcolor=black \
  -Ecolor=black -Efontcolor=black \
  -o t-llvm.svg t-llvm.dot

dot -Tsvg -Gbgcolor=transparent \
  -Gcolor=gray -Gfontcolor=gray \
  -Ncolor=gray -Nfontcolor=gray \
  -Ecolor=gray -Efontcolor=gray \
  -o t-llvm-gray.svg t-llvm.dot

The generated files show this graph:

With this, we conclude the "adding control flow to Kaleidoscope" chapter of the tutorial. In this chapter we added two control flow constructs, and used them to motivate a couple of aspects of the LLVM IR that are important for front-end implementors to know. In the next chapter of our saga, we will get a bit crazier and add user-defined operators to our poor innocent language.

5.5 Download Source Code

https://github.com/alankarmisra/kaleidoscope-mlir-tutorial/

cd kaleidoscope-mlir-tutorial/code/chapter-05

5.6 Full Code Listing

Here is the complete code listing for our running example, enhanced with the if/then/else and for expressions. Here is the CMake configuration:

cmake_minimum_required(VERSION 3.20)

project(kaleidoscope-chapter-05 LANGUAGES C CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED YES)
set(CMAKE_CXX_EXTENSIONS NO)

find_package(MLIR REQUIRED CONFIG)

add_executable(toy toy.cpp)

# Make symbols in the executable available to the JIT for runtime lookup.
set_target_properties(toy PROPERTIES ENABLE_EXPORTS ON)

target_include_directories(toy PRIVATE
  ${LLVM_INCLUDE_DIRS}
  ${MLIR_INCLUDE_DIRS}
)

target_compile_definitions(toy PRIVATE ${LLVM_DEFINITIONS})

if(NOT LLVM_ENABLE_RTTI)
  target_compile_options(toy PRIVATE -fno-rtti)
endif()

llvm_map_components_to_libnames(LLVM_LIBS
  Core
  OrcJIT
  Support
  native
)

target_link_libraries(toy PRIVATE
  MLIRArithToLLVM
  MLIRArithDialect
  MLIRBuiltinToLLVMIRTranslation
  MLIRControlFlowDialect
  MLIRControlFlowToLLVM
  MLIRFuncToLLVM
  MLIRFuncDialect
  MLIRLLVMDialect
  MLIRLLVMToLLVMIRTranslation
  MLIRReconcileUnrealizedCasts
  MLIRSCFDialect
  MLIRSCFToControlFlow
  MLIRTargetLLVMIRExport
  MLIRTransforms
  ${LLVM_LIBS}
)

To build this example, use:

cmake -S . -B build \
  -DMLIR_DIR=/path/to/llvm-project/build/lib/cmake/mlir
cmake --build build
./build/toy

Here is the code:

#include "../include/KaleidoscopeJIT.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Export.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include <cctype>
#include <cstdio>
#include <cstdlib>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>

using namespace mlir;

//===----------------------------------------------------------------------===//
// Lexer
//===----------------------------------------------------------------------===//

// The lexer returns tokens [0-255] if it is an unknown character, otherwise one
// of these for known things.
enum Token {
  tok_eof = -1,

  // commands
  tok_def = -2,
  tok_extern = -3,

  // primary
  tok_identifier = -4,
  tok_number = -5,

  // control
  tok_if = -6,
  tok_then = -7,
  tok_else = -8,
  tok_for = -9,
  tok_in = -10
};

static std::string IdentifierStr; // Filled in for identifiers and keywords
static double NumVal;             // Filled in if tok_number

/// gettok - Return the next token from standard input.
static int gettok() {
  static int LastChar = ' ';

  // Skip any whitespace.
  while (isspace(LastChar))
    LastChar = getchar();

  if (isalpha(LastChar)) { // identifier: [a-zA-Z][a-zA-Z0-9]*
    IdentifierStr = LastChar;
    while (isalnum((LastChar = getchar())))
      IdentifierStr += LastChar;

    if (IdentifierStr == "def")
      return tok_def;
    if (IdentifierStr == "extern")
      return tok_extern;
    if (IdentifierStr == "if")
      return tok_if;
    if (IdentifierStr == "then")
      return tok_then;
    if (IdentifierStr == "else")
      return tok_else;
    if (IdentifierStr == "for")
      return tok_for;
    if (IdentifierStr == "in")
      return tok_in;
    return tok_identifier;
  }

  if (isdigit(LastChar) || LastChar == '.') { // Number: [0-9.]+
    std::string NumStr;
    do {
      NumStr += LastChar;
      LastChar = getchar();
    } while (isdigit(LastChar) || LastChar == '.');

    NumVal = strtod(NumStr.c_str(), nullptr);
    return tok_number;
  }

  if (LastChar == '#') {
    // Comment until end of line.
    do
      LastChar = getchar();
    while (LastChar != EOF && LastChar != '\n' && LastChar != '\r');

    if (LastChar != EOF)
      return gettok();
  }

  // Check for end of file.  Don't eat the EOF.
  if (LastChar == EOF)
    return tok_eof;

  // Otherwise, just return the character as its ascii value.
  int ThisChar = LastChar;
  LastChar = getchar();
  return ThisChar;
}

//===----------------------------------------------------------------------===//
// Abstract Syntax Tree (aka Parse Tree)
//===----------------------------------------------------------------------===//

namespace {

/// ExprAST - Base class for all expression nodes.
class ExprAST {
public:
  virtual ~ExprAST() = default;

  virtual Value codegen() = 0;
};

/// NumberExprAST - Expression class for numeric literals like "1.0".
class NumberExprAST : public ExprAST {
  double Val;

public:
  NumberExprAST(double Val) : Val(Val) {}

  Value codegen() override;
};

/// VariableExprAST - Expression class for referencing a variable, like "a".
class VariableExprAST : public ExprAST {
  std::string Name;

public:
  VariableExprAST(const std::string &Name) : Name(Name) {}

  Value codegen() override;
};

/// BinaryExprAST - Expression class for a binary operator.
class BinaryExprAST : public ExprAST {
  char Op;
  std::unique_ptr<ExprAST> LHS, RHS;

public:
  BinaryExprAST(char Op, std::unique_ptr<ExprAST> LHS,
                std::unique_ptr<ExprAST> RHS)
      : Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}

  Value codegen() override;
};

/// CallExprAST - Expression class for function calls.
class CallExprAST : public ExprAST {
  std::string Callee;
  std::vector<std::unique_ptr<ExprAST>> Args;

public:
  CallExprAST(const std::string &Callee,
              std::vector<std::unique_ptr<ExprAST>> Args)
      : Callee(Callee), Args(std::move(Args)) {}

  Value codegen() override;
};

/// IfExprAST - Expression class for if/then/else.
class IfExprAST : public ExprAST {
  std::unique_ptr<ExprAST> Cond, Then, Else;

public:
  IfExprAST(std::unique_ptr<ExprAST> Cond, std::unique_ptr<ExprAST> Then,
            std::unique_ptr<ExprAST> Else)
      : Cond(std::move(Cond)), Then(std::move(Then)), Else(std::move(Else)) {}

  Value codegen() override;
};

/// ForExprAST - Expression class for for/in.
class ForExprAST : public ExprAST {
  std::string VarName;
  std::unique_ptr<ExprAST> Start, End, Step, Body;

public:
  ForExprAST(const std::string &VarName, std::unique_ptr<ExprAST> Start,
             std::unique_ptr<ExprAST> End, std::unique_ptr<ExprAST> Step,
             std::unique_ptr<ExprAST> Body)
      : VarName(VarName), Start(std::move(Start)), End(std::move(End)),
        Step(std::move(Step)), Body(std::move(Body)) {}

  Value codegen() override;
};

/// PrototypeAST - This class represents the "prototype" for a function,
/// which captures its name, and its argument names (thus implicitly the number
/// of arguments the function takes).
class PrototypeAST {
  std::string Name;
  std::vector<std::string> Args;

public:
  PrototypeAST(const std::string &Name, std::vector<std::string> Args)
      : Name(Name), Args(std::move(Args)) {}

  func::FuncOp codegen();
  const std::string &getName() const { return Name; }
  const std::vector<std::string> &getArgs() const { return Args; }
};

/// FunctionAST - This class represents a function definition itself.
class FunctionAST {
  std::unique_ptr<PrototypeAST> Proto;
  std::unique_ptr<ExprAST> Body;

public:
  FunctionAST(std::unique_ptr<PrototypeAST> Proto,
              std::unique_ptr<ExprAST> Body)
      : Proto(std::move(Proto)), Body(std::move(Body)) {}

  func::FuncOp codegen();
};

} // end anonymous namespace

//===----------------------------------------------------------------------===//
// Parser
//===----------------------------------------------------------------------===//

/// CurTok/getNextToken - Provide a simple token buffer.  CurTok is the current
/// token the parser is looking at.  getNextToken reads another token from the
/// lexer and updates CurTok with its results.
static int CurTok;
static int getNextToken() { return CurTok = gettok(); }

/// BinopPrecedence - This holds the precedence for each binary operator that is
/// defined.
static std::map<char, int> BinopPrecedence;

/// GetTokPrecedence - Get the precedence of the pending binary operator token.
static int GetTokPrecedence() {
  if (!isascii(CurTok))
    return -1;

  // Make sure it's a declared binop.
  int TokPrec = BinopPrecedence[CurTok];
  if (TokPrec <= 0)
    return -1;
  return TokPrec;
}

/// LogError* - These are little helper functions for error handling.
std::unique_ptr<ExprAST> LogError(const char *Str) {
  fprintf(stderr, "Error: %s\n", Str);
  return nullptr;
}
std::unique_ptr<PrototypeAST> LogErrorP(const char *Str) {
  LogError(Str);
  return nullptr;
}

static std::unique_ptr<ExprAST> ParseExpression();

/// numberexpr ::= number
static std::unique_ptr<ExprAST> ParseNumberExpr() {
  auto Result = std::make_unique<NumberExprAST>(NumVal);
  getNextToken(); // consume the number
  return std::move(Result);
}

/// parenexpr ::= '(' expression ')'
static std::unique_ptr<ExprAST> ParseParenExpr() {
  getNextToken(); // eat (.
  auto V = ParseExpression();
  if (!V)
    return nullptr;

  if (CurTok != ')')
    return LogError("expected ')'");
  getNextToken(); // eat ).
  return V;
}

/// identifierexpr
///   ::= identifier
///   ::= identifier '(' expression* ')'
static std::unique_ptr<ExprAST> ParseIdentifierExpr() {
  std::string IdName = IdentifierStr;

  getNextToken(); // eat identifier.

  if (CurTok != '(') // Simple variable ref.
    return std::make_unique<VariableExprAST>(IdName);

  // Call.
  getNextToken(); // eat (
  std::vector<std::unique_ptr<ExprAST>> Args;
  if (CurTok != ')') {
    while (true) {
      if (auto Arg = ParseExpression())
        Args.push_back(std::move(Arg));
      else
        return nullptr;

      if (CurTok == ')')
        break;

      if (CurTok != ',')
        return LogError("Expected ')' or ',' in argument list");
      getNextToken();
    }
  }

  // Eat the ')'.
  getNextToken();

  return std::make_unique<CallExprAST>(IdName, std::move(Args));
}

/// ifexpr ::= 'if' expression 'then' expression 'else' expression
static std::unique_ptr<ExprAST> ParseIfExpr() {
  getNextToken(); // eat the if.

  auto Cond = ParseExpression();
  if (!Cond)
    return nullptr;

  if (CurTok != tok_then)
    return LogError("expected then");
  getNextToken(); // eat the then.

  auto Then = ParseExpression();
  if (!Then)
    return nullptr;

  if (CurTok != tok_else)
    return LogError("expected else");
  getNextToken(); // eat the else.

  auto Else = ParseExpression();
  if (!Else)
    return nullptr;

  return std::make_unique<IfExprAST>(std::move(Cond), std::move(Then),
                                     std::move(Else));
}

/// forexpr ::= 'for' identifier '=' expr ',' expr (',' expr)? 'in' expression
static std::unique_ptr<ExprAST> ParseForExpr() {
  getNextToken(); // eat the for.

  if (CurTok != tok_identifier)
    return LogError("expected identifier after for");

  std::string IdName = IdentifierStr;
  getNextToken(); // eat identifier.

  if (CurTok != '=')
    return LogError("expected '=' after for");
  getNextToken(); // eat '='.

  auto Start = ParseExpression();
  if (!Start)
    return nullptr;
  if (CurTok != ',')
    return LogError("expected ',' after for start value");
  getNextToken();

  auto End = ParseExpression();
  if (!End)
    return nullptr;

  // The step value is optional.
  std::unique_ptr<ExprAST> Step;
  if (CurTok == ',') {
    getNextToken();
    Step = ParseExpression();
    if (!Step)
      return nullptr;
  }

  if (CurTok != tok_in)
    return LogError("expected 'in' after for");
  getNextToken(); // eat the in.

  auto Body = ParseExpression();
  if (!Body)
    return nullptr;

  return std::make_unique<ForExprAST>(IdName, std::move(Start), std::move(End),
                                      std::move(Step), std::move(Body));
}

/// primary
///   ::= identifierexpr
///   ::= numberexpr
///   ::= parenexpr
///   ::= ifexpr
///   ::= forexpr
static std::unique_ptr<ExprAST> ParsePrimary() {
  switch (CurTok) {
  default:
    return LogError("unknown token when expecting an expression");
  case tok_identifier:
    return ParseIdentifierExpr();
  case tok_number:
    return ParseNumberExpr();
  case '(':
    return ParseParenExpr();
  case tok_if:
    return ParseIfExpr();
  case tok_for:
    return ParseForExpr();
  }
}

/// binoprhs
///   ::= ('+' primary)*
static std::unique_ptr<ExprAST> ParseBinOpRHS(int ExprPrec,
                                              std::unique_ptr<ExprAST> LHS) {
  // If this is a binop, find its precedence.
  while (true) {
    int TokPrec = GetTokPrecedence();

    // If this is a binop that binds at least as tightly as the current binop,
    // consume it, otherwise we are done.
    if (TokPrec < ExprPrec)
      return LHS;

    // Okay, we know this is a binop.
    int BinOp = CurTok;
    getNextToken(); // eat binop

    // Parse the primary expression after the binary operator.
    auto RHS = ParsePrimary();
    if (!RHS)
      return nullptr;

    // If BinOp binds less tightly with RHS than the operator after RHS, let
    // the pending operator take RHS as its LHS.
    int NextPrec = GetTokPrecedence();
    if (TokPrec < NextPrec) {
      RHS = ParseBinOpRHS(TokPrec + 1, std::move(RHS));
      if (!RHS)
        return nullptr;
    }

    // Merge LHS/RHS.
    LHS =
        std::make_unique<BinaryExprAST>(BinOp, std::move(LHS), std::move(RHS));
  }
}

/// expression
///   ::= primary binoprhs
///
static std::unique_ptr<ExprAST> ParseExpression() {
  auto LHS = ParsePrimary();
  if (!LHS)
    return nullptr;

  return ParseBinOpRHS(0, std::move(LHS));
}

/// prototype
///   ::= id '(' id* ')'
static std::unique_ptr<PrototypeAST> ParsePrototype() {
  if (CurTok != tok_identifier)
    return LogErrorP("Expected function name in prototype");

  std::string FnName = IdentifierStr;
  getNextToken();

  if (CurTok != '(')
    return LogErrorP("Expected '(' in prototype");

  std::vector<std::string> ArgNames;
  while (getNextToken() == tok_identifier)
    ArgNames.push_back(IdentifierStr);
  if (CurTok != ')')
    return LogErrorP("Expected ')' in prototype");

  // success.
  getNextToken(); // eat ')'.

  return std::make_unique<PrototypeAST>(FnName, std::move(ArgNames));
}

/// definition ::= 'def' prototype expression
static std::unique_ptr<FunctionAST> ParseDefinition() {
  getNextToken(); // eat def.
  auto Proto = ParsePrototype();
  if (!Proto)
    return nullptr;

  if (auto E = ParseExpression())
    return std::make_unique<FunctionAST>(std::move(Proto), std::move(E));
  return nullptr;
}

/// toplevelexpr ::= expression
static std::unique_ptr<FunctionAST> ParseTopLevelExpr() {
  if (auto E = ParseExpression()) {
    // Make an anonymous proto.
    auto Proto = std::make_unique<PrototypeAST>("__anon_expr",
                                                std::vector<std::string>());
    return std::make_unique<FunctionAST>(std::move(Proto), std::move(E));
  }
  return nullptr;
}

/// external ::= 'extern' prototype
static std::unique_ptr<PrototypeAST> ParseExtern() {
  getNextToken(); // eat extern.
  return ParsePrototype();
}

//===----------------------------------------------------------------------===//
// Code Generation
//===----------------------------------------------------------------------===//

static std::unique_ptr<MLIRContext> TheContext;
static OwningOpRef<ModuleOp> TheModule;
static std::unique_ptr<OpBuilder> TheBuilder;
static std::unique_ptr<PassManager> ThePM;
static std::map<std::string, Value> NamedValues;
static std::unique_ptr<llvm::orc::KaleidoscopeJIT> TheJIT;
static std::map<std::string, std::unique_ptr<PrototypeAST>> FunctionProtos;
static llvm::ExitOnError ExitOnErr;
static llvm::cl::opt<bool> DumpMLIR(
    "dump-mlir", llvm::cl::desc("Print generated MLIR"),
    llvm::cl::init(false));
static llvm::cl::opt<bool>
    DumpLLVMIR("dump-llvm-ir",
             llvm::cl::desc("Print LLVM IR before adding it to the JIT"),
             llvm::cl::init(false));

static Location getLocation() { return TheBuilder->getUnknownLoc(); }

Value LogErrorV(const char *Str) {
  LogError(Str);
  return {};
}

func::FuncOp getFunction(const std::string &Name) {
  // First, see if the function has already been added to the current module.
  if (auto Function = TheModule->lookupSymbol<func::FuncOp>(Name))
    return Function;

  // If not, codegen the declaration from an existing prototype.
  auto It = FunctionProtos.find(Name);
  if (It != FunctionProtos.end()) {
    auto Function = It->second->codegen();
    Function.setPrivate();
    return Function;
  }

  return {};
}

Value NumberExprAST::codegen() {
  return TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(Val));
}

Value VariableExprAST::codegen() {
  // Look this variable up in the function.
  auto It = NamedValues.find(Name);
  if (It == NamedValues.end())
    return LogErrorV("Unknown variable name");
  return It->second;
}

Value BinaryExprAST::codegen() {
  Value L = LHS->codegen();
  Value R = RHS->codegen();
  if (!L || !R)
    return {};

  switch (Op) {
  case '+':
    return TheBuilder->create<arith::AddFOp>(getLocation(), L, R);
  case '-':
    return TheBuilder->create<arith::SubFOp>(getLocation(), L, R);
  case '*':
    return TheBuilder->create<arith::MulFOp>(getLocation(), L, R);
  case '<': {
    Value Comparison = TheBuilder->create<arith::CmpFOp>(
        getLocation(), arith::CmpFPredicate::ULT, L, R);
    // Convert bool 0/1 to double 0.0 or 1.0.
    return TheBuilder->create<arith::UIToFPOp>(
        getLocation(), TheBuilder->getF64Type(), Comparison);
  }
  default:
    return LogErrorV("invalid binary operator");
  }
}

Value CallExprAST::codegen() {
  // Look up the name in the global module table.
  auto CalleeF = getFunction(Callee);
  if (!CalleeF)
    return LogErrorV("Unknown function referenced");

  // If argument mismatch error.
  if (CalleeF.getNumArguments() != Args.size())
    return LogErrorV("Incorrect # arguments passed");

  std::vector<Value> ArgsV;
  for (auto &Arg : Args) {
    ArgsV.push_back(Arg->codegen());
    if (!ArgsV.back())
      return {};
  }

  return TheBuilder->create<func::CallOp>(getLocation(), CalleeF, ArgsV)
      .getResult(0);
}

Value IfExprAST::codegen() {
  Value CondV = Cond->codegen();
  if (!CondV)
    return {};

  // Convert the condition to a boolean by comparing it with 0.0.
  Value Zero = TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(0.0));
  CondV = TheBuilder->create<arith::CmpFOp>(
      getLocation(), arith::CmpFPredicate::ONE, CondV, Zero);

  bool CodegenFailed = false;
  auto IfOp = TheBuilder->create<scf::IfOp>(
      getLocation(), CondV,
      [&](OpBuilder &Builder, Location Loc) {
        Value ThenV = Then->codegen();
        if (!ThenV) {
          CodegenFailed = true;
          ThenV = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }
        Builder.create<scf::YieldOp>(Loc, ThenV);
      },
      [&](OpBuilder &Builder, Location Loc) {
        Value ElseV = Else->codegen();
        if (!ElseV) {
          CodegenFailed = true;
          ElseV = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }
        Builder.create<scf::YieldOp>(Loc, ElseV);
      });

  if (CodegenFailed)
    return {};
  return IfOp.getResult(0);
}

Value ForExprAST::codegen() {
  // Emit the start value before putting the loop variable in scope.
  Value StartVal = Start->codegen();
  if (!StartVal)
    return {};

  auto OldValue = NamedValues.find(VarName);
  bool HadOldValue = OldValue != NamedValues.end();
  Value SavedValue = HadOldValue ? OldValue->second : Value();
  bool CodegenFailed = false;

  // The "before" region tests the loop condition. The "after" region emits
  // the body and step, then carries the next induction value back to be tested.
  TheBuilder->create<scf::WhileOp>(
      getLocation(), TypeRange{TheBuilder->getF64Type()}, ValueRange{StartVal},
      [&](OpBuilder &Builder, Location Loc, ValueRange Args) {
        NamedValues[VarName] = Args.front();

        Value EndCond = End->codegen();
        if (!EndCond) {
          CodegenFailed = true;
          EndCond = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }

        Value Zero = Builder.create<arith::ConstantOp>(
            Loc, Builder.getF64FloatAttr(0.0));
        EndCond = Builder.create<arith::CmpFOp>(Loc, arith::CmpFPredicate::ONE,
                                                EndCond, Zero);
        Builder.create<scf::ConditionOp>(Loc, EndCond, Args.front());
      },
      [&](OpBuilder &Builder, Location Loc, ValueRange Args) {
        NamedValues[VarName] = Args.front();

        if (!Body->codegen())
          CodegenFailed = true;

        Value StepVal;
        if (Step)
          StepVal = Step->codegen();
        else
          StepVal = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(1.0));
        if (!StepVal) {
          CodegenFailed = true;
          StepVal = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(1.0));
        }

        Value NextVar =
            Builder.create<arith::AddFOp>(Loc, Args.front(), StepVal);
        Builder.create<scf::YieldOp>(Loc, NextVar);
      });

  // Restore any variable shadowed by the loop induction variable.
  if (HadOldValue)
    NamedValues[VarName] = SavedValue;
  else
    NamedValues.erase(VarName);

  if (CodegenFailed)
    return {};

  // A for expression always returns 0.0.
  return TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(0.0));
}

func::FuncOp PrototypeAST::codegen() {
  // Make the function type: double(double, double), etc.
  std::vector<Type> Doubles(Args.size(), TheBuilder->getF64Type());
  auto FunctionType =
      TheBuilder->getFunctionType(Doubles, {TheBuilder->getF64Type()});

  auto Function = func::FuncOp::create(getLocation(), Name, FunctionType);
  TheModule->push_back(Function);
  return Function;
}

func::FuncOp FunctionAST::codegen() {
  // Save the prototype so declarations can be emitted in later modules.
  auto &P = *Proto;
  FunctionProtos[Proto->getName()] = std::move(Proto);
  auto TheFunction = getFunction(P.getName());

  if (!TheFunction)
    return {};

  if (!TheFunction.isDeclaration()) {
    LogError("Function cannot be redefined.");
    return {};
  }

  // A definition is visible outside the module, even if an earlier extern
  // declaration created the function with private symbol visibility.
  TheFunction.setPublic();

  // Create a new basic block to start insertion into.
  Block *EntryBlock = TheFunction.addEntryBlock();
  TheBuilder->setInsertionPointToStart(EntryBlock);

  // Record the function arguments in the NamedValues map.
  NamedValues.clear();
  unsigned Index = 0;
  for (BlockArgument Argument : TheFunction.getArguments())
    NamedValues[P.getArgs()[Index++]] = Argument;

  if (Value RetVal = Body->codegen()) {
    // Finish off the function.
    TheBuilder->create<func::ReturnOp>(getLocation(), RetVal);

    // Validate the generated code, checking for consistency.
    if (succeeded(verify(TheFunction))) {
      // Run the optimizer on the module.
      if (failed(ThePM->run(*TheModule))) {
        LogError("Could not optimize function.");
        TheFunction.erase();
        return {};
      }
      return TheFunction;
    }
  }

  // Error reading body, remove function.
  TheFunction.erase();
  return {};
}

//===----------------------------------------------------------------------===//
// Top-Level parsing and JIT Driver
//===----------------------------------------------------------------------===//

static void InitializeModuleAndManagers() {
  // Destroy objects that refer to the old context before replacing it.
  ThePM.reset();
  TheBuilder.reset();
  TheModule = OwningOpRef<ModuleOp>();
  TheContext.reset();

  // Open a new context and module.
  TheContext = std::make_unique<MLIRContext>();
  TheContext->loadDialect<arith::ArithDialect, cf::ControlFlowDialect,
                          func::FuncDialect, scf::SCFDialect>();
  TheModule = ModuleOp::create(UnknownLoc::get(TheContext.get()));

  // Create a new builder for the module.
  TheBuilder = std::make_unique<OpBuilder>(TheContext.get());

  // Create a pass manager and add a couple of simple optimizations.
  ThePM = std::make_unique<PassManager>(TheContext.get());
  ThePM->addNestedPass<func::FuncOp>(createCanonicalizerPass());
  ThePM->addNestedPass<func::FuncOp>(createCSEPass());
}

static llvm::Expected<llvm::orc::ThreadSafeModule> lowerToLLVM() {
  // Lower the high-level MLIR operations to the LLVM dialect.
  PassManager LoweringPM(TheContext.get());
  LoweringPM.addPass(createSCFToControlFlowPass());
  LoweringPM.addPass(createConvertFuncToLLVMPass());
  LoweringPM.addPass(createArithToLLVMConversionPass());
  LoweringPM.addPass(createConvertControlFlowToLLVMPass());

  // Clean up any temporary casts introduced by dialect conversion.
  LoweringPM.addPass(createReconcileUnrealizedCastsPass());
  if (failed(LoweringPM.run(*TheModule)))
    return llvm::make_error<llvm::StringError>(
        "could not lower module to the LLVM dialect",
        llvm::inconvertibleErrorCode());

  // Register the translations from MLIR's LLVM dialect to LLVM IR.
  registerBuiltinDialectTranslation(*TheContext);
  registerLLVMDialectTranslation(*TheContext);

  // Translate the lowered MLIR module into an LLVM IR module. The LLVM
  // context is kept with the module because the JIT may compile it later.
  auto LLVMContext = std::make_unique<llvm::LLVMContext>();
  auto LLVMModule = translateModuleToLLVMIR(*TheModule, *LLVMContext);
  if (!LLVMModule)
    return llvm::make_error<llvm::StringError>(
        "could not translate the LLVM dialect to LLVM IR",
        llvm::inconvertibleErrorCode());

  // Match the module's data layout to the target selected by the JIT.
  LLVMModule->setDataLayout(TheJIT->getDataLayout());

  if (DumpLLVMIR) {
    LLVMModule->print(llvm::errs(), nullptr);
    llvm::errs() << '\n';
  }

  // ThreadSafeModule transfers ownership of both objects to the ORC JIT.
  return llvm::orc::ThreadSafeModule(std::move(LLVMModule),
                                     std::move(LLVMContext));
}

static void HandleDefinition() {
  if (auto FnAST = ParseDefinition()) {
    if (auto FnIR = FnAST->codegen()) {
      if (DumpMLIR) {
        fprintf(stderr, "Read function definition:\n");
        FnIR.print(llvm::errs(), OpPrintingFlags().assumeVerified());
        fprintf(stderr, "\n");
      }

      ExitOnErr(TheJIT->addModule(ExitOnErr(lowerToLLVM())));
      InitializeModuleAndManagers();
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

static void HandleExtern() {
  if (auto ProtoAST = ParseExtern()) {
    if (auto FnIR = ProtoAST->codegen()) {
      FnIR.setPrivate();
      if (DumpMLIR) {
        fprintf(stderr, "Read extern:\n");
        FnIR.print(llvm::errs(), OpPrintingFlags().assumeVerified());
        fprintf(stderr, "\n");
      }
      FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

static void HandleTopLevelExpression() {
  // Evaluate a top-level expression into an anonymous function.
  if (auto FnAST = ParseTopLevelExpr()) {
    if (auto FnIR = FnAST->codegen()) {
      if (DumpMLIR) {
        fprintf(stderr, "Read top-level expression:\n");
        FnIR.print(llvm::errs(), OpPrintingFlags().assumeVerified());
        fprintf(stderr, "\n");
      }

      auto RT = TheJIT->getMainJITDylib().createResourceTracker();
      ExitOnErr(TheJIT->addModule(ExitOnErr(lowerToLLVM()), RT));
      InitializeModuleAndManagers();

      auto ExprSymbol = ExitOnErr(TheJIT->lookup("__anon_expr"));
      double (*FP)() = ExprSymbol.getAddress().toPtr<double (*)()>();
      fprintf(stderr, "Evaluated to %f\n", FP());

      ExitOnErr(RT->remove());
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

//===----------------------------------------------------------------------===//
// "Library" functions that can be "extern'd" from user code.
//===----------------------------------------------------------------------===//

#ifdef _WIN32
#define DLLEXPORT __declspec(dllexport)
#else
#define DLLEXPORT
#endif

/// putchard - putchar that takes a double and returns 0.
extern "C" DLLEXPORT double putchard(double X) {
  fputc((char)X, stderr);
  return 0;
}

/// printd - printf that takes a double, prints it as "%f\n", and returns 0.
extern "C" DLLEXPORT double printd(double X) {
  fprintf(stderr, "%f\n", X);
  return 0;
}

/// top ::= definition | external | expression | ';'
static void MainLoop() {
  while (true) {
    switch (CurTok) {
    case tok_eof:
      return;
    case ';': // ignore top-level semicolons.
      fprintf(stderr, "ready> ");
      getNextToken();
      continue;
    case tok_def:
      HandleDefinition();
      break;
    case tok_extern:
      HandleExtern();
      break;
    default:
      HandleTopLevelExpression();
      break;
    }
    if (CurTok != tok_eof && CurTok != ';')
      fprintf(stderr, "ready> ");
  }
}

//===----------------------------------------------------------------------===//
// Main driver code.
//===----------------------------------------------------------------------===//

int main(int argc, char **argv) {
  llvm::cl::ParseCommandLineOptions(argc, argv, "Kaleidoscope JIT\n");

  llvm::InitializeNativeTarget();
  llvm::InitializeNativeTargetAsmPrinter();
  llvm::InitializeNativeTargetAsmParser();

  // Install standard binary operators.
  // 1 is lowest precedence.
  BinopPrecedence['<'] = 10;
  BinopPrecedence['+'] = 20;
  BinopPrecedence['-'] = 20;
  BinopPrecedence['*'] = 40; // highest.

  // Prime the first token.
  fprintf(stderr, "ready> ");
  getNextToken();

  TheJIT = ExitOnErr(llvm::orc::KaleidoscopeJIT::Create());

  // Make the first module, which holds newly generated code.
  InitializeModuleAndManagers();

  // Run the main "interpreter loop" now.
  MainLoop();

  return 0;
}

< 4. Kaleidoscope: Adding JIT and Optimizer Support

6. Kaleidoscope: Extending the Language: User-defined Operators >