[llvm] [libsycl] USM Aligned allocation functions (PR #213468)

[Kseniya Tikhomirova via llvm-commits]

================
@@ -0,0 +1,238 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 <mock/helpers.hpp>
+
+#include <sycl/__impl/device.hpp>
+#include <sycl/__impl/queue.hpp>
+#include <sycl/__impl/usm_functions.hpp>
+
+#include <detail/device_impl.hpp>
+#include <detail/queue_impl.hpp>
+
+#include <cstddef>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+using namespace sycl;
+using namespace ::testing;
+
+constexpr size_t NumBytes = 1024;
+constexpr size_t Alignment = 256;
+
+TEST(USMFunctions, DeviceAllocation) {
+  mock::MockWrapper Mock;
+  queue Q;
+  device Dev = Q.get_device();
+  context Ctx = Q.get_context();
+  ol_device_handle_t OLDev = detail::getSyclObjImpl(Dev)->getOLHandle();
+
+  EXPECT_CALL(Mock.get(), olMemAlloc(OLDev, OL_ALLOC_TYPE_DEVICE, NumBytes, _))
+      .Times(1);
+  void *Ptr1 = malloc_device(NumBytes, Dev, Ctx);
+  EXPECT_NE(Ptr1, nullptr);
+
+  EXPECT_CALL(Mock.get(), olMemFree(Ptr1)).Times(1);
+  free(Ptr1, Ctx);
+
+  EXPECT_CALL(Mock.get(), olMemAllocAligned(OLDev, OL_ALLOC_TYPE_DEVICE,
+                                            NumBytes, Alignment, _))
+      .Times(1);
+  void *Ptr2 = aligned_alloc_device(Alignment, NumBytes, Q);
+  EXPECT_NE(Ptr2, nullptr);
+
+  EXPECT_CALL(Mock.get(), olMemFree(Ptr2)).Times(1);
+  free(Ptr2, Q);
+}
+
+TEST(USMFunctions, HostAllocation) {
+  mock::MockWrapper Mock;
+  queue Q;
+  context Ctx = Q.get_context();
+  ol_device_handle_t OLDev =
+      detail::getSyclObjImpl(Q.get_device())->getOLHandle();
+
+  EXPECT_CALL(Mock.get(), olMemAllocHost(OLDev, NumBytes, _)).Times(1);
+  void *Ptr1 = malloc_host(NumBytes, Ctx);
+  EXPECT_NE(Ptr1, nullptr);
+
+  EXPECT_CALL(Mock.get(), olMemFree(Ptr1)).Times(1);
+  free(Ptr1, Ctx);
+
+  EXPECT_CALL(Mock.get(), olMemAllocAlignedHost(OLDev, NumBytes, Alignment, _))
+      .Times(1);
+  void *Ptr2 = aligned_alloc_host(Alignment, NumBytes, Ctx);
+  EXPECT_NE(Ptr2, nullptr);
+
+  EXPECT_CALL(Mock.get(), olMemFree(Ptr2)).Times(1);
+  free(Ptr2, Ctx);
+}
+
+TEST(USMFunctions, SharedAllocation) {
+  mock::MockWrapper Mock;
+  queue Q;
+  device Dev = Q.get_device();
+  context Ctx = Q.get_context();
+  ol_device_handle_t OLDev = detail::getSyclObjImpl(Dev)->getOLHandle();
+
+  EXPECT_CALL(Mock.get(), olMemAlloc(OLDev, OL_ALLOC_TYPE_MANAGED, NumBytes, _))
+      .Times(1);
+  void *Ptr1 = malloc_shared(NumBytes, Dev, Ctx);
+  EXPECT_NE(Ptr1, nullptr);
+
+  EXPECT_CALL(Mock.get(), olMemFree(Ptr1)).Times(1);
+  free(Ptr1, Ctx);
+
+  EXPECT_CALL(Mock.get(), olMemAllocAligned(OLDev, OL_ALLOC_TYPE_MANAGED,
+                                            NumBytes, Alignment, _))
+      .Times(1);
+  void *Ptr2 = aligned_alloc_shared(Alignment, NumBytes, Q);
+  EXPECT_NE(Ptr2, nullptr);
+
+  EXPECT_CALL(Mock.get(), olMemFree(Ptr2)).Times(1);
+  free(Ptr2, Q);
+}
+
+TEST(USMFunctions, ZeroByteAllocation) {
+  mock::MockWrapper Mock;
+  queue Q;
+  device Dev = Q.get_device();
+  context Ctx = Q.get_context();
+
+  EXPECT_CALL(Mock.get(), olMemAlloc(_, _, _, _)).Times(0);
+  EXPECT_CALL(Mock.get(), olMemAllocHost(_, _, _)).Times(0);
+  EXPECT_CALL(Mock.get(), olMemAllocAligned(_, _, _, _, _)).Times(0);
+  EXPECT_CALL(Mock.get(), olMemAllocAlignedHost(_, _, _, _)).Times(0);
+
+  EXPECT_EQ(malloc_device(0, Dev, Ctx), nullptr);
+  EXPECT_EQ(malloc_shared(0, Dev, Ctx), nullptr);
+  EXPECT_EQ(malloc_host(0, Ctx), nullptr);
+}
+
+TEST(USMFunctions, InvalidAlignment) {
+  mock::MockWrapper Mock;
+  queue Q;
+  device Dev = Q.get_device();
+  context Ctx = Q.get_context();
+  ol_device_handle_t OLDev = detail::getSyclObjImpl(Dev)->getOLHandle();
+
+  constexpr size_t NonPowerOf2Alignment = 3;
+
+  EXPECT_CALL(Mock.get(), olMemAllocAligned(OLDev, OL_ALLOC_TYPE_DEVICE,
+                                            NumBytes, NonPowerOf2Alignment, _))
+      .Times(1)
+      .WillOnce(Return(mock::getMockLiboffload().makeEmptyStrError(
+          OL_ERRC_INVALID_ARGUMENT)));
+  EXPECT_EQ(aligned_alloc_device(NonPowerOf2Alignment, NumBytes, Dev, Ctx),
+            nullptr);
+
+  EXPECT_CALL(Mock.get(),
+              olMemAllocAlignedHost(OLDev, NumBytes, NonPowerOf2Alignment, _))
+      .Times(1)
+      .WillOnce(Return(mock::getMockLiboffload().makeEmptyStrError(
----------------
KseniyaTikhomirova wrote:

this overriding of default action can be removed once you fix my previous comment

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