fix(whl/mgb/imperative): fix symbols conflict runtime crash

GitOrigin-RevId: d60c4052c1

imperative/CMakeLists.txt

@@ -52,15 +52,15 @@ add_custom_target(_version_ld SOURCES ${VERSION_SCRIPT})
 
 add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/pybind11 ${PROJECT_BINARY_DIR}/third_party/pybind11)
 pybind11_add_module(${MODULE_NAME} NO_EXTRAS ${SRCS})
-if (APPLE OR MSVC OR WIN32)
+if (APPLE)
+    target_link_libraries(${MODULE_NAME} PRIVATE megengine_export)
+elseif (MSVC OR WIN32)
+    # Windows does not support implicitly importing data members from DLL.
     target_link_libraries(${MODULE_NAME} PRIVATE megbrain megdnn)
 else()
-    target_link_libraries(${MODULE_NAME} PRIVATE megbrain megdnn -Wl,--version-script=${VERSION_SCRIPT})
-endif()
-if (MGE_WITH_DISTRIBUTED)
-    message("Imperative configured to link megray")
-    target_link_libraries(${MODULE_NAME} PRIVATE megray)
+    target_link_libraries(${MODULE_NAME} PRIVATE megengine_export -Wl,--version-script=${VERSION_SCRIPT})
 endif()
+
 target_include_directories(${MODULE_NAME} PUBLIC src/include PRIVATE ${PYTHON_INCLUDE_DIRS} ${NUMPY_INCLUDE_DIR})
 target_compile_definitions(${MODULE_NAME} PRIVATE MODULE_NAME=${MODULE_NAME})
 target_compile_options(${MODULE_NAME} PRIVATE -Wno-unused-parameter)

imperative/python/test/run.sh

@@ -8,7 +8,7 @@ if [[ "$TEST_PLAT" == cpu ]]; then
 elif [[ "$TEST_PLAT" == cuda ]]; then
     echo "test both cpu and gpu pytest"
 else
-    log "Argument must cpu or cuda"
+    echo "Argument must cpu or cuda"
     exit 1
 fi
 

scripts/cmake-build/utils/utils.sh

