D3D12: Bucket CPU descriptor allocators on the device.

Move the CPU descriptor allocators to the device and bucket them to
ensure only kMaxBindingsPerGroup exist rather than create them per BGL.

Also, renames NonShaderVisible => Staging.

BUG=dawn:155

Change-Id: If6dae368e7e2a2b349343bdf898041a049159038
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/19001
Commit-Queue: Bryan Bernhart <bryan.bernhart@intel.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn_native/BUILD.gn b/src/dawn_native/BUILD.gn
index 246410a..73f0ab2 100644
--- a/src/dawn_native/BUILD.gn
+++ b/src/dawn_native/BUILD.gn
@@ -302,8 +302,6 @@
       "d3d12/HeapD3D12.h",
       "d3d12/NativeSwapChainImplD3D12.cpp",
       "d3d12/NativeSwapChainImplD3D12.h",
-      "d3d12/NonShaderVisibleDescriptorAllocatorD3D12.cpp",
-      "d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h",
       "d3d12/PipelineLayoutD3D12.cpp",
       "d3d12/PipelineLayoutD3D12.h",
       "d3d12/PlatformFunctions.cpp",
@@ -328,6 +326,8 @@
       "d3d12/ShaderVisibleDescriptorAllocatorD3D12.h",
       "d3d12/StagingBufferD3D12.cpp",
       "d3d12/StagingBufferD3D12.h",
+      "d3d12/StagingDescriptorAllocatorD3D12.cpp",
+      "d3d12/StagingDescriptorAllocatorD3D12.h",
       "d3d12/SwapChainD3D12.cpp",
       "d3d12/SwapChainD3D12.h",
       "d3d12/TextureCopySplitter.cpp",
diff --git a/src/dawn_native/CMakeLists.txt b/src/dawn_native/CMakeLists.txt
index 1350c66..95e04f3 100644
--- a/src/dawn_native/CMakeLists.txt
+++ b/src/dawn_native/CMakeLists.txt
@@ -191,8 +191,6 @@
         "d3d12/HeapD3D12.h"
         "d3d12/NativeSwapChainImplD3D12.cpp"
         "d3d12/NativeSwapChainImplD3D12.h"
-        "d3d12/NonShaderVisibleDescriptorAllocatorD3D12.cpp"
-        "d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h"
         "d3d12/PipelineLayoutD3D12.cpp"
         "d3d12/PipelineLayoutD3D12.h"
         "d3d12/PlatformFunctions.cpp"
@@ -217,6 +215,8 @@
         "d3d12/ShaderVisibleDescriptorAllocatorD3D12.h"
         "d3d12/StagingBufferD3D12.cpp"
         "d3d12/StagingBufferD3D12.h"
+        "d3d12/StagingDescriptorAllocatorD3D12.cpp"
+        "d3d12/StagingDescriptorAllocatorD3D12.h"
         "d3d12/SwapChainD3D12.cpp"
         "d3d12/SwapChainD3D12.h"
         "d3d12/TextureCopySplitter.cpp"
diff --git a/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp b/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp
index 0ee36d1..b8012d1 100644
--- a/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp
+++ b/src/dawn_native/d3d12/BindGroupLayoutD3D12.cpp
@@ -17,7 +17,7 @@
 #include "common/BitSetIterator.h"
 #include "dawn_native/d3d12/BindGroupD3D12.h"
 #include "dawn_native/d3d12/DeviceD3D12.h"
-#include "dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h"
+#include "dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h"
 
 namespace dawn_native { namespace d3d12 {
     namespace {
@@ -42,9 +42,6 @@
         }
     }  // anonymous namespace
 
-    // TODO(dawn:155): Figure out this value.
-    static constexpr uint16_t kDescriptorHeapSize = 1024;
-
     BindGroupLayout::BindGroupLayout(Device* device, const BindGroupLayoutDescriptor* descriptor)
         : BindGroupLayoutBase(device, descriptor),
           mDescriptorCounts{},
@@ -133,19 +130,9 @@
             mBindingOffsets[bindingIndex] += descriptorOffsets[descriptorType];
         }
 
