LLVM 是一个用于构建编译器和相关工具的开源项目。它最初是由克里斯·拉特纳（Chris Lattner）在2000年作为他的硕士学位项目在伊利诺伊大学厄巴纳-香槟分校（UIUC）开始的。LLVM 的名字最初是 “Low Level Virtual Machine” 的缩写，但现在它已经不再是一个缩写，而是一个独立的名称。它提供了强大的代码分析和优化，已经是许多现代编译器的基础。JIT 指的是即时编译，其主要思想是在运行时生成机器指令，让程序可以直接运行这些指令而无需额外操作。本文主要是介绍一下 LLVM JIT 的优势和使用 LLVM JIT 实现的一个执行引擎 jitfusion 的实现。
什么是JIT
首先我们来回想一下程序是怎么运行的。
有一串很简单的代码。
int sum(int a, int b) { return a + b; }
它的汇编如下。
.section __TEXT,__text,regular,pure_instructions
.build_version macos, 14, 0 sdk_version 15, 0
.globl __Z3sumii ; -- Begin function _Z3sumii
.p2align 2
__Z3sumii: ; @_Z3sumii
.cfi_startproc
; %bb.0:
sub sp, sp, #16
.cfi_def_cfa_offset 16
str w0, [sp, #12]
str w1, [sp, #8]
ldr w8, [sp, #12]
ldr w9, [sp, #8]
add w0, w8, w9
add sp, sp, #16
ret
.cfi_endproc
; -- End function
.subsections_via_symbols
进程的本质实际上是将一系列机器指令（这些指令实际上是一串01数据，汇编语言是其更高级的抽象）加载到进程空间的 .text 段中，然后依次执行这些指令的结果。
而 JIT（即时编译）则是在运行时将这些机器指令编译好，并放置在进程空间的某个位置，然后获取到这些指令的入口地址，通常是一个函数指针。以一个 sum 函数为例，假设通过 JIT 编译后得到了一个函数指针 s。在执行时，直接调用 s(1, 2) 就可以得到结果，这与静态编译的使用方式并无区别。
LLVM JIT的优势
一个显而易见的优势是可以利用 LLVM 的优化能力。众所周知，Clang 的优化非常强大，而 Clang 正是由 LLVM 提供的一个编译器前端。具体来说，LLVM 可以进行以下优化：
* 常量折叠：将编译时已知的常量表达式计算出来。
* 死代码消除：移除不会被执行的代码。
* 循环优化：包括循环展开、循环分割等。
* 内联：将小函数的代码直接插入到调用点，减少函数调用的开销。
* CPU 指令集优化：可以用 CPU 特有的指令集做优化，比如向量化操作。
不过，这些优化仅限于 JIT 编译时内部可见的函数。如果编译单元依赖了外部的函数，LLVM 只能看到其声明，此时无法进行更深入的优化。
相比解释执行（例如常见的 DAG 实现），JIT 编译的一个显著优势是编译完成后可以重复利用，而 DAG 每次运行都需要重新遍历甚至重新构建图。
除此之外，使用 JIT 相比 DAG 实现还有其他优化优势。
Type Inference
通常情况下，自行实现的 DAG（有向无环图）会在上层封装一个通用的数据结构，但底层必须兼容不同的数据类型。因此，在算子内部必然会存在类似于以下的操作。以标准库的 std::variant 为例：
struct ExampleStruct {
ValueType operator()(int8_t /*v*/) {}
ValueType operator()(int16_t /*v*/) {}
ValueType operator()(int32_t /*v*/) {}
ValueType operator()(int64_t /*v*/) {}
ValueType operator()(uint8_t /*v*/) {}
ValueType operator()(uint16_t /*v*/) {}
ValueType operator()(uint32_t /*v*/) {}
ValueType operator()(uint64_t /*v*/) {}
ValueType operator()(float /*v*/) {}
ValueType operator()(double /*v*/) {}
ValueType operator()(const std::string& /*v*/) {}
};
using ExampleType =
std::variant<int8_t, int16_t, int32_t, int64_t, uint8_t, uint16_t, uint32_t, uint64_t, float, double, std::string>;
void Example(ExampleType &x) {
std::visit(ExampleStruct(), x);
}
需要根据实际的类型转换数据结构并进入相应的分支。这样不仅会增加额外的跳转指令，还会妨碍编译器进行更深入的优化。而使用 LLVM JIT 时，类型可以在编译期确定，不需要再进行分支跳转，LLVM 也能够为你进行最大程度的优化。同时因为这个特性，在编译期还能帮你校验类型是否正确。
Short-Circuit Evaluation
通常情况下，DAG（有向无环图）是基于拓扑排序来实现的。在拓扑排序中，只有入度为0的节点才可以执行，这意味着所有依赖节点都需要完成后，当前节点才能执行。请思考以下这种情况：
假设一个 if 节点依赖于三个节点，其中 condition 是一个条件操作，返回 true 或 false。op1 是在 condition 为 true 时执行的节点，而 op2 则是在 condition 为 false 时执行的节点。然而，由于无法提前知道 condition 的结果，因此需要同时计算 op1 和 op2 的结果，然后再根据 condition 的结果选择其中一个。这会导致额外计算一个分支。当然，这并不是无法避免的。可以通过将拓扑排序的实现修改为深度优先搜索（DFS），根据节点的类型执行不同的遍历方式来解决这个问题。使用 JIT 时，可以先判断 condition，再执行相应的路径，从而实现短路求值。
Inlining & Register Allocation Optimization
在执行 DAG 图时，由于图是动态的，编译器无法预先知道节点之间的运行关系，因此无法进行优化。因此，每个节点在执行前需要重新加载其依赖节点的数据，并且在执行完后需要将结果写入内存的某个位置，以供下一个节点使用。如果节点都是单值运算，这通常会引入 CPU 和内存之间的 I/O 瓶颈。例如，假设需要计算 (a + b) * b，其中加法和乘法是两个不同的操作符。那么，对应需要执行的机器指令可能如下所示：
load a
load b
add a b
store a
load a
load b
mul a b
store a
在这种情况下，你需要执行 4 次内存读取操作和 2 次内存写入操作，而实际的计算操作只有 2 次。然而，CPU 的速度远远快于内存，这样就很可能会触发 I/O 瓶颈。当然，如果操作符不是单值计算，情况可能会有所改善。
使用 LLVM JIT 时，当你通过 LLVM 的 API 实现相应的函数时，LLVM 可以利用其强大的分析能力对程序进行充分优化。通过函数内联，LLVM 可以充分利用 CPU 寄存器，使得你实现的操作符无需反复读写内存。优化后的机器指令可能如下所示：
load a
load b
add a b
mul a b
store a
这样一来，你只需要进行 2 次内存读取操作和 1 次内存写入操作即可完成任务。
尽管上述优化带来了显著的性能提升，但实际上要享受这些优化也存在一些缺点。最显著的缺点之一是调试的难度大大增加。此外，为了实现最大的优化，函数实现需要完全通过 API 来完成。如果函数实现是外部链接到 C 函数，优化器将无法获取到函数的具体实现，只能基于函数声明进行优化。
jitfusion实现
jitfusion 是完全基于 LLVM JIT 实现的一个执行引擎。在 jitfusion 里，最重要的三个组件是 ExecNode，FunctionRegistry，ExecEngine。
ExecNode
ExecNode 用于表示执行流中的每一个执行节点。
首先，让我们思考一下，需要多少种类型的节点才能完整地描述一个函数。请看下面这个函数，它基本上可以涵盖所有可能的情况。
class Test {
public:
int test1(int a) {
a = a + 1;
a = test2(a, b);
if (a > 0) {
a = !a
}
}
private:
int b;
}
为了能够表达这样的一个函数，我设计了以下几种类型的节点。
* EntryArgumentNode：入参节点，用于获取由使用者传递的入口参数。对应上面用于获取 a 变量。
* ExecContextNode：执行引擎的上下文节点，用于获取引擎执行时固定传递的一个上下文变量。对应上面用于获取 b 变量
* ConstValueNode：常量节点，用于表示常量。对应上面 a + 1 中的 1。
* ConstListValueNode：常量节点，用于表示列表类型的常量。
* UnaryOPNode：一元运算节点。对应上面 !a 的操作。
* BinaryOPNode ：二元运算节点。对应上面 a + 1 的操作。
* FunctionNode：函数节点。对应上面 test2(a, b) 的调用。
* IfNode：if 节点。对应上面的 if 语句。
* SwitchNode：switch 节点，用于表达更复杂的 if-else 逻辑。
* NoOPNode：无操作节点，用于表示无操作。
FunctionRegistry
FunctionRegistry 的作用是管理函数节点的实现，用户可以通过 FunctionRegistry 注册自己实现的函数。在设计上，函数通过一个 FunctionSignature 结构来描述其声明。其结构的大概如下：
struct FunctionSignature {
std::string func_name_;
std::vector<ValueType> param_types_;
ValueType ret_type_;
};
该结构大致包括：函数名、参数类型和返回类型。同时，jitfusion 对函数重载的处理方式与 C++ 相同，采用函数名加参数类型来识别一个函数。这意味着函数名相同但参数类型不同是可以接受的，但函数名相同且参数类型相同而返回类型不同则是不被接受的。
函数内容则使用 FunctionStructure 来描述，其结构如下：
struct FunctionStructure {
FunctionType func_type;
void *c_func_ptr;
std::function<llvm::Value *(const FunctionSignature &, const std::vector<llvm::Type *> &,
const std::vector<llvm::Value *> &, IRCodeGenContext &)>
codegen_func;
};
func_type 仅支持两种类型：LLVMIntrinsicFunc 和 CFunc。对于 LLVMIntrinsicFunc 类型，需要实现 codegen_func，利用 LLVM 的 API 实现整个函数，以最大程度地利用 LLVM 的优化。而 CFunc 则直接传递一个 C 函数指针，jitfusion 会帮助你将其映射到编译好的模块的符号表中。
ExecEngine
ExecEngine 是 jitfusion 的执行引擎类，对外提供两个接口：Compile 和 Execute。
Status Compile(const std::unique_ptr<ExecNode>& exec_node, const std::unique_ptr<FunctionRegistry>& func_registry);
Status Execute(void* entry_arguments, RetType* result);
Compile 接口通过传入编排好的 ExecNode 根节点以及一个 FunctionRegistry 对象，实时编译一个模块并将其存储在内存中。用户随后可以通过 Execute 接口传入自己的参数并执行该模块以获得结果。
例子
通过 ExecNode、FunctionRegistry 和 ExecEngine 这三者的抽象，基本上可以适配大部分的场景。举个简单的例子，假设你有两个 std::vector<uint32_t> 类型的变量 a 和 b，你希望将 a 中的每个数值加 1 后写入到 b 中。其执行流程图大致如下：
实际构建执行引擎的代码如下：
#include <memory>
#include "exec_engine.h"
#include "exec_node.h"
#include "function_registry.h"
#include "type.h"
struct TestStruct {
std::vector<uint32_t> a;
std::vector<uint32_t> b;
};
jitfusion::LLVMComplexStruct load_u32s(int64_t entry_args) {
auto *text_struct = reinterpret_cast<TestStruct *>(entry_args);
jitfusion::LLVMComplexStruct ret{reinterpret_cast<int64_t>(text_struct->a.data()),
static_cast<uint32_t>(text_struct->a.size())};
return ret;
}
jitfusion::LLVMComplexStruct add(jitfusion::LLVMComplexStruct a, int64_t exec_context) {
auto *data = reinterpret_cast<uint32_t *>(a.data);
jitfusion::LLVMComplexStruct b;
auto *new_data = reinterpret_cast<uint32_t *>(
reinterpret_cast<jitfusion::ExecContext *>(exec_context)->arena.Allocate(a.len * sizeof(uint32_t)));
b.len = a.len;
for (uint32_t i = 0; i < a.len; i++) {
new_data[i] = data[i] + 1;
}
b.data = reinterpret_cast<int64_t>(new_data);
return b;
}
uint8_t store_u32s(int64_t entry_args, jitfusion::LLVMComplexStruct a) {
auto *text_struct = reinterpret_cast<TestStruct *>(entry_args);
auto *write_data = reinterpret_cast<uint32_t *>(a.data);
text_struct->b.insert(text_struct->b.end(), write_data, write_data + a.len);
return 0;
}
int main() {
std::unique_ptr<jitfusion::FunctionRegistry> func_registry;
jitfusion::FunctionRegistryFactory::CreateFunctionRegistry(&func_registry);
jitfusion::FunctionSignature load_sign("load_u32s", {jitfusion::ValueType::kI64}, jitfusion::ValueType::kU32List);
jitfusion::FunctionStructure load_func_struct = {jitfusion::FunctionType::kCFunc, reinterpret_cast<void *>(load_u32s),
nullptr};
func_registry->RegisterFunc(load_sign, load_func_struct);
jitfusion::FunctionSignature store_sign("store_u32s", {jitfusion::ValueType::kI64, jitfusion::ValueType::kU32List},
jitfusion::ValueType::kU8);
jitfusion::FunctionStructure store_func_struct = {jitfusion::FunctionType::kCFunc,
reinterpret_cast<void *>(store_u32s), nullptr};
func_registry->RegisterFunc(store_sign, store_func_struct);
jitfusion::FunctionSignature add_sign("add", {jitfusion::ValueType::kU32List, jitfusion::ValueType::kI64},
jitfusion::ValueType::kU32List);
jitfusion::FunctionStructure add_func_struct = {jitfusion::FunctionType::kCFunc, reinterpret_cast<void *>(add),
nullptr};
func_registry->RegisterFunc(add_sign, add_func_struct);
auto entry_node1 = std::unique_ptr<jitfusion::ExecNode>(new jitfusion::EntryArgumentNode);
std::vector<std::unique_ptr<jitfusion::ExecNode>> load_func_args;
load_func_args.emplace_back(std::move(entry_node1));
auto load_node =
std::unique_ptr<jitfusion::ExecNode>(new jitfusion::FunctionNode("load_u32s", std::move(load_func_args)));
auto exec_ctx_node = std::unique_ptr<jitfusion::ExecNode>(new jitfusion::ExecContextNode);
std::vector<std::unique_ptr<jitfusion::ExecNode>> add_func_args;
add_func_args.emplace_back(std::move(load_node));
add_func_args.emplace_back(std::move(exec_ctx_node));
auto add_node = std::unique_ptr<jitfusion::ExecNode>(new jitfusion::FunctionNode("add", std::move(add_func_args)));
auto entry_node2 = std::unique_ptr<jitfusion::ExecNode>(new jitfusion::EntryArgumentNode);
std::vector<std::unique_ptr<jitfusion::ExecNode>> store_func_args;
store_func_args.emplace_back(std::move(entry_node2));
store_func_args.emplace_back(std::move(add_node));
auto store_node =
std::unique_ptr<jitfusion::ExecNode>(new jitfusion::FunctionNode("store_u32s", std::move(store_func_args)));
jitfusion::ExecEngine exec_engine;
exec_engine.Compile(store_node, func_registry);
TestStruct text_st;
text_st.a = {1, 2, 3, 4};
jitfusion::RetType result;
exec_engine.Execute(&text_st, &result);
for (const auto &data : text_st.b) {
std::cout << data << std::endl;
}
}
像这样，你只需使用 ExecNode 来表达整个过程，然后通过 ExecEngine 进行编译和执行即可。
结语
目前库已经开源：https://github.com/viktorika/jitfusion。更多细节请看源码，欢迎大家使用。
扫一扫
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
