MegEngine Windows Python wheel 包减肥之路

写在之前

1.1 本文的目的

通过讲述在支持 MegEngine Windows Python wheel 过程中遇到的问题以及解决问题的流程，此文最后的解决方法可能不是最优，欢迎留言指正。

过程中顺便科普一些关于 MegEngine 的构建以及构建时用到的基础东西，当然这些基础知识我相信是工程之道经常会用到的，包括但不限于：

CMake
编译、链接、符号隐藏，符号 export 等。此处先推荐一本 “老书”《程序员的自我修养》，自然它没有 xxx 四库全书让人绞尽脑汁，但是它里面的基础知识依然是目前我们和计算机“交流”中经常遇到的。
Python wheel 包构建

1.2 MegEngine 各平台支持情况

cpp 推理支持情况

TEE：https://en.wikipedia.org/wiki/Trusted_execution_environment

训练

Python 侧

目前官方发布的 wheel 包，只有 Windows-X64-CPU-CUDA，many Linux 64bit -X64-CPU-CUDA，MacOS-X64-CPU，其他的可自己编译，或者社区提单索取。
cpp 侧训练支持情况和上面的 cpp 推理情况一致。

从上面的情况，可看见 MegEngine 无论训练还是推理，还是各种硬件，还是各种 OS 都支持的非常全面，如有需求，不妨试用！！！！

遇到的问题

为了全面的支持上面提到的 MegEngine 各个平台，各个 OS，期间或多或少会遇到一些问题，比如 Windows 平台上 Python wheel 包体积过大。

先看一下目前 MegEngine wheel 包体积大小，摘自 1.7 版本 pypi（https://pypi.org/project/MegEngine/1.7.0/#files）

其中因为 Linux 和 Windows 支持了 CUDA，所以包体积在 900MB 左右，这是一个正常的 size。

在之前 Windows CUDA 包体积在 1.7G 左右: 这就是后面尝试分析和修复的问题。

2.1 先对问题 MECE 一下

MECE 是 Mutually Exclusive Collectively Exhaustive 缩写，中文意思是“相互独立，完全穷尽”。也就是对于一个问题的议题，能够做到不重叠、不遗漏的分类，而且能够藉此有效把握问题的核心，并解决问题的方法。强调两点:

各部分之间相互独立（MutuallyExclusive）

化简后，感觉就是分析问题时，可能的方法要尽可能的独立，尽量不要有交集

所有部分完全穷尽（CollectivelyExhaustive）

化简后，感觉就是分析问题时，尽可能的要把方法想全，尽量不要有遗漏

为啥会选择一条鱼呢：鱼头附近一般都比较大(胖)，而越往鱼尾走，会越来越小(瘦)，从而希望通过这个流程“Windows wheel 减肥”之路，达到减肥的目的。

2.2 问题的影响

这样的体积会有什么问题呢（毕竟人太胖也会有些副作用 xxx etc.）

首先就是体积超过了我们申请的 pypi 上单个文件最大体积限制

给用户体验不好，为什么相同的版本， Window 比 Linux 大这么多呢

Windows 上显存比 Linux 占用大很多（估计提到这点，大家已经猜测到问题所在）

2.3 解决这个问题可能需要的相关知识

问题给人的第一印象是：

编译构建相关的
Python wheel 打包相关的
Windows OS 独有的

先就 MegEngine 如下基础知识做一些基础补充（减肥前总得有一些科普吧，到底是吃药还是锻炼，或者具体到吃什么药吧）

MegEngine CMake 构建流程

MegEngine 构建依赖 CMake 和 Ninja

其中 CMake 描述主要在：

顶层 CMakeList

此文件包含很多的 option，主要用于控制是否编译一些模块，比如是否打开 MGE_BUILD_WITH_ASAN 用于调试内存问题
此文件包含对各种 ARCH 的适配控制，比如编译 X86_64 还是 AARCH64 等
此文件包含对各种 OS 的适配，比如是编译 Linux，Android 还是 Windows 等
以及一些杂项配置，比如有优化等级，比如 CUDA SM 的配置等

src CMakelist

此文件包含了 MegEngine 核心代码 MegBrain 层所有源代码的编译管理

dnn CMakelist

此文件包含了 MegEngine 核心代码 dnn（主要实现各种 backends）层所有源代码的编译管理

一些杂项 CMakelist

各种 example，比如 lite example
各种 test，比如 megbrain test
各种 helper，见 helper module

其中 Ninja 提供丰富的可视化调试功能，下面列举如何通过 Ninja debug server 来看 MegEngine 部分模块的构建依赖

执行 host_build.sh 来进行 host 编译，同时它会在 build dir 生成整个构建依赖描述文件 build.ninja

更多的编译支持请参考 BUILD_README.md
build.ninja 文件功能类似 GNU Makefile

