consumer.hpp

Go to the documentation of this file.

/*
 *   Copyright 2025 Huawei Technologies Co., Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <queue>
#include <hicr/core/definitions.hpp>
#include <hicr/core/exceptions.hpp>
#include <hicr/frontends/channel/fixedSize/base.hpp>
#include <hicr/frontends/channel/fixedSize/spsc/consumer.hpp>
 
namespace HiCR::channel::fixedSize::MPSC::nonlocking
{
 
class Consumer
{
  public:
 
  Consumer(CommunicationManager                          &communicationManager,
           std::vector<std::shared_ptr<GlobalMemorySlot>> tokenBuffers,
           std::vector<std::shared_ptr<LocalMemorySlot>>  internalCoordinationBuffers,
           std::vector<std::shared_ptr<GlobalMemorySlot>> producerCoordinationBuffers,
           const size_t                                   tokenSize,
           const size_t                                   capacity)
    : _tokenBuffers(tokenBuffers),
      _communicationManager(&communicationManager)
  {
    // make sure producer and consumer sides provide p elements, equalling
    // the number of producers
    assert(internalCoordinationBuffers.size() == producerCoordinationBuffers.size());
    assert(internalCoordinationBuffers.size() == tokenBuffers.size());
    // create p (= number of producers) SPSC channels
    for (size_t i = 0; i < internalCoordinationBuffers.size(); i++)
    {
      std::shared_ptr<fixedSize::SPSC::Consumer> consumerPtr(
        new fixedSize::SPSC::Consumer(communicationManager, tokenBuffers[i], internalCoordinationBuffers[i], producerCoordinationBuffers[i], tokenSize, capacity));
      _spscList.push_back(consumerPtr);
      _depths.push_back(0);
    }
  }
 
  ~Consumer() = default;
 
  __INLINE__ std::array<size_t, 2> peek(const size_t pos = 0)
  {
    std::array<size_t, 2> ret = {0};
    // @ToDo: to support pos > 0, we need to modify _channelPushes to
    // be of type std::vector instead of std::queue
    if (pos > 0) HICR_THROW_LOGIC("Nonblocking MPSC not yet implemented for peek with n!=0");
 
    _communicationManager->flushReceived();
    updateDepth();
    if (_channelPushes.empty()) HICR_THROW_RUNTIME("Attempting to peek position (%lu) but supporting queue has size (%lu)", pos, _channelPushes.size());
 
    size_t channelId = _channelPushes.front(); // front() returns the first (i.e. oldest) element
    if (channelId >= _spscList.size()) { HICR_THROW_LOGIC("channelId (%lu) >= _spscList.size() (%lu)", channelId, _spscList.size()); }
    ret[0] = channelId;
    ret[1] = _spscList[channelId]->peek();
 
    return ret;
  }
 
  __INLINE__ size_t getDepth()
  {
    size_t totalDepth = 0;
    for (auto d : _depths) { totalDepth += d; }
    //    if (totalDepth != _channelPushes.size()) {
    //      HICR_THROW_LOGIC("Helper FIFO and channels are out of sync, implemenation issue! getDepth (%lu) != _channelPushes.size() (%lu)", totalDepth, _channelPushes.size());
    //    }
    return totalDepth;
  }
 
  __INLINE__ bool isEmpty() { return (getDepth() == 0); }
 
  __INLINE__ void pop(const size_t n = 1)
  {
    updateDepth();
    // If the exchange buffer does not have enough tokens, reject operation
    if (n > getDepth())
      HICR_THROW_RUNTIME("Attempting to pop (%lu) tokens, which is more than the number of current tokens in the channel (%lu)", n, getDepth());
    else
    {
      size_t channelFirstPushed = _channelPushes.front();
      if (channelFirstPushed >= _spscList.size())
        HICR_THROW_LOGIC("Index of latest push channel incorrect!");
      else
      {
        // pop n elements from the SPSCs in the order recorded in the helper
        // FIFO _channelPushes, and also update the FIFO itself
        for (size_t i = 0; i < n; i++)
        {
          _spscList[channelFirstPushed]->pop();
          _depths[channelFirstPushed]--;
          _channelPushes.pop();
        }
      }
    }
  }
 
  __INLINE__ void updateDepth()
  {
    std::vector<size_t> newDepths(_spscList.size());
    // Note that after calling updateDepth() on each SPSC channel,
    // we must accept this state as a new temporary snapshot in newDepths
    // It is possible that during our iterating through newDepths, producers have
    // sent more elements already, which will be handled in later
    // updateDepth calls
    for (size_t i = 0; i < _spscList.size(); i++)
    {
      _spscList[i]->updateDepth();
      newDepths[i] = _spscList[i]->getDepth();
    }
 
    for (size_t i = 0; i < _spscList.size(); i++)
    {
      for (size_t j = _depths[i]; j < newDepths[i]; j++) { _channelPushes.push(i); }
    }
    std::swap(_depths, newDepths);
    if (getDepth() != _channelPushes.size()) { HICR_THROW_LOGIC("getDepth (%lu) != _channelPushes.size() (%lu)", getDepth(), _channelPushes.size()); }
  }
 
  [[nodiscard]] __INLINE__ std::vector<std::shared_ptr<GlobalMemorySlot>> getTokenBuffers() const { return _tokenBuffers; }
 
  private:
 
  const std::vector<std::shared_ptr<HiCR::GlobalMemorySlot>> _tokenBuffers;
 
  std::vector<std::shared_ptr<channel::fixedSize::SPSC::Consumer>> _spscList;
  std::queue<size_t> _channelPushes;
 
  std::vector<size_t> _depths;
  CommunicationManager *const _communicationManager;
};
 
} // namespace HiCR::channel::fixedSize::MPSC::nonlocking