kse-01/tensorflow/third_party/toolchains/preconfig/generate/generate.sh
github-classroom[bot] 1122cdd8b0
Initial commit
2023-10-09 11:37:31 +00:00

110 lines
3 KiB (Stored with Git LFS)
Bash

#!/bin/bash
#
# Copyright 2018 The TensorFlow 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.
# ==============================================================================
TARGET="$1"
OUTPUT="$2"
if [[ -z "${TARGET}" || -z "${OUTPUT}" ]]; then
echo "Usage:"
echo "$0 <target> <output>"
exit 1
fi
TEMPDIR="$(mktemp -d)"
ROOT="${PWD}"
PKG="third_party/toolchains/preconfig"
IFS='-' read -ra PLATFORM <<< "${TARGET}"
OS="${PLATFORM[0]}"
PY_VERSION="${PLATFORM[1]}"
COMPILER="${PLATFORM[2]}"
GPU_VERSION="${PLATFORM[3]}"
CUDNN_VERSION="${PLATFORM[4]}"
TENSORRT_VERSION="${PLATFORM[5]}"
# TODO(klimek): Put this into the name.
if [[ "${GPU_VERSION}" == "rocm" ]]; then
COMPILER="${COMPILER}"
elif [[ -n "${GPU_VERSION}" ]]; then
if [[ "${COMPILER}" == gcc* ]]; then
COMPILER="${COMPILER}-nvcc-${GPU_VERSION}"
fi
# Currently we create a special toolchain for clang when compiling with
# cuda enabled. We can get rid of this once the default toolchain bazel
# provides supports cuda.
if [[ "${COMPILER}" == clang* ]]; then
COMPILER="${COMPILER}-${GPU_VERSION}"
fi
fi
echo "OS: ${OS}"
echo "Python: ${PY_VERSION}"
echo "Compiler: ${COMPILER}"
echo "CUDA/ROCm: ${GPU_VERSION}"
echo "CUDNN: ${CUDNN_VERSION}"
echo "TensorRT: ${TENSORRT_VERSION}"
bazel build --host_force_python=PY2 --define=mount_project="${PWD}" \
"${PKG}/generate:${TARGET}"
cd "${TEMPDIR}"
tar xvf "${ROOT}/bazel-bin/${PKG}/generate/${TARGET}_outputs.tar"
# Delete all empty files: configurations leave empty files around when they are
# unnecessary.
find . -empty -delete
# We build up the following directory structure with preconfigured packages:
# <OS>/
# <CUDA>-<CUDNN>/
# <COMPILER>/
# <PYTHON>/
# <TENSORRT>/
# Create our toplevel output directory for the OS.
mkdir "${OS}"
# Python:
mv local_config_python "${OS}/${PY_VERSION}"
if [[ "${GPU_VERSION}" == "rocm" ]]; then
# Compiler:
mv local_config_rocm/crosstool "${OS}/${COMPILER}-${GPU_VERSION}"
# ROCm:
mv local_config_rocm "${OS}/${GPU_VERSION}"
elif [[ -n "${GPU_VERSION}" ]]; then
# Compiler:
mv local_config_cuda/crosstool "${OS}/${COMPILER}"
# CUDA:
mv local_config_cuda "${OS}/${GPU_VERSION}-${CUDNN_VERSION}"
# TensorRT:
mv local_config_tensorrt "${OS}/${TENSORRT_VERSION}"
else
# Compiler:
mv local_config_cc "${OS}/${COMPILER}"
fi
# Cleanup for copybara.
find "${OS}" -name '*.h' |xargs clang-format -i
find "${OS}" -name 'BUILD' -o -name '*.bzl' |xargs buildifier
find "${OS}" -name 'BUILD' -o -name '*.bzl' |xargs -I {} mv {} {}.oss
# Tar it up:
tar cvf "${OUTPUT}" "${OS}"