Just-in-time compilation (JIT)#

我们在 MegEngine basic concepts 这篇教程中曾提到过这样一个概念:

“在深度学习领域,任何复杂的深度神经网络模型本质上都可以用一个计算图表示出来。”

../../../_images/computing_graph.png

In MegEngine, Tensor is the data node, and Operator is the computing node#

>>> y = w * x + b
>>> p = w * x
>>> y = p + b

在默认情况下,MegEngine 中的指令将像 Python 之类的解释型语言一样,动态地解释执行。 我们将这种执行模式称为 “动态图” 模式,此时完整的计算图信息其实并不存在; 而与之相对的是 “静态图” 模式,在执行之前能够拿到完整的计算图结构, 能够根据全局图信息进行一定的优化,加快执行速度。

在 MegEngine 中,通过使用即时编译技术(JIT),将动态图编译成静态图,并支持序列化。

接下来的内容将对相关概念和原理进行一定的解释,不了解这些细节并不影响基本使用。

动态图 Vs. 静态图#

动态图

MegEngine uses dynamic calculation graph by default, and its core feature is that the construction and calculation of the calculation graph occur at the same time (Define by run).

  • Principle: When defining a: py:class:~megengine.Tensor in the calculation graph, its value has been calculated and determined.

  • Advantage: This mode is more convenient when debugging the model, and can get the value of the intermediate result in real time.

  • Disadvantages: But because all nodes need to be saved, this makes it difficult for us to optimize the entire calculation graph.

借助即时编译技术,MegEngine 中的动态图可通过 trace 接口转换成静态图。

静态图

MegEngine also supports the static calculation graph mode, which separates the construction of the calculation graph from the actual calculation (Define and run).

  • Principle: In the construction phase, MegEngine optimizes and adjusts the original calculation graph according to the complete calculation process to obtain a calculation graph that saves more memory and has less calculation. This process is called “compilation”. The structure of the graph does not change after compilation, which is the so-called “static”. In the calculation phase, MegEngine executes the compiled calculation graph according to the input data to obtain the calculation result.

  • Advantages: Compared with dynamic images, static images have a richer grasp of global information and can do more optimizations.

  • Disadvantages: But the intermediate process is a black box for the user, and the intermediate calculation result cannot be obtained at any time like a dynamic graph.

什么是即时编译#

即时编译(Just-in-time compilation)是源自编译(Compiling)中的概念。

以传统的 C/C++ 语言为例,我们写完代码之后, 一般会通过编译器编译生成可执行文件,然后再执行该可执行文件获得执行结果。 如果我们将从源代码编译生成可执行文件袋过称为 build 阶段, 将执行可执行文件叫做 runtime 阶段的话,JIT 是没有 build 阶段的,只存在于 runtime 阶段。 JIT 一般被用在解释执行的语言如 Python 中,JIT 会在代码执行的过程中检测热点函数(HotSpot), 随后对热点函数进行重编译,下次运行时遇到热点函数则直接执行编译结果即可。这样做可以显著加快代码执行的速度。

See also

维基百科: Just-in-time compilation

Examples of static graph compilation optimization#

Below we illustrate in FIG static memory and computing compiler optimization process may be:

../../../_images/op_fuse.png

In the calculation diagram on the left side of the above figure, in order to store the five variables of x, w, p, b, y, the dynamic graph needs 40 Bytes (assuming that each variable occupies 8 bytes of memory). In the static diagram, since we only need to know the result y, we can make y reuse the memory of the intermediate variable p to achieve “inplace” modification. In this way, the memory occupied by the static image is reduced to 32 bytes.

See also

更多相关解释可参考 MegEngine 官方博客 《 JIT in MegEngine