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

#include "opencl/test/unit_test/fixtures/hello_world_fixture.h"

using namespace NEO;

struct KernelSubGroupInfoKhrFixture : HelloWorldFixture<HelloWorldFixtureFactory> {
    typedef HelloWorldFixture<HelloWorldFixtureFactory> ParentClass;

    void setUp() {
        ParentClass::setUp();
        maxSimdSize = static_cast<size_t>(pKernel->getKernelInfo().getMaxSimdSize());
        ASSERT_GE(maxSimdSize, 8u);
        maxWorkDim = static_cast<size_t>(pClDevice->getDeviceInfo().maxWorkItemDimensions);
        ASSERT_EQ(maxWorkDim, 3u);
    }

    void tearDown() {
        ParentClass::tearDown();
    }

    size_t inputValue[3];
    size_t paramValue;
    size_t paramValueSizeRet;
    size_t maxSimdSize;
    size_t calculatedWGS;
    size_t maxWorkDim;
};

namespace ULT {

typedef Test<KernelSubGroupInfoKhrFixture> KernelSubGroupInfoKhrTest;

template <typename ParamType>
struct KernelSubGroupInfoKhrParamFixture : KernelSubGroupInfoKhrFixture,
                                           ::testing::TestWithParam<ParamType> {
    void SetUp() override {
        KernelSubGroupInfoKhrFixture::setUp();
    }

    void TearDown() override {
        KernelSubGroupInfoKhrFixture::tearDown();
    }
};

struct TestParam {
    size_t gwsX;
    size_t gwsY;
    size_t gwsZ;
} kernelSubGroupInfoKhrWGS[] = {
    {0, 0, 0},
    {1, 1, 1},
    {1, 5, 1},
    {8, 1, 1},
    {16, 1, 1},
    {32, 1, 1},
    {64, 1, 1},
    {1, 190, 1},
    {1, 510, 1},
    {512, 1, 1}};

typedef KernelSubGroupInfoKhrParamFixture<TestParam> KernelSubGroupInfoKhrReturnSizeTest;

INSTANTIATE_TEST_SUITE_P(wgs,
                         KernelSubGroupInfoKhrReturnSizeTest,
                         ::testing::ValuesIn(kernelSubGroupInfoKhrWGS));

TEST_P(KernelSubGroupInfoKhrReturnSizeTest, GivenLwsParameterWhenGettingMaxSubGroupSizeThenCorrectValueIsReturned) {
    paramValueSizeRet = 0;

    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        CL_KERNEL_MAX_SUB_GROUP_SIZE_FOR_NDRANGE,
        sizeof(size_t) * 3,
        inputValue,
        sizeof(size_t),
        &paramValue,
        &paramValueSizeRet);

    EXPECT_EQ(retVal, CL_SUCCESS);
    EXPECT_EQ(paramValueSizeRet, sizeof(size_t));

    EXPECT_EQ(paramValue, maxSimdSize);
}

typedef KernelSubGroupInfoKhrParamFixture<TestParam> KernelSubGroupInfoKhrReturnCountTest;

INSTANTIATE_TEST_SUITE_P(wgs,
                         KernelSubGroupInfoKhrReturnCountTest,
                         ::testing::ValuesIn(kernelSubGroupInfoKhrWGS));

TEST_P(KernelSubGroupInfoKhrReturnCountTest, GivenLwsParameterWhenGettingSubGroupCountThenCorrectValueIsReturned) {
    paramValueSizeRet = 0;
    inputValue[0] = GetParam().gwsX;
    inputValue[1] = GetParam().gwsY;
    inputValue[2] = GetParam().gwsZ;
    calculatedWGS = inputValue[0] * inputValue[1] * inputValue[2];

    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        CL_KERNEL_SUB_GROUP_COUNT_FOR_NDRANGE,
        sizeof(size_t) * 3,
        inputValue,
        sizeof(size_t),
        &paramValue,
        &paramValueSizeRet);

    EXPECT_EQ(retVal, CL_SUCCESS);
    EXPECT_EQ(paramValueSizeRet, sizeof(size_t));

    if (calculatedWGS % maxSimdSize == 0) {
        EXPECT_EQ(paramValue, calculatedWGS / maxSimdSize);
    } else {
        EXPECT_EQ(paramValue, (calculatedWGS / maxSimdSize) + 1);
    }
}

typedef KernelSubGroupInfoKhrParamFixture<TestParam> KernelSubGroupInfoKhrReturnCompileSizeTest;

TEST_F(KernelSubGroupInfoKhrReturnCompileSizeTest, GivenKernelWhenGettingRequiredSubGroupSizeThenCorrectValueIsReturned) {

    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        CL_KERNEL_COMPILE_SUB_GROUP_SIZE_INTEL,
        0,
        nullptr,
        sizeof(size_t),
        &paramValue,
        &paramValueSizeRet);

    EXPECT_EQ(CL_SUCCESS, retVal);
    EXPECT_EQ(paramValueSizeRet, sizeof(size_t));
    EXPECT_EQ(pKernel->getKernelInfo().kernelDescriptor.kernelMetadata.requiredSubGroupSize, paramValue);
}

