Files
TheRock/rocm-systems/emulation/mirage/tests/run_tiny_torch_mi350.sh
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

82 lines
3.1 KiB
Bash
Executable File

#!/usr/bin/env bash
#
# Run tests/fixtures/ml/tiny_torch.py on an *emulated* MI350X through
# mirage, using a Python venv populated from the gfx950-dcgpu (cdna4 /
# MI350) ROCm nightly. No physical GPU is required: rocjitsu emulates
# the device.
#
# What this does:
# 1. Create/reuse a venv.
# 2. Install torch + the ROCm SDK (which ships rocjitsu) from the
# gfx950-dcgpu nightly index.
# 4. Run tiny_torch.py via `mirage run --profile mi350x` with
# the venv's python3, and assert it prints `tiny_torch_ok`.
#
# Env knobs:
# VENV venv location (default: <mirage>/.venv-mi350)
# INDEX_URL pip index (default: gfx950-dcgpu nightly)
# SKIP_INSTALL set to 1 to reuse an already-populated venv
# MIRAGE_SKIP_BUILD set to 1 to reuse an already-built mirage prefix
# (otherwise scripts/mirage.sh rebuilds via Docker)
#
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MIRAGE_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
FIXTURE="$MIRAGE_DIR/tests/fixtures/ml/tiny_torch.py"
PROFILE="mi350x"
VENV="${VENV:-$MIRAGE_DIR/.venv-mi350}"
INDEX_URL="${INDEX_URL:-https://rocm.nightlies.amd.com/v2/gfx950-dcgpu/}"
log() { printf '==> %s\n' "$*"; }
fail() { printf 'FAIL: %s\n' "$*" >&2; exit 1; }
[[ -f "$FIXTURE" ]] || fail "fixture not found: $FIXTURE"
# 1. Create the venv (idempotent).
if [[ ! -x "$VENV/bin/python3" ]]; then
log "creating venv at $VENV"
python3 -m venv "$VENV"
fi
# shellcheck disable=SC1091
source "$VENV/bin/activate"
VENV_PY="$VENV/bin/python3"
# 2. Install the gfx950-dcgpu nightly (torch + ROCm SDK, which carries rocjitsu).
if [[ "${SKIP_INSTALL:-0}" != "1" ]]; then
log "installing torch + ROCm SDK from $INDEX_URL"
"$VENV_PY" -m pip install --upgrade pip
"$VENV_PY" -m pip install --index-url "$INDEX_URL" "rocm[libraries,devel]" torch numpy
else
log "SKIP_INSTALL=1: reusing existing venv packages"
fi
# 2b. Expand the devel package. rocm[devel] ships its contents (headers,
# libs incl. librocjitsu.so) compressed; `rocm-sdk init` unpacks
# them into site-packages/_rocm_sdk_devel. Safe to re-run.
if [[ -x "$VENV/bin/rocm-sdk" ]]; then
log "expanding ROCm devel contents (rocm-sdk init)"
"$VENV/bin/rocm-sdk" init
fi
# 3. Run the fixture through mirage on the emulated MI350X. Use the venv's
# python3 explicitly so the workload imports the nightly torch. Run mirage
# via scripts/mirage.sh, which builds the glibc-portable mirage + rocjitsu
# prefix (build/manylinux) and execs the installed binary; that binary's
# sibling librocjitsu*.so provides the KMD discovery and preload wiring.
# Set MIRAGE_SKIP_BUILD=1 to reuse an already-built prefix.
cd "$MIRAGE_DIR"
log "running tiny_torch.py via mirage --profile $PROFILE"
set +e
OUTPUT="$("$MIRAGE_DIR/scripts/mirage.sh" run --profile "$PROFILE" \
-- "$VENV_PY" "$FIXTURE" 2>&1)"
STATUS=$?
set -e
printf '%s\n' "$OUTPUT"
# 5. Assert success.
[[ $STATUS -eq 0 ]] || fail "mirage run exited with status $STATUS"
grep -q 'tiny_torch_ok' <<<"$OUTPUT" || fail "tiny_torch.py did not report success"
log "PASS: tiny_torch ran on emulated $PROFILE using rocjitsu from the venv"