File: test_aot_compile.sh

package info (click to toggle)
pytorch-cuda 2.6.0%2Bdfsg-7
  • links: PTS, VCS
  • area: contrib
  • in suites: forky, sid, trixie
  • size: 161,620 kB
  • sloc: python: 1,278,832; cpp: 900,322; ansic: 82,710; asm: 7,754; java: 3,363; sh: 2,811; javascript: 2,443; makefile: 597; ruby: 195; xml: 84; objc: 68
file content (27 lines) | stat: -rwxr-xr-x 893 bytes parent folder | download | duplicates (4)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
#!/bin/bash

set -e -o pipefail

TORCH_INSTALL_DIR=$(python -c "import site; print(site.getsitepackages()[0])")/torch
TORCH_BIN_DIR="$TORCH_INSTALL_DIR"/bin
CURRENT_DIR="$(dirname "${BASH_SOURCE[0]}")"

MODEL=aot_test_model.pt
COMPILED_MODEL=aot_test_model.compiled.pt
COMPILED_CODE=aot_test_model.compiled.ll
TMP_DIR=$(mktemp -d -t build_XXX)

test_aot_model_compiler() {
  python "$CURRENT_DIR"/aot_test_model.py
  mv "$MODEL" "$TMP_DIR"/
  pushd "$TMP_DIR"
  "$TORCH_BIN_DIR"/aot_model_compiler_test --model "$MODEL" --model_name=aot_test_model --output_llvm=$COMPILED_CODE --model_version=v1 --input_dims="2,2,2"
  if [ ! -f "$COMPILED_MODEL" ] || [ ! -f "$COMPILED_CODE" ]; then
    echo "AOT model compiler failed to generate $COMPILED_MODEL and $COMPILED_CODE"
    exit 1
  fi
  popd
}

# Temporarily disable the test since NNC backend is no longer available.
# test_aot_model_compiler