feature: add experimental extension to verify memory in aub mode

Related-To: NEO-14153, NEO-17038

Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>

level_zero/api/driver_experimental/public/zex_cmdlist.cpp

@@ -176,6 +176,32 @@ zexCommandListSetCleanupCallback(ze_command_list_handle_t hCommandList, zex_comm
     return ZE_RESULT_SUCCESS;
 }
 
+ze_result_t ZE_APICALL
+zexCommandListVerifyMemory(ze_command_list_handle_t hCommandList,
+                           const void *allocationPtr,
+                           const void *expectedData,
+                           size_t sizeOfComparison,
+                           zex_verify_memory_compare_type_t comparisonMode) {
+    auto cmdList = L0::CommandList::fromHandle(hCommandList);
+
+    if (!cmdList) {
+        return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
+    }
+
+    if (!cmdList->isImmediateType()) {
+        return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+
+    if (!allocationPtr || !expectedData) {
+        return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
+    }
+    if (sizeOfComparison == 0) {
+        return ZE_RESULT_ERROR_INVALID_SIZE;
+    }
+
+    return cmdList->verifyMemory(allocationPtr, expectedData, sizeOfComparison, comparisonMode) ? ZE_RESULT_SUCCESS : ZE_RESULT_ERROR_UNKNOWN;
+}
+
 } // namespace L0
 
 extern "C" {
@@ -253,4 +279,14 @@ zexCommandListSetCleanupCallback(ze_command_list_handle_t hCommandList, zex_comm
     return L0::zexCommandListSetCleanupCallback(hCommandList, pfnCallback, pUserData, pNext);
 }
 
+ZE_APIEXPORT ze_result_t ZE_APICALL
+zexCommandListVerifyMemory(ze_command_list_handle_t hCommandList,
+                           const void *allocationPtr,
+                           const void *expectedData,
+                           size_t sizeOfComparison,
+                           zex_verify_memory_compare_type_t comparisonMode) {
+
+    return L0::zexCommandListVerifyMemory(hCommandList, allocationPtr, expectedData, sizeOfComparison, comparisonMode);
+}
+
 } // extern "C"

level_zero/core/source/cmdlist/cmdlist.cpp

@@ -299,4 +299,11 @@ void CommandList::executeCleanupCallbacks() {
     }
 }
 
+bool CommandList::verifyMemory(const void *allocationPtr,
+                               const void *expectedData,
+                               size_t sizeOfComparison,
+                               uint32_t comparisonMode) const {
+    return getCsr(false)->expectMemory(allocationPtr, expectedData, sizeOfComparison, comparisonMode);
+}
+
 } // namespace L0

level_zero/core/source/cmdlist/cmdlist.h

@@ -557,6 +557,10 @@ struct CommandList : _ze_command_list_handle_t {
     }
 
     void executeCleanupCallbacks();
+    bool verifyMemory(const void *allocationPtr,
+                      const void *expectedData,
+                      size_t sizeOfComparison,
+                      uint32_t comparisonMode) const;
 
   protected:
     using CleanupCallbackT = std::pair<zex_command_list_cleanup_callback_fn_t, void *>;

level_zero/core/source/driver/extension_function_address.cpp

