Optimizing GoAWK with a bytecode compiler and virtual machine

Ben Hoyt

使用字节码编译器和虚拟机优化 GoAWK

原文由 Ben Hoyt 发布,订阅该博客

摘要:我最近通过将树遍历解释器替换为带有虚拟机解释器的字节码编译器来加速了 GoAWK。本文将讨论它为何更快以及新解释器的工作原理。

更新:GoAWK 现已原生支持 CSV 文件

几年前我用 Go 编写了 GoAWK,一个 AWK 解释器,并写了一篇文章介绍它的工作原理、测试方法以及优化过程。

GoAWK 是一个很有趣的业余项目,并至少在一个颇具规模的开源项目——Benthos 流处理器中得到应用。它甚至帮我获得了现在在 Canonical 的工作。

GoAWK 之前使用的是树遍历解释器:执行代码块时,它会递归遍历已解析的语法树。这种方式非常简单,但速度并不快。我一直想换成字节码编译器加虚拟机解释器的方案,最近终于实现了。

我早期的一个编程项目是 Third,一个运行在 DOS 上的 Forth 编译器。包括 Third 在内的大多数 Forth 编译器都是简单的编译器,使用一种字节码形式——在 Forth 领域被称为线程代码(threaded code)。所以可以说我对虚拟机感兴趣已有 25 年了……这算得上是真正的极客,还是仅仅说明我老了?

为什么虚拟机比树遍历更快?

为什么编译成虚拟指令再用虚拟机执行,会比对语法树求值(“树遍历”)更快,这一点并不直观。

实际上前期的工作反而更多了:原来只需词法分析和解析生成语法树,现在还多了一个编译步骤。不过,虚拟机编译器(包括 GoAWK 的)通常都非常简单且不做优化,所以这一步很快。

执行更快的原因之一在于:RAM——即随机存取存储器——在现代处理器上其实并非真正的“随机访问”。内存块会按需加载到高速的 CPU 缓存中,因此当需要访问新的内存块时,耗时约为已在缓存中时的 10 倍。Peter Norvig 关于典型 CPU 上各种操作耗时的表格显示,从一级缓存取数据大约需要 0.5 纳秒,从二级缓存取则要 14 倍时间,而从主内存取又要再多 14 倍!

以此为前提进行编程被称为“面向数据的设计”。当我观看 Andrew Kelley 的精彩演讲 《应用面向数据设计的实践指南》时,再次体会到这种设计带来的巨大影响。Andrew 是 Zig 编程语言的创造者,他在演讲中讲述了如何通过应用面向数据的设计技术显著提升 Zig 编译器的速度。正是这次演讲促使我开始在 GoAWK 上思考这个问题。但还是回到为什么虚拟机比树遍历更快……

语法树是由大量指向其他节点的节点结构组成的。它们分散在内存各处,因此要求值子节点时,必须沿着指针在内存中来回跳转,可能会把缓存中已有的数据挤掉。

下图是 GoAWK 中表达式 print $1+$2 的语法树,每个节点名称上方显示了其十六进制内存地址:

'print $1+$2' 的语法树

PrintStmtBinaryExpr 仅相距 48 字节,但左侧的 FieldExpr 却与之相距 8KB,而它的 NumExpr 子节点又与 FieldExpr 相距近 120KB。缓存块通常为 64 字节,因此其中每一个很可能都需要从主内存额外加载一个缓存块。显然不够缓存友好。

使用虚拟机解释器时,指令位于一个整齐的线性操作码(指令编号)数组中,很可能会一次性全部加载到一个缓存块中。在内存中的跳转要少得多。下面是同一个程序对应的 GoAWK 虚拟机指令(你可以用新增的调试选项 goawk -da 看到这种“汇编清单”):

$ echo 3 4 | goawk -da '{ print $1+$2 }'
        // { body }
0000    FieldInt 1
0002    FieldInt 2
0004    Add
0005    Print 1    // 1 is the number of values to print

7

GoAWK 编译器所做的(为数不多的)优化之一就体现在这里:当 i 为整数常量时,它会将 $i 编译为单条 FieldInt i 指令,而不是先 Num iField 的两条指令序列。这意味着大多数情况下,字段查找只需经过一次操作码分发循环,而不是两次。

