[Offload] Add olGetMemInfo with platform-less API

[RossBrunton]

No description provided.

[llvmbot]

@llvm/pr-subscribers-offload

Author: Ross Brunton (RossBrunton)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/159581.diff

5 Files Affected:

• (modified) offload/liboffload/API/Memory.td (+50)
• (modified) offload/liboffload/src/OffloadImpl.cpp (+54)
• (modified) offload/unittests/OffloadAPI/CMakeLists.txt (+3-1)
• (added) offload/unittests/OffloadAPI/memory/olGetMemInfo.cpp (+130)
• (added) offload/unittests/OffloadAPI/memory/olGetMemInfoSize.cpp (+63)

diff --git a/offload/liboffload/API/Memory.td b/offload/liboffload/API/Memory.td
index debda165d2b23..3e47b586edd23 100644
--- a/offload/liboffload/API/Memory.td
+++ b/offload/liboffload/API/Memory.td
@@ -45,6 +45,56 @@ def olMemFree : Function {
   let returns = [];
 }
 
+def ol_mem_info_t : Enum {
+  let desc = "Supported memory info.";
+  let is_typed = 1;
+  let etors = [
+    TaggedEtor<"DEVICE", "ol_device_handle_t", "The handle of the device associated with the allocation.">,
+    TaggedEtor<"BASE", "void *", "Base address of this allocation.">,
+    TaggedEtor<"SIZE", "size_t", "Size of this allocation in bytes.">,
+    TaggedEtor<"TYPE", "ol_alloc_type_t", "Type of this allocation.">,
+  ];
+}
+
+def olGetMemInfo : Function {
+  let desc = "Queries the given property of a memory allocation allocated with olMemAlloc.";
+  let details = [
+    "`olGetMemInfoSize` can be used to query the storage size required for the given query.",
+    "The provided pointer can point to any location inside the allocation.",
+  ];
+  let params = [
+    Param<"const void *", "Ptr", "pointer to the allocated memory", PARAM_IN>,
+    Param<"ol_mem_info_t", "PropName", "type of the info to retrieve", PARAM_IN>,
+    Param<"size_t", "PropSize", "the number of bytes pointed to by PropValue.", PARAM_IN>,
+    TypeTaggedParam<"void*", "PropValue", "array of bytes holding the info. "
+      "If Size is not equal to or greater to the real number of bytes needed to return the info "
+      "then the OL_ERRC_INVALID_SIZE error is returned and pPlatformInfo is not used.", PARAM_OUT,
+      TypeInfo<"PropName" , "PropSize">>
+  ];
+  let returns = [
+    Return<"OL_ERRC_INVALID_SIZE", [
+      "`PropSize == 0`",
+      "If `PropSize` is less than the real number of bytes needed to return the info."
+    ]>,
+    Return<"OL_ERRC_NOT_FOUND", ["memory was not allocated by this platform"]>
+  ];
+}
+
+def olGetMemInfoSize : Function {
+  let desc = "Returns the storage size of the given queue query.";
+  let details = [
+    "The provided pointer can point to any location inside the allocation.",
+  ];
+  let params = [
+    Param<"const void *", "Ptr", "pointer to the allocated memory", PARAM_IN>,
+    Param<"ol_mem_info_t", "PropName", "type of the info to query", PARAM_IN>,
+    Param<"size_t*", "PropSizeRet", "pointer to the number of bytes required to store the query", PARAM_OUT>
+  ];
+  let returns = [
+    Return<"OL_ERRC_NOT_FOUND", ["memory was not allocated by this platform"]>
+  ];
+}
+
 def olMemcpy : Function {
     let desc = "Enqueue a memcpy operation.";
     let details = [
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 7e4a309966c9d..20048e2eda07e 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -701,6 +701,60 @@ Error olMemFree_impl(void *Address) {
   return Error::success();
 }
 
+Error olGetMemInfoImplDetail(const void *Ptr, ol_mem_info_t PropName,
+                             size_t PropSize, void *PropValue,
+                             size_t *PropSizeRet) {
+  InfoWriter Info(PropSize, PropValue, PropSizeRet);
+  std::lock_guard<std::mutex> Lock(OffloadContext::get().AllocInfoMapMutex);
+
+  auto &AllocBases = OffloadContext::get().AllocBases;
+  auto &AllocInfoMap = OffloadContext::get().AllocInfoMap;
+  const AllocInfo *Alloc = nullptr;
+  if (AllocInfoMap.contains(Ptr)) {
+    // Fast case, we have been given the base pointer directly
+    Alloc = &AllocInfoMap.at(Ptr);
+  } else {
+    // Slower case, we need to look up the base pointer first
+    // Find the first memory allocation whose end is after the target pointer,
+    // and then check to see if it is in range
+    auto Loc = std::lower_bound(AllocBases.begin(), AllocBases.end(), Ptr,
+                                [&](const void *Iter, const void *Val) {
+                                  return AllocInfoMap.at(Iter).End <= Val;
+                                });
+    if (Loc == AllocBases.end() || Ptr < AllocInfoMap.at(*Loc).Start)
+      return Plugin::error(ErrorCode::NOT_FOUND,
+                           "allocated memory information not found");
+    Alloc = &AllocInfoMap.at(*Loc);
+  }
+
+  switch (PropName) {
+  case OL_MEM_INFO_DEVICE:
+    return Info.write<ol_device_handle_t>(Alloc->Device);
+  case OL_MEM_INFO_BASE:
+    return Info.write<void *>(Alloc->Start);
+  case OL_MEM_INFO_SIZE:
+    return Info.write<size_t>(static_cast<char *>(Alloc->End) -
+                              static_cast<char *>(Alloc->Start));
+  case OL_MEM_INFO_TYPE:
+    return Info.write<ol_alloc_type_t>(Alloc->Type);
+  default:
+    return createOffloadError(ErrorCode::INVALID_ENUMERATION,
+                              "olGetMemInfo enum '%i' is invalid", PropName);
+  }
+
+  return Error::success();
+}
+
+Error olGetMemInfo_impl(const void *Ptr, ol_mem_info_t PropName,
+                        size_t PropSize, void *PropValue) {
+  return olGetMemInfoImplDetail(Ptr, PropName, PropSize, PropValue, nullptr);
+}
+
+Error olGetMemInfoSize_impl(const void *Ptr, ol_mem_info_t PropName,
+                            size_t *PropSizeRet) {
+  return olGetMemInfoImplDetail(Ptr, PropName, 0, nullptr, PropSizeRet);
+}
+
 Error olCreateQueue_impl(ol_device_handle_t Device, ol_queue_handle_t *Queue) {
   auto CreatedQueue = std::make_unique<ol_queue_impl_t>(nullptr, Device);
 
diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt
index ba35c1ee87aac..50c99a5d5b639 100644
--- a/offload/unittests/OffloadAPI/CMakeLists.txt
+++ b/offload/unittests/OffloadAPI/CMakeLists.txt
@@ -27,7 +27,9 @@ add_offload_unittest("memory"
     memory/olMemAlloc.cpp
     memory/olMemFill.cpp
     memory/olMemFree.cpp
-    memory/olMemcpy.cpp)
+    memory/olMemcpy.cpp
+    memory/olGetMemInfo.cpp
+    memory/olGetMemInfoSize.cpp)
 
 add_offload_unittest("platform"
     platform/olGetPlatformInfo.cpp
diff --git a/offload/unittests/OffloadAPI/memory/olGetMemInfo.cpp b/offload/unittests/OffloadAPI/memory/olGetMemInfo.cpp
new file mode 100644
index 0000000000000..a4b382ff298ad
--- /dev/null
+++ b/offload/unittests/OffloadAPI/memory/olGetMemInfo.cpp
@@ -0,0 +1,130 @@
+//===------- Offload API tests - olGetMemInfo -----------------------------===//
+//
+// 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>
+
+constexpr size_t SIZE = 1024;
+
+struct olGetMemInfoBaseTest : OffloadDeviceTest {
+  void *OffsetPtr() { return &reinterpret_cast<char *>(Ptr)[123]; }
+
+  void *Ptr;
+};
+
+template <ol_alloc_type_t AllocType>
+struct olGetMemInfoTest : olGetMemInfoBaseTest {
+  void SetUp() override {
+    RETURN_ON_FATAL_FAILURE(OffloadDeviceTest::SetUp());
+    ASSERT_SUCCESS(olMemAlloc(Device, AllocType, SIZE, &Ptr));
+  }
+
+  void TearDown() override {
+    ASSERT_SUCCESS(olMemFree(Ptr));
+    RETURN_ON_FATAL_FAILURE(OffloadDeviceTest::TearDown());
+  }
+};
+using olGetMemInfoDeviceTest = olGetMemInfoTest<OL_ALLOC_TYPE_DEVICE>;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olGetMemInfoDeviceTest);
+using olGetMemInfoManagedTest = olGetMemInfoTest<OL_ALLOC_TYPE_MANAGED>;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olGetMemInfoManagedTest);
+using olGetMemInfoHostTest = olGetMemInfoTest<OL_ALLOC_TYPE_HOST>;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olGetMemInfoHostTest);
+
+#define PER_ALLOC_TEST(FUNCTION)                                               \
+  TEST_P(olGetMemInfoDeviceTest, FUNCTION) {                                   \
+    FUNCTION(this, Ptr, OL_ALLOC_TYPE_DEVICE);                                 \
+  }                                                                            \
+  TEST_P(olGetMemInfoManagedTest, FUNCTION) {                                  \
+    FUNCTION(this, Ptr, OL_ALLOC_TYPE_MANAGED);                                \
+  }                                                                            \
+  TEST_P(olGetMemInfoHostTest, FUNCTION) {                                     \
+    FUNCTION(this, OffsetPtr(), OL_ALLOC_TYPE_HOST);                           \
+  }                                                                            \
+  TEST_P(olGetMemInfoDeviceTest, FUNCTION##Offset) {                           \
+    FUNCTION(this, Ptr, OL_ALLOC_TYPE_DEVICE);                                 \
+  }                                                                            \
+  TEST_P(olGetMemInfoManagedTest, FUNCTION##Offset) {                          \
+    FUNCTION(this, OffsetPtr(), OL_ALLOC_TYPE_MANAGED);                        \
+  }                                                                            \
+  TEST_P(olGetMemInfoHostTest, FUNCTION##Offset) {                             \
+    FUNCTION(this, OffsetPtr(), OL_ALLOC_TYPE_HOST);                           \
+  }
+
+void SuccessDevice(olGetMemInfoBaseTest *Fixture, void *Ptr,
+                   ol_alloc_type_t Type) {
+  ol_device_handle_t RetrievedDevice;
+  ASSERT_SUCCESS(olGetMemInfo(Fixture->Ptr, OL_MEM_INFO_DEVICE,
+                              sizeof(RetrievedDevice), &RetrievedDevice));
+  ASSERT_EQ(RetrievedDevice, Fixture->Device);
+}
+PER_ALLOC_TEST(SuccessDevice);
+
+void SuccessBase(olGetMemInfoBaseTest *Fixture, void *Ptr,
+                 ol_alloc_type_t Type) {
+  void *RetrievedBase;
+  ASSERT_SUCCESS(olGetMemInfo(Fixture->Ptr, OL_MEM_INFO_BASE,
+                              sizeof(RetrievedBase), &RetrievedBase));
+  ASSERT_EQ(RetrievedBase, Fixture->Ptr);
+}
+PER_ALLOC_TEST(SuccessBase);
+
+void SuccessSize(olGetMemInfoBaseTest *Fixture, void *Ptr,
+                 ol_alloc_type_t Type) {
+  size_t RetrievedSize;
+  ASSERT_SUCCESS(olGetMemInfo(Fixture->Ptr, OL_MEM_INFO_SIZE,
+                              sizeof(RetrievedSize), &RetrievedSize));
+  ASSERT_EQ(RetrievedSize, SIZE);
+}
+PER_ALLOC_TEST(SuccessSize);
+
+void SuccessType(olGetMemInfoBaseTest *Fixture, void *Ptr,
+                 ol_alloc_type_t Type) {
+  ol_alloc_type_t RetrievedType;
+  ASSERT_SUCCESS(olGetMemInfo(Fixture->Ptr, OL_MEM_INFO_TYPE,
+                              sizeof(RetrievedType), &RetrievedType));
+  ASSERT_EQ(RetrievedType, Type);
+}
+PER_ALLOC_TEST(SuccessType);
+
+TEST_P(olGetMemInfoDeviceTest, InvalidNotFound) {
+  // Assuming that we aren't unlucky and happen to get 0x1234 as a random
+  // pointer
+  void *RetrievedBase;
+  ASSERT_ERROR(OL_ERRC_NOT_FOUND,
+               olGetMemInfo(reinterpret_cast<void *>(0x1234), OL_MEM_INFO_BASE,
+                            sizeof(RetrievedBase), &RetrievedBase));
+}
+
+TEST_P(olGetMemInfoDeviceTest, InvalidNullPtr) {
+  ol_device_handle_t RetrievedDevice;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olGetMemInfo(nullptr, OL_MEM_INFO_DEVICE,
+                            sizeof(RetrievedDevice), &RetrievedDevice));
+}
+
+TEST_P(olGetMemInfoDeviceTest, InvalidSizeZero) {
+  ol_device_handle_t RetrievedDevice;
+  ASSERT_ERROR(OL_ERRC_INVALID_SIZE,
+               olGetMemInfo(Ptr, OL_MEM_INFO_DEVICE, 0, &RetrievedDevice));
+}
+
+TEST_P(olGetMemInfoDeviceTest, InvalidSizeSmall) {
+  ol_device_handle_t RetrievedDevice;
+  ASSERT_ERROR(OL_ERRC_INVALID_SIZE,
+               olGetMemInfo(Ptr, OL_MEM_INFO_DEVICE,
+                            sizeof(RetrievedDevice) - 1, &RetrievedDevice));
+}
+
+TEST_P(olGetMemInfoDeviceTest, InvalidNullPointerPropValue) {
+  ol_device_handle_t RetrievedDevice;
+  ASSERT_ERROR(
+      OL_ERRC_INVALID_NULL_POINTER,
+      olGetMemInfo(Ptr, OL_MEM_INFO_DEVICE, sizeof(RetrievedDevice), nullptr));
+}
diff --git a/offload/unittests/OffloadAPI/memory/olGetMemInfoSize.cpp b/offload/unittests/OffloadAPI/memory/olGetMemInfoSize.cpp
new file mode 100644
index 0000000000000..f1a1e790fb22f
--- /dev/null
+++ b/offload/unittests/OffloadAPI/memory/olGetMemInfoSize.cpp
@@ -0,0 +1,63 @@
+//===------- Offload API tests - olGetMemInfoSize -------------------------===//
+//
+// 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 <OffloadAPI.h>
+
+#include "../common/Fixtures.hpp"
+
+struct olGetMemInfoSizeTest : OffloadDeviceTest {
+  void *OffsetPtr() { return &reinterpret_cast<char *>(Ptr)[123]; }
+
+  void SetUp() override {
+    RETURN_ON_FATAL_FAILURE(OffloadDeviceTest::SetUp());
+    ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_DEVICE, 0x1024, &Ptr));
+  }
+
+  void TearDown() override {
+    ASSERT_SUCCESS(olMemFree(Ptr));
+    RETURN_ON_FATAL_FAILURE(OffloadDeviceTest::TearDown());
+  }
+
+  void *Ptr;
+};
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olGetMemInfoSizeTest);
+
+TEST_P(olGetMemInfoSizeTest, SuccessDevice) {
+  size_t Size = 0;
+  ASSERT_SUCCESS(olGetMemInfoSize(Ptr, OL_MEM_INFO_DEVICE, &Size));
+  ASSERT_EQ(Size, sizeof(ol_device_handle_t));
+}
+
+TEST_P(olGetMemInfoSizeTest, SuccessBase) {
+  size_t Size = 0;
+  ASSERT_SUCCESS(olGetMemInfoSize(Ptr, OL_MEM_INFO_BASE, &Size));
+  ASSERT_EQ(Size, sizeof(void *));
+}
+
+TEST_P(olGetMemInfoSizeTest, SuccessSize) {
+  size_t Size = 0;
+  ASSERT_SUCCESS(olGetMemInfoSize(Ptr, OL_MEM_INFO_SIZE, &Size));
+  ASSERT_EQ(Size, sizeof(size_t));
+}
+
+TEST_P(olGetMemInfoSizeTest, SuccessType) {
+  size_t Size = 0;
+  ASSERT_SUCCESS(olGetMemInfoSize(Ptr, OL_MEM_INFO_TYPE, &Size));
+  ASSERT_EQ(Size, sizeof(ol_alloc_type_t));
+}
+
+TEST_P(olGetMemInfoSizeTest, InvalidSymbolInfoEnumeration) {
+  size_t Size = 0;
+  ASSERT_ERROR(OL_ERRC_INVALID_ENUMERATION,
+               olGetMemInfoSize(Ptr, OL_MEM_INFO_FORCE_UINT32, &Size));
+}
+
+TEST_P(olGetMemInfoSizeTest, InvalidNullPointer) {
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olGetMemInfoSize(Ptr, OL_MEM_INFO_DEVICE, nullptr));
+}