xref: /aosp_15_r20/external/pytorch/benchmarks/sparse/test_csr.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1OUTFILE=spmm-no-mkl-test.txt
2PYTORCH_HOME=$1
3
4cd $PYTORCH_HOME
5
6echo "" >> $OUTFILE
7echo "----- USE_MKL=1 -----" >> $OUTFILE
8rm -rf build
9
10export USE_MKL=1
11export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
12python setup.py build --cmake-only
13ccmake build  # or cmake-gui build
14
15python setup.py install
16
17cd benchmarks
18echo "!! SPARSE SPMM TIME BENCHMARK!! " >> $OUTFILE
19for dim0 in 1000 5000 10000; do
20    for nnzr in 0.01 0.05 0.1 0.3; do
21        python -m sparse.spmm --format csr --m $dim0 --n $dim0 --k $dim0 --nnz-ratio $nnzr --outfile $OUTFILE
22        # python -m sparse.spmm --format coo --m $dim0 --n $dim0 --k $dim0 --nnz-ratio $nnzr --outfile $OUTFILE
23    done
24done
25echo "----------------------" >> $OUTFILE
26
27cd $PYTORCH_HOME
28echo "----- USE_MKL=0 ------" >> $OUTFILE
29rm -rf build
30
31export USE_MKL=0
32python setup.py install
33
34cd benchmarks
35for dim0 in 1000 5000 10000; do
36    for nnzr in 0.01 0.05 0.1 0.3; do
37        python -m sparse.spmv --format csr --m $dim0 --nnz-ratio $nnzr --outfile $OUTFILE
38        python -m sparse.spmv --format coo --m $dim0 --nnz-ratio $nnzr --outfile $OUTFILE
39    done
40done
41echo "----------------------" >> $OUTFILE
42