producer.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::SPSC
{
 
class Producer : public variableSize::Base
{
  public:
 
  Producer(CommunicationManager                   &coordinationCommunicationManager,
           CommunicationManager                   &payloadCommunicationManager,
           std::shared_ptr<LocalMemorySlot>        sizeInfoBuffer,
           std::shared_ptr<GlobalMemorySlot>       payloadBuffer,
           std::shared_ptr<GlobalMemorySlot>       tokenBuffer,
           const std::shared_ptr<LocalMemorySlot> &internalCoordinationBufferForCounts,
           const std::shared_ptr<LocalMemorySlot> &internalCoordinationBufferForPayloads,
           std::shared_ptr<GlobalMemorySlot>       consumerCoordinationBufferForCounts,
           std::shared_ptr<GlobalMemorySlot>       consumerCoordinationBufferForPayloads,
           const size_t                            payloadCapacity,
           const size_t                            payloadSize,
           const size_t                            capacity)
    : variableSize::Base(coordinationCommunicationManager,
                         payloadCommunicationManager,
                         internalCoordinationBufferForCounts,
                         internalCoordinationBufferForPayloads,
                         capacity,
                         payloadCapacity),
      _payloadBuffer(std::move(payloadBuffer)),
      _sizeInfoBuffer(std::move(sizeInfoBuffer)),
      _payloadSize(payloadSize),
      _tokenBuffer(std::move(tokenBuffer)),
      _consumerCoordinationBufferForCounts(std::move(consumerCoordinationBufferForCounts)),
      _consumerCoordinationBufferForPayloads(std::move(consumerCoordinationBufferForPayloads))
  {}
 
  ~Producer() = default;
 
  __INLINE__ void push(const std::shared_ptr<LocalMemorySlot> &sourceSlot, const size_t n = 1)
  {
    if (n != 1) HICR_THROW_RUNTIME("HiCR currently has no implementation for n != 1 with push(sourceSlot, n) for variable size version.");
 
    // Updating depth of token (message sizes) and payload buffers
    updateDepth();
 
    // Get bytes required to push the token
    size_t requiredPayloadBufferSize = sourceSlot->getSize();
 
    // Throw exception if the token can not be pushed
    if (isFull(requiredPayloadBufferSize) == true)
    {
      HICR_THROW_RUNTIME("Attempting to push a token while the channel is full.\nChannel depth: %lu capacity: %lu\nPayload depth: %lu capacity: %lu",
                         getCircularBufferForCounts()->getDepth(),
                         getCircularBufferForCounts()->getCapacity(),
                         getCircularBufferForPayloads()->getDepth(),
                         getCircularBufferForPayloads()->getCapacity());
    }
 
    // Get communication managers
    auto payloadCommunicationManager      = getPayloadCommunicationManager();
    auto coordinationCommunicationManager = getCoordinationCommunicationManager();
 
    //  Payload copy. Just push to the channel, we know there is enough space
    payloadCommunicationManager->memcpy(_payloadBuffer, getPayloadHeadPosition(), sourceSlot, 0, requiredPayloadBufferSize);
    payloadCommunicationManager->fence(sourceSlot, 1, 0);
 
    // Advance the head in the payload buffer
    getCircularBufferForPayloads()->advanceHead(requiredPayloadBufferSize);
 
    // update the consumer coordination buffers (consumer does not update its own coordination head positions)
    coordinationCommunicationManager->memcpy(_consumerCoordinationBufferForPayloads,
                                             _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                             getCoordinationBufferForPayloads(),
                                             _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                             sizeof(size_t));
    coordinationCommunicationManager->fence(getCoordinationBufferForPayloads(), 1, 0);
 
    /*
     * Part 2: Copy the message size
     */
    auto *sizeInfoBufferPtr = static_cast<size_t *>(_sizeInfoBuffer->getPointer());
    sizeInfoBufferPtr[0]    = requiredPayloadBufferSize;
 
    coordinationCommunicationManager->memcpy(_tokenBuffer,                                                     /* destination */
                                             getTokenSize() * getCircularBufferForCounts()->getHeadPosition(), /* dst_offset */
                                             _sizeInfoBuffer,                                                  /* source */
                                             0,                                                                /* src_offset */
                                             getTokenSize());                                                  /* size */
    coordinationCommunicationManager->fence(_sizeInfoBuffer, 1, 0);
    getCircularBufferForCounts()->advanceHead(1);
 
    coordinationCommunicationManager->memcpy(_consumerCoordinationBufferForCounts,
                                             _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                             getCoordinationBufferForCounts(),
                                             _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                             sizeof(size_t));
    coordinationCommunicationManager->fence(getCoordinationBufferForCounts(), 1, 0);
  }
 
  __INLINE__ void updateDepth() {}
 
  [[nodiscard]] __INLINE__ size_t getPayloadHeadPosition() const noexcept { return getCircularBufferForPayloads()->getHeadPosition(); }
 
  __INLINE__ size_t getPayloadSize() { return _payloadSize; }
 
  __INLINE__ size_t getPayloadDepth() { return getCircularBufferForPayloads()->getDepth(); }
 
  __INLINE__ size_t getPayloadCapacity() { return getCircularBufferForPayloads()->getCapacity(); }
 
  size_t getCoordinationDepth() { return getCircularBufferForCounts()->getDepth(); }
 
  bool isEmpty() { return getCoordinationDepth() == 0; }
 
  bool isFull(const size_t requiredBufferSize)
  {
    // Check if we can push one more token
    auto coordinationCircularBuffer = getCircularBufferForCounts();
    if (coordinationCircularBuffer->getDepth() == coordinationCircularBuffer->getCapacity()) { return true; }
 
    // Check if there is enough space in the payload buffer. If
    auto payloadCircularBuffer = getCircularBufferForPayloads();
    if (payloadCircularBuffer->getDepth() + requiredBufferSize > payloadCircularBuffer->getCapacity()) { return true; }
 
    return false;
  }
 
  private:
 
  std::shared_ptr<GlobalMemorySlot> _payloadBuffer;
  std::shared_ptr<LocalMemorySlot> _sizeInfoBuffer;
  size_t _payloadSize;
 
  const std::shared_ptr<GlobalMemorySlot> _tokenBuffer;
 
  const std::shared_ptr<GlobalMemorySlot> _consumerCoordinationBufferForCounts;
 
  const std::shared_ptr<GlobalMemorySlot> _consumerCoordinationBufferForPayloads;
};
 
} // namespace HiCR::channel::variableSize::SPSC