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

Code before:
}

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;
    }

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

    // Note: this allocation updates `currentOffset` and `segments` within
    // the parent stage. When destroyed, it will update `segments`.

Review comment (suggestion):
note - there should only be a handful of blocks to search through, i don't anticipate this being a very long search. also, since alignment is only upto 16 bytes for textures / inputs, we're searching for buffers that are within 16 bytes of the required space. that means that most unsuitable buffers will already have been ruled out on line 73.

Code after:
}

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));
    }

    // Note: this allocation updates `currentOffset` and `segments` within
    // the parent stage. When destroyed, it will update `segments`.
    fvkmemory::resource_ptr<VulkanStage::Segment> pSegment =
            pStage->acquireSegment(mResManager, segmentOffset, numBytes);

    // Update the stage's metadata, and reinsert it with the remaining segment