Implement internal storage for buffer usage and buffer binding type

In the timestamp internal pipeline, the ResolveQuery buffer cannnot be
binded as Storage buffer in binding group layout due to it has not
Storage usage.
Add InternalStorageBuffer for buffer usage and
InternalStorageBufferBinding for buffer binding type, make the
QueryResolve buffer implicitly get InternalStorageBuffer and only
compatible with InternalStorageBufferBinding in BGL, not Storage buffer
binding type.

Bug: dawn:797
Change-Id: I286339e703e26d3786c706ded03f850ca17355fb
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/54400
Reviewed-by: Jiawei Shao <jiawei.shao@intel.com>
Commit-Queue: Hao Li <hao.x.li@intel.com>
diff --git a/src/dawn_native/BindGroup.cpp b/src/dawn_native/BindGroup.cpp
index 6078eaa..a51b9ff 100644
--- a/src/dawn_native/BindGroup.cpp
+++ b/src/dawn_native/BindGroup.cpp
@@ -54,6 +54,10 @@
                     requiredUsage = wgpu::BufferUsage::Storage;
                     maxBindingSize = std::numeric_limits<uint64_t>::max();
                     break;
+                case kInternalStorageBufferBinding:
+                    requiredUsage = kInternalStorageBuffer;
+                    maxBindingSize = std::numeric_limits<uint64_t>::max();
+                    break;
                 case wgpu::BufferBindingType::Undefined:
                     UNREACHABLE();
             }
