// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "precomp.h"
#include "PooledUploadHeap.h"
#include "ExecutionContext.h"

namespace Dml
{
    PooledUploadHeap::PooledUploadHeap(ID3D12Device* device, std::shared_ptr<ExecutionContext> executionContext)
        : m_device(device)
        , m_executionContext(std::move(executionContext))
    {
    }

    static size_t Align(size_t offset, size_t alignment)
    {
        assert(alignment != 0);
        return (offset + alignment - 1) & ~(alignment - 1);
    }

    std::optional<size_t> PooledUploadHeap::FindOffsetForAllocation(const Chunk& chunk, size_t sizeInBytes)
    {
        assert(sizeInBytes != 0);

        if (chunk.capacityInBytes < sizeInBytes)
        {
            // This chunk isn't even big enough to accommodate this allocation
            return std::nullopt;
        }

        if (chunk.allocations.empty())
        {
            // The entire chunk is empty - allocate from the beginning
            return 0;
        }

        // Chunks are used as ring buffers, which means this allocation should go after the most recent previous
        // allocation

        const auto& lastAllocation = chunk.allocations.back();
        size_t newAllocationBegin = lastAllocation.offsetInChunk + lastAllocation.sizeInBytes;
        newAllocationBegin = Align(newAllocationBegin, c_allocationAlignment);

        if (newAllocationBegin + sizeInBytes < newAllocationBegin)
        {
            // Overflow
            return std::nullopt;
        }

        const auto& firstAllocation = chunk.allocations.front();
        if (firstAllocation.offsetInChunk <= lastAllocation.offsetInChunk)
        {
            // This is the case where there's potentially free space at the beginning and end of the chunk, but not
            // the middle:
            // e.g.
            //   |------XXXXYYYZZ------|
            //          ^^^^   ^^
            //          first  last

            if (newAllocationBegin + sizeInBytes <= chunk.capacityInBytes)
            {
                // There's enough space between the end of the last allocation and the end of the chunk
                return newAllocationBegin;
            }
            else
            {
                // Otherwise there's not enough space at the end of the chunk - try the beginning of the chunk instead
                newAllocationBegin = 0;
                if (newAllocationBegin + sizeInBytes <= firstAllocation.offsetInChunk)
                {
                    // There was enough space between the start of the buffer, and the start of the first allocation
                    return newAllocationBegin;
                }
            }
        }
        else
        {
            // This is the case where there's potentially free space in the middle of the chunk, but not at the edges
            // e.g.
            //   |YYYZZ---------XXXX-|
            //       ^^         ^^^^
            //       last       first

            if (newAllocationBegin + sizeInBytes <= firstAllocation.offsetInChunk)
            {
                // There's enough space between the end of the last allocation, and the start of the first one
                return newAllocationBegin;
            }
        }

        // Not enough space in this chunk to accommodate the requested allocation
        return std::nullopt;
    }

    /* static */ PooledUploadHeap::Chunk PooledUploadHeap::CreateChunk(ID3D12Device* device, size_t sizeInBytes)
    {
        ComPtr<ID3D12Resource> uploadBuffer;
        auto heap = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD);
        auto buffer = CD3DX12_RESOURCE_DESC::Buffer(sizeInBytes);

        ORT_THROW_IF_FAILED(device->CreateCommittedResource(
            &heap,
            D3D12_HEAP_FLAG_NONE,
            &buffer,
            D3D12_RESOURCE_STATE_GENERIC_READ,
            nullptr,
            IID_GRAPHICS_PPV_ARGS(uploadBuffer.ReleaseAndGetAddressOf())));

