当 wasm 遇上数据处理

文/ 阿里淘系 F(x) Team - 文赫

背景：JavaScript 中的哪些数学困境

本文主要讲述了，如何快速的将一些已经成熟的算法迁移到 JavaScript 并在生产环境中部署的一些心得和体会。而文中提到的能力也会合入我们的开源项目 -- Datacook 中。 后续，我们也会单独出一篇文章来介绍 Datacook 的前世今生和使用场景～

首先我们来看看， Datacook 到底是个啥？ Datacook 的定位是一个高性能的 Node.js 和浏览器通用的数据处理 / 特征工程 / 数据集制作的工具库。

> 可以参考上方的 README，我们可以通过上面的方法调用 datacook 中的各类方法。

后期也会嵌入到 ODPS 环境作为 JavaScript 进行数据处理的一个组成部分。同时，借助 Datacook，将数据处理链路迁移至 JavaScript 使我们可以在 JavaScript 完成完整的深度学习链路（数据处理 / 特征工程 / 模型训练 / 模型推理 / 模型部署）。

能不能在 JavaScript 中实现一个 Beta 分布的函数？

这几天收到同学提出的一个问题 “能不能在 JavaScript 中实现一个 Beta 分布（Beta Distribution -- 下文简称 Beta）的函数？”。 刚接手这个问题的时候很疑惑，这种统计学常用的分布按理说在 Python 的中都有实现，为啥不能通过直接调用或者使用 Boa 来完成这个能力？究其原因，是因为这个同学想把功能部署在 Serverless 服务上，这就对容器用了较高要求导致直接使用 Python 或者使用 Boa 变得十分困难。 那么，为了让这位同学能够成功的把 Beta 分布部署在 Serverless 上，我开始了复习统计的旅程（大雾 🙃🙃🙃

第一次尝试： JavaScript

去Wiki上简单的查了一下 Beta 分布的定义，发现只要满足下面的公式产生的随机数就可以算是 Beta 分布： 其中 x 是随机数， 。 看到这么简单的公式，我寻思好家伙，这不是10分钟就能整出来的一个函数（然而这却是噩梦的开始） 看到公式的几十分钟后，便有了如下的代码：

const gamma = (n) => {
    let res = 1;
    for (let idx = 1; idx < n-1; idx++) res *= idx;
    return res;
}

const beta = (a, b, x) => {
    const B = gamma(a)*gamma(b)/gamma(a+b);
    return Math.pow(x, a-1)*Math.pow(1-x, b-1)/B;
}

随便试了几个数，okay 符合预期，可以考虑封装封装拿到 Serverless 上部署了。然而，就在这个时候，忽然得知 a 和 b 的取值可能上万。那这，不用试也知道，gamma 函数会溢出到 NaN 导致不可用的；而且那样的运算对资源的消耗也是巨大的，需要想办法去优化算法了。

第二次尝试：tfjs + 数学优化

对于阶乘入参过大导致的溢出，常见的优化便是取对数，即：

对于我们的情况：

同时，为了加速运算和保证精确度，我们使用了 tfjs 来实现数学运算逻辑。

基于上面的公式，我们就可以来实现我们的第二版 Beta 函数了～ （代码因为太长了就不在这边展示了） 然而，因为输入的原因，哪怕使用数学公式上的变换，对于较大或者较小的值来说，结果还是会出现下溢 -- 导致结果不准。

第三次尝试：不要重复造轮子

经过了前两次冲击，深刻意识到了 Don't Repeat Yourself(DRY, 不要重复造轮子)，这个法则的含义。我开始在 Github 上寻找 Beta 分布的实现。然而，这些实现都大多数实现都是基于 python 和 c++ 的。 但是不管 python 还是 c++ 的实现都依赖了大量的第三方运算库和高级数学函数（这俩都是 js 所缺失的）。就在这时忽然发现有很多 c++ 的实现都只依赖了 stdlib 内的 random 函数就可以完成，忽然想到是不是可以通过 WebAssembly 的方法把这些实现直接编译到 JS 运行时使用？

Emscripten

Emscripten 是一个基于 llvm 的用来将 c/c++ 代码编译到 WASM 的工具链。具体的细节可以参考官方文档。这边我们就不再赘述了～ 在本篇中，我们将使用这个工具链来帮助我们快速的生成所需的 WASM 字节码和对应的胶水代码。

C++

我们先来看一段 c++ 的 Beta 函数实现：

#include <random>
#include <stdlib.h>

// refer to https://gist.github.com/sftrabbit/5068941

class beta_distribution
{
  public:
    class param_type
    {
      public:
        typedef beta_distribution distribution_type;

        explicit param_type(double a = 2.0, double b = 2.0)
          : a_param(a), b_param(b) { }

        double a() const { return a_param; }
        double b() const { return b_param; }
  
      private:
        double a_param, b_param;
    };

    explicit beta_distribution(double a = 2.0, double b = 2.0)
      : a_gamma(a), b_gamma(b) { 
        rng.seed(rand());
      }

    void param(double a, double b)
    {
      a_gamma = gamma_dist_type(a);
      b_gamma = gamma_dist_type(b);
    }

    void seed(uint32_t seed) {
      rng.seed(seed);
    }

    double generate() {
      return generate_internal(a_gamma, b_gamma);
    }

    double min() const { return 0.0; }
    double max() const { return 1.0; }

