Get Data back from Compute Shader

Development Rendering
1

Hi everyone.

I’ve created a really simple test program to use a compute shader within UE4 to do some very very basic processing. **I’d like to be able to get the processed data back into the CPU side of things. **I have a simple struct defined from which I create a Structured Buffer.

So, first I create the data:


struct MyStruct{
	int i;
	float f;
};


int dim = 10;
TResourceArray<MyStruct> data;
MyStruct init;
init.i = 0;
init.f = 0.0f;

data.Init(init, dim);

for (int i = 0; i < dim; i++){
	data*.i = i;
}

Onwards to the creation of the StructuredBuffer:



//this can be used to set initial data of the buffer
FRHIResourceCreateInfo CreateInfo;
CreateInfo.ResourceArray = &data;

FStructuredBufferRHIRef StructResource = RHICreateStructuredBuffer(
	sizeof(MyStruct),
	data.Num() * sizeof(MyStruct),
	BUF_UnorderedAccess | BUF_ShaderResource,
	CreateInfo
	);

Afterwards the steps aren’t totally clear to me, but here’s what I figured would work:

I get the UAV for the data:


FUnorderedAccessViewRHIRef StructUAV = RHICreateUnorderedAccessView(StructResource, false, false);

I get my custom compute shader:



FRHICommandListImmediate& RHICmdList = GRHICommandList.GetImmediateCommandList();

TShaderMapRef<MyCS> MyCS(GetGlobalShaderMap());
FComputeShaderRHIParamRef CSRef = MyCS->GetComputeShader();
RHICmdList.SetComputeShader(CSRef);

I set the UAV “into” the shader:



//RWBuffer is a FShaderResourceParameter defined and bound in the custom CS.
RHICmdList.SetUAVParameter(CSRef, MyCS->RWBuffer.GetBaseIndex(), StructUAV);

Finally, I dispatch the compute shader:


DispatchComputeShader(RHICmdList, *MyCS, 2, 1, 1);

It compiles and doesn’t crash. However, I’ve rummaged through the RHI methods and whatnot and haven’t found anything that could let me access the data I processed in the compute shader. Does anyone have an idea on how to achieve this? I know it is usual to do it with output textures, for example, but I imagine there would be a more direct way (DirectX has a CopyResource function, for instance…)

Thanks in advance! :slight_smile:

NOTES:

I am using the ENQUEUE_UNIQUE_RENDER_COMMAND macro to send stuff to the rendering thread. I just haven’t shown it for the sake of brevity.
The compute shader has a DECLARE_SHADER_TYPE and IMPLEMENT_SHADER_TYPE correctly set up (I think). If need be, I may also post the CS code here.

2

I’ve just managed to do this myself only last night.

Here is a snippet of the code I came up with…



class FShaderD
{
    TArray<FVector4> ComputedColors;    // My container in main memory. CPU Accessible.
    FTexture2DRHIRef  TexBuf;           // Structured Buffer/ Texture
    FUnorderedAccessViewRHIRef dUAV;    // Buffer UAV for shader output
    
    virtual void DebugColors();         // Read back from GPU memory.
    virtual void DoCompute();           // Queue shader
    virtual void InitBuffer(); 
    
};

Setup the buffers for CPU and GPU side.



void FShaderD::InitBuffer()
{
        //init TArray
        ComputedColors.Init(256)
        TResourceArray<FVector4> buffData;
        buffData.SetAllowCPUAccess(true);//Needed?
        //Init temp resource buffer
        buffData.Init(FVector4(1.f, 0.f, 0.f, 1.f), 256);

        //uint32  float4_stride = 4 * sizeof(float); //Size of elements
        FRHIResourceCreateInfo ResourceCreateInfo;
        ResourceCreateInfo.ResourceArray = &buffData;
        TexBuf = RHICreateTexture2D(16,16,PF_A32B32G32R32F,1,1,(TexCreate_ShaderResource | TexCreate_UAV),ResourceCreateInfo);
        dUAV = RHICreateUnorderedAccessView(TexBuf);
}

Get the data back from the GPU. Called before the shader because otherwise crashes. Not investigated that yet. The second run gets the data from the first run of the shader.
Note: I think this does not return data right away, it has to run on the render thread.



void FShaderD::DebugColors()
{
    // Texture is valid ?
    if (IsValidRef(TexBuf) && !jobdone)
    {
        //queue render command.  Write the texture to TArray<FVector4> ComputedColors
        ENQUEUE_UNIQUE_RENDER_COMMAND_THREEPARAMETER(
            DoDebugColorsCommand,
            TArray<FVector4>&, dest, ComputedColors,
            FTexture2DRHIRef, source, TexBuf,
            bool&, done, jobdone,
            {
                uint32 stride; //ignored
                // Lock texture for reading
                uint8* TextureData = (uint8*)RHILockTexture2D(source,0,RLM_ReadOnly,stride,false);
                // Reference pointer to first element for our destination ComputedColors
                uint8* ArrayData = (uint8 *)dest.GetData();
                // Copy from GPU to main memory
                FMemory::Memcpy(ArrayData, TextureData, GPixelFormats[PF_A32B32G32R32F].BlockBytes * 256);
                //Unlock texture
                RHIUnlockTexture2D(source,0, false);

                done = true;

            });
    }
}

The shader call



void FShaderD::DoCompute()
{
    //queue render command: compute shader
    ENQUEUE_UNIQUE_RENDER_COMMAND_ONEPARAMETER(
        DColorCommand,
        FUnorderedAccessViewRHIRef&, destUAV, dUAV,
        {
            //test Values
            FVector4 myColor = FVector4(0.3f,0.1f,0.1f,1.f);
            FVector4 myDim = FVector4(15.f,15.f,0.f,0.f);
            FVector4 myInclude = FVector4(0.f,0.f, 15.f,15.f);
            
            //4.5 req GetGlobalShaderMap(GRHIFeatureLevel),  master GRHIFeatureLevel_DEPRECATED; GRHIShaderPlatform_DEPRECATED

            TShaderMapRef<FShaderCS> TestCS(GetGlobalShaderMap());
    
            TestCS->SetParameters(RHICmdList, myColor, myDim, myInclude);
            TestCS->SetOutput(RHICmdList, destUAV); 
            RHICmdList.SetComputeShader(TestCS->GetComputeShader());

    
            RHICmdList.DispatchComputeShader(8, 8, 1);
            TestCS->UnbindBuffers(RHICmdList);
        });

}
3

Hey James! You’re a lifesaver. I had to adapt your code just a little bit to what I wanted and I managed to get it working as I intended. Next steps are things like memory barriers but I should be able to figure it out from now on (FlushRenderingCommands might be want I want, actually).

Thanks so much for your trouble! :smiley: Good work!

4

No trouble at all, we’ve just been working in parallel, though I can’t say there was no trouble getting to this point. I am certain my graphics drivers crashed at least 30 times.

Anyway, glad everything worked out.