[llvm] [Offload] Check HSA allocate status before using the pointer (PR #223337)

via llvm-commits llvm-commits at lists.llvm.org
Tue Sep 22 19:55:45 PDT 2026


https://github.com/StevenYangCC updated https://github.com/llvm/llvm-project/pull/223337

>From 0a4c1ba92308073953e82c34283c3a66f8aaaea6 Mon Sep 17 00:00:00 2001
From: "chengcang.yang" <yangchengcang at gmail.com>
Date: Mon, 14 Sep 2026 17:15:16 +0800
Subject: [PATCH 1/2] [Offload] Check HSA allocate status before using the
 pointer

hsa_amd_memory_pool_allocate may leave the output pointer undefined on
failure. The AMDGPU pool allocator inspected that pointer for alignment
and could free it before checking Status. Classify the allocate result
first, and return a failed deallocation Error instead of dropping it.
---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    | 17 +++----
 .../amdgpu/utils/AMDGPUMemoryPoolAllocate.h   | 48 +++++++++++++++++++
 offload/unittests/OffloadAPI/CMakeLists.txt   |  3 ++
 .../memory/AMDGPUMemoryPoolAllocateTest.cpp   | 44 +++++++++++++++++
 4 files changed, 104 insertions(+), 8 deletions(-)
 create mode 100644 offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h
 create mode 100644 offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 281b9e3795a54d..ac1606873cba0b 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -30,6 +30,7 @@
 #include "Shared/Utils.h"
 #include "Utils/ELF.h"
 
+#include "AMDGPUMemoryPoolAllocate.h"
 #include "GlobalHandler.h"
 #include "OffloadAPI.h"
 #include "OpenMP/OMPT/Callback.h"
@@ -349,19 +350,19 @@ struct AMDGPUMemoryPoolTy {
     hsa_status_t Status =
         hsa_amd_memory_pool_allocate(MemoryPool, Size, 0, PtrStorage);
 
-    if (Alignment > 0 && !isAddrAligned(Align(Alignment), *PtrStorage)) {
-      if (auto FreeErr = deallocate(*PtrStorage)) {
-        return Plugin::error(ErrorCode::UNKNOWN,
-                             "Failure in deallcation of the incorrectly "
-                             "aligned pointer; requested alignemnt: %lu",
-                             Alignment);
-      }
+    MemoryPoolAllocateOutcome Outcome = classifyMemoryPoolAllocate(
+        Status == HSA_STATUS_SUCCESS, PtrStorage, Alignment);
+    if (Outcome == MemoryPoolAllocateOutcome::AllocateFailed)
+      return Plugin::check(Status, "error in hsa_amd_memory_pool_allocate: %s");
 
+    if (Outcome == MemoryPoolAllocateOutcome::PointerMisaligned) {
+      if (auto FreeErr = deallocate(*PtrStorage))
+        return FreeErr;
       return Plugin::error(ErrorCode::UNSUPPORTED,
                            "unsupported alignment size");
     }
 
-    return Plugin::check(Status, "error in hsa_amd_memory_pool_allocate: %s");
+    return Plugin::success();
   }
 
   /// Return memory to the memory pool.
diff --git a/offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h b/offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h
new file mode 100644
index 00000000000000..ee51b734f5924e
--- /dev/null
+++ b/offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h
@@ -0,0 +1,48 @@
+//===- AMDGPUMemoryPoolAllocate.h - HSA pool allocate result ----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OFFLOAD_PLUGINS_NEXTGEN_AMDGPU_UTILS_AMDGPUMEMORYPOOLALLOCATE_H
+#define OFFLOAD_PLUGINS_NEXTGEN_AMDGPU_UTILS_AMDGPUMEMORYPOOLALLOCATE_H
+
+#include "llvm/Support/Alignment.h"
+#include <cstddef>
+
+namespace llvm {
+namespace omp {
+namespace target {
+namespace plugin {
+
+/// Result of an HSA memory-pool allocate after the runtime call returns.
+enum class MemoryPoolAllocateOutcome {
+  AllocateFailed,
+  PointerMisaligned,
+  Success,
+};
+
+/// Classify a pool allocate. \p PointerStorage is read only when
+/// \p AllocateSucceeded is true, because a failed HSA allocate may leave the
+/// output pointer undefined.
+inline MemoryPoolAllocateOutcome
+classifyMemoryPoolAllocate(bool AllocateSucceeded, void *const *PointerStorage,
+                           size_t Alignment) {
+  if (!AllocateSucceeded)
+    return MemoryPoolAllocateOutcome::AllocateFailed;
+
+  if (Alignment > 0 &&
+      !llvm::isAddrAligned(llvm::Align(Alignment), *PointerStorage))
+    return MemoryPoolAllocateOutcome::PointerMisaligned;
+
+  return MemoryPoolAllocateOutcome::Success;
+}
+
+} // namespace plugin
+} // namespace target
+} // namespace omp
+} // namespace llvm
+
+#endif // OFFLOAD_PLUGINS_NEXTGEN_AMDGPU_UTILS_AMDGPUMEMORYPOOLALLOCATE_H
diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt
index 292ea1eb4852f0..51eaefec349a49 100644
--- a/offload/unittests/OffloadAPI/CMakeLists.txt
+++ b/offload/unittests/OffloadAPI/CMakeLists.txt
@@ -32,6 +32,7 @@ add_offload_unittest("kernel"
     kernel/olLaunchKernelCooperative.cpp)
 
 add_offload_unittest("memory"
+    memory/AMDGPUMemoryPoolAllocateTest.cpp
     memory/olMemAlloc.cpp
     memory/olMemAllocAligned.cpp
     memory/olMemFill.cpp
@@ -41,6 +42,8 @@ add_offload_unittest("memory"
     memory/olGetMemInfo.cpp
     memory/olGetMemInfoSize.cpp
     memory/olMemRegister.cpp)
+target_include_directories("memory.unittests" PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../plugins-nextgen/amdgpu/utils)
 
 add_offload_unittest("platform"
     platform/olGetPlatformInfo.cpp
diff --git a/offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp b/offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp
new file mode 100644
index 00000000000000..349fb384dc1d50
--- /dev/null
+++ b/offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp
@@ -0,0 +1,44 @@
+//===------- Offload tests - AMDGPU memory pool allocate ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUMemoryPoolAllocate.h"
+#include <cstdint>
+#include <gtest/gtest.h>
+
+using llvm::omp::target::plugin::classifyMemoryPoolAllocate;
+using llvm::omp::target::plugin::MemoryPoolAllocateOutcome;
+
+TEST(AMDGPUMemoryPoolAllocate, FailedStatusDoesNotInspectPointer) {
+  // A failed HSA allocate may leave this undefined. Using 0x1 would also fail
+  // a 16-byte alignment check, so inspecting it would be the wrong outcome.
+  void *Poison = reinterpret_cast<void *>(static_cast<uintptr_t>(1));
+  EXPECT_EQ(classifyMemoryPoolAllocate(false, &Poison, 16),
+            MemoryPoolAllocateOutcome::AllocateFailed);
+  EXPECT_EQ(classifyMemoryPoolAllocate(false, nullptr, 16),
+            MemoryPoolAllocateOutcome::AllocateFailed);
+}
+
+TEST(AMDGPUMemoryPoolAllocate, SuccessWithAlignedPointer) {
+  alignas(16) char Buffer[16];
+  void *Pointer = Buffer;
+  EXPECT_EQ(classifyMemoryPoolAllocate(true, &Pointer, 16),
+            MemoryPoolAllocateOutcome::Success);
+}
+
+TEST(AMDGPUMemoryPoolAllocate, SuccessWithMisalignedPointer) {
+  alignas(16) char Buffer[16];
+  void *Pointer = Buffer + 1;
+  EXPECT_EQ(classifyMemoryPoolAllocate(true, &Pointer, 16),
+            MemoryPoolAllocateOutcome::PointerMisaligned);
+}
+
+TEST(AMDGPUMemoryPoolAllocate, ZeroAlignmentSkipsAlignmentCheck) {
+  void *Pointer = reinterpret_cast<void *>(static_cast<uintptr_t>(1));
+  EXPECT_EQ(classifyMemoryPoolAllocate(true, &Pointer, 0),
+            MemoryPoolAllocateOutcome::Success);
+}

>From 0dd9914c07a54d97c3ca647a1c7f81535f472f13 Mon Sep 17 00:00:00 2001
From: "chengcang.yang" <yangchengcang at gmail.com>
Date: Tue, 15 Sep 2026 11:30:06 +0800
Subject: [PATCH 2/2] [Offload] Inline the AMDGPU pool allocate status check

Drop the helper header and check Status before using the output
pointer. A failed allocate may leave that pointer undefined.

Add olMemAllocAligned tests for oversized alignment and a failed
allocate that must not inspect the output pointer.
---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    |  9 ++--
 .../amdgpu/utils/AMDGPUMemoryPoolAllocate.h   | 48 -------------------
 offload/unittests/OffloadAPI/CMakeLists.txt   |  3 --
 .../memory/AMDGPUMemoryPoolAllocateTest.cpp   | 44 -----------------
 .../OffloadAPI/memory/olMemAllocAligned.cpp   | 26 ++++++++++
 5 files changed, 29 insertions(+), 101 deletions(-)
 delete mode 100644 offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h
 delete mode 100644 offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index ac1606873cba0b..b7b0420ca484f0 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -30,7 +30,6 @@
 #include "Shared/Utils.h"
 #include "Utils/ELF.h"
 
-#include "AMDGPUMemoryPoolAllocate.h"
 #include "GlobalHandler.h"
 #include "OffloadAPI.h"
 #include "OpenMP/OMPT/Callback.h"
@@ -349,13 +348,11 @@ struct AMDGPUMemoryPoolTy {
 
     hsa_status_t Status =
         hsa_amd_memory_pool_allocate(MemoryPool, Size, 0, PtrStorage);
-
-    MemoryPoolAllocateOutcome Outcome = classifyMemoryPoolAllocate(
-        Status == HSA_STATUS_SUCCESS, PtrStorage, Alignment);
-    if (Outcome == MemoryPoolAllocateOutcome::AllocateFailed)
+    // A failed allocate may leave *PtrStorage undefined; check Status first.
+    if (Status != HSA_STATUS_SUCCESS)
       return Plugin::check(Status, "error in hsa_amd_memory_pool_allocate: %s");
 
-    if (Outcome == MemoryPoolAllocateOutcome::PointerMisaligned) {
+    if (Alignment > 0 && !isAddrAligned(Align(Alignment), *PtrStorage)) {
       if (auto FreeErr = deallocate(*PtrStorage))
         return FreeErr;
       return Plugin::error(ErrorCode::UNSUPPORTED,
diff --git a/offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h b/offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h
deleted file mode 100644
index ee51b734f5924e..00000000000000
--- a/offload/plugins-nextgen/amdgpu/utils/AMDGPUMemoryPoolAllocate.h
+++ /dev/null
@@ -1,48 +0,0 @@
-//===- AMDGPUMemoryPoolAllocate.h - HSA pool allocate result ----*- C++ -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef OFFLOAD_PLUGINS_NEXTGEN_AMDGPU_UTILS_AMDGPUMEMORYPOOLALLOCATE_H
-#define OFFLOAD_PLUGINS_NEXTGEN_AMDGPU_UTILS_AMDGPUMEMORYPOOLALLOCATE_H
-
-#include "llvm/Support/Alignment.h"
-#include <cstddef>
-
-namespace llvm {
-namespace omp {
-namespace target {
-namespace plugin {
-
-/// Result of an HSA memory-pool allocate after the runtime call returns.
-enum class MemoryPoolAllocateOutcome {
-  AllocateFailed,
-  PointerMisaligned,
-  Success,
-};
-
-/// Classify a pool allocate. \p PointerStorage is read only when
-/// \p AllocateSucceeded is true, because a failed HSA allocate may leave the
-/// output pointer undefined.
-inline MemoryPoolAllocateOutcome
-classifyMemoryPoolAllocate(bool AllocateSucceeded, void *const *PointerStorage,
-                           size_t Alignment) {
-  if (!AllocateSucceeded)
-    return MemoryPoolAllocateOutcome::AllocateFailed;
-
-  if (Alignment > 0 &&
-      !llvm::isAddrAligned(llvm::Align(Alignment), *PointerStorage))
-    return MemoryPoolAllocateOutcome::PointerMisaligned;
-
-  return MemoryPoolAllocateOutcome::Success;
-}
-
-} // namespace plugin
-} // namespace target
-} // namespace omp
-} // namespace llvm
-
-#endif // OFFLOAD_PLUGINS_NEXTGEN_AMDGPU_UTILS_AMDGPUMEMORYPOOLALLOCATE_H
diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt
index 51eaefec349a49..292ea1eb4852f0 100644
--- a/offload/unittests/OffloadAPI/CMakeLists.txt
+++ b/offload/unittests/OffloadAPI/CMakeLists.txt
@@ -32,7 +32,6 @@ add_offload_unittest("kernel"
     kernel/olLaunchKernelCooperative.cpp)
 
 add_offload_unittest("memory"
-    memory/AMDGPUMemoryPoolAllocateTest.cpp
     memory/olMemAlloc.cpp
     memory/olMemAllocAligned.cpp
     memory/olMemFill.cpp
@@ -42,8 +41,6 @@ add_offload_unittest("memory"
     memory/olGetMemInfo.cpp
     memory/olGetMemInfoSize.cpp
     memory/olMemRegister.cpp)
-target_include_directories("memory.unittests" PRIVATE
-    ${CMAKE_CURRENT_SOURCE_DIR}/../../plugins-nextgen/amdgpu/utils)
 
 add_offload_unittest("platform"
     platform/olGetPlatformInfo.cpp
diff --git a/offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp b/offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp
deleted file mode 100644
index 349fb384dc1d50..00000000000000
--- a/offload/unittests/OffloadAPI/memory/AMDGPUMemoryPoolAllocateTest.cpp
+++ /dev/null
@@ -1,44 +0,0 @@
-//===------- Offload tests - AMDGPU memory pool allocate ------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "AMDGPUMemoryPoolAllocate.h"
-#include <cstdint>
-#include <gtest/gtest.h>
-
-using llvm::omp::target::plugin::classifyMemoryPoolAllocate;
-using llvm::omp::target::plugin::MemoryPoolAllocateOutcome;
-
-TEST(AMDGPUMemoryPoolAllocate, FailedStatusDoesNotInspectPointer) {
-  // A failed HSA allocate may leave this undefined. Using 0x1 would also fail
-  // a 16-byte alignment check, so inspecting it would be the wrong outcome.
-  void *Poison = reinterpret_cast<void *>(static_cast<uintptr_t>(1));
-  EXPECT_EQ(classifyMemoryPoolAllocate(false, &Poison, 16),
-            MemoryPoolAllocateOutcome::AllocateFailed);
-  EXPECT_EQ(classifyMemoryPoolAllocate(false, nullptr, 16),
-            MemoryPoolAllocateOutcome::AllocateFailed);
-}
-
-TEST(AMDGPUMemoryPoolAllocate, SuccessWithAlignedPointer) {
-  alignas(16) char Buffer[16];
-  void *Pointer = Buffer;
-  EXPECT_EQ(classifyMemoryPoolAllocate(true, &Pointer, 16),
-            MemoryPoolAllocateOutcome::Success);
-}
-
-TEST(AMDGPUMemoryPoolAllocate, SuccessWithMisalignedPointer) {
-  alignas(16) char Buffer[16];
-  void *Pointer = Buffer + 1;
-  EXPECT_EQ(classifyMemoryPoolAllocate(true, &Pointer, 16),
-            MemoryPoolAllocateOutcome::PointerMisaligned);
-}
-
-TEST(AMDGPUMemoryPoolAllocate, ZeroAlignmentSkipsAlignmentCheck) {
-  void *Pointer = reinterpret_cast<void *>(static_cast<uintptr_t>(1));
-  EXPECT_EQ(classifyMemoryPoolAllocate(true, &Pointer, 0),
-            MemoryPoolAllocateOutcome::Success);
-}
diff --git a/offload/unittests/OffloadAPI/memory/olMemAllocAligned.cpp b/offload/unittests/OffloadAPI/memory/olMemAllocAligned.cpp
index a894dde5ec77c7..ac57f16fd9aaaa 100644
--- a/offload/unittests/OffloadAPI/memory/olMemAllocAligned.cpp
+++ b/offload/unittests/OffloadAPI/memory/olMemAllocAligned.cpp
@@ -9,6 +9,7 @@
 #include "../common/Properties.hpp"
 #include <OffloadAPI.h>
 #include <gtest/gtest.h>
+#include <limits>
 
 using olMemAllocAlignedTest = OffloadDeviceTest;
 OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olMemAllocAlignedTest);
@@ -105,6 +106,31 @@ TEST_P(olMemAllocAlignedTest, CudaExceedDefaultAlignment) {
   ASSERT_EQ(Alloc, nullptr);
 }
 
+TEST_P(olMemAllocAlignedTest, AmdgpuExceedPoolAlignment) {
+  if (getPlatformBackend() != OL_PLATFORM_BACKEND_AMDGPU)
+    GTEST_SKIP() << "Test intended for AMDGPU backend";
+
+  void *Alloc = nullptr;
+  ASSERT_ERROR(OL_ERRC_UNSUPPORTED,
+               olMemAllocAligned(Device, OL_ALLOC_TYPE_DEVICE, 1024,
+                                 1024 * 64 * 64 * 64, &Alloc));
+  ASSERT_EQ(Alloc, nullptr);
+}
+
+TEST_P(olMemAllocAlignedTest, AmdgpuFailedAllocateDoesNotInspectPointer) {
+  if (getPlatformBackend() != OL_PLATFORM_BACKEND_AMDGPU)
+    GTEST_SKIP() << "Test intended for AMDGPU backend";
+
+  // A failed hsa_amd_memory_pool_allocate may leave the output pointer
+  // undefined. With a non-zero alignment the old code inspected and freed
+  // that pointer before checking Status.
+  void *Alloc = nullptr;
+  ASSERT_ANY_ERROR(olMemAllocAligned(Device, OL_ALLOC_TYPE_DEVICE,
+                                     std::numeric_limits<size_t>::max(),
+                                     DefaultAlignment, &Alloc));
+  ASSERT_EQ(Alloc, nullptr);
+}
+
 TEST_P(olMemAllocAlignedTypesTest, SuccessAllocDifferentAlignments) {
   void *Alloc = nullptr;
   size_t Alignments[] = {8, 16, 32, 64, 128, 256};



More information about the llvm-commits mailing list