Generate Hi-Z buffer using single PSO and single UAV

Question

I'm currently implementing SSR in my engine and I've created Hi-Z buffer using multiple PSOs in way presented below:

1. Create descriptor table with multiple UAV entries (this case - 3 entries for Mip 0, Mip 1, Mip 2). Each of this will be bound to appropriate UAV in shader
2. Create PSO for each compute shader function depending on which mip is currently being generated
3. In runtime, switch PSO and dispatch CS for each shader
4. In result, each mip of Hi-Z buffer texture is filled with correct data. However, there is PSO switch before each dispatch which seems to be costly

I want to get rid of PSO switch (although, I'm not sure if that's possible because I have two functions - one is generating mip 0 by mapping depth from compute space to view space, other one is standard Hi-Z mip generation) and use single UAV to read data from and save data to.

If it is not possible, I would appreciate tips how to implement cleaner solution, because what I'm currently using seems a little ugly.

---

Relevant code:

A - create PSO for each UAV mip

// Create PSO for mip 0
pipelineStateStream.pRootSignature = m_rootSignatureHiZ.Get();
pipelineStateStream.CS = CD3DX12_SHADER_BYTECODE(computeShader.Get());
{
    D3D12_PIPELINE_STATE_STREAM_DESC pipelineStateStreamDesc = { sizeof(PipelineStateStream), &pipelineStateStream };
    ThrowIfFailed(m_device->CreatePipelineState(&pipelineStateStreamDesc, IID_PPV_ARGS(&m_pipelineStateHiZMipZero))); 
}

// Create PSO for mip 1
Compile_Shader(L"Shaders/CS_HiZ.hlsl", NULL, D3D_COMPILE_STANDARD_FILE_INCLUDE, "generateHiZMip1", "cs_5_1", 0, 0, &computeShader);
//...

// Create PSO for mip 2
Compile_Shader(L"Shaders/CS_HiZ.hlsl", NULL, D3D_COMPILE_STANDARD_FILE_INCLUDE, "generateHiZMip2", "cs_5_1", 0, 0, &computeShader);
//...

B - in runtime, switch PSOs and dispatch compute shaders

// Generate Hi-Z mip 0 (transform depth buffer from Compute-Space (CS) to View-Space (VS))
{
    m_commandList->ResourceBarrier(1, &CD3DX12_RESOURCE_BARRIER::Transition(m_hiZBuffer.Get(), D3D12_RESOURCE_STATE_COMMON, D3D12_RESOURCE_STATE_UNORDERED_ACCESS));
    m_commandList->Dispatch(m_windowSize.x / 16, m_windowSize.y / 16, 1);
    m_commandList->ResourceBarrier(1, &CD3DX12_RESOURCE_BARRIER::Transition(m_hiZBuffer.Get(), D3D12_RESOURCE_STATE_UNORDERED_ACCESS, D3D12_RESOURCE_STATE_COMMON));
}

// Generate additional mips for Hi-Z buffer (in this implementation we use mip 1 and mip 2)
m_commandList->SetPipelineState(m_pipelineStateHiZMipOne.Get());
m_commandList->Dispatch(m_windowSize.x / 32, m_windowSize.y / 32, 1); // Mip 1

m_commandList->SetPipelineState(m_pipelineStateHiZMipTwo.Get());
m_commandList->Dispatch(m_windowSize.x / 64, m_windowSize.y / 64, 1); // Mip 2

C - full compute shader, Hi-Z generation HLSL code

Texture2D g_depthBuffer : register(t0);

RWTexture2D<float> hiZMip0 : register(u0);
RWTexture2D<float> hiZMip1 : register(u1);
RWTexture2D<float> hiZMip2 : register(u2);

[numthreads(16, 16, 1)]
void generateHiZMip0(uint3 index : SV_DispatchThreadID)
{
    float depthCS = g_depthBuffer.Load(int3(index.xy, 0));
    hiZMip0[index.xy] = depthCStoVS(depthCS);
}

[numthreads(16, 16, 1)]
void generateHiZMip1(uint3 index : SV_DispatchThreadID)
{
    float a = hiZMip0.Load(int3(index.xy * 2 + int2(0, 0), 0));
    float b = hiZMip0.Load(int3(index.xy * 2 + int2(1, 0), 0));
    float c = hiZMip0.Load(int3(index.xy * 2 + int2(0, 1), 0));
    float d = hiZMip0.Load(int3(index.xy * 2 + int2(1, 1), 0));
        
    hiZMip1[index.xy] = min4(a.x, b.x, c.x, d.x);
}

[numthreads(16, 16, 1)]
void generateHiZMip2(uint3 index : SV_DispatchThreadID)
{
    float a = hiZMip1.Load(int3(index.xy * 2 + int2(0, 0), 0));
    float b = hiZMip1.Load(int3(index.xy * 2 + int2(1, 0), 0));
    float c = hiZMip1.Load(int3(index.xy * 2 + int2(0, 1), 0));
    float d = hiZMip1.Load(int3(index.xy * 2 + int2(1, 1), 0));
        
    hiZMip2[index.xy] = min4(a.x, b.x, c.x, d.x);
}

Answer 1

The code for mip1 and mip2 (and higher mips if you have them—SSR ray marching can benefit a lot from higher mips) is identical except for which mip is being read and written; you don't need separate PSOs for those. Instead of creating bindings for mip0, mip1, mip2, you could rather create bindings for "source mip" and "dest mip" and set the descriptors for each dispatch to point to the appropriate mips.

Also note that you'll need a UAV barrier between these dispatches since each one will be reading from the mip written by the previous one.

The overall sequence would be like:

switch to PSO for mip0
dispatch mip0 (reads g_depthBuffer, writes mip0)
barrier on mip0
switch to shared PSO for all higher mips
dispatch mip1 (src = mip0, dest = mip1)
barrier on mip1
dispatch mip2 (src = mip1, dest = mip2)
barrier on mip2
...

If the remaining PSO switch is still too expensive, you could also try putting it all in one shader and branching on a constant buffer value based on whether it's mip0 or a higher mip. So the shader code would look something like:

cbuffer { bool isFirstMip; }
void generateHiZ(uint3 index : SV_DispatchThreadID)
{
    [branch] if (isFirstMip)
        generateHiZMip0(index);
    else
        generateHiZHigherMips(index);
}

As this branch is 100% coherent within a dispatch, it should not be very expensive (especially if the bool is in a root constant), and might well be faster than incurring the PSO switch.

---

A further optimization you can make is to generate more than 1 mip per dispatch. With a threadgroup size of 16x16 it is possible to generate up to 5 mips in one pass, using one source and 5 dest UAVs. You would loop over the 5 mips in the shader and do a memory barrier at the end of each iteration, ensuring all threads in the group have finished their UAV writes. The first iteration will read 32x32 and write 16x16 (over the whole threadgroup), then will write 8x8, 4x4, 2x2, and 1x1.

It also might be faster to use threadgroup shared memory after the first iteration, rather than re-reading it from the UAV, but this would need to be profiled to be sure (the UAV reads will be hitting cache, so may be fast enough on their own).

Then, you would only need to invoke one dispatch for each 5 mips to generate. Two dispatches gives you up to 10 mips which is probably enough for SSR.