Files
UnrealEngine/Engine/Plugins/Experimental/LearningAgents/Source/LearningTraining/Public/LearningSharedMemory.h
2025-05-18 13:04:45 +08:00

100 lines
3.2 KiB
C++

// Copyright Epic Games, Inc.All Rights Reserved.
#pragma once
#include "LearningArray.h"
#include "GenericPlatform/GenericPlatformMemory.h"
#include "Misc/Guid.h"
/**
* Shared memory is used by Learning to efficiently share both experience
* and policy data between a Python training sub-process and potentially
* multiple unreal processes (which may be gathering that experience).
*/
namespace UE::Learning
{
/**
* Basic struct containing a view to a shared memory region, its guid,
* and a pointer to that region which is used for deallocation.
*/
template<uint8 DimNum, typename ElementType>
struct TSharedMemoryArrayView
{
FGuid Guid;
TLearningArrayView<DimNum, ElementType> View;
FPlatformMemory::FSharedMemoryRegion* Region = nullptr;
};
namespace SharedMemory
{
/**
* Map a view of a region of shared memory.
*
* @param Guid Guid of the shared memory region
* @param Shape Shape of the shared memory region
* @param bCreate If this shared memory region should be created or if it already exists
* @returns View of the shared memory region, guid, and pointer required for unmapping
*/
template<uint8 DimNum, typename ElementType>
TSharedMemoryArrayView<DimNum, ElementType> Map(const FGuid Guid, const TLearningArrayShape<DimNum>& Shape, const bool bCreate = false)
{
const SIZE_T TotalSize = sizeof(ElementType) * (SIZE_T)Shape.Total();
if (TotalSize == 0)
{
return { FGuid(), TLearningArrayView<DimNum, ElementType>(nullptr, Shape), nullptr };
}
FPlatformMemory::FSharedMemoryRegion* RegionPointer = FPlatformMemory::MapNamedSharedMemoryRegion(
Guid.ToString(PLATFORM_MAC ? EGuidFormats::Short : EGuidFormats::DigitsWithHyphensInBraces),
bCreate,
// I was getting odd issues allocating read-only memory so for now we always allocate as read-write
FPlatformMemory::ESharedMemoryAccess::Read | FPlatformMemory::ESharedMemoryAccess::Write,
TotalSize);
if (ensureMsgf(RegionPointer != nullptr, TEXT("Unable to map shared memory.")))
{
return { Guid, TLearningArrayView<DimNum, ElementType>((ElementType*)RegionPointer->GetAddress(), Shape), RegionPointer };
}
return TSharedMemoryArrayView<DimNum, ElementType>();
}
/*
* Unmaps the view of the region of shared memory
*/
template<uint8 DimNum, typename ElementType>
void Unmap(TSharedMemoryArrayView<DimNum, ElementType>& Memory)
{
if (Memory.Region != nullptr)
{
ensureMsgf(FPlatformMemory::UnmapNamedSharedMemoryRegion(Memory.Region), TEXT("Failed to unmap shared memory."));
}
Memory = TSharedMemoryArrayView<DimNum, ElementType>();
}
/**
* Allocate a region of shared memory.
*
* @param Shape Shape of the shared memory region
* @returns View of the shared memory region, guid, and pointer required for unmapping
*/
template<uint8 DimNum, typename ElementType>
TSharedMemoryArrayView<DimNum, ElementType> Allocate(const TLearningArrayShape<DimNum>& Shape)
{
return Map<DimNum, ElementType>(FGuid::NewGuid(), Shape, true);
}
/*
* Deallocate the region of shared memory
*/
template<uint8 DimNum, typename ElementType>
void Deallocate(TSharedMemoryArrayView<DimNum, ElementType>& Memory)
{
Unmap<DimNum, ElementType>(Memory);
}
}
}