Yorkin Blog notes and ideas about PLDI

从零开始的CPS编译之旅

March 31, 2024, Posted by Yorkin

Continuation很好的捕捉了程序控制状态最基本的概念：下一步该做什么。抽象的概念应该与具体的实现无关，而喜欢用GOTO、指针、PC寄存器作为解释，和万物基于C/C++的人都没有领略到这一点。Continuation在社区中貌似一直和LISP的传说同时出现。作为一门不主流的”失败“的编程语言，经常和它一起出现的东西也被大众们贴上了不好的标签。然而当搜寻和最近PL的热门话题Algebraic effect相关的资料时，就会发现delimited continuation的字眼，甚至还有一篇2018年的论文叫做《Capturing the Future by Replaying the Past》。现实一直很荒诞。

出于对CPS的好奇，2023年我完成毕业设计之后，就计划使用CPS作为IR，写一个函数式的编译器。最早参考的是《Compiling with continuation》这本书，在经历了几个月”使用Haskell还是OCaml来实现“的反复横跳之后，同事告诉我，你看的材料有些过时了，应该使用second class而非first class的continuation在IR中抽象程序的控制流，否则很难做到编译的结果是高效的。

Compiling with Continuation, continued

重新参考标题这篇2007年的论文，将toy ML编译到untyped CPS的规则看起来很简单（同时用k和花体k的部分除外，对近视人士非常坏）：

Naive CPS transformation

在General rewrites表格里还包含了function/continuation inlining，constant folding，dead code elimination通用优化的规则。

rewrite-rule

所有的规则都可以编写在一个pass中，论文宣称可以在线性时间内做完这些。不过算法需要几轮迭代才能达到不动点，因此最坏的情况是O(N^2)的。我们可以维护一个census上下文：尽可能接近地表示变量在程序中的出现次数。精确的census信息能够增加一轮迭代里shrinking reductions操作，以此减少迭代次数。我们可以函数式地实现这个算法，大概形状是：

alt text

此时一名眼疾手快的Haskell批惊呼被论文骗了！因为这个算法的实现大部分是纯的，唯独需要一点点的副作用去维护重写时使用的fresh ident数量。写haskell的时候多爽，写到一半要加上State monad和do notation的时候就有多狼狈。那能怎么办，我又不想用unsafe操作开洞，修改后的代码大概是这样的：


data Value
  = Var Name
  | I32 Int
  | Unit
  | Tuple [Name]
  | Cont Name Term -- env x e
  | Fn Name (Maybe Name) [Name] Term -- k env x e
  deriving (Eq, Show, Ord, Read, Data)

type Cont = Name
type Function = Name
type Argument = Name
type Closure = Name

data Term
  = LetVal Name Value Term
  | LetSel Name Int Name Term 
  | LetCont Name Name Term Term
  | LetConts [(Name, Value)] Term
  | LetFns [(Name, Value)] Term
  | Continue Name Name
  | Apply Function Cont (Maybe Closure) [Argument]
  | LetPrim Name Primitive [Name] Term
  | Switch Name [C.Constant] [Term] (Maybe Term)
  | Halt Name
  deriving (Eq, Show, Ord, Read, Data)

