Kaleidoscope: Extending the Language: Control Flow

Posted on Edited on In

Chapter 5 Introduction

这是该教程的第五章, 前四章讲了 Kaleidoscope 的实现, 并且包括对 JIT 的支持. 在这章我们要给它添加简单的条件分支和循环指令.

If/Then/Else

给 Kaleidoscope 添加 if/then/else 支持非常直接: 将条件分支的处理分别加入 lexer, AST, parser, 和 LLVM IR 代码生成.

首先我们看一下, 要给它加入什么样的扩展功能 (what):

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

在 Kaleidoscope 中, 只有表达式 (expression), 没有语句(statement). 所以 if/then/else 表达式需要被计算出一个值, 根据条件表达式的值不同, 最终的结果可能为 then 表达式值, 也可能为 else 表达式值.

条件表达式中, 如果值为 0.0, 则表示 false, 其它都为 true

Lexer Extension for If/Then/Else

我们先添加关键字:

1
2
3
4
5
    // control flow
    tok_if = -6,
    tok_then = -7,
    tok_else = -8,
};

然后在 gettok() 中进行识别:

1
2
3
4
5
6
if (IdentifierStr == "if")
    return tok_if;
if (IdentifierStr == "then")
    return tok_then;
if (IdentifierStr == "else")
    return tok_else;

AST Extensions for If/Then/Else

然后再添加新的 AST 节点:

1
2
3
4
5
6
7
8
9
10
// IfExprAST - This class represents for if/then/else expression

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)) {}

    llvm::Value *codegen() override;
};

Parser Extension for If/Then/Else

