void PartialRowScatter
( T alpha,
const ElementalMatrix<T>& A,
ElementalMatrix<T>& B )
{
DEBUG_ONLY(CSE cse("axpy_contract::PartialRowScatter"))
AssertSameGrids( A, B );
if( A.Height() != B.Height() || A.Width() != B.Width() )
LogicError("Matrix sizes did not match");
if( !B.Participating() )
return;
if( B.RowAlign() % A.RowStride() == A.RowAlign() )
{
const Int rowStride = B.RowStride();
const Int rowStridePart = B.PartialRowStride();
const Int rowStrideUnion = B.PartialUnionRowStride();
const Int rowRankPart = B.PartialRowRank();
const Int height = B.Height();
const Int width = B.Width();
const Int maxLocalWidth = MaxLength( width, rowStride );
const Int recvSize = mpi::Pad( height*maxLocalWidth );
const Int sendSize = rowStrideUnion*recvSize;
//vector<T> buffer( sendSize );
vector<T> buffer;
buffer.reserve( sendSize );
// Pack
copy::util::PartialRowStridedPack
( height, width,
B.RowAlign(), rowStride,
rowStrideUnion, rowStridePart, rowRankPart,
A.RowShift(),
A.LockedBuffer(), A.LDim(),
buffer.data(), recvSize );
// Communicate
mpi::ReduceScatter( buffer.data(), recvSize, B.PartialUnionRowComm() );
// Unpack our received data
axpy::util::InterleaveMatrixUpdate
( alpha, height, B.LocalWidth(),
buffer.data(), 1, height,
B.Buffer(), 1, B.LDim() );
}
else
LogicError("Unaligned PartialRowScatter not implemented");
}
void Gather
( const ElementalMatrix<T>& A,
DistMatrix<T,CIRC,CIRC>& B )
{
DEBUG_ONLY(CSE cse("copy::Gather"))
AssertSameGrids( A, B );
if( A.DistSize() == 1 && A.CrossSize() == 1 )
{
B.Resize( A.Height(), A.Width() );
if( B.CrossRank() == B.Root() )
Copy( A.LockedMatrix(), B.Matrix() );
return;
}
const Int height = A.Height();
const Int width = A.Width();
B.SetGrid( A.Grid() );
B.Resize( height, width );
// Gather the colShifts and rowShifts
// ==================================
Int myShifts[2];
myShifts[0] = A.ColShift();
myShifts[1] = A.RowShift();
vector<Int> shifts;
const Int crossSize = B.CrossSize();
if( B.CrossRank() == B.Root() )
shifts.resize( 2*crossSize );
mpi::Gather( myShifts, 2, shifts.data(), 2, B.Root(), B.CrossComm() );
// Gather the payload data
// =======================
const bool irrelevant = ( A.RedundantRank()!=0 || A.CrossRank()!=A.Root() );
int totalSend = ( irrelevant ? 0 : A.LocalHeight()*A.LocalWidth() );
vector<int> recvCounts, recvOffsets;
if( B.CrossRank() == B.Root() )
recvCounts.resize( crossSize );
mpi::Gather( &totalSend, 1, recvCounts.data(), 1, B.Root(), B.CrossComm() );
int totalRecv = Scan( recvCounts, recvOffsets );
//vector<T> sendBuf(totalSend), recvBuf(totalRecv);
vector<T> sendBuf, recvBuf;
sendBuf.reserve( totalSend );
recvBuf.reserve( totalRecv );
if( !irrelevant )
copy::util::InterleaveMatrix
( A.LocalHeight(), A.LocalWidth(),
A.LockedBuffer(), 1, A.LDim(),
sendBuf.data(), 1, A.LocalHeight() );
mpi::Gather
( sendBuf.data(), totalSend,
recvBuf.data(), recvCounts.data(), recvOffsets.data(),
B.Root(), B.CrossComm() );
// Unpack
// ======
if( B.Root() == B.CrossRank() )
{
for( Int q=0; q<crossSize; ++q )
{
if( recvCounts[q] == 0 )
continue;
const Int colShift = shifts[2*q+0];
const Int rowShift = shifts[2*q+1];
const Int colStride = A.ColStride();
const Int rowStride = A.RowStride();
const Int localHeight = Length( height, colShift, colStride );
const Int localWidth = Length( width, rowShift, rowStride );
copy::util::InterleaveMatrix
( localHeight, localWidth,
&recvBuf[recvOffsets[q]], 1, localHeight,
B.Buffer(colShift,rowShift), colStride, rowStride*B.LDim() );
}
}
}