    double a() const { return a_gamma.alpha(); }
    double b() const { return b_gamma.alpha(); }

  private:
    typedef std::gamma_distribution<double> gamma_dist_type;

    std::mt19937 rng;

    gamma_dist_type a_gamma, b_gamma;

    double generate_internal(
      gamma_dist_type& x_gamma,
      gamma_dist_type& y_gamma)
    { 
      double x = x_gamma(rng);
      return x / (x + y_gamma(rng));
    }
};

上面的代码便是基于这个 gist 修改得到的。其中核心主要是通过 stdlib 中的 gamma_distribution 来完成高精度的 gamma 函数计算继而实现高精度的 beta 函数。其次就是随机数生成的逻辑，因为 js 的随机数的限制（无法设置 seed / 随机长度不够）我们也使用了 stdlib 中实现的 rand.h 来完成随机数生成。 上面的代码，使我们可以仅通过 stdlib 就生成 Beta 分布的能力。

Emscripten 💗 C++

有了 C++ 的逻辑实现，我们需要仅需要通过 Emscripten 将 Beta 这个类暴露出去就好了。 在之前的代码基础上，我们只需要：

#include <emscripten/bind.h>

class beta_distribution {
 ....    
}

// Binding code
  EMSCRIPTEN_BINDINGS(Beta) {
  class_<beta_distribution>("Beta")
    .constructor<double, double>()
    .constructor<>()
    .function("generate", &beta_distribution::generate)
    .function("setParam", &beta_distribution::param)
    .function("setSeed", &beta_distribution::seed)
    ;
}

通过上面的代码，我们告诉 Emscripten 将 我们定义的 beta_distribution 这个 class 以 Beta 这个名字暴露给 Javascript。并实现两个 constructor：入参为两个 double 或者入参为空。 除此之外，我们也将beta_distribution 实现好的三个函数分布注册到 generate , setParam 和 setSeed  JavaScript 的函数空间中。

编译 & 使用

有了上面的源码以及 Emscripten 的工具链，我们就可以把上面的实现编译到 WASM 了～ 首先，我们先来看看编译参数：

$ emcc 
	--bind 
  -O3 
  -s ALLOW_MEMORY_GROWTH=1 
  -s MODULARIZE=1  
  -o ./beta.js 
  ./beta.cpp

上面的命令有几个部分：

• emcc 就是 Emscripten 版本的 gcc 是我们编译的入口
• --bind 是通知编译器我们需要把 c++ class bind 到 JS 上
• -O3 意味着最高级的代码优化
• -s ALLOW_MEMORY_GROWTH=1 通知编译器内存可以增长
• -s MODULARIZE=1 将编译产物模块化 （这一步是为了防止作用域的污染）

经过上面的编译，我们就可以得到一个 ./beta.js 的胶水代码，其中包含了如何加载 WASM 文件以及将 c++ 的函数/类定义映射到 JS 空间的逻辑；以及一个 ./beta.wasm 的 ByteCode 这里面就是我们具体的 Beta 函数的逻辑了。

接下来我们就可以调用这个 WASM 模块了：

const betaModule = require('./beta')
const Beta = (await new betaModule()).Beta;
const beta = new Beta(1, .5)
console.log(beta.generate())

到此为止，我们就通过 WASM + c++ 的方式将一个较为复杂的数学函数轻松的移植到了 JavaScript 的环境中～ 而且这段 WASM 代码，也在 3.8 大促期间作为 serverless 服务出现。

Review 一下

性能

我们先来看一下 wasm 和 c++ native 的运行性能，我们将在 loop 中循环调用我们的实现来计算运算耗时。 我们会惊奇的发现，在这个场景中， wasm 的运行效率随着调用次数的增加而显著优于 c++。这方面可能的原因是，v8 虚拟机对循环和 cache 的优化。初次之外，c++ 会调用系统的 RNG （Random Number Genertor）这部分作为 system call 的耗时是显著的；而 WASM 并不会通过 system call 来完成这部分逻辑。（对于性能这块欢迎来讨论/分析可能的原因。）

GC

在部署到生产环境的时候，我们也不是一帆风顺的。由于 WASM 作为一个运行在 JavaScript 中的虚拟机，WASM 自己生成出来的东西没有办法被自动回收，作为调用者也需要注意内存的回收。 在线上运行中，该实例的内存占用稳定在了很低的一个值（0.01%）上面。

总结

总结一下这一次的开发体验，借助 Emscripten 的工具链，我们可以快速的把一些在 c/c++ 生态上已有的生态迁移到 JavaScript。总体开发体验很棒，而且产物稍加修改就可以在生产环境中运行；同时 Emscripten 也提供了 asm.js 作为无法运行 WASM 环境的 fallback 解决方案。

但是，不得不说，Emscripten 还是太重了，生成了很多胶水代码。这一部分都是后期可以持续优化的地方；同时，Emscripten 的代码模型，对于多个平行的 WASM 项目还是不太友好（依赖全局变量，不模块化容易互相污染，模块化增加复杂度）。在工程侧来讲，作为同时用较重的 JavaScript 逻辑和 WASM 逻辑的项目，打包 / 构建 / 测试链路都不是很完善。感觉上面的方向都可以搞大事情。

最后打一个广告，本次实现的 Beta 函数会作为 datacook 的一个函数导出供社区使用。同时接下来也会对 Boost 库中常用的数学方法导出到 WASM。也欢迎有兴趣的同学一起来共建学习。

---