Seele/src/Engine/Graphics/Vulkan/Buffer.cpp

#include "Buffer.h"
#include "Command.h"

using namespace Seele;
using namespace Seele::Vulkan;

struct PendingBuffer {
  uint64 offset;
  uint64 size;
  VkBuffer stagingBuffer;
  VmaAllocation allocation;
  Gfx::QueueType prevQueue;
  bool writeOnly;
};

static Map<Vulkan::Buffer *, PendingBuffer> pendingBuffers;

Buffer::Buffer(PGraphics graphics, uint64 size, VkBufferUsageFlags usage,
               Gfx::QueueType &queueType, bool dynamic, std::string name)
    : graphics(graphics), currentBuffer(0), size(size), owner(queueType),
      name(name) {
  if (dynamic) {
    numBuffers = Gfx::numFramesBuffered;
  } else {
    numBuffers = 1;
  }
  VkBufferCreateInfo info = {
      .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
      .pNext = nullptr,
      .size = size,
      .usage = usage,
      .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };
  VmaAllocationCreateInfo allocInfo = {
      .flags = VMA_ALLOCATION_CREATE_HOST_ACCESS_ALLOW_TRANSFER_INSTEAD_BIT |
               VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
      .usage = VMA_MEMORY_USAGE_AUTO,
  };
  for (uint32 i = 0; i < numBuffers; ++i) {
    vmaCreateBuffer(graphics->getAllocator(), &info, &allocInfo,
                    &buffers[i].buffer, &buffers[i].allocation,
                    &buffers[i].info);
    vmaGetAllocationMemoryProperties(graphics->getAllocator(),
                                     buffers[i].allocation,
                                     &buffers[i].properties);
    // std::cout << "Create buffer " << std::hex << (uint64)buffers[i].buffer <<
    // std::dec;
    if (!name.empty()) {
      VkDebugUtilsObjectNameInfoEXT nameInfo = {
          .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
          .pNext = nullptr,
          .objectType = VK_OBJECT_TYPE_BUFFER,
          .objectHandle = (uint64)buffers[i].buffer,
          .pObjectName = this->name.c_str()};
      graphics->vkSetDebugUtilsObjectNameEXT(&nameInfo);
      // std::cout << ": " << name;
    }
    // std::cout << std::endl;
  }
}

Buffer::~Buffer() {
  for (uint32 i = 0; i < numBuffers; ++i) {
    graphics->getDestructionManager()->queueBuffer(
        graphics->getQueueCommands(owner)->getCommands(), buffers[i].buffer,
        buffers[i].allocation);
  }
}

void Buffer::executeOwnershipBarrier(Gfx::QueueType newOwner) {
  Gfx::QueueFamilyMapping mapping = graphics->getFamilyMapping();
  VkBufferMemoryBarrier barrier = {
      .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
      .pNext = nullptr,
      .srcQueueFamilyIndex = mapping.getQueueTypeFamilyIndex(owner),
      .dstQueueFamilyIndex = mapping.getQueueTypeFamilyIndex(newOwner),
      .offset = 0,
      .size = size,
  };
  PCommandPool sourcePool = graphics->getQueueCommands(owner);
  PCommandPool dstPool = nullptr;
  VkPipelineStageFlags srcStage = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
  VkPipelineStageFlags dstStage = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
  assert(barrier.srcQueueFamilyIndex != barrier.dstQueueFamilyIndex);
  if (owner == Gfx::QueueType::TRANSFER) {
    barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    srcStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
  } else if (owner == Gfx::QueueType::COMPUTE) {
    barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
    srcStage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
  } else if (owner == Gfx::QueueType::GRAPHICS) {
    barrier.srcAccessMask = getSourceAccessMask();
    srcStage = VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT;
  }
  if (newOwner == Gfx::QueueType::TRANSFER) {
    barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
    dstStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    dstPool = graphics->getTransferCommands();
  } else if (newOwner == Gfx::QueueType::COMPUTE) {
    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    dstStage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
    dstPool = graphics->getComputeCommands();
  } else if (newOwner == Gfx::QueueType::GRAPHICS) {
    barrier.dstAccessMask = getDestAccessMask();
    dstStage = VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT;
    dstPool = graphics->getGraphicsCommands();
  }
  VkCommandBuffer srcCommand = sourcePool->getCommands()->getHandle();
  VkCommandBuffer dstCommand = dstPool->getCommands()->getHandle();
  VkBufferMemoryBarrier dynamicBarriers[Gfx::numFramesBuffered];
  for (uint32 i = 0; i < numBuffers; ++i) {
    dynamicBarriers[i] = barrier;
    dynamicBarriers[i].buffer = buffers[i].buffer;
  }
  vkCmdPipelineBarrier(srcCommand, srcStage, srcStage, 0, 0, nullptr,
                       numBuffers, dynamicBarriers, 0, nullptr);
  vkCmdPipelineBarrier(dstCommand, dstStage, dstStage, 0, 0, nullptr,
                       numBuffers, dynamicBarriers, 0, nullptr);
  sourcePool->submitCommands();
}

