[llvm] [Offload] Fix dataDelete op for TARGET_ALLOC_HOST memory type (PR #91134)

[Jhonatan Cléto via llvm-commits]

https://github.com/cl3to created https://github.com/llvm/llvm-project/pull/91134

Summary:
The `GenericDeviceTy::dataDelete` method doesn't verify the `TargetAllocTy` of the of the device pointer. Because of this, it can use the `MemoryManager` to free the ptr. However, the `TARGET_ALLOC_HOST` and `TARGET_ALLOC_SHARED` types are not allocated using the `MemoryManager` in the `GenericDeviceTy::dataAlloc` method. Since the `MemoryManager` uses the `DeviceAllocatorTy::free` operation without specifying the type of the ptr, some plugins may use incorrect operations to free ptrs of certain types. In particular, this bug causes the CUDA plugin to use the `cuMemFree` operation on ptrs of type `TARGET_ALLOC_HOST`, resulting in an unchecked error, as shown in the output snippet of the test `offload/test/api/omp_host_pinned_memory_alloc.c`:

```
omptarget --> Notifying about an unmapping: HstPtr=0x00007c6114200000
omptarget --> Call to llvm_omp_target_free_host for device 0 and address 0x00007c6114200000
omptarget --> Call to omp_get_num_devices returning 1
omptarget --> Call to omp_get_initial_device returning 1
PluginInterface --> MemoryManagerTy::free: target memory 0x00007c6114200000.
PluginInterface --> Cannot find its node. Delete it on device directly.
TARGET CUDA RTL --> Failure to free memory: Error in cuMemFree[Host]: invalid argument
omptarget --> omp_target_free deallocated device ptr
```

This patch fixes this by adding the check of the device pointer type before calling the appropriate operation for each type.

>From da8cddbe753154078b8f3e77e53a0694b2ef538e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jhonatan=20Cl=C3=A9to?= <j256444 at dac.unicamp.br>
Date: Sun, 5 May 2024 12:14:14 -0300
Subject: [PATCH] [Offload] Fix dataDelete op for TARGET_ALLOC_HOST memory type

---
 .../common/src/PluginInterface.cpp            | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index b5f3c45c835fdb..8de93ba17a560b 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -1348,13 +1348,27 @@ Error GenericDeviceTy::dataDelete(void *TgtPtr, TargetAllocTy Kind) {
     return Plugin::success();
 
   int Res;
-  if (MemoryManager)
-    Res = MemoryManager->free(TgtPtr);
-  else
+  switch (Kind) {
+  case TARGET_ALLOC_DEFAULT:
+  case TARGET_ALLOC_DEVICE_NON_BLOCKING:
+  case TARGET_ALLOC_DEVICE:
+    if (MemoryManager) {
+      Res = MemoryManager->free(TgtPtr);
+      if (Res)
+        return Plugin::error(
+            "Failure to deallocate device pointer %p via memory manager",
+            TgtPtr);
+      break;
+    }
+    [[fallthrough]];
+  case TARGET_ALLOC_HOST:
+  case TARGET_ALLOC_SHARED:
     Res = free(TgtPtr, Kind);
-
-  if (Res)
-    return Plugin::error("Failure to deallocate device pointer %p", TgtPtr);
+    if (Res)
+      return Plugin::error(
+          "Failure to deallocate device pointer %p via device deallocator",
+          TgtPtr);
+  }
 
   // Unregister deallocated pinned memory buffer if the type is host memory.
   if (Kind == TARGET_ALLOC_HOST)