File size: 1,313 Bytes
9aa90e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | #!/usr/bin/env bash
# SPDX-FileCopyrightText: (c) 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
#
# Builds the optional _planar_concat extension consumed by
# models/tt_dit/utils/planar_concat.py. Absence of the .so is not an error:
# the caller falls back to the torch scatter path.
set -euo pipefail
SRC_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUILD_DIR="${SRC_DIR}/build"
PYTHON="${PYTHON:-python3}"
if ! grep -q '\bavx2\b' /proc/cpuinfo; then
echo "build.sh: this CPU reports no AVX2 support; the extension would fault at runtime" >&2
exit 1
fi
PY_INCLUDE="$("${PYTHON}" -c 'import sysconfig; print(sysconfig.get_paths()["include"])')"
NUMPY_INCLUDE="$("${PYTHON}" -c 'import numpy; print(numpy.get_include())')"
EXT_SUFFIX="$("${PYTHON}" -c 'import sysconfig; print(sysconfig.get_config_var("EXT_SUFFIX"))')"
mkdir -p "${BUILD_DIR}"
OUT="${BUILD_DIR}/_planar_concat${EXT_SUFFIX}"
# -mavx2 covers the streaming stores in planar_concat.cpp as well as the transposes.
${CXX:-g++} -O3 -std=c++17 -fPIC -shared -mavx2 -fvisibility=hidden \
-I"${SRC_DIR}" -I"${PY_INCLUDE}" -I"${NUMPY_INCLUDE}" \
"${SRC_DIR}/transpose_avx2.cpp" \
"${SRC_DIR}/planar_concat.cpp" \
"${SRC_DIR}/planar_concat_bindings.cpp" \
-o "${OUT}"
echo "build.sh: wrote ${OUT}"
|