Fixing residency of extern device functions

Change-Id: Icad696cbf6fb3fc0276f0d0d488bf92091525d9b

Author: jchodor

level_zero/core/source/kernel/kernel.h

@@ -41,6 +41,10 @@ struct KernelImmutableData {
         return residencyContainer;
     }
 
+    std::vector<NEO::GraphicsAllocation *> &getResidencyContainer() {
+        return residencyContainer;
+    }
+
     uint32_t getIsaSize() const;
     NEO::GraphicsAllocation *getIsaGraphicsAllocation() const { return isaGraphicsAllocation.get(); }
 

level_zero/core/source/kernel/kernel_imp.cpp

@@ -698,12 +698,8 @@ ze_result_t KernelImp::initialize(const ze_kernel_desc_t *desc) {
 
     this->setDebugSurface();
 
-    for (auto &alloc : kernelImmData->getResidencyContainer()) {
-        residencyContainer.push_back(alloc);
-    }
-    for (auto &alloc : module->getImportedSymbolAllocations()) {
-        residencyContainer.push_back(alloc);
-    }
+    residencyContainer.insert(residencyContainer.end(), kernelImmData->getResidencyContainer().begin(),
+                              kernelImmData->getResidencyContainer().end());
 
     return ZE_RESULT_SUCCESS;
 }

level_zero/core/source/module/module.h

@@ -45,7 +45,6 @@ struct Module : _ze_module_handle_t {
     virtual const std::vector<std::unique_ptr<KernelImmutableData>> &getKernelImmutableDataVector() const = 0;
     virtual uint32_t getMaxGroupSize() const = 0;
     virtual bool isDebugEnabled() const = 0;
-    virtual const std::set<NEO::GraphicsAllocation *> &getImportedSymbolAllocations() const = 0;
 
     Module() = default;
     Module(const Module &) = delete;

level_zero/core/source/module/module_imp.cpp

@@ -481,6 +481,16 @@ bool ModuleImp::linkBinary() {
     }
     DBG_LOG(PrintRelocations, NEO::constructRelocationsDebugMessage(this->symbols));
     isFullyLinked = true;
+    for (auto &kernImmData : this->kernelImmDatas) {
+        kernImmData->getResidencyContainer().reserve(kernImmData->getResidencyContainer().size() +
+                                                     ((this->exportedFunctionsSurface != nullptr) ? 1 : 0) + this->importedSymbolAllocations.size());
+
+        if (nullptr != this->exportedFunctionsSurface) {
+            kernImmData->getResidencyContainer().push_back(this->exportedFunctionsSurface);
+        }
+        kernImmData->getResidencyContainer().insert(kernImmData->getResidencyContainer().end(), this->importedSymbolAllocations.begin(),
+                                                    this->importedSymbolAllocations.end());
+    }
     return true;
 }
 

level_zero/core/source/module/module_imp.h

@@ -114,8 +114,6 @@ struct ModuleImp : public Module {
         return this->translationUnit.get();
     }
 
-    const std::set<NEO::GraphicsAllocation *> &getImportedSymbolAllocations() const override { return importedSymbolAllocations; }
-
   protected:
     void copyPatchedSegments(const NEO::Linker::PatchableSegments &isaSegmentsForPatching);
     void verifyDebugCapabilities();

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

@@ -23,6 +23,7 @@ struct WhiteBox<::L0::Module> : public ::L0::ModuleImp {
     using BaseClass = ::L0::ModuleImp;
     using BaseClass::BaseClass;
     using BaseClass::device;
+    using BaseClass::exportedFunctionsSurface;
     using BaseClass::isFullyLinked;
     using BaseClass::kernelImmDatas;
     using BaseClass::symbols;
@@ -34,7 +35,6 @@ using Module = WhiteBox<::L0::Module>;
 
 template <>
 struct Mock<Module> : public Module {
-    Mock() = delete;
     Mock(::L0::Device *device, ModuleBuildLog *moduleBuildLog);
 
     MOCK_METHOD(ze_result_t, createKernel, (const ze_kernel_desc_t *desc, ze_kernel_handle_t *phFunction), (override));

level_zero/core/test/unit_tests/sources/module/test_module.cpp

@@ -9,6 +9,7 @@
 #include "shared/test/unit_test/device_binary_format/zebin_tests.h"
 
 #include "opencl/source/program/kernel_info.h"
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
 #include "test.h"
 
 #include "level_zero/core/source/context/context.h"
@@ -180,6 +181,17 @@ HWTEST_F(ModuleSpecConstantsTests, givenSpecializationConstantsSetInDescriptorTh
 
 using ModuleLinkingTest = Test<DeviceFixture>;
 
+HWTEST_F(ModuleLinkingTest, whenExternFunctionsAllocationIsPresentThenItsBeingAddedToResidencyContainer) {
+    Mock<Module> module(device, nullptr);
+    MockGraphicsAllocation alloc;
+    module.exportedFunctionsSurface = &alloc;
+    module.kernelImmDatas.push_back(std::make_unique<L0::KernelImmutableData>());
+    module.translationUnit->programInfo.linkerInput.reset(new NEO::LinkerInput);
+    module.linkBinary();
+    ASSERT_EQ(1U, module.kernelImmDatas[0]->getResidencyContainer().size());
+    EXPECT_EQ(&alloc, module.kernelImmDatas[0]->getResidencyContainer()[0]);
+}
+
 HWTEST_F(ModuleLinkingTest, givenFailureDuringLinkingWhenCreatingModuleThenModuleInitialiationFails) {
     auto mockCompiler = new MockCompilerInterface();
     auto rootDeviceEnvironment = neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[0].get();