void Scatter
( const DistMatrix<T,CIRC,CIRC>& A,
DistMatrix<T,STAR,STAR>& B )
{
DEBUG_CSE
AssertSameGrids( A, B );
const Int height = A.Height();
const Int width = A.Width();
B.Resize( height, width );
if( B.Participating() )
{
const Int pkgSize = mpi::Pad( height*width );
vector<T> buffer;
FastResize( buffer, pkgSize );
// Pack
if( A.Participating() )
util::InterleaveMatrix
( height, width,
A.LockedBuffer(), 1, A.LDim(),
buffer.data(), 1, height );
// Broadcast from the process that packed
mpi::Broadcast( buffer.data(), pkgSize, A.Root(), A.CrossComm() );
// Unpack
util::InterleaveMatrix
( height, width,
buffer.data(), 1, height,
B.Buffer(), 1, B.LDim() );
}
}
static void Func
( DistMatrix<T,STAR,MD>& A, T center, typename Base<T>::type radius )
{
if( A.Participating() )
{
const int m = A.Height();
const int localWidth = A.LocalWidth();
for( int jLocal=0; jLocal<localWidth; ++jLocal )
for( int i=0; i<m; ++i )
A.SetLocal( i, jLocal, center+radius*SampleUnitBall<T>() );
}
}
static void Func
( DistMatrix<T,MD,STAR>& A, T center, typename Base<T>::type radius )
{
if( A.Participating() )
{
const int n = A.Width();
const int localHeight = A.LocalHeight();
for( int j=0; j<n; ++j )
for( int iLocal=0; iLocal<localHeight; ++iLocal )
A.SetLocal( iLocal, j, center+radius*SampleUnitBall<T>() );
}
}
void AllGather
( const DistMatrix<T, U, V >& A,
DistMatrix<T,Collect<U>(),Collect<V>()>& B )
{
EL_DEBUG_CSE
AssertSameGrids( A, B );
const Int height = A.Height();
const Int width = A.Width();
B.SetGrid( A.Grid() );
B.Resize( height, width );
if( A.Participating() )
{
if( A.DistSize() == 1 )
{
Copy( A.LockedMatrix(), B.Matrix() );
}
else
{
const Int colStride = A.ColStride();
const Int rowStride = A.RowStride();
const Int distStride = colStride*rowStride;
const Int maxLocalHeight = MaxLength(height,colStride);
const Int maxLocalWidth = MaxLength(width,rowStride);
const Int portionSize = mpi::Pad( maxLocalHeight*maxLocalWidth );
vector<T> buf;
FastResize( buf, (distStride+1)*portionSize );
T* sendBuf = &buf[0];
T* recvBuf = &buf[portionSize];
// Pack
util::InterleaveMatrix
( A.LocalHeight(), A.LocalWidth(),
A.LockedBuffer(), 1, A.LDim(),
sendBuf, 1, A.LocalHeight() );
// Communicate
mpi::AllGather
( sendBuf, portionSize, recvBuf, portionSize, A.DistComm() );
// Unpack
util::StridedUnpack
( height, width,
A.ColAlign(), colStride,
A.RowAlign(), rowStride,
recvBuf, portionSize,
B.Buffer(), B.LDim() );
}
}
if( A.Grid().InGrid() && A.CrossComm() != mpi::COMM_SELF )
El::Broadcast( B, A.CrossComm(), A.Root() );
}
void Filter
( const DistMatrix<T,Collect<U>(),Collect<V>()>& A,
DistMatrix<T, U, V >& B )
{
DEBUG_CSE
AssertSameGrids( A, B );
B.Resize( A.Height(), A.Width() );
if( !B.Participating() )
return;
const Int colShift = B.ColShift();
const Int rowShift = B.RowShift();
util::InterleaveMatrix
( B.LocalHeight(), B.LocalWidth(),
A.LockedBuffer(colShift,rowShift), B.ColStride(), B.RowStride()*A.LDim(),
B.Buffer(), 1, B.LDim() );
}
inline
DistMatrix<T,MD,STAR,Int>::DistMatrix( const DistMatrix<T,U,V,Int>& A )
: AbstractDistMatrix<T,Int>(0,0,false,false,0,0,
(A.Participating() ? A.ColRank() : 0),0,
0,0,A.Grid()),
diagPath_(A.diagPath_)
{
#ifndef RELEASE
PushCallStack("DistMatrix[MD,* ]::DistMatrix");
#endif
if( MD != U || STAR != V ||
reinterpret_cast<const DistMatrix<T,MD,STAR,Int>*>(&A) != this )
*this = A;
else
throw std::logic_error("Tried to construct [MD,* ] with itself");
#ifndef RELEASE
PopCallStack();
#endif
}
const DistMatrix<T,STAR,STAR>&
DistMatrix<T,STAR,STAR>::operator=( const DistMatrix<T,CIRC,CIRC>& A )
{
#ifndef RELEASE
CallStackEntry entry("[* ,* ] = [o ,o ]");
this->AssertNotLocked();
#endif
const Grid& g = A.Grid();
const Int m = A.Height();
const Int n = A.Width();
this->ResizeTo( A.Height(), A.Width() );
if( this->Participating() )
{
const Int pkgSize = mpi::Pad( m*n );
T* commBuffer = this->auxMemory_.Require( pkgSize );
if( A.Participating() )
{
// Pack
const Int ALDim = A.LDim();
const T* ABuf = A.LockedBuffer();
for( Int j=0; j<n; ++j )
for( Int i=0; i<m; ++i )
commBuffer[i+j*m] = ABuf[i+j*ALDim];
}
// Broadcast from the process that packed
mpi::Broadcast( commBuffer, pkgSize, A.Root(), g.VCComm() );
// Unpack
T* buffer = this->Buffer();
const Int ldim = this->LDim();
for( Int j=0; j<n; ++j )
for( Int i=0; i<m; ++i )
buffer[i+j*ldim] = commBuffer[i+j*m];
}
return *this;
}
void PartialColAllGather
( const DistMatrix<T, U, V>& A,
DistMatrix<T,Partial<U>(),V>& B )
{
EL_DEBUG_CSE
AssertSameGrids( A, B );
const Int height = A.Height();
const Int width = A.Width();
#ifdef EL_VECTOR_WARNINGS
if( width == 1 && A.Grid().Rank() == 0 )
{
cerr <<
"The vector version of PartialColAllGather is not yet written but "
"would only require modifying the vector version of "
"PartialRowAllGather" << endl;
}
#endif
#ifdef EL_CACHE_WARNINGS
if( width && A.Grid().Rank() == 0 )
{
cerr <<
"PartialColAllGather potentially causes a large amount of cache-"
"thrashing. If possible, avoid it by performing the redistribution"
"on the (conjugate-)transpose" << endl;
}
#endif
B.AlignColsAndResize
( Mod(A.ColAlign(),B.ColStride()), height, width, false, false );
if( !A.Participating() )
return;
EL_DEBUG_ONLY(
if( A.LocalWidth() != A.Width() )
LogicError("This routine assumes rows are not distributed");
)
void ColAllToAllPromote
( const DistMatrix<T, U, V >& A,
DistMatrix<T,Partial<U>(),PartialUnionRow<U,V>()>& B )
{
DEBUG_CSE
AssertSameGrids( A, B );
const Int height = A.Height();
const Int width = A.Width();
B.AlignColsAndResize
( Mod(A.ColAlign(),B.ColStride()), height, width, false, false );
if( !B.Participating() )
return;
const Int colStride = A.ColStride();
const Int colStridePart = A.PartialColStride();
const Int colStrideUnion = A.PartialUnionColStride();
const Int colRankPart = A.PartialColRank();
const Int colDiff = B.ColAlign() - Mod(A.ColAlign(),colStridePart);
const Int maxLocalHeight = MaxLength(height,colStride);
const Int maxLocalWidth = MaxLength(width,colStrideUnion);
const Int portionSize = mpi::Pad( maxLocalHeight*maxLocalWidth );
if( colDiff == 0 )
{
if( A.PartialUnionColStride() == 1 )
{
Copy( A.LockedMatrix(), B.Matrix() );
}
else
{
vector<T> buffer;
FastResize( buffer, 2*colStrideUnion*portionSize );
T* firstBuf = &buffer[0];
T* secondBuf = &buffer[colStrideUnion*portionSize];
// Pack
util::RowStridedPack
( A.LocalHeight(), width,
B.RowAlign(), colStrideUnion,
A.LockedBuffer(), A.LDim(),
firstBuf, portionSize );
// Simultaneously Gather in columns and Scatter in rows
mpi::AllToAll
( firstBuf, portionSize,
secondBuf, portionSize, A.PartialUnionColComm() );
// Unpack
util::PartialColStridedUnpack
( height, B.LocalWidth(),
A.ColAlign(), colStride,
colStrideUnion, colStridePart, colRankPart,
B.ColShift(),
secondBuf, portionSize,
B.Buffer(), B.LDim() );
}
}
else
{
#ifdef EL_UNALIGNED_WARNINGS
if( A.Grid().Rank() == 0 )
cerr << "Unaligned PartialColAllToAllPromote" << endl;
#endif
const Int sendColRankPart = Mod( colRankPart+colDiff, colStridePart );
const Int recvColRankPart = Mod( colRankPart-colDiff, colStridePart );
vector<T> buffer;
FastResize( buffer, 2*colStrideUnion*portionSize );
T* firstBuf = &buffer[0];
T* secondBuf = &buffer[colStrideUnion*portionSize];
// Pack
util::RowStridedPack
( A.LocalHeight(), width,
B.RowAlign(), colStrideUnion,
A.LockedBuffer(), A.LDim(),
secondBuf, portionSize );
// Realign the input
mpi::SendRecv
( secondBuf, colStrideUnion*portionSize, sendColRankPart,
firstBuf, colStrideUnion*portionSize, recvColRankPart,
A.PartialColComm() );
// Simultaneously Scatter in columns and Gather in rows
mpi::AllToAll
( firstBuf, portionSize,
secondBuf, portionSize, A.PartialUnionColComm() );
// Unpack
util::PartialColStridedUnpack
( height, B.LocalWidth(),
A.ColAlign(), colStride,
colStrideUnion, colStridePart, recvColRankPart,
B.ColShift(),
secondBuf, portionSize,
B.Buffer(), B.LDim() );
}
}
void TranslateBetweenGrids
( const DistMatrix<T,MC,MR>& A, DistMatrix<T,MC,MR>& B )
{
DEBUG_ONLY(CSE cse("copy::TranslateBetweenGrids [MC,MR]"))
B.Resize( A.Height(), A.Width() );
// Just need to ensure that each viewing comm contains the other team's
// owning comm. Congruence is too strong.
// Compute the number of process rows and columns that each process
// needs to send to.
const Int colStride = B.ColStride();
const Int rowStride = B.RowStride();
const Int colRank = B.ColRank();
const Int rowRank = B.RowRank();
const Int colStrideA = A.ColStride();
const Int rowStrideA = A.RowStride();
const Int colGCD = GCD( colStride, colStrideA );
const Int rowGCD = GCD( rowStride, rowStrideA );
const Int colLCM = colStride*colStrideA / colGCD;
const Int rowLCM = rowStride*rowStrideA / rowGCD;
const Int numColSends = colStride / colGCD;
const Int numRowSends = rowStride / rowGCD;
const Int colAlign = B.ColAlign();
const Int rowAlign = B.RowAlign();
const Int colAlignA = A.ColAlign();
const Int rowAlignA = A.RowAlign();
const bool inBGrid = B.Participating();
const bool inAGrid = A.Participating();
if( !inBGrid && !inAGrid )
return;
const Int maxSendSize =
(A.Height()/(colStrideA*numColSends)+1) *
(A.Width()/(rowStrideA*numRowSends)+1);
// Translate the ranks from A's VC communicator to B's viewing so that
// we can match send/recv communicators. Since A's VC communicator is not
// necessarily defined on every process, we instead work with A's owning
// group and account for row-major ordering if necessary.
const int sizeA = A.Grid().Size();
vector<int> rankMap(sizeA), ranks(sizeA);
if( A.Grid().Order() == COLUMN_MAJOR )
{
for( int j=0; j<sizeA; ++j )
ranks[j] = j;
}
else
{
// The (i,j) = i + j*colStrideA rank in the column-major ordering is
// equal to the j + i*rowStrideA rank in a row-major ordering.
// Since we desire rankMap[i+j*colStrideA] to correspond to process
// (i,j) in A's grid's rank in this viewing group, ranks[i+j*colStrideA]
// should correspond to process (i,j) in A's owning group. Since the
// owning group is ordered row-major in this case, its rank is
// j+i*rowStrideA. Note that setting
// ranks[j+i*rowStrideA] = i+j*colStrideA is *NOT* valid.
for( int i=0; i<colStrideA; ++i )
for( int j=0; j<rowStrideA; ++j )
ranks[i+j*colStrideA] = j+i*rowStrideA;
}
mpi::Translate
( A.Grid().OwningGroup(), sizeA, &ranks[0],
B.Grid().ViewingComm(), &rankMap[0] );
// Have each member of A's grid individually send to all numRow x numCol
// processes in order, while the members of this grid receive from all
// necessary processes at each step.
Int requiredMemory = 0;
if( inAGrid )
requiredMemory += maxSendSize;
if( inBGrid )
requiredMemory += maxSendSize;
vector<T> auxBuf( requiredMemory );
Int offset = 0;
T* sendBuf = &auxBuf[offset];
if( inAGrid )
offset += maxSendSize;
T* recvBuf = &auxBuf[offset];
Int recvRow = 0; // avoid compiler warnings...
if( inAGrid )
recvRow = Mod(Mod(A.ColRank()-colAlignA,colStrideA)+colAlign,colStride);
for( Int colSend=0; colSend<numColSends; ++colSend )
{
Int recvCol = 0; // avoid compiler warnings...
if( inAGrid )
recvCol=Mod(Mod(A.RowRank()-rowAlignA,rowStrideA)+rowAlign,
rowStride);
for( Int rowSend=0; rowSend<numRowSends; ++rowSend )
{
mpi::Request sendRequest;
// Fire off this round of non-blocking sends
if( inAGrid )
{
// Pack the data
Int sendHeight = Length(A.LocalHeight(),colSend,numColSends);
Int sendWidth = Length(A.LocalWidth(),rowSend,numRowSends);
copy::util::InterleaveMatrix
( sendHeight, sendWidth,
A.LockedBuffer(colSend,rowSend),
numColSends, numRowSends*A.LDim(),
sendBuf, 1, sendHeight );
// Send data
const Int recvVCRank = recvRow + recvCol*colStride;
const Int recvViewingRank = B.Grid().VCToViewing( recvVCRank );
mpi::ISend
( sendBuf, sendHeight*sendWidth, recvViewingRank,
B.Grid().ViewingComm(), sendRequest );
}
// Perform this round of recv's
if( inBGrid )
{
const Int sendColOffset = colAlignA;
const Int recvColOffset =
(colSend*colStrideA+colAlign) % colStride;
const Int sendRowOffset = rowAlignA;
const Int recvRowOffset =
(rowSend*rowStrideA+rowAlign) % rowStride;
const Int firstSendRow =
Mod( Mod(colRank-recvColOffset,colStride)+sendColOffset,
colStrideA );
const Int firstSendCol =
Mod( Mod(rowRank-recvRowOffset,rowStride)+sendRowOffset,
rowStrideA );
const Int colShift = Mod( colRank-recvColOffset, colStride );
const Int rowShift = Mod( rowRank-recvRowOffset, rowStride );
const Int numColRecvs = Length( colStrideA, colShift, colStride );
const Int numRowRecvs = Length( rowStrideA, rowShift, rowStride );
// Recv data
// For now, simply receive sequentially. Until we switch to
// nonblocking recv's, we won't be using much of the
// recvBuf
Int sendRow = firstSendRow;
for( Int colRecv=0; colRecv<numColRecvs; ++colRecv )
{
const Int sendColShift = Shift( sendRow, colAlignA, colStrideA ) + colSend*colStrideA;
const Int sendHeight = Length( A.Height(), sendColShift, colLCM );
const Int localColOffset = (sendColShift-B.ColShift()) / colStride;
Int sendCol = firstSendCol;
for( Int rowRecv=0; rowRecv<numRowRecvs; ++rowRecv )
{
const Int sendRowShift = Shift( sendCol, rowAlignA, rowStrideA ) + rowSend*rowStrideA;
const Int sendWidth = Length( A.Width(), sendRowShift, rowLCM );
const Int localRowOffset = (sendRowShift-B.RowShift()) / rowStride;
const Int sendVCRank = sendRow+sendCol*colStrideA;
mpi::Recv
( recvBuf, sendHeight*sendWidth, rankMap[sendVCRank],
B.Grid().ViewingComm() );
// Unpack the data
copy::util::InterleaveMatrix
( sendHeight, sendWidth,
recvBuf, 1, sendHeight,
B.Buffer(localColOffset,localRowOffset),
colLCM/colStride, (rowLCM/rowStride)*B.LDim() );
// Set up the next send col
sendCol = (sendCol + rowStride) % rowStrideA;
}
// Set up the next send row
sendRow = (sendRow + colStride) % colStrideA;
}
}
// Ensure that this round of non-blocking sends completes
if( inAGrid )
{
mpi::Wait( sendRequest );
recvCol = (recvCol + rowStrideA) % rowStride;
}
}
if( inAGrid )
recvRow = (recvRow + colStrideA) % colStride;
}
}
const DistMatrix<T,STAR,STAR>&
DistMatrix<T,STAR,STAR>::operator=( const DistMatrix<T,STAR,MD>& A )
{
#ifndef RELEASE
CallStackEntry entry("[* ,* ] = [* ,MD]");
this->AssertNotLocked();
this->AssertSameGrid( A.Grid() );
#endif
const elem::Grid& g = this->Grid();
this->ResizeTo( A.Height(), A.Width() );
if( !this->Participating() )
return *this;
const Int p = g.Size();
const Int lcm = g.LCM();
const Int ownerPath = A.diagPath_;
const Int ownerPathRank = A.rowAlignment_;
const Int height = this->Height();
const Int width = this->Width();
const Int localWidth = A.LocalWidth();
const Int maxLocalWidth = MaxLength( width, lcm );
const Int portionSize = mpi::Pad( height*maxLocalWidth );
// Since a MD communicator has not been implemented, we will take
// the suboptimal route of 'rounding up' everyone's contribution over
// the VC communicator.
T* buffer = this->auxMemory_.Require( (p+1)*portionSize );
T* sendBuf = &buffer[0];
T* recvBuf = &buffer[portionSize];
// Pack
if( A.Participating() )
{
const Int ALDim = A.LDim();
const T* ABuf = A.LockedBuffer();
PARALLEL_FOR
for( Int jLoc=0; jLoc<localWidth; ++jLoc )
MemCopy( &sendBuf[jLoc*height], &ABuf[jLoc*ALDim], height );
}
// Communicate
mpi::AllGather
( sendBuf, portionSize,
recvBuf, portionSize, g.VCComm() );
// Unpack
T* thisBuf = this->Buffer();
const Int thisLDim = this->LDim();
OUTER_PARALLEL_FOR
for( Int k=0; k<p; ++k )
{
if( g.DiagPath( k ) == ownerPath )
{
const T* data = &recvBuf[k*portionSize];
const Int thisPathRank = g.DiagPathRank( k );
const Int thisRowShift = Shift_( thisPathRank, ownerPathRank, lcm );
const Int thisLocalWidth = Length_( width, thisRowShift, lcm );
INNER_PARALLEL_FOR
for( Int jLoc=0; jLoc<thisLocalWidth; ++jLoc )
MemCopy
( &thisBuf[(thisRowShift+jLoc*lcm)*thisLDim],
&data[jLoc*height], height );
}
}
this->auxMemory_.Release();
return *this;
}
void AllGather
( const DistMatrix<T, U, V >& A,
DistMatrix<T,Collect<U>(),Collect<V>()>& B )
{
DEBUG_ONLY(CSE cse("copy::AllGather"))
AssertSameGrids( A, B );
const Int height = A.Height();
const Int width = A.Width();
B.SetGrid( A.Grid() );
B.Resize( height, width );
if( A.Participating() )
{
const Int colStride = A.ColStride();
const Int rowStride = A.RowStride();
const Int distStride = colStride*rowStride;
const Int maxLocalHeight = MaxLength(height,colStride);
const Int maxLocalWidth = MaxLength(width,rowStride);
const Int portionSize = mpi::Pad( maxLocalHeight*maxLocalWidth );
vector<T> buf( (distStride+1)*portionSize );
T* sendBuf = &buf[0];
T* recvBuf = &buf[portionSize];
// Pack
util::InterleaveMatrix
( A.LocalHeight(), A.LocalWidth(),
A.LockedBuffer(), 1, A.LDim(),
sendBuf, 1, A.LocalHeight() );
// Communicate
mpi::AllGather
( sendBuf, portionSize, recvBuf, portionSize, A.DistComm() );
// Unpack
util::StridedUnpack
( height, width,
A.ColAlign(), colStride,
A.RowAlign(), rowStride,
recvBuf, portionSize,
B.Buffer(), B.LDim() );
}
if( A.Grid().InGrid() && A.CrossComm() != mpi::COMM_SELF )
{
// Pack from the root
const Int BLocalHeight = B.LocalHeight();
const Int BLocalWidth = B.LocalWidth();
vector<T> buf(BLocalHeight*BLocalWidth);
if( A.CrossRank() == A.Root() )
util::InterleaveMatrix
( BLocalHeight, BLocalWidth,
B.LockedBuffer(), 1, B.LDim(),
buf.data(), 1, BLocalHeight );
// Broadcast from the root
mpi::Broadcast
( buf.data(), BLocalHeight*BLocalWidth, A.Root(), A.CrossComm() );
// Unpack if not the root
if( A.CrossRank() != A.Root() )
util::InterleaveMatrix
( BLocalHeight, BLocalWidth,
buf.data(), 1, BLocalHeight,
B.Buffer(), 1, B.LDim() );
}
}
void TransposeDist( const DistMatrix<T,U,V>& A, DistMatrix<T,V,U>& B )
{
DEBUG_ONLY(CSE cse("copy::TransposeDist"))
AssertSameGrids( A, B );
const Grid& g = B.Grid();
B.Resize( A.Height(), A.Width() );
if( !B.Participating() )
return;
const Int colStrideA = A.ColStride();
const Int rowStrideA = A.RowStride();
const Int distSize = A.DistSize();
if( A.DistSize() == 1 && B.DistSize() == 1 )
{
Copy( A.LockedMatrix(), B.Matrix() );
}
else if( A.Width() == 1 )
{
const Int height = A.Height();
const Int maxLocalHeight = MaxLength(height,distSize);
const Int portionSize = mpi::Pad( maxLocalHeight );
const Int colDiff = Shift(A.DistRank(),A.ColAlign(),distSize) -
Shift(B.DistRank(),B.ColAlign(),distSize);
const Int sendRankB = Mod( B.DistRank()+colDiff, distSize );
const Int recvRankA = Mod( A.DistRank()-colDiff, distSize );
const Int recvRankB =
(recvRankA/colStrideA)+rowStrideA*(recvRankA%colStrideA);
vector<T> buffer;
FastResize( buffer, (colStrideA+rowStrideA)*portionSize );
T* sendBuf = &buffer[0];
T* recvBuf = &buffer[colStrideA*portionSize];
if( A.RowRank() == A.RowAlign() )
{
// Pack
// TODO: Use kernel from copy::util
const Int AColShift = A.ColShift();
const T* ABuf = A.LockedBuffer();
EL_PARALLEL_FOR
for( Int k=0; k<rowStrideA; ++k )
{
T* data = &recvBuf[k*portionSize];
const Int shift =
Shift_(A.ColRank()+colStrideA*k,A.ColAlign(),distSize);
const Int offset = (shift-AColShift) / colStrideA;
const Int thisLocalHeight = Length_(height,shift,distSize);
for( Int iLoc=0; iLoc<thisLocalHeight; ++iLoc )
data[iLoc] = ABuf[offset+iLoc*rowStrideA];
}
}
// (e.g., A[VC,STAR] <- A[MC,MR])
mpi::Scatter
( recvBuf, portionSize,
sendBuf, portionSize, A.RowAlign(), A.RowComm() );
// (e.g., A[VR,STAR] <- A[VC,STAR])
mpi::SendRecv
( sendBuf, portionSize, sendRankB,
recvBuf, portionSize, recvRankB, B.DistComm() );
// (e.g., A[MR,MC] <- A[VR,STAR])
mpi::Gather
( recvBuf, portionSize,
sendBuf, portionSize, B.RowAlign(), B.RowComm() );
if( B.RowRank() == B.RowAlign() )
{
// Unpack
// TODO: Use kernel from copy::util
T* bufB = B.Buffer();
EL_PARALLEL_FOR
for( Int k=0; k<colStrideA; ++k )
{
const T* data = &sendBuf[k*portionSize];
const Int shift =
Shift_(B.ColRank()+rowStrideA*k,B.ColAlign(),distSize);
const Int offset = (shift-B.ColShift()) / rowStrideA;
const Int thisLocalHeight = Length_(height,shift,distSize);
for( Int iLoc=0; iLoc<thisLocalHeight; ++iLoc )
bufB[offset+iLoc*colStrideA] = data[iLoc];
}
}
}
void ColAllToAllDemote
( const DistMatrix<T,Partial<U>(),PartialUnionRow<U,V>()>& A,
DistMatrix<T, U, V >& B )
{
DEBUG_ONLY(CallStackEntry cse("copy::ColAllToAllDemote"))
AssertSameGrids( A, B );
const Int height = A.Height();
const Int width = A.Width();
B.AlignColsAndResize( A.ColAlign(), height, width, false, false );
if( !B.Participating() )
return;
const Int colAlign = B.ColAlign();
const Int rowAlignA = A.RowAlign();
const Int colStride = B.ColStride();
const Int colStridePart = B.PartialColStride();
const Int colStrideUnion = B.PartialUnionColStride();
const Int colRankPart = B.PartialColRank();
const Int colDiff = (colAlign%colStridePart) - A.ColAlign();
const Int colShiftA = A.ColShift();
const Int localHeightB = B.LocalHeight();
const Int localWidthA = A.LocalWidth();
const Int maxLocalHeight = MaxLength(height,colStride);
const Int maxLocalWidth = MaxLength(width,colStrideUnion);
const Int portionSize = mpi::Pad( maxLocalHeight*maxLocalWidth );
std::vector<T> buffer( 2*colStrideUnion*portionSize );
T* firstBuf = &buffer[0];
T* secondBuf = &buffer[colStrideUnion*portionSize];
if( colDiff == 0 )
{
// Pack
util::PartialColStridedPack
( height, localWidthA,
colAlign, colStride,
colStrideUnion, colStridePart, colRankPart,
colShiftA,
A.LockedBuffer(), A.LDim(),
firstBuf, portionSize );
// Simultaneously Scatter in columns and Gather in rows
mpi::AllToAll
( firstBuf, portionSize,
secondBuf, portionSize, B.PartialUnionColComm() );
// Unpack
util::RowStridedUnpack
( localHeightB, width,
rowAlignA, colStrideUnion,
secondBuf, portionSize,
B.Buffer(), B.LDim() );
}
else
{
#ifdef EL_UNALIGNED_WARNINGS
if( B.Grid().Rank() == 0 )
std::cerr << "Unaligned ColAllToAllDemote" << std::endl;
#endif
const Int sendColRankPart = Mod( colRankPart+colDiff, colStridePart );
const Int recvColRankPart = Mod( colRankPart-colDiff, colStridePart );
// Pack
util::PartialColStridedPack
( height, localWidthA,
colAlign, colStride,
colStrideUnion, colStridePart, sendColRankPart,
colShiftA,
A.LockedBuffer(), A.LDim(),
secondBuf, portionSize );
// Simultaneously Scatter in columns and Gather in rows
mpi::AllToAll
( secondBuf, portionSize,
firstBuf, portionSize, B.PartialUnionColComm() );
// Realign the result
mpi::SendRecv
( firstBuf, colStrideUnion*portionSize, sendColRankPart,
secondBuf, colStrideUnion*portionSize, recvColRankPart,
B.PartialColComm() );
// Unpack
util::RowStridedUnpack
( localHeightB, width,
rowAlignA, colStrideUnion,
secondBuf, portionSize,
B.Buffer(), B.LDim() );
}
}