浅谈 eDSL 在科学计算和数据分析领域的发展趋势

本文来自 “在科学计算和数据分析领域中，今后 Python、R、Julia 各自会有什么样的发展趋势？- 罗秀哲的回答 - 知乎”，经答主授权后由编程语言 Lab 整理并发出。

首先需要提的一点是 Python 的通用性和 Julia 的高性能在下面文章的语境下都是错误的，因为 Python 的 eDSL 大多无法组合，Julia 反而可以有限地组合一些。Julia 也并不是任何程序都性能好，所以这个题目的描述不是非常好。我主要回答的是题目本身。

在科学计算这个领域里提高嵌入式专用语言（eDSL) 相关的功能是这些通用语言在最近几年里的主要方向之一，也是除了具体领域的进展以外对所有人都有用的一个方向。

科学计算领域与 eDSL

首先解释一下什么是专用语言（Domain Specific Language，后简称 DSL），专用语言一般是为某个特定领域设计的，这样可以提供更多的语法糖（syntax sugar）来提高编程效率和可读性，并且由于约束了语义，可以方便优化器进行更有针对性的优化。

随着我们进入后摩尔时代，异构计算（heterogeneous computing）和专用计算（例如各种最近很火的 quantum programming，probabilistic programming 等等）会越来越普遍，这样的结果就是我们将会有越来越多专用语言发明来降低这些本身就不通用的计算模型的编写难度。

而嵌入式专用语言（embedded DSL，后简称 eDSL）是嵌入在某个通用语言的里的专用语言，这样做的好处是用户可以很容易的去组合专用语言写出的功能和通用性的功能，并且由于一定程度上复用了宿主语言（host language）的语法，用户可以几乎不需要学习新语言。此外，由于是嵌入式的，某些专用语言的语义（例如自动微分和一些异构计算硬件）可以缓慢地覆盖宿主语言从而最终实现对整个宿主语言的支持。

改进现有的 eDSL 和相关基础设施是为了有一天可以让大家不用去学习专门的 DSL 语法，而通过编写更加符合直觉的宿主语言，然后通过改变编译的上下文和编译目标来实现对更多硬件，更多计算方式的支持。而不是在宿主语言里硬造出不 Pythonic 或者不 Julian 的语法来，在 Python 里使用大量的 decorator 或者在 Julia 里使用大量的宏是不提倡的，但是很多时候为了支持某些功能是无法避免的，我们当然希望某天一天能够让编译器更加智能。所以，提高 eDSL 的功能不代表需要用户去学习新语言，而是相反，是为了不让用户学习新语言。

此外下文里描述的各种提高不同 eDSL 甚至是不同语言之间的可组合性的努力也是为了在未来可以让编程语言也更加模块化，例如负责概率编程的语言可以和其它语言自由组合，编译器依然能理解他们组装起来以后的程序是什么意思，并提供优化和错误提示。这样反而避免了发明新语言。

下面以 Python 和 Julia 为例简单谈谈我的看法。

Python 的 eDSL 发展

在 Python 里，这两年比较相关的是围绕 Jax [1] 的生态系统和 PyTorch 的 torchscript [2] 以及早年的 numba [3]，当然还有之前知乎上很火热的 taichi [4] 语言，他们实际上都可以算作是 DSL。

Python 由于其本身的性能问题，通过单独编译某个子集一直是常见的提升性能的方案，而这也使得出现了大量的以 Python 作为宿主语言的 DSL。

Python 里的 DSL 要有两种：

伪装成 class 的专用语言

例如 tensorflow 和 theano，他们重载运算符来实现 tracing，然后再编译（也就是我们常说的静态图方案），这样失去了动态性和 Python 本身的反射，随着功能越来愈多会让调试和开发不太方便。并且无法获得宿主的控制流信息。但是好处是由于使用的是 Python 原生的 class，有一定的可组合性。

直接通过 ast 模块拿走 Python 的语法树并自己编译的 DSL

