[SYCL][Graph] Add unit test/mocking for command handle behaviour (#15806)

- Add a unit test and mocking functions that check the command handle
behaviour changes from #15522

---------

Co-authored-by: Ewan Crawford <[email protected]>

Author: Bensuo

sycl/unittests/Extensions/CommandGraph/Update.cpp

@@ -399,3 +399,38 @@ TEST_F(WholeGraphUpdateTest, EmptyNode) {
   auto GraphExec = Graph.finalize(experimental::property::graph::updatable{});
   GraphExec.update(UpdateGraph);
 }
+
+// Vars and callbacks for tracking how many times mocked functions are called
+static int GetInfoCount = 0;
+static int AppendKernelLaunchCount = 0;
+static ur_result_t redefinedCommandBufferGetInfoExpAfter(void *pParams) {
+  GetInfoCount++;
+  return UR_RESULT_SUCCESS;
+}
+static ur_result_t
+redefinedCommandBufferAppendKernelLaunchExpAfter(void *pParams) {
+  AppendKernelLaunchCount++;
+  return UR_RESULT_SUCCESS;
+}
+
+TEST_F(CommandGraphTest, CheckFinalizeBehavior) {
+  // Check that both finalize with and without updatable property work as
+  // expected
+  auto Node = Graph.add(
+      [&](sycl::handler &cgh) { cgh.single_task<TestKernel<>>([]() {}); });
+  mock::getCallbacks().set_after_callback(
+      "urCommandBufferGetInfoExp", &redefinedCommandBufferGetInfoExpAfter);
+  mock::getCallbacks().set_after_callback(
+      "urCommandBufferAppendKernelLaunchExp",
+      &redefinedCommandBufferAppendKernelLaunchExpAfter);
+
+  ASSERT_NO_THROW(Graph.finalize(experimental::property::graph::updatable{}));
+  // GetInfo and AppendKernelLaunch should be called once each time a node is
+  // added to a command buffer during finalization
+  ASSERT_EQ(GetInfoCount, 1);
+  ASSERT_EQ(AppendKernelLaunchCount, 1);
+
+  ASSERT_NO_THROW(Graph.finalize());
+  ASSERT_EQ(GetInfoCount, 2);
+  ASSERT_EQ(AppendKernelLaunchCount, 2);
+}

sycl/unittests/helpers/UrMock.hpp

@@ -439,6 +439,67 @@ inline ur_result_t mock_urVirtualMemReserve(void *pParams) {
   return UR_RESULT_SUCCESS;
 }
 
+// Create dummy command buffer handle and store the provided descriptor as the
+// data
+inline ur_result_t mock_urCommandBufferCreateExp(void *pParams) {
+  auto params =
+      reinterpret_cast<ur_command_buffer_create_exp_params_t *>(pParams);
+  const ur_exp_command_buffer_desc_t *descPtr = *(params->ppCommandBufferDesc);
+  ur_exp_command_buffer_handle_t *retCmdBuffer = *params->pphCommandBuffer;
+  *retCmdBuffer = mock::createDummyHandle<ur_exp_command_buffer_handle_t>(
+      static_cast<size_t>(sizeof(ur_exp_command_buffer_desc_t)));
+  if (descPtr) {
+    reinterpret_cast<mock::dummy_handle_t>(*retCmdBuffer)
+        ->setDataAs<ur_exp_command_buffer_desc_t>(*descPtr);
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
+inline ur_result_t mock_urCommandBufferGetInfoExp(void *pParams) {
+  auto params =
+      reinterpret_cast<ur_command_buffer_get_info_exp_params_t *>(pParams);
+
+  auto cmdBufferDummyHandle =
+      reinterpret_cast<mock::dummy_handle_t>(*params->phCommandBuffer);
+  switch (*params->ppropName) {
+  case UR_EXP_COMMAND_BUFFER_INFO_DESCRIPTOR: {
+    if (*params->ppPropValue) {
+      ur_exp_command_buffer_desc_t *propValue =
+          reinterpret_cast<ur_exp_command_buffer_desc_t *>(
+              *params->ppPropValue);
+      *propValue =
+          cmdBufferDummyHandle->getDataAs<ur_exp_command_buffer_desc_t>();
+    }
+    if (*params->ppPropSizeRet)
+      **params->ppPropSizeRet = sizeof(ur_exp_command_buffer_desc_t);
+  }
+    return UR_RESULT_SUCCESS;
+  default:
+    return UR_RESULT_SUCCESS;
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+// Checking command handle behaviour only
+inline ur_result_t mock_urCommandBufferAppendKernelLaunchExp(void *pParams) {
+  auto params =
+      reinterpret_cast<ur_command_buffer_append_kernel_launch_exp_params_t *>(
+          pParams);
+
+  auto cmdBufferDummyHandle =
+      reinterpret_cast<mock::dummy_handle_t>(*params->phCommandBuffer);
+  // Requesting a command handle when the command buffer is not updatable is an
+  // error
+  if (*(params->pphCommand) &&
+      cmdBufferDummyHandle->getDataAs<ur_exp_command_buffer_desc_t>()
+              .isUpdatable == false) {
+    return UR_RESULT_ERROR_INVALID_OPERATION;
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
 } // namespace MockAdapter
 
 /// The UrMock<> class sets up UR for adapter mocking with the set of default
@@ -484,6 +545,12 @@ template <sycl::backend Backend = backend::opencl> class UrMock {
     ADD_DEFAULT_OVERRIDE(urUsmP2PPeerAccessGetInfoExp,
                          mock_urUsmP2PPeerAccessGetInfoExp)
     ADD_DEFAULT_OVERRIDE(urVirtualMemReserve, mock_urVirtualMemReserve)
+    ADD_DEFAULT_OVERRIDE(urCommandBufferCreateExp,
+                         mock_urCommandBufferCreateExp);
+    ADD_DEFAULT_OVERRIDE(urCommandBufferAppendKernelLaunchExp,
+                         mock_urCommandBufferAppendKernelLaunchExp);
+    ADD_DEFAULT_OVERRIDE(urCommandBufferGetInfoExp,
+                         mock_urCommandBufferGetInfoExp);
 #undef ADD_DEFAULT_OVERRIDE
 
     ur_loader_config_handle_t UrLoaderConfig = nullptr;