B e] @sddlZddlZddlmZmZddlmZej dddZ ddZ dd d Z dd d Z dddZdddZdddZeredS)N)gemmeinsum) default_timercCsXtj||fdd}|tjdd|j7}tj||fdd}|tjdd|j7}||fS)Nf)dtypegg?)numpyzerosrandomuniformshape)nOnI batch_sizeXWr0/tmp/pip-install-b8evvk6i/blis/blis/benchmark.py create_datas rcCstjjd}|dS)N librariesr)rZ __config__Z blas_opt_info)Z blas_libsrrrget_numpy_blass rc Csp|j\}}|jd}d}tj||fdd}x6t|D]*}tj|||d||7}|dq4Wtd|dS)Nrgr)r)outzTotal:)r rr rangedotsumfillprint) rrnr rrtotalyirrr numpy_gemms   r"c Csn|j\}}|jd}d}tj||fdd}x4t|D](}t|||d||7}|dq4Wtd|dS)Nrgr)r)rzTotal:)r rr rrrrr) rrrr rrrr r!rrr blis_gemm!s   r#c Csr|j\}}|jd}d}tj||fdd}x8t|D],}tjd|||d||7}|dq4Wtd|dS)Nrgr)rz ab,cb->ca)rzTotal:)r rr rrrrr) rrrr rrrr r!rrr numpy_einsum-s   r$c Csp|j\}}|jd}d}tj||fdd}x6t|D]*}td|||d||7}|dq4Wtd|dS)Nrgr)rz ab,cb->ca)rzTotal:)r rr rrrrr) rrrr rrrr r!rrr blis_einsum9s   r%c Cstdjftt}t|||\}}|}|}tdt}t||ddt} | |} td| td|t}t||t} | |} td| tdt}t ||ddt} | |} td| td|t}t ||t} | |} td| dS) NzNSetting up data for gemm. 1000 iters, nO={nO} nI={nI} batch_size={batch_size}z Blis gemm...i)rz %.2f secondszNumpy (%s) gemm...zBlis einsum ab,cb->cazNumpy (%s) einsum ab,cb->ca) rformatlocalsrrcopytimerr#r"r%r$) rr rZ numpy_blasZX1ZW1ZX2ZW2startendZ blis_timeZ numpy_timerrrmainEs<        r.)r)r)r)r)r&r&r')rZ numpy.randompyrrZtimeitrr+r seedrrr"r#r$r%r.__name__rrrrs   "