mirror of
https://github.com/AdaCore/cuda.git
synced 2026-02-12 13:05:54 -08:00
65 lines
1.2 KiB
Bash
65 lines
1.2 KiB
Bash
#!/bin/sh
|
|
set -e
|
|
|
|
usage() {
|
|
echo "usage: $(basename "$0") [--expect-single | -s] [--no-expect-single | -S] [--compute-prefix | -c] [--sm-prefix | -C]"
|
|
}
|
|
|
|
expect_single=0
|
|
compute_prefix=1
|
|
while [ -n "$1" ]; do
|
|
case $1 in
|
|
-c|--compute-prefix)
|
|
compute_prefix=1
|
|
;;
|
|
-C|--sm-prefix)
|
|
compute_prefix=0
|
|
;;
|
|
-s|--expect-single)
|
|
expect_single=1
|
|
;;
|
|
-S|--no-expect-single)
|
|
expect_single=0
|
|
;;
|
|
-h|--help)
|
|
usage
|
|
exit 0
|
|
;;
|
|
*)
|
|
echo "unknown arg $1" >&2
|
|
usage
|
|
exit 2
|
|
;;
|
|
esac
|
|
|
|
shift
|
|
done
|
|
|
|
csv=$(nvidia-smi --query-gpu=compute_cap --format=csv)
|
|
cap=$(echo "$csv" | tail -n +2 | tr -d '. \t')
|
|
if [ -z "$cap" ]; then
|
|
number=0
|
|
else
|
|
number=$(echo "$cap" | wc -l)
|
|
fi
|
|
|
|
# number of GPU
|
|
if [ "$expect_single" -eq 0 ]; then
|
|
echo "$number GPU"
|
|
elif [ "$number" -ne 1 ]; then
|
|
echo "expected a single GPU, found $number" >&2
|
|
exit 1
|
|
fi
|
|
|
|
# capacity
|
|
if [ "$number" -ne 0 ]; then
|
|
echo "$cap" | while read -r c; do
|
|
if [ $compute_prefix -eq 0 ]; then
|
|
prefix="sm_"
|
|
else
|
|
prefix="compute_"
|
|
fi
|
|
echo "$prefix$c"
|
|
done
|
|
fi
|