[llvm] [offload] Document and assert allocation alignment guarantee (PR #223381)

[Weronika Lewandowska via llvm-commits]

https://github.com/wlemkows updated https://github.com/llvm/llvm-project/pull/223381

>From 14b9d00f2ed1444b31d7e36a15e008bdee067389 Mon Sep 17 00:00:00 2001
From: Weronika Lewandowska <weronika.lewandowska at intel.com>
Date: Mon, 14 Sep 2026 13:12:09 +0000
Subject: [PATCH] [offload] Document and assert allocation alignment guarantee

---
 offload/liboffload/API/Memory.td                  |  4 +++-
 offload/liboffload/src/OffloadImpl.cpp            |  6 ++++++
 .../unittests/OffloadAPI/memory/olMemAlloc.cpp    | 15 +++++++++++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/offload/liboffload/API/Memory.td b/offload/liboffload/API/Memory.td
index 17eabeccce9e0..6e89e84b75ab4 100644
--- a/offload/liboffload/API/Memory.td
+++ b/offload/liboffload/API/Memory.td
@@ -37,6 +37,7 @@ def olMemAlloc : Function {
   let desc = "Creates a memory allocation on the specified device.";
   let details = [
       "All liboffload allocations share a single virtual address range. There is no risk of multiple devices returning equal pointers to different memory.",
+      "The returned pointer has word-size alignment (`alignof(uintptr_t)`), regardless of allocation size.",
       "This function can only be used to create device or managed allocations. To create a host allocation use `olMemAllocHost`."
   ];
   let params = [
@@ -58,7 +59,8 @@ def olMemAlloc : Function {
 def olMemAllocHost : Function {
   let desc = "Creates a host memory allocation accessible from the specified device.";
   let details = [
-      "All liboffload allocations share a single virtual address range. There is no risk of multiple devices returning equal pointers to different memory."
+      "All liboffload allocations share a single virtual address range. There is no risk of multiple devices returning equal pointers to different memory.",
+      "The returned pointer has word-size alignment (`alignof(uintptr_t)`), regardless of allocation size."
   ];
   let params = [
     Param<"ol_device_handle_t", "Device", "handle of the device to allocate on", PARAM_IN>,
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 11c07483b38aa..e37f013a96dd4 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -716,6 +716,12 @@ Error olMemAllocImplHelper(ol_device_handle_t Device, ol_alloc_type_t Type,
     if (!NewAlloc)
       return NewAlloc.takeError();
 
+    // Check word-size alignment only when no explicit alignment was requested
+    // (olMemAlloc and olMemAllocHost pass zero for the default alignment).
+    if (Alignment == 0)
+      assert(reinterpret_cast<uintptr_t>(*NewAlloc) % alignof(uintptr_t) == 0 &&
+             "allocation does not meet word-size alignment");
+
     void *NewEnd = &static_cast<char *>(*NewAlloc)[Size];
     auto &AllocBases = OffloadContext::get().AllocBases;
     auto &AllocInfoMap = OffloadContext::get().AllocInfoMap;
diff --git a/offload/unittests/OffloadAPI/memory/olMemAlloc.cpp b/offload/unittests/OffloadAPI/memory/olMemAlloc.cpp
index 3dfd8de06eea6..6940bfae445ca 100644
--- a/offload/unittests/OffloadAPI/memory/olMemAlloc.cpp
+++ b/offload/unittests/OffloadAPI/memory/olMemAlloc.cpp
@@ -35,6 +35,21 @@ TEST_P(olMemAllocAllocTypesTest, Success) {
   olMemFree(Alloc);
 }
 
+TEST_P(olMemAllocAllocTypesTest, SuccessSmallSizesWordSizeAlignment) {
+  const size_t Sizes[] = {1, 2, 3, 4, 5, 6, 7, 8, 14};
+  for (size_t Size : Sizes) {
+    SCOPED_TRACE("size: " + std::to_string(Size));
+    void *Alloc = nullptr;
+    ASSERT_SUCCESS(allocateDeviceOrHost(Size, &Alloc));
+    ASSERT_NE(Alloc, nullptr);
+
+    // Check word-size alignment without requesting explicit alignment.
+    EXPECT_EQ(reinterpret_cast<uintptr_t>(Alloc) % alignof(uintptr_t), 0U);
+
+    ASSERT_SUCCESS(olMemFree(Alloc));
+  }
+}
+
 TEST_P(olMemAllocTest, SuccessAllocMany) {
   std::vector<void *> Allocs;
   Allocs.reserve(TestAllocsNum);