@@ -72,6 +72,7 @@ void *ExtensionFunctionAddressHelper::getExtensionFunctionAddress(const std::str
     RETURN_FUNC_PTR_IF_EXIST(zeCommandListAppendHostFunction);
     RETURN_FUNC_PTR_IF_EXIST(zexCommandListAppendMemoryCopyWithParameters);
     RETURN_FUNC_PTR_IF_EXIST(zexCommandListAppendMemoryFillWithParameters);
+    RETURN_FUNC_PTR_IF_EXIST(zexCommandListVerifyMemory);
 
     // mutable command list extension
     RETURN_FUNC_PTR_IF_EXIST(zexCommandListGetVariable);

level_zero/core/test/unit_tests/experimental/CMakeLists.txt

@@ -6,6 +6,7 @@
 
 target_sources(${TARGET_NAME} PRIVATE
                ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
+               ${CMAKE_CURRENT_SOURCE_DIR}/test_cmdlist_exp.cpp
                ${CMAKE_CURRENT_SOURCE_DIR}/test_graph.h
                ${CMAKE_CURRENT_SOURCE_DIR}/test_graph.cpp
                ${CMAKE_CURRENT_SOURCE_DIR}/test_graph_export.h

level_zero/core/test/unit_tests/experimental/test_cmdlist_exp.cpp

@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2025 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/test/common/mocks/mock_aub_csr.h"
+#include "shared/test/common/test_macros/hw_test.h"
+
+#include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
+#include "level_zero/core/test/unit_tests/mocks/mock_cmdlist.h"
+#include "level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h"
+
+using namespace NEO;
+
+namespace L0 {
+
+namespace ult {
+
+using CommandListExpTest = Test<DeviceFixture>;
+
+HWTEST_F(CommandListExpTest, givenCmdListWithSimulatedCsrWhenVerifyMemoryCalledThenExpectMemoryIsCalled) {
+    ze_command_queue_desc_t desc = {};
+    desc.mode = ZE_COMMAND_QUEUE_MODE_SYNCHRONOUS;
+
+    ze_result_t returnValue;
+    std::unique_ptr<L0::CommandList> commandList(CommandList::createImmediate(productFamily, device, &desc, false, NEO::EngineGroupType::renderCompute, returnValue));
+    auto &commandListImmediate = static_cast<MockCommandListImmediate<FamilyType::gfxCoreFamily> &>(*commandList);
+
+    MockAubCsr<FamilyType> mockCommandStreamReceiver("", true, *neoDevice->executionEnvironment, neoDevice->getRootDeviceIndex(), neoDevice->getDeviceBitfield());
+    mockCommandStreamReceiver.callBaseExpectMemory = false;
+    Mock<CommandQueue> mockCommandQueue(device, &mockCommandStreamReceiver, &desc);
+
+    auto oldCommandQueue = commandListImmediate.cmdQImmediate;
+    commandListImmediate.cmdQImmediate = &mockCommandQueue;
+    void *ptr = reinterpret_cast<void *>(0x1000);
+    char data[10] = {};
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zexCommandListVerifyMemory(commandList->toHandle(), ptr, data, sizeof(data), zex_verify_memory_compare_type_t::ZEX_VERIFY_MEMORY_COMPARE_EQUAL));
+    EXPECT_TRUE(mockCommandStreamReceiver.expectMemoryCalled);
+
+    commandListImmediate.cmdQImmediate = oldCommandQueue;
+}
+
+HWTEST_F(CommandListExpTest, givenCmdListWithSimulatedCsrWhenVerifyMemoryCalledWithInvalidParamsThenErrorIsReturned) {
+    ze_command_queue_desc_t desc = {};
+    desc.mode = ZE_COMMAND_QUEUE_MODE_SYNCHRONOUS;
+
+    ze_result_t returnValue;
+    std::unique_ptr<L0::CommandList> commandList(CommandList::createImmediate(productFamily, device, &desc, false, NEO::EngineGroupType::renderCompute, returnValue));
+    auto &commandListImmediate = static_cast<MockCommandListImmediate<FamilyType::gfxCoreFamily> &>(*commandList);
+
+    MockAubCsr<FamilyType> mockCommandStreamReceiver("", true, *neoDevice->executionEnvironment, neoDevice->getRootDeviceIndex(), neoDevice->getDeviceBitfield());
+    mockCommandStreamReceiver.callBaseExpectMemory = false;
+    Mock<CommandQueue> mockCommandQueue(device, &mockCommandStreamReceiver, &desc);
+
+    auto oldCommandQueue = commandListImmediate.cmdQImmediate;
+    commandListImmediate.cmdQImmediate = &mockCommandQueue;
+    void *ptr = reinterpret_cast<void *>(0x1000);
+    char data[10] = {};
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_NULL_HANDLE, zexCommandListVerifyMemory(nullptr, ptr, data, sizeof(data), zex_verify_memory_compare_type_t::ZEX_VERIFY_MEMORY_COMPARE_EQUAL));
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_NULL_POINTER, zexCommandListVerifyMemory(commandList->toHandle(), nullptr, data, sizeof(data), zex_verify_memory_compare_type_t::ZEX_VERIFY_MEMORY_COMPARE_EQUAL));
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_NULL_POINTER, zexCommandListVerifyMemory(commandList->toHandle(), ptr, nullptr, sizeof(data), zex_verify_memory_compare_type_t::ZEX_VERIFY_MEMORY_COMPARE_EQUAL));
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_SIZE, zexCommandListVerifyMemory(commandList->toHandle(), ptr, data, 0, zex_verify_memory_compare_type_t::ZEX_VERIFY_MEMORY_COMPARE_EQUAL));
+    EXPECT_FALSE(mockCommandStreamReceiver.expectMemoryCalled);
+
+    commandListImmediate.cmdQImmediate = oldCommandQueue;
+}
+
+HWTEST_F(CommandListExpTest, givenRegularCmdListWithSimulatedCsrWhenVerifyMemoryCalledThenExpectMemoryIsCalled) {
+    ze_result_t returnValue;
+    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily, device, NEO::EngineGroupType::renderCompute, 0, returnValue, false));
+    void *ptr = reinterpret_cast<void *>(0x1000);
+    char data[10] = {};
+    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_FEATURE, zexCommandListVerifyMemory(commandList->toHandle(), ptr, data, sizeof(data), zex_verify_memory_compare_type_t::ZEX_VERIFY_MEMORY_COMPARE_EQUAL));
+}
+
+} // namespace ult
+} // namespace L0