void Buffer::executePipelineBarrier(VkAccessFlags srcAccess,
                                    VkPipelineStageFlags srcStage,
                                    VkAccessFlags dstAccess,
                                    VkPipelineStageFlags dstStage) {
  PCommand commandBuffer = graphics->getQueueCommands(owner)->getCommands();
  VkBufferMemoryBarrier barrier = {
      .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
      .pNext = nullptr,
      .srcAccessMask = srcAccess,
      .dstAccessMask = dstAccess,
      .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
      .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
      .offset = 0,
      .size = size,
  };
  VkBufferMemoryBarrier dynamicBarriers[Gfx::numFramesBuffered];
  for (uint32 i = 0; i < numBuffers; ++i) {
    dynamicBarriers[i] = barrier;
    dynamicBarriers[i].buffer = buffers[i].buffer;
  }
  vkCmdPipelineBarrier(commandBuffer->getHandle(), srcStage, dstStage, 0, 0,
                       nullptr, numBuffers, dynamicBarriers, 0, nullptr);
}

void *Buffer::map(bool writeOnly) { return mapRegion(0, size, writeOnly); }

void *Buffer::mapRegion(uint64 regionOffset, uint64 regionSize,
                        bool writeOnly) {
  void *data = nullptr;

  PendingBuffer pending;
  pending.writeOnly = writeOnly;
  pending.prevQueue = owner;
  pending.offset = regionOffset;
  pending.size = regionSize;
  if (writeOnly) {
    if (buffers[currentBuffer].properties &
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) {
      VK_CHECK(vmaMapMemory(graphics->getAllocator(),
                            buffers[currentBuffer].allocation, &data));
    } else {
      VkBufferCreateInfo stagingInfo = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
          .pNext = nullptr,
          .flags = 0,
          .size = regionSize,
          .usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
          .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
      };
      VmaAllocationCreateInfo allocInfo = {
          .usage = VMA_MEMORY_USAGE_AUTO,
          .requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
      };
      VK_CHECK(vmaCreateBuffer(graphics->getAllocator(), &stagingInfo,
                               &allocInfo, &pending.stagingBuffer,
                               &pending.allocation, nullptr));
      vmaMapMemory(graphics->getAllocator(), pending.allocation, &data);
    }
  } else {
    assert(false);
  }
  pendingBuffers[this] = std::move(pending);

  assert(data);
  return data;
}

void Buffer::unmap() {
  auto found = pendingBuffers.find(this);
  if (found != pendingBuffers.end()) {
    PendingBuffer &pending = found->value;
    if (pending.writeOnly) {
      if (buffers[currentBuffer].properties &
          VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) {
        vmaUnmapMemory(graphics->getAllocator(),
                       buffers[currentBuffer].allocation);
      } else {
        vmaFlushAllocation(graphics->getAllocator(), pending.allocation, 0,
                           VK_WHOLE_SIZE);
        vmaUnmapMemory(graphics->getAllocator(), pending.allocation);
        PCommand command = graphics->getQueueCommands(owner)->getCommands();
        VkCommandBuffer cmdHandle = command->getHandle();

        VkBufferCopy region = {
            .srcOffset = 0,
            .dstOffset = pending.offset,
            .size = pending.size,
        };
        VkBufferMemoryBarrier barrier = {
            .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
            .pNext = nullptr,
            .srcAccessMask = VK_ACCESS_MEMORY_READ_BIT,
            .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .buffer = buffers[currentBuffer].buffer,
            .offset = 0,
            .size = size,
        };
        vkCmdPipelineBarrier(cmdHandle, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
                             VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 1,
                             &barrier, 0, nullptr);
        vkCmdCopyBuffer(cmdHandle, pending.stagingBuffer,
                        buffers[currentBuffer].buffer, 1, &region);
        barrier = {
            .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
            .pNext = nullptr,
            .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
            .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .buffer = buffers[currentBuffer].buffer,
            .offset = 0,
            .size = size,
        };
        vkCmdPipelineBarrier(cmdHandle, VK_PIPELINE_STAGE_TRANSFER_BIT,
                             VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr,
                             1, &barrier, 0, nullptr);
        graphics->getDestructionManager()->queueBuffer(
            command, pending.stagingBuffer, pending.allocation);
      }
    }
    // requestOwnershipTransfer(pending.prevQueue);

    pendingBuffers.erase(this);
  }
}

