/*
 * Copyright (C) 2022 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "shared/source/gmm_helper/resource_info.h"
#include "shared/source/memory_manager/unified_memory_manager.h"
#include "shared/test/common/helpers/debug_manager_state_restore.h"
#include "shared/test/common/test_macros/hw_test.h"

#include "opencl/extensions/public/cl_ext_private.h"
#include "opencl/source/api/api.h"
#include "opencl/source/mem_obj/buffer.h"
#include "opencl/source/platform/platform.h"
#include "opencl/test/unit_test/aub_tests/fixtures/aub_fixture.h"
#include "opencl/test/unit_test/aub_tests/fixtures/multicontext_aub_fixture.h"
#include "opencl/test/unit_test/fixtures/simple_arg_kernel_fixture.h"
#include "opencl/test/unit_test/mocks/mock_kernel.h"

using namespace NEO;

struct StatelessCompressionInSBA : public KernelAUBFixture<StatelessCopyKernelFixture>,
                                   public ::testing::Test,
                                   public ::testing::WithParamInterface<uint32_t /*EngineType*/> {
    void SetUp() override {
        DebugManager.flags.EnableStatelessCompression.set(1);
        DebugManager.flags.RenderCompressedBuffersEnabled.set(true);
        DebugManager.flags.RenderCompressedImagesEnabled.set(true);
        DebugManager.flags.EnableLocalMemory.set(true);
        DebugManager.flags.NodeOrdinal.set(GetParam());
        DebugManager.flags.ForceAuxTranslationMode.set(static_cast<int32_t>(AuxTranslationMode::Builtin));
        KernelAUBFixture<StatelessCopyKernelFixture>::SetUp();
    }

    void TearDown() override {
        KernelAUBFixture<StatelessCopyKernelFixture>::TearDown();
    }

    DebugManagerStateRestore debugRestorer;
    cl_int retVal = CL_SUCCESS;
};

