|
| 1 | +#!/bin/bash |
| 2 | + |
| 3 | +set -e |
| 4 | + |
| 5 | +ver() { |
| 6 | + printf "%3d%03d%03d%03d" $(echo "$1" | tr '.' ' '); |
| 7 | +} |
| 8 | + |
| 9 | +# Sets GFX_ARCH to default if not set |
| 10 | +if [ -z "$GFX_ARCH" ]; then |
| 11 | + echo "WARNING: missing env var GFX_ARCH, using default (this will take longer)" |
| 12 | + GFX_ARCHS=("gfx900" "gfx906" "gfx908" "gfx90a" "gfx942" "gfx1030") |
| 13 | +else |
| 14 | + # Convert ; seperated string to array |
| 15 | + IFS=';' read -ra GFX_ARCHS <<< "$GFX_ARCH" |
| 16 | +fi |
| 17 | + |
| 18 | +# Sets ROCM_VERSION to "latest" if not set |
| 19 | +if [ -z "$ROCM_VERSION" ]; then |
| 20 | + echo "WARNING: missing env var ROCM_VERSION, using latest kdb repo (NOT RECOMMENDED)" |
| 21 | + ROCM_VERSION="latest" |
| 22 | +fi |
| 23 | + |
| 24 | +# Set PyTorch version and wheel install path |
| 25 | +TORCH_INSTALL_PATH=$(uv pip show torch | grep Location | cut -d" " -f 2) |
| 26 | + |
| 27 | +# Check if Torch installation path exists |
| 28 | +if [ ! -d "$TORCH_INSTALL_PATH" ]; then |
| 29 | + echo "Error: Torch installation path '$TORCH_INSTALL_PATH' does not exist." |
| 30 | + exit 1 |
| 31 | +fi |
| 32 | + |
| 33 | +# Print variable overview |
| 34 | +echo "ROCM version: $ROCM_VERSION" |
| 35 | +echo "GFX architectures: ${GFX_ARCHS[@]}" |
| 36 | +echo "PyTorch installation path: $TORCH_INSTALL_PATH" |
| 37 | + |
| 38 | +# Create directory for extraction |
| 39 | +EXTRACT_DIR=extract_miopen_dbs |
| 40 | +rm -rf $EXTRACT_DIR |
| 41 | +mkdir -p "$EXTRACT_DIR" && cd "$EXTRACT_DIR" |
| 42 | + |
| 43 | +if [[ -f /etc/lsb-release ]]; then |
| 44 | + # Exit if not 20.04, 22.04, or 24.04 |
| 45 | + source /etc/lsb-release |
| 46 | + echo "DISTRIB_RELEASE: $DISTRIB_RELEASE" |
| 47 | + if [[ "$DISTRIB_RELEASE" != "20.04" && "$DISTRIB_RELEASE" != "22.04" ]]; then |
| 48 | + if [[ "$ROCM_VERSION" != "latest" && $(ver $ROCM_VERSION) -lt $(ver 6.2) && "$DISTRIB_RELEASE" == "24.04" ]]; then |
| 49 | + echo "ERROR: Unsupported Ubuntu version." |
| 50 | + exit 1 |
| 51 | + fi |
| 52 | + fi |
| 53 | + |
| 54 | + for arch in "${GFX_ARCHS[@]}"; do |
| 55 | + # Download MIOpen .kdbs for ROCm version and GPU architecture on ubuntu |
| 56 | + echo "Downloading .kdb files for rocm-$ROCM_VERSION ($arch arch) ..." |
| 57 | + wget -q -r -np -nd -A miopen-hip-$arch*kdb_*$DISTRIB_RELEASE*deb \ |
| 58 | + https://repo.radeon.com/rocm/apt/$ROCM_VERSION/pool/main/m/ |
| 59 | + |
| 60 | + # Check if files were downloaded. No KDB files in repo.radeon will result in error. |
| 61 | + if ! ls miopen-hip-$arch*kdb_*$DISTRIB_RELEASE*deb 1> /dev/null 2>&1; then |
| 62 | + echo -e "ERROR: No MIOpen kernel database files found for $arch\nPlease check https://repo.radeon.com/rocm/apt/$ROCM_VERSION/pool/main/m/ for supported architectures" |
| 63 | + exit 1 |
| 64 | + fi |
| 65 | + done |
| 66 | + |
| 67 | + # Extract all .deb files to local directory |
| 68 | + echo "Extracting deb packages for ${GFX_ARCHS[@]} ..." |
| 69 | + for deb_file in `ls *deb`; do |
| 70 | + echo "Extracting $deb_file..." |
| 71 | + dpkg-deb -xv "$deb_file" . > /dev/null 2>&1 |
| 72 | + done |
| 73 | + |
| 74 | +elif [[ -f /etc/centos-release || -f /etc/redhat-release ]]; then |
| 75 | + # Centos kdbs |
| 76 | + source /etc/os-release && RHEL_VERSION="$VERSION_ID" |
| 77 | + RHEL_MAJOR_VERSION=${RHEL_VERSION%%.*} |
| 78 | + echo "RHEL_VERSION: $RHEL_VERSION; RHEL_MAJOR_VERSION: $RHEL_MAJOR_VERSION" |
| 79 | + if [[ ! "$RHEL_VERSION" =~ ^(8|9) ]]; then |
| 80 | + echo "ERROR: Unsupported CentOS/RHEL release" |
| 81 | + fi |
| 82 | + for arch in "${GFX_ARCHS[@]}"; do |
| 83 | + # Download MIOpen .kdbs for ROCm version and GPU architecture on centos |
| 84 | + echo "Downloading .kdb files for rocm-$ROCM_VERSION ($arch arch) ..." |
| 85 | + wget -q -r -np -nd -A miopen-hip-$arch*kdb-[0-9]*rpm \ |
| 86 | + https://repo.radeon.com/rocm/rhel${RHEL_MAJOR_VERSION}/$ROCM_VERSION/main |
| 87 | + |
| 88 | + # Check if files were downloaded. No KDB files in repo.radeon will result in error. |
| 89 | + if ! ls miopen-hip-$arch*kdb-*rpm 1> /dev/null 2>&1; then |
| 90 | + echo -e "ERROR: No MIOpen kernel database files found for $arch\nPlease check https://repo.radeon.com/rocm/rhel${RHEL_MAJOR_VERSION}/$ROCM_VERSION/main for supported architectures" |
| 91 | + exit 1 |
| 92 | + fi |
| 93 | + done |
| 94 | + |
| 95 | + # Extract all RPM files to current directory |
| 96 | + echo "Extracting rpm packages for ${GFX_ARCHS[@]} ..." |
| 97 | + for rpm_file in `ls *rpm`; do |
| 98 | + echo "Extracting $rpm_file..." |
| 99 | + rpm2cpio "$rpm_file" | cpio -idmv 2> /dev/null |
| 100 | + done |
| 101 | +else |
| 102 | + echo "ERROR: Unsupported operating system." |
| 103 | + exit 1 |
| 104 | +fi |
| 105 | + |
| 106 | +# Copy miopen db files to PyTorch installation path |
| 107 | +echo "Copying kdb files to ${TORCH_INSTALL_PATH}/torch/share" |
| 108 | +cp -ra opt/rocm-*/share/miopen $TORCH_INSTALL_PATH/torch/share |
| 109 | + |
| 110 | +# Remove downloaded files and extract directory |
| 111 | +cd .. && rm -rf $EXTRACT_DIR |
| 112 | +echo "Successfully installed MIOpen kernel database files" |
0 commit comments