Handle implicit arg relocation in L0 module

Related-To: NEO-5081
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>

level_zero/core/source/module/module_imp.cpp

@@ -500,11 +500,12 @@ void ModuleImp::copyPatchedSegments(const NEO::Linker::PatchableSegments &isaSeg
 
 bool ModuleImp::linkBinary() {
     using namespace NEO;
-    if (this->translationUnit->programInfo.linkerInput == nullptr) {
+    auto linkerInput = this->translationUnit->programInfo.linkerInput.get();
+    if (linkerInput == nullptr) {
         isFullyLinked = true;
         return true;
     }
-    Linker linker(*this->translationUnit->programInfo.linkerInput);
+    Linker linker(*linkerInput);
     Linker::SegmentInfo globals;
     Linker::SegmentInfo constants;
     Linker::SegmentInfo exportedFunctions;
@@ -518,15 +519,15 @@ bool ModuleImp::linkBinary() {
         constants.gpuAddress = static_cast<uintptr_t>(constantsForPatching->getGpuAddress());
         constants.segmentSize = constantsForPatching->getUnderlyingBufferSize();
     }
-    if (this->translationUnit->programInfo.linkerInput->getExportedFunctionsSegmentId() >= 0) {
-        auto exportedFunctionHeapId = this->translationUnit->programInfo.linkerInput->getExportedFunctionsSegmentId();
+    if (linkerInput->getExportedFunctionsSegmentId() >= 0) {
+        auto exportedFunctionHeapId = linkerInput->getExportedFunctionsSegmentId();
         this->exportedFunctionsSurface = this->kernelImmDatas[exportedFunctionHeapId]->getIsaGraphicsAllocation();
         exportedFunctions.gpuAddress = static_cast<uintptr_t>(exportedFunctionsSurface->getGpuAddressToPatch());
         exportedFunctions.segmentSize = exportedFunctionsSurface->getUnderlyingBufferSize();
     }
     Linker::PatchableSegments isaSegmentsForPatching;
     std::vector<std::vector<char>> patchedIsaTempStorage;
-    if (this->translationUnit->programInfo.linkerInput->getTraits().requiresPatchingOfInstructionSegments) {
+    if (linkerInput->getTraits().requiresPatchingOfInstructionSegments) {
         patchedIsaTempStorage.reserve(this->kernelImmDatas.size());
         for (const auto &kernelInfo : this->translationUnit->programInfo.kernelInfos) {
             auto &kernHeapInfo = kernelInfo->heapInfo;
@@ -558,7 +559,9 @@ bool ModuleImp::linkBinary() {
     }
     DBG_LOG(PrintRelocations, NEO::constructRelocationsDebugMessage(this->symbols));
     isFullyLinked = true;
-    for (auto &kernImmData : this->kernelImmDatas) {
+    for (auto kernelId = 0u; kernelId < kernelImmDatas.size(); kernelId++) {
+        auto &kernImmData = kernelImmDatas[kernelId];
+
         kernImmData->getResidencyContainer().reserve(kernImmData->getResidencyContainer().size() +
                                                      ((this->exportedFunctionsSurface != nullptr) ? 1 : 0) + this->importedSymbolAllocations.size());
 
@@ -567,6 +570,9 @@ bool ModuleImp::linkBinary() {
         }
         kernImmData->getResidencyContainer().insert(kernImmData->getResidencyContainer().end(), this->importedSymbolAllocations.begin(),
                                                     this->importedSymbolAllocations.end());
+
+        auto &kernelDescriptor = const_cast<KernelDescriptor &>(kernImmData->getDescriptor());
+        kernelDescriptor.kernelAttributes.flags.requiresImplicitArgs = linkerInput->areImplicitArgsRequired(kernelId);
     }
     return true;
 }

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

@@ -7,6 +7,7 @@
 
 #include "shared/source/gmm_helper/gmm.h"
 #include "shared/source/gmm_helper/gmm_helper.h"
+#include "shared/source/kernel/implicit_args.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/mocks/mock_elf.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
@@ -640,7 +641,12 @@ HWTEST_F(ModuleLinkingTest, whenExternFunctionsAllocationIsPresentThenItsBeingAd
     Mock<Module> module(device, nullptr);
     MockGraphicsAllocation alloc;
     module.exportedFunctionsSurface = &alloc;
-    module.kernelImmDatas.push_back(std::make_unique<L0::KernelImmutableData>());
+
+    KernelInfo kernelInfo{};
+    std::unique_ptr<WhiteBox<::L0::KernelImmutableData>> kernelImmData{new WhiteBox<::L0::KernelImmutableData>(this->device)};
+    kernelImmData->initialize(&kernelInfo, device, 0, nullptr, nullptr, false);
+    module.kernelImmDatas.push_back(std::move(kernelImmData));
+
     module.translationUnit->programInfo.linkerInput.reset(new NEO::LinkerInput);
     module.linkBinary();
     ASSERT_EQ(1U, module.kernelImmDatas[0]->getResidencyContainer().size());
@@ -1575,5 +1581,34 @@ TEST_F(ModuleTests, whenCopyingPatchedSegmentsThenAllocationsAreSetWritableForTb
     EXPECT_TRUE(allocation->isTbxWritable(std::numeric_limits<uint32_t>::max()));
     EXPECT_TRUE(allocation->isAubWritable(std::numeric_limits<uint32_t>::max()));
 }
+
+TEST_F(ModuleTests, givenImplicitArgsRelocationWhenLinkingModuleThenSegmentIsPatchedAndImplicitArgsAreRequired) {
+    auto pModule = std::make_unique<Module>(device, nullptr, ModuleType::User);
+
+    char data[64]{};
+    auto kernelInfo = new KernelInfo();
+    kernelInfo->heapInfo.KernelHeapSize = 64;
+    kernelInfo->heapInfo.pKernelHeap = data;
+
+    std::unique_ptr<WhiteBox<::L0::KernelImmutableData>> kernelImmData{new WhiteBox<::L0::KernelImmutableData>(this->device)};
+    kernelImmData->initialize(kernelInfo, device, 0, nullptr, nullptr, false);
+
+    auto isaCpuPtr = reinterpret_cast<char *>(kernelImmData->isaGraphicsAllocation->getUnderlyingBuffer());
+    pModule->kernelImmDatas.push_back(std::move(kernelImmData));
+    pModule->translationUnit->programInfo.kernelInfos.push_back(kernelInfo);
+    auto linkerInput = std::make_unique<::WhiteBox<NEO::LinkerInput>>();
+    linkerInput->traits.requiresPatchingOfInstructionSegments = true;
+    linkerInput->relocations.push_back({{implicitArgsRelocationSymbolName, 0x8, LinkerInput::RelocationInfo::Type::AddressLow, SegmentType::Instructions}});
+    pModule->translationUnit->programInfo.linkerInput = std::move(linkerInput);
+
+    EXPECT_FALSE(kernelInfo->kernelDescriptor.kernelAttributes.flags.requiresImplicitArgs);
+    auto status = pModule->linkBinary();
+    EXPECT_TRUE(status);
+
+    EXPECT_EQ(sizeof(ImplicitArgs), *reinterpret_cast<uint32_t *>(ptrOffset(isaCpuPtr, 0x8)));
+
+    EXPECT_TRUE(kernelInfo->kernelDescriptor.kernelAttributes.flags.requiresImplicitArgs);
+}
+
 } // namespace ult
 } // namespace L0

opencl/source/program/process_device_binary.cpp

@@ -10,7 +10,6 @@
 #include "shared/source/helpers/debug_helpers.h"
 #include "shared/source/helpers/ptr_math.h"
 #include "shared/source/helpers/string.h"
-#include "shared/source/kernel/implicit_args.h"
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
 #include "shared/source/program/program_info.h"
@@ -122,10 +121,7 @@ cl_int Program::linkBinary(Device *pDevice, const void *constantsInitData, const
             MemoryTransferHelper::transferMemoryToAllocation(hwHelper.isBlitCopyRequiredForLocalMemory(hwInfo, *kernelInfo->getGraphicsAllocation()),
                                                              *pDevice, kernelInfo->getGraphicsAllocation(), 0, isaSegmentsForPatching[segmentId].hostPointer,
                                                              static_cast<size_t>(kernHeapInfo.KernelHeapSize));
-            if (linkerInput->getRelocationsInInstructionSegments().size() > kernelId) {
-                const auto &relocations = linkerInput->getRelocationsInInstructionSegments()[kernelId];
-                kernelInfo->kernelDescriptor.kernelAttributes.flags.requiresImplicitArgs = (relocations.end() != std::find_if(relocations.begin(), relocations.end(), [&](const auto &relocation) { return relocation.symbolName == implicitArgsRelocationSymbolName; }));
-            }
+            kernelInfo->kernelDescriptor.kernelAttributes.flags.requiresImplicitArgs = linkerInput->areImplicitArgsRequired(kernelId);
         }
     }
     DBG_LOG(PrintRelocations, NEO::constructRelocationsDebugMessage(this->getSymbols(pDevice->getRootDeviceIndex())));

shared/source/compiler_interface/linker.cpp

@@ -458,4 +458,12 @@ void Linker::applyDebugDataRelocations(const NEO::Elf::Elf<NEO::Elf::EI_CLASS_64
     }
 }
 
+bool LinkerInput::areImplicitArgsRequired(uint32_t instructionsSegmentId) const {
+    if (relocations.size() > instructionsSegmentId) {
+        const auto &segmentRelocations = relocations[instructionsSegmentId];
+        return (segmentRelocations.end() != std::find_if(segmentRelocations.begin(), segmentRelocations.end(), [&](const auto &relocation) { return relocation.symbolName == implicitArgsRelocationSymbolName; }));
+    }
+    return false;
+}
+
 } // namespace NEO

shared/source/compiler_interface/linker.h

@@ -137,6 +137,7 @@ struct LinkerInput {
     bool isValid() const {
         return valid;
     }
+    bool areImplicitArgsRequired(uint32_t instructionsSegmentId) const;
 
     bool undefinedSymbolsAllowed = false;
 

shared/test/unit_test/compiler_interface/linker_tests.cpp

@@ -1875,3 +1875,22 @@ TEST(LinkerTests, givenImplicitArgRelocationThenPatchRelocationWithSizeOfImplici
     EXPECT_EQ(initData, *(addressToPatch - 1));
     EXPECT_EQ(initData, *(addressToPatch + 1));
 }
+
+TEST(LinkerTests, givenImplicitArgRelocationThenImplicitArgsAreRequired) {
+    NEO::LinkerInput linkerInput;
+
+    EXPECT_FALSE(linkerInput.areImplicitArgsRequired(0u));
+
+    vISA::GenRelocEntry reloc = {};
+    std::string relocationName = implicitArgsRelocationSymbolName;
+    memcpy_s(reloc.r_symbol, 1024, relocationName.c_str(), relocationName.size());
+    reloc.r_offset = 8;
+    reloc.r_type = vISA::GenRelocType::R_SYM_ADDR_32;
+
+    vISA::GenRelocEntry relocs[] = {reloc};
+    constexpr uint32_t numRelocations = 1;
+    bool decodeRelocSuccess = linkerInput.decodeRelocationTable(&relocs, numRelocations, 0);
+    EXPECT_TRUE(decodeRelocSuccess);
+
+    EXPECT_TRUE(linkerInput.areImplicitArgsRequired(0u));
+}