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 <hicr/frontends/channel/variableSize/base.hpp>
#include <utility>
 
namespace HiCR::channel::variableSize::MPSC::locking
{
 
class Consumer final : public variableSize::Base
{
  public:
 
  Consumer(CommunicationManager                    &coordinationCommunicationManager,
           CommunicationManager                    &payloadCommunicationManager,
           std::shared_ptr<GlobalMemorySlot>        payloadBuffer,
           std::shared_ptr<GlobalMemorySlot>        tokenBuffer,
           const std::shared_ptr<LocalMemorySlot>  &internalCoordinationBufferForCounts,
           const std::shared_ptr<LocalMemorySlot>  &internalCoordinationBufferForPayloads,
           const std::shared_ptr<GlobalMemorySlot> &consumerCoordinationBufferForCounts,
           std::shared_ptr<GlobalMemorySlot>        consumerCoordinationBufferForPayloads,
           const size_t                             payloadCapacity,
           const size_t                             capacity)
    : variableSize::Base(coordinationCommunicationManager,
                         payloadCommunicationManager,
                         internalCoordinationBufferForCounts,
                         internalCoordinationBufferForPayloads,
                         capacity,
                         payloadCapacity),
      _payloadBuffer(std::move(payloadBuffer)),
      _tokenSizeBuffer(std::move(tokenBuffer)),
      _consumerCoordinationBufferForCounts(consumerCoordinationBufferForCounts),
      _consumerCoordinationBufferForPayloads(std::move(consumerCoordinationBufferForPayloads))
  {
    assert(internalCoordinationBufferForCounts != nullptr);
    assert(internalCoordinationBufferForPayloads != nullptr);
    assert(consumerCoordinationBufferForCounts != nullptr);
    getCoordinationCommunicationManager()->queryMemorySlotUpdates(_tokenSizeBuffer->getSourceLocalMemorySlot());
    getPayloadCommunicationManager()->queryMemorySlotUpdates(_payloadBuffer->getSourceLocalMemorySlot());
  }
 
  __INLINE__ size_t basePeek(const size_t pos = 0)
  {
    // Check if the requested position exceeds the capacity of the channel
    if (pos >= getCircularBufferForCounts()->getCapacity())
      HICR_THROW_LOGIC("Attempting to peek for a token with position (%lu), which is beyond than the channel capacity (%lu)", pos, getCircularBufferForCounts()->getCapacity());
 
    // Check if there are enough tokens in the buffer to satisfy the request
    if (pos >= getCircularBufferForCounts()->getDepth())
      HICR_THROW_RUNTIME("Attempting to peek position (%lu) but not enough tokens (%lu) are in the buffer", pos, getCircularBufferForCounts()->getDepth());
 
    // Calculating buffer position
    const size_t bufferPos = (getCircularBufferForCounts()->getTailPosition() + pos) % getCircularBufferForCounts()->getCapacity();
 
    // Succeeded in pushing the token(s)
    return bufferPos;
  }
 
  __INLINE__ std::array<size_t, 2> peek(const size_t pos = 0)
  {
    getCoordinationCommunicationManager()->flushReceived();
    getPayloadCommunicationManager()->flushReceived();
    std::array<size_t, 2> result{};
    if (pos != 0) { HICR_THROW_FATAL("peek only implemented for n = 0 at the moment!"); }
    if (pos >= getCircularBufferForCounts()->getDepth())
    {
      HICR_THROW_RUNTIME("Attempting to peek position (%lu) but not enough tokens (%lu) are in the buffer", pos, getCircularBufferForCounts()->getDepth());
    }
 
    result[0]              = getCircularBufferForPayloads()->getTailPosition() % getCircularBufferForPayloads()->getCapacity();
    size_t *tokenBufferPtr = static_cast<size_t *>(_tokenSizeBuffer->getSourceLocalMemorySlot()->getPointer());
    auto    tokenPos       = basePeek(pos);
    result[1]              = tokenBufferPtr[tokenPos];
    return result;
  }
 
  __INLINE__ size_t getOldPayloadBytes(size_t n)
  {
    if (n == 0) return 0;
    size_t *tokenBufferPtr = static_cast<size_t *>(_tokenSizeBuffer->getSourceLocalMemorySlot()->getPointer());
 
    size_t payloadBytes = 0;
    for (size_t i = 0; i < n; i++)
    {
      assert(i >= 0);
      size_t pos         = basePeek(i);
      auto   payloadSize = tokenBufferPtr[pos];
      payloadBytes += payloadSize;
    }
    return payloadBytes;
  }
 
  __INLINE__ size_t getNewPayloadBytes(size_t n)
  {
    if (n == 0) return 0;
    size_t *tokenBufferPtr = static_cast<size_t *>(_tokenSizeBuffer->getSourceLocalMemorySlot()->getPointer());
    size_t  payloadBytes   = 0;
 
    for (size_t i = 0; i < n; i++)
    {
      size_t ind = getCircularBufferForCounts()->getDepth() - 1 - i;
      assert(ind >= 0);
      size_t pos         = basePeek(ind);
      auto   payloadSize = tokenBufferPtr[pos];
      payloadBytes += payloadSize;
    }
 
    return payloadBytes;
  }
 
  __INLINE__ bool pop(const size_t n = 1)
  {
    bool successFlag = false;
 
    auto coordinationCommunicationManager = getCoordinationCommunicationManager();
 
    // Locking remote coordination buffer slot
    if (coordinationCommunicationManager->acquireGlobalLock(_consumerCoordinationBufferForCounts) == false) return successFlag;
 
    if (n > getCircularBufferForCounts()->getCapacity())
      HICR_THROW_LOGIC("Attempting to pop (%lu) tokens, which is larger than the channel capacity (%lu)", n, getCircularBufferForCounts()->getCapacity());
    // If the exchange buffer does not have n tokens pushed, reject operation
    if (n > getCircularBufferForCounts()->getDepth())
      HICR_THROW_RUNTIME("Attempting to pop (%lu) tokens, which is more than the number of current tokens in the channel (%lu)", n, getCircularBufferForCounts()->getDepth());
 
    size_t *tokenBufferPtr = static_cast<size_t *>(_tokenSizeBuffer->getSourceLocalMemorySlot()->getPointer());
 
    size_t bytesOldestEntry = tokenBufferPtr[getCircularBufferForCounts()->getTailPosition()];
 
    getCircularBufferForCounts()->advanceTail(n);
    getCircularBufferForPayloads()->advanceTail(bytesOldestEntry);
 
    coordinationCommunicationManager->releaseGlobalLock(_consumerCoordinationBufferForCounts);
    successFlag = true;
    return successFlag;
  }
 
  size_t getDepth() { return getCircularBufferForCounts()->getDepth(); }
 
  bool isEmpty() { return (getDepth() == 0); }
 
  [[nodiscard]] std::shared_ptr<GlobalMemorySlot> getPayloadBufferMemorySlot() const { return _payloadBuffer; }
 
  private:
 
  std::shared_ptr<GlobalMemorySlot> _payloadBuffer;
 
  const std::shared_ptr<GlobalMemorySlot> _tokenSizeBuffer;
 
  const std::shared_ptr<GlobalMemorySlot> _consumerCoordinationBufferForCounts;
 
  const std::shared_ptr<GlobalMemorySlot> _consumerCoordinationBufferForPayloads;
};
 
} // namespace HiCR::channel::variableSize::MPSC::locking