From b81317179730a99e5275f103a514db1ebf4b6388 Mon Sep 17 00:00:00 2001
From: Jim Snow <jim.m.snow@intel.com>
Date: Tue, 14 Dec 2021 06:05:02 +0000
Subject: [PATCH] Add unit test to verify rtDispatchGlobals is correctly
 patched

Related-To: NEO-6529

Signed-off-by: Jim Snow <jim.m.snow@intel.com>
---
 .../test/unit_tests/fixtures/module_fixture.h |  1 +
 .../unit_tests/sources/kernel/test_kernel.cpp | 49 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/level_zero/core/test/unit_tests/fixtures/module_fixture.h b/level_zero/core/test/unit_tests/fixtures/module_fixture.h
index 35a9eaa322..35039c316f 100644
--- a/level_zero/core/test/unit_tests/fixtures/module_fixture.h
+++ b/level_zero/core/test/unit_tests/fixtures/module_fixture.h
@@ -116,6 +116,7 @@ struct ModuleImmutableDataFixture : public DeviceFixture {
     class MockKernel : public WhiteBox<L0::KernelImp> {
       public:
         using KernelImp::crossThreadData;
+        using KernelImp::crossThreadDataSize;
         using KernelImp::kernelArgHandlers;
         using KernelImp::kernelHasIndirectAccess;
         using KernelImp::privateMemoryGraphicsAllocation;
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 26c46615e2..ada991019d 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
@@ -659,6 +659,55 @@ TEST_F(KernelImmutableDataTests, whenHasRTCallsIsTrueAndNoRTDispatchGlobalsIsAll
     EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, kernel->initialize(&kernelDesc));
 }
 
+TEST_F(KernelImmutableDataTests, whenHasRTCallsIsTrueThenCrossThreadDataIsPatched) {
+    KernelDescriptorRTCallsTrue mockDescriptor = {};
+    mockDescriptor.kernelMetadata.kernelName = "rt_test";
+    for (auto i = 0u; i < 3u; i++) {
+        mockDescriptor.kernelAttributes.requiredWorkgroupSize[i] = 0;
+    }
+
+    std::unique_ptr<MockImmutableData> mockKernelImmutableData =
+        std::make_unique<MockImmutableData>(32u);
+    mockKernelImmutableData->kernelDescriptor = &mockDescriptor;
+
+    ModuleBuildLog *moduleBuildLog = nullptr;
+    module = std::make_unique<MockModule>(device,
+                                          moduleBuildLog,
+                                          ModuleType::User,
+                                          32u,
+                                          mockKernelImmutableData.get());
+    module->maxGroupSize = 10;
+
+    std::unique_ptr<ModuleImmutableDataFixture::MockKernel> kernel;
+    kernel = std::make_unique<ModuleImmutableDataFixture::MockKernel>(module.get());
+
+    ze_kernel_desc_t kernelDesc = {};
+    kernelDesc.pKernelName = "rt_test";
+
+    auto immDataVector =
+        const_cast<std::vector<std::unique_ptr<KernelImmutableData>> *>(&module.get()->getKernelImmutableDataVector());
+
+    immDataVector->push_back(std::move(mockKernelImmutableData));
+
+    auto crossThreadData = std::make_unique<uint32_t[]>(4);
+    kernel->crossThreadData.reset(reinterpret_cast<uint8_t *>(crossThreadData.get()));
+    kernel->crossThreadDataSize = sizeof(uint32_t[4]);
+
+    neoDevice->setRTDispatchGlobalsForceAllocation();
+
+    auto result = kernel->initialize(&kernelDesc);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto rtDispatchGlobals = neoDevice->getRTDispatchGlobals(NEO::RayTracingHelper::maxBvhLevels);
+    EXPECT_NE(nullptr, rtDispatchGlobals);
+
+    auto dispatchGlobalsAddressPatched = *reinterpret_cast<uintptr_t *>(crossThreadData.get());
+    auto dispatchGlobalsGpuAddressOffset = static_cast<uintptr_t>(rtDispatchGlobals->getGpuAddressToPatch());
+    EXPECT_EQ(dispatchGlobalsGpuAddressOffset, dispatchGlobalsAddressPatched);
+
+    kernel->crossThreadData.release();
+}
+
 using KernelIndirectPropertiesFromIGCTests = KernelImmutableDataTests;
 
 TEST_F(KernelIndirectPropertiesFromIGCTests, whenInitializingKernelWithNoKernelLoadAndNoStoreAndNoAtomicThenHasIndirectAccessIsSetToFalse) {