Add D3D12 root constants to hold dynamic storage buffer sizes

All of these root constants are set when a bind group holding
dynamic storage buffers is applied. This could be improved by
using reflection data to only set constants that are needed
in the shader. This will require adding a way to store the
reflection information in the blob cache.

Bug: dawn:429
Change-Id: I3afce6b781ec5a82d5d0bafb6720d368b82c1b00
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/68600
Commit-Queue: Austin Eng <enga@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>

src/dawn_native/BindGroupLayout.cpp

@@ -553,4 +553,18 @@ namespace dawn_native {
                 {unverifiedBufferSizes, mBindingCounts.unverifiedBufferCount}};
     }
 
+    bool BindGroupLayoutBase::IsStorageBufferBinding(BindingIndex bindingIndex) const {
+        ASSERT(bindingIndex < GetBufferCount());
+        switch (GetBindingInfo(bindingIndex).buffer.type) {
+            case wgpu::BufferBindingType::Uniform:
+                return false;
+            case kInternalStorageBufferBinding:
+            case wgpu::BufferBindingType::Storage:
+            case wgpu::BufferBindingType::ReadOnlyStorage:
+                return true;
+            case wgpu::BufferBindingType::Undefined:
+                UNREACHABLE();
+        }
+    }
+
 }  // namespace dawn_native

src/dawn_native/BindGroupLayout.h