其中比较有名的当属 Google 的 Jax，PyTorch 的 torchscript，还有 TVM，以及 numba 等，他们分别在 Python 的子集上实现了针对科学计算（或者说机器学习）的可微分编程方案。

尽管他们看起来像是 Python，但是实际上当你使用装饰器（decorator）对相应的函数标记之后，Python 的语法树就会被编译到相应的 DSL 的中间表示上（intermediate representation）。所谓的中间表示就是一种中间语言或者中间数据结构来方便编译器进行分析。比较常用的中间表示一般是 SSA 格式的（Static Single Assignment），在这个表示里所有的变量有且仅有一次赋值，大大简化了数据流分析（data flow analysis)。

这个时候代码会进入到对应的编译器实现里，而不再由编译器运行了，当编译器进行优化并且在内存里产生本地代码（native code）之后，再将这个编译好的二进制代码链接回 Python 包装成一个看起来像是 Python 函数的对象。

但是问题也随之而来，早期大家开发的时候都是各写各的，Jax 的编译器看不懂 PyTorch 的代码，Numba 看不懂 Jax 写的代码，哪怕他们可能都是同一个实现。比如我发现 Jax 里没有复数 SVD（举例，其实是有的），但是 numba 里有，我想在 Jax 里调用 numba 的 SVD，那么对 Jax 来说这就变成了一个无法理解的外部函数调用（foreign call），那么导致的结果就是 Jax 本来可以对这段代码进行分析，比如说进行内联（inline）或者说产生后向模式（reverse mode）的导数（gradient），甚至是编译到 TPU，但是由于是一段用其它中间表示的代码（LLVM IR），导致无法进行分析。这样导致的结果就是每个 Python 的大工程都做了大量的重复劳动。

类似于著名的两语言问题（two language problem），我们可以称之为两中间表示问题（two intermediate representation problem），也就是一个中间表示上写的分析无法给另外一个类似的中间表示使用。

MLIR 的出现

这个问题最早在 TensorFlow 里出现，于是 Google 给出了 MLIR [5] 的方案（multi-layer intermediate representation），虽然一开始的目标类似 Halide 和 TVM，是为了优化机器学习任务并编译到异构硬件上设计的，但是后来随着 MLIR 的发展它可以被用来做更多的事情。

MLIR 最重要的功能是 progressive lowering（这个不知道怎么翻译了）和自定义的方言（dialect）。progressive lowering 的意思是 MLIR 里的每一种方言可以选择性地 lowering 到其它的表示上，这样我们就可以提高不同的指令（intrinsic）的可组合性。而对每一种其它的表示我们也都可以去实现一个对应的 MLIR 方言，这样使得我们最终可以把他们编译到同一个格式下从而让不同的专用语言变得能够互相组装。当然实现细节上如果编译本身是单向的（比如从一个 array IR 编译到 LLVM）这个过程中丢掉了 array shape 的信息，那么也不一定能够组装起来，但是这是另外的问题了。

最近在 MLIR 社区里出现了很多构建从 Python 编译到 MLIR 的基础设施或者某个领域的应用，我觉得是很值得期待的（例如 numba 开始支持 MLIR 了 [6]）。我们姑且可以将其算作是未来的第三种 Python eDSL，假如在不远的将来某种通用的基础设施出现了的话（可是 MLIR 的 C API 暂时还不完善）。

Julia 的 eDSL 发展

由于对宏的支持比较接近 lisp，并且能够在一定程度上做有限的 stage programming，Julia 里的 eDSL 发展历史实际上更久一些。

Julia 里目前主要有三种 eDSL 的实现：

运算符重载

这包括 Julia 的 ForwardDiff [7]，早期的机器学习框架 Flux [8]，Julia 早期对 TPU 的支持 [9]，以及量子计算框架 Yao [10] 等。与 Python 的第一种方案类似，他们本身不会创造任何新的语法，所有的编译工作都是通过运算符（函数）重载之后再进行，并且有一定的可组合性。缺点是完全无法获得宿主语言的控制流（control flow）信息。

基于宏

这其中包括最早的优化框架 JuMP [11]，SciML 的 ModelingToolkit [12]，还有 Soss [13]，Turing [14] 以及 Gen [15] 等概率编程语言。