虚拟机方式更快的另一个原因是函数调用更少,而函数调用相对较慢。在对语法树求值时,eval 函数会递归地再次调用 eval 来求值子节点。而使用虚拟机时,这一切都被扁平化为一个操作码数组,只需循环遍历即可——分发操作码时不需要任何函数调用。

编译器与虚拟机细节

GoAWK 的虚拟机使用 32 位操作码。起初我打算使用 8 位操作码(“字节码”中的“字节”即来源于此),但 32 位操作码速度一样快,而且使用 32 位操作码就无需可变长度的跳转偏移:较大的 AWK 脚本可能需要超过 -128 到 +127 的跳转偏移,而没有人会需要超过 32 位所能提供的二十亿的跳转偏移。64 位操作码则不必要地过大,而且也稍慢一些。

以下是前 10 个操作码(总共有 85 个——完整列表可在 internal/compiler/opcodes.go 中查看):

// Opcode represents a single virtual machine instruction (or argument).
// The comments beside each opcode show any arguments that instruction
// consumes.
type Opcode int32

const (
    Nop Opcode = iota

    // Stack operations
    Num // numIndex
    Str // strIndex
    Dupe
    Drop
    Swap

    // Fetch a field, variable, or array item
    Field
    FieldInt    // index
    Global      // index
    Local       // index
    ...
)

正如你在上面 print $1+$2 的汇编清单中所见,我使用的是基于栈的虚拟机。这更易于实现,因为编译器无需考虑如何分配寄存器,只需压栈和弹栈即可。不过,基于栈的虚拟机可能会稍慢一些——像 Lua 这样非常快的虚拟机则是基于寄存器的。

GoAWK 的编译器相当简单,基本是将语法树直接翻译为指令。我针对不同作用域的变量访问做了一些特化:例如,获取全局变量使用 Global 指令,获取局部变量则使用 Local。(如你所见,我的指令命名极具创造力。)

下面是一个对 1 到 10 求和的简单程序的汇编清单:

$ goawk -da 'BEGIN { for (i=1; i<=10; i++) sum += i; print sum }'
        // BEGIN
0000    Num 1 (0)
0002    AssignGlobal i
0004    Global i
0006    Num 10 (1)
0008    JumpGreater 0x0018
000a    Global i
000c    AugAssignGlobal AugOpAdd sum
000f    IncrGlobal 1 i
0012    Global i
0014    Num 10 (1)
0016    JumpLessOrEqual 0x000a
0018    Global sum
001a    Print 1

55

这里展示了一个我从 Python 抄来的小优化,Python 的解释器在 3.10 版本中加入了它(不过我相信这并不是什么新点子)。要编译 forwhile 循环,最简单的做法是在循环顶部做一次判断,然后在循环底部使用无条件 Jump。但这意味着每次循环都要执行两条跳转指令:顶部一条,底部一条。

取而代之,我们将条件编译两次:一次是在循环前取反(JumpGreater),一次是在循环底部(JumpLessOrEqual)。总体代码会稍多一些,因为条件被重复了,但循环本身——也就是关键部分——少了一条跳转指令。

我们几乎肯定可以进一步改进指令集,或许在事先知道操作类型时,为整数或字符串添加专门的指令。但这会增加复杂性,目前我还是想保持简单。

GoAWK 编译器所做的另一个优化是针对赋值的。AWK 中的赋值是表达式,因此默认情况下会将其值压入栈中……而在大多数情况下又会立刻丢弃。你很少会用到赋值表达式的值。

下面是经过优化的赋值表达式的汇编:

$ ./goawk -da 'BEGIN { x=42; print x }'
        // BEGIN
0000    Num 42 (0)
0002    AssignGlobal x
0004    Global x
0006    Print 1

42

而如果没有这个优化,它看起来会是这样:

0000    Num 42 (0)
0002    Dupe              # unnecessary
0003    AssignGlobal x
0005    Drop              # unnecessary
0006    Global x
0008    Print 1