level_zero/core/test/unit_tests/sources/driver/test_driver.cpp

@@ -1359,6 +1359,7 @@ TEST_F(DriverExperimentalApiTest, whenRetrievingApiFunctionThenExpectProperPoint
     using pfnCommandListAppendMILoadRegImm = decltype(&zexCommandListAppendMILoadRegImm);
     using pfnCommandListAppendMIStoreRegMem = decltype(&zexCommandListAppendMIStoreRegMem);
     using pfnCommandListAppendMIMath = decltype(&zexCommandListAppendMIMath);
+    using pfnCommandListVerifyMemory = decltype(&zexCommandListVerifyMemory);
 
     decltype(&zexDriverImportExternalPointer) expectedImport = zexDriverImportExternalPointer;
     decltype(&zexDriverReleaseImportedPointer) expectedRelease = zexDriverReleaseImportedPointer;
@@ -1405,7 +1406,7 @@ TEST_F(DriverExperimentalApiTest, whenRetrievingApiFunctionThenExpectProperPoint
     pfnCommandListAppendMILoadRegImm expectedCommandListAppendMILoadRegImm = zexCommandListAppendMILoadRegImm;
     pfnCommandListAppendMIStoreRegMem expectedCommandListAppendMIStoreRegMem = zexCommandListAppendMIStoreRegMem;
     pfnCommandListAppendMIMath expectedCommandListAppendMIMath = zexCommandListAppendMIMath;
-
+    pfnCommandListVerifyMemory expectedCommandListVerifyMemory = zexCommandListVerifyMemory;
     // Add EXPECT_EQ tests to verify function pointers
     void *funPtr = nullptr;
 
@@ -1528,6 +1529,9 @@ TEST_F(DriverExperimentalApiTest, whenRetrievingApiFunctionThenExpectProperPoint
 
     EXPECT_EQ(ZE_RESULT_SUCCESS, zeDriverGetExtensionFunctionAddress(driverHandle, "zexCommandListAppendMIMath", &funPtr));
     EXPECT_EQ(expectedCommandListAppendMIMath, reinterpret_cast<pfnCommandListAppendMIMath>(funPtr));
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeDriverGetExtensionFunctionAddress(driverHandle, "zexCommandListVerifyMemory", &funPtr));
+    EXPECT_EQ(expectedCommandListVerifyMemory, reinterpret_cast<pfnCommandListVerifyMemory>(funPtr));
 }
 
 TEST_F(DriverExperimentalApiTest, givenHostPointerApiExistWhenImportingPtrThenExpectProperBehavior) {

level_zero/include/level_zero/driver_experimental/zex_cmdlist.h

@@ -84,6 +84,13 @@ zexCommandListSetCleanupCallback(
     void *pUserData,                                    ///< [in] user specific data that would be passed to function
     const void *pNext);                                 ///< [in][optional] must be null or a pointer to an extension-specific structure
 
+ZE_APIEXPORT ze_result_t ZE_APICALL
+zexCommandListVerifyMemory(ze_command_list_handle_t hCommandList,
+                           const void *allocationPtr,
+                           const void *expectedData,
+                           size_t sizeOfComparison,
+                           zex_verify_memory_compare_type_t comparisonMode);
+
 #if defined(__cplusplus)
 } // extern "C"
 #endif

level_zero/include/level_zero/driver_experimental/zex_common.h

@@ -240,6 +240,12 @@ typedef struct _zex_counter_based_event_external_storage_properties_t {
     uint64_t completionValue;      ///< [in] final completion value, when value under deviceAddress is equal or greater then this value then event is considered as completed
 } zex_counter_based_event_external_storage_properties_t;
 
+typedef enum _zex_verify_memory_compare_type_t {
+    ZEX_VERIFY_MEMORY_COMPARE_EQUAL = 0,                // compare memory for equality
+    ZEX_VERIFY_MEMORY_COMPARE_NOT_EQUAL = 1,            // compare memory for inequality
+    ZEX_VERIFY_MEMORY_COMPARE_FORCE_UINT32 = 0x7fffffff ///< Value marking end of ZEX_VERIFY_MEMORY_COMPARE* ENUMs
+} zex_verify_memory_compare_type_t;
+
 #if defined(__cplusplus)
 } // extern "C"
 #endif

