Use batch to optimize sparse heap map

Learning-Ray-Tracing-From-Zero · Feb 19, 2025 · 501e047 · 501e047
1 parent acaf825
commit 501e047
Show file tree

Hide file tree

Showing 8 changed files with 123 additions and 50 deletions.
diff --git a/src/backends/dx/DXApi/LCDevice.cpp b/src/backends/dx/DXApi/LCDevice.cpp
@@ -732,25 +732,27 @@ void LCDevice::update_sparse_resources(
         LUISA_ERROR("sparse-texture update not allowed in Direct-Storage.");
     }
     auto &queuePtr = static_cast<LCCmdBuffer *>(queue)->queue;
+    UpdateTileTracker tile_tracker;
     for (auto &&i : update_cmds) {
         luisa::visit(
             [&]<typename T>(T const &t) {
                 if constexpr (std::is_same_v<T, SparseTextureMapOperation>) {
                     auto tex = reinterpret_cast<SparseTexture *>(i.handle);
-                    tex->AllocateTile(queuePtr.Queue(), t.start_tile, t.tile_count, t.mip_level, t.allocated_heap);
+                    tex->AllocateTile(t.start_tile, t.tile_count, t.mip_level, t.allocated_heap, &tile_tracker);
                 } else if constexpr (std::is_same_v<T, SparseBufferMapOperation>) {
                     auto buffer = reinterpret_cast<SparseBuffer *>(i.handle);
-                    buffer->AllocateTile(queuePtr.Queue(), t.start_tile, t.tile_count, t.allocated_heap);
+                    buffer->AllocateTile( t.start_tile, t.tile_count, t.allocated_heap, &tile_tracker);
                 } else if constexpr (std::is_same_v<T, SparseTextureUnMapOperation>) {
                     auto tex = reinterpret_cast<SparseTexture *>(i.handle);
-                    tex->DeAllocateTile(queuePtr.Queue(), t.start_tile, t.tile_count, t.mip_level);
+                    tex->DeAllocateTile(t.start_tile, t.tile_count, t.mip_level, &tile_tracker);
                 } else {
                     auto buffer = reinterpret_cast<SparseBuffer *>(i.handle);
-                    buffer->DeAllocateTile(queuePtr.Queue(), t.start_tile, t.tile_count);
+                    buffer->DeAllocateTile( t.start_tile, t.tile_count, &tile_tracker);
                 }
             },
             i.operations);
     }
+    tile_tracker.update(queuePtr.Queue(), D3D12_TILE_MAPPING_FLAG_NONE);
     queuePtr.Signal();
 }
 

diff --git a/src/backends/dx/DXApi/LCDevice.h b/src/backends/dx/DXApi/LCDevice.h
@@ -2,6 +2,7 @@
 #include <luisa/vstl/common.h>
 #include <luisa/runtime/device.h>
 #include <DXRuntime/Device.h>
