From 451f1e7e0de702aa2826527ec3cef597012b1909 Mon Sep 17 00:00:00 2001
From: Jaroslaw Chodor <jaroslaw.chodor@intel.com>
Date: Wed, 18 Nov 2020 12:01:27 +0000
Subject: [PATCH] Fixing residency of global const/var buffers

Signed-off-by: Jaroslaw Chodor <jaroslaw.chodor@intel.com>
---
 level_zero/core/source/kernel/kernel_imp.cpp  |  4 ++++
 .../unit_tests/sources/kernel/test_kernel.cpp | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/level_zero/core/source/kernel/kernel_imp.cpp b/level_zero/core/source/kernel/kernel_imp.cpp
index 44466c9acf..447df9f26d 100644
--- a/level_zero/core/source/kernel/kernel_imp.cpp
+++ b/level_zero/core/source/kernel/kernel_imp.cpp
@@ -178,6 +178,8 @@ void KernelImmutableData::initialize(NEO::KernelInfo *kernelInfo, NEO::MemoryMan
                                  static_cast<uintptr_t>(globalConstBuffer->getGpuAddressToPatch()),
                                  *globalConstBuffer, kernelDescriptor->payloadMappings.implicitArgs.globalConstantsSurfaceAddress, *device);
         this->residencyContainer.push_back(globalConstBuffer);
+    } else if (nullptr != globalConstBuffer) {
+        this->residencyContainer.push_back(globalConstBuffer);
     }
 
     if (NEO::isValidOffset(kernelDescriptor->payloadMappings.implicitArgs.globalVariablesSurfaceAddress.stateless)) {
@@ -187,6 +189,8 @@ void KernelImmutableData::initialize(NEO::KernelInfo *kernelInfo, NEO::MemoryMan
                                  static_cast<uintptr_t>(globalVarBuffer->getGpuAddressToPatch()),
                                  *globalVarBuffer, kernelDescriptor->payloadMappings.implicitArgs.globalVariablesSurfaceAddress, *device);
         this->residencyContainer.push_back(globalVarBuffer);
+    } else if (nullptr != globalVarBuffer) {
+        this->residencyContainer.push_back(globalVarBuffer);
     }
 }
 
diff --git a/level_zero/core/test/unit_tests/sources/kernel/test_kernel.cpp b/level_zero/core/test/unit_tests/sources/kernel/test_kernel.cpp
index 108505fd1b..e6268f2c36 100644
--- a/level_zero/core/test/unit_tests/sources/kernel/test_kernel.cpp
+++ b/level_zero/core/test/unit_tests/sources/kernel/test_kernel.cpp
@@ -6,6 +6,7 @@
  */
 
 #include "shared/test/unit_test/mocks/mock_device.h"
+#include "shared/test/unit_test/mocks/mock_graphics_allocation.h"
 
 #include "opencl/source/program/kernel_info.h"
 #include "test.h"
@@ -551,5 +552,26 @@ TEST_F(KernelIsaTests, givenKernelInfoWhenInitializingImmutableDataWithNonIntern
     EXPECT_EQ(NEO::GraphicsAllocation::AllocationType::KERNEL_ISA, kernelImmutableData.getIsaGraphicsAllocation()->getAllocationType());
 }
 
+TEST_F(KernelIsaTests, givenGlobalBuffersWhenCreatingKernelImmutableDataThenBuffersAreAddedToResidencyContainer) {
+    uint32_t kernelHeap = 0;
+    KernelInfo kernelInfo;
+    kernelInfo.heapInfo.KernelHeapSize = 1;
+    kernelInfo.heapInfo.pKernelHeap = &kernelHeap;
+
+    KernelImmutableData kernelImmutableData(device);
+
+    uint64_t gpuAddress = 0x1200;
+    void *buffer = reinterpret_cast<void *>(gpuAddress);
+    size_t size = 0x1100;
+    NEO::MockGraphicsAllocation globalVarBuffer(buffer, gpuAddress, size);
+    NEO::MockGraphicsAllocation globalConstBuffer(buffer, gpuAddress, size);
+
+    kernelImmutableData.initialize(&kernelInfo, *device->getNEODevice()->getMemoryManager(), device->getNEODevice(), 0,
+                                   &globalConstBuffer, &globalVarBuffer, false);
+    auto &resCont = kernelImmutableData.getResidencyContainer();
+    EXPECT_EQ(1, std::count(resCont.begin(), resCont.end(), &globalVarBuffer));
+    EXPECT_EQ(1, std::count(resCont.begin(), resCont.end(), &globalConstBuffer));
+}
+
 } // namespace ult
 } // namespace L0