pytorch/benchmarks/sparse/test_csr.sh
Sameer Deshmukh 5fb1142702 Add CSR (compressed sparse row) layout for sparse tensors (#50937)
Summary:
Implement compressed sparse row format. Derived from the GCS implementation at https://github.com/pytorch/pytorch/pull/44190

Pull Request resolved: https://github.com/pytorch/pytorch/pull/50937

Reviewed By: mrshenli

Differential Revision: D27439865

Pulled By: ezyang

fbshipit-source-id: 3ba3dcb9679505b980ff6a5f513e913bbae2fb1d
2021-04-12 10:09:12 -07:00

42 lines
1.1 KiB
Bash

OUTFILE=spmm-no-mkl-test.txt
PYTORCH_HOME=$1
cd $PYTORCH_HOME
echo "" >> $OUTFILE
echo "----- USE_MKL=1 -----" >> $OUTFILE
rm -rf build
export USE_MKL=1
export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
python setup.py build --cmake-only
ccmake build # or cmake-gui build
python setup.py install
cd benchmarks
echo "!! SPARSE SPMM TIME BENCHMARK!! " >> $OUTFILE
for dim0 in 1000 5000 10000; do
for nnzr in 0.01 0.05 0.1 0.3; do
python -m sparse.spmm --format csr --m $dim0 --n $dim0 --k $dim0 --nnz_ratio $nnzr --outfile $OUTFILE
# python -m sparse.spmm --format coo --m $dim0 --n $dim0 --k $dim0 --nnz_ratio $nnzr --outfile $OUTFILE
done
done
echo "----------------------" >> $OUTFILE
cd $PYTORCH_HOME
echo "----- USE_MKL=0 ------" >> $OUTFILE
rm -rf build
export USE_MKL=0
python setup.py install
cd benchmarks
for dim0 in 1000 5000 10000; do
for nnzr in 0.01 0.05 0.1 0.3; do
python -m sparse.spmv --format csr --m $dim0 --nnz_ratio $nnzr --outfile $OUTFILE
python -m sparse.spmv --format coo --m $dim0 --nnz_ratio $nnzr --outfile $OUTFILE
done
done
echo "----------------------" >> $OUTFILE