added support for dynamic LLM tags (#245)

* added FLLMTag which can be used to register memory tags outside of core
* changed `UE_MEMSCOPE` -> `ZEN_MEMSCOPE` for consistency
* instrumented some subsystems with dynamic tags

9 files changed, 132 insertions, 31 deletions

diff --git a/src/zencore/include/zencore/memory/llm.h b/src/zencore/include/zencore/memory/llm.h
index 4f1c9de77..ea7f68cc6 100644
--- a/src/zencore/include/zencore/memory/llm.h
+++ b/src/zencore/include/zencore/memory/llm.h
@@ -28,4 +28,20 @@ enum class ELLMTag : uint8_t
 	GenericTagCount
 };
 
+struct FLLMTag
+{
+public:
+	FLLMTag(const char* TagName);
+	FLLMTag(const char* TagName, const FLLMTag& ParentTag);
+
+	inline int32_t GetTag() const { return m_Tag; }
+	inline int32_t GetParentTag() const { return m_ParentTag; }
+
+private:
+	int32_t m_Tag		= -1;
+	int32_t m_ParentTag = -1;
+
+	void AssignAndAnnounceNewTag(const char* TagName);
+};
+
 }  // namespace zen
diff --git a/src/zencore/include/zencore/memory/tagtrace.h b/src/zencore/include/zencore/memory/tagtrace.h
index f51b21466..8b5fc0e67 100644
--- a/src/zencore/include/zencore/memory/tagtrace.h
+++ b/src/zencore/include/zencore/memory/tagtrace.h
@@ -9,6 +9,7 @@
 namespace zen {
 
 enum class ELLMTag : uint8_t;
+struct FLLMTag;
 
 int32_t MemoryTrace_AnnounceCustomTag(int32_t Tag, int32_t ParentTag, const char* Display);
 int32_t MemoryTrace_GetActiveTag();
@@ -42,6 +43,7 @@ class FMemScope
 public:
 	FMemScope();  // Used with SetTagAndActivate
 	FMemScope(int32_t InTag, bool bShouldActivate = true);
+	FMemScope(FLLMTag InTag, bool bShouldActivate = true);
 	FMemScope(ELLMTag InTag, bool bShouldActivate = true);
 	~FMemScope();
 
@@ -75,19 +77,18 @@ private:
 };
 
 ////////////////////////////////////////////////////////////////////////////////
-#	define UE_MEMSCOPE(InTag)				FMemScope PREPROCESSOR_JOIN(MemScope, __LINE__)(InTag);
-#	define UE_MEMSCOPE_PTR(InPtr)			FMemScopePtr PREPROCESSOR_JOIN(MemPtrScope, __LINE__)((uint64)InPtr);
-#	define UE_MEMSCOPE_DEFAULT(InTag)		FDefaultMemScope PREPROCESSOR_JOIN(MemScope, __LINE__)(InTag);
-#	define UE_MEMSCOPE_UNINITIALIZED(Line) FMemScope PREPROCESSOR_JOIN(MemScope, Line);
+#	define ZEN_MEMSCOPE(InTag)				 FMemScope PREPROCESSOR_JOIN(MemScope, __LINE__)(InTag);
+#	define ZEN_MEMSCOPE_PTR(InPtr)			 FMemScopePtr PREPROCESSOR_JOIN(MemPtrScope, __LINE__)((uint64)InPtr);
+#	define ZEN_MEMSCOPE_DEFAULT(InTag)		 FDefaultMemScope PREPROCESSOR_JOIN(MemScope, __LINE__)(InTag);
+#	define ZEN_MEMSCOPE_UNINITIALIZED(Line) FMemScope PREPROCESSOR_JOIN(MemScope, Line);
 
 #else  // UE_MEMORY_TAGS_TRACE_ENABLED
 
 ////////////////////////////////////////////////////////////////////////////////
-#	define UE_MEMSCOPE(...)
-#	define UE_MEMSCOPE_PTR(...)
-#	define UE_MEMSCOPE_DEFAULT(...)
-#	define UE_MEMSCOPE_UNINITIALIZED(...)
-#	define UE_MEMSCOPE_ACTIVATE(...)
+#	define ZEN_MEMSCOPE(...)
+#	define ZEN_MEMSCOPE_PTR(...)
+#	define ZEN_MEMSCOPE_DEFAULT(...)
+#	define ZEN_MEMSCOPE_UNINITIALIZED(...)
 
 #endif	// UE_MEMORY_TAGS_TRACE_ENABLED
 }
