[llvm] [offload] add support for aligned allocations (PR #203353)

[Alex Duran via llvm-commits]

================
@@ -0,0 +1,207 @@
+//===--------------- Offload API tests - olMemAllocAligned ----------------===//
+//
+// 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 "../common/Fixtures.hpp"
+#include <OffloadAPI.h>
+#include <gtest/gtest.h>
+
+using olMemAllocAlignedTest = OffloadDeviceTest;
+
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olMemAllocAlignedTest);
+
+constexpr size_t DefaultAlignment = 16;
+constexpr size_t TestAllocsNum = 1000;
+
+TEST_P(olMemAllocAlignedTest, SuccessAllocMany) {
+  std::vector<void *> Allocs;
+  Allocs.reserve(1000);
+
+  constexpr ol_alloc_type_t TYPES[3] = {
+      OL_ALLOC_TYPE_DEVICE, OL_ALLOC_TYPE_MANAGED, OL_ALLOC_TYPE_HOST};
+
+  for (size_t I = 1; I < TestAllocsNum; I++) {
+    void *Alloc = nullptr;
+    ASSERT_SUCCESS(olMemAllocAligned(Device, TYPES[I % 3], 1024 * I,
+                                     DefaultAlignment, &Alloc));
+    ASSERT_NE(Alloc, nullptr);
+
+    Allocs.push_back(Alloc);
+  }
+
+  for (auto *A : Allocs) {
+    olMemFree(A);
+  }
+}
+
+TEST_P(olMemAllocAlignedTest, InvalidNullDevice) {
+  void *Alloc = nullptr;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_HANDLE,
+               olMemAllocAligned(nullptr, OL_ALLOC_TYPE_DEVICE, 1024,
+                                 DefaultAlignment, &Alloc));
+}
+
+TEST_P(olMemAllocAlignedTest, InvalidNullOutPtr) {
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olMemAllocAligned(Device, OL_ALLOC_TYPE_DEVICE, 1024,
+                                 DefaultAlignment, nullptr));
+}
+
+TEST_P(olMemAllocAlignedTest, InvalidAlignmentZero) {
+  void *Alloc = nullptr;
+
+  ASSERT_ERROR(
+      OL_ERRC_INVALID_ARGUMENT,
+      olMemAllocAligned(Device, OL_ALLOC_TYPE_DEVICE, 1024, 0, &Alloc));
+}
+
+TEST_P(olMemAllocAlignedTest, InvalidAlignmentNotAPowerOfTwo) {
+  void *Alloc = nullptr;
+
+  ASSERT_ERROR(
+      OL_ERRC_INVALID_ARGUMENT,
+      olMemAllocAligned(Device, OL_ALLOC_TYPE_DEVICE, 1024, 3, &Alloc));
+}
+
+TEST_P(olMemAllocAlignedTest, CudaExceedDefaultAlignment) {
+  if (getPlatformBackend() != OL_PLATFORM_BACKEND_CUDA) {
+    GTEST_SKIP() << "Test inteded for CUDA backend";
+  }
+
+  void *Alloc = nullptr;
+  // The default page size for cuda is 64 KB
+  ASSERT_ERROR(OL_ERRC_UNSUPPORTED,
+               olMemAllocAligned(Device, OL_ALLOC_TYPE_DEVICE, 1024,
+                                 1024 * 64 * 64, &Alloc));
+  ASSERT_EQ(Alloc, nullptr);
+}
+
+TEST_P(olMemAllocAlignedTest, SuccessAllocManagedDifferentAlignments) {
+  void *Alloc = nullptr;
+  size_t num_alignments = 6;
+  size_t alignments[] = {8, 16, 32, 64, 128, 256};
+
+  for (size_t i = 0; i < num_alignments; i++) {
+    SCOPED_TRACE("alignment: " + std::to_string(alignments[i]));
+    ASSERT_SUCCESS(olMemAllocAligned(Device, OL_ALLOC_TYPE_MANAGED, 1024,
+                                     alignments[i], &Alloc));
+    ASSERT_NE(Alloc, nullptr);
+    olMemFree(Alloc);
+  }
+}
+
+TEST_P(olMemAllocAlignedTest, SuccessAllocHostDifferentAlignments) {
+  void *Alloc = nullptr;
+  size_t num_alignments = 6;
+  size_t alignments[] = {8, 16, 32, 64, 128, 256};
+
+  for (size_t i = 0; i < num_alignments; i++) {
+    SCOPED_TRACE("alignment: " + std::to_string(alignments[i]));
+    ASSERT_SUCCESS(olMemAllocAligned(Device, OL_ALLOC_TYPE_HOST, 1024,
+                                     alignments[i], &Alloc));
+    ASSERT_NE(Alloc, nullptr);
+    olMemFree(Alloc);
+  }
+}
+
+TEST_P(olMemAllocAlignedTest, SuccessAllocDeviceDifferentAlignments) {
+  void *Alloc = nullptr;
+  size_t num_alignments = 6;
+  size_t alignments[] = {8, 16, 32, 64, 128, 256};
+  size_t alignment;
----------------
adurang wrote:

```suggestion
  size_t Alignment;
```

https://github.com/llvm/llvm-project/pull/203353