下面是编译语句的代码,展示了用于此优化的特殊处理。我还包含了我们如何编译 if 语句的示例,以展示一种截然不同的情况。请注意编译器如何大量使用了 Go 的类型分支:

func (c *compiler) stmt(stmt ast.Stmt) {
    switch s := stmt.(type) {
    case *ast.ExprStmt:
        // Optimize assignment expressions to avoid extra Dupe and Drop
        switch expr := s.Expr.(type) {
        case *ast.AssignExpr:
            c.expr(expr.Right)
            c.assign(expr.Left)
            return

        case *ast.IncrExpr:
            ... // similar optimization for i++ and i--

        case *ast.AugAssignExpr:
            ... // similar optimization for i+=2 (for example)
        }

        // Non-optimized ExprStmt: push value and then drop it
        c.expr(s.Expr)
        c.add(Drop)

    ...

    case *ast.IfStmt:
        if len(s.Else) == 0 {
            jumpOp := c.condition(s.Cond, true)
            ifMark := c.jumpForward(jumpOp)
            c.stmts(s.Body)
            c.patchForward(ifMark)
        } else {
            jumpOp := c.condition(s.Cond, true)
            ifMark := c.jumpForward(jumpOp)
            c.stmts(s.Body)
            elseMark := c.jumpForward(Jump)
            c.patchForward(ifMark)
            c.stmts(s.Else)
            c.patchForward(elseMark)
        }

    ...
    }
}

虚拟机的 execute 函数是一个带大 switch 语句的单层 for 循环——每个操作码对应一个 case。下面是展示指令获取以及处理几个操作码的代码片段:

func (p *interp) execute(code []compiler.Opcode) error {
    for ip := 0; ip < len(code); {
        op := code[ip]
        ip++

        switch op {
        case compiler.Num:
            index := code[ip]
            ip++
            p.push(num(p.nums[index]))

        case compiler.Str:
            index := code[ip]
            ip++
            p.push(str(p.strs[index]))

        case compiler.Dupe:
            v := p.peekTop()
            p.push(v)

        ...

        case compiler.FieldInt:
            index := code[ip]
            ip++
            v, err := p.getField(int(index))
            if err != nil {
                return err
            }
            p.push(v)

        ...
        }
    }
}

Go 的 switch 语句

如上所示,虚拟机被实现为一个大的 switch 语句,每个操作码一个 case(大约 80 个 case)。Go 的 switch 语句目前被实现为在“case 空间”上进行二分搜索。你可以将其理解为编译成了类似这样的代码——为简洁起见,只完整展示了树的少数几个分支:

if op < 40 {
    if op < 20 {
        if op < 10 {
            if op < 5 {
                if op < 2 {
                    if op < 1 {
                        // handle opcode 0
                    } else {
                        // handle opcode 1
                    }
                } else {
                    // cases for opcodes 2-4
                }
            } else {
                // cases for opcodes 5-9
            }
        } else {
            // cases for opcodes 10-19
        }
    } else {
        // cases for opcodes 20-39
    }
} else {
    if op < 60 {
        // cases for opcodes 40-59
    } else {
        // cases for opcodes 60-79
    }
}

如你所见,要定位到感兴趣的 case,需要进行 O(log2 N) 次比较和跳转。对于 80 个操作码来说,解码每条指令就需要 6 到 7 次分支。

随着指令数量的增加,分支数量也会增加(不过好在这种增长是对数级的,而非线性)。当我最初为 GoAWK 编写虚拟机概念验证并仅实现演示所需的 7、8 条指令时,它带来了巨大的性能提升,速度快了近 40%,因为 switch 只有少数几个 case。但现在所有操作码都已就位,速度“仅”快了 18%。

实际上,当我有大约 100 个操作码时,速度比这还要慢。我移除了一些我曾以为能提速的特化,随着操作码减少,意味着二分搜索少了一层分支,平均速度反而快了 12%

如果能有一种方式实现常数时间的指令分发,无论有多少条指令都一样,那该多好。为什么 Go 不能将 switch 实现为一张跳转地址表:直接在表中查找代码地址并跳转过去?事实证明,Go 团队的 Keith Randall 正在做这件事,所以我们可能会在 Go 1.19 中看到它。

