Feat: Support rwStorage buffer

Engine/Source/RHI-DirectX12/Include/RHI/DirectX12/Common.h

@@ -239,12 +239,6 @@ namespace RHI::DirectX12 {
         ECIMPL_ITEM(IndexFormat::uint32, DXGI_FORMAT_R32_UINT)
     ECIMPL_END(DXGI_FORMAT)
 
-    ECIMPL_BEGIN(StorageFormat, DXGI_FORMAT)
-        ECIMPL_ITEM(StorageFormat::float32, DXGI_FORMAT_R32_FLOAT)
-        ECIMPL_ITEM(StorageFormat::uint32, DXGI_FORMAT_R32_UINT)
-        ECIMPL_ITEM(StorageFormat::sint32, DXGI_FORMAT_R32_SINT)
-    ECIMPL_END(DXGI_FORMAT)
-
     ECIMPL_BEGIN(VertexStepMode, D3D12_INPUT_CLASSIFICATION)
         ECIMPL_ITEM(VertexStepMode::perVertex, D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA)
         ECIMPL_ITEM(VertexStepMode::perInstance, D3D12_INPUT_CLASSIFICATION_PER_INSTANCE_DATA)
@@ -256,7 +250,8 @@ namespace RHI::DirectX12 {
         ECIMPL_ITEM(BufferState::copySrc, D3D12_RESOURCE_STATE_COPY_SOURCE)
         ECIMPL_ITEM(BufferState::copyDst, D3D12_RESOURCE_STATE_COPY_DEST)
         ECIMPL_ITEM(BufferState::shaderReadOnly, D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE | D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE)
-        ECIMPL_ITEM(BufferState::storage, D3D12_RESOURCE_STATE_UNORDERED_ACCESS)
+        ECIMPL_ITEM(BufferState::storage, D3D12_RESOURCE_STATE_COMMON)
+        ECIMPL_ITEM(BufferState::rwStorage, D3D12_RESOURCE_STATE_UNORDERED_ACCESS)
     ECIMPL_END(D3D12_RESOURCE_STATES)
 
     ECIMPL_BEGIN(TextureDimension, D3D12_RESOURCE_DIMENSION)

Engine/Source/RHI-DirectX12/Src/BindGroup.cpp

@@ -12,11 +12,14 @@
 namespace RHI::DirectX12 {
     static CD3DX12_CPU_DESCRIPTOR_HANDLE GetDescriptorCpuHandle(const BindGroupEntry& entry)
     {
-        if (entry.binding.type == BindingType::uniformBuffer || entry.binding.type == BindingType::storageBuffer) {
+        if (entry.binding.type == BindingType::uniformBuffer ||
+            entry.binding.type == BindingType::storageBuffer ||
+            entry.binding.type == BindingType::rwStorageBuffer) {
             const auto* bufferView = static_cast<DX12BufferView*>(std::get<BufferView*>(entry.entity));
             return bufferView->GetNativeCpuDescriptorHandle();
         }
-        if (entry.binding.type == BindingType::texture || entry.binding.type == BindingType::storageTexture) {
+        if (entry.binding.type == BindingType::texture ||
+            entry.binding.type == BindingType::storageTexture) {
             const auto* textureView = static_cast<DX12TextureView*>(std::get<TextureView*>(entry.entity));
             return textureView->GetNativeCpuDescriptorHandle();
         }

Engine/Source/RHI-DirectX12/Src/Buffer.cpp

@@ -33,7 +33,7 @@ namespace RHI::DirectX12 {
     static D3D12_RESOURCE_FLAGS GetDX12ResourceFlag(const BufferUsageFlags flag)
     {
         static std::unordered_map<BufferUsageFlags, D3D12_RESOURCE_FLAGS> rules = {
-            { BufferUsageBits::storage, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS },
+            { BufferUsageBits::rwStorage, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS },
         };
         static D3D12_RESOURCE_FLAGS fallback = D3D12_RESOURCE_FLAG_NONE;
 

Engine/Source/RHI-DirectX12/Src/BufferView.cpp

@@ -9,19 +9,6 @@
 #include <RHI/DirectX12/BufferView.h>
 #include <RHI/DirectX12/Common.h>
 
-namespace RHI::DirectX12 {
-    static uint32_t GetStrideOfStorageBuffer(StorageFormat format)
-    {
-        if (format == StorageFormat::float32 ||
-            format == StorageFormat::sint32 ||
-            format == StorageFormat::uint32) {
-            return 4;
-        }
-
-        return -1;
-    }
-}
-
 namespace RHI::DirectX12 {
     DX12BufferView::DX12BufferView(DX12Buffer& inBuffer, const BufferViewCreateInfo& inCreateInfo)
         : BufferView(inCreateInfo), buffer(inBuffer)
@@ -63,12 +50,27 @@ namespace RHI::DirectX12 {
             Assert((bufferUsages & BufferUsageBits::storage) != 0);
             auto storageViewInfo = std::get<StorageBufferViewInfo>(inCreateInfo.extend);
 
-            // TODO: check the uav typed load
+            D3D12_SHADER_RESOURCE_VIEW_DESC desc {};
+            desc.Format = DXGI_FORMAT_UNKNOWN;
+            desc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER;
+            desc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
+            desc.Buffer.FirstElement = inCreateInfo.offset;
+            desc.Buffer.NumElements = inCreateInfo.size / storageViewInfo.stride;
+            desc.Buffer.StructureByteStride = storageViewInfo.stride;
+
+            nativeView = buffer.GetDevice().AllocateCbvSrvUavDescriptor();
+            buffer.GetDevice().GetNative()->CreateShaderResourceView(buffer.GetNative(), &desc, std::get<Common::UniqueRef<DescriptorAllocation>>(nativeView)->GetCpuHandle());
+        } else if (inCreateInfo.type == BufferViewType::rwStorageBinding) {
+            Assert((bufferUsages & BufferUsageBits::rwStorage) != 0);
+            auto storageViewInfo = std::get<StorageBufferViewInfo>(inCreateInfo.extend);
+
+            // TODO: check the uav typed load when it is necessary
             D3D12_UNORDERED_ACCESS_VIEW_DESC desc {};
-            desc.Format = EnumCast<StorageFormat, DXGI_FORMAT>(storageViewInfo.format);
+            desc.Format = DXGI_FORMAT_UNKNOWN;
             desc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER;
             desc.Buffer.FirstElement = inCreateInfo.offset;
-            desc.Buffer.NumElements = inCreateInfo.size / GetStrideOfStorageBuffer(storageViewInfo.format);
+            desc.Buffer.NumElements = inCreateInfo.size / storageViewInfo.stride;
+            desc.Buffer.StructureByteStride = storageViewInfo.stride;
 
             nativeView = buffer.GetDevice().AllocateCbvSrvUavDescriptor();
             buffer.GetDevice().GetNative()->CreateUnorderedAccessView(buffer.GetNative(), nullptr, &desc, std::get<Common::UniqueRef<DescriptorAllocation>>(nativeView)->GetCpuHandle());

Engine/Source/RHI-Vulkan/Include/RHI/Vulkan/Common.h

@@ -189,6 +189,7 @@ namespace RHI::Vulkan {
     ECIMPL_BEGIN(BindingType, VkDescriptorType)
         ECIMPL_ITEM(BindingType::uniformBuffer,  VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER)
         ECIMPL_ITEM(BindingType::storageBuffer,  VK_DESCRIPTOR_TYPE_STORAGE_BUFFER)
+        ECIMPL_ITEM(BindingType::rwStorageBuffer,  VK_DESCRIPTOR_TYPE_STORAGE_BUFFER)
         ECIMPL_ITEM(BindingType::sampler,        VK_DESCRIPTOR_TYPE_SAMPLER)
         ECIMPL_ITEM(BindingType::texture,        VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE)
         ECIMPL_ITEM(BindingType::storageTexture, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE)
@@ -286,6 +287,7 @@ namespace RHI::Vulkan {
         FCIMPL_ITEM(BufferUsageBits::vertex,   VK_BUFFER_USAGE_VERTEX_BUFFER_BIT)
         FCIMPL_ITEM(BufferUsageBits::uniform,  VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT)
         FCIMPL_ITEM(BufferUsageBits::storage,  VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
+        FCIMPL_ITEM(BufferUsageBits::rwStorage,  VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
         FCIMPL_ITEM(BufferUsageBits::indirect, VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT)
     FCIMPL_END(VkBufferUsageFlagBits)
 

Engine/Source/RHI-Vulkan/Src/BindGroup.cpp

@@ -81,9 +81,13 @@ namespace RHI::Vulkan {
         int bufferInfosNum = 0;
         for (int i = 0; i < entryCount; i++) {
             const auto& entry = inCreateInfo.entries[i];
-            if (entry.binding.type == BindingType::uniformBuffer || entry.binding.type == BindingType::storageBuffer) {
+            if (entry.binding.type == BindingType::uniformBuffer
+                || entry.binding.type == BindingType::storageBuffer
+                || entry.binding.type == BindingType::rwStorageBuffer) {
                 bufferInfosNum++;
-            } else if (entry.binding.type == BindingType::sampler || entry.binding.type == BindingType::texture ||entry.binding.type == BindingType::storageTexture) {
+            } else if (entry.binding.type == BindingType::sampler
+                || entry.binding.type == BindingType::texture
+                ||entry.binding.type == BindingType::storageTexture) {
                 imageInfosNum++;
             }
         }
@@ -99,7 +103,9 @@ namespace RHI::Vulkan {
             descriptorWrites[i].descriptorCount = 1;
             descriptorWrites[i].descriptorType = EnumCast<BindingType, VkDescriptorType>(entry.binding.type);
 
-            if (entry.binding.type == BindingType::uniformBuffer || entry.binding.type == BindingType::storageBuffer) {
+            if (entry.binding.type == BindingType::uniformBuffer
+                || entry.binding.type == BindingType::storageBuffer
+                || entry.binding.type == BindingType::rwStorageBuffer) {
                 auto* bufferView = static_cast<VulkanBufferView*>(std::get<BufferView*>(entry.entity));
 
                 bufferInfos.emplace_back();

Engine/Source/RHI-Vulkan/Src/CommandRecorder.cpp

@@ -26,7 +26,8 @@ namespace RHI::Vulkan {
             { BufferState::copySrc, VK_ACCESS_TRANSFER_READ_BIT },
             { BufferState::copyDst, VK_ACCESS_TRANSFER_WRITE_BIT },
             { BufferState::shaderReadOnly, VK_ACCESS_SHADER_READ_BIT },
-            { BufferState::storage, VK_ACCESS_SHADER_WRITE_BIT }
+            { BufferState::storage, VK_ACCESS_SHADER_READ_BIT },
+            { BufferState::rwStorage, VK_ACCESS_SHADER_WRITE_BIT }
         };
         return map.at(inState);
     }
@@ -39,7 +40,9 @@ namespace RHI::Vulkan {
             { BufferState::copySrc, VK_PIPELINE_STAGE_TRANSFER_BIT },
             { BufferState::copyDst, VK_PIPELINE_STAGE_TRANSFER_BIT },
             { BufferState::shaderReadOnly, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT },
-            { BufferState::storage, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT }
+            { BufferState::storage, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT },
+            { BufferState::rwStorage, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT }
+
         };
         return map.at(inState);
     }
@@ -52,7 +55,8 @@ namespace RHI::Vulkan {
             { BufferState::copySrc, VK_PIPELINE_STAGE_TRANSFER_BIT },
             { BufferState::copyDst, VK_PIPELINE_STAGE_TRANSFER_BIT },
             { BufferState::shaderReadOnly, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT },
-            { BufferState::storage, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT }
+            { BufferState::storage, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT },
+            { BufferState::rwStorage, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT }
         };
         return map.at(inState);
     }

Engine/Source/RHI/Include/RHI/BufferView.h

@@ -25,7 +25,9 @@ namespace RHI {
     };
 
     struct StorageBufferViewInfo {
-        StorageFormat format;
+        uint32_t stride;
+
+        explicit StorageBufferViewInfo(uint32_t inStride = 0);
     };
 
     struct BufferViewCreateInfo {
@@ -45,7 +47,7 @@ namespace RHI {
         BufferViewCreateInfo& SetSize(uint32_t inSize);
         BufferViewCreateInfo& SetExtendVertex(uint32_t inStride);
         BufferViewCreateInfo& SetExtendIndex(IndexFormat inFormat);
-        BufferViewCreateInfo& SetExtendStorage(StorageFormat inFormat);
+        BufferViewCreateInfo& SetExtendStorage(uint32_t inStride);
 
         size_t Hash() const;
     };

Engine/Source/RHI/Include/RHI/Common.h

@@ -196,6 +196,7 @@ namespace RHI {
         index,
         uniformBinding,
         storageBinding,
+        rwStorageBinding,
         max
     };
 
@@ -235,6 +236,7 @@ namespace RHI {
     enum class BindingType : EnumType {
         uniformBuffer,
         storageBuffer,
+        rwStorageBuffer,
         sampler,
         texture,
         storageTexture,
@@ -294,14 +296,6 @@ namespace RHI {
         max
     };
 
-    // TODO: Support more format
-    enum class StorageFormat: EnumType {
-        float32,
-        uint32,
-        sint32,
-        max
-    };
-
     enum class FrontFace : EnumType {
         ccw,
         cw,
@@ -390,6 +384,7 @@ namespace RHI {
         copyDst,
         shaderReadOnly,
         storage,
+        rwStorage,
         max
     };
 
@@ -496,8 +491,9 @@ namespace RHI {
         vertex       = 0x20,
         uniform      = 0x40,
         storage      = 0x80,
-        indirect     = 0x100,
-        queryResolve = 0x200,
+        rwStorage    = 0x100,
+        indirect     = 0x200,
+        queryResolve = 0x400,
         max
     };
     using BufferUsageFlags = Flags<BufferUsageBits>;

Engine/Source/RHI/Src/BufferView.cpp

@@ -15,6 +15,11 @@ namespace RHI {
     {
     }
 
+    StorageBufferViewInfo::StorageBufferViewInfo(uint32_t inStride)
+        : stride(inStride)
+    {
+    }
+
     BufferViewCreateInfo::BufferViewCreateInfo(
         const BufferViewType inType,
         const uint32_t inSize,
@@ -57,9 +62,9 @@ namespace RHI {
         return *this;
     }
 
-    BufferViewCreateInfo& BufferViewCreateInfo::SetExtendStorage(StorageFormat inFormat)
+    BufferViewCreateInfo& BufferViewCreateInfo::SetExtendStorage(uint32_t inStride)
     {
-        extend = StorageBufferViewInfo { inFormat };
+        extend = StorageBufferViewInfo { inStride };
         return *this;
     }
 

Engine/Source/Render/Src/ShaderCompiler.cpp

@@ -41,7 +41,7 @@ namespace Render {
             { D3D_SIT_SAMPLER, RHI::BindingType::sampler },
             { D3D_SIT_UAV_RWTYPED, RHI::BindingType::storageTexture },
             { D3D_SIT_STRUCTURED, RHI::BindingType::storageBuffer },
-            { D3D_SIT_UAV_RWSTRUCTURED, RHI::BindingType::storageBuffer }
+            { D3D_SIT_UAV_RWSTRUCTURED, RHI::BindingType::rwStorageBuffer }
         };
         return map.at(type);
     }
@@ -53,7 +53,7 @@ namespace Render {
             { D3D_SIT_TEXTURE, RHI::HlslBindingRangeType::texture },
             { D3D_SIT_SAMPLER, RHI::HlslBindingRangeType::sampler },
             { D3D_SIT_UAV_RWTYPED, RHI::HlslBindingRangeType::unorderedAccess },
-            { D3D_SIT_STRUCTURED, RHI::HlslBindingRangeType::unorderedAccess },
+            { D3D_SIT_STRUCTURED, RHI::HlslBindingRangeType::texture },
             { D3D_SIT_UAV_RWSTRUCTURED, RHI::HlslBindingRangeType::unorderedAccess }
         };
         return map.at(type);

Sample/RHI-ParallelCompute/Compute.hlsl

@@ -1,14 +1,18 @@
 #include <Platform.h>
 
-VkBinding(0, 0) cbuffer input : register(b0)
-{
-    float4 v[16];
+// spir-v treat vec2 and vec4 as built-in type?
+struct Data {
+    float4 v1;
+    float4 v2;
 };
 
-VkBinding(1, 0) RWStructuredBuffer<float4> output : register(u0);
+VkBinding(0, 0) StructuredBuffer<Data> input : register(t0);
 
-[numthreads(16, 1, 1)]
-void CSMain(int id : SV_DispatchThreadID) {
-    output[id.x] = v[id.x] * v[id.x];
+VkBinding(1, 0) RWStructuredBuffer<Data> output : register(u0);
+
+[numthreads(32, 1, 1)]
+void CSMain(uint3 id : SV_DispatchThreadID) {
+    output[id.x].v1 = input[id.x].v1 * input[id.x].v1;
+    output[id.x].v2 = input[id.x].v2 * input[id.x].v2;
 }