178 lines
9.6 KiB
C++
178 lines
9.6 KiB
C++
/*
|
|
LA: linear algebra C++ interface library
|
|
Copyright (C) 2008 Jiri Pittner <jiri.pittner@jh-inst.cas.cz> or <jiri@pittnerovi.com>
|
|
|
|
This program is free software: you can redistribute it and/or modify
|
|
it under the terms of the GNU General Public License as published by
|
|
the Free Software Foundation, either version 3 of the License, or
|
|
(at your option) any later version.
|
|
|
|
This program is distributed in the hope that it will be useful,
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
GNU General Public License for more details.
|
|
|
|
You should have received a copy of the GNU General Public License
|
|
along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
|
|
#include "mat.h"
|
|
|
|
namespace LA {
|
|
/*Strassen algorithm in the Winograd variant (7 products, 15 additions), recursion on strided quadrants of the original matrices,
|
|
leaf products by BLAS gemm
|
|
V. Strassen, Gaussian elimination is not optimal, Numer. Math. 13 (1969) 354-356
|
|
S. Winograd, On multiplication of 2x2 matrices, Linear Algebra Appl. 4 (1971) 381-388
|
|
schedule with 2 temporaries per level: B. Boyer, J.-G. Dumas, C. Pernet, W. Zhou, Memory efficient scheduling of
|
|
Strassen-Winograd's matrix multiplication algorithm, Proc. ISSAC 2009, 55-62
|
|
odd dimensions by dynamic peeling: S. Huss-Lederman, E. M. Jacobson, J. R. Johnson, A. Tsao, T. Turnbull,
|
|
Implementation of Strassen's algorithm for matrix multiplication, Proc. Supercomputing '96
|
|
*/
|
|
|
|
//device-dependent primitives: leaf gemm, quadrant addition, workspace (GPU version: cublas?gemm, cublas?geam, gpualloc/gpufree)
|
|
static inline CBLAS_TRANSPOSE s_trans(const char t) {return t=='n'?CblasNoTrans:(t=='c'?CblasConjTrans:CblasTrans);}
|
|
|
|
static inline void s_gemm(const char ta, const char tb, const int m, const int n, const int k, const double alpha, const double *a, const int lda, const double *b, const int ldb, const double beta, double *c, const int ldc)
|
|
{cblas_dgemm(CblasRowMajor,s_trans(ta),s_trans(tb),m,n,k,alpha,a,lda,b,ldb,beta,c,ldc);}
|
|
static inline void s_gemm(const char ta, const char tb, const int m, const int n, const int k, const float alpha, const float *a, const int lda, const float *b, const int ldb, const float beta, float *c, const int ldc)
|
|
{cblas_sgemm(CblasRowMajor,s_trans(ta),s_trans(tb),m,n,k,alpha,a,lda,b,ldb,beta,c,ldc);}
|
|
static inline void s_gemm(const char ta, const char tb, const int m, const int n, const int k, const std::complex<double> alpha, const std::complex<double> *a, const int lda, const std::complex<double> *b, const int ldb, const std::complex<double> beta, std::complex<double> *c, const int ldc)
|
|
{cblas_zgemm(CblasRowMajor,s_trans(ta),s_trans(tb),m,n,k,&alpha,a,lda,b,ldb,&beta,c,ldc);}
|
|
static inline void s_gemm(const char ta, const char tb, const int m, const int n, const int k, const std::complex<float> alpha, const std::complex<float> *a, const int lda, const std::complex<float> *b, const int ldb, const std::complex<float> beta, std::complex<float> *c, const int ldc)
|
|
{cblas_cgemm(CblasRowMajor,s_trans(ta),s_trans(tb),m,n,k,&alpha,a,lda,b,ldb,&beta,c,ldc);}
|
|
|
|
//z = x+y or x-y, r x c blocks with leading dimensions; z may coincide with x or y (one pass, memory bound)
|
|
template<typename T>
|
|
static void s_add(const int r, const int c, T *z, const int ldz, const T *x, const int ldx, const T *y, const int ldy, const bool minus)
|
|
{
|
|
for(int i=0; i<r; ++i)
|
|
{
|
|
T *zz=z+(size_t)i*ldz; const T *xx=x+(size_t)i*ldx; const T *yy=y+(size_t)i*ldy;
|
|
if(minus) for(int j=0; j<c; ++j) zz[j]=xx[j]-yy[j];
|
|
else for(int j=0; j<c; ++j) zz[j]=xx[j]+yy[j];
|
|
}
|
|
}
|
|
|
|
template<typename T> static T *s_alloc(const size_t n) {return n? new T[n] : NULL;}
|
|
template<typename T> static void s_free(T *p) {if(p) delete[] p;}
|
|
|
|
//operand op(M) of the product, op = identity ('n'), transposition ('t') or conjugate transposition ('c'),
|
|
//M stored row-major at p with leading dimension ld; the recursion is formulated in op() space and never transposes data:
|
|
//- at(i,j) is the address of element (i,j) of op(M), i.e. M(i,j) for 'n' and M(j,i) for 't','c';
|
|
// a block of op(M) starting at (i,j) is then the block of M starting at at(i,j) with the same ld and flag
|
|
// (quadrant (op M)_{IJ} = op(M_{JI}), the stored quadrant being swapped for 't','c' and conjugated by gemm for 'c')
|
|
//- sums of quadrants are formed elementwise in the stored layout (S stored = M_{JI} + M_{LK} for 't','c'), which is correct
|
|
// because op(X+Y)=op(X)+op(Y), so the temporaries keep the flag of their operand and gemm applies op() at the leaves only
|
|
template<typename T>
|
|
struct s_op
|
|
{
|
|
const T *p; int ld; char t;
|
|
const T *at(const int i, const int j) const {return t=='n' ? p+(size_t)i*ld+j : p+(size_t)j*ld+i;}
|
|
};
|
|
|
|
static inline bool s_split(const int m, const int n, const int k, const int cutoff) {return m>cutoff && n>cutoff && k>cutoff;}
|
|
|
|
//workspace needed for all levels below (m,n,k)
|
|
static size_t s_work(const int m, const int n, const int k, const int cutoff)
|
|
{
|
|
if(!s_split(m,n,k,cutoff)) return 0;
|
|
size_t mh=m/2, nh=n/2, kh=k/2;
|
|
return (mh*kh>mh*nh ? mh*kh : mh*nh) + kh*nh + s_work(mh,nh,kh,cutoff);
|
|
}
|
|
|
|
//C(m x n, ldc) = alpha op(A) op(B)
|
|
template<typename T>
|
|
static void s_rec(const int m, const int n, const int k, const T alpha, const s_op<T> &A, const s_op<T> &B, T *C, const int ldc, T *ws, const int cutoff)
|
|
{
|
|
if(!s_split(m,n,k,cutoff)) {s_gemm(A.t,B.t,m,n,k,alpha,A.p,A.ld,B.p,B.ld,(T)0,C,ldc); return;}
|
|
const int mh=m/2, nh=n/2, kh=k/2;
|
|
//stored dimensions of the quadrants of A and B
|
|
const int ar= A.t=='n'?mh:kh, ac= A.t=='n'?kh:mh;
|
|
const int br= B.t=='n'?kh:nh, bc= B.t=='n'?nh:kh;
|
|
const T *A11=A.at(0,0), *A12=A.at(0,kh), *A21=A.at(mh,0), *A22=A.at(mh,kh);
|
|
const T *B11=B.at(0,0), *B12=B.at(0,nh), *B21=B.at(kh,0), *B22=B.at(kh,nh);
|
|
T *C11=C, *C12=C+nh, *C21=C+(size_t)mh*ldc, *C22=C21+nh;
|
|
const size_t xs= (size_t)mh*kh>(size_t)mh*nh ? (size_t)mh*kh : (size_t)mh*nh;
|
|
T *X=ws, *Y=ws+xs, *wn=Y+(size_t)kh*nh;
|
|
const s_op<T> sx={X,ac,A.t}, sy={Y,bc,B.t}; //X, Y hold sums of quadrants of A, B in their stored layout, contiguous (ld = stored columns), same flag
|
|
const s_op<T> a11={A11,A.ld,A.t}, a12={A12,A.ld,A.t}, a22={A22,A.ld,A.t};
|
|
const s_op<T> b11={B11,B.ld,B.t}, b21={B21,B.ld,B.t}, b22={B22,B.ld,B.t};
|
|
s_add(ar,ac,X,ac,A11,A.ld,A21,A.ld,true); //S3=A11-A21
|
|
s_add(br,bc,Y,bc,B22,B.ld,B12,B.ld,true); //T3=B22-B12
|
|
s_rec(mh,nh,kh,alpha,sx,sy,C21,ldc,wn,cutoff); //P7
|
|
s_add(ar,ac,X,ac,A21,A.ld,A22,A.ld,false); //S1=A21+A22
|
|
s_add(br,bc,Y,bc,B12,B.ld,B11,B.ld,true); //T1=B12-B11
|
|
s_rec(mh,nh,kh,alpha,sx,sy,C22,ldc,wn,cutoff); //P5
|
|
s_add(ar,ac,X,ac,X,ac,A11,A.ld,true); //S2=S1-A11
|
|
s_add(br,bc,Y,bc,B22,B.ld,Y,bc,true); //T2=B22-T1
|
|
s_rec(mh,nh,kh,alpha,sx,sy,C12,ldc,wn,cutoff); //P6
|
|
s_add(ar,ac,X,ac,A12,A.ld,X,ac,true); //S4=A12-S2
|
|
s_rec(mh,nh,kh,alpha,sx,b22,C11,ldc,wn,cutoff); //P3
|
|
s_rec(mh,nh,kh,alpha,a11,b11,X,nh,wn,cutoff); //P1 in X
|
|
s_add(mh,nh,C12,ldc,X,nh,C12,ldc,false); //U2=P1+P6
|
|
s_add(mh,nh,C21,ldc,C12,ldc,C21,ldc,false); //U3=U2+P7
|
|
s_add(mh,nh,C12,ldc,C12,ldc,C22,ldc,false); //U4=U2+P5
|
|
s_add(mh,nh,C22,ldc,C21,ldc,C22,ldc,false); //C22=U3+P5
|
|
s_add(mh,nh,C12,ldc,C12,ldc,C11,ldc,false); //C12=U4+P3
|
|
s_add(br,bc,Y,bc,Y,bc,B21,B.ld,true); //T4=T2-B21
|
|
s_rec(mh,nh,kh,alpha,a22,sy,C11,ldc,wn,cutoff); //P4
|
|
s_add(mh,nh,C21,ldc,C21,ldc,C11,ldc,true); //C21=U3-P4
|
|
s_rec(mh,nh,kh,alpha,a12,b21,C11,ldc,wn,cutoff); //P2
|
|
s_add(mh,nh,C11,ldc,X,nh,C11,ldc,false); //C11=P1+P2
|
|
//odd dimensions
|
|
const int me=2*mh, ne=2*nh, ke=2*kh;
|
|
if(k>ke) s_gemm(A.t,B.t,me,ne,1,alpha,A.at(0,ke),A.ld,B.at(ke,0),B.ld,(T)1,C,ldc);
|
|
if(n>ne) s_gemm(A.t,B.t,me,1,k,alpha,A.p,A.ld,B.at(0,ne),B.ld,(T)0,C+ne,ldc);
|
|
if(m>me) s_gemm(A.t,B.t,1,n,k,alpha,A.at(me,0),A.ld,B.p,B.ld,(T)0,C+(size_t)me*ldc,ldc);
|
|
}
|
|
|
|
//C = beta C + alpha op(A) op(B), row-major m x n
|
|
template<typename T>
|
|
static void s_top(const int m, const int n, const int k, const T alpha, const T *a, const int lda, const char ta, const T *b, const int ldb, const char tb, const T beta, T *c, const int cutoff)
|
|
{
|
|
const s_op<T> A={a,lda,ta}, B={b,ldb,tb};
|
|
if(!s_split(m,n,k,cutoff)) {s_gemm(ta,tb,m,n,k,alpha,a,lda,b,ldb,beta,c,n); return;}
|
|
const size_t w=s_work(m,n,k,cutoff);
|
|
const bool acc= beta!=(T)0;
|
|
T *ws=s_alloc<T>(w+(acc?(size_t)m*n:0));
|
|
T *cc= acc? ws+w : c;
|
|
s_rec(m,n,k,alpha,A,B,cc,n,ws,cutoff);
|
|
if(acc) for(size_t i=0; i<(size_t)m*n; ++i) c[i]=beta*c[i]+cc[i];
|
|
s_free(ws);
|
|
}
|
|
|
|
static inline char s_flag(const char t, const bool cmplx) {char tt=tolower(t); return tt=='n'?'n':(tt=='c'&&cmplx?'c':'t');}
|
|
|
|
template<typename T>
|
|
static void s_strassen(NRMat<T> &c, const T &beta, const NRMat<T> &a, const char transa, const NRMat<T> &b, const char transb, const T &alpha, const int cutoff)
|
|
{
|
|
const char ta=s_flag(transa,LA_traits<T>::is_complex()), tb=s_flag(transb,LA_traits<T>::is_complex());
|
|
int k(ta=='n'?a.ncols():a.nrows());
|
|
#ifdef DEBUG
|
|
int l(ta=='n'?a.nrows():a.ncols());
|
|
int kk(tb=='n'?b.nrows():b.ncols());
|
|
int ll(tb=='n'?b.ncols():b.nrows());
|
|
if (l!=c.nrows() || ll!=c.ncols() || k!=kk) laerror("incompatible matrices in NRMat<T>::strassen(...)");
|
|
if(b.ncols() <=0 || c.ncols()<=0) laerror("illegal matrix dimension in strassen");
|
|
#endif
|
|
SAME_LOC3(c, a, b);
|
|
NOT_GPU(c);
|
|
if(cutoff<1) laerror("illegal cutoff in strassen");
|
|
if(alpha==(T)0 && beta==(T)1) return;
|
|
c.copyonwrite();
|
|
s_top(c.nrows(),c.ncols(),k,alpha,(const T *)a,a.ncols(),ta,(const T *)b,b.ncols(),tb,beta,(T *)c,cutoff);
|
|
}
|
|
|
|
#define INSTANTIZE(T) \
|
|
template<> \
|
|
void NRMat<T>::strassen(const T &beta, const NRMat<T> &a, const char transa, const NRMat<T> &b, const char transb, const T &alpha, const int cutoff) \
|
|
{s_strassen(*this,beta,a,transa,b,transb,alpha,cutoff);}
|
|
|
|
INSTANTIZE(double)
|
|
INSTANTIZE(float)
|
|
INSTANTIZE(std::complex<double>)
|
|
INSTANTIZE(std::complex<float>)
|
|
|
|
}//namespace
|