我在 GoAWK 上尝试了 Keith 的分支(目前仅支持 int64 类型),它将一个简单微基准测试的速度提升了 10%。所以我非常期待 Go 编译器学会“跳转表”。

我们自己能做这个优化吗?用一个函数数组怎么样?我尝试过,将分发循环改成了如下形式:

func (p *interp) execute(code []compiler.Opcode) error {
    for ip := 0; ip < len(code); {
        op := code[ip]
        ip++

        n, err := vmFuncs[op](p, code, ip)
        if err != nil {
            return err
        }
        ip += n
    }
    return nil
}

// Type of function called for each instruction. Each function returns
// the number of arguments the instruction read from code[ip:].
type vmFunc func(p *interp, code []compiler.Opcode, ip int) (int, error)

var vmFuncs [compiler.EndOpcode]vmFunc

func init() {
    vmFuncs = [compiler.EndOpcode]vmFunc{
        compiler.Nop: vmNop,
        compiler.Num: vmNum,
        compiler.Str: vmStr,
        ...
    }
}

func vmNop(p *interp, code []compiler.Opcode, ip int) (int, error) {
    return 0, nil
}

func vmNum(p *interp, code []compiler.Opcode, ip int) (int, error) {
    index := code[ip]
    p.push(num(p.nums[index]))
    return 1, nil
}

func vmStr(p *interp, code []compiler.Opcode, ip int) (int, error) {
    index := code[ip]
    p.push(str(p.strs[index]))
    return 1, nil
}

这在 GoAWK 的微基准测试上仅带来了 1-2% 的速度提升(查看结果和代码)。最终我决定还是坚持更简单的 switch 代码,转而寻找其他提速方法。而且当 Go 编译器支持 switch 的跳转表时,我什么都不用做就能获得 10% 的提升!

gcc 编译器有一个名为“计算跳转(computed goto)”的非标准特性,它允许你在每条操作码代码的末尾写类似 goto *dispatch_table[code[ip++]] 的语句,直接跳转到下一条操作码的代码。Eli Bendersky 写过一篇关于计算跳转的优秀文章,所以这里就不再赘述。大多数用 C 编写的虚拟机都使用了这一技术,包括 CPython 等众多项目。可惜 Go 没有计算跳转,不过话说回来,当 switch 被编译为跳转表时,也就实现了目标的一半。

如果你对编译器如何优化 switch 的更学术性内容感兴趣,可以阅读 Roger Sayle 的论文《多路分支代码生成中超优化器分析 [PDF]》,该论文发表于 2008 年 GCC 开发者峰会。

其他优化(以及一个负优化)

除了从树遍历改为虚拟机之外,我最近还加入了一些其他优化:

我做的一个有问题的修复是将 GoAWK 的字符串函数,如 length()substr(),改为使用 Unicode 字符索引而非字节索引。我知道这会将这些操作从 O(1) 变为 O(N)(N 为字符串长度),但我以为影响不会太大,因为“N 通常很小”。

结果证明这个假设并不成立:Volodymyr Gubarkov 的 gron.awk 脚本处理一个大型 JSON 文件的时间从 1 秒飙升至 8 分多钟——意外变成了二次方复杂度。这无法接受,所以我决定暂时回退该修复,未来再想办法以 O(1) 的方式解决这个问题。Gawk 的长期维护者 Arnold Robbins 评论了 Gawk 为实现高效字符串处理所做的大量工作。

我希望未来能进一步优化 GoAWK,并已开设了一个总览 issue 来跟踪未来的性能工作。以下是一些想法:

虚拟机改进。以下是我考虑用来加速虚拟机的几件事:

  • 优化或减少栈操作。interp.push 方法特别慢,原因是其中的 append 检查(而在正常的 AWK 代码中几乎从不需要 append)。如果你有关于如何提前确定最大栈大小的好主意,请告诉我。在可能存在递归函数调用的情况下,这是否可行?
  • 有没有可以添加的专门操作码,例如推送其参数中整数常量的 Int 指令?添加 Int 将省去一次对 interp.nums 切片的内存查找。
  • 想必 JumpLess 及类似的操作码很少在字符串上使用。是否最好将其替换为 JumpLessNum,以避免对至少一个操作数进行类型检查?(对于字符串,我们会使用更长的指令序列。)

