# 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"
# Fused-MoE kernel configs tuned for this GPU by scripts/tune_moe.py (vLLM's defaults when
# the files are missing).
export VLLM_TUNED_CONFIG_FOLDER=$PRAIN_ROOT/build/moe-configs
unset _site _dev _core