有了 build.ninja 后，便可进行一些调试

$ ninja -t listninja subtools:browse  browse dependency graph in a web browserclean  clean built filescommands  list all commands required to rebuild given targetsdeps  show dependencies stored in the deps loggraph  output graphviz dot file for targetsquery  show inputs/outputs for a pathtargets  list targets by their rule or depth in the DAGcompdb  dump JSON compilation database to stdoutrecompact  recompacts ninja-internal data structuresrestat  restats all outputs in the build logrules  list all rulescleandead  clean built files that are no longer produced by the manifest

可通过如下命令来启动一个可视化 server（当然熟悉后，也可通过其他命令行参数来调试，自然也可以直接看 CMakeLists.txt 来找关系）

ninja -t browse --port you_port --hostname you_ip --no-browser

然后通过在浏览器输入: you_ip:you_port 进行可视化浏览了
还可以通过 dot 生成 png 来查看

ninja -t graph > 1.dot dot -Tpng 1.dot > output.png

比如 megenginelite 一个最上层的目标：liblite_shared_whl.so 它的依赖图如下

依赖一些 lite/CMakeFiles 下面的 .o
依赖一些第三方的 .a，比如 cpuinfo.a
依赖 libmegbrain_shared.so 此库包含了 megbrain/dnn 所有的编译输出，当然还可以鼠标在 ninja 起的 server 任意点击展开任一一个目标来查看它的依赖情况

MegEngine wheel 构建流程

有了上面 Ninja 编译出来的各种库后，我们就可以将它们和 MegEngine 中的 py src 一起进行打包，最后生成可安装，可分发的 Python wheel 包了

构建流程主要说明在 BUILD_PYTHON_WHL_README

描述了目前 MegEngine python wheel 的支持状态
自己本地构建需要的一些 env 准备
一些使用说明
MegEngine wheel 包遵从 pep-0571

包 setup 入口在 python wheel setup
调用 setup 前各个 OS 的准备差异化在 wheel scripts
有人可能会问，为什么不使用 auditwheel 来自动管理 wheel 包中 so 依赖，有两个原因

auditwheel 不支持所有的操作系统，比如 Windows
auditwheel 不支持依赖库使用 dlopen 的情况
auditwheel 不支持 subpackage 的 wheel 包

当你执行 python3 -m pip install megengine -f megengine.org.cn/whl/mg 后，可以import megengine，也可以import megenginelite，是因为 megengine 和 megenginelite 均会存在安装的包中，且他们会复用 megengine_shared 这一体积超大的库

MegEngine 构建上如何适配 Windows

上面介绍了 MegEngine 基于 CMake 的构建基础和使用 Ninja 自带的调试功能以及帮我们从宏观了解了一下 project 的编译依赖和进行一些常规调试，下面再介绍一下 MegEngine 是如何适配 Windows 平台的。

首先 MegEngine 大部分的源代码都是c++，且cpp 推理要求是c++14，编译 Python 训练要求是 c++17

各家编译器其实对这些标准实现不是完全一致的，抛开和系统相关的，比如POSIX 外，其实还有比较多基本的上层用法各家编译器其实时不太兼容的，特别是明显的是 gcc和clang 能编译过的代码，Windows cl.exe 其实是编译不过的。

为了解决上面提到的两个问题

尽可能的抛弃 cl.exe，Windows 上使用 llvm-clang-cl 进行构建

cmake ...
-DCMAKE_C_COMPILER=clang-cl.exe
-DCMAKE_CXX_COMPILER=clang-cl.exe
...

当然因为 CUDA 的原因，目前不可能完全抛弃 cl.exe，在编译

CUDA host 代码时，依然使用 cl.exe

区分开系统相关的函数实现，所以你会在 MegEngine 代码中看到不

少如下类似的代码

#if WIN32
....
#else
....
#endif

再加上，上面提到的 CMake, Ninja 本身是跨平台的，这样一组合，MegEngine 便原生支持了Windows，注意不是基于 WSL 的哦

2.4 问题简单的分析

在上面“MegEngine CMake 构建流程”小节中，我们提到了 Ninja debug server 能够帮忙可视化整个构建组件的依赖关系，下面我们补充一下在问题修复前 Windows 和Linux 下 imperative 依赖的可视化结果

Linux 下

Windows 下

Windows 和 Linux 下最大的差异化如下:

Linux 下 imperative 是依赖的 libmegengine_shared.so

Windows 下 imperative 是依赖的 megbrain 和 megdnn，又因为 megbrain 和 dnn 在 CMake 这边其实一个 OBJECT，所以相当于直接依赖他们的 .obj 了

初步结论:

MegEngine wheel 包，有两个 Python module 接口

MegEngine 用于python侧训练的基础接口