字符串拼接也由于在拼接超过两个字符串时存在多余的分配和拷贝而开销过大。目前,像 first_name " " last_name 这样的多重拼接表达式会被编译为两条二元 Concat 指令:

Global first_name
Str " "
Concat
Global last_name
Concat

让编译器检测到这种情况并输出一条新的 Concat numArgs 指令会更高效,例如:

Global first_name
Str " "
Global last_name
Concat 3

这少了一条指令,但更重要的是,它将避免分配一个临时字符串,然后又不得不重新分配一个新字符串并拷贝字节。在 AWK 中拼接超过两个值的情况相当常见,而且拼接的值越多,这种优化的效果就越好。

正则表达式也非常需要加速。GoAWK 目前使用 Go 的 regexp 包,但不幸的是它相当慢。这使得大量使用正则表达式的 AWK 脚本速度约为 Gawk 的一半,几乎只有 Mawk 的四分之一。

有两种改进方法:

  1. 自己编写正则表达式引擎(可能会尝试将 Mawk 的引擎直接移植到 Go)。这可能是大量工作,而且由于 Go 的边界检查和较少的编译器优化,可能仍然不会很快。
  2. 提升 Go 正则表达式引擎的速度。这无疑是更好的方式,因为所有使用 Go regexp 包的人都会受益。这也可能相当困难。我可能会把这留给比我更聪明的人——也许这些 issue 中的一些会随着时间得到修复。

虚拟机结果

那么虚拟机解释器到底快了多少?微基准测试——诚然,其中大多并非你会在 AWK 中编写的那类脚本——总体快了约 18%。这些是耗时,因此越小越好(你可以看到原始结果,或使用 benchmark.sh 自行测量,然后用 benchstat.sh 来展示这些差异):

name                    old time/op  new time/op  delta
NativeFunc-8            10.7µs ± 0%  10.8µs ± 0%   +0.67%
BuiltinGsub-8           16.2µs ± 0%  16.2µs ± 0%   +0.36%
BuiltinGsubAmpersand-8  16.2µs ± 0%  16.2µs ± 0%   +0.29%
BuiltinSub-8            13.6µs ± 0%  13.6µs ± 0%     ~   
BuiltinSubAmpersand-8   13.5µs ± 0%  13.6µs ± 0%     ~   
SimplePattern-8          133ns ± 1%   134ns ± 0%     ~   
ConcatLarge-8           8.43ms ± 1%  8.35ms ± 2%     ~   
BuiltinSplitRegex-8     87.9µs ± 0%  87.7µs ± 0%   -0.21%
BuiltinSplitSpace-8     35.4µs ± 0%  35.1µs ± 0%   -0.70%
GetField-8               445ns ± 1%   435ns ± 2%   -2.42%
FuncCall-8              2.84µs ± 0%  2.76µs ± 2%   -2.65%
BuiltinSprintf-8        9.67µs ± 0%  9.23µs ± 0%   -4.58%
RecursiveFunc-8         15.7µs ± 0%  14.9µs ± 0%   -4.95%
ConcatSmall-8            735ns ± 0%   691ns ± 1%   -5.98%
BuiltinMatch-8          2.91µs ± 0%  2.71µs ± 1%   -7.02%
BuiltinIndex-8          1.23µs ± 1%  1.11µs ± 1%   -9.51%
RegexMatch-8            1.24µs ± 1%  1.11µs ± 4%  -10.07%
SetField-8               905ns ± 0%   810ns ± 0%  -10.45%
ForInLoop-8             2.04µs ± 2%  1.78µs ± 4%  -12.86%
ArrayOperations-8        657ns ± 0%   565ns ± 0%  -13.94%
BinaryOperators-8        493ns ± 0%   413ns ± 0%  -16.15%
BuiltinSubstr-8          975ns ± 0%   765ns ± 0%  -21.50%
Comparisons-8            417ns ± 0%   321ns ± 0%  -22.98%
SimpleBuiltins-8        1.00µs ± 0%  0.75µs ± 0%  -25.61%
CondExpr-8               203ns ± 0%   151ns ± 0%  -25.62%
BuiltinLength-8          607ns ± 0%   429ns ± 0%  -29.34%
IfStatement-8            219ns ± 0%   152ns ± 0%  -30.65%
AugAssign-8             1.50µs ± 0%  0.98µs ± 0%  -34.74%
LocalVars-8              479ns ± 0%   300ns ± 2%  -37.32%
Assign-8                 446ns ± 0%   261ns ± 0%  -41.55%
ForLoop-8               4.34µs ± 0%  2.50µs ± 0%  -42.39%
GlobalVars-8             468ns ± 0%   269ns ± 1%  -42.48%
IncrDecr-8               448ns ± 0%   148ns ± 0%  -66.87%
[Geo mean]              2.36µs       1.94µs       -17.90%

