You've already forked UnrealEngineUWP
mirror of
https://github.com/izzy2lost/UnrealEngineUWP.git
synced 2026-03-26 18:15:20 -07:00
61447ea9f5
Problem: DDC bucket are assume to be stable per system, having asset based bucket is making some stats tracking harder and creates some noise in storage. Solution: Use the RuntimeName as a bucket instead of the FileId from the related asset. Note: Fix proposed by Stefan.Boberg via 28219309 Test: NNE editor 64 : style transfer demo + nne all test NNE standalone 64 : style transfer demo + nne smoke test Compiled Unity. #rb Stefan.Boberg, nico.ranieri #rnx [CL 28266828 by florent guinier in ue5-main branch]
491 lines
15 KiB
C++
491 lines
15 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#include "NNEModelData.h"
|
|
|
|
#include "EditorFramework/AssetImportData.h"
|
|
#include "NNE.h"
|
|
#include "NNEAttributeMap.h"
|
|
#include "NNEModelOptimizerInterface.h"
|
|
#include "NNERuntimeFormat.h"
|
|
#include "Serialization/CustomVersion.h"
|
|
#include "UObject/WeakInterfacePtr.h"
|
|
|
|
#if WITH_EDITOR
|
|
#include "Containers/StringFwd.h"
|
|
#include "DerivedDataCache.h"
|
|
#include "DerivedDataCacheKey.h"
|
|
#include "DerivedDataRequestOwner.h"
|
|
#include "Memory/CompositeBuffer.h"
|
|
#endif // WITH_EDITOR
|
|
|
|
namespace UE::NNE::ModelData
|
|
{
|
|
enum Version : uint32
|
|
{
|
|
V0 = 0, // Initial
|
|
V1 = 1, // TargetRuntimes and AssetImportData
|
|
V2 = 2, // Re-arrange fields and store only ModelData in cooked assets
|
|
// New versions can be added above this line
|
|
VersionPlusOne,
|
|
Latest = VersionPlusOne - 1
|
|
};
|
|
|
|
const FGuid GUID(0x9513202e, 0xeba1b279, 0xf17fe5ba, 0xab90c3f2);
|
|
FCustomVersionRegistration NNEModelDataVersion(GUID, Version::Latest, TEXT("NNEModelDataVersion"));// Always save with the latest version
|
|
|
|
const uint32 DDCAssetVersion = 1; // Increase this value to force rebuilding cache entries
|
|
|
|
FString GetRuntimesAsString(TArrayView<const FString> Runtimes)
|
|
{
|
|
if (Runtimes.Num() == 0)
|
|
{
|
|
return TEXT("All");
|
|
}
|
|
|
|
FString RuntimesAsOneString;
|
|
bool bIsFirstRuntime = true;
|
|
|
|
for (const FString& Runtime : Runtimes)
|
|
{
|
|
if (!bIsFirstRuntime)
|
|
{
|
|
RuntimesAsOneString += TEXT(", ");
|
|
}
|
|
RuntimesAsOneString += Runtime;
|
|
bIsFirstRuntime = false;
|
|
}
|
|
return RuntimesAsOneString;
|
|
}
|
|
|
|
#if WITH_EDITOR
|
|
|
|
inline FString GetDDCRequestId(const FString& FileId, const FString& RuntimeName, const FString& ModelDataIdentifier)
|
|
{
|
|
// RuntimeName and FileId are embedded to the id to ensure no potential collision between the runtime/assets
|
|
return RuntimeName + "-" + FileId + "-DDCv" + FString::FromInt(DDCAssetVersion) + "-" + ModelDataIdentifier;
|
|
}
|
|
|
|
inline UE::DerivedData::FCacheKey CreateCacheKey(const FString& RequestId)
|
|
{
|
|
return { UE::DerivedData::FCacheBucket(TEXT("NNEModelData")), FIoHash::HashBuffer(MakeMemoryView(FTCHARToUTF8(RequestId))) };
|
|
}
|
|
|
|
inline void PutIntoDDC(const FGuid& FileId, const FString& RuntimeName, const FString& ModelDataIdentifier, TSharedPtr<UE::NNE::FSharedModelData> Data)
|
|
{
|
|
check(Data.IsValid());
|
|
check(Data->GetView().Num() > 0);
|
|
|
|
FString FileIdStr = FileId.ToString(EGuidFormats::Digits);
|
|
FString RequestId = GetDDCRequestId(FileIdStr, RuntimeName, ModelDataIdentifier);
|
|
|
|
TArray<UE::DerivedData::FCachePutValueRequest> Requests;
|
|
Requests.SetNum(1);
|
|
|
|
Requests[0].Name = FString("Put-") + RequestId;
|
|
Requests[0].Key = CreateCacheKey(RequestId);
|
|
Requests[0].Value = UE::DerivedData::FValue::Compress(FCompositeBuffer(MakeSharedBufferFromArray(TArray<uint32>({ Data->GetMemoryAlignment() })), FSharedBuffer::MakeView(Data->GetView().GetData(), Data->GetView().Num())));
|
|
|
|
UE::DerivedData::FRequestOwner BlockingPutOwner(UE::DerivedData::EPriority::Blocking);
|
|
UE::DerivedData::GetCache().PutValue(Requests, BlockingPutOwner);
|
|
BlockingPutOwner.Wait();
|
|
}
|
|
|
|
inline TSharedPtr<UE::NNE::FSharedModelData> GetFromDDC(const FGuid& FileId, const FString& RuntimeName, const FString& ModelDataIdentifier)
|
|
{
|
|
FString FileIdStr = FileId.ToString(EGuidFormats::Digits);
|
|
FString RequestId = GetDDCRequestId(FileIdStr, RuntimeName, ModelDataIdentifier);
|
|
|
|
TArray<UE::DerivedData::FCacheGetValueRequest> Requests;
|
|
Requests.SetNum(1);
|
|
|
|
Requests[0].Name = FString("Get-") + RequestId;
|
|
Requests[0].Key = CreateCacheKey(RequestId);
|
|
|
|
TSharedPtr<UE::NNE::FSharedModelData> Result;
|
|
UE::DerivedData::FRequestOwner BlockingGetOwner(UE::DerivedData::EPriority::Blocking);
|
|
UE::DerivedData::GetCache().GetValue(Requests, BlockingGetOwner, [&Result](UE::DerivedData::FCacheGetValueResponse&& Response)
|
|
{
|
|
if (Response.Value.HasData() && Response.Value.GetRawSize() > sizeof(uint32))
|
|
{
|
|
FCompressedBufferReader Reader(Response.Value.GetData());
|
|
uint32 MemoryAlignment = ((uint32*)Reader.Decompress(0, sizeof(uint32)).GetData())[0];
|
|
uint64 DataSize = Response.Value.GetRawSize() - sizeof(uint32);
|
|
|
|
void* Data = FMemory::Malloc(DataSize, MemoryAlignment);
|
|
if (Reader.TryDecompressTo(FMutableMemoryView(Data, DataSize), sizeof(uint32)))
|
|
{
|
|
Result = MakeShared<UE::NNE::FSharedModelData>(FSharedBuffer::TakeOwnership(Data, DataSize, FMemory::Free), MemoryAlignment);
|
|
}
|
|
else
|
|
{
|
|
FMemory::Free(Data);
|
|
}
|
|
}
|
|
});
|
|
BlockingGetOwner.Wait();
|
|
return Result;
|
|
}
|
|
|
|
#endif // WITH_EDITOR
|
|
|
|
inline TSharedPtr<UE::NNE::FSharedModelData> CreateModelData(const FString& RuntimeName, FString FileType, const TArray<uint8>& FileData, FGuid FileId, const ITargetPlatform* TargetPlatform)
|
|
{
|
|
TWeakInterfacePtr<INNERuntime> NNERuntime = UE::NNE::GetRuntime<INNERuntime>(RuntimeName);
|
|
if (NNERuntime.IsValid())
|
|
{
|
|
if (NNERuntime->CanCreateModelData(FileType, FileData, FileId, TargetPlatform))
|
|
{
|
|
return NNERuntime->CreateModelData(FileType, FileData, FileId, TargetPlatform);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("UNNEModelData: No runtime '%s' found. Valid runtimes are: "), *RuntimeName);
|
|
TArrayView<TWeakInterfacePtr<INNERuntime>> Runtimes = UE::NNE::GetAllRuntimes();
|
|
for (int32 i = 0; i < Runtimes.Num(); i++)
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("- %s"), *Runtimes[i]->GetRuntimeName());
|
|
}
|
|
}
|
|
return TSharedPtr<UE::NNE::FSharedModelData>();
|
|
}
|
|
|
|
} // UE::NNE::ModelData
|
|
|
|
namespace UE::NNE
|
|
{
|
|
FSharedModelData::FSharedModelData(FSharedBuffer InData, uint32 InMemoryAlignment) : Data(InData), MemoryAlignment(InMemoryAlignment)
|
|
{
|
|
checkf(Data.IsOwned(), TEXT("InData data must be ownned!"));
|
|
checkf(MemoryAlignment <= 1 || (((uintptr_t)(const void*)(InData.GetData())) % MemoryAlignment == 0), TEXT("InData must be aligned with InMemoryAlignment!"))
|
|
}
|
|
|
|
FSharedModelData::FSharedModelData()
|
|
{
|
|
MemoryAlignment = 0;
|
|
}
|
|
|
|
TConstArrayView<uint8> FSharedModelData::GetView() const
|
|
{
|
|
return MakeArrayView(static_cast<const uint8*>(Data.GetData()), Data.GetSize());
|
|
}
|
|
|
|
uint32 FSharedModelData::GetMemoryAlignment() const
|
|
{
|
|
return MemoryAlignment;
|
|
}
|
|
} // UE::NNE
|
|
|
|
void UNNEModelData::GetAssetRegistryTags(TArray<FAssetRegistryTag>& OutTags) const
|
|
{
|
|
OutTags.Add(FAssetRegistryTag("TargetRuntimes", UE::NNE::ModelData::GetRuntimesAsString(GetTargetRuntimes()), FAssetRegistryTag::TT_Alphabetical));
|
|
Super::GetAssetRegistryTags(OutTags);
|
|
}
|
|
|
|
void UNNEModelData::Serialize(FArchive& Ar)
|
|
{
|
|
// Store the asset version (no effect in load)
|
|
Ar.UsingCustomVersion(UE::NNE::ModelData::GUID);
|
|
|
|
if (Ar.IsSaving() || Ar.IsCountingMemory()) {
|
|
bool bWriteModelData = true;
|
|
if (Ar.IsCooking())
|
|
{
|
|
// Optimize storage: FileData is not required anymore because we have the model and can cook it for every runtime
|
|
TArray<FString> TmpTargetRuntimes;
|
|
Ar << TmpTargetRuntimes;
|
|
FString TmpFileType;
|
|
Ar << TmpFileType;
|
|
TArray<uint8> TmpFileData;
|
|
Ar << TmpFileData;
|
|
Ar << FileId;
|
|
|
|
// Cooking must recreate all model data but only if file data is still available
|
|
if (FileData.Num() > 0)
|
|
{
|
|
ModelData.Reset();
|
|
|
|
// No target runtime means all currently registered ones
|
|
TArray<FString, TInlineAllocator<10>> CookRuntimeNames;
|
|
if (GetTargetRuntimes().IsEmpty())
|
|
{
|
|
for (const TWeakInterfacePtr<INNERuntime>& Runtime : UE::NNE::GetAllRuntimes())
|
|
{
|
|
CookRuntimeNames.Add(Runtime->GetRuntimeName());
|
|
}
|
|
}
|
|
else
|
|
{
|
|
CookRuntimeNames.Append(GetTargetRuntimes());
|
|
}
|
|
|
|
for (const FString& RuntimeName : CookRuntimeNames)
|
|
{
|
|
TSharedPtr<UE::NNE::FSharedModelData> CreatedData = UE::NNE::ModelData::CreateModelData(RuntimeName, FileType, FileData, FileId, Ar.GetArchiveState().CookingTarget());
|
|
if (CreatedData.IsValid() && CreatedData->GetView().Num() > 0)
|
|
{
|
|
ModelData.Add(RuntimeName, CreatedData);
|
|
#if WITH_EDITOR
|
|
TWeakInterfacePtr<INNERuntime> NNERuntime = UE::NNE::GetRuntime<INNERuntime>(RuntimeName);
|
|
if (NNERuntime.IsValid())
|
|
{
|
|
FString ModelDataIdentifier = NNERuntime->GetModelDataIdentifier(FileType, FileData, FileId, Ar.GetArchiveState().CookingTarget());
|
|
if (ModelDataIdentifier.Len() > 0)
|
|
{
|
|
UE::NNE::ModelData::PutIntoDDC(FileId, RuntimeName, ModelDataIdentifier, CreatedData);
|
|
}
|
|
else
|
|
{
|
|
UE_LOG(LogNNE, Warning, TEXT("UNNEModelData: Runtime '%s' returned an empty string as a ModelDataIdentifier while cooking. GetModelDataIdentifier should always return a valid identifier."), *RuntimeName);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
UE_LOG(LogNNE, Warning, TEXT("UNNEModelData: Runtime '%s' is among the cooked runtimes but instance is invalid."), *RuntimeName);
|
|
}
|
|
#endif //WITH_EDITOR
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
// Only cooked assets optimize storage
|
|
Ar << TargetRuntimes;
|
|
Ar << FileType;
|
|
Ar << FileData;
|
|
Ar << FileId;
|
|
|
|
#if WITH_EDITOR
|
|
// In editor (when not cooking), no model data is stored as model data can always be recreated and unnecessary data in subversion control should be avoided
|
|
bWriteModelData = false;
|
|
#endif //WITH_EDITOR
|
|
}
|
|
|
|
if (bWriteModelData)
|
|
{
|
|
TArray<FString> RuntimeNames;
|
|
ModelData.GetKeys(RuntimeNames);
|
|
int32 NumItems = RuntimeNames.Num();
|
|
|
|
Ar << NumItems;
|
|
for (int32 i = 0; i < NumItems; i++)
|
|
{
|
|
Ar << RuntimeNames[i];
|
|
|
|
uint32 MemoryAlignment = ModelData[RuntimeNames[i]]->GetMemoryAlignment();
|
|
Ar << MemoryAlignment;
|
|
|
|
uint64 DataSize = ModelData[RuntimeNames[i]]->GetView().Num();
|
|
Ar << DataSize;
|
|
|
|
Ar.Serialize((void*)ModelData[RuntimeNames[i]]->GetView().GetData(), DataSize);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
int32 NumItems = 0;
|
|
Ar << NumItems;
|
|
}
|
|
}
|
|
else if (Ar.IsLoading())
|
|
{
|
|
TObjectPtr<class UAssetImportData> AssetImportData;
|
|
int32 NumItems = 0;
|
|
FString Name;
|
|
uint32 MemoryAlignment = 0;
|
|
uint64 DataSize = 0;
|
|
TArray<uint8> Data;
|
|
void* RawData = nullptr;
|
|
int32 Index = 0;
|
|
|
|
switch (Ar.CustomVer(UE::NNE::ModelData::GUID))
|
|
{
|
|
case UE::NNE::ModelData::Version::V0:
|
|
TargetRuntimes.Empty();
|
|
Ar << FileType;
|
|
Ar << FileData;
|
|
Ar << FileId;
|
|
Ar << NumItems;
|
|
for (Index = 0; Index < NumItems; Index++)
|
|
{
|
|
Ar << Name;
|
|
Ar << Data;
|
|
ModelData.Add(Name, MakeShared<UE::NNE::FSharedModelData>(MakeSharedBufferFromArray(MoveTemp(Data)), 0));
|
|
}
|
|
UE_LOG(LogNNE, Warning, TEXT("[DEPRECATION] UNNEModelData: The asset %s (v0) is deprecated. Please right-click the asset and select 'Save' to update it to the latest version."), *this->GetName());
|
|
break;
|
|
|
|
case UE::NNE::ModelData::Version::V1:
|
|
TargetRuntimes.Empty();
|
|
if (!Ar.IsLoadingFromCookedPackage())
|
|
{
|
|
Ar << TargetRuntimes;
|
|
Ar << AssetImportData;
|
|
}
|
|
Ar << FileType;
|
|
Ar << FileData;
|
|
Ar << FileId;
|
|
Ar << NumItems;
|
|
for (Index = 0; Index < NumItems; Index++)
|
|
{
|
|
Ar << Name;
|
|
Ar << Data;
|
|
ModelData.Add(Name, MakeShared<UE::NNE::FSharedModelData>(MakeSharedBufferFromArray(MoveTemp(Data)), 0));
|
|
}
|
|
UE_LOG(LogNNE, Warning, TEXT("[DEPRECATION] UNNEModelData: The asset %s (v1) is deprecated. Please right-click the asset and select 'Save' to update it to the latest version."), *this->GetName());
|
|
break;
|
|
|
|
case UE::NNE::ModelData::Version::V2:
|
|
Ar << TargetRuntimes;
|
|
Ar << FileType;
|
|
Ar << FileData;
|
|
Ar << FileId;
|
|
Ar << NumItems;
|
|
for (Index = 0; Index < NumItems; Index++)
|
|
{
|
|
Ar << Name;
|
|
Ar << MemoryAlignment;
|
|
Ar << DataSize;
|
|
RawData = FMemory::Malloc(DataSize, MemoryAlignment);
|
|
Ar.Serialize(RawData, DataSize);
|
|
ModelData.Add(Name, MakeShared<UE::NNE::FSharedModelData>(FSharedBuffer::TakeOwnership(RawData, DataSize, FMemory::Free), MemoryAlignment));
|
|
}
|
|
break;
|
|
|
|
default:
|
|
UE_LOG(LogNNE, Error, TEXT("UNNEModelData: Unknown asset version %d: Deserialisation failed, please reimport the original model."), Ar.CustomVer(UE::NNE::ModelData::GUID));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
void UNNEModelData::Init(const FString& Type, TConstArrayView<uint8> Buffer)
|
|
{
|
|
TargetRuntimes.Empty();
|
|
FileType = Type;
|
|
FileData = Buffer;
|
|
FPlatformMisc::CreateGuid(FileId);
|
|
ModelData.Empty();
|
|
}
|
|
|
|
TArrayView<const FString> UNNEModelData::GetTargetRuntimes() const
|
|
{
|
|
return TargetRuntimes;
|
|
}
|
|
|
|
void UNNEModelData::SetTargetRuntimes(TArrayView<const FString> RuntimeNames)
|
|
{
|
|
TargetRuntimes = RuntimeNames;
|
|
|
|
if (RuntimeNames.Num() > 0)
|
|
{
|
|
TArray<FString, TInlineAllocator<10>> CookedRuntimes;
|
|
ModelData.GetKeys(CookedRuntimes);
|
|
for (const FString& Runtime : CookedRuntimes)
|
|
{
|
|
if (!TargetRuntimes.Contains(Runtime))
|
|
{
|
|
ModelData.Remove(Runtime);
|
|
}
|
|
}
|
|
ModelData.Compact();
|
|
}
|
|
}
|
|
|
|
FString UNNEModelData::GetFileType() const
|
|
{
|
|
return FileType;
|
|
}
|
|
|
|
TConstArrayView<uint8> UNNEModelData::GetFileData() const
|
|
{
|
|
return FileData;
|
|
}
|
|
|
|
void UNNEModelData::ClearFileDataAndFileType()
|
|
{
|
|
FileType = "";
|
|
FileData.Empty();
|
|
}
|
|
|
|
FGuid UNNEModelData::GetFileId() const
|
|
{
|
|
return FileId;
|
|
}
|
|
|
|
TSharedPtr<UE::NNE::FSharedModelData> UNNEModelData::GetModelData(const FString& RuntimeName)
|
|
{
|
|
// Check model data is supporting the requested target runtime
|
|
TArrayView<const FString> TargetRuntimesNames = GetTargetRuntimes();
|
|
if (!TargetRuntimesNames.IsEmpty() && !TargetRuntimesNames.Contains(RuntimeName))
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("UNNEModelData: Runtime '%s' is not among the target runtimes. Target runtimes are: "), *RuntimeName);
|
|
for (const FString& TargetRuntimesName : TargetRuntimesNames)
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("- %s"), *TargetRuntimesName);
|
|
}
|
|
return TSharedPtr<UE::NNE::FSharedModelData>();
|
|
}
|
|
|
|
// Check if we have a local cache hit
|
|
TSharedPtr<UE::NNE::FSharedModelData>* LocalDataPtr = ModelData.Find(RuntimeName);
|
|
if (LocalDataPtr)
|
|
{
|
|
return *LocalDataPtr;
|
|
}
|
|
|
|
// After this point FileData is required to either get the cache id or recreate it from scratch
|
|
if (FileData.Num() < 1)
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("UNNEModelData: Cannot create model data from empty file data."));
|
|
return TSharedPtr<UE::NNE::FSharedModelData>();
|
|
}
|
|
|
|
#if WITH_EDITOR
|
|
TWeakInterfacePtr<INNERuntime> NNERuntime = UE::NNE::GetRuntime<INNERuntime>(RuntimeName);
|
|
if (!NNERuntime.IsValid())
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("UNNEModelData: Runtime '%s' is among the target runtimes but instance is invalid."), *RuntimeName);
|
|
return TSharedPtr<UE::NNE::FSharedModelData>();
|
|
}
|
|
|
|
FString ModelDataIdentifier = NNERuntime->GetModelDataIdentifier(FileType, FileData, FileId, nullptr);
|
|
if (ModelDataIdentifier.Len() == 0)
|
|
{
|
|
UE_LOG(LogNNE, Error, TEXT("UNNEModelData: Runtime '%s' returned an empty string as a ModelDataIdentifier. GetModelDataIdentifier should always return a valid identifier."), *RuntimeName);
|
|
return TSharedPtr<UE::NNE::FSharedModelData>();
|
|
}
|
|
|
|
// Check if we have a DDC cache hit
|
|
TSharedPtr<UE::NNE::FSharedModelData> RemoteData = UE::NNE::ModelData::GetFromDDC(FileId, RuntimeName, ModelDataIdentifier);
|
|
if (RemoteData.IsValid() && RemoteData->GetView().Num() > 0)
|
|
{
|
|
ModelData.Add(RuntimeName, RemoteData);
|
|
return RemoteData;
|
|
}
|
|
#endif //WITH_EDITOR
|
|
|
|
// Try to create the model
|
|
TSharedPtr<UE::NNE::FSharedModelData> CreatedData = UE::NNE::ModelData::CreateModelData(RuntimeName, FileType, FileData, FileId, nullptr);
|
|
if (!CreatedData.IsValid() || CreatedData->GetView().Num() < 1)
|
|
{
|
|
return TSharedPtr<UE::NNE::FSharedModelData>();
|
|
}
|
|
|
|
// Cache the model
|
|
ModelData.Add(RuntimeName, CreatedData);
|
|
|
|
#if WITH_EDITOR
|
|
// And put it into DDC
|
|
UE::NNE::ModelData::PutIntoDDC(FileId, RuntimeName, ModelDataIdentifier, CreatedData);
|
|
#endif //WITH_EDITOR
|
|
|
|
return CreatedData;
|
|
}
|
|
|
|
void UNNEModelData::ClearModelData()
|
|
{
|
|
ModelData.Empty();
|
|
}
|