#!/bin/bash
# Build prain's stack on the GPU host, with clang, tuned for this machine:
#   wrappers  clang/clang++ wrappers in build/bin (libstdc++ + libgcc, lld)
#   sdk       venv + AMD's ROCm SDK wheels for $PRAIN_GPU_ARCH
#   fetch     clone PyTorch and vLLM into build/
#   torch     PyTorch from source -> build/wheels, installed into the venv
#   vllm      vLLM from source (adds $PRAIN_GPU_ARCH to its HIP arch list), installed editable
#   prain     this package, installed editable, and its test suite
# usage: scripts/build.sh STEP...      Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VLLM_REF
set -euo pipefail
cd "$(dirname "$0")/.."
ROCM_INDEX=${PRAIN_ROCM_INDEX:-https://stable.repo.amd.com/rocm/whl-next/}
ROCM_VERSION=${PRAIN_ROCM_VERSION:-10.1.0}
TORCH_REF=${PRAIN_TORCH_REF:-v2.13.0}
VLLM_REF=${PRAIN_VLLM_REF:-v0.32.0}
PYTHON=${PRAIN_PYTHON:-python3.13}

wrappers() {
    mkdir -p build/bin
    local gccdir
    gccdir=$(gcc -print-search-dirs | sed -n 's/^install: //p')
    for c in clang clang++; do
        cat > build/bin/$c <<EOF
#!/bin/sh
case "\$*" in
    # clang 23 spins for hours on PyTorch's baseline (non-AVX) linear algebra kernel with
    # -march=native. This CPU runs the AVX2 variant, so the baseline one is built with -O1.
    *LinearAlgebraKernel.cpp.DEFAULT.cpp*) set -- "\$@" -O1 ;;
esac
exec $(command -v $c) --no-default-config --gcc-install-dir=${gccdir%/} -stdlib=libstdc++ -rtlib=libgcc -unwindlib=libgcc -fuse-ld=lld -Wno-unused-command-line-argument "\$@"
EOF
        chmod +x build/bin/$c
    done
    printf '#include <iostream>\nint main() { std::cout << "ok" << std::endl; }\n' > build/bin/probe.cpp
    build/bin/clang++ -march=native -O2 build/bin/probe.cpp -o build/bin/probe
    build/bin/probe
    ldd build/bin/probe | grep -E "c\+\+|gcc|unwind"
}

sdk() {
    [ -x .venv/bin/python ] || $PYTHON -m venv .venv
    .venv/bin/python -m pip install -U pip wheel 'setuptools<80' 'cmake<4' ninja
    .venv/bin/python -m pip install --index-url "$ROCM_INDEX" \
        "rocm[libraries,devel]==$ROCM_VERSION" "rocm-sdk-device-$PRAIN_GPU_ARCH==$ROCM_VERSION"
    .venv/bin/rocm-sdk init
    .venv/bin/rocm-sdk version
    # libstdc++ 16's <format> has [[__gnu__::__noinline__]], and HIP's host_defines.h empties
    # __noinline__ for host compilers: parse <format> first, like Gentoo's hip ebuild does
    # (https://github.com/ROCm/rocm-systems/issues/1874).
    python - "$(echo .venv/lib/python3*/site-packages)/_rocm_sdk_devel/include/hip/amd_detail/host_defines.h" <<'EOF'
import sys
path = sys.argv[1]
text = open(path).read()
fix = ("#if defined(__cplusplus) && defined(__has_include)\n#  if __has_include(<__config>)\n"
       "#    include <__config>\n#  else\n#    include <format>\n#  endif\n#endif\n")
if fix not in text:
    assert text.count("\n#define __noinline__\n") == 1, path
    open(path, "w").write(text.replace("\n#define __noinline__\n", "\n" + fix + "#define __noinline__\n"))
print("host_defines.h: <format> before __noinline__")
EOF
}

fetch() {
    [ -d build/pytorch/.git ] || git clone --depth 1 --branch "$TORCH_REF" --recurse-submodules \
        --shallow-submodules -j 8 https://github.com/pytorch/pytorch build/pytorch
    [ -d build/vllm/.git ] || git clone --depth 1 --branch "$VLLM_REF" \
        https://github.com/vllm-project/vllm build/vllm
}

