使用 weval 免费的编译器

上个月,Chris Fallin来到东北大学编程研究实验室进行了一次演讲。他讲述了自己在一个名为weval的新项目上的工作,这是一个 WebAssembly 部分求值器(之后还帮助作者写了这篇文章)。

部分求值很巧妙。简而言之,就是获取一个现有程序,将其修改为将一些输入作为常量保存,然后让编译器/优化器对其进行优化。结果仍然是一个程序——而不是一个值——并且通常比原始程序更快。

通常的小例子是幂函数。如果有一个函数接受两个参数xy并返回x^y

int power(int x, int y) {
  int result = 1;
  for (int i = 0; i < y; i++) {
    result *= x;
  }
  return result;
}

如果在y = 5时对该函数进行部分求值,就会得到一个新函数,它接受一个参数x并返回x^5

int power_5(int x) {
  int result = 1;
  for (int i = 0; i < 5; i++) {
    result *= x;
  }
  return result;
}

现在,对你们来说,这看起来可能与原始函数没有太大区别。但对于优化器来说,这是一个新的机会世界。优化器可以展开循环并删除条件:

int power_5(int x) {
  return x * x * x * x * x;
}

weval 可以对整个 WebAssembly 模块进行这样的操作。WebAssembly 模块通常比一个小的power函数大得多。例如,如果你的 WebAssembly 模块是一个解释器,你可能会想要使用它。想象一下这样一个世界:你已经将诸如 SpiderMonkey 或 CPython 这样的运行时编译为 WebAssembly。然后你可以在 WebAssembly 运行时上运行你的 Python 或 JavaScript 程序,但它们会比直接编译为 WebAssembly 要慢。即使你直接将 JS/Python 编译为 Wasm,除非你的编译器进行了一些复杂的静态分析,否则它可能也会很慢。这就是 weval 的用武之地。

进入 weval

SpiderMonkey 和 CPython 都非常大。相反,我们将展示一个我和 Chris 编写的小解释器的演示。我们的解释器没有做太多事情——局部变量、累加器、算术和分支。但它足以展示 weval 带来的性能提升。

#define FOR_EACH_INSTRUCTION(V)                                                \
  V(LOAD_IMMEDIATE)                                                            \
  V(STORE_LOCAL)                                                               \
  V(LOAD_LOCAL)                                                                \
  V(PRINT)                                                                     \
  V(PRINTI)                                                                    \
  V(JMPNZ)                                                                     \
  V(INC)                                                                       \
  V(DEC)                                                                       \
  V(ADD)                                                                       \
  V(HALT)

它被设计为一个小循环,读取下一条指令并通过switch进行调度。这不是最快的设计1,但没关系。

uword Execute(uword *program) {
  //...
  while (true) {
    Instruction op = (Instruction)program[pc++];
    switch (op) {
    case LOAD_IMMEDIATE: {
      uword value = program[pc++];
      accumulator = (Object)value;
      break;
    }
    case STORE_LOCAL: {
      uword idx = program[pc++];
      LOCAL_AT_PUT(idx, accumulator);
      break;
    }
    case LOAD_LOCAL: {
      uword idx = program[pc++];
      accumulator = LOCAL_AT(idx);
      break;
    }
    case PRINT: {
      const char *msg = (const char *)program[pc++];
      printf("%s", msg);
      break;
    }
    case PRINTI: {
      printf("%" PRIu64, accumulator);
      break;
    }
    case HALT: {
      return accumulator;
    }
    case JMPNZ: {
      uword offset = program[pc++];
      if (accumulator!= 0) {
        pc = offset;
      }
      break;
    }
    case INC: {
      accumulator++;
      break;
    }
    case DEC: {
      accumulator--;
      break;
    }
    case ADD: {
      uword idx1 = program[pc++];
      uword idx2 = program[pc++];
      accumulator = LOCAL_AT(idx1) + LOCAL_AT(idx2);
      break;
    }
    //...
    }
  }
}

使用这种字节码,我们可以编写一个简单的程序,将从 1 到 1 亿的所有数字相加:

enum {
  result = 0,
  loopc = 1,
};
uword program[] = {
  // result = 0
  LOAD_IMMEDIATE, 0,
  STORE_LOCAL, result,
  // loopc = 100_000_000
  LOAD_IMMEDIATE, 100000000,
  STORE_LOCAL, loopc,

  // loop:
  // result += loopc
  ADD, result, loopc,
  STORE_LOCAL, result,
  // loopc--
  LOAD_LOCAL, loopc,
  DEC,
  STORE_LOCAL, loopc,
  // if loopc!= 0, jump to loop
  JMPNZ, 8,

  // print result
  PRINT, (uword)"Result: ",
  LOAD_LOCAL, result,
  PRINTI,
  PRINT, (uword)"\n",
  HALT,
};

