DeviceContextD3D11: fixed issue with PS UAVs being incorrectly unbound

Author: TheMostDiligent

Graphics/GraphicsEngineD3D11/src/DeviceContextD3D11Impl.cpp

@@ -434,18 +434,24 @@ void DeviceContextD3D11Impl::BindShaderResources(Uint32 BindSRBMask)
         PixelShaderUAVBindMode::Clear :
         PixelShaderUAVBindMode::Keep;
 
+    Uint32 UpToDateSRBMask = m_BindInfo.ActiveSRBMask & ~BindSRBMask;
     while (BindSRBMask != 0)
     {
         Uint32 SignBit = ExtractLSB(BindSRBMask);
-        Uint32 sign    = PlatformMisc::GetLSB(SignBit);
-        VERIFY_EXPR(sign < m_pPipelineState->GetResourceSignatureCount());
-        const D3D11ShaderResourceCounters& BaseBindings = m_pPipelineState->GetBaseBindings(sign);
+        Uint32 SignIdx = PlatformMisc::GetLSB(SignBit);
+        VERIFY_EXPR(SignIdx < m_pPipelineState->GetResourceSignatureCount());
+        const D3D11ShaderResourceCounters& BaseBindings = m_pPipelineState->GetBaseBindings(SignIdx);
 
 #ifdef DILIGENT_DEVELOPMENT
-        m_BindInfo.BaseBindings[sign] = BaseBindings;
+        m_BindInfo.BaseBindings[SignIdx] = BaseBindings;
 #endif
-        const ShaderResourceCacheD3D11* pResourceCache = m_BindInfo.ResourceCaches[sign];
-        DEV_CHECK_ERR(pResourceCache != nullptr, "Shader resource cache at index ", sign, " is null.");
+        const ShaderResourceCacheD3D11* pResourceCache = m_BindInfo.ResourceCaches[SignIdx];
+        if (pResourceCache == nullptr)
+        {
+            DEV_ERROR("Shader resource cache at index ", SignIdx, " is null.");
+            continue;
+        }
+
         if (m_BindInfo.StaleSRBMask & SignBit)
         {
             // Bind all cache resources
@@ -467,6 +473,30 @@ void DeviceContextD3D11Impl::BindShaderResources(Uint32 BindSRBMask)
             BindDynamicCBs(*pResourceCache, BaseBindings);
         }
     }
+
+    if (PsUavBindMode == PixelShaderUAVBindMode::Clear)
+    {
+        // Check if SRBs that are not in the BindSRBMask contain UAVs.
+        // In this case we should keep UAVs, not clear them.
+        while (UpToDateSRBMask != 0)
+        {
+            Uint32 SignBit = ExtractLSB(UpToDateSRBMask);
+            Uint32 SignIdx = PlatformMisc::GetLSB(SignBit);
+            VERIFY_EXPR(SignIdx < m_pPipelineState->GetResourceSignatureCount());
+            const ShaderResourceCacheD3D11* pResourceCache = m_BindInfo.ResourceCaches[SignIdx];
+            if (pResourceCache == nullptr)
+            {
+                DEV_ERROR("Shader resource cache at index ", SignIdx, " is null.");
+                continue;
+            }
+            if (pResourceCache->GetUAVCount(PSInd) > 0)
+            {
+                PsUavBindMode = PixelShaderUAVBindMode::Keep;
+                break;
+            }
+        }
+    }
+
     m_BindInfo.StaleSRBMask &= ~m_BindInfo.ActiveSRBMask;
 
     if (PsUavBindMode == PixelShaderUAVBindMode::Bind)

Tests/DiligentCoreAPITest/include/InlineShaders/ComputeShaderTestHLSL.h

@@ -106,6 +106,28 @@ void main(in VSOutput VSOut)
 )"
 };
 
