Repository: google/filament
File: filament/backend/src/vulkan/VulkanStagePool.cpp (C++)

Code before:
fvkmemory::resource_ptr<VulkanStage::Segment> VulkanStage::acquireSegment(
        fvkmemory::ResourceManager* resManager, uint32_t segmentOffset, uint32_t numBytes) {
    auto segment = fvkmemory::resource_ptr<Segment>::construct(resManager, this, numBytes,
            segmentOffset, [this](uint32_t offset) { mSegments.erase(offset); });
    mSegments.insert({ segmentOffset, segment.get() });
    mCurrentOffset = segmentOffset + numBytes;
    return segment;
}

VulkanStagePool::VulkanStagePool(VmaAllocator allocator, fvkmemory::ResourceManager* resManager,
        VulkanCommands* commands, const VkPhysicalDeviceLimits* deviceLimits)
    : mAllocator(allocator),
      mResManager(resManager),
      mCommands(commands),
      mDeviceLimits(deviceLimits) {}

fvkmemory::resource_ptr<VulkanStage::Segment> VulkanStagePool::acquireStage(uint32_t numBytes,
        std::optional<uint32_t> alignment) {
    // Apply alignment to the byte count to ensure that, when we later flush
    // data written by the host, we only flush the atoms that we modified, and
    // no adjacent atoms.
    numBytes = alignToNonCoherentAtomSize(numBytes);

    // First check if a stage segment exists whose capacity is greater than or
    // equal to the requested size.
    auto iter = mStages.lower_bound(numBytes);

    // If we have potential buffer candidates, loop through and try to find one
    // that works with our alignment. Also, set the offset for the buffer.
    // If we have no buffer, create a new one.
    VulkanStage* pStage = nullptr;
    uint32_t segmentOffset = 0;
    while (iter != mStages.end()) {
        segmentOffset = iter->second->currentOffset();
        if (alignment.has_value()) {
            // If aligned, adjust the segment offset.
            segmentOffset = alignValue(segmentOffset, *alignment);
        }

        // Check for overflow + if there's space.
        if (segmentOffset >= iter->second->currentOffset() &&
                iter->second->capacity() - numBytes >= segmentOffset) {
            // This buffer has enough space.
            pStage = iter->second;
            mStages.erase(iter);
            break;
        }

        // This buffer wasn't suitable. Check the next larger one.
        ++iter;
    }

Review comment (suggestion):
I haven't thought too deeply about this, but I wonder if you could just do `lower_bound(numBytes + alignmentOffsetVal)` and not have to do the loop below.

Code after:
fvkmemory::resource_ptr<VulkanStage::Segment> VulkanStage::acquireSegment(
        fvkmemory::ResourceManager* resManager, uint32_t segmentOffset, uint32_t numBytes) {
    auto segment = fvkmemory::resource_ptr<Segment>::construct(resManager, this, numBytes,
            segmentOffset, [this](uint32_t offset) { mSegments.erase(offset); });
    mSegments.insert({ segmentOffset, segment.get() });
    mCurrentOffset = segmentOffset + numBytes;
    return segment;
}

VulkanStagePool::VulkanStagePool(VmaAllocator allocator, fvkmemory::ResourceManager* resManager,
        VulkanCommands* commands, const VkPhysicalDeviceLimits* deviceLimits)
    : mAllocator(allocator),
      mResManager(resManager),
      mCommands(commands),
      mDeviceLimits(deviceLimits) {}

fvkmemory::resource_ptr<VulkanStage::Segment> VulkanStagePool::acquireStage(uint32_t numBytes,
        uint32_t alignment) {
    // Apply alignment to the byte count to ensure that, when we later flush
    // data written by the host, we only flush the atoms that we modified, and
    // no adjacent atoms.
    numBytes = alignToNonCoherentAtomSize(numBytes);

    // First check if a stage segment exists whose capacity is greater than or
    // equal to the requested size.
    auto iter = mStages.lower_bound(numBytes);

    // If we have potential buffer candidates, loop through and try to find one
    // that works with our alignment. Also, set the offset for the buffer.
    // If we have no buffer, create a new one.
    VulkanStage* pStage = nullptr;
    uint32_t segmentOffset = 0;
    while (iter != mStages.end()) {
        segmentOffset = alignValue(iter->second->currentOffset(), alignment);

        // Check for overflow + if there's space.
        if (segmentOffset >= iter->second->currentOffset() &&
                iter->second->capacity() - numBytes >= segmentOffset) {
            // This buffer has enough space.
            pStage = iter->second;
            mStages.erase(iter);
            break;
        }

        // This buffer wasn't suitable. Check the next larger one.
        ++iter;
    }

    if (!pStage) {
        pStage = allocateNewStage(std::max(numBytes, STAGE_SIZE));
    }