Add ByteAddressBuffer/RWByteAddressBuffer support

andrejnau · andrejnau · commit cda46bd077b8 · 2025-12-13T12:46:17.000+01:00
diff --git a/assets/shaders/BufferViewTest/PixelShader.hlsl b/assets/shaders/BufferViewTest/PixelShader.hlsl
@@ -4,7 +4,7 @@ struct VsOutput {
 };
 
 static const uint kIndexCount = 64;
-static const uint kRows = 4;
+static const uint kRows = 6;
 static const uint kColumns = 4;
 static const float4 kPass = float4(0.0, 1.0, 0.0, 1.0);
 static const float4 kFail = float4(1.0, 0.0, 0.0, 1.0);
@@ -180,6 +180,28 @@ float4 RunRWStructuredBufferTest4() {
     return kPass;
 }
 
+ByteAddressBuffer byte_address_buffer;
+
+float4 RunByteAddressBufferTest() {
+    for (uint u = 0; u < kIndexCount; ++u) {
+        if (byte_address_buffer.Load(u * 4) != MakeU32(u)) {
+            return kFail;
+        }
+    }
+    return kPass;
+}
+
+RWByteAddressBuffer rwbyte_address_buffer;
+
+float4 RunRWByteAddressBufferTest() {
+    for (uint u = 0; u < kIndexCount; ++u) {
+        if (rwbyte_address_buffer.Load(u * 4) != MakeU32(u)) {
+            return kFail;
+        }
+    }
+    return kPass;
+}
+
 float4 RunTest(uint index) {
     switch (index) {
     case 0:
@@ -214,6 +236,16 @@ float4 RunTest(uint index) {
         return RunRWStructuredBufferTest3();
     case 15:
         return RunRWStructuredBufferTest4();
+    case 16:
+    case 17:
+    case 18:
+    case 19:
+        return RunByteAddressBufferTest();
+    case 20:
+    case 21:
+    case 22:
+    case 23:
+        return RunRWByteAddressBufferTest();
     default:
         return kIgnore;
     }
diff --git a/src/Apps/BufferViewTest/main.cpp b/src/Apps/BufferViewTest/main.cpp
@@ -3,6 +3,7 @@
 #include "AppSettings/ArgsParser.h"
 #include "Instance/Instance.h"
 #include "Utilities/Asset.h"
+#include "Utilities/Check.h"
 #include "Utilities/NotReached.h"
 
 #include <array>
@@ -33,13 +34,20 @@ std::string GetBufferPrefix(ViewType view_type)
         return "structured_buffer";
     case ViewType::kRWStructuredBuffer:
         return "rwstructured_buffer";
+    case ViewType::kByteAddressBuffer:
+        return "byte_address_buffer";
+    case ViewType::kRWByteAddressBuffer:
+        return "rwbyte_address_buffer";
     default:
         NOTREACHED();
     }
 }
 
 bool IsSupported(ViewType view_type, gli::format format)
 {
+    if (view_type == ViewType::kByteAddressBuffer || view_type == ViewType::kRWByteAddressBuffer) {
+        return format == gli::format::FORMAT_R32_UINT_PACK32;
+    }
     if (format == gli::format::FORMAT_RGB32_UINT_PACK32) {
         return view_type != ViewType::kBuffer && view_type != ViewType::kRWBuffer;
     }
@@ -110,6 +118,8 @@ BufferViewTestRenderer::BufferViewTestRenderer(const Settings& settings)
         ViewType::kRWBuffer,
         ViewType::kStructuredBuffer,
         ViewType::kRWStructuredBuffer,
+        ViewType::kByteAddressBuffer,
+        ViewType::kRWByteAddressBuffer,
     });
     const auto formats = std::to_array({
         gli::format::FORMAT_R32_UINT_PACK32,
@@ -131,7 +141,8 @@ BufferViewTestRenderer::BufferViewTestRenderer(const Settings& settings)
             if (buffer_desc.size % structure_stride != 0) {
                 buffer_desc.size += structure_stride - (buffer_desc.size % structure_stride);
             }
-            if (view_type == ViewType::kBuffer || view_type == ViewType::kStructuredBuffer) {
+            if (view_type == ViewType::kBuffer || view_type == ViewType::kStructuredBuffer ||
+                view_type == ViewType::kByteAddressBuffer) {
                 buffer_desc.usage = BindFlag::kShaderResource;
             } else {
                 buffer_desc.usage = BindFlag::kCopySource;
@@ -144,7 +155,8 @@ BufferViewTestRenderer::BufferViewTestRenderer(const Settings& settings)
             }
             buffer->Unmap();
 
-            if (view_type == ViewType::kRWBuffer || view_type == ViewType::kRWStructuredBuffer) {
+            if (view_type == ViewType::kRWBuffer || view_type == ViewType::kRWStructuredBuffer ||
+                view_type == ViewType::kRWByteAddressBuffer) {
                 upload_buffers_.push_back(std::move(buffer));
                 buffer_desc.usage = BindFlag::kUnorderedAccess | BindFlag::kCopyDest;
                 buffer = device_->CreateBuffer(MemoryType::kDefault, buffer_desc);
@@ -168,9 +180,12 @@ BufferViewTestRenderer::BufferViewTestRenderer(const Settings& settings)
             }
             std::shared_ptr<View> buffer_view = device_->CreateView(buffer, buffer_view_desc);
 
-            std::string bind_key_name = GetBufferPrefix(view_type) + "_uint";
-            bind_key_name += std::to_string(structure_stride / 4);
+            std::string bind_key_name = GetBufferPrefix(view_type);
+            if (view_type != ViewType::kByteAddressBuffer && view_type != ViewType::kRWByteAddressBuffer) {
+                bind_key_name += "_uint" + std::to_string(structure_stride / 4);
+            }
             BindKey bind_key = pixel_shader_->GetBindKey(bind_key_name);
+            DCHECK(bind_key.view_type == view_type);
             binding_set_layout_desc.bind_keys.push_back(bind_key);
             write_bindings_desc.bindings.emplace_back(bind_key, buffer_view);
 
diff --git a/src/FlyCube/BindingSet/MTBindingSet.mm b/src/FlyCube/BindingSet/MTBindingSet.mm
@@ -37,6 +37,8 @@ void SetView(id<MTL4ArgumentTable> argument_table, const std::shared_ptr<MTView>
     case ViewType::kConstantBuffer:
     case ViewType::kStructuredBuffer:
     case ViewType::kRWStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
+    case ViewType::kRWByteAddressBuffer:
         SetBuffer(argument_table, view->GetBuffer(), view->GetViewDesc().offset, index);
         break;
     case ViewType::kSampler:
diff --git a/src/FlyCube/BindingSetLayout/DXBindingSetLayout.cpp b/src/FlyCube/BindingSetLayout/DXBindingSetLayout.cpp
@@ -45,11 +45,13 @@ D3D12_DESCRIPTOR_RANGE_TYPE GetRangeType(ViewType view_type)
     case ViewType::kTexture:
     case ViewType::kBuffer:
     case ViewType::kStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
     case ViewType::kAccelerationStructure:
         return D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
     case ViewType::kRWTexture:
     case ViewType::kRWBuffer:
     case ViewType::kRWStructuredBuffer:
+    case ViewType::kRWByteAddressBuffer:
         return D3D12_DESCRIPTOR_RANGE_TYPE_UAV;
     case ViewType::kConstantBuffer:
         return D3D12_DESCRIPTOR_RANGE_TYPE_CBV;
diff --git a/src/FlyCube/BindingSetLayout/VKBindingSetLayout.cpp b/src/FlyCube/BindingSetLayout/VKBindingSetLayout.cpp
@@ -24,6 +24,10 @@ vk::DescriptorType GetDescriptorType(ViewType view_type)
         return vk::DescriptorType::eStorageBuffer;
     case ViewType::kRWStructuredBuffer:
         return vk::DescriptorType::eStorageBuffer;
+    case ViewType::kByteAddressBuffer:
+        return vk::DescriptorType::eStorageBuffer;
+    case ViewType::kRWByteAddressBuffer:
+        return vk::DescriptorType::eStorageBuffer;
     case ViewType::kAccelerationStructure:
         return vk::DescriptorType::eAccelerationStructureKHR;
     default:
diff --git a/src/FlyCube/CPUDescriptorPool/DXCPUDescriptorPool.cpp b/src/FlyCube/CPUDescriptorPool/DXCPUDescriptorPool.cpp
@@ -31,6 +31,8 @@ DXCPUDescriptorPoolTyped& DXCPUDescriptorPool::SelectHeap(ViewType view_type)
     case ViewType::kRWBuffer:
     case ViewType::kStructuredBuffer:
     case ViewType::kRWStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
+    case ViewType::kRWByteAddressBuffer:
         return resource_;
     case ViewType::kSampler:
         return sampler_;
diff --git a/src/FlyCube/Instance/BaseTypes.h b/src/FlyCube/Instance/BaseTypes.h
@@ -98,6 +98,8 @@ enum class ViewType {
     kRWBuffer,
     kStructuredBuffer,
     kRWStructuredBuffer,
+    kByteAddressBuffer,
+    kRWByteAddressBuffer,
     kAccelerationStructure,
     kShadingRateSource,
     kRenderTarget,
diff --git a/src/FlyCube/ShaderReflection/DXReflection.cpp b/src/FlyCube/ShaderReflection/DXReflection.cpp
@@ -62,6 +62,10 @@ ViewType GetViewType(const D3D12_SHADER_INPUT_BIND_DESC& bind_desc)
             return ViewType::kRWTexture;
         }
     }
+    case D3D_SIT_BYTEADDRESS:
+        return ViewType::kByteAddressBuffer;
+    case D3D_SIT_UAV_RWBYTEADDRESS:
+        return ViewType::kRWByteAddressBuffer;
     default:
         NOTREACHED();
     }
diff --git a/src/FlyCube/ShaderReflection/SPIRVReflection.cpp b/src/FlyCube/ShaderReflection/SPIRVReflection.cpp
@@ -156,11 +156,18 @@ ViewType GetViewType(const spirv_cross::Compiler& compiler, const spirv_cross::S
         if (type.storage == spv::StorageClassStorageBuffer) {
             spirv_cross::Bitset flags = compiler.get_buffer_block_flags(resource_id);
             bool is_readonly = flags.get(spv::DecorationNonWritable);
-            if (is_readonly) {
-                return ViewType::kStructuredBuffer;
+            bool is_byte_address_buffer = false;
+            if (compiler.has_decoration(resource_id, spv::DecorationUserTypeGOOGLE)) {
+                decltype(auto) user_type = compiler.get_decoration_string(resource_id, spv::DecorationUserTypeGOOGLE);
+                is_byte_address_buffer = user_type.find("byteaddressbuffer") != std::string::npos;
             } else {
-                return ViewType::kRWStructuredBuffer;
+                is_byte_address_buffer =
+                    compiler.get_name(type.parent_type).find("ByteAddressBuffer") != std::string::npos;
+            }
+            if (is_byte_address_buffer) {
+                return is_readonly ? ViewType::kByteAddressBuffer : ViewType::kRWByteAddressBuffer;
             }
+            return is_readonly ? ViewType::kStructuredBuffer : ViewType::kRWStructuredBuffer;
         } else if (type.storage == spv::StorageClassPushConstant || type.storage == spv::StorageClassUniform) {
             return ViewType::kConstantBuffer;
         }
diff --git a/src/FlyCube/View/DXBindlessTypedViewPool.cpp b/src/FlyCube/View/DXBindlessTypedViewPool.cpp
@@ -16,6 +16,8 @@ DXBindlessTypedViewPool::DXBindlessTypedViewPool(DXDevice& device, ViewType view
     case ViewType::kRWBuffer:
     case ViewType::kStructuredBuffer:
     case ViewType::kRWStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
+    case ViewType::kRWByteAddressBuffer:
     case ViewType::kAccelerationStructure: {
         range_ = std::make_shared<DXGPUDescriptorPoolRange>(
             device.GetGPUDescriptorPool().Allocate(D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV, view_count));
diff --git a/src/FlyCube/View/DXView.cpp b/src/FlyCube/View/DXView.cpp
@@ -44,6 +44,7 @@ void DXView::CreateView()
     case ViewType::kTexture:
     case ViewType::kBuffer:
     case ViewType::kStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
         CreateSRV();
         break;
     case ViewType::kAccelerationStructure:
@@ -52,6 +53,7 @@ void DXView::CreateView()
     case ViewType::kRWTexture:
     case ViewType::kRWBuffer:
     case ViewType::kRWStructuredBuffer:
+    case ViewType::kRWByteAddressBuffer:
         CreateUAV();
         break;
     case ViewType::kConstantBuffer:
@@ -150,10 +152,14 @@ void DXView::CreateSRV()
         if (view_desc_.view_type == ViewType::kBuffer) {
             srv_desc.Format = static_cast<DXGI_FORMAT>(gli::dx().translate(view_desc_.buffer_format).DXGIFormat.DDS);
             stride = gli::detail::bits_per_pixel(view_desc_.buffer_format) / 8;
-        } else {
-            assert(view_desc_.view_type == ViewType::kStructuredBuffer);
+        } else if (view_desc_.view_type == ViewType::kStructuredBuffer) {
             srv_desc.Buffer.StructureByteStride = view_desc_.structure_stride;
             stride = srv_desc.Buffer.StructureByteStride;
+        } else {
+            assert(view_desc_.view_type == ViewType::kByteAddressBuffer);
+            srv_desc.Format = DXGI_FORMAT_R32_TYPELESS;
+            srv_desc.Buffer.Flags = D3D12_BUFFER_SRV_FLAG_RAW;
+            stride = 4;
         }
         uint64_t size = std::min(resource_->GetResourceDesc().Width - view_desc_.offset, view_desc_.buffer_size);
         DCHECK(view_desc_.offset % stride == 0);
@@ -221,14 +227,19 @@ void DXView::CreateUAV()
         if (view_desc_.view_type == ViewType::kRWBuffer) {
             uav_desc.Format = static_cast<DXGI_FORMAT>(gli::dx().translate(view_desc_.buffer_format).DXGIFormat.DDS);
             stride = gli::detail::bits_per_pixel(view_desc_.buffer_format) / 8;
-        } else {
-            assert(view_desc_.view_type == ViewType::kRWStructuredBuffer);
+        } else if (view_desc_.view_type == ViewType::kRWStructuredBuffer) {
             uav_desc.Buffer.StructureByteStride = view_desc_.structure_stride;
             stride = uav_desc.Buffer.StructureByteStride;
+        } else {
+            assert(view_desc_.view_type == ViewType::kRWByteAddressBuffer);
+            uav_desc.Format = DXGI_FORMAT_R32_TYPELESS;
+            uav_desc.Buffer.Flags = D3D12_BUFFER_UAV_FLAG_RAW;
+            stride = 4;
         }
-        uint64_t size = std::min(resource_->GetResourceDesc().Width, view_desc_.buffer_size);
+        uint64_t size = std::min(resource_->GetResourceDesc().Width - view_desc_.offset, view_desc_.buffer_size);
+        DCHECK(view_desc_.offset % stride == 0);
         uav_desc.Buffer.FirstElement = view_desc_.offset / stride;
-        uav_desc.Buffer.NumElements = (size - view_desc_.offset) / (stride);
+        uav_desc.Buffer.NumElements = size / stride;
         break;
     }
     default: {
diff --git a/src/FlyCube/View/MTView.mm b/src/FlyCube/View/MTView.mm
@@ -163,6 +163,8 @@ MTLTextureType ConvertTextureType(ViewDimension dimension)
     case ViewType::kRWBuffer:
     case ViewType::kStructuredBuffer:
     case ViewType::kRWStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
+    case ViewType::kRWByteAddressBuffer:
         return GetBuffer();
     case ViewType::kSampler:
         return {};
@@ -189,7 +191,9 @@ MTLTextureType ConvertTextureType(ViewDimension dimension)
     case ViewType::kBuffer:
     case ViewType::kRWBuffer:
     case ViewType::kStructuredBuffer:
-    case ViewType::kRWStructuredBuffer: {
+    case ViewType::kRWStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
+    case ViewType::kRWByteAddressBuffer: {
         return [GetBuffer() gpuAddress] + GetViewDesc().offset;
     }
     case ViewType::kSampler: {
diff --git a/src/FlyCube/View/VKView.cpp b/src/FlyCube/View/VKView.cpp
@@ -85,7 +85,9 @@ void VKView::CreateView()
     }
     case ViewType::kConstantBuffer:
     case ViewType::kStructuredBuffer:
-    case ViewType::kRWStructuredBuffer: {
+    case ViewType::kRWStructuredBuffer:
+    case ViewType::kByteAddressBuffer:
+    case ViewType::kRWByteAddressBuffer: {
         uint64_t size = std::min(resource_->GetWidth() - view_desc_.offset, view_desc_.buffer_size);
         descriptor_buffer_.buffer = resource_->GetBuffer();
         descriptor_buffer_.offset = view_desc_.offset;

Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,10 @@ ViewType GetViewType(const D3D12_SHADER_INPUT_BIND_DESC& bind_desc)`
`62`	`62`	`return ViewType::kRWTexture;`
`63`	`63`	`}`
`64`	`64`	`}`
	`65`	`+ case D3D_SIT_BYTEADDRESS:`
	`66`	`+ return ViewType::kByteAddressBuffer;`
	`67`	`+ case D3D_SIT_UAV_RWBYTEADDRESS:`
	`68`	`+ return ViewType::kRWByteAddressBuffer;`
`65`	`69`	`default:`
`66`	`70`	`NOTREACHED();`
`67`	`71`	`}`