Register zebin module for all kernel allocations

Related-To: NEO-6637

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

level_zero/core/source/debugger/debugger_l0.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2021 Intel Corporation
+ * Copyright (C) 2020-2022 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -90,6 +90,9 @@ class DebuggerL0 : public NEO::Debugger, NEO::NonCopyableOrMovableClass {
     virtual size_t getSbaTrackingCommandsSize(size_t trackedAddressCount) = 0;
     virtual void programSbaTrackingCommands(NEO::LinearStream &cmdStream, const SbaAddresses &sba) = 0;
 
+    MOCKABLE_VIRTUAL bool attachZebinModuleToSegmentAllocations(const StackVec<NEO::GraphicsAllocation *, 32> &kernelAlloc, uint32_t &moduleHandle);
+    MOCKABLE_VIRTUAL bool removeZebinModule(uint32_t moduleHandle);
+
   protected:
     static bool isAnyTrackedAddressChanged(SbaAddresses sba) {
         return sba.GeneralStateBaseAddress != 0 ||

level_zero/core/source/debugger/linux/debugger_l0_linux.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2021 Intel Corporation
+ * Copyright (C) 2020-2022 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -32,4 +32,30 @@ void DebuggerL0::registerElf(NEO::DebugData *debugData, NEO::GraphicsAllocation
         static_cast<NEO::DrmAllocation *>(isaAllocation)->linkWithRegisteredHandle(handle);
     }
 }
+
+bool DebuggerL0::attachZebinModuleToSegmentAllocations(const StackVec<NEO::GraphicsAllocation *, 32> &allocs, uint32_t &moduleHandle) {
+    if (device->getRootDeviceEnvironment().osInterface == nullptr) {
+        return false;
+    }
+    auto drm = device->getRootDeviceEnvironment().osInterface->getDriverModel()->as<NEO::Drm>();
+    uint32_t segmentCount = static_cast<uint32_t>(allocs.size());
+    moduleHandle = drm->registerResource(NEO::Drm::ResourceClass::L0ZebinModule, &segmentCount, sizeof(uint32_t));
+
+    for (auto &allocation : allocs) {
+        auto drmAllocation = static_cast<NEO::DrmAllocation *>(allocation);
+        drmAllocation->linkWithRegisteredHandle(moduleHandle);
+    }
+
+    return true;
+}
+
+bool DebuggerL0::removeZebinModule(uint32_t moduleHandle) {
+    if (device->getRootDeviceEnvironment().osInterface == nullptr) {
+        return false;
+    }
+    auto drm = device->getRootDeviceEnvironment().osInterface->getDriverModel()->as<NEO::Drm>();
+
+    drm->unregisterResource(moduleHandle);
+    return true;
+}
 } // namespace L0

level_zero/core/source/debugger/windows/debugger_l0_windows.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2021 Intel Corporation
+ * Copyright (C) 2020-2022 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -19,4 +19,12 @@ bool DebuggerL0::initDebuggingInOs(NEO::OSInterface *osInterface) {
 void DebuggerL0::registerElf(NEO::DebugData *debugData, NEO::GraphicsAllocation *isaAllocation) {
 }
 
+bool DebuggerL0::attachZebinModuleToSegmentAllocations(const StackVec<NEO::GraphicsAllocation *, 32> &kernelAlloc, uint32_t &moduleHandle) {
+    return false;
+}
+
+bool DebuggerL0::removeZebinModule(uint32_t moduleHandle) {
+    return false;
+}
+
 } // namespace L0

level_zero/core/source/module/module_imp.cpp

