Yorkin Blog notes and ideas about PLDI

从零开始的CPS编译之旅

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 = valuelet 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不相关。