基于宏的 eDSL 由于工作在语法树上，缺乏类型信息，所以设计者往往会通过增加新的关键字来获取相关的信息，于是这一类 eDSL 的好坏往往非常取决于设计者的编程水平和品味。水平比较高的开发者知道怎么直接在语法树上做数据流分析（甚至自己做类型推导），或者知道应该砍掉哪些功能，往往可以做到只需要一个宏标记一下 eDSL 的入口即可，中等水平的开发者可能还是会需要增加几个宏作为关键字来获取信息，而水平不足的开发者或者支持的功能过多则有可能会做出满屏宏的 eDSL。

所以会出现 Turing 或者 Gen 开发的 eDSL 还没有简单的运算符重载好用的情况。实际上当初在我们开发 Yao 的时候也讨论过宏的方案，后来被 Leo（Yao 的另一位作者刘金国博士） 一票否决了。

社区里的元老级人物 Steven G Johnson [16]（也是著名的 FFTW [17] 的作者）有过一句名言：

如果你没有 Jeff（Julia 语言的创造者之一）聪明，那就别用宏。

基于生成函数和抽象解释器（abstract interpreter）框架

这套方案来自于 Julia 起初在对 CUDA 和自动微分进行支持的时候发展了生成函数用来在类型推导时期插入自定义的 SSA IR 生成，也就是 Julia 的 CodeInfo [18]。然后 Julia 在 1.6 版本中进一步增加了抽象解释器的接口，使得在这层 IR 上我们可以获得大部分类型信息，甚至可以自己有限地构造 type lattice，并且复用大量的 Julia 编译器的优化与分析的工具。

而起初为了支持 CUDA 而设计的 GPUCompiler [19] 在支持更多加速卡（AMD，Intel one 等）的过程里也变得越来越可复用，于是到了 2021 年我们已经可以利用这套设施来编译到任意 LLVM 支持的 target 了。例如目前最快的自动微分编译器 Enzyme [20] 就通过这种方式支持了 Julia。

这套方案的好处就是可以做到完全不改变表层语法，不需要额外的宏，可以支持几乎所有的 Julia 本身的语法，像 CUDA 只需要用 device 宏标记一下就可以在 GPU 上运行，或者 Zygote 连宏都不用，调用一下 gradient 函数就可以求导就是用了这一套流水线。当然 Zygote 本身还有很多问题，因为虽然它使用了生成函数和 SSA IR 但是它没有用到抽象解释器，它很快会被使用了全套编译工具的下一代自动微分引擎 Diffractor [21] 代替。

我们最近开启了一个新的社区合作：Julia 编译器插件组织 —— JuliaCompilerPlugins [22]，目标是为了这套流水线开发相应的支持工具。其中一个应用就是新的 YaoCompiler，在 YaoCompiler 里我们实现了一套长得和普通 Julia 代码差不多的量子计算专用语言，通过这套框架我们可以将看起来是 Julia 代码的函数编译到 OpenQASM，IBM 和亚马逊的硬件，或者是 LLVM IR。

这套框架相比 Python 各写各的好处是由于所有的 eDSL 只是修改了 Julia 的 SSA IR，而不是使用一个完全新的 IR，我们依然保留一定程度的可组合性。

但是问题也是有的，由于 Julia 的中间表示本身不是为了这套流水线设计的，想要支持更复杂和更灵活的自定义 type lattice，或者是代码优化逻辑我们还需要去对 Julia 编译器本身进行大量的改进，这也是为什么你会发现目前 Julia 里没有人在这套中间表示上做编译优化，或者某个专有领域的优化。因为在这套表示里将自定义的指令作为第一公民来优化都还是不够稳定的。而 Julia 里目前对线性代数的自动优化也是远远不如 TVM，JAX 等以 Python 为主要支持的编译器的。

而这件事的难度和能够完成的时间都还不可知。因为想要做出一个足够一般的 abstract interpreter 还要保证这个过程中不破坏任何已有代码这件事本身并不简单。

