[OFFLOAD] Add memory data locking API for libomptarget migration (#173138)

Add liboffload memory data locking API for libomptarget migration

This PR adds liboffload memory data locking API that needed to make
libomptarget to use liboffload

offload/liboffload/API/Memory.td

@@ -19,6 +19,11 @@ def ol_alloc_type_t : Enum {
   ];
 }
 
+def ol_memory_register_flags_t : Typedef {
+  let desc = "Memory registering/locking flags";
+  let value = "uint32_t";
+}
+
 def olMemAlloc : Function {
   let desc = "Creates a memory allocation on the specified device.";
   let details = [
@@ -131,3 +136,43 @@ def olMemFill : Function {
     Return<"OL_ERRC_INVALID_SIZE", ["`FillSize % PatternSize != 0`"]>
   ];
 }
+
+def olMemRegister : Function {
+  let desc = "Register and page-lock host memory so it can be accessible by the device.";
+  let details = [
+        "Pins host memory to optimize transfers and returns the device accessible",
+        "stable pointer that devices should use for memory transfers involving the host",
+        "pinned allocation. If the buffer intersects with other existing buffer,",
+        "a new user will be registered. A partial overlapping is not allowed.",
+        "The pinned pointer can be accessed both on host and device and",
+        "no guarantees are made about consistency.",
+        "The pinned pointer should be used to execute memory transfers",
+        "as it is a stable pointer for memory access."  
+  ];
+  let params = [
+    Param<"ol_device_handle_t", "Device", "handle of the device", PARAM_IN>,
+    Param<"void *", "Ptr", "host pointer", PARAM_IN>,
+    Param<"size_t", "Size", "size of the memory in bytes", PARAM_IN>,
+    Param<"ol_memory_register_flags_t", "Flags", "flags. Reserved for future use", PARAM_IN>,
+    Param<"void**", "PinnedPtr", "pointer to the pinned memory", PARAM_OUT>
+  ];
+  let returns = [
+    Return<"OL_ERRC_INVALID_SIZE", [
+      "`Size == 0`"
+    ]>
+  ];
+}
+
+def olMemUnregister : Function {
+  let desc = "Unregister and page-unlock host memory.";
+  let details = [
+        "Unpins host memory that was previously pinned or unregister the buffer",
+        "if other users are still using the buffer. If no users are using the buffer",
+        "the memory is unlocked."
+  ];
+  let params = [
+    Param<"ol_device_handle_t", "Device", "handle of the device", PARAM_IN>,
+    Param<"void *", "Ptr", "host pointer", PARAM_IN>
+  ];
+  let returns = [];
+}

offload/liboffload/src/OffloadImpl.cpp

@@ -1216,5 +1216,20 @@ Error olLaunchHostFunction_impl(ol_queue_handle_t Queue,
                                                 Queue->AsyncInfo);
 }
 
+Error olMemRegister_impl(ol_device_handle_t Device, void *Ptr, size_t Size,
+                         ol_memory_register_flags_t flags, void **LockedPtr) {
+  Expected<void *> LockedPtrOrErr = Device->Device->dataLock(Ptr, Size);
+  if (!LockedPtrOrErr)
+    return LockedPtrOrErr.takeError();
+
+  *LockedPtr = *LockedPtrOrErr;
+
+  return Error::success();
+}
+
+Error olMemUnregister_impl(ol_device_handle_t Device, void *Ptr) {
+  return Device->Device->dataUnlock(Ptr);
+}
+
 } // namespace offload
 } // namespace llvm

offload/plugins-nextgen/level_zero/include/L0Device.h

@@ -556,9 +556,10 @@ public:
                             TargetAllocTy Kind) override;
   Error free(void *TgtPtr, TargetAllocTy Kind = TARGET_ALLOC_DEFAULT) override;
 
+  /// This plugin does nothing to lock buffers. Do not return an error, just
+  /// return the same pointer as the device pointer.
   Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) override {
-    return Plugin::error(error::ErrorCode::UNKNOWN,
-                         "dataLockImpl not supported");
+    return HstPtr;
   }
   Error dataUnlockImpl(void *HstPtr) override { return Plugin::success(); }
 