我们可以使用任何 C 或 C++编译器编译这个解释器程序,将解释器喂入字节码,它将在大约 350 毫秒后打印结果:

$ c++ -O2 peval.cc -o peval.out
$./peval.out
Result: 5000000050000000
$

但是假设你想在 WebAssembly 中沙盒这个程序。幸运的是,有一个名为wasi-sdk的项目,它提供了针对 WebAssembly 的接近即插即用的 Clang 替代品。我们可以使用 wasi-sdk 编译解释器,并使用wasmtime或任何其他提供 WASI 垫片的 WebAssembly 运行时来运行它2。这大约需要 530 毫秒:

$ /opt/wasi-sdk/bin/clang++ -O2 peval.cc -o peval.normal.wasm
$ wasmtime peval.normal.wasm
Result: 5000000050000000
$

但实际上我们一直想要的是在沙盒中部署程序——而不是解释器本身,只是程序——我们可以通过使用 weval 将字节码和解释器组合在一起来实现这一点。这大约需要 40 毫秒:

$ /opt/wasi-sdk/bin/clang++ -O2 -DDO_WEVAL -I include peval.cc -o peval.wasm
$ weval weval -i peval.wasm -o peval.wevaled.wasm -w
$ wasmtime peval.wevaled.wasm
Result: 5000000050000000
$

首先,让我们退一步。我们有一个用 C++编写的解释器,运行需要 350 毫秒。通过将其编译为 WebAssembly,它变得稍微慢了一些(530 毫秒)。然后,通过使用 weval,我们获得了8.5 倍的加速。这太疯狂了。这可能接近我们如果为字节码机器手动编写一个小编译器所获得的加速,但我不需要编写编译器。

方法时间(毫秒)
C++350
WASI530
WASI+weval40 (!!)

如果你想感受这种差异,可以查看一个asciicast

你可能会注意到我在那里添加了一些偷偷摸摸的标志,如-DDO_WEVAL-I include。这是怎么回事?

专门化解释器

总体思路:给解释器函数提供对常量字节码的访问。

虽然 weval 在任何 WebAssembly 模块及其输入的组合上都能正常工作,但当你给它一些帮助并告诉它哪些数据是常量时,它的效果最好。为了做到这一点,我们使用一个名为wizer的项目预先初始化 WebAssembly 模块。它为程序员提供了在将运行状态转换回 WebAssembly 模块之前设置一些内存的钩子。让我们看一下当前情况的示意图:

现在,在运行时,解释器加载字节码并运行它。字节码事先是不知道的,所以解释器必须是通用的。

为了专门化解释器,我们在 WebAssembly 模块初始化时执行三个步骤:

  1. 加载字节码
  2. 创建一个具有常量参数的专门化解释器函数版本
  3. 在函数上运行常量传播和其他编译器传递

在这个例子中,我们知道字节码是常量的。我们可以通过使用它的一个辅助内部函数来告诉 weval 这一点。在这种情况下,我们创建一个带有常量参数(program)的Execute函数的副本。现在我们有两个函数:ExecuteExecuteSpecialized。所有这些都在init函数中发生:

template <bool IsSpecialized>
uword Execute(uword *program) {
    //...
}


Object (*ExecuteSpecialized)(uword *) = 0;

#ifdef DO_WEVAL
void init() {
  uword result = 0;
  uword loopc = 1;
  weval::weval(&ExecuteSpecialized, &Execute<true>, /*func_id=*/123,
               weval::SpecializeMemory<uword *>(program, sizeof program));
}

WIZER_INIT(init);
WEVAL_DEFINE_GLOBALS();
#endif

int main(int argc, char **argv) {
  if (ExecuteSpecialized) {
    ExecuteSpecialized(nullptr);
  } else {
    Execute<false>(program);
  }
}

现在代码已经被加载并标记为常量,情况看起来更像这样:

虽然代码是常量的,但 weval 不是魔法。默认情况下,它不会修改控制流,除了简化变成常量的分支(它甚至不知道什么是解释器!)。

为了开始使ExecuteSpecialized更快,我们必须在解释器中为 weval 提供一些提示让其获取。我们实际上想要专门化控制流——使字节码中的控制流成为新函数本身的控制流——所以我们告诉 weval 关于 PC 的信息,让它展开代码。

修改解释器

总体思路:通过在程序计数器上进行专门化来展开循环。

