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 <array>
#include <numeric>
#include <cassert>
#include <hicr/frontends/channel/variableSize/base.hpp>
#include <utility>
 
namespace HiCR::channel::variableSize::SPSC
{
 
class Consumer final : public variableSize::Base
{
  public:
 
  Consumer(CommunicationManager                    &communicationManager,
           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> &producerCoordinationBufferForCounts,
           std::shared_ptr<GlobalMemorySlot>        producerCoordinationBufferForPayloads,
           const size_t                             payloadCapacity,
           const size_t                             capacity)
    : variableSize::Base(communicationManager, internalCoordinationBufferForCounts, internalCoordinationBufferForPayloads, capacity, payloadCapacity),
 
      _payloadBuffer(std::move(payloadBuffer)),
      _tokenSizeBuffer(std::move(tokenBuffer)),
      _producerCoordinationBufferForCounts(producerCoordinationBufferForCounts),
      _producerCoordinationBufferForPayloads(std::move(producerCoordinationBufferForPayloads))
  {
    assert(internalCoordinationBufferForCounts != nullptr);
    assert(internalCoordinationBufferForPayloads != nullptr);
    assert(producerCoordinationBufferForCounts != nullptr);
    assert(producerCoordinationBufferForCounts != nullptr);
  }
 
  __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());
 
    // Updating channel depth
    updateDepth();
 
    // 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)
  {
    if (pos != 0) { HICR_THROW_FATAL("peek only implemented for n = 0 at the moment!"); }
    updateDepth();
 
    if (pos >= getCircularBufferForCounts()->getDepth())
    {
      HICR_THROW_RUNTIME("Attempting to peek position (%lu) but not enough tokens (%lu) are in the buffer", pos, getCircularBufferForCounts()->getDepth());
    }
 
    std::array<size_t, 2> result{};
    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;
  }
 
  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;
  }
 
  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__ void pop(const size_t n = 1)
  {
    if (n > getCircularBufferForCounts()->getCapacity())
      HICR_THROW_LOGIC("Attempting to pop (%lu) tokens, which is larger than the channel capacity (%lu)", n, getCircularBufferForCounts()->getCapacity());
 
    // Updating channel depth
    updateDepth();
 
    // 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());
    auto payloadBytes = getOldPayloadBytes(n);
    getCircularBufferForCounts()->advanceTail(n);
    getCircularBufferForPayloads()->advanceTail(payloadBytes);
 
    const auto coordBuffElemSize = sizeof(_HICR_CHANNEL_COORDINATION_BUFFER_ELEMENT_TYPE);
    // Notifying producer(s) of buffer liberation
    getCommunicationManager()->memcpy(_producerCoordinationBufferForCounts, /* destination */
                                      _HICR_CHANNEL_TAIL_ADVANCE_COUNT_IDX * coordBuffElemSize,
                                      getCoordinationBufferForCounts(),
                                      _HICR_CHANNEL_TAIL_ADVANCE_COUNT_IDX * coordBuffElemSize,
                                      coordBuffElemSize);
 
    getCommunicationManager()->memcpy(_producerCoordinationBufferForPayloads, /* destination */
                                      _HICR_CHANNEL_TAIL_ADVANCE_COUNT_IDX * coordBuffElemSize,
                                      getCoordinationBufferForPayloads(), /* source */
                                      _HICR_CHANNEL_TAIL_ADVANCE_COUNT_IDX * coordBuffElemSize,
                                      coordBuffElemSize);
 
    getCommunicationManager()->fence(getCoordinationBufferForCounts(), 1, 0);
    getCommunicationManager()->fence(getCoordinationBufferForPayloads(), 1, 0);
  }
 
  __INLINE__ void updateDepth() {}
 
  size_t getDepth() { return getCircularBufferForCounts()->getDepth(); }
 
  size_t getPayloadDepth() { return getCircularBufferForPayloads()->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> _producerCoordinationBufferForCounts;
 
  const std::shared_ptr<GlobalMemorySlot> _producerCoordinationBufferForPayloads;
};
 
} // namespace HiCR::channel::variableSize::SPSC