Merge branch 'release_candidate'

1 files changed, 45 insertions, 9 deletions

diff --git a/webui.sh b/webui.sh
index cff43327..f116376f 100755
--- a/webui.sh
+++ b/webui.sh
@@ -133,7 +133,7 @@ case "$gpu_info" in
             if [[ $(bc <<< "$pyv <= 3.10") -eq 1 ]] 
             then
                 # Navi users will still use torch 1.13 because 2.0 does not seem to work.
-                export TORCH_COMMAND="pip install torch==1.13.1+rocm5.2 torchvision==0.14.1+rocm5.2 --index-url https://download.pytorch.org/whl/rocm5.2"
+                export TORCH_COMMAND="pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm5.6"
             else
                 printf "\e[1m\e[31mERROR: RX 5000 series GPUs must be using at max python 3.10, aborting...\e[0m"
                 exit 1
@@ -143,8 +143,7 @@ case "$gpu_info" in
     *"Navi 2"*) export HSA_OVERRIDE_GFX_VERSION=10.3.0
     ;;
     *"Navi 3"*) [[ -z "${TORCH_COMMAND}" ]] && \
-         export TORCH_COMMAND="pip install torch torchvision --index-url https://download.pytorch.org/whl/test/rocm5.6"
-        # Navi 3 needs at least 5.5 which is only on the torch 2.1.0 release candidates right now
+         export TORCH_COMMAND="pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm5.7"
     ;;
     *"Renoir"*) export HSA_OVERRIDE_GFX_VERSION=9.0.0
         printf "\n%s\n" "${delimiter}"
@@ -159,6 +158,10 @@ then
     if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
     then
         export TORCH_COMMAND="pip install torch==2.0.1+rocm5.4.2 torchvision==0.15.2+rocm5.4.2 --index-url https://download.pytorch.org/whl/rocm5.4.2"
+    elif echo "$gpu_info" | grep -q "Huawei" && [[ -z "${TORCH_COMMAND}" ]]
+    then
+        export TORCH_COMMAND="pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; pip install torch_npu"
+    
     fi
 fi
 
@@ -223,12 +226,45 @@ fi
 # Try using TCMalloc on Linux
 prepare_tcmalloc() {
     if [[ "${OSTYPE}" == "linux"* ]] && [[ -z "${NO_TCMALLOC}" ]] && [[ -z "${LD_PRELOAD}" ]]; then
-        TCMALLOC="$(PATH=/sbin:$PATH ldconfig -p | grep -Po "libtcmalloc(_minimal|)\.so\.\d" | head -n 1)"
-        if [[ ! -z "${TCMALLOC}" ]]; then
-            echo "Using TCMalloc: ${TCMALLOC}"
-            export LD_PRELOAD="${TCMALLOC}"
-        else
-            printf "\e[1m\e[31mCannot locate TCMalloc (improves CPU memory usage)\e[0m\n"
+        # check glibc version
+        LIBC_VER=$(echo $(ldd --version | awk 'NR==1 {print $NF}') | grep -oP '\d+\.\d+')
+        echo "glibc version is $LIBC_VER"
+        libc_vernum=$(expr $LIBC_VER)
+        # Since 2.34 libpthread is integrated into libc.so
+        libc_v234=2.34
+        # Define Tcmalloc Libs arrays
+        TCMALLOC_LIBS=("libtcmalloc(_minimal|)\.so\.\d" "libtcmalloc\.so\.\d")
+        # Traversal array
+        for lib in "${TCMALLOC_LIBS[@]}"
+        do
+            # Determine which type of tcmalloc library the library supports
+            TCMALLOC="$(PATH=/usr/sbin:$PATH ldconfig -p | grep -P $lib | head -n 1)"
+            TC_INFO=(${TCMALLOC//=>/})
+            if [[ ! -z "${TC_INFO}" ]]; then
+                echo "Check TCMalloc: ${TC_INFO}"
+                # Determine if the library is linked to libpthread and resolve undefined symbol: pthread_key_create
+                if [ $(echo "$libc_vernum < $libc_v234" | bc) -eq 1 ]; then
+                    # glibc < 2.34 pthread_key_create into libpthread.so. check linking libpthread.so...
+                    if ldd ${TC_INFO[2]} | grep -q 'libpthread'; then
+                        echo "$TC_INFO is linked with libpthread,execute LD_PRELOAD=${TC_INFO[2]}"
+                        # set fullpath LD_PRELOAD (To be on the safe side)
+                        export LD_PRELOAD="${TC_INFO[2]}"
+                        break
+                    else
+                        echo "$TC_INFO is not linked with libpthread will trigger undefined symbol: pthread_Key_create error"
+                    fi
+                else
+                    # Version 2.34 of libc.so (glibc) includes the pthread library IN GLIBC. (USE ubuntu 22.04 and modern linux system and WSL)
+                    # libc.so(glibc) is linked with a library that works in ALMOST ALL Linux userlands. SO NO CHECK!
+                    echo "$TC_INFO is linked with libc.so,execute LD_PRELOAD=${TC_INFO[2]}"
+                    # set fullpath LD_PRELOAD (To be on the safe side)
+                    export LD_PRELOAD="${TC_INFO[2]}"
+                    break
+                fi
+            fi
+        done
+        if [[ -z "${LD_PRELOAD}" ]]; then
+            printf "\e[1m\e[31mCannot locate TCMalloc. Do you have tcmalloc or google-perftool installed on your system? (improves CPU memory usage)\e[0m\n"
         fi
     fi
 }