test: shrd sys alloc tests for *MemoryPrefetch (#215)

Related to VLCLJ-2515

Signed-off-by: Bagria, Narendra <[email protected]>

Author: narenbagria

conformance_tests/core/test_cmdlist/src/test_cmdlist_immediate.cpp

@@ -971,24 +971,49 @@ TEST_P(zeImmediateCommandListExecutionTests,
   lzt::free_memory(memory);
 }
 
+TEST_P(
+    zeImmediateCommandListExecutionTests,
+    GivenImmediateCommandListWhenAppendMemoryPrefetchThenSuccessIsReturnedWithSharedSystemAllocator) {
+  SKIP_IF_SHARED_SYSTEM_ALLOC_UNSUPPORTED();
+  size_t size = 4096;
+  void *memory = lzt::aligned_malloc(size, 1);
+  EXPECT_EQ(ZE_RESULT_SUCCESS,
+            zeCommandListAppendMemoryPrefetch(cmdlist_immediate, memory, size));
+  EXPECT_EQ(ZE_RESULT_SUCCESS,
+            zeCommandListHostSynchronize(cmdlist_immediate, timeout));
+  lzt::aligned_free(memory);
+}
+
 static void RunAppendMemoryPrefetch(ze_command_list_handle_t cmdlist_immediate,
-                                    const uint64_t timeout) {
+                                    const uint64_t timeout,
+                                    bool is_shared_system) {
   size_t size = 4096;
-  void *memory = lzt::allocate_shared_memory(size);
+  void *memory = lzt::allocate_shared_memory_with_allocator_selector(
+      size, is_shared_system);
   EXPECT_EQ(ZE_RESULT_SUCCESS,
             zeCommandListAppendMemoryPrefetch(cmdlist_immediate, memory, size));
   EXPECT_EQ(ZE_RESULT_SUCCESS,
             zeCommandListHostSynchronize(cmdlist_immediate, timeout));
-  lzt::free_memory(memory);
+  lzt::free_memory_with_allocator_selector(memory, is_shared_system);
 }
 
 TEST_F(
     zeImmediateCommandListInOrderExecutionTests,
     GivenInOrderImmediateCommandListWhenAppendMemoryPrefetchThenSuccessIsReturned) {
   const uint64_t timeout = UINT64_MAX - 1;
-  RunAppendMemoryPrefetch(cmdlist_immediate_default_mode, timeout);
-  RunAppendMemoryPrefetch(cmdlist_immediate_sync_mode, 0);
-  RunAppendMemoryPrefetch(cmdlist_immediate_async_mode, timeout);
+  RunAppendMemoryPrefetch(cmdlist_immediate_default_mode, timeout, false);
+  RunAppendMemoryPrefetch(cmdlist_immediate_sync_mode, 0, false);
+  RunAppendMemoryPrefetch(cmdlist_immediate_async_mode, timeout, false);
+}
+
+TEST_F(
+    zeImmediateCommandListInOrderExecutionTests,
+    GivenInOrderImmediateCommandListWhenAppendMemoryPrefetchThenSuccessIsReturnedWithSharedSystemAllocator) {
+  SKIP_IF_SHARED_SYSTEM_ALLOC_UNSUPPORTED();
+  const uint64_t timeout = UINT64_MAX - 1;
+  RunAppendMemoryPrefetch(cmdlist_immediate_default_mode, timeout, true);
+  RunAppendMemoryPrefetch(cmdlist_immediate_sync_mode, 0, true);
+  RunAppendMemoryPrefetch(cmdlist_immediate_async_mode, timeout, true);
 }
 
 TEST_P(zeImmediateCommandListExecutionTests,

conformance_tests/core/test_copy/src/test_copy.cpp

@@ -1349,10 +1349,12 @@ TEST_F(
 
 class zeCommandListAppendMemoryPrefetchTests : public ::testing::Test {
 protected:
-  void RunGivenMemoryPrefetchWhenWritingFromDeviceTest(bool is_immediate) {
+  void RunGivenMemoryPrefetchWhenWritingFromDeviceTest(bool is_immediate,
+                                                       bool is_shared_system) {
     const size_t size = 16;
     const uint8_t value = 0x55;
-    void *memory = allocate_shared_memory(size);
+    void *memory = lzt::allocate_shared_memory_with_allocator_selector(
+        size, is_shared_system);
     memset(memory, 0xaa, size);
     auto cmd_bundle = lzt::create_command_bundle(is_immediate);
 
@@ -1366,20 +1368,34 @@ class zeCommandListAppendMemoryPrefetchTests : public ::testing::Test {
       ASSERT_EQ(value, ((uint8_t *)memory)[i]);
     }
 
-    free_memory(memory);
+    free_memory_with_allocator_selector(memory, is_shared_system);
     lzt::destroy_command_bundle(cmd_bundle);
   }
 };
 
 TEST_F(zeCommandListAppendMemoryPrefetchTests,
        GivenMemoryPrefetchWhenWritingFromDeviceThenDataisCorrectFromHost) {
-  RunGivenMemoryPrefetchWhenWritingFromDeviceTest(false);
+  RunGivenMemoryPrefetchWhenWritingFromDeviceTest(false, false);
 }
 
 TEST_F(
     zeCommandListAppendMemoryPrefetchTests,
     GivenMemoryPrefetchWhenWritingFromDeviceOnImmediateCmdListThenDataisCorrectFromHost) {
-  RunGivenMemoryPrefetchWhenWritingFromDeviceTest(true);
+  RunGivenMemoryPrefetchWhenWritingFromDeviceTest(true, false);
+}
+
+TEST_F(
+    zeCommandListAppendMemoryPrefetchTests,
+    GivenMemoryPrefetchWhenWritingFromDeviceThenDataisCorrectFromHostWithSharedSystemAllocator) {
+  SKIP_IF_SHARED_SYSTEM_ALLOC_UNSUPPORTED();
+  RunGivenMemoryPrefetchWhenWritingFromDeviceTest(false, true);
+}
+
+TEST_F(
+    zeCommandListAppendMemoryPrefetchTests,
+    GivenMemoryPrefetchWhenWritingFromDeviceOnImmediateCmdListThenDataisCorrectFromHostWithSharedSystemAllocator) {
+  SKIP_IF_SHARED_SYSTEM_ALLOC_UNSUPPORTED();
+  RunGivenMemoryPrefetchWhenWritingFromDeviceTest(true, true);
 }
 
 class zeCommandListAppendMemAdviseTests