98 lines
3.0 KiB
Bash
98 lines
3.0 KiB
Bash
#!/usr/bin/env bash
|
|
# build_bridge_fused_ar.sh — Compile ix_full_bridge_fused_ar.so with linear_allreduce support
|
|
#
|
|
# Usage (on BI-V100 machine):
|
|
# bash ex_engine/build_bridge_fused_ar.sh
|
|
#
|
|
# Produces: ex_engine/prebuilt/ix_full_bridge_fused_ar.so
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
SRC="${SCRIPT_DIR}/csrc/ix_full_bridge_fused_ar.cpp"
|
|
OUTPUT_DIR="${SCRIPT_DIR}/prebuilt"
|
|
mkdir -p "$OUTPUT_DIR"
|
|
OUTPUT="${OUTPUT_DIR}/ix_full_bridge_fused_ar.so"
|
|
|
|
COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
|
|
PYTHON="${PYTHON:-python3}"
|
|
|
|
# --- Compiler ---
|
|
CLANGXX="${COREX_ROOT}/bin/clang++"
|
|
if [[ ! -x "$CLANGXX" ]]; then
|
|
CLANGXX=$(command -v clang++ 2>/dev/null || true)
|
|
fi
|
|
if [[ -z "$CLANGXX" ]]; then
|
|
echo "[fused_ar] ERROR: clang++ not found" >&2
|
|
exit 1
|
|
fi
|
|
|
|
# --- torch include/lib paths ---
|
|
TORCH_INC=$($PYTHON -c "from torch.utils.cpp_extension import include_paths; print(' '.join(['-I'+p for p in include_paths()]))")
|
|
TORCH_LIB=$($PYTHON -c "from torch.utils.cpp_extension import library_paths; print(' '.join(['-L'+p for p in library_paths()]))")
|
|
PYTHON_INC=$($PYTHON -c "from sysconfig import get_paths; print('-I' + get_paths()['include'])")
|
|
|
|
# --- Find _ixformer_torch.so (has ixformer_linear_allreduce) ---
|
|
IX_TORCH_SO=""
|
|
RPATH_DIRS=""
|
|
for d in \
|
|
"${COREX_ROOT}/lib/python3/dist-packages/ixformer" \
|
|
"${COREX_ROOT}/lib64/python3/dist-packages/ixformer" \
|
|
"/usr/local/lib/python3.10/dist-packages/ixformer"; do
|
|
if [[ -d "$d" ]]; then
|
|
for so in "$d"/_ixformer_torch*.so; do
|
|
if [[ -f "$so" ]]; then
|
|
IX_TORCH_SO="$so"
|
|
RPATH_DIRS="${RPATH_DIRS} -Wl,-rpath,$d"
|
|
break 2
|
|
fi
|
|
done
|
|
fi
|
|
done
|
|
|
|
if [[ -z "$IX_TORCH_SO" ]]; then
|
|
echo "[fused_ar] ERROR: _ixformer_torch.so not found — cannot link linear_allreduce" >&2
|
|
exit 1
|
|
fi
|
|
|
|
# Also link libixformer.so for allreduce internals
|
|
IX_LIBS="$IX_TORCH_SO"
|
|
for so in "${COREX_ROOT}/lib64"/libixformer*.so; do
|
|
if [[ -f "$so" ]]; then
|
|
IX_LIBS="${IX_LIBS} $so"
|
|
RPATH_DIRS="${RPATH_DIRS} -Wl,-rpath,${COREX_ROOT}/lib64"
|
|
fi
|
|
done
|
|
|
|
# --- torch lib rpath (so libc10.so / libtorch.so are found at runtime) ---
|
|
TORCH_LIB_DIR=$($PYTHON -c "from torch.utils.cpp_extension import library_paths; print(library_paths()[0])")
|
|
RPATH_DIRS="${RPATH_DIRS} -Wl,-rpath,${TORCH_LIB_DIR}"
|
|
|
|
echo "[fused_ar] Source: $SRC"
|
|
echo "[fused_ar] Compiler: $CLANGXX"
|
|
echo "[fused_ar] Link: $IX_LIBS"
|
|
echo "[fused_ar] Rpath: $RPATH_DIRS"
|
|
echo "[fused_ar] Output: $OUTPUT"
|
|
|
|
$CLANGXX \
|
|
-shared -fPIC -O2 -std=c++17 \
|
|
-DTORCH_EXTENSION_NAME=ix_full_bridge_fused_ar \
|
|
-D_GLIBCXX_USE_CXX11_ABI=0 \
|
|
$PYTHON_INC \
|
|
$TORCH_INC \
|
|
$TORCH_LIB \
|
|
-L"${TORCH_LIB_DIR}" \
|
|
-ltorch -ltorch_cpu -ltorch_python -lc10 \
|
|
${IX_LIBS} \
|
|
${RPATH_DIRS} \
|
|
-o "$OUTPUT" \
|
|
"$SRC"
|
|
|
|
echo "[fused_ar] ✓ Built: $OUTPUT"
|
|
ls -lh "$OUTPUT"
|
|
|
|
# Quick symbol check
|
|
echo "[fused_ar] Exported symbols:"
|
|
nm -D "$OUTPUT" | c++filt | grep ' T ' | grep -v __device
|
|
|
|
echo "[fused_ar] Done" |