        return Chunk{ sizeInBytes, std::move(uploadBuffer) };
    }

    std::pair<PooledUploadHeap::Chunk*, size_t> PooledUploadHeap::Reserve(size_t sizeInBytes)
    {
        // Try to find a chunk with enough free space to accommodate the requested allocation size
        for (Chunk& chunk : m_chunks)
        {
            std::optional<size_t> offsetForAllocation = FindOffsetForAllocation(chunk, sizeInBytes);
            if (offsetForAllocation)
            {
                // There's enough space in this chunk - return
                return std::make_pair(&chunk, *offsetForAllocation);
            }
        }

        // No chunks were able to accommodate the allocation - create a new chunk and return that instead

        // At least double the capacity of the pool
        const size_t newChunkSize = std::max({ m_totalCapacity, c_minChunkSize, sizeInBytes });
        m_chunks.push_back(CreateChunk(m_device.Get(), newChunkSize));
        m_totalCapacity += newChunkSize;

        // Allocate from the beginning of the new chunk
        return std::make_pair(&m_chunks.back(), 0);
    }

    void PooledUploadHeap::ReclaimAllocations()
    {
        for (Chunk& chunk : m_chunks)
        {
            auto* allocs = &chunk.allocations;

            // Remove all allocations which have had their fences signaled - this indicates that they are no longer
            // being used by the GPU. We can stop as soon as we find an allocation which is still in use, because we
            // only use a single command queue and executions always complete in the order they were submitted.
            while (!allocs->empty() && allocs->front().doneEvent.IsSignaled())
            {
                allocs->pop_front();
            }
        }
    }

    GpuEvent PooledUploadHeap::BeginUploadToGpu(
        ID3D12Resource* dst,
        uint64_t dstOffset,
        D3D12_RESOURCE_STATES dstState,
        gsl::span<const std::byte> src)
    {
        assert(!src.empty());
        assert(dst->GetDesc().Dimension == D3D12_RESOURCE_DIMENSION_BUFFER);

        InvariantChecker checker(this);

        ReclaimAllocations();

        // Allocate space from the upload heap
        Chunk* chunk = nullptr;
        size_t offsetInChunk = 0;
        std::tie(chunk, offsetInChunk) = Reserve(src.size());

        assert(chunk != nullptr);
        assert(offsetInChunk + src.size() <= chunk->capacityInBytes);

        // Map the upload heap and copy the source data into it at the specified offset
        void* uploadHeapData = nullptr;
        ORT_THROW_IF_FAILED(chunk->resource->Map(0, nullptr, &uploadHeapData));
        memcpy(static_cast<byte*>(uploadHeapData) + offsetInChunk, src.data(), src.size());
        chunk->resource->Unmap(0, nullptr);

        // Copy from the upload heap into the destination resource
        m_executionContext->CopyBufferRegion(
            dst,
            dstOffset,
            dstState,
            chunk->resource.Get(),
            offsetInChunk,
            D3D12_RESOURCE_STATE_GENERIC_READ,
            src.size());

        GpuEvent doneEvent = m_executionContext->GetCurrentCompletionEvent();

        // Add an allocation entry to the chunk
        chunk->allocations.push_back(Allocation{ static_cast<size_t>(src.size()), offsetInChunk, doneEvent });

        return doneEvent;
    }

    void PooledUploadHeap::Trim()
    {
        InvariantChecker checker(this);

        ReclaimAllocations();

        // Release any chunks which have no allocations
        auto it = std::remove_if(m_chunks.begin(), m_chunks.end(), [](const Chunk& c) {
            return c.allocations.empty();
        });
        m_chunks.erase(it, m_chunks.end());

        // Re-calculate total capacity
        m_totalCapacity = 0;
        for (const auto& chunk : m_chunks)
        {
            m_totalCapacity += chunk.capacityInBytes;
        }
    }

    void PooledUploadHeap::AssertInvariants()
    {
    #ifdef _DEBUG

        auto chunkCapacityComparer = [](const Chunk& lhs, const Chunk& rhs) {
            return lhs.capacityInBytes < rhs.capacityInBytes;
        };

        // Chunks should be sorted by ascending capacity
        assert(std::is_sorted(m_chunks.begin(), m_chunks.end(), chunkCapacityComparer));

        // Allocations in a chunk should be sorted by ascending fence value
        for (const auto& chunk : m_chunks)
        {
            auto allocFenceValueComparer = [](const Allocation& lhs, const Allocation& rhs) {
                return lhs.doneEvent.fenceValue < rhs.doneEvent.fenceValue;
            };
            assert(std::is_sorted(chunk.allocations.begin(), chunk.allocations.end(), allocFenceValueComparer));
        }

        // Validate chunk properties
        for (const auto& chunk : m_chunks)
        {
            assert(chunk.resource != nullptr);
            assert(chunk.capacityInBytes == chunk.resource->GetDesc().Width);
        }

        // Validate allocation properties
        for (const auto& chunk : m_chunks)
        {
            for (const auto& alloc : chunk.allocations)
            {
                assert(alloc.offsetInChunk + alloc.sizeInBytes <= chunk.capacityInBytes);
                assert(alloc.offsetInChunk % c_allocationAlignment == 0); // Validate alignment
            }
        }

        // Validate no overlapping allocations
        for (const auto& chunk : m_chunks)
        {
            auto allocOffsetComparer = [](const Allocation& lhs, const Allocation& rhs) {
                return lhs.offsetInChunk < rhs.offsetInChunk;
            };

            std::vector<Allocation> allocationsSortedByOffset(chunk.allocations.begin(), chunk.allocations.end());
            std::sort(allocationsSortedByOffset.begin(), allocationsSortedByOffset.end(), allocOffsetComparer);

            for (size_t i = 1; i < allocationsSortedByOffset.size(); ++i)
            {
                const auto& alloc = allocationsSortedByOffset[i - 1];
                const auto& nextAlloc = allocationsSortedByOffset[i];
                assert(alloc.offsetInChunk + alloc.sizeInBytes <= nextAlloc.offsetInChunk);
            }
        }

        // Validate total capacity of pool
        size_t calculatedCapacity = 0;
        for (const auto& chunk : m_chunks)
        {
            calculatedCapacity += chunk.capacityInBytes;
        }
        assert(calculatedCapacity == m_totalCapacity);

    #endif // #ifdef _DEBUG
    }
} // namespace Dml
