VirtualAllocPageStatus.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
#pragma once
 
#include "HAL/LowLevelMemTracker.h"
 
// FVirtualAllocPageTracker is only referenced by LLM-specific code
#define UE_VIRTUALALLOC_PAGE_STATUS_ENABLED ENABLE_LOW_LEVEL_MEM_TRACKER
 
#if UE_VIRTUALALLOC_PAGE_STATUS_ENABLED
 
#include "HAL/Platform.h"
#include "HAL/PlatformMutex.h"
#include "Math/NumericLimits.h"
 
struct FHashMapLinearProbingVAlloc
{
public:
    void Init(uint64 PageSize);
    ~FHashMapLinearProbingVAlloc();
 
    uint64* Find(uint64 Key);
    uint64& FindOrAdd(uint64 Key, uint64 ValueIfMissing);
 
    uint64 RemoveAndGetValue(uint64 Key, uint64 ValueIfMissing);
 
private:
    enum class EValueType
    {
        Unallocated,
        Tombstone,
        Active, // Any Key value <= MaxKey is an Active address.
    };
 
    struct FBucketEntry
    {
        uint64 Key;
        uint64 Value;
 
        EValueType GetValueType();
        void SetAsActive(uint64 InKey);
        void SetAsUnallocated();
        void SetAsTombstone();
    };
 
private:
    void* Malloc(size_t Size);
    void Free(void* Ptr, size_t Size);
    void Realloc(int64 InNum);
    FBucketEntry* FindBucket(uint64 Key);
 
private:
    FBucketEntry* Buckets = nullptr;
    int64 NumBuckets = 0;
    int64 NumActive = 0;
    static constexpr float DesiredPopulation = 0.3f;
    static constexpr uint64 MaxKey = MAX_uint64 - static_cast<uint64>(EValueType::Active);
    static constexpr uint64 CollisionResolutionDeltaFraction = 32;
};
 
struct FVirtualAllocPageStatus
{
    FVirtualAllocPageStatus();
 
    int64 MarkChangedAndReturnDeltaSize(void* InStartAddress, SIZE_T InSize, bool bCommitted);
    void AddReservationSize(void* InReservationAddress, SIZE_T InSize, SIZE_T& OutOldSize);
    SIZE_T GetAndRemoveReservationSize(void* InReservationAddress);
 
private:
    UE::FPlatformRecursiveMutex Mutex;
    FHashMapLinearProbingVAlloc GroupToPageBitsMap;
    FHashMapLinearProbingVAlloc PageToReservationSizeMap;
    uint64 PageSize = 1;
    uint64 ReservationAlignment = 1;
    uint64 AccumulatedSize = 0;
 
    static constexpr uint64 PagesPerGroup = sizeof(uint64) * 8;
};
 
#endif // UE_VIRTUALALLOC_PAGE_STATUS_ENABLED