Univerxel/src/client/render/vk/api/Buffers.cpp

#include "Buffers.hpp"
#include "../Allocator.hpp"
#include "../../../../core/utils/logger.hpp"
#include <algorithm>

using namespace render::vk;

Buffer::~Buffer() {
    vkDestroyBuffer(Allocator::GetDefault()->getDevice(), ref, ALLOC);
    //NOTE: memory_ptr self destroy
}

memory::ptr render::vk::createBuffer(VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags properties, const render::data_view view, Buffer::info &out) {
    auto alloc = Allocator::GetDefault();
    auto device = alloc->getDevice();

    VkBufferCreateInfo bufferInfo{};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    bufferInfo.size = size;
    bufferInfo.usage = usage;
    if (view)
        bufferInfo.usage |= VK_BUFFER_USAGE_TRANSFER_DST_BIT;

    if (vkCreateBuffer(device, &bufferInfo, ALLOC, &out.ref) != VK_SUCCESS) {
        LOG_E("Failed to create buffer");
        return memory::GetNull();
    }
    out.offset = 0;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(device, out.ref, &memRequirements);

    auto memory = alloc->allocate(memRequirements, properties);
    if (!memory || vkBindBufferMemory(device, out.ref, memory->ref, memory->offset) != VK_SUCCESS) {
        LOG_E("Failed to allocate buffer memory");
        return memory::GetNull();
    }

    if (view) {
        if (memory->ptr != nullptr) {
            memory->write(view.ptr, view.size);
        } else if(auto staging = WritableBuffer::Create(size)) {
            staging->write(view, 0);
            alloc->copyBuffer(staging->getRef(), out.ref, size);
        } else {
            FATAL("Cannot allocate staging memory");
            return memory::GetNull();
        }
    }
    return memory;
}
memory::ptr render::vk::createBuffers(const std::vector<Buffer::requirement>& requirements, VkMemoryPropertyFlags properties, std::vector<Buffer::info>& out) {
    assert(!requirements.empty());
    out.resize(requirements.size()+1);

    auto alloc = Allocator::GetDefault();
    auto device = alloc->getDevice();

    // Create buffers
    VkMemoryRequirements memRequirements = {0, 0, UINT32_MAX};
    std::vector<VkDeviceSize> sizes;
    sizes.resize(requirements.size());
    for (size_t i = 0; i < requirements.size(); i++) {
        VkBufferCreateInfo bufferInfo{};
        bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
        bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
        bufferInfo.size = requirements[i].size;
        bufferInfo.usage = static_cast<VkBufferUsageFlags>(requirements[i].usage);
        if (requirements[i].view)
            bufferInfo.usage |= VK_BUFFER_USAGE_TRANSFER_DST_BIT;

        if (vkCreateBuffer(device, &bufferInfo, ALLOC, &out[i].ref) != VK_SUCCESS) {
            LOG_E("Failed to create buffer");
            return memory::GetNull();
        }

        VkMemoryRequirements individualMemRequirements;
        vkGetBufferMemoryRequirements(device, out[i].ref, &individualMemRequirements);
        memRequirements.alignment = std::max(memRequirements.alignment, individualMemRequirements.alignment);
        memRequirements.memoryTypeBits &= individualMemRequirements.memoryTypeBits;
        sizes[i] = individualMemRequirements.size;
    }

    // Align blocks
    auto aligned = [&](VkDeviceSize offset) {
        if (offset % memRequirements.alignment == 0)
            return offset;
        return offset + memRequirements.alignment - (offset % memRequirements.alignment);
    };
    out[0].offset = 0;
    for (size_t i = 1; i < out.size(); i++) {
        out[i].offset = aligned(out[i-1].offset + sizes[i-1]);
    }
    memRequirements.size = out.back().offset;
    out.pop_back();

    // Bind memory
    auto memory = alloc->allocate(memRequirements, properties);
    if (!memory) {
        LOG_E("Failed to allocate buffers");
        return memory::GetNull();
    }
    for (size_t i = 0; i < requirements.size(); i++) {
        if (vkBindBufferMemory(device, out[i].ref, memory->ref, memory->offset + out[i].offset) != VK_SUCCESS) {
            LOG_E("Failed to bind buffer");
            return memory::GetNull();
        }
    }

    VkDeviceSize stagingSize = 0;
    for (auto& requirement: requirements)
        if (requirement.view)
            stagingSize = std::max<VkDeviceSize>(stagingSize, requirement.size);

    // Copy views
    // MAYBE: allow single copy
    if (stagingSize != 0) {
        if (memory->ptr != nullptr) {
            for (size_t i = 0; i < requirements.size(); i++) {
                if (requirements[i].view) {
                    assert(requirements[i].view.size <= requirements[i].size);
                    memory->write(requirements[i].view.ptr, requirements[i].view.size, out[i].offset);
                }
            }
        } else if(auto staging = WritableBuffer::Create(stagingSize)) {
            for (size_t i = 0; i < requirements.size(); i++) {
                if (requirements[i].view) {
                    assert(requirements[i].view.size <= requirements[i].size);
                    staging->write(requirements[i].view, 0);
                    alloc->copyBuffer(staging->getRef(), out[i].ref, requirements[i].size);
                }
            }
        } else {
            FATAL("Cannot allocate staging memory");
            return memory::GetNull();
        }
    }

    return memory;
}

