#!/bin/bash
# =============================================================
# PyTorch 2.5.1 Build for AMD BC-250 — native gfx1010 target
# ROCm 7.2 on CachyOS
# =============================================================
set -euo pipefail

echo "=========================================="
echo " PyTorch Build — gfx1010 for BC-250"
echo "=========================================="

# ========== PHASE 1: Clone ==========
if [ ! -d ~/pytorch-build ]; then
    echo "[1/5] Cloning PyTorch v2.5.1 (full recursive)..."
    cd ~
    git clone --recursive --branch v2.5.1 https://github.com/pytorch/pytorch.git pytorch-build
else
    echo "[1/5] Source already exists, skipping clone"
fi

cd ~/pytorch-build

# ========== PHASE 2: Venv + deps ==========
echo "[2/5] Installing build dependencies..."
source ~/ComfyUI/venv/bin/activate
pip install -q cmake ninja pyyaml typing-extensions numpy setuptools wheel cffi

# ========== PHASE 3: Hipify ==========
echo "[3/5] Running hipify (CUDA → HIP conversion)..."
python tools/amd_build/build_amd.py

# ========== PHASE 4: Environment ==========
echo "[4/5] Configuring build environment..."

# ROCm paths
export ROCM_PATH=/opt/rocm
export HIP_PATH=/opt/rocm
export ROCM_HOME=/opt/rocm
export PATH="/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH"
export LD_LIBRARY_PATH="/opt/rocm/lib:${LD_LIBRARY_PATH:-}"

# Target gfx1010 (RDNA1 — closest supported arch for BC-250 gfx1013)
export PYTORCH_ROCM_ARCH="gfx1010"
export AMDGPU_TARGETS="gfx1010"

# Build config
export USE_ROCM=1
export USE_CUDA=0
export USE_CUDNN=0
export USE_NCCL=0
export USE_DISTRIBUTED=0
export USE_MKLDNN=1
export USE_OPENMP=1
export BUILD_TEST=0
export USE_FLASH_ATTENTION=0
export USE_MEM_EFF_ATTENTION=0
export REL_WITH_DEB_INFO=0
export CMAKE_BUILD_TYPE=Release
export CMAKE_POLICY_VERSION_MINIMUM=3.5

# Parallelism — 12 cores, 14GB RAM, be conservative
export MAX_JOBS=4
export NINJA_STATUS="[%f/%t %e] "

# Compiler
export CC=gcc
export CXX=g++

echo "  PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
echo "  MAX_JOBS=$MAX_JOBS"
echo "  ROCM_PATH=$ROCM_PATH"
echo "  Python: $(python --version)"

# ========== PHASE 5: Build ==========
echo "[5/5] Building PyTorch wheel..."
echo "  Start: $(date)"

python setup.py bdist_wheel

echo ""
echo "=========================================="
echo " Build complete: $(date)"
echo "=========================================="
ls -lh dist/*.whl 2>/dev/null || echo "ERROR: No wheel found!"