+
+const std::string FillTexturePS2{
+R"(
+struct VSOutput
+{
+    float4 Pos : SV_Position;
+};
+
+cbuffer Constants
+{
+	float4 g_Color;
+};
+
+RWTexture2D</*format=rgba8*/ float4> g_tex2DUAV : register(u0);
+
+void main(in VSOutput VSOut)
+{
+	g_tex2DUAV[uint2(VSOut.Pos.xy)] = float4(float2(uint2(VSOut.Pos.xy) % 256u) / 256.0, 0.0, 1.0) * g_Color;
+}
+)"
+};
+
 // clang-format on
 
 } // namespace HLSL

Tests/DiligentCoreAPITest/src/ComputeShaderTest.cpp

@@ -463,4 +463,150 @@ TEST(ComputeShaderTest, FillTexturePS_InRenderPass)
     TestFillTexturePS(true);
 }
 
+TEST(ComputeShaderTest, FillTexturePS_Signatures)
+{
+    GPUTestingEnvironment*  pEnv       = GPUTestingEnvironment::GetInstance();
+    IRenderDevice*          pDevice    = pEnv->GetDevice();
+    const RenderDeviceInfo& DeviceInfo = pDevice->GetDeviceInfo();
+    if (!DeviceInfo.Features.ComputeShaders)
+    {
+        GTEST_SKIP() << "Compute shaders are not supported by this device";
+    }
+
+    ISwapChain*     pSwapChain = pEnv->GetSwapChain();
+    IDeviceContext* pContext   = pEnv->GetDeviceContext();
+
+    const SwapChainDesc& SCDesc = pSwapChain->GetDesc();
+
+    GPUTestingEnvironment::ScopedReset EnvironmentAutoReset;
+
+    RefCntAutoPtr<ITestingSwapChain> pTestingSwapChain{pSwapChain, IID_TestingSwapChain};
+    if (!pTestingSwapChain)
+    {
+        GTEST_SKIP() << "Compute shader test requires testing swap chain";
+    }
+
+    RefCntAutoPtr<ITextureView> pDummyRTV;
+    if (DeviceInfo.IsWebGPUDevice())
+    {
+        // WebGPU does not support render passes without attachments (https://github.com/gpuweb/gpuweb/issues/503)
+        RefCntAutoPtr<ITexture> pDummyTex = pEnv->CreateTexture("Dummy render target", SCDesc.ColorBufferFormat, BIND_RENDER_TARGET, SCDesc.Width, SCDesc.Height);
+        ASSERT_TRUE(pDummyTex != nullptr);
+        pDummyRTV = pDummyTex->GetDefaultView(TEXTURE_VIEW_RENDER_TARGET);
+    }
+
+    pContext->Flush();
+    pContext->InvalidateState();
+
+    ComputeShaderReference(pSwapChain);
+
+    ShaderCreateInfo ShaderCI;
+    ShaderCI.SourceLanguage = SHADER_SOURCE_LANGUAGE_HLSL;
+    ShaderCI.ShaderCompiler = pEnv->GetDefaultCompiler(ShaderCI.SourceLanguage);
+
+    ShaderCI.Desc       = {"Compute shader test - FillTextureVS", SHADER_TYPE_VERTEX, true};
+    ShaderCI.EntryPoint = "main";
+    ShaderCI.Source     = HLSL::FillTextureVS.c_str();
+    RefCntAutoPtr<IShader> pVS;
+    pDevice->CreateShader(ShaderCI, &pVS);
+    ASSERT_NE(pVS, nullptr);
+
+    ShaderCI.Desc       = {"Compute shader test - FillTexturePS", SHADER_TYPE_PIXEL, true};
+    ShaderCI.EntryPoint = "main";
+    ShaderCI.Source     = HLSL::FillTexturePS2.c_str();
+    RefCntAutoPtr<IShader> pPS;
+    pDevice->CreateShader(ShaderCI, &pPS);
+    ASSERT_NE(pPS, nullptr);
+
+    RefCntAutoPtr<IPipelineResourceSignature> pSignature0;
+    {
+        PipelineResourceSignatureDescX SignDesc{"ComputeShaderTest.FillTexturePS_InRenderPass - Signature 0"};
+        SignDesc.AddResource(SHADER_TYPE_PIXEL, "Constants", 1u, SHADER_RESOURCE_TYPE_CONSTANT_BUFFER, SHADER_RESOURCE_VARIABLE_TYPE_MUTABLE);
+        pDevice->CreatePipelineResourceSignature(SignDesc, &pSignature0);
+        ASSERT_NE(pSignature0, nullptr);
+    }
+
+    RefCntAutoPtr<IPipelineResourceSignature> pSignature1;
+    {
+        PipelineResourceSignatureDescX SignDesc{"ComputeShaderTest.FillTexturePS_InRenderPass - Signature 1"};
+        SignDesc.AddResource(SHADER_TYPE_PIXEL, "g_tex2DUAV", 1u, SHADER_RESOURCE_TYPE_TEXTURE_UAV, SHADER_RESOURCE_VARIABLE_TYPE_MUTABLE, PIPELINE_RESOURCE_FLAG_NONE,
+                             WebGPUResourceAttribs{WEB_GPU_BINDING_TYPE_WRITE_ONLY_TEXTURE_UAV, RESOURCE_DIM_TEX_2D, TEX_FORMAT_RGBA8_UNORM});
+        SignDesc.SetBindingIndex(1);
+        pDevice->CreatePipelineResourceSignature(SignDesc, &pSignature1);
+        ASSERT_NE(pSignature1, nullptr);
+    }
+
+    GraphicsPipelineStateCreateInfoX PSOCreateInfo{"Compute shader test - output from PS"};
+
+    PSOCreateInfo.AddSignature(pSignature0);
+    PSOCreateInfo.AddSignature(pSignature1);
+
+    PSOCreateInfo.GraphicsPipeline.RasterizerDesc.CullMode      = CULL_MODE_NONE;
+    PSOCreateInfo.GraphicsPipeline.DepthStencilDesc.DepthEnable = False;
+
+    PSOCreateInfo.pVS = pVS;
+    PSOCreateInfo.pPS = pPS;
+
+    if (DeviceInfo.IsWebGPUDevice())
+    {
+        PSOCreateInfo.GraphicsPipeline.NumRenderTargets                                 = 1;
+        PSOCreateInfo.GraphicsPipeline.RTVFormats[0]                                    = SCDesc.ColorBufferFormat;
+        PSOCreateInfo.GraphicsPipeline.BlendDesc.RenderTargets[0].RenderTargetWriteMask = COLOR_MASK_NONE;
+    }
+
+    RefCntAutoPtr<IPipelineState> pPSO;
+    pDevice->CreateGraphicsPipelineState(PSOCreateInfo, &pPSO);
+    ASSERT_NE(pPSO, nullptr);
+
+    const float4 Zero{0};
+    const float4 One{1};
+
+    RefCntAutoPtr<IBuffer> pBuffer0 = pEnv->CreateBuffer({"ComputeShaderTest.FillTexturePS_InRenderPass - Buffer 0", sizeof(Zero), BIND_UNIFORM_BUFFER}, &Zero);
+    ASSERT_NE(pBuffer0, nullptr);
+    RefCntAutoPtr<IBuffer> pBuffer1 = pEnv->CreateBuffer({"ComputeShaderTest.FillTexturePS_InRenderPass - Buffer 1", sizeof(One), BIND_UNIFORM_BUFFER}, &One);
+    ASSERT_NE(pBuffer1, nullptr);
+
+    RefCntAutoPtr<IShaderResourceBinding> pSRB0[2];
+    pSignature0->CreateShaderResourceBinding(&pSRB0[0], true);
+    ASSERT_NE(pSRB0[0], nullptr);
+    pSRB0[0]->GetVariableByName(SHADER_TYPE_PIXEL, "Constants")->Set(pBuffer0);
+
+    pSignature0->CreateShaderResourceBinding(&pSRB0[1], true);
+    ASSERT_NE(pSRB0[1], nullptr);
+    pSRB0[1]->GetVariableByName(SHADER_TYPE_PIXEL, "Constants")->Set(pBuffer1);
+
+
+    RefCntAutoPtr<IShaderResourceBinding> pSRB1;
+    pSignature1->CreateShaderResourceBinding(&pSRB1, true);
+    ASSERT_NE(pSRB1, nullptr);
+    pSRB1->GetVariableByName(SHADER_TYPE_PIXEL, "g_tex2DUAV")->Set(pTestingSwapChain->GetCurrentBackBufferUAV());
+
+    ITextureView* pRTVs[] = {pSwapChain->GetCurrentBackBufferRTV()};
+    pContext->SetRenderTargets(1, pRTVs, pSwapChain->GetDepthBufferDSV(), RESOURCE_STATE_TRANSITION_MODE_TRANSITION);
+
+    pContext->SetPipelineState(pPSO);
+    pContext->CommitShaderResources(pSRB0[0], RESOURCE_STATE_TRANSITION_MODE_TRANSITION);
+    pContext->CommitShaderResources(pSRB1, RESOURCE_STATE_TRANSITION_MODE_TRANSITION);
+
+    if (DeviceInfo.IsWebGPUDevice())
+    {
+        ITextureView* pDummyRTVs[] = {pDummyRTV};
+        pContext->SetRenderTargets(1, pDummyRTVs, nullptr, RESOURCE_STATE_TRANSITION_MODE_TRANSITION);
+    }
+    else
+    {
+        pContext->SetRenderTargets(0, nullptr, nullptr, RESOURCE_STATE_TRANSITION_MODE_NONE);
+    }
+
+    Viewport VP{SCDesc};
+    pContext->SetViewports(1, &VP, SCDesc.Width, SCDesc.Height);
+
+    pContext->Draw(DrawAttribs{3, DRAW_FLAG_VERIFY_ALL});
+
+    pContext->CommitShaderResources(pSRB0[1], RESOURCE_STATE_TRANSITION_MODE_TRANSITION);
+    pContext->Draw(DrawAttribs{3, DRAW_FLAG_VERIFY_ALL});
+
+    pSwapChain->Present();
+}
+
 } // namespace