shared/source/command_stream/tbx_command_stream_receiver_hw.inl

@@ -313,6 +313,7 @@ bool TbxCommandStreamReceiverHw<GfxFamily>::expectMemory(const void *gfxAddress,
         hardwareContextController->readMemory((uint64_t)gfxAddress, readMemory.get(), length, this->getMemoryBankForGtt(), MemoryConstants::pageSize64k);
         auto isMemoryEqual = (memcmp(readMemory.get(), srcAddress, length) == 0);
         auto isEqualMemoryExpected = (compareOperation == aub_stream::CompareOperationValues::CompareEqual);
+        hardwareContextController->expectMemory(reinterpret_cast<uint64_t>(gfxAddress), srcAddress, length, compareOperation);
         return (isMemoryEqual == isEqualMemoryExpected);
     }
 

shared/test/common/mocks/mock_aub_csr.h

@@ -99,6 +99,14 @@ struct MockAubCsr : public AUBCommandStreamReceiverHw<GfxFamily> {
         skipTaskCountCheckForCompletionPoll = skipTaskCountCheck;
     }
 
+    bool expectMemory(const void *gfxAddress, const void *srcAddress, size_t length, uint32_t compareOperation) override {
+        expectMemoryCalled = true;
+        if (callBaseExpectMemory) {
+            return AUBCommandStreamReceiverHw<GfxFamily>::expectMemory(gfxAddress, srcAddress, length, compareOperation);
+        }
+        return true;
+    }
+
     bool expectMemoryEqual(void *gfxAddress, const void *srcAddress, size_t length) override {
         expectMemoryEqualCalled = true;
         return AUBCommandStreamReceiverHw<GfxFamily>::expectMemoryEqual(gfxAddress, srcAddress, length);
@@ -139,6 +147,7 @@ struct MockAubCsr : public AUBCommandStreamReceiverHw<GfxFamily> {
     bool writeMMIOCalled = false;
     bool submitBatchBufferCalled = false;
     bool pollForCompletionCalled = false;
+    bool expectMemoryCalled = false;
     bool expectMemoryEqualCalled = false;
     bool expectMemoryNotEqualCalled = false;
     bool expectMemoryCompressedCalled = false;
@@ -146,6 +155,7 @@ struct MockAubCsr : public AUBCommandStreamReceiverHw<GfxFamily> {
     bool dumpAllocationCalled = false;
     bool skipTaskCountCheckForCompletionPoll = false;
     bool lockStreamCalled = false;
+    bool callBaseExpectMemory = true;
 
     std::unique_lock<std::mutex> lockStream() override {
         lockStreamCalled = true;