@@ -1005,9 +1005,21 @@ void ModuleImp::registerElfInDebuggerL0() {
         NEO::DebugData debugData; // pass debug zebin in vIsa field
         debugData.vIsa = reinterpret_cast<const char *>(translationUnit->debugData.get());
         debugData.vIsaSize = static_cast<uint32_t>(translationUnit->debugDataSize);
+
+        StackVec<NEO::GraphicsAllocation *, 32> segmentAllocs;
         for (auto &kernImmData : kernelImmDatas) {
             device->getL0Debugger()->registerElf(&debugData, kernImmData->getIsaGraphicsAllocation());
+            segmentAllocs.push_back(kernImmData->getIsaGraphicsAllocation());
         }
+
+        if (translationUnit->globalVarBuffer) {
+            segmentAllocs.push_back(translationUnit->globalVarBuffer);
+        }
+        if (translationUnit->globalConstBuffer) {
+            segmentAllocs.push_back(translationUnit->globalConstBuffer);
+        }
+
+        device->getL0Debugger()->attachZebinModuleToSegmentAllocations(segmentAllocs, debugModuleHandle);
     } else {
         for (auto &kernImmData : kernelImmDatas) {
             if (kernImmData->getKernelInfo()->kernelDescriptor.external.debugData.get()) {

level_zero/core/source/module/module_imp.h

@@ -85,7 +85,12 @@ struct ModuleImp : public Module {
     ~ModuleImp() override;
 
     ze_result_t destroy() override {
+        auto tempHandle = debugModuleHandle;
+        auto tempDevice = device;
         delete this;
+        if (tempDevice->getL0Debugger()) {
+            tempDevice->getL0Debugger()->removeZebinModule(tempHandle);
+        }
         return ZE_RESULT_SUCCESS;
     }
 
@@ -157,6 +162,7 @@ struct ModuleImp : public Module {
     ModuleType type;
     NEO::Linker::UnresolvedExternals unresolvedExternalsInfo{};
     std::set<NEO::GraphicsAllocation *> importedSymbolAllocations{};
+    uint32_t debugModuleHandle = 0;
 };
 
 bool moveBuildOption(std::string &dstOptionsSet, std::string &srcOptionSet, NEO::ConstStringRef dstOptionName, NEO::ConstStringRef srcOptionName);

level_zero/core/test/unit_tests/mocks/mock_l0_debugger.h

@@ -49,11 +49,28 @@ class MockDebuggerL0Hw : public L0::DebuggerL0Hw<GfxFamily> {
         L0::DebuggerL0Hw<GfxFamily>::registerElf(debugData, isaAllocation);
     }
 
+    bool attachZebinModuleToSegmentAllocations(const StackVec<NEO::GraphicsAllocation *, 32> &allocs, uint32_t &moduleHandle) override {
+        segmentCountWithAttachedModuleHandle = static_cast<uint32_t>(allocs.size());
+        if (std::numeric_limits<uint32_t>::max() != moduleHandleToReturn) {
+            moduleHandle = moduleHandleToReturn;
+            return true;
+        }
+        return L0::DebuggerL0Hw<GfxFamily>::attachZebinModuleToSegmentAllocations(allocs, moduleHandle);
+    }
+    bool removeZebinModule(uint32_t moduleHandle) override {
+        removedZebinModuleHandle = moduleHandle;
+        return L0::DebuggerL0Hw<GfxFamily>::removeZebinModule(moduleHandle);
+    }
+
     uint32_t captureStateBaseAddressCount = 0;
     uint32_t programSbaTrackingCommandsCount = 0;
     uint32_t getSbaTrackingCommandsSizeCount = 0;
     uint32_t registerElfCount = 0;
     const char *lastReceivedElf = nullptr;
+
+    uint32_t segmentCountWithAttachedModuleHandle = 0;
+    uint32_t removedZebinModuleHandle = 0;
+    uint32_t moduleHandleToReturn = std::numeric_limits<uint32_t>::max();
 };
 
 template <uint32_t productFamily, typename GfxFamily>

level_zero/core/test/unit_tests/mocks/mock_module.h

@@ -74,6 +74,7 @@ struct MockModuleTranslationUnit : public L0::ModuleTranslationUnit {
 
 struct MockModule : public L0::ModuleImp {
     using ModuleImp::debugEnabled;
+    using ModuleImp::debugModuleHandle;
     using ModuleImp::kernelImmDatas;
     using ModuleImp::translationUnit;
 

level_zero/core/test/unit_tests/sources/debugger/linux/test_l0_debugger_linux.cpp

@@ -14,6 +14,7 @@
 
 #include "level_zero/core/test/unit_tests/sources/debugger/l0_debugger_fixture.h"
 
+#include <algorithm>
 #include <memory>
 
 using namespace NEO;
@@ -177,5 +178,46 @@ TEST_F(L0DebuggerLinuxTest, givenNoOSInterfaceThenRegisterElfDoesNothing) {
     neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[0]->osInterface.reset(OSInterface_tmp);
 }
 
+TEST_F(L0DebuggerLinuxTest, givenAllocationsWhenAttachingZebinModuleThenAllAllocationsHaveRegisteredHandle) {
+    MockDrmAllocation isaAllocation(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockBufferObject bo(drmMock, 0, 0, 1);
+    isaAllocation.bufferObjects[0] = &bo;
+
+    MockDrmAllocation isaAllocation2(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockBufferObject bo2(drmMock, 0, 0, 1);
+    isaAllocation2.bufferObjects[0] = &bo2;
+
+    uint32_t handle = 0;
+
+    StackVec<NEO::GraphicsAllocation *, 32> kernelAllocs;
+    kernelAllocs.push_back(&isaAllocation);
+    kernelAllocs.push_back(&isaAllocation2);
+
+    drmMock->registeredDataSize = 0;
+    drmMock->registeredClass = NEO::Drm::ResourceClass::MaxSize;
+
+    EXPECT_TRUE(device->getL0Debugger()->attachZebinModuleToSegmentAllocations(kernelAllocs, handle));
+
+    EXPECT_EQ(sizeof(uint32_t), drmMock->registeredDataSize);
+    EXPECT_EQ(NEO::Drm::ResourceClass::L0ZebinModule, drmMock->registeredClass);
+
+    const auto containsModuleHandle = [handle](const auto &bufferObject) {
+        const auto &bindExtHandles = bufferObject.getBindExtHandles();
+        return std::find(bindExtHandles.begin(), bindExtHandles.end(), handle) != bindExtHandles.end();
+    };
+
+    EXPECT_TRUE(containsModuleHandle(bo));
+    EXPECT_TRUE(containsModuleHandle(bo2));
+}
+
+TEST_F(L0DebuggerLinuxTest, givenModuleHandleWhenRemoveZebinModuleIsCalledThenHandleIsUnregistered) {
+    uint32_t handle = 20;
+
+    EXPECT_TRUE(device->getL0Debugger()->removeZebinModule(handle));
+
+    EXPECT_EQ(1u, drmMock->unregisterCalledCount);
+    EXPECT_EQ(20u, drmMock->unregisteredHandle);
+}
+
 } // namespace ult
 } // namespace L0

level_zero/core/test/unit_tests/sources/debugger/test_module_with_debug.cpp

@@ -600,6 +600,50 @@ HWTEST_F(ModuleWithZebinAndL0DebuggerTest, GivenZebinNoDebugDataWhenInitializing
     EXPECT_EQ(0u, getMockDebuggerL0Hw<FamilyType>()->registerElfCount);
 }
 
+HWTEST_F(ModuleWithZebinAndL0DebuggerTest, GivenZebinWhenModuleIsInitializedAndDestroyedThenModuleHandleIsAttachedAndRemoved) {
+    NEO::MockCompilerEnableGuard mock(true);
+    auto cip = new NEO::MockCompilerInterfaceCaptureBuildOptions();
+    neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[neoDevice->getRootDeviceIndex()]->compilerInterface.reset(cip);
+    uint8_t binary[10];
+    ze_module_desc_t moduleDesc = {};
+    moduleDesc.format = ZE_MODULE_FORMAT_IL_SPIRV;
+    moduleDesc.pInputModule = binary;
+    moduleDesc.inputSize = 10;
+
+    uint32_t kernelHeap = 0;
+    auto kernelInfo = std::make_unique<KernelInfo>();
+    kernelInfo->heapInfo.KernelHeapSize = 1;
+    kernelInfo->heapInfo.pKernelHeap = &kernelHeap;
+
+    auto kernelImmutableData = ::std::make_unique<KernelImmutableData>(device);
+    kernelImmutableData->initialize(kernelInfo.get(), device, 0, nullptr, nullptr, false);
+    std::unique_ptr<MockModule> moduleMock = std::make_unique<MockModule>(device, nullptr, ModuleType::User);
+    moduleMock->translationUnit = std::make_unique<MockModuleTranslationUnit>(device);
+    moduleMock->kernelImmDatas.push_back(std::move(kernelImmutableData));
+
+    auto zebin = ZebinTestData::ValidEmptyProgram();
+    moduleMock->translationUnit = std::make_unique<MockModuleTranslationUnit>(device);
+    moduleMock->translationUnit->unpackedDeviceBinarySize = zebin.storage.size();
+    moduleMock->translationUnit->unpackedDeviceBinary.reset(new char[zebin.storage.size()]);
+    memcpy_s(moduleMock->translationUnit->unpackedDeviceBinary.get(), moduleMock->translationUnit->unpackedDeviceBinarySize,
+             zebin.storage.data(), zebin.storage.size());
+
+    getMockDebuggerL0Hw<FamilyType>()->moduleHandleToReturn = 6;
+    EXPECT_TRUE(moduleMock->initialize(&moduleDesc, neoDevice));
+
+    auto expectedSegmentAllocationCount = 1u;
+    expectedSegmentAllocationCount += moduleMock->translationUnit->globalConstBuffer != nullptr ? 1 : 0;
+    expectedSegmentAllocationCount += moduleMock->translationUnit->globalVarBuffer != nullptr ? 1 : 0;
+
+    EXPECT_EQ(expectedSegmentAllocationCount, getMockDebuggerL0Hw<FamilyType>()->segmentCountWithAttachedModuleHandle);
+    EXPECT_EQ(getMockDebuggerL0Hw<FamilyType>()->moduleHandleToReturn, moduleMock->debugModuleHandle);
+
+    moduleMock->destroy();
+    moduleMock.release();
+
+    EXPECT_EQ(6u, getMockDebuggerL0Hw<FamilyType>()->removedZebinModuleHandle);
+}
+
 using NotifyModuleLoadTest = Test<ModuleFixture>;
 
 HWTEST_F(NotifyModuleLoadTest, givenDebuggingEnabledWhenModuleIsCreatedAndFullyLinkedThenIsaAllocationsAreCopiedAndResident) {