#!/bin/bash
set -e

echo "============================================"
echo " PyTorch 2.5.1 Build for gfx1010 (BC-250)"
echo " NO FLASH/MEM_EFF ATTENTION = NO AOTRITON"
echo "============================================"

cd ~/pytorch-build

# --- Phase 1: Clean ---
echo "[Phase 1] Clean build dir..."
rm -rf build dist
echo "[Phase 1] Done."

# --- Phase 2: Environment ---
echo "[Phase 2] Setting environment..."
export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0
export ROCM_PATH=/opt/rocm
export PATH="/opt/rocm/bin:/home/fabian/.local/bin:$PATH"
export LD_LIBRARY_PATH="/opt/rocm/lib:$LD_LIBRARY_PATH"

# PyTorch build config
export USE_ROCM=1
export USE_CUDA=0
export USE_NCCL=0
export USE_DISTRIBUTED=0
export USE_MKLDNN=0
export USE_FBGEMM=0
export USE_NNPACK=0
export USE_QNNPACK=0
export USE_XNNPACK=0
export USE_KINETO=0
export BUILD_TEST=0
export BUILD_CAFFE2=0
export USE_NUMPY=1
# Disable Flash/MemEff attention -> no aotriton dependency
export USE_FLASH_ATTENTION=0
export USE_MEM_EFF_ATTENTION=0
export PYTORCH_ROCM_ARCH="gfx1010"
export HIP_PATH=/opt/rocm
export ROCM_SOURCE_DIR=/opt/rocm
export CMAKE_PREFIX_PATH=/opt/rocm
export MAX_JOBS=8
export CMAKE_POLICY_VERSION_MINIMUM=3.5

echo "  PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
echo "  USE_FLASH_ATTENTION=0"
echo "  USE_MEM_EFF_ATTENTION=0"
echo "  MAX_JOBS=$MAX_JOBS"
echo "[Phase 2] Done."

# --- Phase 3: Build ---
echo "[Phase 3] Starting build at $(date)..."
python3.11 setup.py bdist_wheel 2>&1 | tee ~/pytorch-build-progress.log

echo ""
echo "============================================"
echo " BUILD FINISHED at $(date)"
echo "============================================"
ls -lh dist/*.whl 2>/dev/null || echo "ERROR: No wheel produced!"