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 <RHI/AliasedHeap.h>
  #include <RHI/AliasingBarrierTracker.h>
  #include <RHI/Image.h>
  #include <RHI/Buffer.h>
  #include <RHI/Device.h>
  #include <RHI/Scope.h>
  #include <Atom/RHI.Reflect/TransientBufferDescriptor.h>
  #include <Atom/RHI.Reflect/TransientImageDescriptor.h>
  #include <Atom/RHI/MemoryStatisticsBuilder.h>
  #include <AzCore/Debug/EventTrace.h>
  #include <AzCore/std/sort.h>
  
  namespace AZ
  {
      namespace DX12
      {
          static size_t CalculateHeapAlignment(D3D12_HEAP_FLAGS heapFlags)
          {
              /**
              * Heap alignment is the alignment of the actual heap we are allocating, not the base alignment of
              * of sub-allocations from the heap. That is confusing in the D3D12 docs. The heap itself is
              * required to be 4MB aligned if it holds MSAA textures. Therefore, this simple metric just
              * forces 4MB alignment of the heap for all textures, because our chances of having an MSAA target
              * across the whole frame is high, and the amount of internal fragmentation is low relative to the full
              * heap size.
              */
              if (D3D12_HEAP_FLAG_ALLOW_ALL_BUFFERS_AND_TEXTURES == heapFlags ||
                  D3D12_HEAP_FLAG_ALLOW_ONLY_RT_DS_TEXTURES == heapFlags ||
                  D3D12_HEAP_FLAG_ALLOW_ONLY_NON_RT_DS_TEXTURES == heapFlags)
              {
                  return D3D12_DEFAULT_MSAA_RESOURCE_PLACEMENT_ALIGNMENT;
              }
              else
              {
                  return D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT;
              }
          }
  
          RHI::Ptr<AliasedHeap> AliasedHeap::Create()
          {
              return aznew AliasedHeap;
          }
  
          AZStd::unique_ptr<RHI::AliasingBarrierTracker> AliasedHeap::CreateBarrierTrackerInternal()
          {
              return AZStd::make_unique<AliasingBarrierTracker>();
          }
  
          RHI::ResultCode AliasedHeap::InitInternal(
              RHI::Device& rhiDevice,
              const RHI::AliasedHeapDescriptor& descriptor)
          {
              Device& device = static_cast<Device&>(rhiDevice);
              m_descriptor = static_cast<const Descriptor&>(descriptor);
              
              D3D12_HEAP_DESC heapDesc;
              heapDesc.Alignment = CalculateHeapAlignment(m_descriptor.m_heapFlags);
              // Even the dx12 document said 'but non-aligned SizeInBytes is also supported', but it may lead to TDR issue with some graphics cards
              // Such as nvidia 2070, 2080
              heapDesc.SizeInBytes = RHI::AlignUp<size_t>(descriptor.m_budgetInBytes, heapDesc.Alignment);
              heapDesc.Properties = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT);
              heapDesc.Flags = m_descriptor.m_heapFlags;
  
              RHI::HeapMemoryUsage& heapMemoryUsage = m_memoryUsage.GetHeapMemoryUsage(RHI::HeapMemoryLevel::Device);
              heapMemoryUsage.m_residentInBytes = heapDesc.SizeInBytes;
              heapMemoryUsage.m_reservedInBytes = heapDesc.SizeInBytes;
  
              Microsoft::WRL::ComPtr<ID3D12Heap> heap;
              AssertSuccess(device.GetDevice()->CreateHeap(&heapDesc, IID_GRAPHICS_PPV_ARGS(heap.GetAddressOf())));
              m_heap = heap.Get();
              
              return m_heap ? RHI::ResultCode::Success : RHI::ResultCode::Fail;
          }
  
          void AliasedHeap::ShutdownInternal()
          {
              m_heap = nullptr;
              Base::ShutdownInternal();
          }
  
          void AliasedHeap::ShutdownResourceInternal(RHI::Resource& resource)
          {
              Device& device = GetDX12RHIDevice();
              if (Buffer* buffer = azrtti_cast <Buffer*>(&resource))
              {
                  device.QueueForRelease(buffer->GetMemoryView());
                  buffer->m_memoryView = {};
              }
              else if (Image* image = azrtti_cast < Image * > (&resource))
              {
                  device.QueueForRelease(image->GetMemoryView());
                  image->m_memoryView = {};
              }
          }
  
          RHI::ResultCode AliasedHeap::InitBufferInternal(const RHI::BufferInitRequest& request, size_t heapOffset)
          {
              const RHI::BufferDescriptor& descriptor = request.m_descriptor;
              Buffer* buffer = static_cast<Buffer*>(request.m_buffer);
  
              MemoryView memoryView =
                  GetDX12RHIDevice().CreateBufferPlaced(
                      descriptor,
                      D3D12_RESOURCE_STATE_COMMON,
                      m_heap.get(),
                      heapOffset);
  
              if (!memoryView.IsValid())
              {
                  RHI::ResultCode::OutOfMemory;
              }
  
              buffer->SetDescriptor(descriptor);
              buffer->m_memoryView = BufferMemoryView{ AZStd::move(memoryView), BufferMemoryType::Unique };
              return RHI::ResultCode::Success;
          }
  
          RHI::ResultCode AliasedHeap::InitImageInternal(const RHI::ImageInitRequest& request, size_t heapOffset)
          {
              const RHI::ImageDescriptor& descriptor = request.m_descriptor;
              Image* image = static_cast<Image*>(request.m_image);
  
              image->SetDescriptor(descriptor);
  
              /// Ownership is managed by the cache.
              MemoryView memoryView =
                  GetDX12RHIDevice().CreateImagePlaced(
                      descriptor,
                      request.m_optimizedClearValue,
                      image->GetInitialResourceState(),
                      m_heap.get(),
                      heapOffset);
  
              image->GenerateSubresourceLayouts();
              image->m_residentSizeInBytes = memoryView.GetSize();
              image->m_memoryView = AZStd::move(memoryView);
              return RHI::ResultCode::Success;
          }
  
          Device& AliasedHeap::GetDX12RHIDevice()
          {
              return static_cast<Device&>(Base::GetDevice());
          }
  
          const AliasedHeap::Descriptor& AliasedHeap::GetDescriptor() const
          {
              return m_descriptor;
          }
      }
  }
  