From 823586b6fbbdcf3ed3480a46acdf7a1879c50e60 Mon Sep 17 00:00:00 2001
From: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>
Date: Tue, 7 Jul 2020 17:28:56 -0700
Subject: [PATCH] Correct address in setBufferSurfaceState

Change-Id: I17c6766f63aa2cde526db297b8334428a29ec85f
Signed-off: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>
---
 level_zero/core/source/kernel/kernel_hw.h     |  2 +-
 .../unit_tests/sources/module/test_module.cpp | 42 +++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/level_zero/core/source/kernel/kernel_hw.h b/level_zero/core/source/kernel/kernel_hw.h
index 9a173d001e..5a41014f9d 100644
--- a/level_zero/core/source/kernel/kernel_hw.h
+++ b/level_zero/core/source/kernel/kernel_hw.h
@@ -25,7 +25,7 @@ struct KernelHw : public KernelImp {
     using GfxFamily = typename NEO::GfxFamilyMapper<gfxCoreFamily>::GfxFamily;
 
     void setBufferSurfaceState(uint32_t argIndex, void *address, NEO::GraphicsAllocation *alloc) override {
-        uintptr_t baseAddress = static_cast<uintptr_t>(alloc->getGpuAddress());
+        uint64_t baseAddress = castToUint64(address);
         auto sshAlignmentMask = NEO::EncodeSurfaceState<GfxFamily>::getSurfaceBaseAddressAlignmentMask();
 
         // Remove misalligned bytes, accounted for in in bufferOffset patch token
diff --git a/level_zero/core/test/unit_tests/sources/module/test_module.cpp b/level_zero/core/test/unit_tests/sources/module/test_module.cpp
index acc0993faf..658911b027 100644
--- a/level_zero/core/test/unit_tests/sources/module/test_module.cpp
+++ b/level_zero/core/test/unit_tests/sources/module/test_module.cpp
@@ -7,6 +7,7 @@
 
 #include "test.h"
 
+#include "level_zero/core/source/kernel/kernel_imp.h"
 #include "level_zero/core/source/module/module_imp.h"
 #include "level_zero/core/test/unit_tests/fixtures/module_fixture.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_module.h"
@@ -43,6 +44,47 @@ HWTEST_F(ModuleTest, givenKernelCreateReturnsSuccess) {
     Kernel::fromHandle(kernelHandle)->destroy();
 }
 
+using ModuleTestSupport = IsWithinProducts<IGFX_SKYLAKE, IGFX_TIGERLAKE_LP>;
+
+HWTEST2_F(ModuleTest, givenNonPatchedTokenThenSurfaceBaseAddressIsCorrectlySet, ModuleTestSupport) {
+    using RENDER_SURFACE_STATE = typename FamilyType::RENDER_SURFACE_STATE;
+    ze_kernel_handle_t kernelHandle;
+
+    ze_kernel_desc_t kernelDesc = {};
+    kernelDesc.version = ZE_KERNEL_DESC_VERSION_CURRENT;
+    kernelDesc.flags = ZE_KERNEL_FLAG_NONE;
+    kernelDesc.pKernelName = kernelName.c_str();
+
+    ze_result_t res = module->createKernel(&kernelDesc, &kernelHandle);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    auto kernelImp = reinterpret_cast<L0::KernelImp *>(L0::Kernel::fromHandle(kernelHandle));
+
+    void *devicePtr = nullptr;
+    res = device->getDriverHandle()->allocDeviceMem(device->toHandle(),
+                                                    ZE_DEVICE_MEM_ALLOC_FLAG_DEFAULT,
+                                                    16384u,
+                                                    0u,
+                                                    &devicePtr);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, res);
+
+    auto gpuAlloc = device->getDriverHandle()->getSvmAllocsManager()->getSVMAllocs()->get(devicePtr)->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
+    ASSERT_NE(nullptr, gpuAlloc);
+
+    uint32_t argIndex = 0u;
+    kernelImp->setBufferSurfaceState(argIndex, devicePtr, gpuAlloc);
+
+    auto argInfo = kernelImp->getImmutableData()->getDescriptor().payloadMappings.explicitArgs[argIndex].as<NEO::ArgDescPointer>();
+    auto surfaceStateAddressRaw = ptrOffset(kernelImp->getSurfaceStateHeapData(), argInfo.bindful);
+    auto surfaceStateAddress = reinterpret_cast<RENDER_SURFACE_STATE *>(const_cast<unsigned char *>(surfaceStateAddressRaw));
+    EXPECT_EQ(reinterpret_cast<void *>(surfaceStateAddress->getSurfaceBaseAddress()), devicePtr);
+
+    Kernel::fromHandle(kernelHandle)->destroy();
+
+    device->getDriverHandle()->freeMem(devicePtr);
+}
+
 HWTEST_F(ModuleTest, givenKernelCreateWithIncorrectKernelNameReturnsFailure) {
     ze_kernel_handle_t kernelHandle;