vtkBatch.h

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#ifndef vtkBatch_h
#define vtkBatch_h

#include "vtkCommonCoreModule.h" // For export macro
#include "vtkSMPTools.h"
#include "vtkType.h"

#include <algorithm>
#include <functional>
#include <vector>

VTK_ABI_NAMESPACE_BEGIN
//-----------------------------------------------------------------------------
template <typename TBatchData>
struct vtkBatch
{
  vtkIdType BeginId;
  vtkIdType EndId;
  TBatchData Data;
};

//-----------------------------------------------------------------------------
template <typename TBatchData>
class vtkBatches
{
private:
  using vtkTBatch = vtkBatch<TBatchData>;
  using vector = std::vector<vtkTBatch>;
  using size_type = typename vector::size_type;
  using reference = typename vector::reference;
  using const_reference = typename vector::const_reference;
  using iterator = typename vector::iterator;
  using const_iterator = typename vector::const_iterator;

public:
  vtkBatches()
    : BatchSize(0)
  {
  }
  ~vtkBatches() = default;

  void Initialize(vtkIdType numberOfElements, unsigned int batchSize = 1000)
  {
    this->BatchSize = batchSize;

    const auto batchSizeSigned = static_cast<vtkIdType>(batchSize);
    const auto numberOfBatches = ((numberOfElements - 1) / batchSizeSigned) + 1;
    this->Batches.resize(static_cast<size_t>(numberOfBatches));
    const auto lastBatchId = numberOfBatches - 1;

    vtkSMPTools::For(0, numberOfBatches,
      [&](vtkIdType beginBatchId, vtkIdType endBatchId)
      {
        vtkIdType endIdValues[2] = { -1, numberOfElements };
        for (vtkIdType batchId = beginBatchId; batchId < endBatchId; ++batchId)
        {
          auto& batch = this->Batches[batchId];
          batch.BeginId = batchId * batchSizeSigned;
          endIdValues[0] = (batchId + 1) * batchSizeSigned;
          batch.EndId = endIdValues[batchId == lastBatchId];
        }
      });
  }

  void TrimBatches(const std::function<bool(const vtkTBatch&)> shouldRemoveBatch)
  {
    const vtkIdType numberOfBatches = this->GetNumberOfBatches();
    if (numberOfBatches == 0)
    {
      return;
    }
    const vtkIdType numberOfThreads =
      std::min(static_cast<vtkIdType>(vtkSMPTools::GetEstimatedNumberOfThreads()), numberOfBatches);
    const vtkIdType lastThreadId = numberOfThreads - 1;
    const vtkIdType numberOfBatchesPerThread = numberOfBatches / numberOfThreads;

    std::vector<vtkIdType> tlInterEndBatchId;
    tlInterEndBatchId.resize(static_cast<size_t>(numberOfThreads));

    // perform batch trimming of batch segments
    vtkSMPTools::For(0, numberOfThreads,
      [&](vtkIdType beginThreadId, vtkIdType endThreadId)
      {
        for (vtkIdType threadId = beginThreadId; threadId < endThreadId; ++threadId)
        {
          const vtkIdType beginBatchId = threadId * numberOfBatchesPerThread;
          const vtkIdType endBatchId =
            threadId != lastThreadId ? (threadId + 1) * numberOfBatchesPerThread : numberOfBatches;
          auto beginBatchIter = this->Batches.begin() + beginBatchId;
          auto endBatchIter = this->Batches.begin() + endBatchId;
          auto newEndBatchIter = std::remove_if(beginBatchIter, endBatchIter, shouldRemoveBatch);
          tlInterEndBatchId[threadId] =
            beginBatchId + std::distance(beginBatchIter, newEndBatchIter);
        }
      });
    // compute the new end batch ids;
    std::vector<vtkIdType> tlNewEndBatchId;
    tlNewEndBatchId.resize(static_cast<size_t>(numberOfThreads));
    tlNewEndBatchId[0] = tlInterEndBatchId[0];
    for (vtkIdType threadId = 1; threadId < numberOfThreads; ++threadId)
    {
      const vtkIdType beginOldBatchId = threadId * numberOfBatchesPerThread;
      const auto& prevNewEndBatchId = tlNewEndBatchId[threadId - 1];
      const vtkIdType distance = tlInterEndBatchId[threadId] - beginOldBatchId;
      tlNewEndBatchId[threadId] = prevNewEndBatchId + distance;
      // Only move if the destination is different from the source
      if (prevNewEndBatchId != beginOldBatchId)
      {
        std::move(this->Batches.begin() + beginOldBatchId,
          this->Batches.begin() + tlInterEndBatchId[threadId],
          this->Batches.begin() + prevNewEndBatchId);
      }
    }
    this->Batches.erase(this->Batches.begin() + tlNewEndBatchId[lastThreadId], this->Batches.end());
  }