我们可以首先告诉 weval 用哪个变量作为专门化上下文。在这种情况下,因为我们知道字节码是常量的,所以我们可以在pc——程序计数器上进行专门化。这让 weval 完全展开解释器循环。

uword Execute(uword *program) {
  while (true) {
    //...
    switch (op) {
      //...
    }
    weval::update_context(pc);
  }
}

在对捆绑模块运行 weval 并让 weval 展开循环后,情况看起来像这样:

这意味着从这一点开始,我们使用 weval 将我们的解释器变成了一个编译器。到目前为止只有两个优化——展开循环和常量传播——但它们非常有效。结果是一个完全的 WebAssembly 模块,没有解释器,也没有字节码。

Weval 的编译器传递不是魔法。它们是任何编译器在你的代码上都会运行的相同传递。它们可以展开解释器循环并将字节码转换为直线 WebAssembly 代码。但该代码仍然有局部变量写入推送和局部变量读取以及解释器的所有其他开销。所以还有更多的工作要做……

但如果我们更多地修改它呢?

总体思路:通过告诉 weval 局部变量的位置,将解释器局部变量展开为 WebAssembly 局部变量。

在编译器中处理内存可能很困难。weval 不是一个整个程序的优化编译器,可能无法证明一个内存位置(在这种情况下,是局部变量数组)永远不会逃逸或与其他东西别名。但我们,作为解释器的作者,知道这一点。所以我们可以添加更多的提示。

现在,LOCAL_ATLOCAL_AT_PUT是读取和写入局部变量数组的宏:

#define LOCAL_AT(idx) (locals[idx])
#define LOCAL_AT_PUT(idx, val) (locals[idx] = val)

这对解释器来说很好,但对编译后的代码来说不是很好。我们真正想要的是给 weval 能力,将每个内存位置——每个局部索引——分别作为一个 SSA 值进行推理。

为了做到这一点,我们使用 weval 内部函数:weval_read_regweval_write_reg。为了获得最大的灵活性,我们有几个切换这两个的宏:

#ifdef DO_WEVAL
#define LOCAL_AT(idx) (IsSpecialized? weval_read_reg(idx) : locals[idx])
#define LOCAL_AT_PUT(idx, val)                                                 \
  if (IsSpecialized) {                                                         \
    weval_write_reg(idx, val);                                                 \
  } else {                                                                     \
    locals[idx] = val;                                                         \
  }
#else
#define LOCAL_AT(idx) (locals[idx])
#define LOCAL_AT_PUT(idx, val) (locals[idx] = val)
#endif

现在,weval 可以分别对每个局部变量进行推理,它们最终将被编译为普通的 WebAssembly 局部变量。

总结

这里的主要想法是通过在常量数据上进行专门化,然后进行正常的编译器传递,有可能逐步将一个解释器分解为一个编译器。在构建时能够专门化的越多,生成的代码就会越快。

可以查看作者在已过时的 weval 分支中的代码。它还包括不同 JS 运行时中 Wasm JIT 的意外基准测试!

展望未来

Chris 在这次演讲(PDF)中更详细地介绍了 weval 的工作原理,包括解释器函数如何与字节码实际组合的描述。主要思想是在上下文敏感的数据流分析中使用 PC 值作为“上下文”,这样常规的常量传播将看到 PC 值和一个解释器循环迭代的操作码,而不是它们的并集(就像静态分析通常会做的那样)。有很多繁琐的细节需要使其工作良好,Chris 还计划不久后写一篇关于 weval 及其在 SpiderMonkey 中的应用的博客文章。

此外,我们的解释器很小,本身并不太有趣。它只是用于向你解释一些 weval 概念。但同样的原则也适用于更大的解释器!有 SpiderMonkey,是的,同样的事情也可能适用于 CPython,主要的 Python 运行时。CPython 已经支持被编译为 WebAssembly 了!

想象一下直接将 Python 编译为 WebAssembly……也许很快就会实现?

更疯狂的想法

CPython 已经支持 vectorcall 函数指针。这是一种以可移植的方式添加 JIT 编译器的方法。我们也可以使用它将 weval 变成 CPython 的 Wasm JIT。

类似的项目

BuildIt是一个类似的 C++项目,采用库的方法。

  1. 许多解释器使用计算 goto 来调度指令。这比 switch 语句稍微快一些,但也不具有可移植性。它甚至不被所有编译器(例如,WASI)支持,因为它不是合法的 C。
  2. 这非常直接,意味着程序甚至可以在浏览器中运行。但这不是这篇文章的重点。请查看存储库以获取使其成为可能的小 JS 包装器。
阅读 7
0 条评论