Use ScopedCommandList to get thread local CL

Same as the CUDA implementation. This means that any CommandList
obtained through urQueueGetNativeHandle will be the same CommmandList
that is synchronized with before the interop func call.

Author: Hugh Delaney

source/adapters/level_zero/device.cpp

@@ -877,7 +877,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
   }
   case UR_DEVICE_INFO_ENQUEUE_NATIVE_COMMAND_SUPPORT_EXP: {
     // L0 doesn't support enqueueing native work through the urNativeEnqueueExp
-    return ReturnValue(static_cast<ur_bool_t>(false));
+    return ReturnValue(static_cast<ur_bool_t>(true));
   }
 
   case UR_DEVICE_INFO_ESIMD_SUPPORT: {

source/adapters/level_zero/enqueue_native.cpp

@@ -43,6 +43,7 @@ ur_result_t ur_queue_handle_legacy_t_::enqueueNativeCommandExp(
   // support is added
   UR_CALL(Queue->Context->getAvailableCommandList(
       Queue, CommandList, UseCopyEngine, NumEventsInWaitList, phEventList));
+  ScopedCommandList Active{Queue, CommandList->first};
 
   // TODO: do we need to create a unique command type for this?
   ze_event_handle_t ZeEvent = nullptr;

source/adapters/level_zero/queue.cpp

@@ -705,6 +705,15 @@ ur_result_t ur_queue_handle_legacy_t_::queueGetNativeHandle(
 ) {
   auto Queue = this;
 
+  // Needed for EnqueueNativeCommandExp, so that the native queue 'got' in the
+  // interop func is the as the native queue used to manage dependencies
+  // before the interop func invocation
+  if (Queue->getThreadLocalCommandList() != ze_command_list_handle_t{0}) {
+    auto ZeCmdList = ur_cast<ze_command_list_handle_t *>(NativeQueue);
+    *ZeCmdList = Queue->getThreadLocalCommandList();
+    return UR_RESULT_SUCCESS;
+  }
+
   // Lock automatically releases when this goes out of scope.
   std::shared_lock<ur_shared_mutex> lock(Queue->Mutex);
 

source/adapters/level_zero/queue.hpp

@@ -423,6 +423,12 @@ struct ur_queue_handle_legacy_t_ : _ur_object, public ur_queue_handle_t_ {
                           uint32_t, const ur_event_handle_t *,
                           ur_event_handle_t *) override;
 
+  // Thread local stream will be used if ScopedStream is active
+  static ze_command_list_handle_t &getThreadLocalCommandList() {
+    static thread_local ze_command_list_handle_t CommandList{0};
+    return CommandList;
+  }
+
   using queue_type = ur_device_handle_t_::queue_group_info_t::type;
   // PI queue is in general a one to many mapping to L0 native queues.
   struct ur_queue_group_t {
@@ -941,3 +947,23 @@ ur_result_t setSignalEvent(ur_queue_handle_legacy_t Queue, bool UseCopyEngine,
 ur_result_t CleanupEventListFromResetCmdList(
     std::vector<ur_event_handle_t> &EventListToCleanup,
     bool QueueLocked = false);
+
+// RAII object to make hQueue command list getter methods all return the same
+// command list within the lifetime of this object.
+//
+// This is useful for urEnqueueNativeCommandExp where we want guarantees that
+// the user submitted native calls will be dispatched to a known command list,
+// which must be "got" within the user submitted fuction.
+class ScopedCommandList {
+  ur_queue_handle_legacy_t hQueue;
+
+public:
+  ScopedCommandList(ur_queue_handle_legacy_t hQueue,
+                    ze_command_list_handle_t CommandList)
+      : hQueue{hQueue} {
+    hQueue->getThreadLocalCommandList() = CommandList;
+  }
+  ~ScopedCommandList() {
+    hQueue->getThreadLocalCommandList() = ze_command_list_handle_t{0};
+  }
+};