prain: build scripts (ROCm SDK venv, PyTorch and vLLM from source with clang)

scripts/build.sh

@@ -0,0 +1,121 @@
+#!/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

scripts/env.sh

@@ -0,0 +1,26 @@
+# prain environment: the project venv with AMD's ROCm SDK wheels (TheRock) and a clang host
+# toolchain tuned for this CPU. Source it from bash:   . scripts/env.sh
+# Overrides: PRAIN_GPU_ARCH (default gfx1032), PRAIN_JOBS (parallel build jobs, default 16).
+PRAIN_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
+export PRAIN_ROOT
+export PRAIN_GPU_ARCH=${PRAIN_GPU_ARCH:-gfx1032}
+export PRAIN_JOBS=${PRAIN_JOBS:-16}
+export VIRTUAL_ENV=$PRAIN_ROOT/.venv
+case ":$PATH:" in *":$VIRTUAL_ENV/bin:"*) ;; *) PATH=$VIRTUAL_ENV/bin:$PATH ;; esac
+_site=$(echo "$VIRTUAL_ENV"/lib/python3*/site-packages)
+if [ -d "$_site/_rocm_sdk_devel" ]; then
+    _dev=$_site/_rocm_sdk_devel
+    _core=$_site/_rocm_sdk_core
+    export ROCM_PATH=$_dev ROCM_HOME=$_dev HIP_PATH=$_dev CMAKE_PREFIX_PATH=$_dev
+    export HIP_CLANG_PATH=$_dev/lib/llvm/bin HIP_DEVICE_LIB_PATH=$_core/lib/llvm/amdgcn/bitcode
+    case ":${LD_LIBRARY_PATH:-}:" in *":$_dev/lib:"*) ;; *)
+        export LD_LIBRARY_PATH=$_dev/lib:$_dev/lib/rocm_sysdeps/lib:$_core/lib:$_core/lib/rocm_sysdeps/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} ;;
+    esac
+    case ":$PATH:" in *":$_dev/bin:"*) ;; *) PATH=$_dev/bin:$PATH ;; esac
+fi
+export PATH
+# Host C/C++: the system clang through wrappers (scripts/build.sh wrappers) that drop Gentoo's
+# libc++/compiler-rt defaults: everything links libstdc++/libgcc like the ROCm libraries.
+export CC=$PRAIN_ROOT/build/bin/clang CXX=$PRAIN_ROOT/build/bin/clang++
+export CFLAGS="-march=native" CXXFLAGS="-march=native"
+unset _site _dev _core