UniformBuffer::UniformBuffer(PGraphics graphics,
                             const UniformBufferCreateInfo &createInfo)
    : Gfx::UniformBuffer(graphics->getFamilyMapping(), createInfo.sourceData),
      Vulkan::Buffer(graphics, createInfo.sourceData.size,
                     VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, currentOwner,
                     createInfo.dynamic, createInfo.name) {
  if (createInfo.sourceData.data != nullptr) {
    void *data = map();
    std::memcpy(data, createInfo.sourceData.data, createInfo.sourceData.size);
    unmap();
  }
}

UniformBuffer::~UniformBuffer() {}

bool UniformBuffer::updateContents(const DataSource &sourceData) {
  if (!Gfx::UniformBuffer::updateContents(sourceData)) {
    // no update was performed, skip
    return false;
  }
  void *data = map();
  std::memcpy(data, sourceData.data, sourceData.size);
  unmap();
  return true;
}

void UniformBuffer::beginFrame() { Vulkan::Buffer::advanceBuffer(); }

void UniformBuffer::requestOwnershipTransfer(Gfx::QueueType newOwner) {
  Gfx::QueueOwnedResource::transferOwnership(newOwner);
}

void UniformBuffer::executeOwnershipBarrier(Gfx::QueueType newOwner) {
  Vulkan::Buffer::executeOwnershipBarrier(newOwner);
}

void UniformBuffer::executePipelineBarrier(VkAccessFlags srcAccess,
                                           VkPipelineStageFlags srcStage,
                                           VkAccessFlags dstAccess,
                                           VkPipelineStageFlags dstStage) {
  Vulkan::Buffer::executePipelineBarrier(srcAccess, srcStage, dstAccess,
                                         dstStage);
}

VkAccessFlags UniformBuffer::getSourceAccessMask() {
  return VK_ACCESS_MEMORY_WRITE_BIT;
}

VkAccessFlags UniformBuffer::getDestAccessMask() {
  return VK_ACCESS_UNIFORM_READ_BIT;
}

ShaderBuffer::ShaderBuffer(PGraphics graphics,
                           const ShaderBufferCreateInfo &sourceData)
    : Gfx::ShaderBuffer(graphics->getFamilyMapping(), sourceData.numElements,
                        sourceData.sourceData),
      Vulkan::Buffer(graphics, sourceData.sourceData.size,
                     VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, currentOwner,
                     sourceData.dynamic, sourceData.name) {
  if (sourceData.sourceData.data != nullptr) {
    void *data = map();
    std::memcpy(data, sourceData.sourceData.data, sourceData.sourceData.size);
    unmap();
  }
}

ShaderBuffer::~ShaderBuffer() {}

bool ShaderBuffer::updateContents(const DataSource &sourceData) {
  assert(sourceData.size <= getSize());
  Gfx::ShaderBuffer::updateContents(sourceData);
  // We always want to update, as the contents could be different on the GPU
  void *data = map();
  std::memcpy(data, sourceData.data, sourceData.size);
  unmap();
  return true;
}

void ShaderBuffer::beginFrame() { Vulkan::Buffer::advanceBuffer(); }

void ShaderBuffer::requestOwnershipTransfer(Gfx::QueueType newOwner) {
  Gfx::QueueOwnedResource::transferOwnership(newOwner);
}

void ShaderBuffer::executeOwnershipBarrier(Gfx::QueueType newOwner) {
  Vulkan::Buffer::executeOwnershipBarrier(newOwner);
}

void ShaderBuffer::executePipelineBarrier(VkAccessFlags srcAccess,
                                          VkPipelineStageFlags srcStage,
                                          VkAccessFlags dstAccess,
                                          VkPipelineStageFlags dstStage) {
  Vulkan::Buffer::executePipelineBarrier(srcAccess, srcStage, dstAccess,
                                         dstStage);
}