void Buffer::MakeDefault() { }

std::unique_ptr<WritableBuffer> WritableBuffer::Create(size_t size, Usage usage, const render::data_view write) {
    vk::Buffer::info tmp;
    auto mem = createBuffer(size, static_cast<int>(usage), memory::HOST_EASILY_WRITABLE, write, tmp);
    return std::unique_ptr<WritableBuffer>(new WritableBuffer(tmp.ref, std::move(mem), tmp.offset));
}
void WritableBuffer::write(const render::data_view view, size_t offset) {
    memory->write(view.ptr, view.size, memOffset + offset);
}
void WritableBuffer::read(render::data_ref ref, size_t offset) {
    memory->read(ref.ptr, ref.size, memOffset + offset);
}

static std::vector<std::vector<std::unique_ptr<ShortIndexedVertexBuffer>>> usedVertexBuffers;
ShortIndexedVertexBuffer::~ShortIndexedVertexBuffer() {
    if (lastUseImage == UINT32_MAX) {
        vkDestroyBuffer(Allocator::GetDefault()->getDevice(), vertex, ALLOC);
        vkDestroyBuffer(Allocator::GetDefault()->getDevice(), index, ALLOC);
        //NOTE: memory_ptr self destroy
    } else {
        if (lastUseImage >= usedVertexBuffers.size()) {
            usedVertexBuffers.resize(lastUseImage+1);
        }
        usedVertexBuffers.at(lastUseImage).emplace_back(new ShortIndexedVertexBuffer(vertex, index, std::move(memory)));
    }
}
void ShortIndexedVertexBuffer::ClearUnused(uint32_t image) {
    if (image < usedVertexBuffers.size()) {
        usedVertexBuffers.at(image).clear();
    }
}

std::unique_ptr<ShortIndexedVertexBuffer> ShortIndexedVertexBuffer::Create(const render::data_view vertices, const render::data_view indices) {
    std::vector<vk::Buffer::info> tmp;
    auto mem = createBuffers({{vertices.size, Usage::VERTEX, vertices}, {indices.size, Usage::INDEX, indices}}, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, tmp);
    return std::unique_ptr<ShortIndexedVertexBuffer>(new ShortIndexedVertexBuffer(tmp.at(0).ref, tmp.at(1).ref, std::move(mem)));
}

BufferGroup::~BufferGroup() {
    free();
}

void BufferGroup::allocate(const std::vector<Buffer::requirement> & reqs, bool writable) {
    free();
    memory = createBuffers(reqs, writable ? memory::HOST_EASILY_WRITABLE : (VkMemoryPropertyFlags)VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, refs);
}
void BufferGroup::free() {
    for (const auto& buffer: refs) {
        vkDestroyBuffer(Allocator::GetDefault()->getDevice(), buffer.ref, ALLOC);
    }
    refs.clear();
    memory = memory::GetNull();
}
void BufferGroup::write(size_t i, const render::data_view view) {
    memory->write(view.ptr, view.size, refs.at(i).offset);
}