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                   &communicationManager,
           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(communicationManager, 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 updateDepth() {}
 
  __INLINE__ void advancePayloadTail(const size_t n = 1) { getCircularBufferForPayloads()->advanceTail(n); }
 
  [[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(); }
 
  __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.");
 
    // Make sure source slot is beg enough to satisfy the operation
    size_t requiredBufferSize     = sourceSlot->getSize();
    size_t providedBufferCapacity = getPayloadCapacity();
 
    // Updating depth of token (message sizes) and payload buffers
    updateDepth();
    auto currentCountsDepth  = getCircularBufferForCounts()->getDepth();
    auto currentPayloadDepth = getCircularBufferForPayloads()->getDepth();
    auto currentDepth        = getDepth();
 
    /*
     * Part 1: Copy the payload data
     */
    if (currentPayloadDepth + requiredBufferSize > providedBufferCapacity)
      HICR_THROW_RUNTIME("Attempting to push (%lu) bytes while the channel currently has payload depth (%lu). This would exceed capacity (%lu).\n",
                         requiredBufferSize,
                         currentPayloadDepth,
                         providedBufferCapacity);
 
    /*
     * Payload copy:
     *  - We have checked (requiredBufferSize  <= depth)
     *  that the payload fits into available circular buffer,
     *  but it is possible it spills over the end into the
     *  beginning. Cover this corner case below
     *
     */
    getCircularBufferForPayloads()->setCachedDepth(currentPayloadDepth);
    if (requiredBufferSize + getPayloadHeadPosition() > getPayloadCapacity())
    {
      size_t first_chunk  = getPayloadCapacity() - getPayloadHeadPosition();
      size_t second_chunk = requiredBufferSize - first_chunk;
 
      // copy first part to end of buffer
      getCommunicationManager()->memcpy(_payloadBuffer,           /* destination */
                                        getPayloadHeadPosition(), /* dst_offset */
                                        sourceSlot,               /* source */
                                        0,                        /* src_offset */
                                        first_chunk);             /* size */
      // copy second part to beginning of buffer
      getCommunicationManager()->memcpy(_payloadBuffer, /* destination */
                                        0,              /* dst_offset */
                                        sourceSlot,     /* source */
                                        first_chunk,    /* src_offset */
                                        second_chunk);  /* size */
      getCommunicationManager()->fence(sourceSlot, 2, 0);
    }
    else
    {
      getCommunicationManager()->memcpy(_payloadBuffer, getPayloadHeadPosition(), sourceSlot, 0, requiredBufferSize);
      getCommunicationManager()->fence(sourceSlot, 1, 0);
    }
 
    getCircularBufferForPayloads()->advanceHead(requiredBufferSize, true);
 
    // update the consumer coordination buffers (consumer does not update
    // its own coordination head positions)
    getCommunicationManager()->memcpy(_consumerCoordinationBufferForPayloads,
                                      _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                      getCoordinationBufferForPayloads(),
                                      _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                      sizeof(size_t));
    getCommunicationManager()->fence(getCoordinationBufferForPayloads(), 1, 0);
 
    /*
     * Part 2: Copy the message size
     */
 
    auto *sizeInfoBufferPtr = static_cast<size_t *>(_sizeInfoBuffer->getPointer());
    sizeInfoBufferPtr[0]    = requiredBufferSize;
 
    // If the exchange buffer does not have n free slots, reject the operation
    if (currentDepth + 1 > getCircularBufferForCounts()->getCapacity())
      HICR_THROW_RUNTIME("Attempting to push with (%lu) tokens while the channel has (%lu) tokens and this would exceed capacity (%lu).\n",
                         1,
                         getDepth(),
                         getCircularBufferForCounts()->getCapacity());
 
    getCircularBufferForCounts()->setCachedDepth(currentCountsDepth);
    /*
     * Advance head, as we have added new elements.
     * It is important to do the advanceHead and copy together,
     * or else issues such as advancing head index too early or too late might occur!
     */
 
    getCommunicationManager()->memcpy(_tokenBuffer,                                                     /* destination */
                                      getTokenSize() * getCircularBufferForCounts()->getHeadPosition(), /* dst_offset */
                                      _sizeInfoBuffer,                                                  /* source */
                                      0,                                                                /* src_offset */
                                      getTokenSize());                                                  /* size */
    getCommunicationManager()->fence(_sizeInfoBuffer, 1, 0);
    getCircularBufferForCounts()->advanceHead(1, true);
 
    getCommunicationManager()->memcpy(_consumerCoordinationBufferForCounts,
                                      _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                      getCoordinationBufferForCounts(),
                                      _HICR_CHANNEL_HEAD_ADVANCE_COUNT_IDX * sizeof(size_t),
                                      sizeof(size_t));
    getCommunicationManager()->fence(getCoordinationBufferForCounts(), 1, 0);
  }
 
  size_t getDepth() { return getCircularBufferForCounts()->getDepth(); }
 
  bool isEmpty() { return getDepth() == 0; }
 
  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