ShaderSerialization.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "Compression/CompressedBuffer.h"
#include "Memory/CompositeBuffer.h"
#include "Memory/SharedBuffer.h"
#include "Templates/Function.h"
#include "Templates/UniquePtr.h"
#include "UObject/NameTypes.h" 
 
class FMemoryReaderView;
class FMemoryWriter64;
class FShaderCodeResource;
 
#if WITH_EDITOR
namespace UE::DerivedData
{
    class FCacheRecord;
    struct FCacheKey;
}
#endif
 
/* Context object used for storing state and serialization parameterization related to shader/shadermap serialization. */
struct FShaderSerializeContext
{
    /* Default constructor, if called it's expected the archive pointer will be set by derived implementations. */
    FShaderSerializeContext() = default;
 
    /* Constructor which accepts an FArchive reference, used for simple serialization cases */
    FShaderSerializeContext(FArchive& InArchive) : Ar(&InArchive) {}
 
    virtual ~FShaderSerializeContext() = default;
 
    /* If this is overridden to return true, SerializeCode will be called to serialize shader code separately from the rest of the serialized object. */
    virtual bool EnableCustomCodeSerialize() { return false; }
 
    /* Optional function which must be implemented if EnableCustomCodeSerialize returns true; use to serialize shader code separately from the main object */
    virtual void SerializeCode(FShaderCodeResource& Resource, int32 Index) {};
 
    /* Optional function that can be used to reserve space for the given number of code objects by derived classes. */
    virtual void ReserveCode(int32 Count) {};
 
    /* Archive pointer which should be used for serializing the object, possibly excluding shader code if the EnableCustomCodeSerialize returns true */
    FArchive* Ar = nullptr;
 
    /* Flag indicating whether this serialization is a cooked load; used to change serialization behaviour for cooked data vs cached data */
    bool bLoadingCooked = false;
 
    /* FName of the asset which triggered the serialization, this is used only for diagnostic messages */
    FName SerializingAsset = NAME_None;
 
    /* Helper function to retrieve an FArchive reference; exists for convenience and validation purposes. */
    FArchive& GetMainArchive()
    {
        check(Ar);
        return *Ar;
    }
};
 
 
/* FShaderSerializeContext implementation used for saving to or loading from caches (either DDC or in-memory job cache) 
 * Note that this is just a base class of the below and should not be used directly.
 */
struct FShaderCacheSerializeContext : public FShaderSerializeContext
{
    FShaderCacheSerializeContext() : FShaderSerializeContext() 
    {
        ShaderCode = OwnedShaderCode;
        ShaderSymbols = OwnedShaderSymbols;
 
        checkf(ShaderCode.Num() == ShaderSymbols.Num(), TEXT("It is required to serialize a (possibly empty, but non-null) symbols buffer for every code buffer."));
    }
 
    virtual ~FShaderCacheSerializeContext() = default;
 
    /* Buffer which stores the main object data for a cache entry - i.e. a shadermap or job structure */
    FSharedBuffer ShaderObjectData;
 
    /* View on array of buffers which store the bytecode objects for the object, one per shader/stage.
     * Note that this may or may not point to the OwnedShaderCode array below, depending on usage.
     */
    TArrayView<FCompositeBuffer> ShaderCode;
 
    /* View on array of compressed buffers which store the symbols for the object, one per shader/stage.
     * Note that this may or may not point to the OwnedShaderSymbols array below, depending on usage.
     */
    TArrayView<FCompressedBuffer> ShaderSymbols;
 
    /* Array of code buffers actually owned by this context object; it is valid for this to be empty in cases
     * where the array of ShaderCode buffers is stored externally.
     */
    TArray<FCompositeBuffer> OwnedShaderCode;
 
    /* Array of symbol buffers actually owned by this context object; it is valid for this to be empty in cases
     * where the array of ShaderSymbols buffers is stored externally.
     */
    TArray<FCompressedBuffer> OwnedShaderSymbols;
 
    /* Subclasses of this type must implement custom code serialization function */
    virtual bool EnableCustomCodeSerialize() override { return true; }
 
    /* Get the total serialized size of data for this context; note that this will return 0 if called prior to the FSharedBuffers
     * being set (this is done in the derived implementations, see below).
     */
    int64 GetSerializedSize() const
    {
        int64 Size = 0u;
        if (ShaderObjectData)
        {
            Size += ShaderObjectData.GetSize();
 
            for (FCompositeBuffer& CodeBuf : ShaderCode)
            {
                Size += CodeBuf.GetSize();
            }
        }
        return Size;
    }
 
