SequencePacker.cpp
//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE.md file in the project root for full license information.
//
#define _CRT_SECURE_NO_WARNINGS
#define _SCL_SECURE_NO_WARNINGS
#include <numeric>
#include <inttypes.h>
#include "SequencePacker.h"
#include "ElementTypeUtils.h"
namespace Microsoft { namespace MSR { namespace CNTK {
MBLayoutPtr SequencePacker::CreateMBLayout(const StreamBatch& batch)
{
vector<MBLayout::SequenceInfo> infos;
for (size_t index = 0; index < batch.size(); ++index)
{
MBLayout::SequenceInfo info;
info.seqId = index;
info.tBegin = 0;
info.tEnd = batch[index]->m_numberOfSamples;
infos.push_back(info);
}
vector<pair<size_t, size_t>> placement;
vector<size_t> rowAllocations;
// Creating the minibatch layout.
MBLayoutPtr pMBLayout = make_shared<MBLayout>();
pMBLayout->InitAsPackedSequences(infos, placement, rowAllocations);
return pMBLayout;
}
Minibatch SequencePacker::ReadMinibatch()
{
auto sequences = m_transformer->GetNextSequences(m_minibatchSize);
const auto& batch = sequences.m_data;
Minibatch minibatch(sequences.m_endOfEpoch);
if (batch.empty())
{
return minibatch;
}
assert(m_outputStreamDescriptions.size() == batch.size());
for (int streamIndex = 0; streamIndex < batch.size(); ++streamIndex)
{
const auto& streamBatch = batch[streamIndex];
const auto& type = m_outputStreamDescriptions[streamIndex]->m_storageType;
auto pMBLayout = (type == StorageType::dense) ?
PackDenseStream(streamBatch, streamIndex) : PackSparseStream(streamBatch, streamIndex);
auto& buffer = m_streamBuffers[streamIndex];
auto streamMinibatch = std::make_shared<StreamMinibatch>();
streamMinibatch->m_data = buffer.m_data.get();
streamMinibatch->m_layout = pMBLayout;
minibatch.m_data.push_back(streamMinibatch);
}
return minibatch;
}
MBLayoutPtr SequencePacker::PackDenseStream(const StreamBatch& batch, size_t streamIndex)
{
assert(m_outputStreamDescriptions[streamIndex]->m_storageType == StorageType::dense);
const auto& stream = m_inputStreamDescriptions[streamIndex];
auto& buffer = m_streamBuffers[streamIndex];
size_t sampleSize = GetSampleSize(stream);
auto pMBLayout = CreateMBLayout(batch);
size_t requiredSize = pMBLayout->GetNumCols() * sampleSize;
if (buffer.m_size < requiredSize)
{
buffer.Resize(requiredSize);
}
auto elementSize = GetSizeByType(stream->m_elementType);
const auto& sequenceInfos = pMBLayout->GetAllSequences();
// Iterate over sequences in the layout, copy samples from the
// source sequences into the buffer (at appropriate offsets).
for (const auto& sequenceInfo : sequenceInfos)
{
// skip gaps
if (sequenceInfo.seqId == GAP_SEQUENCE_ID)
{
continue;
}
const auto& sequence = batch[sequenceInfo.seqId];
size_t numSamples = sequence->m_numberOfSamples;
assert(numSamples == sequenceInfo.GetNumTimeSteps());
char* bufferPtr = buffer.m_data.get();
// Iterate over all samples in the sequence, keep track of the sample offset (which is especially
// important for sparse input, where offset == number of preceding nnz elements).
for (size_t sampleIndex = 0, sampleOffset = 0; sampleIndex < numSamples; ++sampleIndex)
{
// Compute the offset into the destination buffer, using the layout information
// to get the column index corresponding to the given sample.
auto destinationOffset = pMBLayout->GetColumnIndex(sequenceInfo, sampleIndex) * sampleSize;
// verify that there's enough space left in the buffer to fit a full sample.
assert(destinationOffset <= buffer.m_size - sampleSize);
auto* destination = bufferPtr + destinationOffset;
if (stream->m_storageType == StorageType::dense)
{
// verify that the offset (an invariant for dense).
assert(sampleOffset == sampleIndex * sampleSize);
PackDenseSample(destination, sequence, sampleOffset, sampleSize);
sampleOffset += sampleSize;
}
else if (stream->m_storageType == StorageType::sparse_csc)
{
// TODO: make type casts members of the SparseSequenceData
SparseSequenceDataPtr sparseSequence = static_pointer_cast<SparseSequenceData>(sequence);
// make sure that the sequence meta-data is correct.
assert(numSamples == sparseSequence->m_nnzCounts.size());
PackSparseSampleAsDense(destination, sparseSequence, sampleIndex, sampleOffset, sampleSize, elementSize);
// move the offset by nnz count of the sample.
sampleOffset += sparseSequence->m_nnzCounts[sampleIndex];
// verify that the offset is within the bounds (less or equal
// to the total nnz count of the sequence).
assert(sampleOffset <= sparseSequence->m_totalNnzCount);
}
else
{
RuntimeError("Storage type %d is not supported.", (int)stream->m_storageType);
}
}
}
return pMBLayout;
}
MBLayoutPtr SequencePacker::PackSparseStream(const StreamBatch& batch, size_t streamIndex)
{
assert(m_outputStreamDescriptions[streamIndex]->m_storageType == StorageType::sparse_csc);
// compute the aggregate nnz count of all the sequence in the batch.
size_t nnzCount = 0;
for (const auto& sequence : batch)
{
SparseSequenceDataPtr sparseSequence = static_pointer_cast<SparseSequenceData>(sequence);
nnzCount += sparseSequence->m_totalNnzCount;
}
if (nnzCount > numeric_limits<IndexType>::max())
{
RuntimeError("Minibatch NNZ count (%" PRIu64 ") exceeds the maximum allowed "
"value (%" PRIu64 ")\n", nnzCount, (size_t)numeric_limits<IndexType>::max());
}
const auto& stream = m_inputStreamDescriptions[streamIndex];
assert(stream->m_storageType == StorageType::sparse_csc);
auto elementSize = GetSizeByType(stream->m_elementType);
auto indexSize = sizeof(IndexType);
auto pMBLayout = CreateMBLayout(batch);
// Compute the required buffer size:
// size of nnz type + nnz * (size of the element type) + nnz * (size of the row index type) +
// (number of columns + 1) * (size of the column index type).
size_t requiredSize =
sizeof(nnzCount) +
nnzCount * (elementSize + indexSize) +
indexSize * (pMBLayout->GetNumCols() + 1);
auto& buffer = m_streamBuffers[streamIndex];
if (buffer.m_size < requiredSize)
{
buffer.Resize(requiredSize);
}
auto* destination = buffer.m_data.get();
// insert the nnzCount as the first element in the buffer.
memcpy(destination, &nnzCount, sizeof(nnzCount));
// create two pointers to the memory blocks inside the buffer,
// one for data portion and anther -- for indices.
auto* dataDst = destination + sizeof(nnzCount);
auto* indicesDst = dataDst + elementSize* nnzCount;
// column index for the current sample (= number of nnz value packed so far).
IndexType columnOffset = 0;
// a vector to store column index for each sample in the resulting (packed) matrix.
vector<IndexType> sparseColumnIndices;
// a vector to keep track of the offsets into each input sequence,
// there an offset is the number of nnz values packed so far. Current sample
// values/indices start of the offset position in the sequence data/index array
vector<IndexType> sequenceOffsets(batch.size(), 0);
vector<MBLayout::SequenceInfo> sequenceInfos(pMBLayout->GetAllSequences());
// sort the vector in ascending order of the parallel sequence index.
sort(sequenceInfos.begin(), sequenceInfos.end(),
[](const MBLayout::SequenceInfo& a, const MBLayout::SequenceInfo& b){ return a.s < b.s; });
// Iterate over the all time steps in the layout (total number of samples/columns
// in a parallel sequence), traversing the layout in horizontal direction.
for (auto timeStep = 0; timeStep < pMBLayout->GetNumTimeSteps(); ++timeStep)
{
// For each time step, iterate over all sequences in the minibatch,
// traversing the layout in vertical direction.
for (const auto& sequenceInfo : sequenceInfos)
{
// skip the sequence if it does not intersect with the time step
if (timeStep < sequenceInfo.tBegin || timeStep >= sequenceInfo.tEnd)
{
continue;
}
// store the offset of the current column )...
sparseColumnIndices.push_back(columnOffset);
auto seqId = sequenceInfo.seqId;
if (seqId == GAP_SEQUENCE_ID)
{
continue;
}
// compute the index of the sample inside the sequence.
size_t sampleIndex = timeStep - sequenceInfo.tBegin;
const auto& sequence = batch[seqId];
// make sure the index less than the sequence length in samples.
assert(sampleIndex < sequence->m_numberOfSamples);
auto& sequenceOffset = sequenceOffsets[seqId];
SparseSequenceDataPtr sparseSequence = static_pointer_cast<SparseSequenceData>(sequence);
IndexType nnz = sparseSequence->m_nnzCounts[sampleIndex];
// compute the sample offset in bytes.
size_t sampleOffset = sequenceOffset * elementSize;
// copy all nzz values from source sequence into the buffer.
const auto* dataSrc = reinterpret_cast<const char*>(sequence->m_data) + sampleOffset;
memcpy(dataDst, dataSrc, nnz * elementSize);
dataDst += nnz * elementSize; // advance the destination pointer
// copy all nzz value indices from source sequence into the buffer.
const auto* indicesSrc = sparseSequence->m_indices + sequenceOffset;
memcpy(indicesDst, indicesSrc, nnz * indexSize);
indicesDst += nnz * indexSize; // advance the destination pointer
sequenceOffset += nnz;
columnOffset += nnz;
}
}
// at this point each element in sequenceOffsets should be equal to the total
// nnz count of the respective sequence and the sum of all elements - to the
// overall nnz count.
assert(accumulate(sequenceOffsets.begin(), sequenceOffsets.end(), 0) == nnzCount);
// check the distance between data and index destination pointers.
assert(indicesDst == dataDst + nnzCount * indexSize);
// after we packed all samples, the column offset must be equal to the total nnz count.
assert(columnOffset == nnzCount);
sparseColumnIndices.push_back(columnOffset);
// check that the number of column indices == N + 1 (where N is the number of
// column in the packed matrix)
assert((pMBLayout->GetNumCols() + 1) == sparseColumnIndices.size());
// verify that there's enough space in the buffer for the array of column indices.
assert(indicesDst + sparseColumnIndices.size()*indexSize <= destination + requiredSize);
// copy column indices into the buffer.
memcpy(indicesDst, sparseColumnIndices.data(), sparseColumnIndices.size() * indexSize);
return pMBLayout;
}
}}}