| # 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}" | |