o3de/Gems/AtomTressFX/Code/Rendering/SharedBuffer.cpp

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <Atom/RHI.Reflect/ShaderSemantic.h>
#include <Atom/RHI/RHIUtils.h>
#include <Atom/RHI/Factory.h>

#include <Atom/RPI.Reflect/Buffer/BufferAssetCreator.h>
#include <Atom/RPI.Reflect/ResourcePoolAssetCreator.h>

#include <Rendering/SharedBuffer.h>
#include <Rendering/HairCommon.h>

#include <numeric>

namespace AZ::Render
{
    //! Setting the constructor as private will create compile error to remind the developer to set
    //! the buffer Init in the FeatureProcessor and initialize properly

    SharedBuffer::SharedBuffer()
    {
        AZ_Warning("SharedBuffer", false, "Missing information to properly create SharedBuffer. Init is required");
    }

    SharedBuffer::SharedBuffer(AZStd::string bufferName, AZStd::vector<SrgBufferDescriptor>& buffersDescriptors)
    {
        m_bufferName = bufferName;
        Init(bufferName, buffersDescriptors);
    }

    SharedBuffer::~SharedBuffer()
    {
        m_bufferAsset = {};
    }

    //! Crucial method that will ensure that the alignment for the BufferViews is always kept.
    //! This is important when requesting a BufferView as the offset needs to be aligned according
    //!  to the element type of the buffer.
    void SharedBuffer::CalculateAlignment(AZStd::vector<SrgBufferDescriptor>& buffersDescriptors)
    {
        m_alignment = 1;
        for (uint8_t bufferIndex = 0; bufferIndex < buffersDescriptors.size() ; ++bufferIndex)
        {
            // Using the least common multiple enables resource views to be typed and ensures they can get
            // an offset in bytes that is a multiple of an element count
            m_alignment = std::lcm(m_alignment, buffersDescriptors[bufferIndex].m_elementSize);
        }
    }

    void SharedBuffer::InitAllocator()
    {
        RHI::FreeListAllocator::Descriptor allocatorDescriptor;
        allocatorDescriptor.m_alignmentInBytes = m_alignment;
        allocatorDescriptor.m_capacityInBytes = m_sizeInBytes;
        allocatorDescriptor.m_policy = RHI::FreeListAllocatorPolicy::BestFit;
        allocatorDescriptor.m_garbageCollectLatency = 0;
        m_freeListAllocator.Init(allocatorDescriptor);
    }

    void SharedBuffer::CreateBuffer()
    {
        SrgBufferDescriptor  descriptor = SrgBufferDescriptor(
            RPI::CommonBufferPoolType::ReadWrite, RHI::Format::Unknown,
            sizeof(float), uint32_t(m_sizeInBytes / sizeof(float)),
            Name{ "HairSharedDynamicBuffer" }, Name{ "m_skinnedHairSharedBuffer" }, 0, 0
        );
        m_buffer = Hair::UtilityClass::CreateBuffer("Hair Gem", descriptor, nullptr);
    }

    void SharedBuffer::CreateBufferAsset()
    {
        // Create the shared buffer pool
        {
            auto bufferPoolDesc = AZStd::make_unique<RHI::BufferPoolDescriptor>();
            // Output buffers are both written to during skinning and used as input assembly buffers
            bufferPoolDesc->m_bindFlags = RHI::BufferBindFlags::ShaderReadWrite | RHI::BufferBindFlags::Indirect;
            bufferPoolDesc->m_heapMemoryLevel = RHI::HeapMemoryLevel::Device;
            bufferPoolDesc->m_hostMemoryAccess = RHI::HostMemoryAccess::Write;

            RPI::ResourcePoolAssetCreator creator;
            creator.Begin(Uuid::CreateRandom());
            creator.SetPoolDescriptor(AZStd::move(bufferPoolDesc));
            creator.SetPoolName("SharedBufferPool");
            creator.End(m_bufferPoolAsset);
        }

        // Create the shared buffer
        {
            RPI::BufferAssetCreator creator;
            Uuid uuid = Uuid::CreateRandom();
            creator.Begin(uuid);
            creator.SetBufferName(m_bufferName);
            creator.SetPoolAsset(m_bufferPoolAsset);

            RHI::BufferDescriptor bufferDescriptor;
            bufferDescriptor.m_bindFlags = RHI::BufferBindFlags::ShaderReadWrite | RHI::BufferBindFlags::Indirect;
            bufferDescriptor.m_byteCount = m_sizeInBytes;
            bufferDescriptor.m_alignment = m_alignment;
            creator.SetBuffer(nullptr, 0, bufferDescriptor);

            RHI::BufferViewDescriptor viewDescriptor;
            viewDescriptor.m_elementFormat = RHI::Format::Unknown;

            // [To Do] - set this as AZ::Vector4 for offset approach shader code optimization
            viewDescriptor.m_elementSize = sizeof(float);
            viewDescriptor.m_elementCount = aznumeric_cast<uint32_t>(m_sizeInBytes) / sizeof(float);
            viewDescriptor.m_elementOffset = 0;
            creator.SetBufferViewDescriptor(viewDescriptor);

            creator.End(m_bufferAsset);
        }
    }

