feat(lite/whl): merge lite whl to main package

GitOrigin-RevId: 27c7e50207

imperative/python/setup.py

@@ -78,6 +78,16 @@ megengine_data += [
     for f in pathlib.Path('megengine', 'core', 'lib').glob('**/*')
 ]
 
+megenginelite_data = [
+    str(f.relative_to('megenginelite'))
+    for f in pathlib.Path('megenginelite').glob('**/*')
+]
+
+if platform.system() == 'Windows':
+    megenginelite_data.remove('libs\\liblite_shared_whl.pyd')
+else:
+    megenginelite_data.remove('libs/liblite_shared_whl.so')
+
 with open('requires.txt') as f:
     requires = f.read().splitlines()
 with open('requires-style.txt') as f:
@@ -86,6 +96,7 @@ with open('requires-test.txt') as f:
     requires_test = f.read().splitlines()
 
 prebuild_modules=[PrecompiledExtesion('megengine.core._imperative_rt')]
+prebuild_modules.append(PrecompiledExtesion('megenginelite.libs.liblite_shared_whl'))
 setup_kwargs = dict(
     name=package_name,
     version=__version__,
@@ -96,6 +107,7 @@ setup_kwargs = dict(
     packages=packages,
     package_data={
         'megengine': megengine_data,
+        'megenginelite': megenginelite_data,
     },
     ext_modules=prebuild_modules,
     install_requires=requires,

lite/CMakeLists.txt

@@ -75,12 +75,30 @@ if(ANDROID)
     target_link_libraries(lite_shared PRIVATE log)
 endif()
 
+# define a shared lib for whl
+add_library(lite_shared_whl SHARED $<TARGET_OBJECTS:lite_static>)
+if(LITE_BUILD_WITH_MGE)
+    if (MSVC OR WIN32 OR IOS)
+        # TODO: this will lead whl size increase on Windows, caused by
+        # Windows does not support implicitly importing data members from DLL.
+        target_link_libraries(lite_shared_whl PRIVATE megbrain megdnn ${MGE_CUDA_LIBS})
+    else()
+        target_link_libraries(lite_shared_whl PRIVATE megengine_export)
+    endif()
+endif()
+if(ANDROID)
+    target_link_libraries(lite_shared_whl PRIVATE log)
+endif()
+
+
 if(LITE_BUILD_WITH_MGE AND LITE_WITH_CUDA AND NOT WIN32)
     # FXIME third_party cpp redis do not support build with clang-cl
     target_include_directories(lite_static PRIVATE ${PROJECT_SOURCE_DIR}/third_party/cpp_redis/includes)
     target_include_directories(lite_static PRIVATE ${PROJECT_SOURCE_DIR}/third_party/tacopie/includes)
     target_include_directories(lite_shared PRIVATE ${PROJECT_SOURCE_DIR}/third_party/cpp_redis/includes)
     target_include_directories(lite_shared PRIVATE ${PROJECT_SOURCE_DIR}/third_party/tacopie/includes)
+    target_include_directories(lite_shared_whl PRIVATE ${PROJECT_SOURCE_DIR}/third_party/cpp_redis/includes)
+    target_include_directories(lite_shared_whl PRIVATE ${PROJECT_SOURCE_DIR}/third_party/tacopie/includes)
 endif()
 set(LITE_VERSION_SCRIPT ${PROJECT_SOURCE_DIR}/lite/src/version_lite.ld CACHE INTERNAL "Path to linker version script")
 add_custom_target(_lite_version_ld SOURCES ${LITE_VERSION_SCRIPT})
@@ -92,6 +110,8 @@ endif()
 if (UNIX AND NOT APPLE)
     target_link_options(lite_shared PRIVATE -Wl,--version-script=${LITE_VERSION_SCRIPT})
     set_target_properties(lite_shared PROPERTIES LINK_DEPENDS ${LITE_VERSION_SCRIPT})
+    target_link_options(lite_shared_whl PRIVATE -Wl,--version-script=${LITE_VERSION_SCRIPT})
+    set_target_properties(lite_shared_whl PROPERTIES LINK_DEPENDS ${LITE_VERSION_SCRIPT})
 endif()
 
 # config install

lite/pylite/megenginelite/base.py

@@ -15,7 +15,7 @@ import sys
 from ctypes import *
 
 if sys.platform == "win32":
-    lib_path = os.path.join(os.path.dirname(__file__), "libs")
+    lib_path = os.path.join(os.path.dirname(__file__), "../megengine/core/lib")
     dll_paths = list(filter(os.path.exists, [lib_path,]))
     assert len(dll_paths) > 0
 

lite/pylite/requires.txt

@@ -1 +0,0 @@
-numpy>=1.18

lite/pylite/setup.py

@@ -1,130 +0,0 @@
-# -*- coding: utf-8 -*-
-# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
-#
-# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-
-import os
-import re
-import pathlib
-import platform
-from distutils.file_util import copy_file
-from setuptools import setup, find_packages, Extension
-from setuptools.command.build_ext import build_ext as _build_ext
-
-class PrecompiledExtesion(Extension):
-    def __init__(self, name):
-        super().__init__(name, sources=[])
-
-class build_ext(_build_ext):
-
-    def build_extension(self, ext):
-        if not isinstance(ext, PrecompiledExtesion):
-            return super().build_extension(ext)
-
-        if not self.inplace:
-            fullpath = self.get_ext_fullpath(ext.name)
-            extdir = pathlib.Path(fullpath)
-            extdir.parent.mkdir(parents=True, exist_ok=True)
-
-            modpath = self.get_ext_fullname(ext.name).split('.')
-            if platform.system() == 'Windows':
-                modpath[-1] += '.dll'
-            elif platform.system() == 'Darwin':
-                modpath[-1] += '.dylib'
-            else:
-                modpath[-1] += '.so'
-            modpath = str(pathlib.Path(*modpath).resolve())
-
-            copy_file(modpath, fullpath, verbose=self.verbose, dry_run=self.dry_run)
-
-v = {}
-with open("megenginelite/version.py") as fp:
-    exec(fp.read(), v)
-__version__ = v['__version__']
-
-email = 'megengine@megvii.com'
-# https://www.python.org/dev/peps/pep-0440
-# Public version identifiers: [N!]N(.N)*[{a|b|rc}N][.postN][.devN]
-# Local version identifiers: <public version identifier>[+<local version label>]
-# PUBLIC_VERSION_POSTFIX use to handle rc or dev info
-public_version_postfix = os.environ.get('PUBLIC_VERSION_POSTFIX')
-if public_version_postfix:
-    __version__ = '{}{}'.format(__version__, public_version_postfix)
-
-local_version = []
-strip_sdk_info = os.environ.get('STRIP_SDK_INFO', 'False').lower()
-sdk_name = os.environ.get('SDK_NAME', 'cpu')
-if 'true' == strip_sdk_info:
-    print('wheel version strip sdk info')
-else:
-    local_version.append(sdk_name)
-local_postfix = os.environ.get('LOCAL_VERSION')
-if local_postfix:
-    local_version.append(local_postfix)
-if len(local_version):
-    __version__ = '{}+{}'.format(__version__, '.'.join(local_version))
-
-packages = find_packages()
-megenginelite_data = [
-    str(f.relative_to('megenginelite'))
-    for f in pathlib.Path('megenginelite').glob('**/*')
-]
-
-if platform.system() == 'Windows':
-    megenginelite_data.remove('libs\\liblite_shared.dll')
-elif platform.system() == 'Darwin':
-    megenginelite_data.remove('libs/liblite_shared.dylib')
-else:
-    megenginelite_data.remove('libs/liblite_shared.so')
-
-with open('requires.txt') as f:
-    requires = f.read().splitlines()
-
-prebuild_modules=[PrecompiledExtesion('megenginelite.libs.liblite_shared')]
-setup_kwargs = dict(
-    name=package_name,
-    version=__version__,
-    description='Inference Framework for MegEngine',
-    author='Megvii Engine Team',
-    author_email=email,
-    packages=packages,
-    package_data={
-        'megenginelite': megenginelite_data,
-    },
-    ext_modules=prebuild_modules,
-    install_requires=requires,
-    cmdclass={'build_ext': build_ext},
-)
-setup_kwargs.update(dict(
-    classifiers=[
-    'Development Status :: 3 - Alpha',
-    'Intended Audience :: Developers',
-    'Intended Audience :: Education',
-    'Intended Audience :: Science/Research',
-    'License :: OSI Approved :: Apache Software License',
-    'Programming Language :: C++',
-    'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.5',
-    'Programming Language :: Python :: 3.6',
-    'Programming Language :: Python :: 3.7',
-    'Programming Language :: Python :: 3.8',
-    'Topic :: Scientific/Engineering',
-    'Topic :: Scientific/Engineering :: Mathematics',
-    'Topic :: Scientific/Engineering :: Artificial Intelligence',
-    'Topic :: Software Development',
-    'Topic :: Software Development :: Libraries',
-    'Topic :: Software Development :: Libraries :: Python Modules',
-    ],
-    license='Apache 2.0',
-    keywords='megengine deep learning',
-    data_files = [("megengine", [
-        "../LICENSE",
-        "../ACKNOWLEDGMENTS",
-    ])]
-))
-
-setup(**setup_kwargs)

scripts/whl/macos/macos_build_whl.sh

@@ -199,6 +199,17 @@ function do_build() {
         mkdir ${DEPEND_LIB}
         depend_real_copy ${DEPEND_LIB}
 
+        # handle megenginelite
+        cd ${BUILD_DIR}
+        mkdir -p staging/megenginelite
+        cp ${SRC_DIR}/lite/pylite/megenginelite/* staging/megenginelite/
+        mkdir -p ${BUILD_DIR}/staging/megenginelite/libs
+        LITE_LIB=${BUILD_DIR}/staging/megenginelite/libs/liblite_shared_whl.so
+        cp ${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_OFF/Release/build/lite/liblite_shared_whl.dylib ${LITE_LIB}
+        llvm-strip -s ${LITE_LIB}
+        #handle dlopen path
+        install_name_tool -change @rpath/libmegengine_export.dylib @loader_path/../../megengine/core/lib/libmegengine_export.dylib ${LITE_LIB}
+
         cd ${BUILD_DIR}/staging
         ${PYTHON_DIR}/bin/python3 setup.py bdist_wheel
         cd ${BUILD_DIR}/staging/dist/
@@ -209,35 +220,6 @@ function do_build() {
         echo "comapt whl name: ${compat_whl_name}"
         cp ${BUILD_DIR}/staging/dist/Meg*.whl ${MACOS_WHL_HOME}/${compat_whl_name}
 
-        # handle megenginelite
-        cd ${BUILD_DIR}
-        rm -rf lite_staging
-        mkdir -p lite_staging/megenginelite
-        cp ${SRC_DIR}/lite/pylite/megenginelite/* lite_staging/megenginelite/
-        cp ${SRC_DIR}/lite/pylite/setup.py lite_staging/
-        cp ${SRC_DIR}/lite/pylite/requires.txt lite_staging/
-        VER_FILE=${SRC_DIR}/imperative/python/megengine/version.py
-        if [ -f ${VER_FILE} ];then
-            cp ${VER_FILE} lite_staging/megenginelite
-        else
-            echo "ERROR: can not find version file"
-            exit -1
-        fi
-        mkdir -p ${BUILD_DIR}/lite_staging/megenginelite/libs
-        LITE_LIB=${BUILD_DIR}/lite_staging/megenginelite/libs/liblite_shared.dylib
-        cp ${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_OFF/Release/build/lite/liblite_shared.dylib ${LITE_LIB}
-        llvm-strip -s ${LITE_LIB}
-
-        cd ${BUILD_DIR}/lite_staging/
-        ${PYTHON_DIR}/bin/python3 setup.py bdist_wheel
-        cd ${BUILD_DIR}/lite_staging/dist/
-        org_whl_name=`ls Meg*.whl`
-        index=`awk -v a="${org_whl_name}" -v b="-macosx" 'BEGIN{print index(a,b)}'`
-        compat_whl_name=`echo ${org_whl_name} |cut -b -$index`macosx_10_14_x86_64.whl
-        echo "megenginelite org whl name: ${org_whl_name}"
-        echo "megenginelite comapt whl name: ${compat_whl_name}"
-        cp ${BUILD_DIR}/lite_staging/dist/Meg*.whl ${MACOS_WHL_HOME}/${compat_whl_name}
-
         cd ${SRC_DIR}
         echo ""
         echo "##############################################################################################"

scripts/whl/manylinux2014/do_build_common.sh

@@ -60,6 +60,17 @@ function patch_elf_depend_lib_mgb_mge() {
     handle_copy_cuda_libs ${LIBS_DIR}
 }
 
+function patch_elf_depend_lib_megenginelite() {
+    echo "handle common depend lib for megenginelite"
+    LIBS_DIR=${BUILD_DIR}/staging/megenginelite/libs
+    mkdir -p ${LIBS_DIR}
+
+    cp ${BUILD_DIR}/lite/liblite_shared_whl.so ${LIBS_DIR}/
+    patchelf --remove-rpath ${LIBS_DIR}/liblite_shared_whl.so
+    patchelf --force-rpath --set-rpath '$ORIGIN/../../megengine/core/lib' ${LIBS_DIR}/liblite_shared_whl.so
+    handle_strip ${LIBS_DIR}/liblite_shared_whl.so
+}
+
 SRC_DIR=$(readlink -f "`dirname $0`/../../../")
 source ${SRC_DIR}/scripts/whl/utils/utils.sh
 
@@ -144,42 +155,21 @@ do
     mkdir -p lib/ucx
     patch_elf_depend_lib_mgb_mge
 
-    cd ${BUILD_DIR}/staging/
-    ${PYTHON_DIR}/bin/python setup.py bdist_wheel
-    cd /home/output
-    mkdir -p ${SRC_DIR}/scripts/whl/manylinux2014/output/wheelhouse/${SDK_NAME}
-    cd ${BUILD_DIR}/staging/dist/
-    org_whl_name=`ls Meg*${ver}*.whl`
-    compat_whl_name=`echo ${org_whl_name} | sed 's/linux/manylinux2014/'`
-    echo "org whl name: ${org_whl_name}"
-    echo "comapt whl name: ${compat_whl_name}"
-    mv ${org_whl_name} ${SRC_DIR}/scripts/whl/manylinux2014/output/wheelhouse/${SDK_NAME}/${compat_whl_name}
-
     # handle megenginelite
     cd ${BUILD_DIR}
-    rm -rf lite_staging
-    mkdir -p lite_staging/megenginelite
-    cp ${SRC_DIR}/lite/pylite/megenginelite/* lite_staging/megenginelite/
-    cp ${SRC_DIR}/lite/pylite/setup.py lite_staging/
-    cp ${SRC_DIR}/lite/pylite/requires.txt lite_staging/
-    VER_FILE=${SRC_DIR}/imperative/python/megengine/version.py
-    if [ -f ${VER_FILE} ];then
-        cp ${VER_FILE} lite_staging/megenginelite
-    else
-        echo "ERROR: can not find version file"
-        exit -1
-    fi
+    mkdir -p staging/megenginelite
+    cp ${SRC_DIR}/lite/pylite/megenginelite/* staging/megenginelite/
     patch_elf_depend_lib_megenginelite
 
-    cd ${BUILD_DIR}/lite_staging/
+    cd ${BUILD_DIR}/staging/
     ${PYTHON_DIR}/bin/python setup.py bdist_wheel
     cd /home/output
     mkdir -p ${SRC_DIR}/scripts/whl/manylinux2014/output/wheelhouse/${SDK_NAME}
-    cd ${BUILD_DIR}/lite_staging/dist/
+    cd ${BUILD_DIR}/staging/dist/
     org_whl_name=`ls Meg*${ver}*.whl`
     compat_whl_name=`echo ${org_whl_name} | sed 's/linux/manylinux2014/'`
-    echo "megenginelite org whl name: ${org_whl_name}"
-    echo "megenginelite comapt whl name: ${compat_whl_name}"
+    echo "org whl name: ${org_whl_name}"
+    echo "comapt whl name: ${compat_whl_name}"
     mv ${org_whl_name} ${SRC_DIR}/scripts/whl/manylinux2014/output/wheelhouse/${SDK_NAME}/${compat_whl_name}
 
     cd /home/output

scripts/whl/windows/windows_build_whl.sh

@@ -108,13 +108,7 @@ function copy_more_dll() {
 }
 
 function lite_copy_more_dll() {
-    # for python whl real use
-    echo "config megenginelite core lib dir"
-    CP_WHL_DST_IMP=${BUILD_DIR}/lite_staging/megenginelite/libs
-
     if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
-        echo "copy nvidia lib to whl use...."
-        depend_real_copy ${CP_WHL_DST_IMP}
         if [ ${IN_CI} = "true" ]; then
             echo "copy lib for lite for ci test"
             IMP_TEST_DST=${SRC_DIR}/build_dir/host/build/lite/test/
@@ -208,36 +202,22 @@ function do_build() {
         mv ${rt_file} _imperative_rt.pyd
 
         copy_more_dll
-        cd ${BUILD_DIR}/staging
-        echo "call setup.py now"
-        ${PYTHON_DIR}/python3 setup.py bdist_wheel
-        cp ${BUILD_DIR}/staging/dist/Meg*.whl ${WINDOWS_WHL_HOME}/
 
         # handle megenginelite
         cd ${BUILD_DIR}
-        rm -rf lite_staging
-        mkdir -p lite_staging/megenginelite
-        cp ${SRC_DIR}/lite/pylite/megenginelite/* lite_staging/megenginelite/
-        cp ${SRC_DIR}/lite/pylite/setup.py lite_staging/
-        cp ${SRC_DIR}/lite/pylite/requires.txt lite_staging/
-        VER_FILE=${SRC_DIR}/imperative/python/megengine/version.py
-        if [ -f ${VER_FILE} ];then
-            cp ${VER_FILE} lite_staging/megenginelite
-        else
-            echo "ERROR: can not find version file"
-            exit -1
-        fi
-
-        LITE_CORE_LIB_DIR=${BUILD_DIR}/lite_staging/megenginelite/libs/
+        mkdir -p staging/megenginelite
+        cp ${SRC_DIR}/lite/pylite/megenginelite/* staging/megenginelite/
+        LITE_CORE_LIB_DIR=${BUILD_DIR}/staging/megenginelite/libs/
         mkdir -p ${LITE_CORE_LIB_DIR}
         cd ${LITE_CORE_LIB_DIR}
-        cp ${BUILD_DIR}/lite/lite_shared.dll liblite_shared.dll
-        llvm-strip -s liblite_shared.dll
+        cp ${BUILD_DIR}/lite/lite_shared_whl.dll liblite_shared_whl.pyd
+        llvm-strip -s liblite_shared_whl.pyd
         lite_copy_more_dll
 
-        cd ${BUILD_DIR}/lite_staging/
+        cd ${BUILD_DIR}/staging
+        echo "call setup.py now"
         ${PYTHON_DIR}/python3 setup.py bdist_wheel
-        cp ${BUILD_DIR}/lite_staging/dist/Meg*.whl ${WINDOWS_WHL_HOME}/
+        cp ${BUILD_DIR}/staging/dist/Meg*.whl ${WINDOWS_WHL_HOME}/
 
         echo ""
         echo "##############################################################################################"

src/CMakeLists.txt

@@ -202,17 +202,15 @@ set (_VER_FILE ${PROJECT_SOURCE_DIR}/src/version.ld)
 # depends on megdnn/megbrain target, refs to sdk/load-and-run/CMakeLists.txt
 # depends on megengine lite_share or lite_static
 if(NOT WIN32)
-    if(MGE_BUILD_IMPERATIVE_RT
-            )
-        message(VERBOSE "create a export SHARED lib for python use")
-        add_library(megengine_export SHARED)
-        target_link_libraries(megengine_export PUBLIC megbrain megdnn)
-        target_link_libraries(megengine_export PRIVATE ${MGE_CUDA_LIBS})
-        if (MGE_WITH_DISTRIBUTED)
-            message(VERBOSE "megengine_export configured to link megray")
-            target_link_libraries(megengine_export PUBLIC megray)
-        endif()
+    message(VERBOSE "create a export SHARED lib for python use")
+    add_library(megengine_export SHARED)
+    target_link_libraries(megengine_export PUBLIC megbrain megdnn)
+    target_link_libraries(megengine_export PRIVATE ${MGE_CUDA_LIBS})
+    if (MGE_WITH_DISTRIBUTED)
+        message(VERBOSE "megengine_export configured to link megray")
+        target_link_libraries(megengine_export PUBLIC megray)
     endif()
+
     # Build as SHARED or STATIC depending on BUILD_SHARED_LIBS=ON/OFF
     add_library(megengine)
     target_link_libraries(megengine PRIVATE ${MGE_CUDA_LIBS})