diff --git a/src/zencore/iobuffer.cpp b/src/zencore/iobuffer.cpp
index 02f85a9d9..3d9d6706a 100644
--- a/src/zencore/iobuffer.cpp
+++ b/src/zencore/iobuffer.cpp
@@ -39,7 +39,7 @@ namespace zen {
 void
 IoBufferCore::AllocateBuffer(size_t InSize, size_t Alignment) const
 {
-	UE_MEMSCOPE(ELLMTag::IoBufferMemory);
+	ZEN_MEMSCOPE(ELLMTag::IoBufferMemory);
 
 	void* Ptr = Memory::Alloc(InSize, Alignment);
 
@@ -64,7 +64,7 @@ IoBufferCore::FreeBuffer()
 void*
 IoBufferCore::operator new(size_t Size)
 {
-	UE_MEMSCOPE(ELLMTag::IoBufferCore);
+	ZEN_MEMSCOPE(ELLMTag::IoBufferCore);
 	return Memory::Malloc(Size);
 }
 
diff --git a/src/zencore/logging.cpp b/src/zencore/logging.cpp
index 7bd500b3b..5885587ad 100644
--- a/src/zencore/logging.cpp
+++ b/src/zencore/logging.cpp
@@ -68,7 +68,7 @@ static_assert(offsetof(spdlog::source_loc, funcname) == offsetof(SourceLocation,
 void
 EmitLogMessage(LoggerRef& Logger, int LogLevel, const std::string_view Message)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 	const spdlog::level::level_enum InLevel = (spdlog::level::level_enum)LogLevel;
 	Logger.SpdLogger->log(InLevel, Message);
 	if (IsErrorLevel(LogLevel))
@@ -83,7 +83,7 @@ EmitLogMessage(LoggerRef& Logger, int LogLevel, const std::string_view Message)
 void
 EmitLogMessage(LoggerRef& Logger, int LogLevel, std::string_view Format, fmt::format_args Args)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 	zen::logging::LoggingContext LogCtx;
 	fmt::vformat_to(fmt::appender(LogCtx.MessageBuffer), Format, Args);
 	zen::logging::EmitLogMessage(Logger, LogLevel, LogCtx.Message());
@@ -92,7 +92,7 @@ EmitLogMessage(LoggerRef& Logger, int LogLevel, std::string_view Format, fmt::fo
 void
 EmitLogMessage(LoggerRef& Logger, const SourceLocation& InLocation, int LogLevel, const std::string_view Message)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 	const spdlog::source_loc&		Location = *reinterpret_cast<const spdlog::source_loc*>(&InLocation);
 	const spdlog::level::level_enum InLevel	 = (spdlog::level::level_enum)LogLevel;
 	Logger.SpdLogger->log(Location, InLevel, Message);
@@ -108,7 +108,7 @@ EmitLogMessage(LoggerRef& Logger, const SourceLocation& InLocation, int LogLevel
 void
 EmitLogMessage(LoggerRef& Logger, const SourceLocation& InLocation, int LogLevel, std::string_view Format, fmt::format_args Args)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 	zen::logging::LoggingContext LogCtx;
 	fmt::vformat_to(fmt::appender(LogCtx.MessageBuffer), Format, Args);
 	zen::logging::EmitLogMessage(Logger, InLocation, LogLevel, LogCtx.Message());
@@ -117,7 +117,7 @@ EmitLogMessage(LoggerRef& Logger, const SourceLocation& InLocation, int LogLevel
 void
 EmitConsoleLogMessage(int LogLevel, const std::string_view Message)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 	const spdlog::level::level_enum InLevel = (spdlog::level::level_enum)LogLevel;
 	ConsoleLog().SpdLogger->log(InLevel, Message);
 }
@@ -125,7 +125,7 @@ EmitConsoleLogMessage(int LogLevel, const std::string_view Message)
 void
 EmitConsoleLogMessage(int LogLevel, std::string_view Format, fmt::format_args Args)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 	zen::logging::LoggingContext LogCtx;
 	fmt::vformat_to(fmt::appender(LogCtx.MessageBuffer), Format, Args);
 	zen::logging::EmitConsoleLogMessage(LogLevel, LogCtx.Message());
@@ -200,7 +200,7 @@ std::string LogLevels[level::LogLevelCount];
 void
 ConfigureLogLevels(level::LogLevel Level, std::string_view Loggers)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	RwLock::ExclusiveLockScope _(LogLevelsLock);
 	LogLevels[Level] = Loggers;
@@ -209,7 +209,7 @@ ConfigureLogLevels(level::LogLevel Level, std::string_view Loggers)
 void
 RefreshLogLevels(level::LogLevel* DefaultLevel)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	spdlog::details::registry::log_levels Levels;
 
@@ -287,7 +287,7 @@ Default()
 void
 SetDefault(std::string_view NewDefaultLoggerId)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	auto NewDefaultLogger = spdlog::get(std::string(NewDefaultLoggerId));
 	ZEN_ASSERT(NewDefaultLogger);
@@ -307,7 +307,7 @@ ErrorLog()
 void
 SetErrorLog(std::string_view NewErrorLoggerId)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	if (NewErrorLoggerId.empty())
 	{
@@ -326,7 +326,7 @@ SetErrorLog(std::string_view NewErrorLoggerId)
 LoggerRef
 Get(std::string_view Name)
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	std::shared_ptr<spdlog::logger> Logger = spdlog::get(std::string(Name));
 
@@ -357,7 +357,7 @@ SuppressConsoleLog()
 LoggerRef
 ConsoleLog()
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	std::call_once(ConsoleInitFlag, [&] {
 		if (!ConLogger)
@@ -375,7 +375,7 @@ ConsoleLog()
 void
 InitializeLogging()
 {
-	UE_MEMSCOPE(ELLMTag::Logging);
+	ZEN_MEMSCOPE(ELLMTag::Logging);
 
 	TheDefaultLogger = *spdlog::default_logger_raw();
 }
diff --git a/src/zencore/memory/llm.cpp b/src/zencore/memory/llm.cpp
new file mode 100644
index 000000000..fe4853d49
--- /dev/null
+++ b/src/zencore/memory/llm.cpp
@@ -0,0 +1,77 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include <zencore/memory/llm.h>
+
+#include <zencore/string.h>
+#include <zencore/thread.h>
+
+#include <atomic>
+
+namespace zen {
+
+static std::atomic<int32_t> CustomTagCounter = 257;	 // NOTE: hard-coded TRACE_TAG = 257
+
+static const int32_t TagNamesBaseIndex	 = 256;
+static const int32_t TrackedTagNameCount = 256;
+static const char*	 TagNames[TrackedTagNameCount];
+static uint32_t		 TagNameHashes[TrackedTagNameCount];
+
+static RwLock TableLock;
+
+FLLMTag::FLLMTag(const char* TagName)
+{
+	// NOTE: should add verification to prevent multiple definitions of same name?
+
+	AssignAndAnnounceNewTag(TagName);
+}
+
+FLLMTag::FLLMTag(const char* TagName, const FLLMTag& ParentTag)
+{
+	// NOTE: should add verification to prevent multiple definitions of same name?
+
+	m_ParentTag = ParentTag.GetTag();
+
+	AssignAndAnnounceNewTag(TagName);
+}
+
+void
+FLLMTag::AssignAndAnnounceNewTag(const char* TagName)
+{
+	const uint32_t TagNameHash = HashStringDjb2(TagName);
+
+	{
+		RwLock::ExclusiveLockScope _(TableLock);
+
+		const int32_t CurrentMaxTagIndex = CustomTagCounter - TagNamesBaseIndex;
+
+		for (int TagIndex = 0; TagIndex <= CurrentMaxTagIndex; ++TagIndex)
+		{
+			if (TagNameHashes[TagIndex] == TagNameHash)
+			{
+				m_Tag = TagIndex + TagNamesBaseIndex;
+				// could verify the string matches here to catch hash collisions
+
+				// return early, no need to announce the tag as it is already known
+				return;
+			}
+		}
+
+		m_Tag = ++CustomTagCounter;
+
+		const int TagIndex = m_Tag - TagNamesBaseIndex;
+
+		if (TagIndex < TrackedTagNameCount)
+		{
+			TagNameHashes[TagIndex] = TagNameHash;
+			TagNames[TagIndex]		= TagName;
+		}
+		else
+		{
+			// should really let user know there's an overflow
+		}
+	}
+
+	MemoryTrace_AnnounceCustomTag(m_Tag, m_ParentTag, TagName);
+}
+
+}  // namespace zen
diff --git a/src/zencore/memtrack/memorytrace.cpp b/src/zencore/memtrack/memorytrace.cpp
index b147aee91..7089c356a 100644
--- a/src/zencore/memtrack/memorytrace.cpp
+++ b/src/zencore/memtrack/memorytrace.cpp
@@ -755,7 +755,7 @@ FTraceMalloc::Malloc(SIZE_T Count, uint32_t Alignment)
 {
 #if UE_MEMORY_TRACE_ENABLED
 	// UE_TRACE_METADATA_CLEAR_SCOPE();
-	UE_MEMSCOPE(TRACE_TAG);
+	ZEN_MEMSCOPE(TRACE_TAG);
 
 	void* NewPtr;
 	{
@@ -783,7 +783,7 @@ FTraceMalloc::Realloc(void* Original, SIZE_T Count, uint32_t Alignment)
 {
 #if UE_MEMORY_TRACE_ENABLED
 	// UE_TRACE_METADATA_CLEAR_SCOPE();
-	UE_MEMSCOPE(TRACE_TAG);
+	ZEN_MEMSCOPE(TRACE_TAG);
 
 	UE_TRACE_LOG(Memory, ReallocFree, MemAllocChannel)
 		<< ReallocFree.Address(uint64_t(Original)) << ReallocFree.RootHeap(uint8(EMemoryTraceRootHeap::SystemMemory));
diff --git a/src/zencore/memtrack/moduletrace.cpp b/src/zencore/memtrack/moduletrace.cpp
index 51280ff3a..cf37c5932 100644
--- a/src/zencore/memtrack/moduletrace.cpp
+++ b/src/zencore/memtrack/moduletrace.cpp
@@ -228,7 +228,7 @@ FModuleTrace::OnDllLoaded(const UNICODE_STRING& Name, UPTRINT Base)
 
 #	if UE_MEMORY_TRACE_ENABLED
 	{
-		UE_MEMSCOPE(ELLMTag::ProgramSize);
+		ZEN_MEMSCOPE(ELLMTag::ProgramSize);
 		MemoryTrace_Alloc(Base, OptionalHeader.SizeOfImage, 4 * 1024, EMemoryTraceRootHeap::SystemMemory);
 		MemoryTrace_MarkAllocAsHeap(Base, ProgramHeapId);
 		MemoryTrace_Alloc(Base, OptionalHeader.SizeOfImage, 4 * 1024, EMemoryTraceRootHeap::SystemMemory);
diff --git a/src/zencore/memtrack/tagtrace.cpp b/src/zencore/memtrack/tagtrace.cpp
index 15ba78ae4..797da0fab 100644
--- a/src/zencore/memtrack/tagtrace.cpp
+++ b/src/zencore/memtrack/tagtrace.cpp
@@ -51,7 +51,15 @@ FMemScope::FMemScope(ELLMTag InTag, bool bShouldActivate /*= true*/)
 {
 	if (UE_TRACE_CHANNELEXPR_IS_ENABLED(MemAllocChannel) & bShouldActivate)
 	{
-		ActivateScope(static_cast<int32>(InTag));
+		ActivateScope(static_cast<int32_t>(InTag));
+	}
+}
+
+FMemScope::FMemScope(FLLMTag InTag, bool bShouldActivate /*= true*/)
+{
+	if (UE_TRACE_CHANNELEXPR_IS_ENABLED(MemAllocChannel) & bShouldActivate)
+	{
+		ActivateScope(static_cast<int32_t>(InTag.GetTag()));
 	}
 }
 
@@ -212,7 +220,6 @@ int32_t
 MemoryTrace_AnnounceCustomTag(int32_t Tag, int32_t ParentTag, const char* Display)
 {
 #if UE_MEMORY_TAGS_TRACE_ENABLED && UE_TRACE_ENABLED
-	// todo: How do we check if tag trace is active?
 	if (GTagTrace)
 	{
 		return GTagTrace->AnnounceCustomTag(Tag, ParentTag, Display);
diff --git a/src/zencore/stats.cpp b/src/zencore/stats.cpp
index 6be16688b..8a424c5ad 100644
--- a/src/zencore/stats.cpp
+++ b/src/zencore/stats.cpp
@@ -226,7 +226,7 @@ thread_local xoshiro256 ThreadLocalRng;
 
 UniformSample::UniformSample(uint32_t ReservoirSize)
 {
-	UE_MEMSCOPE(ELLMTag::Metrics);
+	ZEN_MEMSCOPE(ELLMTag::Metrics);
 	m_Values = std::vector<std::atomic<int64_t>>(ReservoirSize);
 }
 
@@ -277,7 +277,7 @@ UniformSample::Update(int64_t Value)
 SampleSnapshot
 UniformSample::Snapshot() const
 {
-	UE_MEMSCOPE(ELLMTag::Metrics);
+	ZEN_MEMSCOPE(ELLMTag::Metrics);
 
 	uint64_t			ValuesSize = Size();
 	std::vector<double> Values(ValuesSize);