+#include <DXRuntime/UpdateTileTracker.h>
 namespace lc::dx {
 using namespace luisa;
 using namespace luisa::compute;

diff --git a/src/backends/dx/DXRuntime/UpdateTileTracker.cpp b/src/backends/dx/DXRuntime/UpdateTileTracker.cpp
@@ -0,0 +1,63 @@
+#include "UpdateTileTracker.h"
+
+namespace lc::dx {
+void UpdateTileTracker::record(
+    ID3D12Heap *heap,
+    ID3D12Resource *resource,
+    D3D12_TILED_RESOURCE_COORDINATE const &ResourceRegionStartCoordinate,
+    D3D12_TILE_REGION_SIZE const &ResourceRegionSize,
+    D3D12_TILE_RANGE_FLAGS RangeFlag,
+    UINT HeapRangeStartOffset,
+    UINT RangeTileCount) {
+    std::pair<ID3D12Heap *, ID3D12Resource *> key{
+        heap,
+        resource};
+    auto iter = map.emplace(key);
+    auto &v = iter.value();
+    v.ResourceRegionStartCoordinates.emplace_back(ResourceRegionStartCoordinate);
+    v.ResourceRegionSizes.emplace_back(ResourceRegionSize);
+    v.RangeFlags.emplace_back(RangeFlag);
+    v.HeapRangeStartOffsets.emplace_back(HeapRangeStartOffset);
+    v.RangeTileCounts.emplace_back(RangeTileCount);
+}
+void UpdateTileTracker::update(
+    ID3D12CommandQueue *queue,
+    D3D12_TILE_MAPPING_FLAGS Flags) {
+    for (auto &kv : map) {
+        queue->UpdateTileMappings(
+            kv.first.second,
+            kv.second.ResourceRegionStartCoordinates.size(),
+            kv.second.ResourceRegionStartCoordinates.data(),
+            kv.second.ResourceRegionSizes.data(),
+            kv.first.first, kv.second.RangeFlags.size(),
+            kv.second.RangeFlags.data(),
+            kv.second.HeapRangeStartOffsets.data(),
+            kv.second.RangeTileCounts.data(),
+            Flags);
+    }
+    for (auto &kv : disp_map) {
+        queue->UpdateTileMappings(
+            kv.first,
+            kv.second.ResourceRegionStartCoordinates.size(),
+            kv.second.ResourceRegionStartCoordinates.data(),
+            kv.second.ResourceRegionSizes.data(),
+            nullptr,
+            1,
+            vstd::get_rval_ptr(D3D12_TILE_RANGE_FLAG_NULL),
+            nullptr,
+            nullptr,
+            Flags);
+    }
+    map.clear();
+    disp_map.clear();
+}
+void UpdateTileTracker::deallocate(
+    ID3D12Resource *resource,
+    D3D12_TILED_RESOURCE_COORDINATE const &ResourceRegionStartCoordinate,
+    D3D12_TILE_REGION_SIZE const &ResourceRegionSize) {
+    auto iter = disp_map.emplace(resource);
+    auto &v = iter.value();
+    v.ResourceRegionStartCoordinates.emplace_back(ResourceRegionStartCoordinate);
+    v.ResourceRegionSizes.emplace_back(ResourceRegionSize);
+}
+}// namespace lc::dx
diff --git a/src/backends/dx/DXRuntime/UpdateTileTracker.h b/src/backends/dx/DXRuntime/UpdateTileTracker.h
@@ -0,0 +1,36 @@
+#pragma once
+#include <d3d12.h>
+#include <luisa/vstl/common.h>
+namespace lc::dx {
+class UpdateTileTracker {
+    struct UpdateTileCommand {
+        vstd::fixed_vector<D3D12_TILED_RESOURCE_COORDINATE, 4>ResourceRegionStartCoordinates;
+        vstd::fixed_vector<D3D12_TILE_REGION_SIZE, 4>ResourceRegionSizes;
+        vstd::fixed_vector<D3D12_TILE_RANGE_FLAGS, 4>RangeFlags;
+        vstd::fixed_vector<UINT, 4>HeapRangeStartOffsets;
+        vstd::fixed_vector<UINT, 4>RangeTileCounts;
+    };
+    struct DisposeTileTracker {
+        vstd::fixed_vector<D3D12_TILED_RESOURCE_COORDINATE, 4>ResourceRegionStartCoordinates;
+        vstd::fixed_vector<D3D12_TILE_REGION_SIZE, 4>ResourceRegionSizes;
+    };
+    vstd::HashMap<std::pair<ID3D12Heap *, ID3D12Resource *>, UpdateTileCommand> map;
+    vstd::HashMap<ID3D12Resource *, DisposeTileTracker> disp_map;
+public:
+    void record(
+        ID3D12Heap *heap,
+        ID3D12Resource *resource,
+        D3D12_TILED_RESOURCE_COORDINATE const &ResourceRegionStartCoordinate,
+        D3D12_TILE_REGION_SIZE const &ResourceRegionSize,
+        D3D12_TILE_RANGE_FLAGS RangeFlag,
+        UINT HeapRangeStartOffset,
+        UINT RangeTileCount);
+    void deallocate(
+        ID3D12Resource *resource,
+        D3D12_TILED_RESOURCE_COORDINATE const &ResourceRegionStartCoordinate,
+        D3D12_TILE_REGION_SIZE const &ResourceRegionSize);
+    void update(
+        ID3D12CommandQueue *queue,
+        D3D12_TILE_MAPPING_FLAGS Flags);
+};
+}// namespace lcdx
diff --git a/src/backends/dx/Resource/SparseBuffer.cpp b/src/backends/dx/Resource/SparseBuffer.cpp
@@ -2,7 +2,7 @@
 #include <Resource/GpuAllocator.h>
 #include <luisa/core/logging.h>
 #include "SparseHeap.h"
-
+#include <DXRuntime/UpdateTileTracker.h>
 namespace lc::dx {
 SparseBuffer::SparseBuffer(
     Device *device,
@@ -32,7 +32,7 @@ vstd::optional<D3D12_UNORDERED_ACCESS_VIEW_DESC> SparseBuffer::GetColorUavDesc(u
 }
 SparseBuffer::~SparseBuffer() {
 }
-void SparseBuffer::DeAllocateTile(ID3D12CommandQueue *queue, uint coord, uint size) const {
+void SparseBuffer::DeAllocateTile(uint coord, uint size, UpdateTileTracker *tile_tracker) const {
     D3D12_TILED_RESOURCE_COORDINATE tileCoord{
         .X = coord,
         .Subresource = 0};
@@ -42,36 +42,19 @@ void SparseBuffer::DeAllocateTile(ID3D12CommandQueue *queue, uint coord, uint si
         .Width = static_cast<uint>(size),
         .Height = 1,
         .Depth = 1};
-    queue->UpdateTileMappings(
-        resource.Get(), 1,
-        &tileCoord,
-        &tileSize,
-        nullptr, 1,
-        vstd::get_rval_ptr(D3D12_TILE_RANGE_FLAG_NULL),
-        nullptr,
-        nullptr,
-        D3D12_TILE_MAPPING_FLAG_NONE);
+    tile_tracker->deallocate(resource.Get(), tileCoord, tileSize);
 }
-void SparseBuffer::AllocateTile(ID3D12CommandQueue *queue, uint coord, uint size, uint64 alloc) const {
+void SparseBuffer::AllocateTile(uint coord, uint size, uint64 alloc, UpdateTileTracker *tile_tracker) const {
     auto heap = reinterpret_cast<SparseHeap const *>(alloc);
     D3D12_TILED_RESOURCE_COORDINATE tileCoord{
         .X = coord,
         .Subresource = 0};
-    if (heap->size_bytes < size * D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES) [[unlikely]]{
+    if (heap->size_bytes < size * D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES) [[unlikely]] {
         LUISA_ERROR("Map size out of range. Required size: {}, heap size: {}", size * D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES, heap->size_bytes);
     }
     D3D12_TILE_REGION_SIZE tileSize{
         .NumTiles = size};
     uint tileOffset = heap->offset / D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES;
-    queue->UpdateTileMappings(
-        resource.Get(), 1,
-        &tileCoord,
-        &tileSize,
-        heap->heap,
-        1,
-        vstd::get_rval_ptr(D3D12_TILE_RANGE_FLAG_NONE),
-        &tileOffset,
-        &size,
-        D3D12_TILE_MAPPING_FLAG_NONE);
+    tile_tracker->record(heap->heap, resource.Get(), tileCoord, tileSize, D3D12_TILE_RANGE_FLAG_NONE, tileOffset, size);
 }
 }// namespace lc::dx
diff --git a/src/backends/dx/Resource/SparseBuffer.h b/src/backends/dx/Resource/SparseBuffer.h
@@ -1,6 +1,7 @@
 #pragma once
 #include <Resource/Buffer.h>
 namespace lc::dx {
+class UpdateTileTracker;
 class SparseBuffer final : public Buffer {
 private:
     ComPtr<ID3D12Resource> resource;
@@ -20,8 +21,8 @@ class SparseBuffer final : public Buffer {
     ID3D12Resource *GetResource() const override { return resource.Get(); }
     D3D12_GPU_VIRTUAL_ADDRESS GetAddress() const override { return resource->GetGPUVirtualAddress(); }
     uint64 GetByteSize() const override { return byteSize; }
-    void AllocateTile(ID3D12CommandQueue *queue, uint coord, uint size, uint64 alloc) const;
-    void DeAllocateTile(ID3D12CommandQueue *queue, uint coord, uint size) const;
+    void AllocateTile(uint coord, uint size, uint64 alloc, UpdateTileTracker* tile_tracker) const;
+    void DeAllocateTile(uint coord, uint size, UpdateTileTracker* tile_tracker) const;
     SparseBuffer(
         Device *device,
         uint64 byteSize,

diff --git a/src/backends/dx/Resource/SparseTexture.cpp b/src/backends/dx/Resource/SparseTexture.cpp
@@ -2,6 +2,7 @@
 #include <Resource/DescriptorHeap.h>
 #include <luisa/core/logging.h>
 #include "SparseHeap.h"
+#include <DXRuntime/UpdateTileTracker.h>
 namespace lc::dx {
 SparseTexture::SparseTexture(
     Device *device,
@@ -65,7 +66,7 @@ uint SparseTexture::GetGlobalSRVIndex(uint mipOffset) const {
 uint SparseTexture::GetGlobalUAVIndex(uint mipLevel) const {
     return GetGlobalUAVIndexBase(mipLevel, allocMtx, uavIdcs);
 }
-void SparseTexture::AllocateTile(ID3D12CommandQueue *queue, uint3 coord, uint3 size, uint mipLevel, uint64 alloc) const {
+void SparseTexture::AllocateTile(uint3 coord, uint3 size, uint mipLevel, uint64 alloc, UpdateTileTracker *tile_tracker) const {
     auto heap = reinterpret_cast<SparseHeap const *>(alloc);
     if (heap->size_bytes < (size.x * size.y * size.z) * D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES) [[unlikely]] {
         LUISA_ERROR("Map size out of range. Required size: {}, heap size: {}", (size.x * size.y * size.z) * D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES, heap->size_bytes);
@@ -83,17 +84,9 @@ void SparseTexture::AllocateTile(ID3D12CommandQueue *queue, uint3 coord, uint3 s
         .Depth = static_cast<uint16_t>(size.z)};
     uint rangeTileCount = tileSize.NumTiles;
     uint offsetTile = heap->offset / D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES;
-    queue->UpdateTileMappings(
-        resource.Get(), 1,
-        &tileCoord,
-        &tileSize,
-        heap->heap, 1,
-        vstd::get_rval_ptr(D3D12_TILE_RANGE_FLAG_NONE),
-        &offsetTile,
-        &rangeTileCount,
-        D3D12_TILE_MAPPING_FLAG_NONE);
+    tile_tracker->record(heap->heap, resource.Get(), tileCoord, tileSize, D3D12_TILE_RANGE_FLAG_NONE, offsetTile, rangeTileCount);
 }
-void SparseTexture::DeAllocateTile(ID3D12CommandQueue *queue, uint3 coord, uint3 size, uint mipLevel) const {
+void SparseTexture::DeAllocateTile(uint3 coord, uint3 size, uint mipLevel, UpdateTileTracker *tile_tracker) const {
     D3D12_TILED_RESOURCE_COORDINATE tileCoord{
         .X = coord.x,
         .Y = coord.y,
@@ -105,13 +98,6 @@ void SparseTexture::DeAllocateTile(ID3D12CommandQueue *queue, uint3 coord, uint3
         .Width = size.x,
         .Height = static_cast<uint16_t>(size.y),
         .Depth = static_cast<uint16_t>(size.z)};
-    queue->UpdateTileMappings(
-        resource.Get(), 1, &tileCoord,
-        &tileSize,
-        nullptr, 1,
-        vstd::get_rval_ptr(D3D12_TILE_RANGE_FLAG_NULL),
-        nullptr,
-        nullptr,
-        D3D12_TILE_MAPPING_FLAG_NONE);
+    tile_tracker->deallocate(resource.Get(), tileCoord, tileSize);
 }
 }// namespace lc::dx
diff --git a/src/backends/dx/Resource/SparseTexture.h b/src/backends/dx/Resource/SparseTexture.h
@@ -1,6 +1,7 @@
 #pragma once
 #include <Resource/TextureBase.h>
 namespace lc::dx {
+class UpdateTileTracker;
 class SparseTexture final : public TextureBase {
 private:
     ComPtr<ID3D12Resource> resource;
@@ -34,7 +35,7 @@ class SparseTexture final : public TextureBase {
     uint GetGlobalSRVIndex(uint mipOffset = 0) const override;
     D3D12_SHADER_RESOURCE_VIEW_DESC GetColorSrvDesc(uint mipOffset = 0) const override;
     uint GetGlobalUAVIndex(uint mipLevel) const override;
-    void AllocateTile(ID3D12CommandQueue *queue, uint3 coord, uint3 size, uint mipLevel, uint64 alloc) const;
-    void DeAllocateTile(ID3D12CommandQueue *queue, uint3 coord, uint3 size, uint mipLevel) const;
+    void AllocateTile(uint3 coord, uint3 size, uint mipLevel, uint64 alloc, UpdateTileTracker* tile_tracker) const;
+    void DeAllocateTile(uint3 coord, uint3 size, uint mipLevel, UpdateTileTracker* tile_tracker) const;
 };
 }// namespace lc::dx