新的 Parser 函数:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
// ifexpr ::= 'if' expr 'then' expr 'else' expr;
std::unique_ptr<ExprAST> ParseIfExpr() {
    // eat 'if'
    getNextToken();
    // condition
    auto Cond = ParseExpression();
    if (!Cond) {
        return nullptr;
    }
    if (CurTok != tok_then) {
        return LogError("Excepted then");
    }

    // eat 'then'
    getNextToken();
    auto Then = ParseExpression();
    if (!Then) {
        return nullptr;
    }
    if (CurTok != tok_else) {
        return LogError("Expected else");
    }

    // eat 'else'
    getNextToken();
    auto Else = ParseExpression();
    if (!Else) {
        return nullptr;
    }

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

将这个 ParseIfExpr() 添加到 ParsePrimary() 中:

1
2
3
4
5
6
7
8
9
10
11
12
13
switch (CurTok)
{
case tok_identifier:
    return ParseIdentifierExpr();
case tok_number:
    return ParseNumberExpr();
case '(':
    return ParseParenExpr();
case tok_if:
    return ParseIfExpr();
default:
    return LogError("unknown token when expecting an expression");
}

LLVM IR for If/Then/Else

在有了 AST 和 Parser 之后, 我们就应该实现 codegen() 方法了, 我们要在这部分介绍新的内容.

我们首先来看一段 LLVM IR:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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:       ; preds = %entry
  %calltmp = call double @foo()
  br label %ifcont

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

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

通过使用 llvm-as 和 opt 工具, 我们可以画出这段代码 (t.ll) 的 CFG 图, 运行 llvm-as-7 < t.ll| opt-7 -analyze -view-cfg 之后:

其它生成 CFG 图的方法还有在代码中调用或者在调试器中调用 F->viewCFG() 或者 F->viewCFGOnly(), 其中 F, 是一个 Function*.

不管是 LLVM IR, 还是 CFG 图, 我们都可以发现这些信息: entry block 中将 x 的值和 0 进行 one 类型的 fcmp 比较. one 的意思是 Ordered and not equal, Ordered 的解释. fcmp 表示是浮点数的比较. br 表示的是根据 %ifcond 的值跳转到 label %then 或者 %label %else.

%then 和 %else 两个 basic block 的指令都差不多, 调用其它函数, 然后跳转到 %ifcont. 所以, 在最后一个 basic block 中, 我们需要确定控制流是从哪个 block 来的, 从而确定整个 if/then/else 表达式的值.

所以这就是 phi 指令的作用, 如果控制流是从 %then 过来, 那么它的值为 %calltmp; 如果控制流是从 %else 过来, 那么它的值为 %calltmp1.

Code Generation for If/Then/Else

为了给 If/Then/Else 生成代码, 我们就要实现 Value *IfExprAST::codegen().

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
llvm::Value *IfExprAST::codegen() {
    llvm::Value * CondV = Cond->codegen();
    if (!CondV) {
        return nullptr;
    }

    // convert condition to a bool by comparing non-equal to 0.0
    CondV = Builder.CreateFCmpONE(CondV, llvm::ConstantFP::get(TheContext, llvm::APFloat(0.0)), "ifcond");

    llvm::Function *TheFunction = Builder.GetInsertBlock()->getParent();

    // Create blocks for the then and else 
    // cases. Insert the 'then' block at the 
    // end of the function
    llvm::BasicBlock *ThenBB = llvm::BasicBlock::Create(TheContext, "then", TheFunction);
    llvm::BasicBlock *ElseBB = llvm::BasicBlock::Create(TheContext, "else");
    llvm::BasicBlock *MergeBB = llvm::BasicBlock::Create(TheContext, "ifcont");
    Builder.CreateCondBr(CondV, ThenBB, ElseBB);

    // Emit then Value
    Builder.SetInsertPoint(ThenBB);
    llvm::Value *ThenV = Then->codegen();
    if (!ThenV) {
        return nullptr;
    }

    Builder.CreateBr(MergeBB);
    ThenBB = Builder.GetInsertBlock();

    // Emit the 'else' block
    TheFunction->getBasicBlockList().push_back(ElseBB);
    Builder.SetInsertPoint(ElseBB);

    llvm::Value *ElseV = Else->codegen();
    if(!ElseV) {
        return nullptr;
    }

    Builder.CreateBr(MergeBB);
    ElseBB = Builder.GetInsertBlock();

    // Emit the 'merge' block
    TheFunction->getBasicBlockList().push_back(MergeBB);
    Builder.SetInsertPoint(MergeBB);
    llvm::PHINode *PN = Builder.CreatePHI(llvm::Type::getDoubleTy(TheContext), 2, "iftmp");

    PN->addIncoming(ThenV, ThenBB);
    PN->addIncoming(ElseV, ElseBB);
    return PN;
}

首先将 CondV 生成代码, 然后使用 Builder->CreateFCmpONE() 生成条件判断代码.

然后通过 Builder->GetInsertBlock()->getParent() 获取整个函数的指针, 用来创建这个函数中其它的 BasicBlock, 也就是 then block, else block 和最后的 merge block. 注意这里的 else block 是直接插入到函数后端的. then block 和 merge block 还没有插入到函数中. 然后使用 Builder.CreateCondBr() 来创建跳转指令, 根据条件判断的值跳转到不同的 Basic Block.

然后使用 Builder.SetInsertPoint(ThenBB) 将当前代码的插入点设置为 then block 的最后, 在后面进行类似于 append 的操作. 然后递归调用 Then 的 codegen() 生成代码, 生成完毕之后创建一个无条件跳转 CreateBr, 将控制流转移到 merge block. (LLVM 中, 所有的 basic block 必须使用控制流指令来终结, 即使是 fall-through 也需要一个无条件的跳转)

最重要的是, 这里我们又需要一个 ThenBB = Builder.getInsetBlock(). 在我们设置 Phi node 时, 我们要创建 basic block 和 value 的 key/value 对, 所以Phi node 需要获取和它直接相邻的前面两个 block 的信息. 然而, 在对 Else block 和 Then block 进行代码生成时, 如果它们中嵌套了 if/then/else, 那么之前的 Then block 和 Else block 就不是和 Phi node 直接相连了. 所以我们要在代码生成完毕之后再获取一次当前的 basic block, 来获取直接与 Phi node 相连的前驱 block.

然后是将 Else block 添加到函数中, 再调用 Builder.SetInsetPoint() 设置需要生成代码的位置, 递归生成代码, 创建跳转, 更新 Phi node 前驱 block.

然后是 merge block, 前面也是一样的操作, 将 basic block 加入函数, 并设置代码插入点. 然后使用 Builder.CreatePHI() 创建一个 Phi node, 来处理分支的合并. 然后再加上进入分支的信息就可以.

最后, 整个 If/Then/Else 表达式的值就是 Phi node 的值.

for loop expression

我们想添加一个如下的 for loop:

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

这个表达式定义了一个新的变量 i, 从 1 开始迭代, 当条件 i < n 成立时, 执行循环体中的 putchard(42), 并将自身加 1. 目前我们将循环表达式的值设置为 0.

Lexer Extension for the for Loop

看代码

AST Extension for the for Loop

看代码

Parser Extension for the for loop

看代码

LLVM IR for the for loop

当我们生成代码时, 应该是如下所示的结构:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
declare double @putchard(double)

define double @printstar(double %n) {
entry:
    ; initial value = 1.0 (inlined into phi)
    br label %loop

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

    ; termination test
    %cmptmp = fcmp ult double %i, %n
    %booltmp = uitofp i1 %cmptmp to double
    %loopcond = fcmp one double %booltmp, 0.000000e+00
    br i1 %loopcond, label %loop, label %afterloop

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

Code Generation for the for Loop

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
llvm::Value *ForExprAST::codegen () {
    // Emit the start code first, without 'variable' in scope
    llvm::Value *StartVal = Start->codegen();
    if (!StartVal) {
        return nullptr;
    }

    // Make the new basic block for the loop header, 
    // inserting after current block
    llvm::Function *TheFunction = Builder.GetInsertBlock()->getParent();
    llvm::BasicBlock *PreheaderBB = Builder.GetInsertBlock();
    llvm::BasicBlock *LoopBB = llvm::BasicBlock::Create(TheContext, "loop", TheFunction);

    // Inserting an explict fall-through from current (PreheaderBB) to LoopBB.
    Builder.CreateBr(LoopBB);

    // Start insertion to LoopBB
    Builder.SetInsertPoint(LoopBB);
    // Start the PHI node with one Incoming basic block
    llvm::PHINode *Variable = Builder.CreatePHI(llvm::Type::getDoubleTy(TheContext), 2, VarName.c_str());
    Variable->addIncoming(StartVal, PreheaderBB);

    // Within the loop, the variable is defined equal to 
    // the PHI node, If it shadows an existing variable, 
    // we have to restore it, so save it now.
    llvm::Value *OldVal = NamedValues[VarName];
    NamedValues[VarName] = Variable;

    // Emit the body of the loop. This, like any other 
    // expr, can change the current BB. Note that we 
    // ignore the value computed by the body, but don't 
    // allow an error.
    if (!Body->codegen()) {
        return nullptr;
    }

    // Emit the step value
    llvm::Value *StepVal = nullptr;
    if (Step) {
        StepVal = Step->codegen();
        if (!StepVal) {
            return nullptr;
        }
    } else {
        // if not specified, use 1.0
        StepVal = llvm::ConstantFP::get(TheContext, llvm::APFloat(1.0));
    }

    // add the step value
    llvm::Value *NextVal = Builder.CreateFAdd(Variable, StepVal, "nextvar");

    // Compute the end condition
    llvm::Value *EndCond = End->codegen();
    if (!EndCond) {
        return nullptr;
    }

    // Convert condition to a bool by comparing non-equal to 0.0
    EndCond = Builder.CreateFCmpONE(EndCond, llvm::ConstantFP::get(TheContext, llvm::APFloat(0.0)), "loopcond");

    // Create the "after loop" block and insert it
    llvm::BasicBlock *LoopEndBB = Builder.GetInsertBlock();
    llvm::BasicBlock *AfterBB = llvm::BasicBlock::Create(TheContext, "afterloop", TheFunction);

    // Insert the conditional branch into the end of LoopEndBB
    Builder.CreateCondBr(EndCond, LoopBB, AfterBB);

    // Any new code will be inserted in AfterBB
    Builder.SetInsertPoint(AfterBB);

    // Add a new entry to the PHI node for the backage
    Variable->addIncoming(NextVal, LoopEndBB);

    // Restore the unshadowed variable
    if (OldVal) {
        NamedValues[VarName] = OldVal;
    } else {
        NamedValues.erase(VarName);
    }

    // Always returns 0.0
    return llvm::Constant::getNullValue(llvm::Type::getDoubleTy(TheContext));
}

首先是生成 start 部分, 对应 parser 来看就是循环变量 i 被赋予的初始值, 所以先计算这部分. 下一步就是生成循环体的代码, 首先获取当前 basic block 所在的函数 Function *TheFunction = Builder->GetInsertBlock()->getParent(), 然后把当前的 basic block 标记为 BasicBlock *PreheaderBB. 创建循环体 basic block (BasicBlock *LoopBB), 并把它添加到函数末尾. 然后在 PreheaderBB 最后添加一条 fall-through 到 LoopBB 的Br 指令.

这部分代码和之前看到的 if/then/else 代码生成的方式很相似. 由于循环体所在的 basic block 有两个 incoming 的控制流边 (fall-through 和循环尾), 我们要使用 Phi node 来标记控制流是从哪来的. 目前我们创建的只有 PreheaderBB, 所以现在只添加一个 incoming 即可. 所以, 先把函数插入点移到 LoopBB, 然后用 Builder 生成 Phi node 即可.

在对循环变量作处理时, 循环变量可能和函数参数重名, 导致冲突. 所以我们使用一个 shadow 变量先把变量符号表 (NamedValues) 中的同名值存起来, 这样就实现了作用域的概念. 然后我们给循环体生成代码, 递归的调用 Body->codegen() 即可. 如果我们设置了 StepVal, 那么就生成 Step 的代码, 没有的话默认就是 1. 然后再创建一个 NextVar 用作下一次循环时的循环变量.

最后计算结束条件 EndCond, 首先生成代码. 再使用生成的代码去创建一个 FCmpOne, 用作后续 CondBr 指令的条件判断. 然后使用 BasicBlock *LoopEndBB = Builder.GetInsertBlock() 去更新最新的 basic block, 作为循环体的最后一个 block, 并创建一个 BasicBlock *AfterBB 使得之后的代码生成都是在循环结束之后. 最后就是给 Phi node 添加一个来自于 BasicBlock *LoopEndBB 的边, 并从 shadow object 中恢复原来变量的值, 最后直接返回 0.