[SYCL] Reset PiFunctionTable in PiMock::~PiMock (#6128)

Author: aelovikov-intel

sycl/unittests/helpers/PiMock.hpp

@@ -35,6 +35,7 @@
 #include <detail/platform_impl.hpp>
 
 #include <functional>
+#include <optional>
 
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
@@ -133,14 +134,26 @@ class PiMock {
     MPiPluginMockPtr = &NewPluginPtr->getPiPlugin();
     // Save a copy of the platform resource
     MPlatform = OriginalPlatform;
+    OrigFuncTable = OriginalPiPlugin.getPiPlugin().PiFunctionTable;
   }
 
   /// Explicit construction from a host_selector is forbidden.
   PiMock(const cl::sycl::host_selector &HostSelector) = delete;
 
+  PiMock(PiMock &&Other) {
+    MPlatform = std::move(Other.MPlatform);
+    OrigFuncTable = std::move(Other.OrigFuncTable);
+    Other.OrigFuncTable = {}; // Move above doesn't reset the optional.
+    MPiPluginMockPtr = std::move(Other.MPiPluginMockPtr);
+  }
   PiMock(const PiMock &) = delete;
   PiMock &operator=(const PiMock &) = delete;
-  ~PiMock() = default;
+  ~PiMock() {
+    if (!OrigFuncTable)
+      return;
+
+    MPiPluginMockPtr->PiFunctionTable = *OrigFuncTable;
+  }
 
   /// Returns a handle to the SYCL platform instance.
   ///
@@ -184,6 +197,7 @@ class PiMock {
 
 private:
   cl::sycl::platform MPlatform;
+  std::optional<pi_plugin::FunctionPointers> OrigFuncTable;
   // Extracted at initialization for convenience purposes. The resource
   // itself is owned by the platform instance.
   RT::PiPlugin *MPiPluginMockPtr;

sycl/unittests/queue/EventClear.cpp

@@ -82,11 +82,11 @@ pi_result redefinedEventRelease(pi_event event) {
   return PI_SUCCESS;
 }
 
-bool preparePiMock(platform &Plt) {
+std::optional<unittest::PiMock> preparePiMock(platform &Plt) {
   if (Plt.is_host()) {
     std::cout << "Not run on host - no PI events created in that case"
               << std::endl;
-    return false;
+    return {};
   }
 
   unittest::PiMock Mock{Plt};
@@ -98,14 +98,15 @@ bool preparePiMock(platform &Plt) {
   Mock.redefine<detail::PiApiKind::piEventGetInfo>(redefinedEventGetInfo);
   Mock.redefine<detail::PiApiKind::piEventRetain>(redefinedEventRetain);
   Mock.redefine<detail::PiApiKind::piEventRelease>(redefinedEventRelease);
-  return true;
+  return std::move(Mock);
 }
 
 // Check that the USM events are cleared from the queue upon call to wait(),
 // so that they are not waited for multiple times.
 TEST(QueueEventClear, ClearOnQueueWait) {
   platform Plt{default_selector()};
-  if (!preparePiMock(Plt))
+  auto Mock = preparePiMock(Plt);
+  if (!Mock)
     return;
 
   context Ctx{Plt.get_devices()[0]};
@@ -126,7 +127,8 @@ TEST(QueueEventClear, ClearOnQueueWait) {
 // exceeds a threshold.
 TEST(QueueEventClear, CleanupOnThreshold) {
   platform Plt{default_selector()};
-  if (!preparePiMock(Plt))
+  auto Mock = preparePiMock(Plt);
+  if (!Mock)
     return;
 
   context Ctx{Plt.get_devices()[0]};

sycl/unittests/queue/GetProfilingInfo.cpp

@@ -17,6 +17,8 @@
 #include <helpers/PiImage.hpp>
 #include <helpers/PiMock.hpp>
 
+#include <detail/context_impl.hpp>
+
 class InfoTestKernel;
 
 __SYCL_INLINE_NAMESPACE(cl) {
@@ -85,14 +87,14 @@ TEST(GetProfilingInfo, normal_pass_without_exception) {
   Mock.redefine<sycl::detail::PiApiKind::piEventGetProfilingInfo>(
       redefinedPiEventGetProfilingInfo);
   const sycl::device Dev = Plt.get_devices()[0];
+  sycl::context Ctx{Dev};
   static sycl::unittest::PiImage DevImage_1 =
       generateTestImage<InfoTestKernel>();
 
   static sycl::unittest::PiImageArray<1> DevImageArray = {&DevImage_1};
   auto KernelID_1 = sycl::get_kernel_id<InfoTestKernel>();
   sycl::queue Queue{
-      Dev, sycl::property_list{sycl::property::queue::enable_profiling{}}};
-  const sycl::context Ctx = Queue.get_context();
+      Ctx, Dev, sycl::property_list{sycl::property::queue::enable_profiling{}}};
   auto KernelBundle = sycl::get_kernel_bundle<sycl::bundle_state::input>(
       Ctx, {Dev}, {KernelID_1});
 
@@ -139,13 +141,13 @@ TEST(GetProfilingInfo, command_exception_check) {
       redefinedPiEventGetProfilingInfo);
 
   const sycl::device Dev = Plt.get_devices()[0];
+  sycl::context Ctx{Dev};
   static sycl::unittest::PiImage DevImage_1 =
       generateTestImage<InfoTestKernel>();
 
   static sycl::unittest::PiImageArray<1> DevImageArray = {&DevImage_1};
   auto KernelID_1 = sycl::get_kernel_id<InfoTestKernel>();
-  sycl::queue Queue{Dev};
-  const sycl::context Ctx = Queue.get_context();
+  sycl::queue Queue{Ctx, Dev};
   auto KernelBundle = sycl::get_kernel_bundle<sycl::bundle_state::input>(
       Ctx, {Dev}, {KernelID_1});
   const int globalWIs{512};
@@ -219,6 +221,7 @@ TEST(GetProfilingInfo, check_if_now_dead_queue_property_set) {
   Mock.redefine<sycl::detail::PiApiKind::piEventGetProfilingInfo>(
       redefinedPiEventGetProfilingInfo);
   const sycl::device Dev = Plt.get_devices()[0];
+  sycl::context Ctx{Dev};
   static sycl::unittest::PiImage DevImage_1 =
       generateTestImage<InfoTestKernel>();
 
@@ -228,8 +231,8 @@ TEST(GetProfilingInfo, check_if_now_dead_queue_property_set) {
   cl::sycl::event event;
   {
     sycl::queue Queue{
-        Dev, sycl::property_list{sycl::property::queue::enable_profiling{}}};
-    const sycl::context Ctx = Queue.get_context();
+        Ctx, Dev,
+        sycl::property_list{sycl::property::queue::enable_profiling{}}};
     auto KernelBundle = sycl::get_kernel_bundle<sycl::bundle_state::input>(
         Ctx, {Dev}, {KernelID_1});
     event = Queue.submit([&](sycl::handler &cgh) {
@@ -274,6 +277,7 @@ TEST(GetProfilingInfo, check_if_now_dead_queue_property_not_set) {
   Mock.redefine<sycl::detail::PiApiKind::piEventGetProfilingInfo>(
       redefinedPiEventGetProfilingInfo);
   const sycl::device Dev = Plt.get_devices()[0];
+  sycl::context Ctx{Dev};
   static sycl::unittest::PiImage DevImage_1 =
       generateTestImage<InfoTestKernel>();
 
@@ -282,8 +286,7 @@ TEST(GetProfilingInfo, check_if_now_dead_queue_property_not_set) {
   const int globalWIs{512};
   cl::sycl::event event;
   {
-    sycl::queue Queue{Dev};
-    const sycl::context Ctx = Queue.get_context();
+    sycl::queue Queue{Ctx, Dev};
     auto KernelBundle = sycl::get_kernel_bundle<sycl::bundle_state::input>(
         Ctx, {Dev}, {KernelID_1});
     event = Queue.submit([&](sycl::handler &cgh) {
@@ -325,4 +328,6 @@ TEST(GetProfilingInfo, check_if_now_dead_queue_property_not_set) {
                              "'enable_profiling' queue property");
     }
   }
+  // The test passes without this, but keep it still, just in case.
+  sycl::detail::getSyclObjImpl(Ctx)->getKernelProgramCache().reset();
 }

sycl/unittests/queue/USM.cpp

@@ -54,12 +54,18 @@ pi_result redefinedUSMEnqueueMemset(pi_queue, void *, pi_int32, size_t,
 }
 
 pi_result redefinedEventRelease(pi_event) { return PI_SUCCESS; }
+pi_result redefinedEventsWait(pi_uint32 /* num_events */,
+                              const pi_event * /* event_list  */) {
+  return PI_SUCCESS;
+}
 
-bool preparePiMock(platform &Plt) {
+// Check that zero-length USM memset/memcpy use piEnqueueEventsWait.
+TEST(USM, NoOpPreservesDependencyChain) {
+  platform Plt{default_selector()};
   if (Plt.is_host()) {
     std::cout << "Not run on host - no PI events created in that case"
               << std::endl;
-    return false;
+    return;
   }
 
   unittest::PiMock Mock{Plt};
@@ -70,14 +76,7 @@ bool preparePiMock(platform &Plt) {
   Mock.redefine<detail::PiApiKind::piextUSMEnqueueMemset>(
       redefinedUSMEnqueueMemset);
   Mock.redefine<detail::PiApiKind::piEventRelease>(redefinedEventRelease);
-  return true;
-}
-
-// Check that zero-length USM memset/memcpy use piEnqueueEventsWait.
-TEST(USM, NoOpPreservesDependencyChain) {
-  platform Plt{default_selector()};
-  if (!preparePiMock(Plt))
-    return;
+  Mock.redefine<detail::PiApiKind::piEventsWait>(redefinedEventsWait);
 
   context Ctx{Plt.get_devices()[0]};
   queue Q{Ctx, default_selector()};
@@ -102,6 +101,6 @@ TEST(USM, NoOpPreservesDependencyChain) {
 
   free(Src, Q);
   free(Dst, Q);
+  TestContext.Deps.clear();
 }
-
 } // namespace

sycl/unittests/queue/Wait.cpp

@@ -86,11 +86,12 @@ pi_result redefinedEventRelease(pi_event event) {
   return PI_SUCCESS;
 }
 
-bool preparePiMock(platform &Plt) {
+TEST(QueueWait, QueueWaitTest) {
+  platform Plt{default_selector()};
   if (Plt.is_host()) {
     std::cout << "Not run on host - no PI events created in that case"
               << std::endl;
-    return false;
+    return;
   }
 
   unittest::PiMock Mock{Plt};
@@ -105,13 +106,6 @@ bool preparePiMock(platform &Plt) {
   Mock.redefine<detail::PiApiKind::piEventGetInfo>(redefinedEventGetInfo);
   Mock.redefine<detail::PiApiKind::piEventRetain>(redefinedEventRetain);
   Mock.redefine<detail::PiApiKind::piEventRelease>(redefinedEventRelease);
-  return true;
-}
-
-TEST(QueueWait, QueueWaitTest) {
-  platform Plt{default_selector()};
-  if (!preparePiMock(Plt))
-    return;
   context Ctx{Plt.get_devices()[0]};
   queue Q{Ctx, default_selector()};
 

sycl/unittests/scheduler/RequiredWGSize.cpp

@@ -238,6 +238,5 @@ TEST(RequiredWGSize, NoRequiredSize) {
 TEST(RequiredWGSize, HasRequiredSize) {
   reset();
   RequiredLocalSize = {1, 2, 3};
-  return; // FIXME: Resolve post-commit failures.
   performChecks();
 }