# PyTorch preloads these SDK libraries on import (torch/_rocm_init.py), so no LD_LIBRARY_PATH.
TORCH_PRELOADS='"amd_comgr", "amdhip64", "hiprtc", "roctx64", "hipblas", "hipblaslt", "hipfft",
    "hiprand", "hipsparse", "hipsparselt", "hipsolver", "miopen"'

torch() {
    local src=build/pytorch
    python -m pip install -r $src/requirements-build.txt filelock 'fsspec>=0.8.5' jinja2 \
        'networkx>=2.5.1' 'optree>=0.13.0' psutil 'sympy>=1.13.3'
    (cd $src && python tools/amd_build/build_amd.py > ../logs/hipify.log)
    printf 'def initialize():\n    import rocm_sdk\n    rocm_sdk.initialize_process(\n        preload_shortnames=[%s],\n        check_version="%s")\n' \
        "$TORCH_PRELOADS" "$ROCM_VERSION" > $src/torch/_rocm_init.py
    # Only what prain needs on one RDNA2 GPU: no CK/aotriton attention (not built for gfx103x),
    # no NCCL (vLLM falls back to gloo on one GPU), no profiler, no mobile/quantized CPU backends.
    (cd $src && env USE_ROCM=1 USE_CUDA=0 USE_XPU=0 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" \
        USE_ROCM_CK_GEMM=0 USE_ROCM_CK_SDPA=0 USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 \
        USE_MSLK=0 USE_NCCL=0 USE_MPI=0 USE_KINETO=0 USE_FBGEMM=0 USE_NNPACK=0 \
        USE_PYTORCH_QNNPACK=0 USE_XNNPACK=0 USE_ITT=0 USE_MAGMA=0 BUILD_TEST=0 BLAS=OpenBLAS \
        MAX_JOBS="$PRAIN_JOBS" nice -n 19 python -m pip wheel --no-build-isolation --no-deps -v \
        -w "$PRAIN_ROOT/build/wheels" .)
    python -m pip install --no-deps --force-reinstall build/wheels/torch-*.whl
    env -u LD_LIBRARY_PATH python -c 'import torch; x = torch.ones(4, device="cuda"); print(torch.__version__, torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName, (x * 2).sum().item(), torch.backends.openmp.is_available(), torch.get_num_threads())'
}

vllm() {
    local src=build/vllm site
    site=$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
    # vLLM lists gfx1030 but not its siblings gfx1031/gfx1032: add ours.
    grep -q "$PRAIN_GPU_ARCH" $src/CMakeLists.txt ||
        sed -i "s/;gfx1030;/;gfx1030;$PRAIN_GPU_ARCH;/" $src/CMakeLists.txt
    # Keep our torch and AMD's Triton for this ROCm release: pip fails rather than swap them.
    printf 'torch==%s\ntriton==3.8.0+git669b31ac.rocm%s\n' \
        "$(python -c 'import torch; print(torch.__version__)')" "$ROCM_VERSION" > build/constraints.txt
    python -m pip install -c build/constraints.txt --extra-index-url "$ROCM_INDEX" triton \
        -r $src/requirements/common.txt 'setuptools-scm>=9.2.0' 'setuptools-rust>=1.9.0' 'jinja2>=3.1.6'
    # The SDK's own amdsmi package (vLLM's ROCm platform imports it), not PyPI's.
    echo "$site/_rocm_sdk_core/share/amd_smi" > "$site/rocm-sdk-amdsmi.pth"
    (cd $src && env VLLM_TARGET_DEVICE=rocm PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" \
        MAX_JOBS="$PRAIN_JOBS" nice -n 19 python -m pip install --no-build-isolation --no-deps -v -e .)
    env -u LD_LIBRARY_PATH python -c 'import vllm, triton, amdsmi; print(vllm.__version__, triton.__version__)'
}

prain() {
    python -m pip install --no-deps -e .
    python -m unittest discover -s tests -t .
}

for step in "$@"; do
    . scripts/env.sh
    mkdir -p build/logs
    echo "== $step $(date -Is)"
    "$step" 2>&1 | tee "build/logs/$step.log"
done