XE_HP_CORE_TEST_P(StatelessCompressionInSBA, GENERATEONLY_givenCompressedBuffersWhenStatelessAccessToLocalMemoryThenEnableCompressionInSBA) {
    const size_t bufferSize = 2048;
    uint8_t writePattern[bufferSize];
    std::fill(writePattern, writePattern + sizeof(writePattern), 1);

    device->getGpgpuCommandStreamReceiver().overrideDispatchPolicy(DispatchMode::BatchedDispatch);

    auto unCompressedBuffer = std::unique_ptr<Buffer>(Buffer::create(context, CL_MEM_UNCOMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto unCompressedAllocation = unCompressedBuffer->getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_EQ(AllocationType::BUFFER, unCompressedAllocation->getAllocationType());
    EXPECT_FALSE(unCompressedAllocation->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, unCompressedAllocation->getMemoryPool());

    auto compressedBuffer1 = std::unique_ptr<Buffer>(Buffer::create(context, CL_MEM_COMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto compressedAllocation1 = compressedBuffer1->getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_TRUE(compressedAllocation1->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedAllocation1->getMemoryPool());
    EXPECT_TRUE(compressedAllocation1->getDefaultGmm()->isCompressionEnabled);

    auto compressedBuffer2 = std::unique_ptr<Buffer>(Buffer::create(context, CL_MEM_COMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto compressedAllocation2 = compressedBuffer2->getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_TRUE(compressedAllocation2->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedAllocation2->getMemoryPool());
    EXPECT_TRUE(compressedAllocation2->getDefaultGmm()->isCompressionEnabled);

    retVal = pCmdQ->enqueueWriteBuffer(compressedBuffer1.get(), CL_FALSE, 0, bufferSize, writePattern, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(0, compressedBuffer1.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(1, compressedBuffer2.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = pCmdQ->enqueueKernel(kernel, 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    pCmdQ->finish();

    expectNotEqualMemory<FamilyType>(AUBFixture::getGpuPointer(compressedAllocation1), writePattern, bufferSize);

    expectNotEqualMemory<FamilyType>(AUBFixture::getGpuPointer(compressedAllocation2), writePattern, bufferSize);

    retVal = pCmdQ->enqueueCopyBuffer(compressedBuffer2.get(), unCompressedBuffer.get(), 0, 0, bufferSize, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    pCmdQ->finish();

    expectMemory<FamilyType>(AUBFixture::getGpuPointer(unCompressedAllocation), writePattern, bufferSize);
}

XE_HP_CORE_TEST_P(StatelessCompressionInSBA, GENERATEONLY_givenCompressedDeviceMemoryWhenAccessedStatelesslyThenEnableCompressionInSBA) {
    const size_t bufferSize = 2048;
    uint8_t writePattern[bufferSize];
    std::fill(writePattern, writePattern + sizeof(writePattern), 1);

    device->getGpgpuCommandStreamReceiver().overrideDispatchPolicy(DispatchMode::BatchedDispatch);

    auto compressedDeviceMemAllocPtr1 = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    EXPECT_NE(nullptr, compressedDeviceMemAllocPtr1);
    auto compressedDeviceMemAlloc1 = context->getSVMAllocsManager()->getSVMAllocs()->get(compressedDeviceMemAllocPtr1)->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_NE(nullptr, compressedDeviceMemAlloc1);
    EXPECT_TRUE(compressedDeviceMemAlloc1->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedDeviceMemAlloc1->getMemoryPool());
    EXPECT_TRUE(compressedDeviceMemAlloc1->getDefaultGmm()->isCompressionEnabled);

    auto compressedDeviceMemAllocPtr2 = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    EXPECT_NE(nullptr, compressedDeviceMemAllocPtr2);
    auto compressedDeviceMemAlloc2 = context->getSVMAllocsManager()->getSVMAllocs()->get(compressedDeviceMemAllocPtr2)->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_NE(nullptr, compressedDeviceMemAlloc2);
    EXPECT_TRUE(compressedDeviceMemAlloc2->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedDeviceMemAlloc2->getMemoryPool());
    EXPECT_TRUE(compressedDeviceMemAlloc2->getDefaultGmm()->isCompressionEnabled);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAllocPtr1, writePattern, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 0, compressedDeviceMemAllocPtr1);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 1, compressedDeviceMemAllocPtr2);
    EXPECT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = clEnqueueNDRangeKernel(pCmdQ, multiDeviceKernel.get(), 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clFinish(pCmdQ);
    EXPECT_EQ(CL_SUCCESS, retVal);

    expectNotEqualMemory<FamilyType>(compressedDeviceMemAllocPtr2, writePattern, bufferSize);

    expectNotEqualMemory<FamilyType>(compressedDeviceMemAllocPtr2, writePattern, bufferSize);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAllocPtr1);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAllocPtr2);
    EXPECT_EQ(CL_SUCCESS, retVal);
}

XE_HP_CORE_TEST_P(StatelessCompressionInSBA, givenUncompressibleBufferInHostMemoryWhenAccessedStatelesslyThenDisableCompressionInSBA) {
    const size_t bufferSize = 2048;
    uint8_t writePattern[bufferSize];
    std::fill(writePattern, writePattern + sizeof(writePattern), 1);

    device->getGpgpuCommandStreamReceiver().overrideDispatchPolicy(DispatchMode::BatchedDispatch);

    auto unCompressedBuffer = std::unique_ptr<Buffer>(Buffer::create(context, CL_MEM_UNCOMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto unCompressedAllocation = unCompressedBuffer->getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_EQ(AllocationType::BUFFER, unCompressedAllocation->getAllocationType());
    EXPECT_FALSE(unCompressedAllocation->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, unCompressedAllocation->getMemoryPool());

    auto compressedBuffer = std::unique_ptr<Buffer>(Buffer::create(context, CL_MEM_COMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto compressedAllocation = compressedBuffer->getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_TRUE(compressedAllocation->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedAllocation->getMemoryPool());
    EXPECT_TRUE(compressedAllocation->getDefaultGmm()->isCompressionEnabled);

    auto uncompressibleBufferInHostMemory = std::unique_ptr<Buffer>(Buffer::create(context, CL_MEM_FORCE_HOST_MEMORY_INTEL, bufferSize, nullptr, retVal));
    auto uncompressibleAllocationInHostMemory = uncompressibleBufferInHostMemory->getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_EQ(AllocationType::BUFFER_HOST_MEMORY, uncompressibleAllocationInHostMemory->getAllocationType());
    EXPECT_TRUE(MemoryPoolHelper::isSystemMemoryPool(uncompressibleAllocationInHostMemory->getMemoryPool()));

    retVal = pCmdQ->enqueueWriteBuffer(compressedBuffer.get(), CL_FALSE, 0, bufferSize, writePattern, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(0, compressedBuffer.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(1, uncompressibleBufferInHostMemory.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = pCmdQ->enqueueKernel(kernel, 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);
    pCmdQ->finish();

    expectNotEqualMemory<FamilyType>(AUBFixture::getGpuPointer(compressedAllocation), writePattern, bufferSize);

    expectMemory<FamilyType>(AUBFixture::getGpuPointer(uncompressibleAllocationInHostMemory), writePattern, bufferSize);

    retVal = pCmdQ->enqueueCopyBuffer(uncompressibleBufferInHostMemory.get(), unCompressedBuffer.get(), 0, 0, bufferSize, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    pCmdQ->finish();

    expectMemory<FamilyType>(AUBFixture::getGpuPointer(unCompressedAllocation), writePattern, bufferSize);
}

XE_HP_CORE_TEST_P(StatelessCompressionInSBA, givenUncompressibleHostMemoryAllocationWhenAccessedStatelesslyThenDisableCompressionInSBA) {
    const size_t bufferSize = 2048;
    uint8_t writePattern[bufferSize];
    std::fill(writePattern, writePattern + sizeof(writePattern), 1);

    device->getGpgpuCommandStreamReceiver().overrideDispatchPolicy(DispatchMode::BatchedDispatch);

    auto compressedDeviceMemAllocPtr = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    EXPECT_NE(nullptr, compressedDeviceMemAllocPtr);
    auto compressedDeviceMemAlloc = context->getSVMAllocsManager()->getSVMAllocs()->get(compressedDeviceMemAllocPtr)->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_NE(nullptr, compressedDeviceMemAlloc);
    EXPECT_TRUE(compressedDeviceMemAlloc->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedDeviceMemAlloc->getMemoryPool());
    EXPECT_TRUE(compressedDeviceMemAlloc->getDefaultGmm()->isCompressionEnabled);

    auto uncompressibleHostMemAllocPtr = clHostMemAllocINTEL(context, nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    EXPECT_NE(nullptr, uncompressibleHostMemAllocPtr);
    auto uncompressibleHostMemAlloc = context->getSVMAllocsManager()->getSVMAllocs()->get(uncompressibleHostMemAllocPtr)->gpuAllocations.getGraphicsAllocation(device->getRootDeviceIndex());
    EXPECT_NE(nullptr, uncompressibleHostMemAlloc);
    EXPECT_EQ(AllocationType::BUFFER_HOST_MEMORY, uncompressibleHostMemAlloc->getAllocationType());
    EXPECT_TRUE(MemoryPoolHelper::isSystemMemoryPool(uncompressibleHostMemAlloc->getMemoryPool()));

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAllocPtr, writePattern, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 0, compressedDeviceMemAllocPtr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 1, uncompressibleHostMemAllocPtr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = clEnqueueNDRangeKernel(pCmdQ, multiDeviceKernel.get(), 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clFinish(pCmdQ);
    EXPECT_EQ(CL_SUCCESS, retVal);

    expectNotEqualMemory<FamilyType>(compressedDeviceMemAllocPtr, writePattern, bufferSize);
    expectMemory<FamilyType>(uncompressibleHostMemAllocPtr, writePattern, bufferSize);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAllocPtr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, uncompressibleHostMemAllocPtr);
    EXPECT_EQ(CL_SUCCESS, retVal);
}

INSTANTIATE_TEST_CASE_P(,
                        StatelessCompressionInSBA,
                        ::testing::Values(aub_stream::ENGINE_RCS,
                                          aub_stream::ENGINE_CCS));

struct UmStatelessCompressionInSBA : public KernelAUBFixture<StatelessKernelWithIndirectAccessFixture>,
                                     public ::testing::Test,
                                     public ::testing::WithParamInterface<uint32_t /*EngineType*/> {
    void SetUp() override {
        DebugManager.flags.EnableStatelessCompression.set(1);
        DebugManager.flags.RenderCompressedBuffersEnabled.set(true);
        DebugManager.flags.RenderCompressedImagesEnabled.set(true);
        DebugManager.flags.EnableLocalMemory.set(true);
        DebugManager.flags.NodeOrdinal.set(GetParam());
        DebugManager.flags.ForceAuxTranslationMode.set(static_cast<int32_t>(AuxTranslationMode::Builtin));
        KernelAUBFixture<StatelessKernelWithIndirectAccessFixture>::SetUp();
        EXPECT_TRUE(multiDeviceKernel->getKernel(rootDeviceIndex)->getKernelInfo().hasIndirectStatelessAccess);
    }

    void TearDown() override {
        KernelAUBFixture<StatelessKernelWithIndirectAccessFixture>::TearDown();
    }

    DebugManagerStateRestore debugRestorer;
    cl_int retVal = CL_SUCCESS;
};

XE_HP_CORE_TEST_P(UmStatelessCompressionInSBA, GENERATEONLY_givenStatelessKernelWhenItHasIndirectDeviceAccessThenEnableCompressionInSBA) {
    const size_t bufferSize = MemoryConstants::kiloByte;
    uint8_t bufferData[bufferSize] = {};

    auto uncompressibleHostMemAlloc = clHostMemAllocINTEL(context, nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, uncompressibleHostMemAlloc);

    auto compressedDeviceMemAlloc1 = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, compressedDeviceMemAlloc1);

    auto compressedDeviceMemAlloc2 = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, compressedDeviceMemAlloc2);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAlloc2, bufferData, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    reinterpret_cast<uint64_t *>(bufferData)[0] = reinterpret_cast<uint64_t>(compressedDeviceMemAlloc2);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAlloc1, bufferData, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 0, compressedDeviceMemAlloc1);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelExecInfo(multiDeviceKernel.get(), CL_KERNEL_EXEC_INFO_USM_PTRS_INTEL, sizeof(compressedDeviceMemAlloc2), &compressedDeviceMemAlloc2);
    EXPECT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = clEnqueueNDRangeKernel(pCmdQ, multiDeviceKernel.get(), 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clFinish(pCmdQ);
    EXPECT_EQ(CL_SUCCESS, retVal);

    reinterpret_cast<uint64_t *>(bufferData)[0] = 1;
    expectNotEqualMemory<FamilyType>(compressedDeviceMemAlloc2, bufferData, bufferSize);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, uncompressibleHostMemAlloc, compressedDeviceMemAlloc2, bufferSize, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    expectMemory<FamilyType>(uncompressibleHostMemAlloc, bufferData, bufferSize);

    retVal = clMemFreeINTEL(context, uncompressibleHostMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAlloc1);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAlloc2);
    EXPECT_EQ(CL_SUCCESS, retVal);
}

XE_HP_CORE_TEST_P(UmStatelessCompressionInSBA, GENERATEONLY_givenKernelExecInfoWhenItHasIndirectDeviceAccessThenEnableCompressionInSBA) {
    const size_t bufferSize = MemoryConstants::kiloByte;
    uint8_t bufferData[bufferSize] = {};

    auto uncompressibleHostMemAlloc = clHostMemAllocINTEL(context, nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, uncompressibleHostMemAlloc);

    auto compressedDeviceMemAlloc1 = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, compressedDeviceMemAlloc1);

    auto compressedDeviceMemAlloc2 = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, compressedDeviceMemAlloc2);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAlloc2, bufferData, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    reinterpret_cast<uint64_t *>(bufferData)[0] = reinterpret_cast<uint64_t>(compressedDeviceMemAlloc2);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAlloc1, bufferData, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 0, compressedDeviceMemAlloc1);
    EXPECT_EQ(CL_SUCCESS, retVal);

    cl_bool enableIndirectDeviceAccess = CL_TRUE;
    retVal = clSetKernelExecInfo(multiDeviceKernel.get(), CL_KERNEL_EXEC_INFO_INDIRECT_DEVICE_ACCESS_INTEL, sizeof(cl_bool), &enableIndirectDeviceAccess);
    EXPECT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = clEnqueueNDRangeKernel(pCmdQ, multiDeviceKernel.get(), 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clFinish(pCmdQ);
    EXPECT_EQ(CL_SUCCESS, retVal);

    reinterpret_cast<uint64_t *>(bufferData)[0] = 1;
    expectNotEqualMemory<FamilyType>(compressedDeviceMemAlloc2, bufferData, bufferSize);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, uncompressibleHostMemAlloc, compressedDeviceMemAlloc2, bufferSize, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    expectMemory<FamilyType>(uncompressibleHostMemAlloc, bufferData, bufferSize);

    retVal = clMemFreeINTEL(context, uncompressibleHostMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAlloc1);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAlloc2);
    EXPECT_EQ(CL_SUCCESS, retVal);
}

XE_HP_CORE_TEST_P(UmStatelessCompressionInSBA, givenStatelessKernelWhenItHasIndirectHostAccessThenDisableCompressionInSBA) {
    const size_t bufferSize = MemoryConstants::kiloByte;
    uint8_t bufferData[bufferSize] = {};

    auto compressedDeviceMemAlloc = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, compressedDeviceMemAlloc);

    auto uncompressibleHostMemAlloc = clHostMemAllocINTEL(context, nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, uncompressibleHostMemAlloc);

    memset(uncompressibleHostMemAlloc, 0, bufferSize);

    reinterpret_cast<uint64_t *>(bufferData)[0] = reinterpret_cast<uint64_t>(uncompressibleHostMemAlloc);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAlloc, bufferData, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 0, compressedDeviceMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelExecInfo(multiDeviceKernel.get(), CL_KERNEL_EXEC_INFO_USM_PTRS_INTEL, sizeof(uncompressibleHostMemAlloc), &uncompressibleHostMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = clEnqueueNDRangeKernel(pCmdQ, multiDeviceKernel.get(), 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clFinish(pCmdQ);
    EXPECT_EQ(CL_SUCCESS, retVal);

    reinterpret_cast<uint64_t *>(bufferData)[0] = 1;
    expectMemory<FamilyType>(uncompressibleHostMemAlloc, bufferData, bufferSize);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, uncompressibleHostMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);
}

XE_HP_CORE_TEST_P(UmStatelessCompressionInSBA, givenKernelExecInfoWhenItHasIndirectHostAccessThenDisableCompressionInSBA) {
    const size_t bufferSize = MemoryConstants::kiloByte;
    uint8_t bufferData[bufferSize] = {};

    auto compressedDeviceMemAlloc = clDeviceMemAllocINTEL(context, device.get(), nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, compressedDeviceMemAlloc);

    auto uncompressibleHostMemAlloc = clHostMemAllocINTEL(context, nullptr, bufferSize, 0, &retVal);
    EXPECT_EQ(CL_SUCCESS, retVal);
    ASSERT_NE(nullptr, uncompressibleHostMemAlloc);

    memset(uncompressibleHostMemAlloc, 0, bufferSize);

    reinterpret_cast<uint64_t *>(bufferData)[0] = reinterpret_cast<uint64_t>(uncompressibleHostMemAlloc);

    retVal = clEnqueueMemcpyINTEL(pCmdQ, true, compressedDeviceMemAlloc, bufferData, bufferSize, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clSetKernelArgSVMPointer(multiDeviceKernel.get(), 0, compressedDeviceMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    cl_bool enableIndirectHostAccess = CL_TRUE;
    retVal = clSetKernelExecInfo(multiDeviceKernel.get(), CL_KERNEL_EXEC_INFO_INDIRECT_HOST_ACCESS_INTEL, sizeof(cl_bool), &enableIndirectHostAccess);
    EXPECT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = clEnqueueNDRangeKernel(pCmdQ, multiDeviceKernel.get(), 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clFinish(pCmdQ);
    EXPECT_EQ(CL_SUCCESS, retVal);

    reinterpret_cast<uint64_t *>(bufferData)[0] = 1;
    expectMemory<FamilyType>(uncompressibleHostMemAlloc, bufferData, bufferSize);

    retVal = clMemFreeINTEL(context, compressedDeviceMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);

    retVal = clMemFreeINTEL(context, uncompressibleHostMemAlloc);
    EXPECT_EQ(CL_SUCCESS, retVal);
}

INSTANTIATE_TEST_CASE_P(,
                        UmStatelessCompressionInSBA,
                        ::testing::Values(aub_stream::ENGINE_RCS,
                                          aub_stream::ENGINE_CCS));

struct StatelessCompressionInSBAWithBCS : public MulticontextAubFixture,
                                          public StatelessCopyKernelFixture,
                                          public ::testing::Test {
    void SetUp() override {
        DebugManager.flags.EnableStatelessCompression.set(1);
        DebugManager.flags.ForceAuxTranslationMode.set(static_cast<int32_t>(AuxTranslationMode::Blit));
        DebugManager.flags.EnableBlitterOperationsSupport.set(true);
        MulticontextAubFixture::SetUp(1, EnabledCommandStreamers::Single, true);
        StatelessCopyKernelFixture::SetUp(tileDevices[0], context.get());
        if (!tileDevices[0]->getHardwareInfo().featureTable.flags.ftrLocalMemory) {
            GTEST_SKIP();
        }
    }

    void TearDown() override {
        MulticontextAubFixture::TearDown();
        StatelessCopyKernelFixture::TearDown();
    }

    DebugManagerStateRestore debugRestorer;
    cl_int retVal = CL_SUCCESS;
};

XE_HP_CORE_TEST_F(StatelessCompressionInSBAWithBCS, GENERATEONLY_givenCompressedBufferInDeviceMemoryWhenAccessedStatelesslyThenEnableCompressionInSBA) {
    const size_t bufferSize = 2048;
    uint8_t writePattern[bufferSize];
    std::fill(writePattern, writePattern + sizeof(writePattern), 1);

    auto unCompressedBuffer = std::unique_ptr<Buffer>(Buffer::create(context.get(), CL_MEM_UNCOMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto unCompressedAllocation = unCompressedBuffer->getGraphicsAllocation(tileDevices[0]->getRootDeviceIndex());
    EXPECT_EQ(AllocationType::BUFFER, unCompressedAllocation->getAllocationType());
    EXPECT_FALSE(unCompressedAllocation->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, unCompressedAllocation->getMemoryPool());

    auto compressedBuffer = std::unique_ptr<Buffer>(Buffer::create(context.get(), CL_MEM_COMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto compressedAllocation = compressedBuffer->getGraphicsAllocation(tileDevices[0]->getRootDeviceIndex());
    EXPECT_TRUE(compressedAllocation->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedAllocation->getMemoryPool());
    EXPECT_TRUE(compressedAllocation->getDefaultGmm()->isCompressionEnabled);

    retVal = commandQueues[0][0]->enqueueWriteBuffer(compressedBuffer.get(), CL_FALSE, 0, bufferSize, writePattern, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(0, compressedBuffer.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(1, unCompressedBuffer.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = commandQueues[0][0]->enqueueKernel(kernel, 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    commandQueues[0][0]->finish();

    expectMemoryNotEqual<FamilyType>(AUBFixture::getGpuPointer(compressedAllocation), writePattern, bufferSize, 0, 0);

    expectMemory<FamilyType>(AUBFixture::getGpuPointer(unCompressedAllocation), writePattern, bufferSize, 0, 0);
}

XE_HP_CORE_TEST_F(StatelessCompressionInSBAWithBCS, givenUncompressibleBufferInHostMemoryWhenAccessedStatelesslyThenDisableCompressionInSBA) {
    const size_t bufferSize = 2048;
    uint8_t writePattern[bufferSize];
    std::fill(writePattern, writePattern + sizeof(writePattern), 1);

    auto compressedBuffer = std::unique_ptr<Buffer>(Buffer::create(context.get(), CL_MEM_COMPRESSED_HINT_INTEL, bufferSize, nullptr, retVal));
    auto compressedAllocation = compressedBuffer->getGraphicsAllocation(tileDevices[0]->getRootDeviceIndex());
    EXPECT_TRUE(compressedAllocation->isCompressionEnabled());
    EXPECT_EQ(MemoryPool::LocalMemory, compressedAllocation->getMemoryPool());
    EXPECT_TRUE(compressedAllocation->getDefaultGmm()->isCompressionEnabled);

    auto uncompressibleBufferInHostMemory = std::unique_ptr<Buffer>(Buffer::create(context.get(), CL_MEM_FORCE_HOST_MEMORY_INTEL, bufferSize, nullptr, retVal));
    auto uncompressibleAllocationInHostMemory = uncompressibleBufferInHostMemory->getGraphicsAllocation(tileDevices[0]->getRootDeviceIndex());
    EXPECT_EQ(AllocationType::BUFFER_HOST_MEMORY, uncompressibleAllocationInHostMemory->getAllocationType());
    EXPECT_TRUE(MemoryPoolHelper::isSystemMemoryPool(uncompressibleAllocationInHostMemory->getMemoryPool()));

    retVal = commandQueues[0][0]->enqueueWriteBuffer(compressedBuffer.get(), CL_FALSE, 0, bufferSize, writePattern, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(0, compressedBuffer.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = kernel->setArg(1, uncompressibleBufferInHostMemory.get());
    ASSERT_EQ(CL_SUCCESS, retVal);

    size_t globalWorkSize[3] = {bufferSize, 1, 1};
    retVal = commandQueues[0][0]->enqueueKernel(kernel, 1, nullptr, globalWorkSize, nullptr, 0, nullptr, nullptr);
    ASSERT_EQ(CL_SUCCESS, retVal);

    commandQueues[0][0]->finish();

    expectMemoryNotEqual<FamilyType>(AUBFixture::getGpuPointer(compressedAllocation), writePattern, bufferSize, 0, 0);

    expectMemory<FamilyType>(AUBFixture::getGpuPointer(uncompressibleAllocationInHostMemory), writePattern, bufferSize, 0, 0);
}