Tests/GPUTestFramework/include/GPUTestingEnvironment.hpp

@@ -136,6 +136,8 @@ class GPUTestingEnvironment : public TestingEnvironment
 
     RefCntAutoPtr<ITexture> CreateTexture(const char* Name, TEXTURE_FORMAT Fmt, BIND_FLAGS BindFlags, Uint32 Width, Uint32 Height, const void* pInitData = nullptr);
 
+    RefCntAutoPtr<IBuffer> CreateBuffer(const BufferDesc& Desc, const void* pInitData = nullptr);
+
     RefCntAutoPtr<ISampler> CreateSampler(const SamplerDesc& Desc);
 
     void            SetDefaultCompiler(SHADER_COMPILER compiler);

Tests/GPUTestFramework/src/GPUTestingEnvironment.cpp

@@ -747,6 +747,17 @@ RefCntAutoPtr<ITexture> GPUTestingEnvironment::CreateTexture(const char* Name, T
     return pTexture;
 }
 
+RefCntAutoPtr<IBuffer> GPUTestingEnvironment::CreateBuffer(const BufferDesc& Desc, const void* pInitData)
+{
+    RefCntAutoPtr<IBuffer> pBuffer;
+
+    BufferData InitData{pInitData, Desc.Size};
+    m_pDevice->CreateBuffer(Desc, pInitData ? &InitData : nullptr, &pBuffer);
+    VERIFY_EXPR(pBuffer != nullptr);
+
+    return pBuffer;
+}
+
 RefCntAutoPtr<ISampler> GPUTestingEnvironment::CreateSampler(const SamplerDesc& Desc)
 {
     RefCntAutoPtr<ISampler> pSampler;