consumer.hpp

/*
 *   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/frontends/channel/variableSize/spsc/consumer.hpp>
 
namespace HiCR::channel::variableSize::MPSC::nonlocking
{
 
class Consumer
{
  public:
 
  Consumer(CommunicationManager                                 &communicationManager,
           const std::vector<std::shared_ptr<GlobalMemorySlot>> &payloadBuffers,
           const std::vector<std::shared_ptr<GlobalMemorySlot>> &tokenBuffers,
           const std::vector<std::shared_ptr<LocalMemorySlot>>  &internalCoordinationBufferForCounts,
           const std::vector<std::shared_ptr<LocalMemorySlot>>  &internalCoordinationBufferForPayloads,
           const std::vector<std::shared_ptr<GlobalMemorySlot>> &producerCoordinationBufferForCounts,
           const std::vector<std::shared_ptr<GlobalMemorySlot>> &producerCoordinationBufferForPayloads,
           const size_t                                          payloadCapacity,
           const size_t                                          payloadSize,
           const size_t                                          capacity)
    : _communicationManager(&communicationManager)
  {
    // make sure producer and consumer sides have the same element size
    // the size is hopefully the producer count
    assert(!internalCoordinationBufferForCounts.empty());
    auto producerCount = internalCoordinationBufferForCounts.size();
    assert(producerCount == internalCoordinationBufferForPayloads.size());
    assert(producerCount == producerCoordinationBufferForCounts.size());
    assert(producerCount == producerCoordinationBufferForPayloads.size());
 
    // create p (= number of producers) SPSC channels
    for (size_t i = 0; i < producerCount; i++)
    {
      std::shared_ptr<variableSize::SPSC::Consumer> consumerPtr(new variableSize::SPSC::Consumer(communicationManager,
                                                                                                 payloadBuffers[i],
                                                                                                 tokenBuffers[i],
                                                                                                 internalCoordinationBufferForCounts[i],
                                                                                                 internalCoordinationBufferForPayloads[i],
                                                                                                 producerCoordinationBufferForCounts[i],
                                                                                                 producerCoordinationBufferForPayloads[i],
                                                                                                 payloadCapacity,
                                                                                                 capacity));
      _spscList.push_back(consumerPtr);
 
      /*
      * Note that it is important to record messages that might already have been received
      * immediately upon creation of the SPSC channel. Therefore we do not reset
      * _depths to zero, and check for "early" received messages
      */
      _depths.push_back(consumerPtr->getDepth());
      for (size_t j = 0; j < _depths.back(); j++) { _channelPushes.push(i); }
    }
  }
 
  ~Consumer() = default;
 
  __INLINE__ std::array<size_t, 3> peek(const size_t pos = 0)
  {
    std::array<size_t, 3> 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()[0];
    ret[2] = _spscList[channelId]->peek()[1];
 
    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++)
    {
      assert(_depths[i] <= newDepths[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()); }
  }
 
  private:
 
  std::vector<std::shared_ptr<channel::variableSize::SPSC::Consumer>> _spscList;
  std::queue<size_t> _channelPushes;
 
  std::vector<size_t> _depths;
  CommunicationManager *const _communicationManager;
};
 
} // namespace HiCR::channel::variableSize::MPSC::nonlocking