Merge pull request #1397 from GeorgeWeb/georgi/check-allocation-error…

…-on-event-from-native-handle [CUDA][HIP] Catch and report bad_alloc errors for event object creation
oneapi-src · May 31, 2024 · 5083f4f · 5083f4f
2 parents a97eed1 + 4c3f9ab
commit 5083f4f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/source/adapters/cuda/event.cpp b/source/adapters/cuda/event.cpp
@@ -280,8 +280,17 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventCreateWithNativeHandle(
 
   std::unique_ptr<ur_event_handle_t_> EventPtr{nullptr};
 
-  *phEvent = ur_event_handle_t_::makeWithNative(
-      hContext, reinterpret_cast<CUevent>(hNativeEvent));
+  try {
+    EventPtr =
+        std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeWithNative(
+            hContext, reinterpret_cast<CUevent>(hNativeEvent)));
+  } catch (const std::bad_alloc &) {
+    return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+  } catch (...) {
+    return UR_RESULT_ERROR_UNKNOWN;
+  }
+
+  *phEvent = EventPtr.release();
 
   return UR_RESULT_SUCCESS;
 }
diff --git a/source/adapters/hip/event.cpp b/source/adapters/hip/event.cpp
@@ -327,8 +327,19 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventCreateWithNativeHandle(
     ur_event_handle_t *phEvent) {
   std::ignore = pProperties;
 
-  *phEvent = ur_event_handle_t_::makeWithNative(
-      hContext, reinterpret_cast<hipEvent_t>(hNativeEvent));
+  std::unique_ptr<ur_event_handle_t_> EventPtr{nullptr};
+
+  try {
+    EventPtr =
+        std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeWithNative(
+            hContext, reinterpret_cast<hipEvent_t>(hNativeEvent)));
+  } catch (const std::bad_alloc &) {
+    return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+  } catch (...) {
+    return UR_RESULT_ERROR_UNKNOWN;
+  }
+
+  *phEvent = EventPtr.release();
 
   return UR_RESULT_SUCCESS;
 }