25倍性能提升，三大优化

# 25倍性能提升，三项优化 来源：https://maplant.com/2025-04-20-25x-Performance,-Three-Optimizations.html 作者：Matthew Plant 2024年12月31日，我决定将我的Scheme实现（https://www.github.com/maplant/scheme-rs）从AST解释器转换为基于CPS的JIT编译器。这样做的主要原因是为了使续延的创建变得机械化。我认为在实现支持`call-with-current-continuation`（https://en.wikipedia.org/wiki/Call-with-current-continuation）的语言时，有两种“正确”的方法： - 将你的程序转换为某种可在虚拟机中运行的字节码。这使得实现`call/cc`变得相当简单；捕获当前续延类似于保存栈和当前程序计数器以便稍后恢复。使用字节码虚拟机，这个过程可以变得安全且高效。 - 将你的程序转换为**续延传递风格**（https://en.wikipedia.org/wiki/Continuation-passing_style）。显式且机械地将程序转换为CPS会明确描述程序的续延。完成之后，`call/cc`的实现就非常简单了：只需编写一个内置函数，克隆传入续延的环境并将其传递给提供的thunk（单参函数）即可。 我选择采用CPS方式，因为我认为这将是一次更有趣的学习经历，并且符合我让`scheme-rs`尽可能快的目标；CPS是一种适合进行分析和优化的中间表示（IR），因此这似乎是让`scheme-rs`尽可能快的正确方向。 1. https://maplant.com/2025-04-20-25x-Performance,-Three-Optimizations.html#1 那么首先，让我们对我们解释器版本的scheme进行一些粗略的估算。我们的基准测试函数将是斐波那契函数，我们将调用它来计算`fib(10000)`： ```scheme (define (fib n) (define (iter a b count) (if (<= count 0) a (iter b (+ a b) (- count 1)))) (iter 0 1 n)) (fib 10000) ``` 这是一个很好的优化目标函数，因为它代表了可以最大化优化的函数：它是纯粹的、进行数学计算、迭代式的，并且广为人知。因此，我们即将实施的大多数优化的成功与否都会对这个函数产生某种影响，尤其是在编译器生命周期的早期阶段进行的优化。在最终解释器版本的`scheme-rs`中，在我的机器上，调用这个函数平均需要**40毫秒**。作为对比，`Guile`平均需要**3毫秒**，因此快了超过10倍。`Guile`可能是安装在*最多*计算机上的Scheme实现，所以这似乎是一个足够好的比较对象。我们将尝试与之竞争。 于是我着手实现了CPS。`scheme-rs`现在接收AST，将其转换为CPS，再将*后者*转换为`LLVM SSA`，然后对*最终产物*进行JIT编译。我运行了我的斐波那契基准测试： ``` fib 10000 time: [356.36 ms 357.57 ms 358.82 ms] change: [+772.31% +778.10% +783.54%] (p = 0.00 < 0.05) Performance has regressed. ``` 天哪。是的，我猜你以为这是一个关于让`scheme-rs`*与其他scheme实现相比超级快*的故事。不是！这是一个关于让`scheme-rs`*回到同一竞争水平*的故事！ 但确实*通过三项优化*，我们回到了同一水平。尽管我说是三项，但第一项优化极其重要，我认为它实际上是续延传递转换算法的一个*必要*部分，那就是： ## Beta归约 CPS转换算法通常会产生许多看起来很容易通过某些分析消除的无用闭包。例如，这是对于非常简单的`(+ 1 2 3 4)`最初的CPS输出： Scheme等效输出： ```scheme (define (k3 arg k) (+ arg 4 k)) (define (k2 arg k) (+ arg 3 k)) (define (k1 arg k) (+ arg 2 k)) (k1 1 (lambda (arg) (k2 arg (lambda (arg) (k3 arg (halt)))))) ``` CPS IR： ``` compiling: TopLevelExpr { body: Closure { args: ClosureArgs { args: [ %1288, ], variadic: true, continuation: None, }, body: ReturnValues( %1288, ), val: %1287, cexp: Closure { args: ClosureArgs { args: [ %1290, ], variadic: false, continuation: None, }, body: Closure { args: ClosureArgs { args: [ %1294, ], variadic: false, continuation: None, }, body: App( %1294, [ $+, ], ), val: %1293, cexp: Closure { args: ClosureArgs { args: [ %1291, ], variadic: false, continuation: None, }, body: Closure { args: ClosureArgs { args: [ %1296, ], variadic: false, continuation: None, }, body: Closure { args: ClosureArgs { args: [ %1298, ], variadic: false, continuation: None, }, body: Closure { args: ClosureArgs { args: [ %1300, ], variadic: false, continuation: None, }, body: Closure { args: ClosureArgs { args: [ %1302, ], variadic: false, continuation: None, }, body: App( %1291, [ %1296, %1298, %1300, %1302, %1290, ], ), val: %1301, cexp: Closure { args: ClosureArgs { args: [ %1304, ], variadic: false, continuation: None, }, body: App( %1304, [ $4, ], ), val: %1303, cexp: App( %1303, [ %1301, ], ) } }, val: %1299, cexp: Closure { args: ClosureArgs { args: [ %1306, ], variadic: false, continuation: None, }, body: App( %1306, [ $3, ], ), val: %1305, cexp: App( %1305, [ %1299, ], ) } }, val: %1297, cexp: Closure { args: ClosureArgs { args: [ %1308, ], variadic: false, continuation: None, }, body: App( %1308, [ $2, ], ), val: %1307, cexp: App( %1307, [ %1297, ], ) } }, val: %1295, cexp: Closure { args: ClosureArgs { args: [ %1310, ], variadic: false, continuation: None, }, body: App( %1310, [ $1, ], ), val: %1309, cexp: App( %1309, [ %1295, ], ) } }, val: %1292, cexp: App( %1293, [ %1292, ], ) } }, val: %1289, cexp: App( %1289, [ %1287, ], ) }, }, } ``` 这个输出的问题在于存在大量无用的函数。这很糟糕，这些都会带来开销。据我所知，有两种方法可以产生良好的CPS：要么从一开始就产生好的输出，要么在输出上运行优化过程直到它变好。我采用了一个相对简单的*Beta归约*步骤。这实际上是函数内联。我们将函数调用替换为函数体的输出，并将函数调用的参数代入函数体。如果在进行若干次此操作后，该函数不再被调用，我们就可以完全消除它。它被称为Beta归约，因为在λ演算中函数应用就叫这个名字。 有许多种方法可以实现这一点，但我们将使用一个极其简单的启发式方法：如果一个函数在其续延表达式中仅被*使用一次*且是非递归的，就用一次Beta归约替换该使用： ```rust impl Cps { pub(super) fn reduce(self) -> Self { self.beta_reduction(&mut HashMap::default()) .beta_reduction(&mut HashMap::default()) } fn beta_reduction(self, uses_cache: &mut HashMap<Local, usize>) -> Self { match self { Cps::PrimOp(prim_op, values, result, cexp) => Cps::PrimOp( prim_op, values, result, Box::new(cexp.beta_reduction(uses_cache)), ), Cps::If(cond, success, failure) => Cps::If( cond, Box::new(success.beta_reduction(uses_cache)), Box::new(failure.beta_reduction(uses_cache)), ), Cps::Closure { args, body, val, cexp, debug, } => { let body = body.beta_reduction(uses_cache); let mut cexp = cexp.beta_reduction(uses_cache); let is_recursive = body.uses(uses_cache).contains_key(&val); let uses = cexp.uses(uses_cache).get(&val).copied().unwrap_or(0); if !args.variadic && !is_recursive && uses == 1 { let reduced = cexp.reduce_function(val, &args, &body, uses_cache); if reduced { uses_cache.remove(&val); return cexp; } } Cps::Closure { args, body: Box::new(body), val, cexp: Box::new(cexp), debug, } } cexp => cexp, } } fn reduce_function( &mut self, func: Local, args: &ClosureArgs, func_body: &Cps, uses_cache: &mut HashMap<Local, usize>, ) -> bool { let new = match self { Cps::PrimOp(_, _, _, cexp) => { return cexp.reduce_function(func, args, func_body, uses_cache) } Cps::If(_, succ, fail) => { return succ.reduce_function(func, args, func_body, uses_cache) || fail.reduce_function(func, args, func_body, uses_cache) } Cps::Closure { val, body, cexp, .. } => { let reduced = body.reduce_function(func, args, func_body, uses_cache) || cexp.reduce_function(func, args, func_body, uses_cache); if reduced { uses_cache.remove(val); } return reduced; } Cps::App(Value::Var(Var::Local(operator)), applied, _) if *operator == func => { let substitutions: HashMap<_, _> = args .to_vec() .into_iter() .zip(applied.iter().cloned()) .collect(); let mut body = func_body.clone(); body.substitute(&substitutions); body } Cps::App(_, _, _) | Cps::Forward(_, _) | Cps::Halt(_) => return false, }; *self = new; true } } ``` 这项优化*极其有效*，我们瞬间收回了所有性能损失： ``` Running benches/fib.rs (target/release/deps/fib-08cbbdc89a1c73d1) fib 10000 time: [42.418 ms 43.993 ms 45.544 ms] change: [-88.135% -87.697% -87.241%] (p = 0.00 < 0.05) Performance has improved. ``` 现在我们对`(+ 1 2 3 4)`的输出不再那么疯狂了： ``` compiling: TopLevelExpr { body: Closure { args: ClosureArgs { args: [ %2142, ], variadic: true, continuation: None, }, body: Halt( %2142, ), val: %2141, cexp: App( $+, [ $1, $2, $3, $4, %2141, ], ), }, } ``` 我应该注意，我最初实现的这项优化存在微妙的错误；我已用你在这里看到的更正确的版本替换了它。如前所述，这是迄今为止最有效的优化。在几百行代码内，我们重新获得了CPS转换所损失的所有性能。让我们再添加几项优化： ## 原始操作符 有一些函数比其他函数更易于理解，无论是因为在我们的运行时系统上下文中更易理解（例如列表操作），还是仅仅因为它们更常见且经过更深入的研究（例如算术函数）。能够识别这些函数的调用以便优化它们的使用非常重要。我们的斐波那契函数中的原始操作符是算术运算符`+ - * /`和比较运算符`> < >= <= =`。因此，我们将优化这些运算符。如果我们为这些运算符提供专用的运行时函数，而不是简单的scheme或rust函数，我们应该能够获得一些性能提升。 首先，我们需要将表达式映射到原始操作符。我们通过检查全局变量的函数指针是否等于已知函数来实现这一点： ```rust impl Expression { pub fn to_primop(&self) -> Option<PrimOp> { use crate::{ num::{ add_builtin_wrapper, div_builtin_wrapper, equal_builtin_wrapper, greater_builtin_wrapper, greater_equal_builtin_wrapper, lesser_builtin_wrapper, lesser_equal_builtin_wrapper, mul_builtin_wrapper, sub_builtin_wrapper, }, proc::{Closure, FuncPtr::Bridge}, }; if let Expression::Var(Var::Global(global)) = self { let val = global.value_ref().read().clone(); let val: Gc<Closure> = val.try_into().ok()?; let val_read = val.read(); match val_read.func { Bridge(ptr) if ptr == add_builtin_wrapper => Some(PrimOp::Add), Bridge(ptr) if ptr == sub_builtin_wrapper => Some(PrimOp::Sub), Bridge(ptr) if ptr == mul_builtin_wrapper => Some(PrimOp::Mul), Bridge(ptr) if ptr == div_builtin_wrapper => Some(PrimOp::Div), Bridge(ptr) if ptr == equal_builtin_wrapper => Some(PrimOp::Equal), Bridge(ptr) if ptr == greater_builtin_wrapper => Some(PrimOp::Greater), Bridge(ptr) if ptr == greater_equal_builtin_wrapper => Some(PrimOp::GreaterEqual), Bridge(ptr) if ptr == lesser_builtin_wrapper => Some(PrimOp::Lesser), Bridge(ptr) if ptr == lesser_equal_builtin_wrapper => Some(PrimOp::LesserEqual), _ => None, } } else { None } } } ``` 当将AST编译为续延表达式时，我们可以检查它是否是原始操作符，并单独编译： ```rust fn compile_apply( operator: &Expression, args: &[Expression], call_site_id: CallSiteId, mut meta_cont: Box<dyn Fn(Value) -> Cps + '_>, ) -> Cps { let k1 = Local::gensym(); let k2 = Local::gensym(); let k3 = Local::gensym(); let k4 = Local::gensym(); Cps::Closure { args: ClosureArgs::new(vec![k2], false, None), body: Box::new(if let Some(primop) = operator.to_primop() { compile_primop(Value::from(k2), primop, Vec::new(), args) } else { Cps::App( Value::Var(Var::Global(operator.clone())), args.iter() .map(|arg| { let k = Local::gensym(); Cps::Closure { args: ClosureArgs::new(vec![k], false, None), body: Box::new(Cps::App( Value::Var(Var::Global(operator.clone())), args.iter() .map(|_| Value::Var(Var::Local(k))) .collect(), )), val: Local::gensym(), cexp: Box::new(arg.compile(Box::new(|result| { Cps::App( Value::Var(Var::Local(k)), vec![result], ) }))), debug: None, } }) .collect(), ) }), val: k1, cexp: Box::new(meta_cont(Value::from(k1))), debug: None, } } fn compile_primop( cont: Value, primop: PrimOp, mut collected_args: Vec<Value>, remaining_args: &[Expression], ) -> Cps { let (arg, tail) = match remaining_args { [] => { let val = Local::gensym(); return Cps::PrimOp( primop, collected_args, val, Box::new(Cps::App(cont, vec![Value::from(val)])), ); } [arg, tail @ ..] => (arg, tail), }; let k1 = Local::gensym(); let k2 = Local::gensym(); Cps::Closure { args: ClosureArgs::new(vec![k2], false, None), body: Box::new({ collected_args.push(Value::from(k2)); compile_primop(cont, primop, collected_args, tail) }), val: k1, cexp: Box::new(arg.compile(Box::new(|result| Cps::App(result, vec![Value::from(k1)])))), } } ``` 拥有这样的信息对于众多优化非常有益，但即使只是避免调用用户函数而调用已知函数，也是一种非常有效的优化，因为它避免了使用续延表达式： ``` Running benches/fib.rs (target/release/deps/fib-08cbbdc89a1c73d1) fib 10000 time: [16.564 ms 18.038 ms 19.554 ms] change: [-62.996% -58.997% -55.266%] (p = 0.00 < 0.05) Performance has improved. ``` ## 标记指针 在继续功能开发之前，我想要做的最后一个主要优化是修复`scheme-rs`的值类型。最初，它是一个巨大的枚举（enum），单元格（cells）通过具有内部可变性的引用计数智能指针来引用它。这相当糟糕。原因是，使用枚举我没有对值类型的最终表示形式进行足够的控制。如果我发布了`scheme-rs`，人们将开始对枚举进行模式匹配，而我将必须提供该枚举作为稳定的接口。到那时更改它将是不可能的。 因此，如果我要更改值类型的表示形式，必须在发布之前完成。我决定采用相当标准的标记指针（https://en.wikipedia.org/wiki/Tagged_pointer）表示法。我选定了一个四位的标签，为自己提供了16个可能的标签值。这相当大。大多数实现使用三位标签。增加标签大小意味着增加了分配指针的对齐要求，这可能会增加内存碎片。好消息是我可以更改它！我现在将我的Value类型设为一个不透明结构体，API消费者不能直接使用它，必须通过方法进行检查。如果我希望将来更改表示形式，我可以，只要我能够继续实现提供的函数。 人们可能会认为这会使接口更难使用、更笨重，但事实上，由于迫使我在检查和消费的方法上进行改进，最终得到的接口反而更容易使用。16个标签值并非严格必要。事实上，它甚至没有覆盖我自己的所有需求。最后一个值类型表示“其他”，指示该值指向一个`OtherData`结构体，其中包含一些不够重要的类型： ```rust #[repr(u64)] #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] pub enum ValueType { Undefined = 0, Null = 1, Boolean = 2, Character = 3, Number = 4, String = 5, Symbol = 6, Vector = 7, ByteVector = 8, Syntax = 9, Closure = 10, Record = 11, Condition = 12, Pair = 13, HashTable = 14, Other = 15, } #[derive(Clone, Trace)] #[repr(align(16))] pub enum OtherData { CapturedEnv(CapturedEnv), Transformer(Transformer), Future(Future), RecordType(RecordType), UserData(Arc<dyn Any>), } ``` 基本上，我选择这些值是为了使它们...