Revert "Revert "using lld linker" (#1008)"

docker/dev/Dockerfile

@@ -20,10 +20,11 @@ RUN apt-get update -y && apt-get install -y git curl vim libssl-dev wget unzip o
 RUN bash /install/scripts/install-cmake.sh
 RUN bash /install/scripts/install-bazel.sh
 RUN bash /install/scripts/install-python.sh
+RUN bash /install/scripts/install-llvm_toolchain.sh
 
 ARG DEVICE=cu110
 RUN bash /install/scripts/install-tensorrt.sh
 RUN bash /install/scripts/install-cudnn.sh
 
-ENV PATH="/opt/cmake/bin:${PATH}"
-ENV LD_LIBRARY_PATH="/usr/local/TensorRT/lib/:/usr/local/cuda/lib64/:${LD_LIBRARY_PATH}"
+ENV PATH="/usr/local/llvm_toolchain/bin:/opt/cmake/bin:${PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/llvm_toolchain/lib:/usr/local/TensorRT/lib/:/usr/local/cuda/lib64/:${LD_LIBRARY_PATH}"

docker/dev/Dockerfile.aarch64

@@ -84,7 +84,8 @@ RUN apt-get -y update && \
       zip \
       zlib1g-dev \
       openjdk-11-jdk \
-      patchelf
+      patchelf \
+      lld
 
 RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 1 && \
     update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-10 1 && \

docker/scripts/install-llvm_toolchain.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Copyright 2022 The BladeDISC Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+function install_llvm_toolchain() {
+    llvm_version=clang+llvm-10.0.1-x86_64-linux-gnu-ubuntu-16.04
+    llvm_tgz=${llvm_version}".tar.xz"
+    local_tgz="/install/${llvm_tgz}"
+    llvm_url="https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.1/${llvm_tgz}"
+
+    echo "Download url: ${llvm_url}"
+
+    wget -nv ${llvm_url} -O ${local_tgz}
+    if [[ "$?" != "0" ]]; then
+        echo "Download clang+llvm-12.0.1 failed."
+        exit -1
+    fi
+
+    tar -xf ${local_tgz} --skip-old-files -C /usr/local
+    mv /usr/local/${llvm_version} /usr/local/llvm_toolchain
+    rm -f ${local_tgz}
+}
+
+get_arch=`arch`
+if [[ $get_arch =~ "x86_64" ]]; then
+    install_llvm_toolchain
+else
+    exit 0
+fi

pytorch_blade/bazel_build.py

@@ -153,6 +153,8 @@ def __init__(self, *args, **kwargs):
         if running_on_ci():
             self.configs += ["--config=ci_build"]
 