offload/unittests/OffloadAPI/CMakeLists.txt

@@ -29,7 +29,8 @@ add_offload_unittest("memory"
     memory/olMemFree.cpp
     memory/olMemcpy.cpp
     memory/olGetMemInfo.cpp
-    memory/olGetMemInfoSize.cpp)
+    memory/olGetMemInfoSize.cpp
+    memory/olMemRegister.cpp)
 
 add_offload_unittest("platform"
     platform/olGetPlatformInfo.cpp

offload/unittests/OffloadAPI/memory/olMemRegister.cpp

@@ -0,0 +1,84 @@
+//===------- Offload API tests - olMemRegister ----------------------------===//
+//
+// 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 olMemRegisterTest = OffloadDeviceTest;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olMemRegisterTest);
+
+TEST_P(olMemRegisterTest, SuccessRegister) {
+  int Arr[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  ASSERT_SUCCESS(olMemRegister(Device, Arr, sizeof(Arr), Flags, &PinnedPtr));
+  ASSERT_NE(PinnedPtr, nullptr);
+  ASSERT_SUCCESS(olMemUnregister(Device, PinnedPtr));
+}
+
+TEST_P(olMemRegisterTest, SuccessMultipleRegister) {
+  int Arr[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  void *PinnedPtr1 = nullptr;
+  ASSERT_SUCCESS(olMemRegister(Device, Arr, sizeof(Arr), Flags, &PinnedPtr));
+  ASSERT_NE(PinnedPtr, nullptr);
+  ASSERT_SUCCESS(olMemRegister(Device, Arr, sizeof(Arr), Flags, &PinnedPtr1));
+  ASSERT_NE(PinnedPtr1, nullptr);
+  ASSERT_SUCCESS(olMemUnregister(Device, PinnedPtr));
+  ASSERT_SUCCESS(olMemUnregister(Device, PinnedPtr1));
+}
+
+TEST_P(olMemRegisterTest, InvalidSizeRegister) {
+  int Arr[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  ASSERT_ERROR(OL_ERRC_INVALID_SIZE,
+               olMemRegister(Device, Arr, 0, Flags, &PinnedPtr));
+}
+
+TEST_P(olMemRegisterTest, InvalidPtrRegister) {
+  int Arr[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olMemRegister(Device, nullptr, sizeof(Arr), Flags, &PinnedPtr));
+}
+
+TEST_P(olMemRegisterTest, InvalidPtrUnRegister) {
+  int Arr[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  ASSERT_SUCCESS(olMemRegister(Device, Arr, sizeof(Arr), Flags, &PinnedPtr));
+  ASSERT_NE(PinnedPtr, nullptr);
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER, olMemUnregister(Device, nullptr));
+  ASSERT_SUCCESS(olMemUnregister(Device, PinnedPtr));
+}
+
+TEST_P(olMemRegisterTest, UnregisteredPtrUnRegister) {
+  int Arr[50];
+  int Arr1[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  ASSERT_SUCCESS(olMemRegister(Device, Arr, sizeof(Arr), Flags, &PinnedPtr));
+  ASSERT_NE(PinnedPtr, nullptr);
+  ASSERT_ERROR(OL_ERRC_INVALID_ARGUMENT, olMemUnregister(Device, Arr1));
+  ASSERT_SUCCESS(olMemUnregister(Device, PinnedPtr));
+}
+
+TEST_P(olMemRegisterTest, PartialOverlapPtrRegister) {
+  int Arr[50];
+  ol_memory_register_flags_t Flags = {};
+  void *PinnedPtr = nullptr;
+  ASSERT_SUCCESS(olMemRegister(Device, Arr, sizeof(Arr), Flags, &PinnedPtr));
+  ASSERT_NE(PinnedPtr, nullptr);
+  ASSERT_ERROR(OL_ERRC_INVALID_ARGUMENT,
+               olMemRegister(Device, Arr + 2, sizeof(Arr), Flags, &PinnedPtr));
+  ASSERT_SUCCESS(olMemUnregister(Device, PinnedPtr));
+}