fix: pass USM alloc to external CB event if possible

Signed-off-by: Bartosz Dunajski <[email protected]>

Author: BartoszDunajski

level_zero/core/source/event/event.cpp

@@ -673,6 +673,15 @@ void Event::resetInOrderTimestampNode(NEO::TagNodeBase *newNode) {
     inOrderTimestampNode = newNode;
 }
 
+NEO::GraphicsAllocation *Event::getExternalCounterAllocationFromAddress(uint64_t *address) const {
+    NEO::SvmAllocationData *allocData = nullptr;
+    if (!address || !device->getDriverHandle()->findAllocationDataForRange(address, sizeof(uint64_t), allocData)) {
+        return nullptr;
+    }
+
+    return allocData->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
+}
+
 ze_result_t Event::enableExtensions(const EventDescriptor &eventDescriptor) {
     bool interruptMode = false;
     bool kmdWaitMode = false;
@@ -699,31 +708,31 @@ ze_result_t Event::enableExtensions(const EventDescriptor &eventDescriptor) {
                 return ZE_RESULT_ERROR_INVALID_ARGUMENT;
             }
 
-            NEO::SvmAllocationData *externalHostAllocData = nullptr;
-            if (!device->getDriverHandle()->findAllocationDataForRange(externalSyncAllocProperties->hostAddress, sizeof(uint64_t), externalHostAllocData)) {
+            auto deviceAlloc = getExternalCounterAllocationFromAddress(externalSyncAllocProperties->deviceAddress);
+            auto hostAlloc = getExternalCounterAllocationFromAddress(externalSyncAllocProperties->hostAddress);
+
+            if (!hostAlloc) {
                 return ZE_RESULT_ERROR_INVALID_ARGUMENT;
             }
 
-            auto allocation = externalHostAllocData->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
-            auto inOrderExecInfo = NEO::InOrderExecInfo::createFromExternalAllocation(*device->getNEODevice(), nullptr, castToUint64(externalSyncAllocProperties->deviceAddress),
-                                                                                      allocation, externalSyncAllocProperties->hostAddress, externalSyncAllocProperties->completionValue, 1, 1);
+            auto inOrderExecInfo = NEO::InOrderExecInfo::createFromExternalAllocation(*device->getNEODevice(), deviceAlloc, castToUint64(externalSyncAllocProperties->deviceAddress),
+                                                                                      hostAlloc, externalSyncAllocProperties->hostAddress, externalSyncAllocProperties->completionValue, 1, 1);
             updateInOrderExecState(inOrderExecInfo, externalSyncAllocProperties->completionValue, 0);
         } else if (extendedDesc->stype == ZEX_STRUCTURE_COUNTER_BASED_EVENT_EXTERNAL_STORAGE_ALLOC_PROPERTIES) {
             auto externalStorageProperties = reinterpret_cast<const zex_counter_based_event_external_storage_properties_t *>(extendedDesc);
 
-            NEO::SvmAllocationData *externalDeviceAllocData = nullptr;
-            if (!externalStorageProperties->deviceAddress || !device->getDriverHandle()->findAllocationDataForRange(externalStorageProperties->deviceAddress, sizeof(uint64_t), externalDeviceAllocData) ||
-                externalStorageProperties->incrementValue == 0) {
+            auto deviceAlloc = getExternalCounterAllocationFromAddress(externalStorageProperties->deviceAddress);
+
+            if (!deviceAlloc || externalStorageProperties->incrementValue == 0) {
                 return ZE_RESULT_ERROR_INVALID_ARGUMENT;
             }
 
-            auto allocation = externalDeviceAllocData->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
-            auto offset = ptrDiff(externalStorageProperties->deviceAddress, allocation->getGpuAddress());
+            auto offset = ptrDiff(externalStorageProperties->deviceAddress, deviceAlloc->getGpuAddress());
 
-            auto hostAddress = ptrOffset(device->getNEODevice()->getMemoryManager()->lockResource(allocation), offset);
+            auto hostAddress = ptrOffset(device->getNEODevice()->getMemoryManager()->lockResource(deviceAlloc), offset);
 
-            auto inOrderExecInfo = NEO::InOrderExecInfo::createFromExternalAllocation(*device->getNEODevice(), allocation, castToUint64(externalStorageProperties->deviceAddress),
-                                                                                      allocation, reinterpret_cast<uint64_t *>(hostAddress), externalStorageProperties->completionValue, 1, 1);
+            auto inOrderExecInfo = NEO::InOrderExecInfo::createFromExternalAllocation(*device->getNEODevice(), deviceAlloc, castToUint64(externalStorageProperties->deviceAddress),
+                                                                                      deviceAlloc, reinterpret_cast<uint64_t *>(hostAddress), externalStorageProperties->completionValue, 1, 1);
 
             updateInOrderExecState(inOrderExecInfo, externalStorageProperties->completionValue, 0);
 

level_zero/core/source/event/event.h

@@ -341,6 +341,7 @@ struct Event : _ze_event_handle_t {
     Event(int index, Device *device) : device(device), index(index) {}
 
     ze_result_t enableExtensions(const EventDescriptor &eventDescriptor);
+    NEO::GraphicsAllocation *getExternalCounterAllocationFromAddress(uint64_t *address) const;
 
     void unsetCmdQueue();
     void releaseTempInOrderTimestampNodes();

level_zero/core/test/unit_tests/sources/cmdlist/test_in_order_cmdlist_1.cpp

@@ -5237,19 +5237,51 @@ HWTEST2_F(InOrderCmdListTests, givenCorrectInputParamsWhenCreatingCbEvent2ThenRe
     EXPECT_EQ(counterValue, eventObj->getInOrderExecInfo()->getCounterValue());
     EXPECT_EQ(hostAddress, eventObj->getInOrderExecInfo()->getBaseHostAddress());
     EXPECT_EQ(castToUint64(gpuAddress), eventObj->getInOrderExecInfo()->getBaseDeviceAddress());
+    EXPECT_EQ(nullptr, eventObj->getInOrderExecInfo()->getDeviceCounterAllocation());
 
-    uint64_t addresss = 0;
+    uint64_t address = 0;
     uint64_t value = 0;
-    zexEventGetDeviceAddress(handle, &value, &addresss);
+    zexEventGetDeviceAddress(handle, &value, &address);
 
-    EXPECT_EQ(addresss, eventObj->getInOrderExecInfo()->getBaseDeviceAddress());
+    EXPECT_EQ(address, eventObj->getInOrderExecInfo()->getBaseDeviceAddress());
     EXPECT_EQ(value, counterValue);
 
     zeEventDestroy(handle);
 
     context->freeMem(hostAddress);
 }
 
+HWTEST2_F(InOrderCmdListTests, givenUsmDeviceAllocationWhenCreatingCbEventFromExternalMemoryThenAssignGraphicsAllocation, MatchAny) {
+    auto hostAddress = reinterpret_cast<uint64_t *>(allocHostMem(sizeof(uint64_t)));
+    auto deviceAddress = reinterpret_cast<uint64_t *>(allocDeviceMem(sizeof(uint64_t)));
+
+    zex_counter_based_event_external_sync_alloc_properties_t externalSyncAllocProperties = {ZEX_STRUCTURE_COUNTER_BASED_EVENT_EXTERNAL_SYNC_ALLOC_PROPERTIES}; // NOLINT(clang-analyzer-optin.core.EnumCastOutOfRange), NEO-12901
+    externalSyncAllocProperties.completionValue = 2;
+    externalSyncAllocProperties.deviceAddress = deviceAddress;
+    externalSyncAllocProperties.hostAddress = hostAddress;
+
+    zex_counter_based_event_desc_t counterBasedDesc = {ZEX_STRUCTURE_COUNTER_BASED_EVENT_DESC}; // NOLINT(clang-analyzer-optin.core.EnumCastOutOfRange), NEO-12901
+    counterBasedDesc.flags = ZEX_COUNTER_BASED_EVENT_FLAG_IMMEDIATE | ZEX_COUNTER_BASED_EVENT_FLAG_NON_IMMEDIATE;
+    counterBasedDesc.pNext = &externalSyncAllocProperties;
+
+    ze_event_handle_t handle = nullptr;
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zexCounterBasedEventCreate2(context, device, &counterBasedDesc, &handle));
+
+    auto eventObj = Event::fromHandle(handle);
+
+    ASSERT_NE(nullptr, eventObj);
+    ASSERT_NE(nullptr, eventObj->getInOrderExecInfo().get());
+
+    EXPECT_EQ(castToUint64(deviceAddress), eventObj->getInOrderExecInfo()->getBaseDeviceAddress());
+    EXPECT_NE(nullptr, eventObj->getInOrderExecInfo()->getDeviceCounterAllocation());
+
+    zeEventDestroy(handle);
+
+    context->freeMem(hostAddress);
+    context->freeMem(deviceAddress);
+}
+
 HWTEST2_F(InOrderCmdListTests, givenExternalSyncStorageWhenCreatingCounterBasedEventThenSetAllParams, MatchAny) {
     uint64_t counterValue = 4;
     uint64_t incValue = 2;