objectStore.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 <atomic>
 
#include <hicr/core/globalMemorySlot.hpp>
#include <hicr/core/instance.hpp>
#include <hicr/core/communicationManager.hpp>
#include <hicr/core/memoryManager.hpp>
#include <hicr/frontends/objectStore/dataObject.hpp>
 
enum
{
  OBJECT_STORE_KEY_INSTANCE_ID_BITS = 32
};
 
namespace HiCR::objectStore
{
 
using compoundId_t = uint64_t;
 
class ObjectStore
{
  public:
 
  ObjectStore(CommunicationManager        &communicationManager,
              GlobalMemorySlot::tag_t      tag,
              MemoryManager               &memoryManager,
              std::shared_ptr<MemorySpace> memorySpace,
              Instance::instanceId_t       instanceId)
    : _memoryManager(memoryManager),
      _communicationManager(communicationManager),
      _tag(tag),
      _memorySpace(memorySpace),
      _instanceId(instanceId)
  {}
 
  ~ObjectStore() = default;
 
  [[nodiscard]] std::shared_ptr<MemorySpace> getMemorySpace() const { return _memorySpace; }
 
  [[nodiscard]] __INLINE__ std::shared_ptr<DataObject> createObject(void *ptr, size_t size, blockId id)
  {
    // Register the given allocation as a memory slot
    auto slot = _memoryManager.registerLocalMemorySlot(_memorySpace, ptr, size);
 
    return std::make_shared<DataObject>(_instanceId, id, slot);
  }
 
  [[nodiscard]] __INLINE__ std::shared_ptr<DataObject> createObject(std::shared_ptr<LocalMemorySlot> slot, blockId id)
  {
    return std::make_shared<DataObject>(_instanceId, id, slot);
  }
 
  __INLINE__ void publish(std::shared_ptr<DataObject> dataObject)
  {
    blockId objectId = dataObject->_id;
 
    if (dataObject->_globalSlot)
    {
      _communicationManager.destroyPromotedGlobalMemorySlot(dataObject->_globalSlot);
      HICR_THROW_LOGIC("Trying to publish a block that has already been published."); // FIXME: What do we do with republishing?
    }
    // One-sided publish of the global slot
    auto globalSlot = _communicationManager.promoteLocalMemorySlot(dataObject->_localSlot, _tag);
 
    dataObject->_globalSlot = globalSlot;
    _globalObjects.insert({dataObject->_instanceId << OBJECT_STORE_KEY_INSTANCE_ID_BITS | objectId, dataObject});
  }
 
  __INLINE__ std::shared_ptr<LocalMemorySlot> get(DataObject &dataObject)
  {
    // Deduce globally unique block ID for the GlobalMemorySlot key, by combining the given instance and block IDs
    compoundId_t compoundId = dataObject._instanceId << OBJECT_STORE_KEY_INSTANCE_ID_BITS | dataObject._id;
 
    _globalObjects[compoundId] = std::make_shared<DataObject>(dataObject);
 
    // If the data object has never been fetched before, allocate a local memory slot for it
    if (dataObject._localSlot == nullptr)
    {
      // Allocate a local memory slot for the block, associate it with the global slot and add it to the blocks of the given instance
      auto newSlot          = _memoryManager.allocateLocalMemorySlot(_memorySpace, dataObject._size);
      dataObject._localSlot = newSlot;
    }
 
    if (dataObject._globalSlot == nullptr) HICR_THROW_LOGIC("Trying to get a block that has not been properly transferred.");
 
    // Initiate the data movement from the owner to the current worker
    _communicationManager.memcpy(dataObject._localSlot, 0, dataObject._globalSlot, 0, dataObject._size);
    return dataObject._localSlot;
 
    // Notes: Consider a variant (maybe hidden under another layer/wrapper) that allows to get portion
    // of the block, using offset and size
  }
 
  __INLINE__ void destroy(DataObject &dataObject)
  {
    // Only free the local memory slot if it is a fetched object, which means the local slot allocation
    // was done by the object store during a call to #get
    if (dataObject._instanceId != _instanceId)
    {
      if (dataObject._localSlot) _memoryManager.freeLocalMemorySlot(dataObject._localSlot);
    }
 
    if (dataObject._globalSlot) _communicationManager.destroyPromotedGlobalMemorySlot(dataObject._globalSlot);
 
    compoundId_t compoundId = dataObject._instanceId << OBJECT_STORE_KEY_INSTANCE_ID_BITS | dataObject._id;
    auto         it         = _globalObjects.find(compoundId);
 
    if (it != _globalObjects.end()) { _globalObjects.erase(it); }
  }
 
  __INLINE__ handle serialize(DataObject &dataObject)
  {
    auto ret                      = handle{};
    ret.instanceId                = dataObject._instanceId;
    ret.id                        = dataObject._id;
    ret.size                      = dataObject._size;
    uint8_t *serializedGlobalSlot = _communicationManager.serializeGlobalMemorySlot(dataObject._globalSlot); // implicitly produces allocation
    std::memcpy(ret.serializedGlobalSlot, serializedGlobalSlot, sizeof(ret.serializedGlobalSlot));
 
    // free the allocation produced by serialize() TODO: dangerous, remove
    delete[] serializedGlobalSlot;
    return ret;
  }
 
  __INLINE__ std::shared_ptr<DataObject> deserialize(const handle &handle)
  {
    auto dataObject   = std::make_shared<DataObject>(handle.instanceId, handle.id, nullptr);
    dataObject->_size = handle.size;
    dataObject->_globalSlot =
      _communicationManager.deserializeGlobalMemorySlot(reinterpret_cast<uint8_t *>(const_cast<uint8_t(*)[28 + sizeof(size_t)]>(&handle.serializedGlobalSlot)), _tag);
    return dataObject;
  }
 
  __INLINE__ bool fence()
  {
    _communicationManager.fence(_tag);
    return true;
  }
 
  __INLINE__ bool fence(const std::shared_ptr<DataObject> &dataObject)
  {
    // Use the zero-cost variant of fence with (sent,recv) = (0,1)
    _communicationManager.fence(dataObject->_localSlot, 0, 1);
    return true;
  }
 
#ifdef HICR_ENABLE_NONBLOCKING_FENCE // Dummy flag for when we implement non-blocking fence
  static bool test_fence(const Tag &tag);
#endif
 
  private:
 
  MemoryManager &_memoryManager;
 
  CommunicationManager &_communicationManager;
 
  const GlobalMemorySlot::tag_t _tag;
 
  const std::shared_ptr<MemorySpace> _memorySpace;
 
  std::map<compoundId_t, std::shared_ptr<DataObject>> _globalObjects{};
 
  const Instance::instanceId_t _instanceId;
};
 
} // namespace HiCR::objectStore