[L0 Adapter] Mutable command buffers

Author: againull

source/adapters/cuda/command_buffer.hpp

@@ -298,7 +298,7 @@ struct ur_exp_command_buffer_handle_t_ {
   // Device associated with this command buffer
   ur_device_handle_t Device;
   // Whether commands in the command-buffer can be updated
-  bool IsUpdatable;
+  bool IsUpdatable = false;
   // Cuda Graph handle
   CUgraph CudaGraph;
   // Cuda Graph Exec handle

source/adapters/level_zero/command_buffer.cpp

@@ -78,4 +78,21 @@ struct ur_exp_command_buffer_handle_t_ : public _ur_object {
   // Event which a command-buffer waits on until the main command-list event
   // have been reset.
   ur_event_handle_t AllResetEvent = nullptr;
+  // Indicates if command-buffer commands can be updated after it is closed.
+  bool IsUpdatable = false;
+  // Indicates if command buffer was finalized.
+  bool IsFinalized = false;
+};
+
+struct ur_exp_command_buffer_command_handle_t_ : public _ur_object {
+  ur_exp_command_buffer_command_handle_t_(ur_exp_command_buffer_handle_t,
+                                          uint64_t, ur_kernel_handle_t);
+
+  ~ur_exp_command_buffer_command_handle_t_();
+
+  // Command-buffer of this command.
+  ur_exp_command_buffer_handle_t CommandBuffer;
+
+  uint64_t CommandId;
+  ur_kernel_handle_t Kernel;
 };

source/adapters/level_zero/command_buffer.hpp

@@ -58,6 +58,8 @@ ur_result_t ze2urResult(ze_result_t ZeResult) {
     return UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY;
   case ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY:
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+  case ZE_RESULT_ERROR_UNSUPPORTED_FEATURE:
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   default:
     return UR_RESULT_ERROR_UNKNOWN;
   }
@@ -171,6 +173,35 @@ template <> ze_structure_type_t getZeStructureType<ze_fence_desc_t>() {
 template <> ze_structure_type_t getZeStructureType<ze_command_list_desc_t>() {
   return ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC;
 }
+template <>
+ze_structure_type_t getZeStructureType<ze_mutable_command_list_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_COMMAND_LIST_EXP_DESC;
+}
+template <>
+ze_structure_type_t getZeStructureType<ze_mutable_command_id_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_COMMAND_ID_EXP_DESC;
+}
+template <>
+ze_structure_type_t getZeStructureType<ze_mutable_group_count_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_GROUP_COUNT_EXP_DESC;
+}
+template <>
+ze_structure_type_t getZeStructureType<ze_mutable_group_size_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_GROUP_SIZE_EXP_DESC;
+}
+template <>
+ze_structure_type_t getZeStructureType<ze_mutable_global_offset_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_GLOBAL_OFFSET_EXP_DESC;
+}
+template <>
+ze_structure_type_t
+getZeStructureType<ze_mutable_kernel_argument_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_KERNEL_ARGUMENT_EXP_DESC;
+}
+template <>
+ze_structure_type_t getZeStructureType<ze_mutable_commands_exp_desc_t>() {
+  return ZE_STRUCTURE_TYPE_MUTABLE_COMMANDS_EXP_DESC;
+}
 template <> ze_structure_type_t getZeStructureType<ze_context_desc_t>() {
   return ZE_STRUCTURE_TYPE_CONTEXT_DESC;
 }

source/adapters/level_zero/common.cpp

@@ -918,7 +918,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
   case UR_DEVICE_INFO_COMMAND_BUFFER_SUPPORT_EXP:
     return ReturnValue(true);
   case UR_DEVICE_INFO_COMMAND_BUFFER_UPDATE_SUPPORT_EXP:
-    return ReturnValue(false);
+    return ReturnValue(Device->Platform->ZeMutableCmdListExt.Supported);
   case UR_DEVICE_INFO_BINDLESS_IMAGES_SUPPORT_EXP:
     return ReturnValue(true);
   case UR_DEVICE_INFO_BINDLESS_IMAGES_SHARED_USM_SUPPORT_EXP:

source/adapters/level_zero/device.cpp

@@ -206,6 +206,39 @@ ur_result_t ur_platform_handle_t_::initialize() {
   // If yes, then set up L0 API pointers if the platform supports it.
   ZeUSMImport.setZeUSMImport(this);
 
+  // Check if mutable command list extension is supported and initialize
+  // function pointers.
+  ZeMutableCmdListExt.Supported |=
+      (ZE_CALL_NOCHECK(
+           zeDriverGetExtensionFunctionAddress,
+           (ZeDriver, "zeCommandListGetNextCommandIdExp",
+            reinterpret_cast<void **>(
+                &ZeMutableCmdListExt.zexCommandListGetNextCommandIdExp))) == 0);
+
+  ZeMutableCmdListExt.Supported &=
+      (ZE_CALL_NOCHECK(zeDriverGetExtensionFunctionAddress,
+                       (ZeDriver, "zeCommandListUpdateMutableCommandsExp",
+                        reinterpret_cast<void **>(
+                            &ZeMutableCmdListExt
+                                 .zexCommandListUpdateMutableCommandsExp))) ==
+       0);
+
+  ZeMutableCmdListExt.Supported &=
+      (ZE_CALL_NOCHECK(
+           zeDriverGetExtensionFunctionAddress,
+           (ZeDriver, "zeCommandListUpdateMutableCommandSignalEventExp",
+            reinterpret_cast<void **>(
+                &ZeMutableCmdListExt
+                     .zexCommandListUpdateMutableCommandSignalEventExp))) == 0);
+
+  ZeMutableCmdListExt.Supported &=
+      (ZE_CALL_NOCHECK(
+           zeDriverGetExtensionFunctionAddress,
+           (ZeDriver, "zeCommandListUpdateMutableCommandWaitEventsExp",
+            reinterpret_cast<void **>(
+                &ZeMutableCmdListExt
+                     .zexCommandListUpdateMutableCommandWaitEventsExp))) == 0);
+
   return UR_RESULT_SUCCESS;
 }
 

