From c4306a13b2adb791339e9b6e942518877d0f9970 Mon Sep 17 00:00:00 2001
From: Kamil Kopryk <kamil.kopryk@intel.com>
Date: Wed, 11 Sep 2024 11:21:09 +0000
Subject: [PATCH] fix: set debugSurface in subdevices

Related-To: NEO-10681

Signed-off-by: Kamil Kopryk <kamil.kopryk@intel.com>
---
 shared/source/device/device.cpp               | 30 ++++++++++++-------
 shared/source/device/device.h                 |  1 +
 shared/source/device/root_device.cpp          |  7 +++++
 shared/test/common/mocks/mock_device.h        |  7 +++++
 .../unit_test/device/neo_device_tests.cpp     | 15 ++++++++++
 5 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/shared/source/device/device.cpp b/shared/source/device/device.cpp
index 921847122a..12b6dab459 100644
--- a/shared/source/device/device.cpp
+++ b/shared/source/device/device.cpp
@@ -272,15 +272,9 @@ void Device::initializeCommonResources() {
         debugSurfaceSize = NEO::SipKernel::getSipKernel(*this, nullptr).getStateSaveAreaSize(this);
     }
 
-    const bool allocateDebugSurface = getL0Debugger();
-    if (allocateDebugSurface) {
-        debugSurface = getMemoryManager()->allocateGraphicsMemoryWithProperties(
-            {getRootDeviceIndex(), true,
-             debugSurfaceSize,
-             NEO::AllocationType::debugContextSaveArea,
-             false,
-             false,
-             getDeviceBitfield()});
+    const bool isDebugSurfaceRequired = getL0Debugger();
+    if (isDebugSurfaceRequired) {
+        allocateDebugSurface(debugSurfaceSize);
     }
 
     if (ApiSpecificConfig::isDeviceUsmPoolingEnabled() &&
@@ -467,7 +461,23 @@ void Device::createSecondaryContexts(const EngineControl &primaryEngine, Seconda
     }
 
     primaryEngine.osContext->setContextGroup(true);
-};
+}
+
+void Device::allocateDebugSurface(size_t debugSurfaceSize) {
+    this->debugSurface = getMemoryManager()->allocateGraphicsMemoryWithProperties(
+        {getRootDeviceIndex(), true,
+         debugSurfaceSize,
+         NEO::AllocationType::debugContextSaveArea,
+         false,
+         false,
+         getDeviceBitfield()});
+
+    for (auto &subdevice : this->subdevices) {
+        if (subdevice) {
+            subdevice->debugSurface = this->debugSurface;
+        }
+    }
+}
 
 void Device::addEngineToEngineGroup(EngineControl &engine) {
     auto &hardwareInfo = this->getHardwareInfo();
diff --git a/shared/source/device/device.h b/shared/source/device/device.h
index 8ae4ae5a61..ed4f68b143 100644
--- a/shared/source/device/device.h
+++ b/shared/source/device/device.h
@@ -255,6 +255,7 @@ class Device : public ReferenceTrackedObject<Device> {
     void setAsEngineInstanced();
     void finalizeRayTracing();
     void createSecondaryContexts(const EngineControl &primaryEngine, SecondaryContexts &secondaryEnginesForType, uint32_t contextCount, uint32_t regularPriorityCount, uint32_t highPriorityContextCount);
+    void allocateDebugSurface(size_t debugSurfaceSize);
 
     DeviceInfo deviceInfo = {};
 
diff --git a/shared/source/device/root_device.cpp b/shared/source/device/root_device.cpp
index 5467097993..5502437268 100644
--- a/shared/source/device/root_device.cpp
+++ b/shared/source/device/root_device.cpp
@@ -29,6 +29,13 @@ RootDevice::RootDevice(ExecutionEnvironment *executionEnvironment, uint32_t root
 
 RootDevice::~RootDevice() {
     if (getDebugSurface()) {
+
+        for (auto *subDevice : this->getSubDevices()) {
+            if (subDevice) {
+                subDevice->setDebugSurface(nullptr);
+            }
+        }
+
         getMemoryManager()->freeGraphicsMemory(debugSurface);
         debugSurface = nullptr;
     }
diff --git a/shared/test/common/mocks/mock_device.h b/shared/test/common/mocks/mock_device.h
index 484e041507..14cdc437cd 100644
--- a/shared/test/common/mocks/mock_device.h
+++ b/shared/test/common/mocks/mock_device.h
@@ -17,6 +17,8 @@
 #include "shared/test/common/helpers/variable_backup.h"
 #include "shared/test/common/mocks/mock_memory_operations_handler.h"
 
+#include "gtest/gtest.h"
+
 namespace NEO {
 class CommandStreamReceiver;
 class DriverInfo;
@@ -37,6 +39,10 @@ struct MockSubDevice : public SubDevice {
     using SubDevice::getGlobalMemorySize;
     using SubDevice::SubDevice;
 
+    ~MockSubDevice() override {
+        EXPECT_EQ(nullptr, this->getDebugSurface());
+    }
+
     std::unique_ptr<CommandStreamReceiver> createCommandStreamReceiver() const override {
         return std::unique_ptr<CommandStreamReceiver>(createCommandStreamReceiverFunc(*executionEnvironment, getRootDeviceIndex(), getDeviceBitfield()));
     }
@@ -50,6 +56,7 @@ class MockDevice : public RootDevice {
   public:
     using Device::addEngineToEngineGroup;
     using Device::allEngines;
+    using Device::allocateDebugSurface;
     using Device::commandStreamReceivers;
     using Device::createDeviceInternals;
     using Device::createEngine;
diff --git a/shared/test/unit_test/device/neo_device_tests.cpp b/shared/test/unit_test/device/neo_device_tests.cpp
index 7ea69cc92c..939010f60e 100644
--- a/shared/test/unit_test/device/neo_device_tests.cpp
+++ b/shared/test/unit_test/device/neo_device_tests.cpp
@@ -1153,6 +1153,21 @@ TEST(Device, givenDifferentEngineTypesWhenIsSecondaryContextEngineTypeCalledThen
     }
 }
 
+TEST(Device, whenAllocateDebugSurfaceIsCalledThenEachSubDeviceContainsCorrectDebugSurface) {
+
+    DebugManagerStateRestore dbgRestorer;
+    debugManager.flags.CreateMultipleSubDevices.set(4);
+    auto device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(defaultHwInfo.get()));
+
+    size_t size = 8u;
+    device->allocateDebugSurface(size);
+    auto *debugSurface = device->getDebugSurface();
+
+    for (auto *subDevice : device->getSubDevices()) {
+        EXPECT_EQ(debugSurface, subDevice->getDebugSurface());
+    }
+}
+
 HWCMDTEST_F(IGFX_XE_HP_CORE, DeviceTests, givenCCSEngineAndContextGroupSizeEnabledWhenCreatingEngineThenItsContextHasContextGroupFlagSet) {
     DebugManagerStateRestore dbgRestorer;
     const uint32_t contextGroupSize = 8;