-        const uint32_t viewDescriptorCount = GetCbvUavSrvDescriptorCount();
-        if (viewDescriptorCount > 0) {
-            mViewAllocator = std::make_unique<NonShaderVisibleDescriptorAllocator>(
-                device, viewDescriptorCount, kDescriptorHeapSize,
-                D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
-        }
-
-        const uint32_t samplerDescriptorCount = GetSamplerDescriptorCount();
-        if (samplerDescriptorCount > 0) {
-            mSamplerAllocator = std::make_unique<NonShaderVisibleDescriptorAllocator>(
-                device, samplerDescriptorCount, kDescriptorHeapSize,
-                D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER);
-        }
+        mViewAllocator = device->GetViewStagingDescriptorAllocator(GetCbvUavSrvDescriptorCount());
+        mSamplerAllocator =
+            device->GetSamplerStagingDescriptorAllocator(GetSamplerDescriptorCount());
     }
 
     ResultOrError<BindGroup*> BindGroupLayout::AllocateBindGroup(
diff --git a/src/dawn_native/d3d12/BindGroupLayoutD3D12.h b/src/dawn_native/d3d12/BindGroupLayoutD3D12.h
index e3e19d3..d04ab75 100644
--- a/src/dawn_native/d3d12/BindGroupLayoutD3D12.h
+++ b/src/dawn_native/d3d12/BindGroupLayoutD3D12.h
@@ -23,9 +23,9 @@
 namespace dawn_native { namespace d3d12 {
 
     class BindGroup;
-    class Device;
-    class NonShaderVisibleDescriptorAllocator;
     class CPUDescriptorHeapAllocation;
+    class Device;
+    class StagingDescriptorAllocator;
 
     class BindGroupLayout final : public BindGroupLayoutBase {
       public:
@@ -61,9 +61,8 @@
 
         SlabAllocator<BindGroup> mBindGroupAllocator;
 
-        // TODO(dawn:155): Store and bucket allocators by size on the device.
-        std::unique_ptr<NonShaderVisibleDescriptorAllocator> mSamplerAllocator;
-        std::unique_ptr<NonShaderVisibleDescriptorAllocator> mViewAllocator;
+        StagingDescriptorAllocator* mSamplerAllocator = nullptr;
+        StagingDescriptorAllocator* mViewAllocator = nullptr;
     };
 
 }}  // namespace dawn_native::d3d12
diff --git a/src/dawn_native/d3d12/DeviceD3D12.cpp b/src/dawn_native/d3d12/DeviceD3D12.cpp
index c08c06b..f1f9baa 100644
--- a/src/dawn_native/d3d12/DeviceD3D12.cpp
+++ b/src/dawn_native/d3d12/DeviceD3D12.cpp
@@ -37,11 +37,15 @@
 #include "dawn_native/d3d12/ShaderModuleD3D12.h"
 #include "dawn_native/d3d12/ShaderVisibleDescriptorAllocatorD3D12.h"
 #include "dawn_native/d3d12/StagingBufferD3D12.h"
+#include "dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h"
 #include "dawn_native/d3d12/SwapChainD3D12.h"
 #include "dawn_native/d3d12/TextureD3D12.h"
 
 namespace dawn_native { namespace d3d12 {
 
+    // TODO(dawn:155): Figure out this value.
+    static constexpr uint16_t kStagingDescriptorHeapSize = 1024;
+
     // static
     ResultOrError<Device*> Device::Create(Adapter* adapter, const DeviceDescriptor* descriptor) {
         Ref<Device> device = AcquireRef(new Device(adapter, descriptor));
@@ -83,6 +87,17 @@
             std::make_unique<ShaderVisibleDescriptorAllocator>(this);
         DAWN_TRY(mShaderVisibleDescriptorAllocator->Initialize());
 
+        // Zero sized allocator is never requested and does not need to exist.
+        for (uint32_t countIndex = 1; countIndex < kNumOfStagingDescriptorAllocators;
+             countIndex++) {
+            mViewAllocators[countIndex] = std::make_unique<StagingDescriptorAllocator>(
+                this, countIndex, kStagingDescriptorHeapSize,
+                D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
+
+            mSamplerAllocators[countIndex] = std::make_unique<StagingDescriptorAllocator>(
+                this, countIndex, kStagingDescriptorHeapSize, D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER);
+        }
+
         mMapRequestTracker = std::make_unique<MapRequestTracker>(this);
         mResidencyManager = std::make_unique<ResidencyManager>(this);
         mResourceAllocatorManager = std::make_unique<ResourceAllocatorManager>(this);
@@ -458,4 +473,16 @@
     ShaderVisibleDescriptorAllocator* Device::GetShaderVisibleDescriptorAllocator() const {
         return mShaderVisibleDescriptorAllocator.get();
     }
+
+    StagingDescriptorAllocator* Device::GetViewStagingDescriptorAllocator(
+        uint32_t descriptorCount) const {
+        ASSERT(descriptorCount < kNumOfStagingDescriptorAllocators);
+        return mViewAllocators[descriptorCount].get();
+    }
+
+    StagingDescriptorAllocator* Device::GetSamplerStagingDescriptorAllocator(
+        uint32_t descriptorCount) const {
+        ASSERT(descriptorCount < kNumOfStagingDescriptorAllocators);
+        return mSamplerAllocators[descriptorCount].get();
+    }
 }}  // namespace dawn_native::d3d12
diff --git a/src/dawn_native/d3d12/DeviceD3D12.h b/src/dawn_native/d3d12/DeviceD3D12.h
index 406ce39..364951e 100644
--- a/src/dawn_native/d3d12/DeviceD3D12.h
+++ b/src/dawn_native/d3d12/DeviceD3D12.h
@@ -17,6 +17,7 @@
 
 #include "dawn_native/dawn_platform.h"
 
+#include "common/Constants.h"
 #include "common/SerialQueue.h"
 #include "dawn_native/Device.h"
 #include "dawn_native/d3d12/CommandRecordingContext.h"
@@ -30,11 +31,12 @@
 
     class CommandAllocatorManager;
     class DescriptorHeapAllocator;
-    class ShaderVisibleDescriptorAllocator;
     class MapRequestTracker;
     class PlatformFunctions;
-    class ResourceAllocatorManager;
     class ResidencyManager;
+    class ResourceAllocatorManager;
+    class ShaderVisibleDescriptorAllocator;
+    class StagingDescriptorAllocator;
 
 #define ASSERT_SUCCESS(hr)            \
     {                                 \
@@ -101,6 +103,13 @@
 
         ShaderVisibleDescriptorAllocator* GetShaderVisibleDescriptorAllocator() const;
 
+        // Returns nullptr when descriptor count is zero.
+        StagingDescriptorAllocator* GetViewStagingDescriptorAllocator(
+            uint32_t descriptorCount) const;
+
+        StagingDescriptorAllocator* GetSamplerStagingDescriptorAllocator(
+            uint32_t descriptorCount) const;
+
         TextureBase* WrapSharedHandle(const ExternalImageDescriptor* descriptor,
                                       HANDLE sharedHandle,
                                       uint64_t acquireMutexKey,
@@ -170,6 +179,15 @@
         std::unique_ptr<ResourceAllocatorManager> mResourceAllocatorManager;
         std::unique_ptr<ResidencyManager> mResidencyManager;
         std::unique_ptr<ShaderVisibleDescriptorAllocator> mShaderVisibleDescriptorAllocator;
+
+        // Index corresponds to the descriptor count in the range [0, kMaxBindingsPerGroup].
+        static constexpr uint32_t kNumOfStagingDescriptorAllocators = kMaxBindingsPerGroup + 1;
+
+        std::array<std::unique_ptr<StagingDescriptorAllocator>, kNumOfStagingDescriptorAllocators>
+            mViewAllocators;
+
+        std::array<std::unique_ptr<StagingDescriptorAllocator>, kNumOfStagingDescriptorAllocators>
+            mSamplerAllocators;
     };
 
 }}  // namespace dawn_native::d3d12
diff --git a/src/dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.cpp b/src/dawn_native/d3d12/StagingDescriptorAllocatorD3D12.cpp
similarity index 83%
rename from src/dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.cpp
rename to src/dawn_native/d3d12/StagingDescriptorAllocatorD3D12.cpp
index 0bb60f7..d7e0653 100644
--- a/src/dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.cpp
+++ b/src/dawn_native/d3d12/StagingDescriptorAllocatorD3D12.cpp
@@ -16,15 +16,14 @@
 
 #include "dawn_native/d3d12/D3D12Error.h"
 #include "dawn_native/d3d12/DeviceD3D12.h"
-#include "dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h"
+#include "dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h"
 
 namespace dawn_native { namespace d3d12 {
 
-    NonShaderVisibleDescriptorAllocator::NonShaderVisibleDescriptorAllocator(
-        Device* device,
-        uint32_t descriptorCount,
-        uint32_t heapSize,
-        D3D12_DESCRIPTOR_HEAP_TYPE heapType)
+    StagingDescriptorAllocator::StagingDescriptorAllocator(Device* device,
+                                                           uint32_t descriptorCount,
+                                                           uint32_t heapSize,
+                                                           D3D12_DESCRIPTOR_HEAP_TYPE heapType)
         : mDevice(device),
           mSizeIncrement(device->GetD3D12Device()->GetDescriptorHandleIncrementSize(heapType)),
           mBlockSize(descriptorCount * mSizeIncrement),
@@ -35,7 +34,7 @@
         ASSERT(descriptorCount <= heapSize);
     }
 
-    NonShaderVisibleDescriptorAllocator::~NonShaderVisibleDescriptorAllocator() {
+    StagingDescriptorAllocator::~StagingDescriptorAllocator() {
         const Index freeBlockIndicesSize = GetFreeBlockIndicesSize();
         for (auto& buffer : mPool) {
             ASSERT(buffer.freeBlockIndices.size() == freeBlockIndicesSize);
@@ -44,7 +43,7 @@
     }
 
     ResultOrError<CPUDescriptorHeapAllocation>
-    NonShaderVisibleDescriptorAllocator::AllocateCPUDescriptors() {
+    StagingDescriptorAllocator::AllocateCPUDescriptors() {
         if (mAvailableHeaps.empty()) {
             DAWN_TRY(AllocateCPUHeap());
         }
@@ -70,7 +69,7 @@
         return CPUDescriptorHeapAllocation{baseCPUDescriptor, heapIndex};
     }
 
-    MaybeError NonShaderVisibleDescriptorAllocator::AllocateCPUHeap() {
+    MaybeError StagingDescriptorAllocator::AllocateCPUHeap() {
         D3D12_DESCRIPTOR_HEAP_DESC heapDescriptor;
         heapDescriptor.Type = mHeapType;
         heapDescriptor.NumDescriptors = mHeapSize;
@@ -98,7 +97,7 @@
         return {};
     }
 
-    void NonShaderVisibleDescriptorAllocator::Deallocate(CPUDescriptorHeapAllocation* allocation) {
+    void StagingDescriptorAllocator::Deallocate(CPUDescriptorHeapAllocation* allocation) {
         const uint32_t heapIndex = allocation->GetHeapIndex();
 
         ASSERT(heapIndex < mPool.size());
@@ -125,12 +124,11 @@
         allocation->Invalidate();
     }
 
-    uint32_t NonShaderVisibleDescriptorAllocator::GetSizeIncrement() const {
+    uint32_t StagingDescriptorAllocator::GetSizeIncrement() const {
         return mSizeIncrement;
     }
 
-    NonShaderVisibleDescriptorAllocator::Index
-    NonShaderVisibleDescriptorAllocator::GetFreeBlockIndicesSize() const {
+    StagingDescriptorAllocator::Index StagingDescriptorAllocator::GetFreeBlockIndicesSize() const {
         return ((mHeapSize * mSizeIncrement) / mBlockSize);
     }
 
diff --git a/src/dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h b/src/dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h
similarity index 76%
rename from src/dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h
rename to src/dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h
index 8cbb7e2..292fb10 100644
--- a/src/dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h
+++ b/src/dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef DAWNNATIVE_D3D12_NONSHADERVISIBLEDESCRIPTORALLOCATOR_H_
-#define DAWNNATIVE_D3D12_NONSHADERVISIBLEDESCRIPTORALLOCATOR_H_
+#ifndef DAWNNATIVE_D3D12_STAGINGDESCRIPTORALLOCATOR_H_
+#define DAWNNATIVE_D3D12_STAGINGDESCRIPTORALLOCATOR_H_
 
 #include "dawn_native/Error.h"
 
@@ -21,7 +21,7 @@
 
 #include <vector>
 
-// |NonShaderVisibleDescriptorAllocator| allocates a fixed-size block of descriptors from a CPU
+// |StagingDescriptorAllocator| allocates a fixed-size block of descriptors from a CPU
 // descriptor heap pool.
 // Internally, it manages a list of heaps using a fixed-size block allocator. The fixed-size
 // block allocator is backed by a list of free blocks (free-list). The heap is in one of two
@@ -34,14 +34,14 @@
 
     class Device;
 
-    class NonShaderVisibleDescriptorAllocator {
+    class StagingDescriptorAllocator {
       public:
-        NonShaderVisibleDescriptorAllocator() = default;
-        NonShaderVisibleDescriptorAllocator(Device* device,
-                                            uint32_t descriptorCount,
-                                            uint32_t heapSize,
-                                            D3D12_DESCRIPTOR_HEAP_TYPE heapType);
-        ~NonShaderVisibleDescriptorAllocator();
+        StagingDescriptorAllocator() = default;
+        StagingDescriptorAllocator(Device* device,
+                                   uint32_t descriptorCount,
+                                   uint32_t heapSize,
+                                   D3D12_DESCRIPTOR_HEAP_TYPE heapType);
+        ~StagingDescriptorAllocator();
 
         ResultOrError<CPUDescriptorHeapAllocation> AllocateCPUDescriptors();
 
@@ -75,4 +75,4 @@
 
 }}  // namespace dawn_native::d3d12
 
-#endif  // DAWNNATIVE_D3D12_NONSHADERVISIBLEDESCRIPTORALLOCATOR_H_
\ No newline at end of file
+#endif  // DAWNNATIVE_D3D12_STAGINGDESCRIPTORALLOCATOR_H_
\ No newline at end of file
diff --git a/src/tests/white_box/D3D12DescriptorHeapTests.cpp b/src/tests/white_box/D3D12DescriptorHeapTests.cpp
index 665bddd..455b454 100644
--- a/src/tests/white_box/D3D12DescriptorHeapTests.cpp
+++ b/src/tests/white_box/D3D12DescriptorHeapTests.cpp
@@ -17,8 +17,8 @@
 #include "dawn_native/Toggles.h"
 #include "dawn_native/d3d12/BindGroupLayoutD3D12.h"
 #include "dawn_native/d3d12/DeviceD3D12.h"
-#include "dawn_native/d3d12/NonShaderVisibleDescriptorAllocatorD3D12.h"
 #include "dawn_native/d3d12/ShaderVisibleDescriptorAllocatorD3D12.h"
+#include "dawn_native/d3d12/StagingDescriptorAllocatorD3D12.h"
 #include "utils/ComboRenderPipelineDescriptor.h"
 #include "utils/WGPUHelpers.h"
 
@@ -94,11 +94,11 @@
     wgpu::ShaderModule mSimpleFSModule;
 };
 
-class DummyNonShaderVisibleDescriptorAllocator {
+class DummyStagingDescriptorAllocator {
   public:
-    DummyNonShaderVisibleDescriptorAllocator(Device* device,
-                                             uint32_t descriptorCount,
-                                             uint32_t allocationsPerHeap)
+    DummyStagingDescriptorAllocator(Device* device,
+                                    uint32_t descriptorCount,
+                                    uint32_t allocationsPerHeap)
         : mAllocator(device,
                      descriptorCount,
                      allocationsPerHeap * descriptorCount,
@@ -116,7 +116,7 @@
     }
 
   private:
-    NonShaderVisibleDescriptorAllocator mAllocator;
+    StagingDescriptorAllocator mAllocator;
 };
 
 // Verify the shader visible heaps switch over within a single submit.
@@ -719,8 +719,7 @@
 TEST_P(D3D12DescriptorHeapTests, Single) {
     constexpr uint32_t kDescriptorCount = 4;
     constexpr uint32_t kAllocationsPerHeap = 3;
-    DummyNonShaderVisibleDescriptorAllocator allocator(mD3DDevice, kDescriptorCount,
-                                                       kAllocationsPerHeap);
+    DummyStagingDescriptorAllocator allocator(mD3DDevice, kDescriptorCount, kAllocationsPerHeap);
 
     CPUDescriptorHeapAllocation allocation = allocator.AllocateCPUDescriptors();
     EXPECT_EQ(allocation.GetHeapIndex(), 0u);
@@ -735,8 +734,7 @@
 TEST_P(D3D12DescriptorHeapTests, Sequential) {
     constexpr uint32_t kDescriptorCount = 4;
     constexpr uint32_t kAllocationsPerHeap = 3;
-    DummyNonShaderVisibleDescriptorAllocator allocator(mD3DDevice, kDescriptorCount,
-                                                       kAllocationsPerHeap);
+    DummyStagingDescriptorAllocator allocator(mD3DDevice, kDescriptorCount, kAllocationsPerHeap);
 
     // Allocate |kNumOfHeaps| worth.
     constexpr uint32_t kNumOfHeaps = 2;
@@ -766,8 +764,7 @@
 TEST_P(D3D12DescriptorHeapTests, ReuseFreedHeaps) {
     constexpr uint32_t kDescriptorCount = 4;
     constexpr uint32_t kAllocationsPerHeap = 25;
-    DummyNonShaderVisibleDescriptorAllocator allocator(mD3DDevice, kDescriptorCount,
-                                                       kAllocationsPerHeap);
+    DummyStagingDescriptorAllocator allocator(mD3DDevice, kDescriptorCount, kAllocationsPerHeap);
 
     constexpr uint32_t kNumofHeaps = 10;
 
@@ -810,8 +807,7 @@
 TEST_P(D3D12DescriptorHeapTests, AllocateDeallocateMany) {
     constexpr uint32_t kDescriptorCount = 4;
     constexpr uint32_t kAllocationsPerHeap = 25;
-    DummyNonShaderVisibleDescriptorAllocator allocator(mD3DDevice, kDescriptorCount,
-                                                       kAllocationsPerHeap);
+    DummyStagingDescriptorAllocator allocator(mD3DDevice, kDescriptorCount, kAllocationsPerHeap);
 
     std::list<CPUDescriptorHeapAllocation> list3;
     std::list<CPUDescriptorHeapAllocation> list5;