#!/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"