MegEngine Windows Python wheel 包减肥之路

时间:2022-12-19 18:03:28

作者:张浩龙 | 旷视科技 MegEngine 架构师

写在之前

本文的目的

  • 通过讲述在支持 MegEngine Windows Python wheel 过程中遇到的问题以及解决问题的流程,此文最后的解决方法可能不是最优,欢迎留言指正。
  • 过程中顺便科普一些关于 MegEngine 的构建以及构建时用到的基础东西,当然这些基础知识我相信是工程之道经常会用到的,包括但不限于:
    • CMake
    • 编译、链接、符号隐藏,符号 export 等。此处先推荐一本 “老书”《程序员的自我修养》,自然它没有 xxx 四库全书让人绞尽脑汁,但是它里面的基础知识依然是目前我们和计算机“交流”中经常遇到的。
    • Python wheel 包构建

MegEngine 各平台支持情况

  • cpp 推理支持情况:

MegEngine Windows Python wheel 包减肥之路

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

  • 训练:

    • Python侧: MegEngine Windows Python wheel 包减肥之路
    • 目前官方发布的 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

MegEngine Windows Python wheel 包减肥之路

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

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

先对问题 MECE 一下

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

  • 各部分之间相互独立(MutuallyExclusive)
    • 化简后, 感觉就是分析问题时,可能的方法要尽可能的独立,尽量不要有交集
  • 所有部分完全穷尽(CollectivelyExhaustive)
    • 化简后, 感觉就是分析问题时,尽可能的要把方法想全,尽量不要有遗漏

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

MegEngine Windows Python wheel 包减肥之路

问题的影响

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

  • 首先就是体积超过了我们申请的 pypi 上单个文件最大体积限制
  • 给用户体验不好, 为什么相同的版本, Window 比 Linux大这么多呢
  • Windows 上显存比 Linux 占用大很多(估计提到这点,大家已经猜测到问题所在)

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

问题给人的第一印象是:

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

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

MegEngine CMake 构建流程

  • MegEngine 构建依赖 CMakeNinja

  • 其中 CMake 描述主要在: MegEngine Windows Python wheel 包减肥之路

    • 顶层 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
  • 其中 Ninja 提供丰富的可视化调试功能,下面列举如何通过 Ninja debug server 来看 MegEngine 部分模块的构建依赖

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

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

      $ ninja -t list
      ninja subtools:
          browse  browse dependency graph in a web browser
           clean  clean built files
        commands  list all commands required to rebuild given targets
            deps  show dependencies stored in the deps log
           graph  output graphviz dot file for targets
           query  show inputs/outputs for a path
         targets  list targets by their rule or depth in the DAG
          compdb  dump JSON compilation database to stdout
       recompact  recompacts ninja-internal data structures
          restat  restats all outputs in the build log
           rules  list all rules
       cleandead  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 它的依赖图如下 MegEngine Windows Python wheel 包减肥之路

      • 依赖一些 lite/CMakeFiles 下面的 .o

      • 依赖一些第三方的.a,比如 cpuinfo.a

      • 依赖 libmegbrain_shared.so 此库包含了 megbrain/dnn 所有的编译输出, 当然还可以鼠标在 ninja 起的 server 任意点击展开任一一个目标来查看它的依赖情况

MegEngine wheel 构建流程

  • 有了上面 Ninja 编译出来的各种库后,我们就可以将它们和 MegEngine 中的 py src 一起进行打包,最后生成可安装,可分发的 Python wheel 包了 MegEngine Windows 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 https://megengine.org.cn/whl/mge.html 后,可以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 的哦

问题简单的分析

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

Linux 下

MegEngine Windows Python wheel 包减肥之路

Windows 下

MegEngine Windows Python wheel 包减肥之路

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 Windows Python wheel 包减肥之路

    • 从下往上依次是:

      • MegEngine 不同 backends 的差异化实现被封装到了 dnn(对应到上图的“硬件层”,对应到 CMakeList 中的 megdnn 模块), 而其中 CUDA backends 因为有大量的 kernel以及对较多的 SM 支持,会对整个库或者可执行程序体积产生大量的体积贡献
      • 图中“硬件抽象层”,部分“核心组件层”,对应 CMakeList 中的 megbrain 模块
      • 在往上“接口层”,对应 CMakeList 中的 imperative 和 MegEngineLite 模块。
    • 由于上述的原因,在 Windows 平台, imperative 模块和 MegEngineLite 模块会同时静态依赖 dnn 和 megbrain 代码,导致体积几乎翻倍。问题修复前的依赖图:

MegEngine Windows Python wheel 包减肥之路

可能的解决方案

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

  • 为什么 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.h
      #ifdef DLL_EXPORT
      #define DECLSPEC_FUC __declspec(dllexport)
      #define DECLSPEC_DATA __declspec(dllexport)
      #else
      #define DECLSPEC_FUC
      #define DECLSPEC_DATA __declspec(dllimport)
      #endif
      DECLSPEC_DATA extern int a;
      DECLSPEC_FUC void func_a();
      /////api.c
      #include "api.h"
      int a = 0;
      void func_a() {}
      # build api.c with define DLL_EXPORT
      
      ////client.c
      #include "api.h"
      int main(){func_a();a = 1;}
      # build client.c without define DLL_EXPORT
      
    • 跨 dll 动态库通过函数访问数据成员方式

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

      ////api.h
      #ifdef DLL_EXPORT
      #define DECLSPEC_FUC __declspec(dllexport)
      #else
      #define DECLSPEC_FUC
      #endif
      extern int a;
      DECLSPEC_FUC void func_a();
      DECLSPEC_FUC int * get_a();
      /////api.c
      #include "api.h"
      int a = 0;
      void func_a() {}
      int * get_a() {return &a;}
      # build api.c with define DLL_EXPORT
      
      ////client.c
      #include "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
      • CMakeList 中目标依赖修改过后的exports.def (让其符号不超过 65536)
      • 结论:不可行
        • 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_HOST
     TensorShape() = 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 scalar
     bool 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 check
     size_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