@@ -113,6 +113,8 @@ namespace dawn_native {
 
         BindingDataPointers ComputeBindingDataPointers(void* dataStart) const;
 
+        bool IsStorageBufferBinding(BindingIndex bindingIndex) const;
+
       protected:
         // Constructor used only for mocking and testing.
         BindGroupLayoutBase(DeviceBase* device);

src/dawn_native/d3d12/AdapterD3D12.cpp

@@ -238,12 +238,12 @@ namespace dawn_native { namespace d3d12 {
         //  - (2 * maxDynamicBuffers)
         //    Each dynamic buffer is a root descriptor
         //  RESERVED:
-        //  - 2 root constants for the baseVertex/baseInstance constants.
+        //  - 3 = max of:
-        //  - 3 root constants for num workgroups X, Y, Z
+        //    - 2 root constants for the baseVertex/baseInstance constants.
-        //  - (1)
+        //    - 3 root constants for num workgroups X, Y, Z
-        //    TODO(crbug.com/dawn/429): Dynamic storage buffers need bounds checks.
+        //  - 4 root constants (kMaxDynamicStorageBuffersPerPipelineLayout) for dynamic storage
-        //    This will probably be 1 CBV (root descriptor table) to store all the lengths.
+        //  buffer lengths.
-        static constexpr uint32_t kReservedSlots = 6;
+        static constexpr uint32_t kReservedSlots = 7;
 
         // Available slots after base limits considered.
         uint32_t availableRootSignatureSlots =

src/dawn_native/d3d12/BindGroupD3D12.cpp

@@ -203,6 +203,20 @@ namespace dawn_native { namespace d3d12 {
                 }
             }
         }
+
+        // Loop through the dynamic storage buffers and build a flat map from the index of the
+        // dynamic storage buffer to its binding size. The index |dynamicStorageBufferIndex|
+        // means that it is the i'th buffer that is both dynamic and storage, in increasing order
+        // of BindingNumber.
+        mDynamicStorageBufferLengths.resize(bgl->GetBindingCountInfo().dynamicStorageBufferCount);
+        uint32_t dynamicStorageBufferIndex = 0;
+        for (BindingIndex bindingIndex(0); bindingIndex < bgl->GetDynamicBufferCount();
+             ++bindingIndex) {
+            if (bgl->IsStorageBufferBinding(bindingIndex)) {
+                mDynamicStorageBufferLengths[dynamicStorageBufferIndex++] =
+                    GetBindingAsBufferBinding(bindingIndex).size;
+            }
+        }
     }
 
     BindGroup::~BindGroup() = default;
@@ -261,4 +275,10 @@ namespace dawn_native { namespace d3d12 {
     void BindGroup::SetSamplerAllocationEntry(Ref<SamplerHeapCacheEntry> entry) {
         mSamplerAllocationEntry = std::move(entry);
     }
+
+    const BindGroup::DynamicStorageBufferLengths& BindGroup::GetDynamicStorageBufferLengths()
+        const {
+        return mDynamicStorageBufferLengths;
+    }
+
 }}  // namespace dawn_native::d3d12

src/dawn_native/d3d12/BindGroupD3D12.h

@@ -16,6 +16,8 @@
 #define DAWNNATIVE_D3D12_BINDGROUPD3D12_H_
 
 #include "common/PlacementAllocated.h"
+#include "common/ityp_span.h"
+#include "common/ityp_stack_vec.h"
 #include "dawn_native/BindGroup.h"
 #include "dawn_native/d3d12/CPUDescriptorHeapAllocationD3D12.h"
 #include "dawn_native/d3d12/GPUDescriptorHeapAllocationD3D12.h"
@@ -45,6 +47,10 @@ namespace dawn_native { namespace d3d12 {
 
         void SetSamplerAllocationEntry(Ref<SamplerHeapCacheEntry> entry);
 
+        using DynamicStorageBufferLengths =
+            ityp::stack_vec<uint32_t, uint32_t, kMaxDynamicStorageBuffersPerPipelineLayout>;
+        const DynamicStorageBufferLengths& GetDynamicStorageBufferLengths() const;
+
       private:
         ~BindGroup() override;
 
@@ -54,6 +60,8 @@ namespace dawn_native { namespace d3d12 {
 
         GPUDescriptorHeapAllocation mGPUViewAllocation;
         CPUDescriptorHeapAllocation mCPUViewAllocation;
+
+        DynamicStorageBufferLengths mDynamicStorageBufferLengths;
     };
 }}  // namespace dawn_native::d3d12
 

src/dawn_native/d3d12/CommandBufferD3D12.cpp

@@ -530,6 +530,24 @@ namespace dawn_native { namespace d3d12 {
                     mBoundRootSamplerTables[index] = baseDescriptor;
                 }
             }
+
+            const auto& dynamicStorageBufferLengths = group->GetDynamicStorageBufferLengths();
+            if (dynamicStorageBufferLengths.size() != 0) {
+                uint32_t parameterIndex =
+                    pipelineLayout->GetDynamicStorageBufferLengthsParameterIndex();
+                uint32_t firstRegisterOffset =
+                    pipelineLayout->GetDynamicStorageBufferLengthInfo()[index].firstRegisterOffset;
+
+                if (mInCompute) {
+                    commandList->SetComputeRoot32BitConstants(
+                        parameterIndex, dynamicStorageBufferLengths.size(),
+                        dynamicStorageBufferLengths.data(), firstRegisterOffset);
+                } else {
+                    commandList->SetGraphicsRoot32BitConstants(
+                        parameterIndex, dynamicStorageBufferLengths.size(),
+                        dynamicStorageBufferLengths.data(), firstRegisterOffset);
+                }
+            }
         }
 
         Device* mDevice;

src/dawn_native/d3d12/PipelineLayoutD3D12.cpp

@@ -26,6 +26,17 @@ using Microsoft::WRL::ComPtr;
 
 namespace dawn_native { namespace d3d12 {
     namespace {
+
+        // Reserve register names for internal use. This registers map to bindings in the shader,
+        // but are not directly related to allocation of the root signature.
+        // In the root signature, it the index of the root parameter where these registers are
+        // used that determines the layout of the root signature.
+        static constexpr uint32_t kRenderOrComputeInternalRegisterSpace = kMaxBindGroups + 1;
+        static constexpr uint32_t kRenderOrComputeInternalBaseRegister = 0;
+
+        static constexpr uint32_t kDynamicStorageBufferLengthsRegisterSpace = kMaxBindGroups + 2;
+        static constexpr uint32_t kDynamicStorageBufferLengthsBaseRegister = 0;
+
         D3D12_SHADER_VISIBILITY ShaderVisibilityType(wgpu::ShaderStage visibility) {
             ASSERT(visibility != wgpu::ShaderStage::None);
 
@@ -54,6 +65,7 @@ namespace dawn_native { namespace d3d12 {
                     UNREACHABLE();
             }
         }
+
     }  // anonymous namespace
 
     ResultOrError<Ref<PipelineLayout>> PipelineLayout::Create(
@@ -160,34 +172,70 @@ namespace dawn_native { namespace d3d12 {
         // |ranges| will have resized and the pointers in the |rootParameter|s will be invalid.
         ASSERT(rangeIndex == rangesCount);
 
-        D3D12_ROOT_PARAMETER indexOffsetConstants{};
+        D3D12_ROOT_PARAMETER renderOrComputeInternalConstants{};
-        indexOffsetConstants.ShaderVisibility = D3D12_SHADER_VISIBILITY_VERTEX;
+        renderOrComputeInternalConstants.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
-        indexOffsetConstants.ParameterType = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
+        renderOrComputeInternalConstants.ParameterType = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
-        // Always allocate 2 constants for vertex_index and instance_index
+        // Always allocate 3 constants for either:
+        //  - vertex_index and instance_index
+        //  - num_workgroups_x, num_workgroups_y and num_workgroups_z
         // NOTE: We should consider delaying root signature creation until we know how many values
         // we need
-        indexOffsetConstants.Constants.Num32BitValues = 2;
+        renderOrComputeInternalConstants.Constants.Num32BitValues = 3;
-        indexOffsetConstants.Constants.RegisterSpace = kReservedRegisterSpace;
+        renderOrComputeInternalConstants.Constants.RegisterSpace =
-        indexOffsetConstants.Constants.ShaderRegister = kFirstOffsetInfoBaseRegister;
+            kRenderOrComputeInternalRegisterSpace;
+        renderOrComputeInternalConstants.Constants.ShaderRegister =
+            kRenderOrComputeInternalBaseRegister;
         mFirstIndexOffsetParameterIndex = rootParameters.size();
+        mNumWorkgroupsParameterIndex = rootParameters.size();
         // NOTE: We should consider moving this entry to earlier in the root signature since offsets
         // would need to be updated often
-        rootParameters.emplace_back(indexOffsetConstants);
+        rootParameters.emplace_back(renderOrComputeInternalConstants);
 
-        // Always allocate 3 constants for num_workgroups_x, num_workgroups_y and num_workgroups_z
+        // Loops over all of the dynamic storage buffer bindings in the layout and build
-        // for Dispatch calls
+        // a mapping from the binding to the next offset into the root constant array where
-        // NOTE: We should consider delaying root signature creation until we know how many values
+        // that dynamic storage buffer's binding size will be stored. The next register offset
-        // we need
+        // to use is tracked with |dynamicStorageBufferLengthsShaderRegisterOffset|.
-        D3D12_ROOT_PARAMETER numWorkgroupsConstants{};
+        // This data will be used by shader translation to emit a load from the root constant
-        numWorkgroupsConstants.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+        // array to use as the binding's size in runtime array calculations.
-        numWorkgroupsConstants.ParameterType = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
+        // Each bind group's length data is stored contiguously in the root constant array,
-        numWorkgroupsConstants.Constants.Num32BitValues = 3;
+        // so the loop also computes the first register offset for each group where the
-        numWorkgroupsConstants.Constants.RegisterSpace = GetNumWorkgroupsRegisterSpace();
+        // data should start.
-        numWorkgroupsConstants.Constants.ShaderRegister = GetNumWorkgroupsShaderRegister();
+        uint32_t dynamicStorageBufferLengthsShaderRegisterOffset = 0;
-        mNumWorkgroupsParameterIndex = rootParameters.size();
+        for (BindGroupIndex group : IterateBitSet(GetBindGroupLayoutsMask())) {
-        // NOTE: We should consider moving this entry to earlier in the root signature since
+            const BindGroupLayoutBase* bgl = GetBindGroupLayout(group);
-        // dispatch sizes would need to be updated often
+
-        rootParameters.emplace_back(numWorkgroupsConstants);
+            mDynamicStorageBufferLengthInfo[group].firstRegisterOffset =
+                dynamicStorageBufferLengthsShaderRegisterOffset;
+            mDynamicStorageBufferLengthInfo[group].bindingAndRegisterOffsets.reserve(
+                bgl->GetBindingCountInfo().dynamicStorageBufferCount);
+
+            for (BindingIndex bindingIndex(0); bindingIndex < bgl->GetDynamicBufferCount();
+                 ++bindingIndex) {
+                if (bgl->IsStorageBufferBinding(bindingIndex)) {
+                    mDynamicStorageBufferLengthInfo[group].bindingAndRegisterOffsets.push_back(
+                        {bgl->GetBindingInfo(bindingIndex).binding,
+                         dynamicStorageBufferLengthsShaderRegisterOffset++});
+                }
+            }
+
+            ASSERT(mDynamicStorageBufferLengthInfo[group].bindingAndRegisterOffsets.size() ==
+                   bgl->GetBindingCountInfo().dynamicStorageBufferCount);
+        }
+        ASSERT(dynamicStorageBufferLengthsShaderRegisterOffset <=
+               kMaxDynamicStorageBuffersPerPipelineLayout);
+
+        D3D12_ROOT_PARAMETER dynamicStorageBufferLengthConstants{};
+        dynamicStorageBufferLengthConstants.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+        dynamicStorageBufferLengthConstants.ParameterType =
+            D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
+        dynamicStorageBufferLengthConstants.Constants.Num32BitValues =
+            dynamicStorageBufferLengthsShaderRegisterOffset;
+        dynamicStorageBufferLengthConstants.Constants.RegisterSpace =
+            kDynamicStorageBufferLengthsRegisterSpace;
+        dynamicStorageBufferLengthConstants.Constants.ShaderRegister =
+            kDynamicStorageBufferLengthsBaseRegister;
+        mDynamicStorageBufferLengthsParameterIndex = rootParameters.size();
+        rootParameters.emplace_back(dynamicStorageBufferLengthConstants);
 
         D3D12_ROOT_SIGNATURE_DESC rootSignatureDescriptor;
         rootSignatureDescriptor.NumParameters = rootParameters.size();
@@ -234,6 +282,11 @@ namespace dawn_native { namespace d3d12 {
         return mRootSignature.Get();
     }
 
+    const PipelineLayout::DynamicStorageBufferLengthInfo&
+    PipelineLayout::GetDynamicStorageBufferLengthInfo() const {
+        return mDynamicStorageBufferLengthInfo;
+    }
+
     uint32_t PipelineLayout::GetDynamicRootParameterIndex(BindGroupIndex group,
                                                           BindingIndex bindingIndex) const {
         ASSERT(group < kMaxBindGroupsTyped);
@@ -245,11 +298,11 @@ namespace dawn_native { namespace d3d12 {
     }
 
     uint32_t PipelineLayout::GetFirstIndexOffsetRegisterSpace() const {
-        return kFirstIndexOffsetRegisterSpace;
+        return kRenderOrComputeInternalRegisterSpace;
     }
 
     uint32_t PipelineLayout::GetFirstIndexOffsetShaderRegister() const {
-        return kFirstOffsetInfoBaseRegister;
+        return kRenderOrComputeInternalBaseRegister;
     }
 
     uint32_t PipelineLayout::GetFirstIndexOffsetParameterIndex() const {
@@ -257,17 +310,29 @@ namespace dawn_native { namespace d3d12 {
     }
 
     uint32_t PipelineLayout::GetNumWorkgroupsRegisterSpace() const {
-        return kNumWorkgroupsRegisterSpace;
+        return kRenderOrComputeInternalRegisterSpace;
     }
 
     uint32_t PipelineLayout::GetNumWorkgroupsShaderRegister() const {
-        return kNumWorkgroupsBaseRegister;
+        return kRenderOrComputeInternalBaseRegister;
     }
 
     uint32_t PipelineLayout::GetNumWorkgroupsParameterIndex() const {
         return mNumWorkgroupsParameterIndex;
     }
 
+    uint32_t PipelineLayout::GetDynamicStorageBufferLengthsRegisterSpace() const {
+        return kDynamicStorageBufferLengthsRegisterSpace;
+    }
+
+    uint32_t PipelineLayout::GetDynamicStorageBufferLengthsShaderRegister() const {
+        return kDynamicStorageBufferLengthsBaseRegister;
+    }
+
+    uint32_t PipelineLayout::GetDynamicStorageBufferLengthsParameterIndex() const {
+        return mDynamicStorageBufferLengthsParameterIndex;
+    }
+
     ID3D12CommandSignature* PipelineLayout::GetDispatchIndirectCommandSignatureWithNumWorkgroups() {
         // mDispatchIndirectCommandSignatureWithNumWorkgroups won't be created until it is needed.
         if (mDispatchIndirectCommandSignatureWithNumWorkgroups.Get() != nullptr) {

src/dawn_native/d3d12/PipelineLayoutD3D12.h

@@ -23,13 +23,6 @@
 
 namespace dawn_native { namespace d3d12 {
 
-    // We reserve a register space that a user cannot use.
-    static constexpr uint32_t kReservedRegisterSpace = kMaxBindGroups + 1;
-    static constexpr uint32_t kFirstOffsetInfoBaseRegister = 0;
-    static constexpr uint32_t kFirstIndexOffsetRegisterSpace = kReservedRegisterSpace;
-    static constexpr uint32_t kNumWorkgroupsRegisterSpace = kReservedRegisterSpace + 1;
-    static constexpr uint32_t kNumWorkgroupsBaseRegister = 0;
-
     class Device;
 
     class PipelineLayout final : public PipelineLayoutBase {
@@ -53,10 +46,37 @@ namespace dawn_native { namespace d3d12 {
         uint32_t GetNumWorkgroupsShaderRegister() const;
         uint32_t GetNumWorkgroupsParameterIndex() const;
 
+        uint32_t GetDynamicStorageBufferLengthsRegisterSpace() const;
+        uint32_t GetDynamicStorageBufferLengthsShaderRegister() const;
+        uint32_t GetDynamicStorageBufferLengthsParameterIndex() const;
+
         ID3D12RootSignature* GetRootSignature() const;
 
         ID3D12CommandSignature* GetDispatchIndirectCommandSignatureWithNumWorkgroups();
 
+        struct PerBindGroupDynamicStorageBufferLengthInfo {
+            // First register offset for a bind group's dynamic storage buffer lengths.
+            // This is the index into the array of root constants where this bind group's
+            // lengths start.
+            uint32_t firstRegisterOffset;
+
+            struct BindingAndRegisterOffset {
+                BindingNumber binding;
+                uint32_t registerOffset;
+            };
+            // Associative list of (BindingNumber,registerOffset) pairs, which is passed into
+            // the shader to map the BindingPoint(thisGroup, BindingNumber) to the registerOffset
+            // into the root constant array which holds the dynamic storage buffer lengths.
+            std::vector<BindingAndRegisterOffset> bindingAndRegisterOffsets;
+        };
+
+        // Flat map from bind group index to the list of (BindingNumber,Register) pairs.
+        // Each pair is used in shader translation to
+        using DynamicStorageBufferLengthInfo =
+            ityp::array<BindGroupIndex, PerBindGroupDynamicStorageBufferLengthInfo, kMaxBindGroups>;
+
+        const DynamicStorageBufferLengthInfo& GetDynamicStorageBufferLengthInfo() const;
+
       private:
         ~PipelineLayout() override = default;
         using PipelineLayoutBase::PipelineLayoutBase;
@@ -67,8 +87,10 @@ namespace dawn_native { namespace d3d12 {
                     ityp::array<BindingIndex, uint32_t, kMaxDynamicBuffersPerPipelineLayout>,
                     kMaxBindGroups>
             mDynamicRootParameterIndices;
+        DynamicStorageBufferLengthInfo mDynamicStorageBufferLengthInfo;
         uint32_t mFirstIndexOffsetParameterIndex;
         uint32_t mNumWorkgroupsParameterIndex;
+        uint32_t mDynamicStorageBufferLengthsParameterIndex;
         ComPtr<ID3D12RootSignature> mRootSignature;
         ComPtr<ID3D12CommandSignature> mDispatchIndirectCommandSignatureWithNumWorkgroups;
     };