@@ -15,26 +15,15 @@ function build_flatc() {
     REMOVE_OLD_BUILD=$2
 
     if [ $REMOVE_OLD_BUILD = "true" ]; then
-        if [ -e $BUILD_DIR ];then
-            echo "clean old dir: $BUILD_DIR"
-            rm -rf $BUILD_DIR
-        fi
-        if [ -e $INSTALL_DIR ];then
-            echo "clean old dir: $INSTALL_DIR"
-            rm -rf $INSTALL_DIR
-        fi
+        echo "remove old build/install dir"
+        rm -rf $INSTALL_DIR
+        rm -rf $BUILD_DIR
     else
         echo "strip remove old build"
     fi
 
-    if [ ! -e $BUILD_DIR ];then
-        echo "create build: $BUILD_DIR"
-        mkdir -p $BUILD_DIR
-    fi
-    if [ ! -e $INSTALL_DIR ];then
-        echo "create install $INSTALL_DIR"
-        mkdir -p $INSTALL_DIR
-    fi
+    mkdir -p $BUILD_DIR
+    mkdir -p $INSTALL_DIR
 
     cd $BUILD_DIR
     cmake -G "$MAKEFILE_TYPE Makefiles" \
@@ -57,14 +46,9 @@ function try_remove_old_build() {
     INSTALL_DIR=$3
 
     if [ $REMOVE_OLD_BUILD = "true" ]; then
-        if [ -e $BUILD_DIR ];then
-            echo "clean old dir: $BUILD_DIR"
-            rm -rf $BUILD_DIR
-        fi
-        if [ -e $INSTALL_DIR ];then
-            echo "clean old dir: $INSTALL_DIR"
-            rm -rf $INSTALL_DIR
-        fi
+        echo "remove old build/install dir"
+        rm -rf ${BUILD_DIR}
+        rm -rf ${INSTALL_DIR}
     else
         echo "strip remove old build"
     fi

scripts/whl/macos/macos_build_whl.sh

@@ -84,6 +84,13 @@ function config_python_env() {
     fi
 }
 
+MEGENGINE_LIB="${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_OFF/Release/build/src/libmegengine_export.dylib"
+function depend_real_copy() {
+    REAL_DST=$1
+    echo "real copy lib to $1"
+    cp "${MEGENGINE_LIB}" ${REAL_DST}
+}
+
 function do_build() {
     for ver in ${ALL_PYTHON}
     do
@@ -142,6 +149,15 @@ function do_build() {
             echo "valid..."
         fi
 
+        #handle dlopen path
+        install_name_tool -change @rpath/libmegengine_export.dylib @loader_path/lib/libmegengine_export.dylib _imperative_rt.so
+
+
+        #copy megbrain_export lib
+        DEPEND_LIB=${BUILD_DIR}/staging/megengine/core/lib/
+        rm -rf ${DEPEND_LIB}
+        mkdir ${DEPEND_LIB}
+        depend_real_copy ${DEPEND_LIB}
 
         cd ${BUILD_DIR}/staging
         ${PYTHON_DIR}/bin/python3 setup.py bdist_wheel

scripts/whl/manylinux2010/do_build.sh

@@ -16,17 +16,22 @@ BUILD_DIR=${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_OFF/Rel
 if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
     BUILD_DIR=${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_ON/MGE_INFERENCE_ONLY_OFF/Release/build/
 fi
-SO_NAME=_imperative_rt
-SO_PATH=megengine/core
 NEW_LIB_PATH=core/lib
 
+function handle_strip() {
+    echo "now handle strip $1"
+    objcopy --only-keep-debug $1 $1.dbg
+    strip -s $1
+    objcopy --add-gnu-debuglink=$1.dbg $1
+    rm $1.dbg
+}
+
 for ver in ${ALL_PYTHON}
 do
     python_ver=${ver:0:2}
     MAJOR=${python_ver:0:1}
     MINOR=${ver:1}
     PYTHON_DIR=/opt/python/cp${python_ver}-cp${ver}/
-    EXT_NAME=${SO_NAME}.cpython-${ver}-x86_64-linux-gnu.so
     export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DCMAKE_BUILD_TYPE=RelWithDebInfo"
     export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DCMAKE_PREFIX_PATH=${PYTHON_DIR}"
     export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_EXECUTABLE=${PYTHON_DIR}/bin/python3"
@@ -51,17 +56,15 @@ do
         ln -sf libcuda.so libcuda.so.1
     fi
 
-    cd ${BUILD_DIR}/staging/${SO_PATH}
-    SO_NAME_EXT=${SO_NAME}.so
-    objcopy --only-keep-debug ${SO_NAME_EXT} ${EXT_NAME}.dbg
-    strip -s ${SO_NAME_EXT}
-    objcopy --add-gnu-debuglink=${EXT_NAME}.dbg ${SO_NAME_EXT}
+    handle_strip ${BUILD_DIR}/src/libmegengine_export.so
+
+    cd ${BUILD_DIR}/staging/megengine/core
+    handle_strip _imperative_rt.so
+
     mkdir -p lib/ucx
 
     if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
         cp -L /usr/local/cuda/lib*/libnvrtc-builtins.so lib
-        cp -L ${BUILD_DIR}/third_party/MegRay/third_party/ucx/lib/ucx/*.so lib/ucx/
-        strip -s lib/ucx/*.so
     fi
 
     cd ${BUILD_DIR}/staging/

src/CMakeLists.txt

@@ -174,6 +174,16 @@ else()
     set (_VER_FILE ${PROJECT_SOURCE_DIR}/imperative/src/version.ld)
 endif()
 
+if(MGE_BUILD_IMPERATIVE_RT
+        )
+    message("-- create a export SHARED lib for python use")
+    add_library(megengine_export SHARED)
+    target_link_libraries(megengine_export PUBLIC megbrain megdnn)
+    if (MGE_WITH_DISTRIBUTED)
+        message("megengine_export configured to link megray")
+        target_link_libraries(megengine_export PUBLIC megray)
+    endif()
+endif()
 # Build as SHARED or STATIC depending on BUILD_SHARED_LIBS=ON/OFF
 add_library(megengine)
 target_link_libraries(megengine PUBLIC megbrain megdnn)