    void SharedBuffer::Init(AZStd::string bufferName, AZStd::vector<SrgBufferDescriptor>& buffersDescriptors)
    {
        m_bufferName = bufferName;
        // m_sizeInBytes = 256u * (1024u * 1024u);
        //
        // [To Do] replace this with max size request for allocation that can be given by the calling function
        // This has the following problems:
        //  1. The need to have this aggregated size in advance
        //  2. The size might grow dynamically between frames
        //  3. Due to having several stream buffers (position, tangent, structured), alignment padding
        //      size calculation must be added.
        // Requirement: the buffer already has an assert on allocation beyond the memory.  In the future it should
        // support greedy memory allocation when memory has reached its end.  This must not invalidate the buffer during
        // the current frame, hence allocation of second buffer, fence and a copy must take place.

        CalculateAlignment(buffersDescriptors);

        InitAllocator();

        CreateBuffer();

        SystemTickBus::Handler::BusConnect();
    }

    AZStd::intrusive_ptr<HairSharedBufferAllocation> SharedBuffer::Allocate(size_t byteCount)
    {
        RHI::VirtualAddress result;
        {
            AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
            result = m_freeListAllocator.Allocate(byteCount, m_alignment);
        }

        if (result.IsValid())
        {
            return aznew HairSharedBufferAllocation(result);
        }

        return nullptr;
    }

    void SharedBuffer::DeAllocate(RHI::VirtualAddress allocation)
    {
        if (allocation.IsValid())
        {
            {
                AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                m_freeListAllocator.DeAllocate(allocation);
            }

            m_memoryWasFreed = true;
            m_broadcastMemoryAvailableEvent = true;
        }
    }

    void SharedBuffer::DeAllocateNoSignal(RHI::VirtualAddress allocation)
    {
        if (allocation.IsValid())
        {
            {
                AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                m_freeListAllocator.DeAllocate(allocation);
            }
            m_memoryWasFreed = true;
        }
    }

    Data::Asset<RPI::BufferAsset> SharedBuffer::GetBufferAsset() const
    {
        return m_bufferAsset;
    }

    Data::Instance<RPI::Buffer> SharedBuffer::GetBuffer()
    {
        if (!m_buffer)
        {
            m_buffer = RPI::Buffer::FindOrCreate(m_bufferAsset);
        }
        return m_buffer;
    }

    //! Update buffer's content with sourceData at an offset of bufferByteOffset
    bool SharedBuffer::UpdateData(const void* sourceData, uint64_t sourceDataSizeInBytes, uint64_t bufferByteOffset)
    {
        AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
        if (m_buffer.get())
        {
            return m_buffer->UpdateData(sourceData, sourceDataSizeInBytes, bufferByteOffset);
        }
        AZ_Assert(false, "SharedBuffer error in data allocation - the buffer doesn't exist yet");
        return false;
    }

    void SharedBuffer::OnSystemTick()
    {
        GarbageCollect();
    }

    void SharedBuffer::GarbageCollect()
    {
        if (m_memoryWasFreed)
        {
            m_memoryWasFreed = false;
            {
                AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                m_freeListAllocator.GarbageCollect();
            }
            if (m_broadcastMemoryAvailableEvent)
            {
                SharedBufferNotificationBus::Broadcast(&SharedBufferNotificationBus::Events::OnSharedBufferMemoryAvailable);
                m_broadcastMemoryAvailableEvent = false;
            }
        }
    }

    //! Utility function to create a resource view of different type than the shared buffer data.
    //! Since this class is sub-buffer container, this method should be used after creating
    //!  a new allocation to be used as a sub-buffer.
    //! Notice the alignment required according to the element size - this might need
    RHI::BufferViewDescriptor SharedBuffer::CreateResourceViewWithDifferentFormat(
        uint32_t offsetInBytes, uint32_t elementCount, uint32_t elementSize,
        RHI::Format format, RHI::BufferBindFlags overrideBindFlags)
    {
        RHI::BufferViewDescriptor viewDescriptor;

        // In the following line I use the element size and not the size based of the
        // element format since in the more interesting case of structured buffer, the
        // size will result in an error.
        uint32_t elementOffset = offsetInBytes / elementSize;
        viewDescriptor.m_elementOffset = elementOffset;
        viewDescriptor.m_elementCount = elementCount;
        viewDescriptor.m_elementFormat = format;
        viewDescriptor.m_elementSize = elementSize;
        viewDescriptor.m_overrideBindFlags = overrideBindFlags;
        return viewDescriptor;
    }

} // namespace AZ::Render