[Unittests] Fix device info query for shared USM (#13801)

The unittests have a default mock implementation of device info queries
which returns success for any query made. However, for cases where the
query does not have a case in the implementation, nothing is done to the
return values. For support queries this may mean that one call to the
query returns true and another returns false. This has been seen for
shared USM allocation support, which caused flaky failures when
unittests would call malloc_shared.

This commit:
* Changes the default implementation of the device info query to return
true for shared USM allocation support, like it already does for host
and device variants.
* Changes the default implementation of the device info query to fill
the return value with 0 and set the return type size to 1 when the info
query is unrecognized. This means unittests should always use
`redefineAfter` for the device info query.
* Fixes tests that assumes device info queries may be device-dependent
in unittests.

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/unittests/Extensions/CommandGraph/InOrderQueue.cpp

@@ -419,9 +419,12 @@ TEST_F(CommandGraphTest, InOrderQueueMemsetAndGraph) {
   experimental::command_graph<experimental::graph_state::modifiable>
       InOrderGraph{InOrderQueue.get_context(), InOrderQueue.get_device()};
 
-  // Check if device has usm shared allocation.
-  if (!InOrderQueue.get_device().has(sycl::aspect::usm_shared_allocations))
-    return;
+  // The mock plugin should return true for shared USM allocation support by
+  // default. If this fails it means this test needs to redefine the device info
+  // query.
+  ASSERT_TRUE(
+      InOrderQueue.get_device().has(sycl::aspect::usm_shared_allocations));
+
   size_t Size = 128;
   std::vector<int> TestDataHost(Size);
   int *TestData = sycl::malloc_shared<int>(Size, InOrderQueue);
@@ -494,9 +497,12 @@ TEST_F(CommandGraphTest, InOrderQueueMemcpyAndGraph) {
   experimental::command_graph<experimental::graph_state::modifiable>
       InOrderGraph{InOrderQueue.get_context(), InOrderQueue.get_device()};
 
-  // Check if device has usm shared allocation.
-  if (!InOrderQueue.get_device().has(sycl::aspect::usm_shared_allocations))
-    return;
+  // The mock plugin should return true for shared USM allocation support by
+  // default. If this fails it means this test needs to redefine the device info
+  // query.
+  ASSERT_TRUE(
+      InOrderQueue.get_device().has(sycl::aspect::usm_shared_allocations));
+
   size_t Size = 128;
   std::vector<int> TestDataHost(Size);
   int *TestData = sycl::malloc_shared<int>(Size, InOrderQueue);

sycl/unittests/helpers/PiMockPlugin.hpp

@@ -204,6 +204,7 @@ inline pi_result mock_piDeviceGetInfo(pi_device device,
   }
   case PI_DEVICE_INFO_USM_HOST_SUPPORT:
   case PI_DEVICE_INFO_USM_DEVICE_SUPPORT:
+  case PI_DEVICE_INFO_USM_SINGLE_SHARED_SUPPORT:
   case PI_DEVICE_INFO_HOST_UNIFIED_MEMORY:
   case PI_DEVICE_INFO_AVAILABLE:
   case PI_DEVICE_INFO_LINKER_AVAILABLE:
@@ -238,9 +239,23 @@ inline pi_result mock_piDeviceGetInfo(pi_device device,
     }
     return PI_SUCCESS;
   }
-  default:
+  default: {
+    // In the default case we fill the return value with 0's. This may not be
+    // valid for all device queries, but it will mean a consistent return value
+    // for the query.
+    // Any tests that need special return values should either add behavior
+    // the this function or use redefineAfter with a function that adds the
+    // intended behavior.
+    if (param_value && param_value_size != 0)
+      std::memset(param_value, 0, param_value_size);
+    // Likewise, if the device info query asks for the size of the return value
+    // we tell it there is a single byte to avoid cases where the runtime tries
+    // to allocate some random amount of memory for the return value.
+    if (param_value_size_ret)
+      *param_value_size_ret = 1;
     return PI_SUCCESS;
   }
+  }
 }
 
 inline pi_result mock_piDeviceRetain(pi_device device) { return PI_SUCCESS; }

sycl/unittests/kernel-and-program/DeviceInfo.cpp

@@ -83,7 +83,7 @@ class DeviceInfoTest : public ::testing::Test {
 
 protected:
   void SetUp() override {
-    Mock.redefineBefore<detail::PiApiKind::piDeviceGetInfo>(
+    Mock.redefineAfter<detail::PiApiKind::piDeviceGetInfo>(
         redefinedDeviceGetInfo);
   }
 
@@ -131,11 +131,7 @@ TEST_F(DeviceInfoTest, GetDeviceUUID) {
 
   device Dev = Ctx.get_devices()[0];
 
-  if (!Dev.has(aspect::ext_intel_device_info_uuid)) {
-    std::clog
-        << "This test is only for the devices with UUID extension support.\n";
-    return;
-  }
+  EXPECT_TRUE(Dev.has(aspect::ext_intel_device_info_uuid));
 
   auto UUID = Dev.get_info<ext::intel::info::device::uuid>();
 
@@ -154,11 +150,7 @@ TEST_F(DeviceInfoTest, GetDeviceFreeMemory) {
 
   device Dev = Ctx.get_devices()[0];
 
-  if (!Dev.has(aspect::ext_intel_free_memory)) {
-    std::clog << "This test is only for the devices with "
-                 "ext_intel_free_memory extension support.\n";
-    return;
-  }
+  EXPECT_TRUE(Dev.has(aspect::ext_intel_free_memory));
 
   auto FreeMemory = Dev.get_info<ext::intel::info::device::free_memory>();