From fecdc458a5f4b7f42789e7549ffa50a00d2c4ce1 Mon Sep 17 00:00:00 2001
From: Yuefeng Geng <maxwellgeng@outlook.com>
Date: Thu, 26 Sep 2024 13:35:59 +0800
Subject: [PATCH] use committed resource in Accel to avoid potential driver bug

---
 src/backends/dx/Resource/BottomAccel.cpp |  4 ++--
 src/backends/dx/Resource/TopAccel.cpp    | 19 +++++++++++--------
 src/backends/dx/Resource/TopAccel.h      |  3 ++-
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/backends/dx/Resource/BottomAccel.cpp b/src/backends/dx/Resource/BottomAccel.cpp
index e153dacca..b12a28bcf 100644
--- a/src/backends/dx/Resource/BottomAccel.cpp
+++ b/src/backends/dx/Resource/BottomAccel.cpp
@@ -129,7 +129,7 @@ size_t BottomAccel::PreProcessStates(
         accelBuffer = vstd::create_unique(new DefaultBuffer(
             device,
             CalcAlign(bottomLevelPrebuildInfo.ResultDataMaxSizeInBytes, 65536),
-            device->defaultAllocator.get(),
+            nullptr,
             D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE,
             false,
             "blas-accel-buffer"));
@@ -164,7 +164,7 @@ bool BottomAccel::CheckAccel(
     auto newAccelBuffer = vstd::create_unique(new DefaultBuffer(
         device,
         CalcAlign(compactSize, 65536),
-        device->defaultAllocator.get(),
+        nullptr,
         D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE,
         false,
         "blas-accel-buffer"));
diff --git a/src/backends/dx/Resource/TopAccel.cpp b/src/backends/dx/Resource/TopAccel.cpp
index 8d97c9246..c0498ef46 100644
--- a/src/backends/dx/Resource/TopAccel.cpp
+++ b/src/backends/dx/Resource/TopAccel.cpp
@@ -62,13 +62,14 @@ bool TopAccel::GenerateNewBuffer(
     char const *name,
     ResourceStateTracker &tracker,
     CommandBufferBuilder &builder,
-    vstd::unique_ptr<DefaultBuffer> &oldBuffer, size_t newSize, bool needCopy, D3D12_RESOURCE_STATES state) {
+    vstd::unique_ptr<DefaultBuffer> &oldBuffer, size_t newSize, bool needCopy, D3D12_RESOURCE_STATES state,
+    GpuAllocator *allocator) {
     if (!oldBuffer) {
         newSize = CalcAlign(newSize, 65536);
         oldBuffer = vstd::create_unique(new DefaultBuffer(
             device,
             newSize,
-            device->defaultAllocator.get(),
+            allocator,
             state,
             false,
             name));
@@ -79,7 +80,7 @@ bool TopAccel::GenerateNewBuffer(
         auto newBuffer = new DefaultBuffer(
             device,
             newSize,
-            device->defaultAllocator.get(),
+            allocator,
             state);
         if (needCopy) {
             tracker.RecordState(
@@ -125,7 +126,7 @@ void TopAccel::PreProcessInst(
     size_t instanceByteCount = size * sizeof(D3D12_RAYTRACING_INSTANCE_DESC);
     if (GenerateNewBuffer(
             "tlas-instance-buffer",
-            tracker, builder, instBuffer, instanceByteCount, true, tracker.ReadState(ResourceReadUsage::AccelBuildSrc))) {
+            tracker, builder, instBuffer, instanceByteCount, true, tracker.ReadState(ResourceReadUsage::AccelBuildSrc), device->defaultAllocator.get())) {
         input.InstanceDescs = instBuffer->GetAddress();
     }
 }
@@ -232,11 +233,11 @@ size_t TopAccel::PreProcess(
     size_t instanceByteCount = size * sizeof(D3D12_RAYTRACING_INSTANCE_DESC);
     if (GenerateNewBuffer(
             "tlas-instance-buffer",
-            tracker, builder, instBuffer, instanceByteCount, true, tracker.ReadState(ResourceReadUsage::AccelBuildSrc))) {
+            tracker, builder, instBuffer, instanceByteCount, true, tracker.ReadState(ResourceReadUsage::AccelBuildSrc), device->defaultAllocator.get())) {
         input.InstanceDescs = instBuffer->GetAddress();
     }
     device->device->GetRaytracingAccelerationStructurePrebuildInfo(&input, &topLevelPrebuildInfo);
-    if (GenerateNewBuffer("tlas-accel-buffer", tracker, builder, accelBuffer, topLevelPrebuildInfo.ResultDataMaxSizeInBytes, false, D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE)) {
+    if (GenerateNewBuffer("tlas-accel-buffer", tracker, builder, accelBuffer, topLevelPrebuildInfo.ResultDataMaxSizeInBytes, false, D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE, nullptr)) {
         update = false;
         topLevelBuildDesc.DestAccelerationStructureData = accelBuffer->GetAddress();
     }
@@ -347,8 +348,10 @@ bool TopAccel::CheckAccel(
     auto newAccelBuffer = vstd::create_unique(new DefaultBuffer(
         device,
         CalcAlign(compactSize, 65536),
-        device->defaultAllocator.get(),
-        D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE));
+        nullptr,
+        D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE,
+        false,
+        "tlas-accel-buffer"));
 
     builder.GetCB()->CmdList()->CopyRaytracingAccelerationStructure(
         newAccelBuffer->GetAddress(),
diff --git a/src/backends/dx/Resource/TopAccel.h b/src/backends/dx/Resource/TopAccel.h
index aa34ffbf0..fe2e6c9bb 100644
--- a/src/backends/dx/Resource/TopAccel.h
+++ b/src/backends/dx/Resource/TopAccel.h
@@ -46,7 +46,8 @@ class TopAccel : public Resource {
         char const *name,
         ResourceStateTracker &tracker,
         CommandBufferBuilder &builder,
-        vstd::unique_ptr<DefaultBuffer> &oldBuffer, size_t newSize, bool needCopy, D3D12_RESOURCE_STATES state);
+        vstd::unique_ptr<DefaultBuffer> &oldBuffer, size_t newSize, bool needCopy, D3D12_RESOURCE_STATES state,
+        GpuAllocator* allocator);
     void InitSetDesc(vstd::span<AccelBuildCommand::Modification const> const &modifications);
     void ProcessSetDesc();
     void ProcessSetMap();