text: /*
   * 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_Metal_precompiled.h"
  
  #include <AzCore/Casting/lossy_cast.h>
  #include <RHI/BufferMemoryAllocator.h>
  #include <RHI/Conversions.h>
  #include <RHI/Device.h>
  
  namespace AZ
  {
      namespace Metal
      {
          void BufferMemoryAllocator::Init(const Descriptor& descriptor)
          {
              m_descriptor = descriptor;
  
              m_usePageAllocator = false;
  
              if (!RHI::CheckBitsAny(descriptor.m_bindFlags, RHI::BufferBindFlags::ShaderWrite | RHI::BufferBindFlags::CopyWrite))
              {
                  m_usePageAllocator = true;
  
                  const RHI::HeapMemoryUsage* heapMemoryUsage = descriptor.m_getHeapMemoryUsageFunction();
  
                  const size_t budgetInBytes = heapMemoryUsage->m_budgetInBytes;
                  if (budgetInBytes)
                  {
                      // The buffer page size should not exceed the budget.
                      m_descriptor.m_pageSizeInBytes = static_cast<uint32_t>(AZStd::min<size_t>(budgetInBytes, m_descriptor.m_pageSizeInBytes));
                  }
  
                  m_pageAllocator.Init(m_descriptor);
  
                  // Constant buffers have stricter alignment requirements.
                  m_subAllocationAlignment = RHI::CheckBitsAny(m_descriptor.m_bindFlags, RHI::BufferBindFlags::Constant)
                      ? Alignment::Constant
                      : Alignment::Buffer;
  
                  MemoryFreeListSubAllocator::Descriptor pageAllocatorDescriptor;
                  pageAllocatorDescriptor.m_alignmentInBytes = m_subAllocationAlignment;
                  pageAllocatorDescriptor.m_garbageCollectLatency = RHI::Limits::Device::FrameCountMax;
                  pageAllocatorDescriptor.m_inactivePageCycles = 1;
                  m_subAllocator.Init(pageAllocatorDescriptor, m_pageAllocator);
              }
          }
  
          void BufferMemoryAllocator::Shutdown()
          {
              if (m_usePageAllocator)
              {
                  m_subAllocator.Shutdown();
                  m_pageAllocator.Shutdown();
              }
          }
  
          void BufferMemoryAllocator::GarbageCollect()
          {
              {
                  AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                  m_subAllocator.GarbageCollect();
              }
  
              m_pageAllocator.Collect();
          }
  
          BufferMemoryView BufferMemoryAllocator::Allocate(size_t sizeInBytes)
          {
              BufferMemoryView bufferMemoryView;
  
              // First attempt to sub-allocate a buffer from the sub-allocator.
              if (m_usePageAllocator)
              {
                  AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
  
                  bufferMemoryView = BufferMemoryView(MemoryView(m_subAllocator.Allocate(sizeInBytes, m_subAllocationAlignment)), BufferMemoryType::SubAllocated);
              }
  
              // Next, try a unique buffer allocation.
              if (!bufferMemoryView.IsValid())
              {
                  RHI::BufferDescriptor bufferDescriptor;
                  bufferDescriptor.m_byteCount = azlossy_caster(sizeInBytes);
                  bufferDescriptor.m_bindFlags = m_descriptor.m_bindFlags;
                  
                  AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                  bufferMemoryView = AllocateUnique(bufferDescriptor);
              }
  
              return bufferMemoryView;
          }
  
          void BufferMemoryAllocator::DeAllocate(const BufferMemoryView& memoryView)
          {
              switch (memoryView.GetType())
              {
              case BufferMemoryType::SubAllocated:
                  {
                      AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                      m_subAllocator.DeAllocate(memoryView.m_memoryAllocation);
                  }
                  break;
  
              case BufferMemoryType::Unique:
                  {
                      AZStd::lock_guard<AZStd::mutex> lock(m_allocatorMutex);
                      DeAllocateUnique(memoryView);
                  }
                  break;
              }
          }
  
          BufferMemoryView BufferMemoryAllocator::AllocateUnique(const RHI::BufferDescriptor& bufferDescriptor)
          {
              AZ_TRACE_METHOD();
              const size_t alignedSize = RHI::AlignUp(bufferDescriptor.m_byteCount, Alignment::Buffer);
              
              RHI::HeapMemoryUsage& heapMemoryUsage = *m_descriptor.m_getHeapMemoryUsageFunction();
              if (!heapMemoryUsage.TryReserveMemory(alignedSize))
              {
                  return BufferMemoryView();
              }
              
              MemoryView memoryView = m_descriptor.m_device->CreateBufferCommitted(bufferDescriptor, m_descriptor.m_heapMemoryLevel);
              if (memoryView.IsValid())
              {
                  heapMemoryUsage.m_residentInBytes += alignedSize;
              }
              else
              {
                  heapMemoryUsage.m_reservedInBytes -= alignedSize;
              }
              return BufferMemoryView(AZStd::move(memoryView), BufferMemoryType::Unique);
          }
  
          void BufferMemoryAllocator::DeAllocateUnique(const BufferMemoryView& memoryView)
          {
              AZ_Assert(memoryView.GetType() == BufferMemoryType::Unique, "This call only supports unique BufferMemoryView allocations.");
              const size_t sizeInBytes = memoryView.GetSize();
              
              RHI::HeapMemoryUsage& heapHemoryUsage = *m_descriptor.m_getHeapMemoryUsageFunction();
              heapHemoryUsage.m_residentInBytes -= sizeInBytes;
              heapHemoryUsage.m_reservedInBytes -= sizeInBytes;
              
              m_descriptor.m_device->QueueForRelease(memoryView);
          }
      }
  }
  