    /* Populates the given code array (transfering ownership) and resets the internal view to point to the new owning array's data */
    void MoveCode(TArray<FCompositeBuffer>& TargetCode, TArray<FCompressedBuffer>& TargetSymbols)
    {
        TargetCode = MoveTemp(OwnedShaderCode);
        ShaderCode = TargetCode;
        TargetSymbols = MoveTemp(OwnedShaderSymbols);
        ShaderSymbols = TargetSymbols;
 
        checkf(ShaderCode.Num() == ShaderSymbols.Num(), TEXT("It is required to serialize a (possibly empty, but non-null) symbols buffer for every code buffer."));
    }
 
    /* Returns true if there is valid serialized data referenced by this context. */
    bool HasData() const { return ShaderObjectData && !ShaderCode.IsEmpty(); }
};
 
/* Implementation of FShaderCacheSerializeContext used for saving data to caches. */
struct FShaderCacheSaveContext : public FShaderCacheSerializeContext
{
    /* Default constructor; sets up base class FArchive pointing to the owned memory writer. */
    RENDERCORE_API FShaderCacheSaveContext();
 
    virtual ~FShaderCacheSaveContext() = default;
 
    /* Converts the raw serialized object data into the ShaderObjectData FSharedBuffer.
     * Note that this is called by BuildCacheRecord as well, calls subsequent to the first will have no effect.
     */
    RENDERCORE_API void Finalize();
 
    /* Overridden code serialize implementation. */
    RENDERCORE_API virtual void SerializeCode(FShaderCodeResource& Resource, int32 Index) override;
 
    /* Overridden code reserve implementation. */
    RENDERCORE_API virtual void ReserveCode(int32 Count) override;
 
#if WITH_EDITOR
    /* Helper function to generate a DDC record from the data serialized using this context. Serialization must
     * be executed prior to calling this function.
     */
    RENDERCORE_API UE::DerivedData::FCacheRecord BuildCacheRecord(const UE::DerivedData::FCacheKey& Key);
#endif
 
    /* Call to reset internal state, this allows re-using internal allocations to store other objects (as an optimization). */
    RENDERCORE_API void Reset();
 
    /* Data that will be written to when serializing */
    TArray64<uint8> ShaderObjectRawData;
 
    /* FArchive that is passed to the base class, which writes to ShaderObjectRawData */
    TUniquePtr<FMemoryWriter64> Writer;
};
 
/* Implementation of FShaderCacheSerializeContext used for loading data from caches. */
struct FShaderCacheLoadContext : public FShaderCacheSerializeContext
{
    /*  Default constructor, use when array of code buffers will be allocated via ReadFromRecord */
    FShaderCacheLoadContext() = default;
 
    /* Constructor which references buffers (and an array of code buffers) owned elsewhere, does not allocate the OwnedShaderCode array */
    RENDERCORE_API FShaderCacheLoadContext(FSharedBuffer ShaderObjectData, TArrayView<FCompositeBuffer> CodeBuffers, TArrayView<FCompressedBuffer> SymbolBuffers);
 
    virtual ~FShaderCacheLoadContext() = default;
 
    /* Resets internal state to point to the given buffers and recreates the owned memoryreader */
    RENDERCORE_API void Reset(FSharedBuffer ShaderObjectData, TArrayView<FCompositeBuffer> CodeBuffers, TArrayView<FCompressedBuffer> SymbolBuffers);
 
    /* Call to reset reader to start position so the same load context can be used to populate multiple objects. */
    RENDERCORE_API void Reuse();
 
    /* Overridden code serialize implementation. */
    RENDERCORE_API virtual void SerializeCode(FShaderCodeResource& Resource, int32 Index) override;
 
#if WITH_EDITOR
    /* Helper function to populate the internal state (FSharedBuffers defined on FShaderCacheSerializeContext) from a DDC record. */
    RENDERCORE_API void ReadFromRecord(const UE::DerivedData::FCacheRecord& CacheRecord, bool bIsPersistent = false);
#endif
 
    /* FArchive passed to the base class; this will be used to read data from the FShaderCacheSerializeContext::ShaderObjectData shared buffer */
    TUniquePtr<FMemoryReaderView> Reader;
};