MLIR 与 Julia

和 Python 类似，MLIR 也会是 Julia 里的 eDSL 的未来。我觉得你甚至可以认为 MLIR 是 IR 领域的 Julia - 它（尝试）解决了多 IR 的问题。

尽管我们认为 Julia 在很多方面相比 Python 更好，但是对于用户和商业公司来说 Python 依然是一个必须要支持的语言，所以使用一个共同的编译后端（像 TVM 就支持了很多种语言作为前端）是一个必然的选择。

实际上我们现在已经可以把整个 Julia 编译到 MLIR 的 Julia dialect 了，MIT 的 Julia Lab 开发了一套新的基于 MLIR 的 Julia 编译器 JuliaLabs/brutus [23] 和 MLIR 的 Julia 封装 vchuravy/MLIR.jl [24]。至于具体的 MLIR 的优势我在上面已经介绍过了就不重复介绍了。

当然由于编译器工程师普遍比较稀缺，我相信这个方向还会发展个至少五六年才能走进千家万户， 而或许等那个时候你可能会发现，虽然你写的是 Python 代码，但是其实编译到了 Julia + 某几个 MLIR 方言，然后运行在某台有 GPU + CPU + 量子硬件的机器上呢？

点击阅读原文跳转到知乎原文

参考

[1] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake Vander-Plas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transfor-mations of Python+NumPy programs, 2018.

[2] TorchScript https://pytorch.org/docs/stable/jit.html

[3] Numba: A High Performance Python Compiler. http://numba.pydata.org/

[4] Taichi: A High-Performance Programming Language for Computer Graphics Applications. https://github.com/taichi-dev/taichi

[5] MLIR https://mlir.llvm.org/

[6] MLIR based backend for Numba https://github.com/numba/numba/issues/6475

[7] ForwardDiff.jl: Forward Mode Automatic Differentiation for Julia https://github.com/JuliaDiff/ForwardDiff.jl

[8] Flux — Elegant ML https://github.com/FluxML/Flux.jl

[9] Keno Fischer and Elliot Saba. Automatic full compilation of julia programsand ML models to cloud tpus.CoRR, abs/1810.09868, 2018. https://arxiv.org/abs/1810.09868

[10] Yao.jlL Extensible, Efficient Quantum Algorithm Design for Humans. https://github.com/QuantumBFS/Yao.jl

[11] JuMP: A Modeling Language and Supporting Packages for Mathematical Optimization in Julia. https://jump.dev/

[12] ModelingToolkit.jl: A Modeling Framework for Automatically Parallelized Scientific Machine Learning (SciML) in Julia.

[13] Soss.jl: Probabilistic Programming via Source Rewriting https://github.com/cscherrer/Soss.jl

[14] Turing.jl: Bayesian Inference with Probabilistic Programming https://github.com/TuringLang/Turing.jl

[15] Gen.jl: A General-Purpose Probabilistic Programming System with Programmable Inference https://github.com/probcomp/Gen.jl

[16] Professor Steven G. Johnson https://math.mit.edu/~stevenj/

[17] FFTW https://www.fftw.org/

[18] Julia CodeInfo https://docs.julialang.org/en/v1/devdocs/ast/#CodeInfo

[19] GPUCompiler.jl: Reusable compiler infrastructure for Julia GPU backends. https://link.zhihu.com/?target=https%3A//github.com/JuliaGPU/GPUCompiler.jl

[20] Enzyme: High-performance automatic differentiation of LLVM. https://github.com/wsmoses/enzyme

[21] Diffractor.jl: Next-Generation AD. https://github.com/JuliaDiff/Diffractor.jl

[22] JuliaCompilerPlugin: Configurable compiler plugins for the Julia language. https://juliacompilerplugins.github.io/

[23] Brutus: A research project that uses MLIR to implement code-generation and optimisations for Julia. https://github.com/JuliaLabs/brutus/

[24] MLIR.jl: Presents high-level tools to manipulate MLIR dialects through the MLIR C API. https://github.com/vchuravy/MLIR.jl