diff --git a/src/dawn_native/BindGroupLayout.cpp b/src/dawn_native/BindGroupLayout.cpp
index 0f322af..875d940 100644
--- a/src/dawn_native/BindGroupLayout.cpp
+++ b/src/dawn_native/BindGroupLayout.cpp
@@ -85,9 +85,15 @@
             if (entry.buffer.type != wgpu::BufferBindingType::Undefined) {
                 bindingMemberCount++;
                 const BufferBindingLayout& buffer = entry.buffer;
-                DAWN_TRY(ValidateBufferBindingType(buffer.type));
 
-                if (buffer.type == wgpu::BufferBindingType::Storage) {
+                // The kInternalStorageBufferBinding is used internally and not a value
+                // in wgpu::BufferBindingType.
+                if (buffer.type != kInternalStorageBufferBinding) {
+                    DAWN_TRY(ValidateBufferBindingType(buffer.type));
+                }
+
+                if (buffer.type == wgpu::BufferBindingType::Storage ||
+                    buffer.type == kInternalStorageBufferBinding) {
                     allowedStages &= ~wgpu::ShaderStage::Vertex;
                 }
 
@@ -96,6 +102,7 @@
                 if (device->IsToggleEnabled(Toggle::DisallowUnsafeAPIs) &&
                     buffer.hasDynamicOffset &&
                     (buffer.type == wgpu::BufferBindingType::Storage ||
+                     buffer.type == kInternalStorageBufferBinding ||
                      buffer.type == wgpu::BufferBindingType::ReadOnlyStorage)) {
                     return DAWN_VALIDATION_ERROR(
                         "Dynamic storage buffers are disallowed because they aren't secure yet. "
diff --git a/src/dawn_native/BindingInfo.cpp b/src/dawn_native/BindingInfo.cpp
index e2facc4..4b6516f 100644
--- a/src/dawn_native/BindingInfo.cpp
+++ b/src/dawn_native/BindingInfo.cpp
@@ -40,6 +40,7 @@
                     break;
 
                 case wgpu::BufferBindingType::Storage:
+                case kInternalStorageBufferBinding:
                 case wgpu::BufferBindingType::ReadOnlyStorage:
                     if (buffer.hasDynamicOffset) {
                         ++bindingCounts->dynamicStorageBufferCount;
diff --git a/src/dawn_native/Buffer.cpp b/src/dawn_native/Buffer.cpp
index cdd9f05..56c6ea6 100644
--- a/src/dawn_native/Buffer.cpp
+++ b/src/dawn_native/Buffer.cpp
@@ -138,12 +138,11 @@
             mUsage |= kReadOnlyStorageBuffer;
         }
 
-        // TODO(crbug.com/dawn/434): This is just a workaround to make QueryResolve buffer pass the
-        // binding group validation when used as an internal resource. Instead the buffer made with
-        // QueryResolve usage would implicitly get StorageInternal usage which is only compatible
-        // with StorageBufferInternal binding type in BGL, not StorageBuffer binding type.
+        // The buffer made with QueryResolve usage implicitly get InternalStorage usage which is
+        // only compatible with InternalStorageBuffer binding type in BGL, not StorageBuffer binding
+        // type.
         if (mUsage & wgpu::BufferUsage::QueryResolve) {
-            mUsage |= wgpu::BufferUsage::Storage;
+            mUsage |= kInternalStorageBuffer;
         }
     }
 
diff --git a/src/dawn_native/PassResourceUsageTracker.cpp b/src/dawn_native/PassResourceUsageTracker.cpp
index 18fd319..3aa90bf 100644
--- a/src/dawn_native/PassResourceUsageTracker.cpp
+++ b/src/dawn_native/PassResourceUsageTracker.cpp
@@ -80,6 +80,9 @@
                         case wgpu::BufferBindingType::Storage:
                             BufferUsedAs(buffer, wgpu::BufferUsage::Storage);
                             break;
+                        case kInternalStorageBufferBinding:
+                            BufferUsedAs(buffer, kInternalStorageBuffer);
+                            break;
                         case wgpu::BufferBindingType::ReadOnlyStorage:
                             BufferUsedAs(buffer, kReadOnlyStorageBuffer);
                             break;
diff --git a/src/dawn_native/QueryHelper.cpp b/src/dawn_native/QueryHelper.cpp
index 179db7b..ca98984 100644
--- a/src/dawn_native/QueryHelper.cpp
+++ b/src/dawn_native/QueryHelper.cpp
@@ -121,10 +121,33 @@
                     DAWN_TRY_ASSIGN(store->timestampCS, device->CreateShaderModule(&descriptor));
                 }
 
+                // Create binding group layout
+                std::array<BindGroupLayoutEntry, 3> entries = {};
+                for (uint32_t i = 0; i < entries.size(); i++) {
+                    entries[i].binding = i;
+                    entries[i].visibility = wgpu::ShaderStage::Compute;
+                }
+                entries[0].buffer.type = kInternalStorageBufferBinding;
+                entries[1].buffer.type = wgpu::BufferBindingType::ReadOnlyStorage;
+                entries[2].buffer.type = wgpu::BufferBindingType::Uniform;
+
+                BindGroupLayoutDescriptor bglDesc;
+                bglDesc.entryCount = static_cast<uint32_t>(entries.size());
+                bglDesc.entries = entries.data();
+                Ref<BindGroupLayoutBase> bgl;
+                DAWN_TRY_ASSIGN(bgl, device->CreateBindGroupLayout(&bglDesc));
+
+                // Create pipeline layout
+                PipelineLayoutDescriptor plDesc;
+                plDesc.bindGroupLayoutCount = 1;
+                plDesc.bindGroupLayouts = &bgl.Get();
+                Ref<PipelineLayoutBase> layout;
+                DAWN_TRY_ASSIGN(layout, device->CreatePipelineLayout(&plDesc));
+
                 // Create ComputePipeline.
                 ComputePipelineDescriptor computePipelineDesc = {};
                 // Generate the layout based on shader module.
-                computePipelineDesc.layout = nullptr;
+                computePipelineDesc.layout = layout.Get();
                 computePipelineDesc.compute.module = store->timestampCS.Get();
                 computePipelineDesc.compute.entryPoint = "main";
 
diff --git a/src/dawn_native/ShaderModule.cpp b/src/dawn_native/ShaderModule.cpp
index 930fcf4..7a5e149 100644
--- a/src/dawn_native/ShaderModule.cpp
+++ b/src/dawn_native/ShaderModule.cpp
@@ -594,10 +594,14 @@
                         // Binding mismatch between shader and bind group is invalid. For example, a
                         // writable binding in the shader with a readonly storage buffer in the bind
                         // group layout is invalid. However, a readonly binding in the shader with a
-                        // writable storage buffer in the bind group layout is valid.
+                        // writable storage buffer in the bind group layout is valid, a storage
+                        // binding in the shader with an internal storage buffer in the bind group
+                        // layout is also valid.
                         bool validBindingConversion =
-                            layoutInfo.buffer.type == wgpu::BufferBindingType::Storage &&
-                            shaderInfo.buffer.type == wgpu::BufferBindingType::ReadOnlyStorage;
+                            (layoutInfo.buffer.type == wgpu::BufferBindingType::Storage &&
+                             shaderInfo.buffer.type == wgpu::BufferBindingType::ReadOnlyStorage) ||
+                            (layoutInfo.buffer.type == kInternalStorageBufferBinding &&
+                             shaderInfo.buffer.type == wgpu::BufferBindingType::Storage);
 
                         if (layoutInfo.buffer.type != shaderInfo.buffer.type &&
                             !validBindingConversion) {
diff --git a/src/dawn_native/d3d12/BindGroupD3D12.cpp b/src/dawn_native/d3d12/BindGroupD3D12.cpp
index 2857d2d..a9f2731 100644
--- a/src/dawn_native/d3d12/BindGroupD3D12.cpp
+++ b/src/dawn_native/d3d12/BindGroupD3D12.cpp
@@ -78,7 +78,8 @@
                                                                  bindingOffsets[bindingIndex]));
                             break;
                         }