  TBatchData BuildOffsetsAndGetGlobalSum()
  {
    const vtkIdType numberOfBatches = this->GetNumberOfBatches();
    if (numberOfBatches == 0)
    {
      return TBatchData{};
    }
    const vtkIdType numberOfThreads =
      std::min(static_cast<vtkIdType>(vtkSMPTools::GetEstimatedNumberOfThreads()), numberOfBatches);
    const vtkIdType lastThreadId = numberOfThreads - 1;
    const vtkIdType numberOfBatchesPerThread = numberOfBatches / numberOfThreads;

    // calculate the thread local sums;
    std::vector<TBatchData> tlSums;
    tlSums.resize(static_cast<size_t>(numberOfThreads));

    vtkSMPTools::For(0, numberOfThreads,
      [&](vtkIdType beginThreadId, vtkIdType endThreadId)
      {
        for (vtkIdType threadId = beginThreadId; threadId < endThreadId; ++threadId)
        {
          const vtkIdType beginBatchId = threadId * numberOfBatchesPerThread;
          const vtkIdType endBatchId =
            threadId != lastThreadId ? (threadId + 1) * numberOfBatchesPerThread : numberOfBatches;

          auto& threadSum = tlSums[threadId];
          for (vtkIdType batchId = beginBatchId; batchId < endBatchId; ++batchId)
          {
            threadSum += this->Batches[batchId].Data;
          }
        }
      });

    // calculate the global sum;
    TBatchData globalSum;
    for (const auto& threadSum : tlSums)
    {
      globalSum += threadSum;
    }

    // calculate the thread local offsets using the thread local sums
    std::vector<TBatchData> tlOffsets;
    tlOffsets.resize(static_cast<size_t>(numberOfThreads));
    for (vtkIdType threadId = 1; threadId < numberOfThreads; ++threadId)
    {
      tlOffsets[threadId] = tlOffsets[threadId - 1] + tlSums[threadId - 1];
    }

    // convert the batch sums to offsets using the thread local offsets
    vtkSMPTools::For(0, numberOfThreads,
      [&](vtkIdType beginThreadId, vtkIdType endThreadId)
      {
        for (vtkIdType threadId = beginThreadId; threadId < endThreadId; ++threadId)
        {
          const vtkIdType beginBatchId = threadId * numberOfBatchesPerThread;
          const vtkIdType endBatchId =
            threadId != lastThreadId ? (threadId + 1) * numberOfBatchesPerThread : numberOfBatches;

          // store the first batch sum
          auto lastBatchData = this->Batches[beginBatchId].Data;
          // convert the first batch sum to an offset
          this->Batches[beginBatchId].Data = tlOffsets[threadId];
          for (vtkIdType batchId = beginBatchId + 1; batchId < endBatchId; ++batchId)
          {
            // store the current batch sum
            const auto currentBatchData = this->Batches[batchId].Data;
            // // Compute the offsets for the current batch
            const auto newData = this->Batches[batchId - 1].Data + lastBatchData;
            // // Convert the current batch sum to the computed offsets
            this->Batches[batchId].Data = newData;
            // store the current batch sum for the next iteration
            lastBatchData = std::move(currentBatchData);
          }
        }
      });

    return globalSum;
  }

  vtkIdType GetNumberOfBatches() const { return static_cast<vtkIdType>(this->Batches.size()); }

  unsigned int GetBatchSize() const { return this->BatchSize; }


  reference operator[](size_type pos) noexcept { return this->Batches[pos]; }
  const_reference operator[](size_type pos) const noexcept { return this->Batches[pos]; }
  iterator begin() noexcept { return this->Batches.begin(); }
  const_iterator begin() const noexcept { return this->Batches.begin(); }
  const_iterator cbegin() const noexcept { return this->Batches.cbegin(); }
  iterator end() noexcept { return this->Batches.end(); }
  const_iterator end() const noexcept { return this->Batches.end(); }
  const_iterator cend() const noexcept { return this->Batches.cend(); }

private:
  vector Batches;
  unsigned int BatchSize;
};
VTK_ABI_NAMESPACE_END

#endif // vtkBatch_h
// VTK-HeaderTest-Exclude: vtkBatch.h