VkAccessFlags ShaderBuffer::getSourceAccessMask() {
  return VK_ACCESS_MEMORY_WRITE_BIT;
}

VkAccessFlags ShaderBuffer::getDestAccessMask() {
  return VK_ACCESS_MEMORY_READ_BIT;
}

VertexBuffer::VertexBuffer(PGraphics graphics,
                           const VertexBufferCreateInfo &sourceData)
    : Gfx::VertexBuffer(graphics->getFamilyMapping(), sourceData.numVertices,
                        sourceData.vertexSize, sourceData.sourceData.owner),
      Vulkan::Buffer(graphics, sourceData.sourceData.size,
                     VK_BUFFER_USAGE_VERTEX_BUFFER_BIT, currentOwner, false,
                     sourceData.name) {
  if (sourceData.sourceData.data != nullptr) {
    void *data = map();
    std::memcpy(data, sourceData.sourceData.data, sourceData.sourceData.size);
    unmap();
  }
}

VertexBuffer::~VertexBuffer() {}

void VertexBuffer::updateRegion(DataSource update) {
  void *data = mapRegion(update.offset, update.size);
  std::memcpy(data, update.data, update.size);
  unmap();
}

void VertexBuffer::download(Array<uint8> &buffer) {
  void *data = map(false);
  buffer.resize(size);
  std::memcpy(buffer.data(), data, size);
  unmap();
}

void VertexBuffer::beginFrame() { Vulkan::Buffer::advanceBuffer(); }

void VertexBuffer::requestOwnershipTransfer(Gfx::QueueType newOwner) {
  Gfx::QueueOwnedResource::transferOwnership(newOwner);
}

void VertexBuffer::executeOwnershipBarrier(Gfx::QueueType newOwner) {
  Vulkan::Buffer::executeOwnershipBarrier(newOwner);
}

void VertexBuffer::executePipelineBarrier(VkAccessFlags srcAccess,
                                          VkPipelineStageFlags srcStage,
                                          VkAccessFlags dstAccess,
                                          VkPipelineStageFlags dstStage) {
  Vulkan::Buffer::executePipelineBarrier(srcAccess, srcStage, dstAccess,
                                         dstStage);
}

VkAccessFlags VertexBuffer::getSourceAccessMask() {
  return VK_ACCESS_MEMORY_WRITE_BIT;
}

VkAccessFlags VertexBuffer::getDestAccessMask() {
  return VK_ACCESS_VERTEX_ATTRIBUTE_READ_BIT;
}

IndexBuffer::IndexBuffer(PGraphics graphics,
                         const IndexBufferCreateInfo &sourceData)
    : Gfx::IndexBuffer(graphics->getFamilyMapping(), sourceData.sourceData.size,
                       sourceData.indexType, sourceData.sourceData.owner),
      Vulkan::Buffer(graphics, sourceData.sourceData.size,
                     VK_BUFFER_USAGE_INDEX_BUFFER_BIT, currentOwner, false,
                     sourceData.name) {
  if (sourceData.sourceData.data != nullptr) {
    void *data = map();
    std::memcpy(data, sourceData.sourceData.data, sourceData.sourceData.size);
    unmap();
  }
}

IndexBuffer::~IndexBuffer() {}

void IndexBuffer::download(Array<uint8> &buffer) {
  void *data = map(false);
  buffer.resize(size);
  std::memcpy(buffer.data(), data, size);
  unmap();
}

void IndexBuffer::beginFrame() { Vulkan::Buffer::advanceBuffer(); }

void IndexBuffer::requestOwnershipTransfer(Gfx::QueueType newOwner) {
  Gfx::QueueOwnedResource::transferOwnership(newOwner);
}

void IndexBuffer::executeOwnershipBarrier(Gfx::QueueType newOwner) {
  Vulkan::Buffer::executeOwnershipBarrier(newOwner);
}

void IndexBuffer::executePipelineBarrier(VkAccessFlags srcAccess,
                                         VkPipelineStageFlags srcStage,
                                         VkAccessFlags dstAccess,
                                         VkPipelineStageFlags dstStage) {
  Vulkan::Buffer::executePipelineBarrier(srcAccess, srcStage, dstAccess,
                                         dstStage);
}

VkAccessFlags IndexBuffer::getSourceAccessMask() {
  return VK_ACCESS_MEMORY_WRITE_BIT;
}

VkAccessFlags IndexBuffer::getDestAccessMask() {
  return VK_ACCESS_INDEX_READ_BIT;
}