-                        case wgpu::BufferBindingType::Storage: {
+                        case wgpu::BufferBindingType::Storage:
+                        case kInternalStorageBufferBinding: {
                             // Since SPIRV-Cross outputs HLSL shaders with RWByteAddressBuffer,
                             // we must use D3D12_BUFFER_UAV_FLAG_RAW when making the
                             // UNORDERED_ACCESS_VIEW_DESC. Using D3D12_BUFFER_UAV_FLAG_RAW requires
diff --git a/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp b/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp
index ce84faa..28719de 100644
--- a/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp
+++ b/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp
@@ -30,6 +30,7 @@
                         case wgpu::BufferBindingType::Uniform:
                             return BindGroupLayout::DescriptorType::CBV;
                         case wgpu::BufferBindingType::Storage:
+                        case kInternalStorageBufferBinding:
                             return BindGroupLayout::DescriptorType::UAV;
                         case wgpu::BufferBindingType::ReadOnlyStorage:
                             return BindGroupLayout::DescriptorType::SRV;
diff --git a/src/dawn_native/d3d12/BufferD3D12.cpp b/src/dawn_native/d3d12/BufferD3D12.cpp
index 0223111..5f73b7e 100644
--- a/src/dawn_native/d3d12/BufferD3D12.cpp
+++ b/src/dawn_native/d3d12/BufferD3D12.cpp
@@ -31,7 +31,7 @@
         D3D12_RESOURCE_FLAGS D3D12ResourceFlags(wgpu::BufferUsage usage) {
             D3D12_RESOURCE_FLAGS flags = D3D12_RESOURCE_FLAG_NONE;
 
-            if (usage & wgpu::BufferUsage::Storage) {
+            if (usage & (wgpu::BufferUsage::Storage | kInternalStorageBuffer)) {
                 flags |= D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS;
             }
 
diff --git a/src/dawn_native/d3d12/CommandBufferD3D12.cpp b/src/dawn_native/d3d12/CommandBufferD3D12.cpp
index 29b3a98..706ab52 100644
--- a/src/dawn_native/d3d12/CommandBufferD3D12.cpp
+++ b/src/dawn_native/d3d12/CommandBufferD3D12.cpp
@@ -392,6 +392,7 @@
                             }
                             break;
                         case wgpu::BufferBindingType::Storage:
+                        case kInternalStorageBufferBinding:
                             if (mInCompute) {
                                 commandList->SetComputeRootUnorderedAccessView(parameterIndex,
                                                                                bufferLocation);
diff --git a/src/dawn_native/d3d12/PipelineLayoutD3D12.cpp b/src/dawn_native/d3d12/PipelineLayoutD3D12.cpp
index 913c845..8565eb0 100644
--- a/src/dawn_native/d3d12/PipelineLayoutD3D12.cpp
+++ b/src/dawn_native/d3d12/PipelineLayoutD3D12.cpp
@@ -45,6 +45,7 @@
                 case wgpu::BufferBindingType::Uniform:
                     return D3D12_ROOT_PARAMETER_TYPE_CBV;
                 case wgpu::BufferBindingType::Storage:
+                case kInternalStorageBufferBinding:
                     return D3D12_ROOT_PARAMETER_TYPE_UAV;
                 case wgpu::BufferBindingType::ReadOnlyStorage:
                     return D3D12_ROOT_PARAMETER_TYPE_SRV;
diff --git a/src/dawn_native/d3d12/ShaderModuleD3D12.cpp b/src/dawn_native/d3d12/ShaderModuleD3D12.cpp
index f7a5a91..ddf4cac 100644
--- a/src/dawn_native/d3d12/ShaderModuleD3D12.cpp
+++ b/src/dawn_native/d3d12/ShaderModuleD3D12.cpp
@@ -229,10 +229,13 @@
                 // storage buffer in the BGL produces the wrong output.
                 // Force read-only storage buffer bindings to be treated as UAV
                 // instead of SRV.
+                // Internal storage buffer is a storage buffer used in the internal pipeline.
                 const bool forceStorageBufferAsUAV =
                     (bindingInfo.buffer.type == wgpu::BufferBindingType::ReadOnlyStorage &&
-                     bgl->GetBindingInfo(bindingIndex).buffer.type ==
-                         wgpu::BufferBindingType::Storage);
+                     (bgl->GetBindingInfo(bindingIndex).buffer.type ==
+                          wgpu::BufferBindingType::Storage ||
+                      bgl->GetBindingInfo(bindingIndex).buffer.type ==
+                          kInternalStorageBufferBinding));
                 if (forceStorageBufferAsUAV) {
                     accessControls.emplace(srcBindingPoint, tint::ast::Access::kReadWrite);
                 }
diff --git a/src/dawn_native/dawn_platform.h b/src/dawn_native/dawn_platform.h
index 4a0ba09..9e75de4 100644
--- a/src/dawn_native/dawn_platform.h
+++ b/src/dawn_native/dawn_platform.h
@@ -35,6 +35,13 @@
     // This currently aliases wgpu::TextureUsage::Present, we would assign it
     // some bit when wgpu::TextureUsage::Present is removed.
     static constexpr wgpu::TextureUsage kPresentTextureUsage = wgpu::TextureUsage::Present;
+
+    // Add an extra buffer usage and an extra binding type for binding the buffers with QueryResolve
+    // usage as storage buffer in the internal pipeline.
+    static constexpr wgpu::BufferUsage kInternalStorageBuffer =
+        static_cast<wgpu::BufferUsage>(0x40000000);
+    static constexpr wgpu::BufferBindingType kInternalStorageBufferBinding =
+        static_cast<wgpu::BufferBindingType>(0xFFFFFFFF);
 }  // namespace dawn_native
 
 #endif  // DAWNNATIVE_DAWNPLATFORM_H_
diff --git a/src/dawn_native/metal/BufferMTL.mm b/src/dawn_native/metal/BufferMTL.mm
index f28021e..f7a59a0 100644
--- a/src/dawn_native/metal/BufferMTL.mm
+++ b/src/dawn_native/metal/BufferMTL.mm
@@ -53,7 +53,8 @@
         // Metal validation layer requires the size of uniform buffer and storage buffer to be no
         // less than the size of the buffer block defined in shader, and the overall size of the
         // buffer must be aligned to the largest alignment of its members.
-        if (GetUsage() & (wgpu::BufferUsage::Uniform | wgpu::BufferUsage::Storage)) {
+        if (GetUsage() &
+            (wgpu::BufferUsage::Uniform | wgpu::BufferUsage::Storage | kInternalStorageBuffer)) {
             if (currentSize >
                 std::numeric_limits<NSUInteger>::max() - kMinUniformOrStorageBufferAlignment) {
                 // Alignment would overlow.
diff --git a/src/dawn_native/opengl/CommandBufferGL.cpp b/src/dawn_native/opengl/CommandBufferGL.cpp
index 085a8f0..b4e98fe 100644
--- a/src/dawn_native/opengl/CommandBufferGL.cpp
+++ b/src/dawn_native/opengl/CommandBufferGL.cpp
@@ -265,6 +265,7 @@
                                     target = GL_UNIFORM_BUFFER;
                                     break;
                                 case wgpu::BufferBindingType::Storage:
+                                case kInternalStorageBufferBinding:
                                 case wgpu::BufferBindingType::ReadOnlyStorage:
                                     target = GL_SHADER_STORAGE_BUFFER;
                                     break;
diff --git a/src/dawn_native/opengl/PipelineLayoutGL.cpp b/src/dawn_native/opengl/PipelineLayoutGL.cpp
index ef2cf7c..a1742a5 100644
--- a/src/dawn_native/opengl/PipelineLayoutGL.cpp
+++ b/src/dawn_native/opengl/PipelineLayoutGL.cpp
@@ -43,6 +43,7 @@
                                 uboIndex++;
                                 break;
                             case wgpu::BufferBindingType::Storage:
+                            case kInternalStorageBufferBinding:
                             case wgpu::BufferBindingType::ReadOnlyStorage:
                                 mIndexInfo[group][bindingIndex] = ssboIndex;
                                 ssboIndex++;
diff --git a/src/dawn_native/vulkan/BindGroupLayoutVk.cpp b/src/dawn_native/vulkan/BindGroupLayoutVk.cpp
index 6ee1d49..1f8514d 100644
--- a/src/dawn_native/vulkan/BindGroupLayoutVk.cpp
+++ b/src/dawn_native/vulkan/BindGroupLayoutVk.cpp
@@ -56,6 +56,7 @@
                         }
                         return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
                     case wgpu::BufferBindingType::Storage:
+                    case kInternalStorageBufferBinding:
                     case wgpu::BufferBindingType::ReadOnlyStorage:
                         if (bindingInfo.buffer.hasDynamicOffset) {
                             return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
diff --git a/src/dawn_native/vulkan/BufferVk.cpp b/src/dawn_native/vulkan/BufferVk.cpp
index efe26c6..15d43a5 100644
--- a/src/dawn_native/vulkan/BufferVk.cpp
+++ b/src/dawn_native/vulkan/BufferVk.cpp
@@ -74,7 +74,7 @@
                 flags |= VK_PIPELINE_STAGE_VERTEX_INPUT_BIT;
             }
             if (usage & (wgpu::BufferUsage::Uniform | wgpu::BufferUsage::Storage |
-                         kReadOnlyStorageBuffer)) {
+                         kInternalStorageBuffer | kReadOnlyStorageBuffer)) {
                 flags |= VK_PIPELINE_STAGE_VERTEX_SHADER_BIT |
                          VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT |
                          VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
@@ -113,7 +113,7 @@
             if (usage & wgpu::BufferUsage::Uniform) {
                 flags |= VK_ACCESS_UNIFORM_READ_BIT;
             }
-            if (usage & wgpu::BufferUsage::Storage) {
+            if (usage & (wgpu::BufferUsage::Storage | kInternalStorageBuffer)) {
                 flags |= VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;
             }
             if (usage & kReadOnlyStorageBuffer) {
diff --git a/src/tests/unittests/validation/BindGroupValidationTests.cpp b/src/tests/unittests/validation/BindGroupValidationTests.cpp
index 69d890f..759b265 100644
--- a/src/tests/unittests/validation/BindGroupValidationTests.cpp
+++ b/src/tests/unittests/validation/BindGroupValidationTests.cpp
@@ -544,8 +544,7 @@
 }
 
 // Check that a resolve buffer with internal storge usage cannot be used as SSBO
-// TODO(hao.x.li@intel.com): Disable until internal storage usage is implemented
-TEST_F(BindGroupValidationTest, DISABLED_BufferUsageQueryResolve) {
+TEST_F(BindGroupValidationTest, BufferUsageQueryResolve) {
     wgpu::BindGroupLayout layout = utils::MakeBindGroupLayout(
         device, {{0, wgpu::ShaderStage::Fragment, wgpu::BufferBindingType::Storage}});
 
diff --git a/src/tests/white_box/QueryInternalShaderTests.cpp b/src/tests/white_box/QueryInternalShaderTests.cpp
index 56e88f9..4930f1e 100644
--- a/src/tests/white_box/QueryInternalShaderTests.cpp
+++ b/src/tests/white_box/QueryInternalShaderTests.cpp
@@ -122,7 +122,7 @@
     wgpu::BufferDescriptor timestampsDesc;
     timestampsDesc.size = kTimestampCount * sizeof(uint64_t);
     timestampsDesc.usage =
-        wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst;
+        wgpu::BufferUsage::QueryResolve | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst;
     wgpu::Buffer timestampsBuffer = device.CreateBuffer(&timestampsDesc);
 
     auto PrepareExpectedResults = [&](uint32_t first, uint32_t count,