imperative：当你安装完成 MegEngine 时，在 Python 中输入 import megengine 时。加载的就是它，我们提供了一些入门的教程供您快速上手 MegEngine

MegEngineLite: 易用的 cpp , Python 推理接口

当你使用 MegEngine 完成训练模型后，可参考部署文档使用 MegEngineLite，快速将你的模型部署落地。

因为有这两个顶层构建目标的存在，且他们在 Windows 和其他 OS 上，依赖底层的目标不同，导致了问题的产生

为什么不同的依赖关系，会产生这么大的体积区别呢，先看一张 MegEngine 的架构图

从下往上依次是：

MegEngine 不同 backends 的差异化实现被封装到了 dnn（对应到上图的“硬件层”，对应到 CMakeList 中的 megdnn 模块），而其中 CUDA backends 因为有大量的 kernel 以及对较多的 SM 支持，会对整个库或者可执行程序体积产生大量的体积贡献
图中“硬件抽象层”，部分“核心组件层”，对应 CMakeList 中的 megbrain 模块
在往上“接口层”，对应 CMakeList 中的 imperative和 MegEngineLite 模块。

由于上述的原因，在 Windows 平台， imperative 模块和 MegEngineLite 模块会同时静态依赖 dnn 和 megbrain 代码，导致体积几乎翻倍。问题修复前的依赖图：

2.5 可能的解决方案

通过上面的分析，问题原因已经找到，再来猜想一下

为什么 Windows 平台上和其他平台目标依赖有差异

Windows class member 不能隐式的被export，需要显式的使用 dllexport 和 dllimport，详细见 Microsoft Specific

dnn，megbrain 层有大量的 data 数据访问并没有抽象成函数，而是需要直接访问数据成员下面举一个栗子来说明跨 dll 动态库访问数据成员方式的差异，主要包含三支文件 api.h 和 api.c 实现函数 func_a 和定义一个变量 a，被编译成动态库 dll client.c 会调用上面 api.c 实现方法和访问变量 a

跨 dll 动态库直接访问数据成员方式

api.hifdef DLL_EXPORTdefine DECLSPEC_FUC __declspec(dllexport)define DECLSPEC_DATA __declspec(dllexport)elsedefine DECLSPEC_FUCdefine DECLSPEC_DATA __declspec(dllimport)endifDECLSPEC_DATA extern int a; DECLSPEC_FUC void func_a(); /api.cinclude "api.h"int a = 0; void func_a() {}build api.c with define DLL_EXPORTclient.cinclude "api.h"int main(){func_a();a = 1;}build client.c without define DLL_EXPORT

跨 dll 动态库通过函数访问数据成员方式

改造上面的 example 代码，把其中变量 a 的访问封装到一个函数

api.hifdef DLL_EXPORTdefine DECLSPEC_FUC __declspec(dllexport)elsedefine DECLSPEC_FUCendifextern int a; DECLSPEC_FUC void func_a(); DECLSPEC_FUC int * get_a(); /api.cinclude "api.h"int a = 0; void func_a() {} int * get_a() {return &a;}build api.c with define DLL_EXPORTclient.cinclude "api.h"int main(){func_a();int a = get_a(); a = 1; return 0;}build client.c without define DLL_EXPORT

可以看见在 Windows 上函数符号和数据成员符号 export 是等价的，但是 import data 要求要严格的多
在 Linux、MacOS 下，函数符号和数据成员符号 export 属性是等价的

由于上面提到种种限制，导致在最初支持 Windows 平台时，所有的上层目标（MegEngine，MegEngineLite）都必须静态依赖 megbrain 和 dnn。

既然问题原因已经找到，需要修复这个问题的目标就变的非常清晰了：让 megengine_shared 动态库 (dll) 在 Windows 平台上可用。

列举一下可能的方案：

方案一：CMake 自带的 WINDOWS_EXPORT_ALL_SYMBOLS

结论：不太适用 MegEngine 这类“大”工程
原因：MegEngine 符号太多，超过了 link.exe max symbols 65536 的限制（使能 CUDA 时，大约有 1.7W 个符号）
分析 CMake WINDOWS_EXPORT_ALL_SYMBOLS 的原理，能否中间加一些hook来过滤不需要export的符号,以达到类似gcc/clang - Wl,--version-script 的效果，cmake 对他的处理逻辑:

(stage a): 生成 CMakeFiles/megengine_export.dir/exports.def.objs，本质是 obj 的集合
(stage b): 插入 PRE_LINK stage 生成 CMakeFiles/megengine_export.dir/exports.def（此文件类似gcc/clang -Wl,--version-script）
(stage c): LINK_FLAG 自动插入 /DEF exports.def
CMake 提供了对 stage a output的 hook，意思是可以修改 exports.def.objs，但是没有机会修改 exports.def