TEST_F(KernelSubGroupInfoKhrTest, GivenNullKernelWhenGettingKernelSubGroupInfoThenInvalidKernelErrorIsReturned) {
    retVal = clGetKernelSubGroupInfoKHR(
        nullptr,
        pClDevice,
        0,
        0,
        nullptr,
        0,
        nullptr,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_KERNEL);
}

TEST_F(KernelSubGroupInfoKhrTest, GivenInvalidDeviceWhenGettingSubGroupInfoFromSingleDeviceKernelThenInvalidDeviceErrorIsReturned) {

    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        reinterpret_cast<cl_device_id>(pKernel),
        CL_KERNEL_COMPILE_SUB_GROUP_SIZE_INTEL,
        0,
        nullptr,
        sizeof(size_t),
        &paramValue,
        &paramValueSizeRet);

    EXPECT_EQ(CL_INVALID_DEVICE, retVal);
}

TEST_F(KernelSubGroupInfoKhrTest, GivenNullDeviceWhenGettingSubGroupInfoFromSingleDeviceKernelThenSuccessIsReturned) {

    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        nullptr,
        CL_KERNEL_COMPILE_SUB_GROUP_SIZE_INTEL,
        0,
        nullptr,
        sizeof(size_t),
        &paramValue,
        &paramValueSizeRet);

    EXPECT_EQ(CL_SUCCESS, retVal);
}

TEST_F(KernelSubGroupInfoKhrTest, GivenNullDeviceWhenGettingSubGroupInfoFromMultiDeviceKernelThenInvalidDeviceErrorIsReturned) {
    cl_int retVal{CL_SUCCESS};
    MockUnrestrictiveContext context;
    auto mockProgram = std::make_unique<MockProgram>(&context, false, context.getDevices());
    std::unique_ptr<MultiDeviceKernel> pMultiDeviceKernel(
        MultiDeviceKernel::create<MockKernel>(mockProgram.get(), this->pMultiDeviceKernel->getKernelInfos(), retVal));
    ASSERT_EQ(CL_SUCCESS, retVal);

    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel.get(),
        nullptr,
        CL_KERNEL_COMPILE_SUB_GROUP_SIZE_INTEL,
        0,
        nullptr,
        sizeof(size_t),
        &paramValue,
        &paramValueSizeRet);

    EXPECT_EQ(CL_INVALID_DEVICE, retVal);
}

TEST_F(KernelSubGroupInfoKhrTest, GivenInvalidParamNameWhenGettingKernelSubGroupInfoThenInvalidValueErrorIsReturned) {
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        0,
        sizeof(size_t),
        inputValue,
        sizeof(size_t),
        &paramValue,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_VALUE);
}

uint32_t /*cl_kernel_sub_group_info_khr*/ kernelSubGroupInfoKhrInputParams[] = {
    CL_KERNEL_MAX_SUB_GROUP_SIZE_FOR_NDRANGE_KHR,
    CL_KERNEL_SUB_GROUP_COUNT_FOR_NDRANGE_KHR};

typedef KernelSubGroupInfoKhrParamFixture<uint32_t /*cl_kernel_sub_group_info_khr*/> KernelSubGroupInfoKhrInputParamsTest;

INSTANTIATE_TEST_SUITE_P(KernelSubGroupInfoKhrInputParams,
                         KernelSubGroupInfoKhrInputParamsTest,
                         ::testing::ValuesIn(kernelSubGroupInfoKhrInputParams));

TEST_P(KernelSubGroupInfoKhrInputParamsTest, GivenInvalidInputWhenGettingKernelSubGroupInfoThenInvalidValueErrorIsReturned) {
    // work dim == 0
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        GetParam(),
        0,
        inputValue,
        0,
        nullptr,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_VALUE);

    // work dim % sizeof(size_t) != 0
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        GetParam(),
        (sizeof(size_t) * maxWorkDim) - 1,
        inputValue,
        0,
        nullptr,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_VALUE);

    // work dim > maxWorkDim
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        GetParam(),
        sizeof(size_t) * (maxWorkDim + 1),
        inputValue,
        0,
        nullptr,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_VALUE);

    // null input_value
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        GetParam(),
        sizeof(size_t) * (maxWorkDim),
        nullptr,
        0,
        nullptr,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_VALUE);
}

TEST_P(KernelSubGroupInfoKhrInputParamsTest, GivenInvalidParamSizeWhenGettingKernelSubGroupInfoThenInvalidValueErrorIsReturned) {
    // param_value_size < sizeof(size_t)
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        GetParam(),
        sizeof(size_t),
        inputValue,
        sizeof(size_t) - 1,
        &paramValue,
        nullptr);

    EXPECT_EQ(retVal, CL_INVALID_VALUE);
}

TEST_P(KernelSubGroupInfoKhrInputParamsTest, GivenNoReturnPointerWhenGettingKernelSubGroupInfoThenSuccessIsReturned) {
    retVal = clGetKernelSubGroupInfoKHR(
        pMultiDeviceKernel,
        pClDevice,
        GetParam(),
        sizeof(size_t),
        inputValue,
        0,
        nullptr,
        nullptr);

    EXPECT_EQ(retVal, CL_SUCCESS);
}
} // namespace ULT
