Fix bind groups not being applied

Fix for crbug.com/dawn/1049, where setting a pipeline without drawing can
prevent bind groups from being applied later. This occurs because the mask
for the pipeline is being saved but not its layout, because the bind
groups are never applied. This changes to only save the mask if the bind
groups are applied (the pipeline is used in a draw or dispatch).

Bug: dawn:1049
Change-Id: I4c7ae1125d1b6a06af90aea49a9dd1e4883f4826
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/60740
Reviewed-by: Austin Eng <enga@chromium.org>
Commit-Queue: Austin Eng <enga@chromium.org>

src/dawn_native/BindGroupTracker.h

@@ -60,13 +60,16 @@ namespace dawn_native {
 
         void OnSetPipeline(PipelineBase* pipeline) {
             mPipelineLayout = pipeline->GetLayout();
+        }
+
+      protected:
+        // The Derived class should call this before it applies bind groups.
+        void BeforeApply() {
             if (mLastAppliedPipelineLayout == mPipelineLayout) {
                 return;
             }
 
-            // Keep track of the bind group layout mask to avoid marking unused bind groups as
-            // dirty. This also allows us to avoid computing the intersection of the dirty bind
-            // groups and bind group layout mask in Draw or Dispatch which is very hot code.
+            // Use the bind group layout mask to avoid marking unused bind groups as dirty.
             mBindGroupLayoutsMask = mPipelineLayout->GetBindGroupLayoutsMask();
 
             // Changing the pipeline layout sets bind groups as dirty. If CanInheritBindGroups,
@@ -88,13 +91,15 @@ namespace dawn_native {
             }
         }
 
-      protected:
-        // The Derived class should call this when it applies bind groups.
-        void DidApply() {
+        // The Derived class should call this after it applies bind groups.
+        void AfterApply() {
             // Reset all dirty bind groups. Dirty bind groups not in the bind group layout mask
             // will be dirtied again by the next pipeline change.
             mDirtyBindGroups.reset();
             mDirtyBindGroupsObjectChangedOrIsDynamic.reset();
+            // Keep track of the last applied pipeline layout. This allows us to avoid computing
+            // the intersection of the dirty bind groups and bind group layout mask in next Draw
+            // or Dispatch (which is very hot code) until the layout is changed again.
             mLastAppliedPipelineLayout = mPipelineLayout;
         }
 

src/dawn_native/d3d12/CommandBufferD3D12.cpp

@@ -338,6 +338,8 @@ namespace dawn_native { namespace d3d12 {
         }
 
         MaybeError Apply(CommandRecordingContext* commandContext) {
+            BeforeApply();
+
             // Bindgroups are allocated in shader-visible descriptor heaps which are managed by a
             // ringbuffer. There can be a single shader-visible descriptor heap of each type bound
             // at any given time. This means that when we switch heaps, all other currently bound
@@ -389,7 +391,7 @@ namespace dawn_native { namespace d3d12 {
                                mDynamicOffsetCounts[index], mDynamicOffsets[index].data());
             }
 
-            DidApply();
+            AfterApply();
 
             return {};
         }

src/dawn_native/metal/CommandBufferMTL.mm

@@ -359,13 +359,14 @@ namespace dawn_native { namespace metal {
 
             template <typename Encoder>
             void Apply(Encoder encoder) {
+                BeforeApply();
                 for (BindGroupIndex index :
                      IterateBitSet(mDirtyBindGroupsObjectChangedOrIsDynamic)) {
                     ApplyBindGroup(encoder, index, ToBackend(mBindGroups[index]),
                                    mDynamicOffsetCounts[index], mDynamicOffsets[index].data(),
                                    ToBackend(mPipelineLayout));
                 }
-                DidApply();
+                AfterApply();
             }
 
           private:

src/dawn_native/opengl/CommandBufferGL.cpp

@@ -225,12 +225,13 @@ namespace dawn_native { namespace opengl {
             }
 
             void Apply(const OpenGLFunctions& gl) {
+                BeforeApply();
                 for (BindGroupIndex index :
                      IterateBitSet(mDirtyBindGroupsObjectChangedOrIsDynamic)) {
                     ApplyBindGroup(gl, index, mBindGroups[index], mDynamicOffsetCounts[index],
                                    mDynamicOffsets[index].data());
                 }
-                DidApply();
+                AfterApply();
             }
 
           private:

src/dawn_native/vulkan/CommandBufferVk.cpp

@@ -130,6 +130,7 @@ namespace dawn_native { namespace vulkan {
             void Apply(Device* device,
                        CommandRecordingContext* recordingContext,
                        VkPipelineBindPoint bindPoint) {
+                BeforeApply();
                 for (BindGroupIndex dirtyIndex :
                      IterateBitSet(mDirtyBindGroupsObjectChangedOrIsDynamic)) {
                     VkDescriptorSet set = ToBackend(mBindGroups[dirtyIndex])->GetHandle();
@@ -141,7 +142,7 @@ namespace dawn_native { namespace vulkan {
                         ToBackend(mPipelineLayout)->GetHandle(), static_cast<uint32_t>(dirtyIndex),
                         1, &*set, mDynamicOffsetCounts[dirtyIndex], dynamicOffset);
                 }
-                DidApply();
+                AfterApply();
             }
         };
 

src/tests/end2end/BindGroupTests.cpp

@@ -755,6 +755,110 @@ TEST_P(BindGroupTests, DrawThenChangePipelineAndBindGroup) {
     EXPECT_PIXEL_RGBA8_EQ(notFilled, renderPass.color, max, max);
 }
 
+// Test for crbug.com/dawn/1049, where setting a pipeline without drawing can prevent
+// bind groups from being applied later
+TEST_P(BindGroupTests, DrawThenChangePipelineTwiceAndBindGroup) {
+    // TODO(crbug.com/dawn/1055) find out why this test fails on Windows Intel D3D12 drivers.
+    DAWN_SUPPRESS_TEST_IF(IsIntel() && IsWindows() && IsD3D12());
+
+    utils::BasicRenderPass renderPass = utils::CreateBasicRenderPass(device, kRTSize, kRTSize);
+
+    // Create a bind group layout which uses a single dynamic uniform buffer.
+    wgpu::BindGroupLayout uniformLayout = utils::MakeBindGroupLayout(
+        device, {{0, wgpu::ShaderStage::Fragment, wgpu::BufferBindingType::Uniform, true}});
+
+    // Create a pipeline with pipeline layout (uniform, uniform, uniform).
+    wgpu::RenderPipeline pipeline0 =
+        MakeTestPipeline(renderPass,
+                         {wgpu::BufferBindingType::Uniform, wgpu::BufferBindingType::Uniform,
+                          wgpu::BufferBindingType::Uniform},
+                         {uniformLayout, uniformLayout, uniformLayout});
+
+    // Create a pipeline with pipeline layout (uniform).
+    wgpu::RenderPipeline pipeline1 = MakeTestPipeline(
+        renderPass, {wgpu::BufferBindingType::Uniform, wgpu::BufferBindingType::Uniform},
+        {uniformLayout, uniformLayout});
+
+    // Prepare color data.
+    // The first draw will use { color0, color1, color2 }.
+    // The second draw will use { color0, color1, color3 }.
+    // The pipeline uses additive color and alpha so the result of two draws should be
+    // { 2 * color0 + 2 * color1 + color2 + color3} = RGBAunorm(1, 1, 1, 1)
+    std::array<float, 4> color0 = {0.501, 0, 0, 0};
+    std::array<float, 4> color1 = {0, 0.501, 0, 0};
+    std::array<float, 4> color2 = {0, 0, 1, 0};
+    std::array<float, 4> color3 = {0, 0, 0, 1};
+
+    size_t color0Offset = 0;
+    size_t color1Offset = Align(color0Offset + sizeof(color0), kMinUniformBufferOffsetAlignment);
+    size_t color2Offset = Align(color1Offset + sizeof(color1), kMinUniformBufferOffsetAlignment);
+    size_t color3Offset = Align(color2Offset + sizeof(color2), kMinUniformBufferOffsetAlignment);
+
+    std::vector<uint8_t> data(color3Offset + sizeof(color3), 0);
+    memcpy(data.data(), color0.data(), sizeof(color0));
+    memcpy(data.data() + color1Offset, color1.data(), sizeof(color1));
+    memcpy(data.data() + color2Offset, color2.data(), sizeof(color2));
+    memcpy(data.data() + color3Offset, color3.data(), sizeof(color3));
+
+    // Create a uniform and storage buffer bind groups to bind the color data.
+    wgpu::Buffer uniformBuffer =
+        utils::CreateBufferFromData(device, data.data(), data.size(), wgpu::BufferUsage::Uniform);
+
+    wgpu::BindGroup uniformBindGroup =
+        utils::MakeBindGroup(device, uniformLayout, {{0, uniformBuffer, 0, 4 * sizeof(float)}});
+
+    wgpu::CommandEncoder encoder = device.CreateCommandEncoder();
+    wgpu::RenderPassEncoder pass = encoder.BeginRenderPass(&renderPass.renderPassInfo);
+
+    // Set the pipeline to (uniform, uniform, storage)
+    pass.SetPipeline(pipeline0);
+
+    // Set the first bind group to color0 in the dynamic uniform buffer.
+    uint32_t dynamicOffset = color0Offset;
+    pass.SetBindGroup(0, uniformBindGroup, 1, &dynamicOffset);
+
+    // Set the first bind group to color1 in the dynamic uniform buffer.
+    dynamicOffset = color1Offset;
+    pass.SetBindGroup(1, uniformBindGroup, 1, &dynamicOffset);
+
+    // Set the first bind group to color2 in the dynamic uniform buffer.
+    dynamicOffset = color2Offset;
+    pass.SetBindGroup(2, uniformBindGroup, 1, &dynamicOffset);
+
+    // This draw will internally apply bind groups for pipeline 0.
+    pass.Draw(3);
+
+    // When we set pipeline 1, which has no bind group at index 2 in its layout, it
+    // should not prevent bind group 2 from being used after reverting to pipeline 0.
+    // More specifically, internally the pipeline 1 layout should not be saved,
+    // because we never applied the bind groups via a Draw or Dispatch.
+    pass.SetPipeline(pipeline1);
+
+    // Set the second bind group to color3 in the dynamic uniform buffer.
+    dynamicOffset = color3Offset;
+    pass.SetBindGroup(2, uniformBindGroup, 1, &dynamicOffset);
+
+    // Revert to pipeline 0
+    pass.SetPipeline(pipeline0);
+
+    // Internally this will not re-apply the bind groups, because we already
+    // drew with this pipeline (setting pipeline 1 did not dirty the bind groups).
+    pass.Draw(3);
+
+    pass.EndPass();
+
+    wgpu::CommandBuffer commands = encoder.Finish();
+    queue.Submit(1, &commands);
+
+    RGBA8 filled(255, 255, 255, 255);
+    RGBA8 notFilled(0, 0, 0, 0);
+    uint32_t min = 1, max = kRTSize - 3;
+    EXPECT_PIXEL_RGBA8_EQ(filled, renderPass.color, min, min);
+    EXPECT_PIXEL_RGBA8_EQ(filled, renderPass.color, max, min);
+    EXPECT_PIXEL_RGBA8_EQ(filled, renderPass.color, min, max);
+    EXPECT_PIXEL_RGBA8_EQ(notFilled, renderPass.color, max, max);
+}
+
 // Regression test for crbug.com/dawn/408 where dynamic offsets were applied in the wrong order.
 // Dynamic offsets should be applied in increasing order of binding number.
 TEST_P(BindGroupTests, DynamicOffsetOrder) {