加入 hook command，把 exports.def.objs 中所有DNN 的 obj删除，想象中应该可以了
但是 imperive 和 megenginelite，不仅仅是和 megbrain 打交到，很多直接使用dnn的接口和数据成员

方案二：“优化”版本的 WINDOWS_EXPORT_ALL_SYMBOLS

如上面分析 WINDOWS_EXPORT_ALL_SYMBOLS 有一定的缺陷，会把所有的 obj 的符号全部 export，那能不能手动修改 WINDOWS_EXPORT_ALL_SYMBOLS 生成的 exports.def

保留必要的 symbols
如上面分析 WINDOWS_EXPORT_ALL_SYMBOLS 有一定的缺陷，会把所有的 obj的符号全部export，那能不能手动修改 WINDOWS_EXPORT_ALL_SYMBOLS生成的 exports.defef
结论：不可行

Windows cl linker.exe不支持 * 通配符，不支持存放一个不存在的符号，导致一旦放了固定的exports.def，稍微更改一个编译参数，或者加点代码，都会编译不过

方案三：到最后发现没有一个“偷懒的”方式来解决这个问题，回退到最 naive 的方式

把 megbrain、dnn、megenginelite 对外暴露的 API 依赖的成员

符号全部显式的加上 declspec(dllexport) 和 declspec(dllimport) 属性描述

修复示例，完整修改见 PR

diff --git a/dnn/include/megdnn/basic_types.h b/dnn/include/megdnn/basic_types.h
index 53f22c9af..44831f6d7 100644
--- a/dnn/include/megdnn/basic_types.h
+++ b/dnn/include/megdnn/basic_types.h
@@ -104,22 +104,22 @@ struct TensorShape {#if MEGDNN_CC_HOSTTensorShape() = default;TensorShape(const TensorShape& rhs) = default;
-    TensorShape(const SmallVector<size_t>& init_shape);
-    TensorShape(std::initializer_list<size_t> init_shape);
-    std::string to_string() const;
+    MGE_WIN_DECLSPEC_FUC TensorShape(const SmallVector<size_t>& init_shape);
+    MGE_WIN_DECLSPEC_FUC TensorShape(std::initializer_list<size_t> init_shape);
+    MGE_WIN_DECLSPEC_FUC std::string to_string() const;#endif//! total number of elements
-    size_t total_nr_elems() const;
+    MGE_WIN_DECLSPEC_FUC size_t total_nr_elems() const;//! check whether two shapes are equal
-    bool eq_shape(const TensorShape& rhs) const;
+    MGE_WIN_DECLSPEC_FUC bool eq_shape(const TensorShape& rhs) const;//! check whether the shape can be treated as a scalarbool is_scalar() const { return ndim == 1 && shape[0] == 1; }//! check whether ndim != 0 and at least one shape is 0
-    bool is_empty() const;
+    MGE_WIN_DECLSPEC_FUC bool is_empty() const;//! access single element, without boundary checksize_t& operator[](size_t i) { return shape[i]; }
@@ -168,8 +168,8 @@ struct TensorLayout : public TensorShape {class ImplBase;#if MEGDNN_CC_HOST
-        Format();
-        Format(DType dtype);
+        MGE_WIN_DECLSPEC_FUC Format();
+        MGE_WIN_DECLSPEC_FUC Format(DType dtype);

想象未来更好的解决方法：

修改 CMake 本身源代码，让 flag WINDOWS_EXPORT_ALL_SYMBOLS 支持用户自定义filter，让其生成的exports.def 本身就是带用户过滤参数的

当然因为 Windows 数据成员在import 部分处还必须显式的加上 dllimport，对这块似乎 CMake 也无能为力

可以考虑工程一开始设计时，API 尽可能的不要存在隐式的数据成员之间的访问，尽可能的将其转换成一个函数 API

点击“阅读原文”，前往知乎技术硬核文章

招聘帖

文章都看完啦~ 是否有兴趣加入到深度学习框架开发中来？

Megengine 团队现正火热招聘中！期待你的加入~
简历投递或详细信息了解可添加微信：duoduo715495

【框架开发工程师（C++）】
职位描述:
1. 负责旷视核心深度框架 MegEngine 的设计，演进，实现，维护和优化
2. 优化 MegEngine 在各个计算平台（CUDA / Arm / x86 等）上的性能
3. 持续探索改进深度学习框架的先进优化方法（例如图优化，自动代码生成，超低 bit 量化，稀疏化等）

技能要求:
1. 1-3 年的性能优化经验（X86，CUDA，ARM，OpenCL 等）
2. 有良好的数据结构与算法功底，能够熟练使用 C++ 语言编写较复杂的算法
3. 对深度学习和深度学习框架（Caffe，Tensorflow，PyTorch 等）有基本了解者优先
4. 有复杂系统设计经验者优先