BLAS for Tensors: What, Why, and How Devin Ma;hews - PowerPoint PPT Presentation

BLAS ¡for ¡Tensors: ¡What, ¡Why, ¡ and ¡How ¡ Devin ¡Ma;hews ¡

Outline ¡ I. The ¡role ¡of ¡interfaces ¡in ¡domain ¡applicaDons ¡ II. BLAS ¡as ¡an ¡example ¡interface ¡ III. BLIS ¡as ¡a ¡be;er ¡BLAS ¡ IV. Tensors ¡ V. Tensor ¡interfaces ¡now ¡ VI. A ¡“Tensor ¡BLAS” ¡

Math ¡to ¡Code: ¡Fundamental ¡ ConsideraDons ¡ What ¡domain ¡scienDsts ¡want ¡(math): ¡ ¡ H, δ T (1) ] , δ T (2) ] | 0 i 1 4 λ ijk (1) H ma (3) t ebc (1) h 0 | δ Λ (1) [[ ¯ ¯ ie abc mjk ¡ ¡ ¡ ¡ ¡ ¡ ¡ ¡ ¡ ¡

Math ¡to ¡Code: ¡Fundamental ¡ ConsideraDons ¡ What ¡domain ¡scienDsts ¡want ¡(math): ¡ ¡ H, δ T (1) ] , δ T (2) ] | 0 i 1 4 λ ijk (1) H ma (3) t ebc (1) h 0 | δ Λ (1) [[ ¯ ¯ ie abc mjk ¡ ¡ What ¡computer ¡(computaDonal) ¡science ¡provides ¡(code): ¡ ¡ DGEMM ¡(BLAS) ¡ SCSRMM ¡(SparseBLAS) ¡ MPI_Reduce_sca;er_block ¡(MPI) ¡ ¡ “\” ¡(Matlab) ¡ ZGGEV ¡(LAPACK) ¡ contract(…)/“*” ¡(i.e. ¡libtensor ¡et ¡al.) ¡ ¡ ¡ ¡ ¡ ¡

Math ¡to ¡Code: ¡Fundamental ¡ ConsideraDons ¡ What ¡domain ¡scienDsts ¡want ¡(math): ¡ ¡ H, δ T (1) ] , δ T (2) ] | 0 i 1 4 λ ijk (1) H ma (3) t ebc (1) h 0 | δ Λ (1) [[ ¯ ¯ ie abc mjk ¡ ¡ What ¡computer ¡(computaDonal) ¡science ¡provides ¡(code): ¡ ¡ DGEMM ¡(BLAS) ¡ SCSRMM ¡(SparseBLAS) ¡ MPI_Reduce_sca;er_block ¡(MPI) ¡ ¡ “\” ¡(Matlab) ¡ ZGGEV ¡(LAPACK) ¡ contract(…)/“*” ¡(i.e. ¡libtensor ¡et ¡al.) ¡ ¡ What ¡no ¡one ¡likes ¡to ¡have ¡to ¡deal ¡with ¡(but ¡is ¡important): ¡ ¡ ¡ Data ¡layout ¡ Performance ¡ Alignment/caching ¡ Numerical ¡representaDon ¡ Architectural ¡opDmizaDons ¡

Math ¡to ¡Code: ¡Fundamental ¡ ConsideraDons ¡ What ¡domain ¡scienDsts ¡want ¡(math): ¡ ¡ H, δ T (1) ] , δ T (2) ] | 0 i 1 4 λ ijk (1) H ma (3) t ebc (1) h 0 | δ Λ (1) [[ ¯ ¯ ie abc mjk ¡ ¡ What ¡computer ¡(computaDonal) ¡science ¡provides ¡(code): ¡ ¡ DGEMM ¡(BLAS) ¡ SCSRMM ¡(SparseBLAS) ¡ MPI_Reduce_sca;er_block ¡(MPI) ¡ ¡ “\” ¡(Matlab) ¡ ZGGEV ¡(LAPACK) ¡ contract(…)/“*” ¡(i.e. ¡libtensor ¡et ¡al.) ¡ ¡ What ¡no ¡one ¡likes ¡to ¡have ¡to ¡deal ¡with ¡(but ¡is ¡important): ¡ ¡ ¡ Data ¡layout ¡ Performance ¡ Alignment/caching ¡ Numerical ¡representaDon ¡ Architectural ¡opDmizaDons ¡

Features ¡of ¡BLAS ¡ User/external ¡code ¡owns ¡the ¡ DGEMM(CHARACTER*1 TRANSA, data. ¡ CHARACTER*1 TRANSB, INTEGER M, INTEGER N, Data ¡and ¡specificaDon ¡are ¡ INTEGER K, separate. ¡Allows ¡easy ¡hacking ¡ REAL*8 ALPHA, (e.g. ¡submatrices, ¡resizing, ¡ REAL*8(*) A, etc.), ¡but ¡can ¡be ¡error-­‑prone. ¡ INTEGER LDA REAL*8(*) B, INTEGER LDB All ¡informaDon ¡about ¡the ¡ REAL*8 BETA operaDon ¡is ¡explicit. ¡Can ¡be ¡ REAL*8(*) C, unwieldy ¡for ¡users ¡but ¡good ¡ INTEGER LDC) for ¡flexibility. ¡

BLAS: ¡the ¡Good ¡and ¡the ¡Bad ¡ (a.k.a. ¡How ¡to ¡Start ¡a ¡Flame ¡War) ¡ Good: ¡ Bad: ¡ ¡ ¡ Easier ¡than ¡wriDng ¡three ¡loops. ¡ FORTRAN ¡dependency. ¡ • • Flexible ¡enough ¡for ¡end ¡users ¡and ¡ Requires ¡unit ¡stride. ¡ • • library ¡writers. ¡ No ¡conjugaDon ¡without ¡ • High-­‑performance ¡ transpose. ¡ • implementaDons ¡exist. ¡ Some ¡obvious ¡missing ¡features. ¡ • Bindings ¡in ¡many ¡languages. ¡ • User ¡must ¡allocate ¡data ¡and ¡deal ¡ • with ¡alignment ¡etc. ¡ Can ¡be ¡wrapped ¡in ¡a ¡more ¡user-­‑ • friendly ¡interface. ¡ Opaque. ¡ • Poor ¡threading ¡control. ¡ • No ¡mixed-­‑precision ¡support. ¡ •

How ¡People ¡Really ¡(Ab)use ¡BLAS ¡ dcopy(n, ¡&constant, ¡0, ¡array, ¡1) ¡ daxpy(n, ¡alpha1, ¡x_1, ¡1, ¡y, ¡1) ¡ (i.e. ¡fill ¡array) ¡ daxpy(n, ¡alpha1, ¡x_2, ¡1, ¡y, ¡1) ¡ ¡ daxpy(n, ¡alpha2, ¡x_3, ¡1, ¡y, ¡1) ¡ … ¡ Actually ¡transpose ¡a ¡matrix ¡ (and ¡other ¡loop ¡code) ¡ dscal(n, ¡beta, ¡y, ¡1) ¡ ddot(n, ¡&one, ¡0, ¡array, ¡1) ¡ daxpy(n, ¡alpha, ¡x, ¡1, ¡y, ¡1) ¡ (i.e. ¡ non-­‑ absolute ¡sum) ¡ (i.e. ¡daxpby) ¡ Lots ¡of ¡copying ¡and ¡data ¡movement ¡ Separate ¡real ¡and ¡imaginary ¡ BLAS ¡call ¡ arrays/matrices ¡ Lots ¡more ¡data ¡movement ¡

BLIS ¡ void bli_dgemm( trans_t transa, trans_t transb, dim_t m, dim_t n, dim_t k, DGEMM(CHARACTER*1 TRANSA, double* alpha, CHARACTER*1 TRANSB, double* a, INTEGER M, inc_t rsa, inc_t csa, INTEGER N, double* b, INTEGER K, inc_t rsb, inc_t csb, REAL*8 ALPHA, double* beta, REAL*8(*) A, double* c, INTEGER LDA inc_t rsc, inc_t csc ); REAL*8(*) B, INTEGER LDB REAL*8 BETA REAL*8(*) C, void bli_gemm( obj_t* alpha, INTEGER LDC) obj_t* a, obj_t* b, obj_t* beta, obj_t* c )

Tensors ¡ • Tensors ¡are ¡essenDally ¡mulD-­‑dimensional ¡arrays: ¡ T ab ij ∈ R n a × n b × n i × n j double T[na][nb][ni][nj]; • Matrices ¡are ¡2-­‑D ¡tensors. ¡ ¡ • Tensor ¡indices ¡are ¡usually ¡explicit. ¡Einstein ¡notaDon ¡is ¡ very ¡helpful: ¡ h ck + P bj t fc t fc z abc ( 1 + P ai t abe v mn ef � 2 ( 1 + P ai t aeb v mn = ck ) 4ˇ ijm ˇ bj ) ˇ ijm ˇ ˇ nk ˇ nk ˇ ijk ef t cf t fc t abe v mn t abe v mn � 2ˇ ijm ˇ ef � 2ˇ ijm ˇ nk ˇ nk ˇ fe ¡ ¡ t cf t fc +( 1 + P ai t aeb v mn ef +( 1 + P ai t aeb v mn bj ) ˇ ijm ˇ bj ) ˇ ijm ˇ nk ˇ nk ˇ fe i t cf t bf t abe v mn fe +( 1 + P ai t aec v mn + ˇ ijm ˇ bj ) ˇ ijm ˇ nk ˇ nk ˇ fe .

Tensors ¡in ¡Chemistry ¡ CollecDon ¡ CollecDon ¡ CollecDon ¡ of ¡various ¡ of ¡various ¡ of ¡various ¡ tensors ¡ tensors ¡ tensors ¡ T ] ˆ ˆ [ ˆ R | 0 i = E ˆ R | 0 i H, e • Sequence ¡of ¡tensor ¡contracDons, ¡summaDons, ¡etc. ¡ • Need ¡to ¡account ¡for ¡physics: ¡spin, ¡spaDal ¡symmetry, ¡and ¡so ¡on. ¡ • Some ¡tensors ¡are ¡very ¡big, ¡some ¡are ¡very ¡small, ¡many ¡different ¡ dimensionaliDes. ¡ • Oien ¡need ¡(distributed), ¡out-­‑of-­‑core ¡storage. ¡

High-­‑level ¡Interfaces: ¡ Chemistry ¡ + 1 ef t ef F m = f m 2 v mn i i in e − 1 F a e = f a 2 v mn ef t af mn + 1 ef t ef W mn = v mn 2 v mn ij ij ij ei + 1 ef t af W ma = v ma 2 v mn ei in z ab ij = v ab ij + P ( ab ) F a e t eb ij − P ( ij ) F m i t ab mj +1 2 W mn ij t ab mn +1 ef t ef 2 v ab ij + P ( ab ) P ( ij ) W ma ei t eb mj

Low-­‑level ¡Interfaces: ¡ Blitz++, ¡Eigen, ¡Boost, ¡BTAS, ¡etc. ¡ • Handle ¡allocaDon, ¡alignment, ¡indexing, ¡etc. ¡“NaDve” ¡efficiency. ¡ Array<double,4> pqrs(np, nq, nr, ns); pqrs[0][1][6][3] = 0.122; • C++ ¡allows ¡for ¡efficient ¡and ¡expressive ¡interfaces ¡(fixed ¡vs. ¡variable ¡ dimensionality, ¡staDc ¡typing, ¡operator ¡overloading, ¡operaDon ¡trees, ¡ etc.). ¡ Array<int,3> A(4, 10, 2); Array<int,1> G = A(2, 7, Range::all()); • Limited ¡computaDonal ¡support ¡(i.e. ¡Level ¡1-­‑type ¡operaDons, ¡exptl. ¡ contracDon ¡support ¡in ¡Eigen). ¡ ¡ ¡ ¡ C(i,j) = sum(A(i,k), B(k,j), k) + Cprime(i,j);

Under ¡the ¡Hood: ¡BLAS! ¡ libtensor, ¡ Eigen, ¡ TCE ¡ BTAS ¡ etc. ¡ Python/NumPy, ¡ Everything ¡else ¡ Matlab ¡ permute permute gemm permute