[llvm] [offload] Respect alignment in memory manager (PR #214753)

[Robert Imschweiler via llvm-commits]

https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/214753

>From fe78ec1ace50271d25a897e7dcd9a9721191c989 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Thu, 6 Aug 2026 11:51:52 -0500
Subject: [PATCH] [offload] Respect alignment in memory manager

Also fix alignment comparison in AMDGPU rtl.cpp.
---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp             | 2 +-
 offload/plugins-nextgen/common/include/MemoryManager.h | 9 +++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 29b5006f0d7b6..0f0de67b20c7c 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -339,7 +339,7 @@ struct AMDGPUMemoryPoolTy {
     // compared with the alignment of the memory allocated using the given pool.
     // If the default alignment is greater than or equal to the alignment
     // requested by the user, it would still meet the user's requirements.
-    if (Alignment > 0 && Alignment >= PoolAllocationAlignment) {
+    if (Alignment > 0 && Alignment > PoolAllocationAlignment) {
       return Plugin::error(ErrorCode::UNSUPPORTED,
                            "requested alignment (%lu) larger than maximum "
                            "supported pool alignment (%lu)",
diff --git a/offload/plugins-nextgen/common/include/MemoryManager.h b/offload/plugins-nextgen/common/include/MemoryManager.h
index 4b57be45e7551..883ac40c269ea 100644
--- a/offload/plugins-nextgen/common/include/MemoryManager.h
+++ b/offload/plugins-nextgen/common/include/MemoryManager.h
@@ -13,6 +13,7 @@
 #ifndef LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H
 #define LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H
 
+#include <algorithm>
 #include <cassert>
 #include <functional>
 #include <list>
@@ -25,6 +26,7 @@
 #include "Shared/Utils.h"
 #include "omptarget.h"
 
+#include "llvm/Support/Alignment.h"
 #include "llvm/Support/Error.h"
 
 using namespace llvm::offload::debug;
@@ -268,9 +270,12 @@ class MemoryManagerTy {
 
       NodeTy TempNode(Size, nullptr);
       std::lock_guard<std::mutex> LG(FreeListLocks[B]);
-      const auto Itr = List.find(TempNode);
+      auto [First, Last] = List.equal_range(TempNode);
 
-      if (Itr != List.end()) {
+      auto Itr = std::find_if(First, Last, [Alignment](const NodeTy &N) {
+        return Alignment == 0 || isAddrAligned(Align(Alignment), N.Ptr);
+      });
+      if (Itr != Last) {
         NodePtr = &Itr->get();
         List.erase(Itr);
       }