source/adapters/level_zero/platform.cpp

@@ -55,4 +55,22 @@ struct ur_platform_handle_t_ : public _ur_platform {
   // in the driver.
   std::list<ur_context_handle_t> Contexts;
   ur_shared_mutex ContextsMutex;
+
+  // Structure with function pointers for mutable command list extension.
+  // Not all drivers may support it, so considering that the platform object is
+  // associated with particular Level Zero driver, store this extension here.
+  struct ZeMutableCmdListExtension {
+    bool Supported = false;
+    ze_result_t (*zexCommandListGetNextCommandIdExp)(
+        ze_command_list_handle_t, const ze_mutable_command_id_exp_desc_t *,
+        uint64_t *) = nullptr;
+    ze_result_t (*zexCommandListUpdateMutableCommandsExp)(
+        ze_command_list_handle_t,
+        const ze_mutable_commands_exp_desc_t *) = nullptr;
+    ze_result_t (*zexCommandListUpdateMutableCommandSignalEventExp)(
+        ze_command_list_handle_t, uint64_t, ze_event_handle_t) = nullptr;
+    ze_result_t (*zexCommandListUpdateMutableCommandWaitEventsExp)(
+        ze_command_list_handle_t, uint64_t, uint32_t,
+        ze_event_handle_t *) = nullptr;
+  } ZeMutableCmdListExt;
 };

source/adapters/level_zero/platform.hpp

@@ -153,7 +153,7 @@ TEST_P(BufferFillCommandTest, UpdateGlobalSize) {
         0,                // numNewPointerArgs
         0,                // numNewValueArgs
         0,                // numNewExecInfos
-        0,                // newWorkDim
+        1,                // newWorkDim
         &new_output_desc, // pNewMemObjArgList
         nullptr,          // pNewPointerArgList
         nullptr,          // pNewValueArgList
@@ -180,7 +180,8 @@ TEST_P(BufferFillCommandTest, SeparateUpdateCalls) {
     ASSERT_SUCCESS(urQueueFinish(queue));
     ValidateBuffer(buffer, sizeof(val) * global_size, val);
 
-    size_t new_global_size = 64;
+    size_t new_global_size =
+        global_size; //64; // Try same value for testing purposes.
     const size_t new_buffer_size = sizeof(val) * new_global_size;
     ASSERT_SUCCESS(urMemBufferCreate(context, UR_MEM_FLAG_READ_WRITE,
                                      new_buffer_size, nullptr, &new_buffer));
@@ -247,7 +248,7 @@ TEST_P(BufferFillCommandTest, SeparateUpdateCalls) {
     ASSERT_SUCCESS(urCommandBufferUpdateKernelLaunchExp(command_handle,
                                                         &input_update_desc));
 
-    ur_exp_command_buffer_update_kernel_launch_desc_t global_size_update_desc = {
+    /*ur_exp_command_buffer_update_kernel_launch_desc_t global_size_update_desc = {
         UR_STRUCTURE_TYPE_EXP_COMMAND_BUFFER_UPDATE_KERNEL_LAUNCH_DESC, // stype
         nullptr,                                                        // pNext
         0,                // numNewMemObjArgs
@@ -265,7 +266,7 @@ TEST_P(BufferFillCommandTest, SeparateUpdateCalls) {
     };
 
     ASSERT_SUCCESS(urCommandBufferUpdateKernelLaunchExp(
-        command_handle, &global_size_update_desc));
+        command_handle, &global_size_update_desc));*/
 
     ASSERT_SUCCESS(urCommandBufferEnqueueExp(updatable_cmd_buf_handle, queue, 0,
                                              nullptr, nullptr));

test/conformance/exp_command_buffer/buffer_fill_kernel_update.cpp

@@ -88,7 +88,7 @@ TEST_P(USMFillCommandTest, UpdateParameters) {
     Validate((uint32_t *)shared_ptr, global_size, val);
 
     // Allocate a new USM pointer of larger size
-    size_t new_global_size = 64;
+    size_t new_global_size = global_size; // 64;
     const size_t new_allocation_size = sizeof(val) * new_global_size;
     ASSERT_SUCCESS(urUSMSharedAlloc(context, device, nullptr, nullptr,
                                     new_allocation_size, &new_shared_ptr));
@@ -128,7 +128,7 @@ TEST_P(USMFillCommandTest, UpdateParameters) {
         &new_input_desc,  // pNewValueArgList
         nullptr,          // pNewExecInfoList
         nullptr,          // pNewGlobalWorkOffset
-        &new_global_size, // pNewGlobalWorkSize
+        nullptr,          //&new_global_size, // pNewGlobalWorkSize
         nullptr,          // pNewLocalWorkSize
     };