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::MPSC::locking
{
 
class Producer final : 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),
      _tokenSizeBuffer(std::move(tokenBuffer)),
      _consumerCoordinationBufferForCounts(std::move(consumerCoordinationBufferForCounts)),
      _consumerCoordinationBufferForPayloads(std::move(consumerCoordinationBufferForPayloads))
  {}
 
  ~Producer() = default;
 
  __INLINE__ void updateDepth() // NOTE: we DO know we have the lock!!!!
  {
    auto coordinationCommunicationManager = getCoordinationCommunicationManager();
    coordinationCommunicationManager->memcpy(getCoordinationBufferForCounts(),                            /* destination */
                                             0,                                                           /* dst_offset */
                                             _consumerCoordinationBufferForCounts,                        /* source */
                                             0,                                                           /* src_offset */
                                             2 * sizeof(_HICR_CHANNEL_COORDINATION_BUFFER_ELEMENT_TYPE)); /* size */
 
    coordinationCommunicationManager->memcpy(getCoordinationBufferForPayloads(),                          /* destination */
                                             0,                                                           /* dst_offset */
                                             _consumerCoordinationBufferForPayloads,                      /* source */
                                             0,                                                           /* src_offset */
                                             2 * sizeof(_HICR_CHANNEL_COORDINATION_BUFFER_ELEMENT_TYPE)); /* size */
 
    coordinationCommunicationManager->fence(getCoordinationBufferForCounts(), 0, 1);
    coordinationCommunicationManager->fence(getCoordinationBufferForPayloads(), 0, 1);
  }
 
  inline size_t getPayloadSize() { return _payloadSize; }
 
  inline size_t getPayloadDepth() { return getCircularBufferForPayloads()->getDepth(); }
 
  __INLINE__ bool 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 big enough to satisfy the operation
    size_t requiredBufferSize = sourceSlot->getSize();
 
    // Flag to record whether the operation was successful or not (it simplifies code by releasing locks only once)
    bool successFlag = false;
 
    auto coordinationCommunicationManager = getCoordinationCommunicationManager();
 
    // Locking remote token and coordination buffer slots
    if (coordinationCommunicationManager->acquireGlobalLock(_consumerCoordinationBufferForCounts) == false) return successFlag;
 
    // Updating depth of token (message sizes) and payload buffers from the consumer
    updateDepth();
 
    // Check if the amount of data we wish to write (requiredBufferSize)
    // would fit in the consumer payload buffer in its current state.
    // If not, reject the operation
    if (getCircularBufferForPayloads()->getDepth() + requiredBufferSize > getCircularBufferForPayloads()->getCapacity())
    {
      coordinationCommunicationManager->releaseGlobalLock(_consumerCoordinationBufferForCounts);
      return successFlag;
    }
 
    auto *sizeInfoBufferPtr = static_cast<size_t *>(_sizeInfoBuffer->getPointer());
    sizeInfoBufferPtr[0]    = requiredBufferSize;
 
    // Check if the consumer buffer has n free slots. If not, reject the operation
    if (getCircularBufferForCounts()->getDepth() + 1 > getCircularBufferForCounts()->getCapacity())
    {
      coordinationCommunicationManager->releaseGlobalLock(_consumerCoordinationBufferForCounts);
      return successFlag;
    }
 
    coordinationCommunicationManager->memcpy(_tokenSizeBuffer,                                                 /* destination */
                                             getTokenSize() * getCircularBufferForCounts()->getHeadPosition(), /* dst_offset */
                                             _sizeInfoBuffer,                                                  /* source */
                                             0,                                                                /* src_offset */
                                             getTokenSize());                                                  /* size */
    coordinationCommunicationManager->fence(_sizeInfoBuffer, 1, 0);
    successFlag = true;
 
    auto payloadCommunicationManager = getPayloadCommunicationManager();
 
    if (requiredBufferSize + getCircularBufferForPayloads()->getHeadPosition() > getCircularBufferForPayloads()->getCapacity())
    {
      size_t first_chunk  = getCircularBufferForPayloads()->getCapacity() - getCircularBufferForPayloads()->getHeadPosition();
      size_t second_chunk = requiredBufferSize - first_chunk;
      // copy first part to end of buffer
      payloadCommunicationManager->memcpy(_payloadBuffer,                                    /* destination */
                                          getCircularBufferForPayloads()->getHeadPosition(), /* dst_offset */
                                          sourceSlot,                                        /* source */
                                          0,                                                 /* src_offset */
                                          first_chunk);                                      /* size */
      // copy second part to beginning of buffer
      payloadCommunicationManager->memcpy(_payloadBuffer, /* destination */
                                          0,              /* dst_offset */
                                          sourceSlot,     /* source */
                                          first_chunk,    /* src_offset */
                                          second_chunk);  /* size */
 
      payloadCommunicationManager->fence(sourceSlot, 2, 0);
    }
    else
    {
      payloadCommunicationManager->memcpy(_payloadBuffer, getCircularBufferForPayloads()->getHeadPosition(), sourceSlot, 0, requiredBufferSize);
      payloadCommunicationManager->fence(sourceSlot, 1, 0);
    }
 
    // Remotely push an element into consumer side, updating consumer head indices
    getCircularBufferForCounts()->advanceHead(1);
    getCircularBufferForPayloads()->advanceHead(requiredBufferSize);
 
    // only update head index at consumer (byte size = one buffer element)
    coordinationCommunicationManager->memcpy(_consumerCoordinationBufferForCounts, 0, getCoordinationBufferForCounts(), 0, sizeof(_HICR_CHANNEL_COORDINATION_BUFFER_ELEMENT_TYPE));
    // only update head index at consumer (byte size = one buffer element)
    coordinationCommunicationManager->memcpy(
      _consumerCoordinationBufferForPayloads, 0, getCoordinationBufferForPayloads(), 0, sizeof(_HICR_CHANNEL_COORDINATION_BUFFER_ELEMENT_TYPE));
    // backend LPF needs this to complete
    coordinationCommunicationManager->fence(getCoordinationBufferForCounts(), 1, 0);
    coordinationCommunicationManager->fence(getCoordinationBufferForPayloads(), 1, 0);
 
    coordinationCommunicationManager->releaseGlobalLock(_consumerCoordinationBufferForCounts);
 
    return successFlag;
  }
 
  __INLINE__ size_t getDepth()
  {
    // Because the current implementation first receives the message size in the token buffer, followed
    // by the message payload, it is possible for the token buffer to have a larged depth (by 1) than the payload buffer.
    // Therefore, we need to return the minimum of the two depths
    return std::min(getCircularBufferForCounts()->getDepth(), getCircularBufferForPayloads()->getDepth() / getPayloadSize());
  }
 
  __INLINE__ bool isEmpty() { return (getCircularBufferForCounts()->getDepth() == 0) && (getCircularBufferForPayloads()->getDepth() == 0); }
 
  private:
 
  std::shared_ptr<GlobalMemorySlot> _payloadBuffer;
 
  std::shared_ptr<LocalMemorySlot> _sizeInfoBuffer;
 
  size_t _payloadSize;
 
  const std::shared_ptr<GlobalMemorySlot> _tokenSizeBuffer;
 
  const std::shared_ptr<GlobalMemorySlot> _consumerCoordinationBufferForCounts;
 
  const std::shared_ptr<GlobalMemorySlot> _consumerCoordinationBufferForPayloads;
};
 
} // namespace HiCR::channel::variableSize::MPSC::locking