simp :: Census -> Env -> Subst -> Term -> CompEnv Term
simp census env s p =
  case p of
    LetVal x v l ->
      if count census x == 0
        then simp census env s l
        else do
          v' <- simpVal census env s v
          LetVal x v' <$> simp census (addEnv env x v') s l
    LetSel x i y l ->
      let y' = applySubst s y
       in case lookup env y' of
            Just (Tuple elems) -> simp census env (extendSubst s x (elems !! i)) l
            _ -> LetSel x i y' <$> simp census env s l
    LetCont k x l m -> do
      l' <- simp census env s l
      case count census k of
        0 -> simp census env s m
        _ -> LetCont k x l' <$> simp census (addEnv env k (Cont x l')) s m
    LetFns fns m -> do
      fns' <- mapM (\(x, b) -> (x,) <$> simpVal census env s b) fns
      m' <- simp census env s m
      pure $ LetFns fns' m'
    Continue k x ->
      let x' = applySubst s x
       in let k' = applySubst s k
           in case lookup env k' of
                Just (Cont y l) -> inst l >>= simp census env (extendSubst s y x')
                _ -> pure $ Continue k' x'
    Apply f k _ xs ->
      let f' = applySubst s f
          k' = applySubst s k
          xs' = map (applySubst s) xs
       in case lookup env f' of
            Just (Fn k1 _ xs1 l) ->
              inst l >>= simp census env (extendSubst (extendSubsts s (zip xs1 xs')) k1 k')
            _ -> pure $ Apply f' k' Nothing xs'
    LetPrim n op ns t ->
      let fallback = simp census env s t
       in if count census n == 0
            then fallback
            else
              let ns' = map (applySubst s) ns
               in case (op, mapM (lookup env) ns') of
                    (Primitive.EQ, Just [I32 a, I32 b]) ->
                      let v = if a == b then I32 1 else I32 0 in LetVal n v <$> simp census (addEnv env n v) s t
                    (Add2, Just [I32 a, I32 b]) ->
                      let v = I32 $ a + b in LetVal n v <$> simp census (addEnv env n v) s t
                    _ -> LetPrim n op ns' <$> fallback

完整代码在effektos仓库。

这篇论文还给出了另一种基于图的O(N)的实现，从基准测试的结果单独来看也很优秀。不过由于数据结构的不同，将CPS IR转换成图也需要开销，放整个编译管线中来看起来没有明显的优势。所以这部分的细节我就跳过了。（再一次说明优化时参考基准测试有多重要！）

CPS，但是Closure Passing Style

下一步当然是闭包转换。这部分参考了USTC的Advanced Topics in Software Technologies讲义。

closure conversion

上图的实现是：

  • 提取function/cont的自由变量成一个事先声明的tuple, 里面塞着自由变量和”函数指针“。这个tuple命名成function/cont的名字去代替它们。
  • 所有的function/cont apply都重写成取出”函数指针“，然后apply这个函数。
  • 在function/cont内部插入取出闭包中的自由变量的代码。

在实践中我把function hoisting在这一步之后做了。经过了closure conversion，function没有自由变量捕获之后也就不必嵌套在scope里。

经过进一步的known function分析可以把一些闭包优化掉（参考《Optimizing Closures in O(0) time》）。但是现在更紧急的是，所有的continuation都带上了闭包，每个控制流的转移都需要传递闭包，这显然是不能接受的。

CPS are SSA done right

如何把continuation的闭包全部优化掉，这部分问题困扰了我一周。目前来看，second class continuation无法被当作函数一样被传递，能不能把cont编译成GOTO label，cont apply编译成Jump？后来群友指路EPFL的课程，Intermediate representations这节课的讲义揭示了modern CPS和SSA的关系：

comparison of IRs

In other words, functional IRs based on continuations, like CPS/L3, are SSA done right, in the sense that they avoid introducing the weird (and ultimetely useless) ϕ-functions.

好好好，现在感觉SSA也非常的函数式，联想到传统编译器爱用SSA，再联想到GCC RTL那一堆括号的中间格式，让我们做个”X也是函数式编程语言“和”X也是lisp“的九宫格，看看C/C++们的骄傲处在”不主流的失败者阵营“九宫格的哪个位置。

那么，当最终的编译目标是x86这样的平台时，只需要像编译SSA那样把CPS编译到RTL（register transfer language），把continuation编译成GOTO label，continuation application编译成goto，把continuation参数的传递翻译成MOV opernad, reg。RTL的抽象已经非常接近常见的汇编，下一步，只需要写个寄存器分配，翻亿下指令集手册把工业垃圾的各种屎山规则纳入考虑，就可以产出高效的编译结果了。

人类友好的CPS IR

虽然CPS是做对了的SSA，但是打印出来的效果不太易读。

不能透露姓名的鸽

针对这个某群群友讨论了几点改善办法:

  1. 从CPS IR的pretty printer改进

    1. 对于仅apply一次的cont，可以打印的时候inline进来，就像OCaml的let*重载和haskell的do notation那样

    2. 把cont放到where后，例如

    letcont k1 x = body in 
  letcont k2 x = body in 
    letcont k3 x = body in expr

    打印成

    body
where k1 x = body
      k2 x = body
      k3 x = body

  2. 从CPS生成改进

    1. 只在控制流处生成cont。let x = value，let x = y primOp z是不需要的，避免缩进层数过深

在我的实践项目effektos的做法是1b和2a的结合，并且将所有的嵌套的cont挪出来拍平。结合function hoisting，这样打印结果最多只有两层缩进：

LetFix 
  f env k x =
     ...
     where
     cont1 x = body
     cont2 x = body
  g env k x =  
     ...
     where
     cont1 x = body
     cont2 x = body
in 
expr

Compiling with Continuation, to be continued

这篇博客只是简单记录一下探索CPS IR遇到的问题和有用的参考。对于我正在的进行的项目effktos还需要阅读更多论文，解决更多的问题才能产出真正可用的函数式编译器（又是一大堆体力活！）。除去已经提及的，这里留下一些正在进行的主题以及参考, 有些和CPS不相关。

Related posts