+        self.configs += ["--linkopt='-fuse-ld=gold -Xlinker --stub-group-size -Xlinker 10000000'"]
+
         root_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir)
         self.configs += [
             build_tao_compiler_add_flags_platform_alibaba_cached(root_dir, ""),

scripts/python/common_setup.py

@@ -180,6 +180,15 @@ def which(cmd):
         raise Exception("failed to find command: " + cmd)
     return found
 
+def extra_link_flags(force_gold=True):
+    from distutils.spawn import find_executable
+    if find_executable("ld.lld"):
+        return " --linkopt='-fuse-ld=lld -Xlinker -lm'"
+    else:
+        if force_gold:
+            return " --linkopt='-fuse-ld=gold -Xlinker --stub-group-size -Xlinker 10000000'"
+        else:
+            return " --linkopt='-Xlinker --stub-group-size -Xlinker 10000000'"
 
 @contextmanager
 def cwd(path):

scripts/python/tao_build.py

@@ -42,6 +42,7 @@
     add_ral_link_if_not_exist,
     logger,
     which,
+    extra_link_flags,
     running_on_ci,
     ci_build_flag,
     remote_cache_token,
@@ -280,6 +281,7 @@ def configure(root, args):
 @time_stage()
 def build_tao_compiler(root, args):
     BAZEL_BUILD_CMD = "bazel build --verbose_failures --experimental_multi_threaded_digest --define framework_shared_object=false" + ci_build_flag()
+    BAZEL_BUILD_CMD += extra_link_flags()
     TARGET_TAO_COMPILER_MAIN = "//decoupling:tao_compiler_main"
     TARGET_DISC_OPT = "//mlir/disc:disc-opt"
     TARGET_DISC_REPLAY = "//mlir/disc/tools/disc-replay:disc-replay-main"
@@ -354,6 +356,8 @@ def test_tao_compiler(root, args):
     BAZEL_TEST_CMD = "bazel test --experimental_multi_threaded_digest --define framework_shared_object=false --test_timeout=600 --javabase=@bazel_tools//tools/jdk:remote_jdk11"
     BAZEL_TEST_CMD += ci_build_flag()
     BAZEL_BUILD_CMD += ci_build_flag()
+    BAZEL_TEST_CMD += extra_link_flags()
+    BAZEL_BUILD_CMD += extra_link_flags()
     if running_on_ci():
         # NOTE: using the lower parallel jobs on CI host to avoid OOM
         BAZEL_TEST_CMD += " --jobs=10 --test_output=errors"
@@ -472,6 +476,7 @@ def tao_bridge_bazel_config(args):
             bazel_config += " --config=disc_mkldnn"
     if args.platform_alibaba:
         bazel_config += " --config=platform_alibaba"
+    bazel_config += extra_link_flags()
     return bazel_config
 
 @time_stage()

tao/.bazelrc

@@ -37,7 +37,7 @@ build:disc --config=release_base
 
 build:disc_cpu --config=disc --cxxopt=-DTAO_CPU_ONLY
 build:disc_x86 --config=disc_cpu --config=release_cpu_linux --cxxopt=-DTAO_X86 --define disc_x86=true
-build:disc_aarch64 --config=disc_cpu --cxxopt=-DTAO_AARCH64 --define disc_aarch64=true --linkopt="-Xlinker --stub-group-size -Xlinker 10000000"
+build:disc_aarch64 --config=disc_cpu --cxxopt=-DTAO_AARCH64 --define disc_aarch64=true
 build:disc_cuda --config=disc --config=cuda
 build:disc_dcu --config=disc --config=dcu
 build:disc_rocm --config=disc --config=rocm

tao_compiler/.bazelrc.user

@@ -3,7 +3,7 @@ build:ci_build --noshow_loading_progress --show_progress_rate_limit=600 --jobs=1
 build:disc --define framework_shared_object=false --experimental_multi_threaded_digest
 build:disc_cpu --config=disc --cxxopt=-DTAO_CPU_ONLY
 build:disc_x86 --config=disc_cpu --config=release_cpu_linux --cxxopt=-DTAO_X86 --define disc_x86=true
-build:disc_aarch64 --config=disc_cpu --config=mkl_aarch64 --cxxopt=-DTAO_AARCH64 --define disc_aarch64=true --linkopt="-Xlinker --stub-group-size -Xlinker 10000000"
+build:disc_aarch64 --config=disc_cpu --config=mkl_aarch64 --cxxopt=-DTAO_AARCH64 --define disc_aarch64=true
 build:disc_cuda --config=disc --config=cuda
 build:disc_dcu --config=disc --config=dcu
 

tensorflow_blade/.bazelrc

@@ -67,7 +67,6 @@ build:disc_aarch64 --config=disc_cpu
 build:disc_aarch64 --config=mkl_aarch64
 build:disc_aarch64 --cxxopt=-DTAO_AARCH64
 build:disc_aarch64 --define disc_aarch64=true
-build:disc_aarch64 --linkopt="-Xlinker --stub-group-size -Xlinker 10000000"
 build:disc_aarch64 --action_env BUILD_WITH_AARCH64=1
 build:disc_cuda --config=disc --config=cuda
 build:disc_dcu --config=disc --config=dcu