自增、自减和复合赋值的速度快得多,是因为虚拟机为它们提供了专门的操作码。变量访问也有了相当大的改进,for 循环、if 语句、二元运算符以及许多其他基准测试亦是如此。

我更“真实”的基准测试套件——其中大部分取自原始 AWK 源码——总体快了 13%。在这个表格中,goawk 是新的虚拟机解释器,orig 是旧的树遍历版本。有些违反直觉的是,这里的数字表示比原始 awk 快多少倍,因此越大越好。

测试goawkorigawkgawkmawk
tt.01 (print)2.021.911.001.662.29
tt.02 (print NR NF)1.591.601.001.772.20
tt.02a (print length)1.541.561.001.732.05
tt.03 (sum length)1.321.271.003.851.83
tt.03a (sum field)1.291.261.004.081.79
tt.04 (printf fields)0.970.801.001.262.74
tt.05 (concat fields)0.950.881.001.612.26
tt.06 (count lengths)1.391.351.002.531.97
tt.07 (even fields)1.241.181.001.461.71
tt.08 (even lengths)1.971.981.001.132.70
tt.09 (regex starts with)2.132.131.002.415.01
tt.10 (regex ends with)0.340.341.001.453.40
tt.10a (regex ends with var)0.320.341.001.303.08
tt.11 (substr)2.202.131.001.153.68
tt.12 (update fields)1.211.191.001.701.78
tt.13 (array ops)3.142.621.003.115.92
tt.13a (array printf)2.251.801.001.864.85
tt.14 (function call)1.171.091.000.641.56
tt.15 (format lines)0.620.611.000.962.21
tt.16 (count words)1.471.211.001.272.12
tt.big (complex program)1.621.411.001.833.82
tt.x1 (mandelbrot)2.251.621.001.343.44
tt.x2 (sum loop)1.761.031.001.152.68
几何平均1.451.281.001.862.32

结论

我确实很喜欢这次获得的性能提升。虽然没有达到我期望的那么多,但 GoAWK 现在在许多 CPU 密集型操作上比 Gawk 更快,这已经很酷了。它仍然始终慢于追求极致性能的 Mawk。而对于 AWK 通常用于的字符串处理和正则表达式,GoAWK 仍有很大的改进空间。

老实说,我不确定这多出的 2500 行代码(对于一个包括测试在内仅有 15000 行代码的项目来说)是否值得。如果有工程经理来监督这个项目,我想他会提出质疑(“这对真实工作负载有帮助吗?”)。不过,GoAWK 始终是一个热情驱动的项目——我享受制作和分享它的过程,这对我来说就足够了。

我已经合并了编译器和虚拟机,并在 GoAWK v1.15.0 中发布。Go API 和 goawk 命令应该保持 100% 向后兼容。它已经通过我的解释器测试以及来自原始 AWK 和 Gawk 相关测试的充分检验,但如果你发现任何问题,请提交一个 issue

希望你喜欢这篇文章或从中学到了东西。如有任何反馈或想法,请随时与我联系。

本文章由 muse-spark-1.2-contributor 进行翻译

评论