void SUMMA_TNA
( Orientation orientA,
T alpha,
const AbstractDistMatrix<T>& APre,
const AbstractDistMatrix<T>& BPre,
AbstractDistMatrix<T>& CPre )
{
DEBUG_CSE
const Int n = CPre.Width();
const Int bsize = Blocksize();
const Grid& g = APre.Grid();
DistMatrixReadProxy<T,T,MC,MR> AProx( APre );
DistMatrixReadProxy<T,T,MC,MR> BProx( BPre );
DistMatrixReadWriteProxy<T,T,MC,MR> CProx( CPre );
auto& A = AProx.GetLocked();
auto& B = BProx.GetLocked();
auto& C = CProx.Get();
// Temporary distributions
DistMatrix<T,MC,STAR> B1_MC_STAR(g);
DistMatrix<T,MR,STAR> D1_MR_STAR(g);
DistMatrix<T,MR,MC > D1_MR_MC(g);
B1_MC_STAR.AlignWith( A );
D1_MR_STAR.AlignWith( A );
for( Int k=0; k<n; k+=bsize )
{
const Int nb = Min(bsize,n-k);
auto B1 = B( ALL, IR(k,k+nb) );
auto C1 = C( ALL, IR(k,k+nb) );
// D1[MR,*] := alpha (A1[MC,MR])^T B1[MC,*]
// = alpha (A1^T)[MR,MC] B1[MC,*]
B1_MC_STAR = B1;
LocalGemm( orientA, NORMAL, alpha, A, B1_MC_STAR, D1_MR_STAR );
// C1[MC,MR] += scattered & transposed D1[MR,*] summed over grid cols
Contract( D1_MR_STAR, D1_MR_MC );
Axpy( T(1), D1_MR_MC, C1 );
}
}
inline void
GemmTTA
( Orientation orientationOfA,
Orientation orientationOfB,
T alpha, const DistMatrix<T>& A,
const DistMatrix<T>& B,
T beta, DistMatrix<T>& C )
{
#ifndef RELEASE
PushCallStack("internal::GemmTTA");
if( A.Grid() != B.Grid() || B.Grid() != C.Grid() )
throw std::logic_error
("{A,B,C} must be distributed over the same grid");
if( orientationOfA == NORMAL || orientationOfB == NORMAL )
throw std::logic_error
("GemmTTA expects A and B to be (Conjugate)Transposed");
if( A.Width() != C.Height() ||
B.Height() != C.Width() ||
A.Height() != B.Width() )
{
std::ostringstream msg;
msg << "Nonconformal GemmTTA: \n"
<< " A ~ " << A.Height() << " x " << A.Width() << "\n"
<< " B ~ " << B.Height() << " x " << B.Width() << "\n"
<< " C ~ " << C.Height() << " x " << C.Width() << "\n";
throw std::logic_error( msg.str().c_str() );
}
#endif
const Grid& g = A.Grid();
// Matrix views
DistMatrix<T> BT(g), B0(g),
BB(g), B1(g),
B2(g);
DistMatrix<T> CL(g), CR(g),
C0(g), C1(g), C2(g);
// Temporary distributions
DistMatrix<T,STAR,MC > B1_STAR_MC(g);
DistMatrix<T,MR, STAR> D1_MR_STAR(g);
DistMatrix<T,MR, MC > D1_MR_MC(g);
DistMatrix<T> D1(g);
B1_STAR_MC.AlignWith( A );
D1_MR_STAR.AlignWith( A );
// Start the algorithm
Scale( beta, C );
LockedPartitionDown
( B, BT,
BB, 0 );
PartitionRight( C, CL, CR, 0 );
while( BB.Height() > 0 )
{
LockedRepartitionDown
( BT, B0,
/**/ /**/
B1,
BB, B2 );
RepartitionRight
( CL, /**/ CR,
C0, /**/ C1, C2 );
D1.AlignWith( C1 );
Zeros( C1.Height(), C1.Width(), D1_MR_STAR );
//--------------------------------------------------------------------//
B1_STAR_MC = B1; // B1[*,MC] <- B1[MC,MR]
// D1[MR,*] := alpha (A[MC,MR])^T (B1[*,MC])^T
// = alpha (A^T)[MR,MC] (B1^T)[MC,*]
LocalGemm
( orientationOfA, orientationOfB,
alpha, A, B1_STAR_MC, T(0), D1_MR_STAR );
// C1[MC,MR] += scattered & transposed D1[MR,*] summed over grid cols
D1_MR_MC.SumScatterFrom( D1_MR_STAR );
D1 = D1_MR_MC;
Axpy( T(1), D1, C1 );
//--------------------------------------------------------------------//
D1.FreeAlignments();
SlideLockedPartitionDown
( BT, B0,
B1,
/**/ /**/
BB, B2 );
SlidePartitionRight
( CL, /**/ CR,
C0, C1, /**/ C2 );
}
#ifndef RELEASE
PopCallStack();
#endif
}
inline void
internal::GemmTNA
( Orientation orientationOfA,
T alpha, const DistMatrix<T,MC,MR>& A,
const DistMatrix<T,MC,MR>& B,
T beta, DistMatrix<T,MC,MR>& C )
{
#ifndef RELEASE
PushCallStack("internal::GemmTNA");
if( A.Grid() != B.Grid() || B.Grid() != C.Grid() )
throw std::logic_error
("{A,B,C} must be distributed over the same grid");
if( orientationOfA == NORMAL )
throw std::logic_error("GemmTNA assumes A is (Conjugate)Transposed");
if( A.Width() != C.Height() ||
B.Width() != C.Width() ||
A.Height() != B.Height() )
{
std::ostringstream msg;
msg << "Nonconformal GemmTNA: \n"
<< " A ~ " << A.Height() << " x " << A.Width() << "\n"
<< " B ~ " << B.Height() << " x " << B.Width() << "\n"
<< " C ~ " << C.Height() << " x " << C.Width() << "\n";
throw std::logic_error( msg.str().c_str() );
}
#endif
const Grid& g = A.Grid();
// Matrix views
DistMatrix<T,MC,MR> BL(g), BR(g),
B0(g), B1(g), B2(g);
DistMatrix<T,MC,MR> CL(g), CR(g),
C0(g), C1(g), C2(g);
// Temporary distributions
DistMatrix<T,MC,STAR> B1_MC_STAR(g);
DistMatrix<T,MR,STAR> D1_MR_STAR(g);
DistMatrix<T,MR,MC > D1_MR_MC(g);
DistMatrix<T,MC,MR > D1(g);
// Start the algorithm
Scal( beta, C );
LockedPartitionRight( B, BL, BR, 0 );
PartitionRight( C, CL, CR, 0 );
while( BR.Width() > 0 )
{
LockedRepartitionRight
( BL, /**/ BR,
B0, /**/ B1, B2 );
RepartitionRight
( CL, /**/ CR,
C0, /**/ C1, C2 );
B1_MC_STAR.AlignWith( A );
D1_MR_STAR.AlignWith( A );
D1_MR_STAR.ResizeTo( C1.Height(), C1.Width() );
D1.AlignWith( C1 );
//--------------------------------------------------------------------//
B1_MC_STAR = B1; // B1[MC,*] <- B1[MC,MR]
// D1[MR,*] := alpha (A1[MC,MR])^T B1[MC,*]
// = alpha (A1^T)[MR,MC] B1[MC,*]
internal::LocalGemm
( orientationOfA, NORMAL, alpha, A, B1_MC_STAR, (T)0, D1_MR_STAR );
// C1[MC,MR] += scattered & transposed D1[MR,*] summed over grid cols
D1_MR_MC.SumScatterFrom( D1_MR_STAR );
D1 = D1_MR_MC;
Axpy( (T)1, D1, C1 );
//--------------------------------------------------------------------//
B1_MC_STAR.FreeAlignments();
D1_MR_STAR.FreeAlignments();
D1.FreeAlignments();
SlideLockedPartitionRight
( BL, /**/ BR,
B0, B1, /**/ B2 );
SlidePartitionRight
( CL, /**/ CR,
C0, C1, /**/ C2 );
}
#ifndef RELEASE
PopCallStack();
#endif
}
