From bb53de1e4edc576292f02d8b2be3d0f173905796 Mon Sep 17 00:00:00 2001
From: Young Jin Yoon <young.jin.yoon@intel.com>
Date: Wed, 28 Apr 2021 00:15:37 +0000
Subject: [PATCH] Fix errors for incorrect work group size

Related-To: LOCI-2084
Signed-off-by: Young Jin Yoon <young.jin.yoon@intel.com>
---
 level_zero/core/source/kernel/kernel_imp.cpp  | 12 +++++++++
 .../unit_tests/sources/kernel/test_kernel.cpp | 26 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/level_zero/core/source/kernel/kernel_imp.cpp b/level_zero/core/source/kernel/kernel_imp.cpp
index e8b83aef94..d2ef2be1dd 100644
--- a/level_zero/core/source/kernel/kernel_imp.cpp
+++ b/level_zero/core/source/kernel/kernel_imp.cpp
@@ -292,6 +292,18 @@ ze_result_t KernelImp::setGroupSize(uint32_t groupSizeX, uint32_t groupSizeY,
     this->groupSize[1] = groupSizeY;
     this->groupSize[2] = groupSizeZ;
     const NEO::KernelDescriptor &kernelDescriptor = kernelImmData->getDescriptor();
+    for (uint32_t i = 0u; i < 3u; i++) {
+        if (kernelDescriptor.kernelAttributes.requiredWorkgroupSize[i] != 0 &&
+            kernelDescriptor.kernelAttributes.requiredWorkgroupSize[i] != this->groupSize[i]) {
+            NEO::printDebugString(NEO::DebugManager.flags.PrintDebugMessages.get(), stderr,
+                                  "Invalid group size {%d, %d, %d} specified, requiredWorkGroupSize = {%d, %d, %d}\n",
+                                  this->groupSize[0], this->groupSize[1], this->groupSize[2],
+                                  kernelDescriptor.kernelAttributes.requiredWorkgroupSize[0],
+                                  kernelDescriptor.kernelAttributes.requiredWorkgroupSize[1],
+                                  kernelDescriptor.kernelAttributes.requiredWorkgroupSize[2]);
+            return ZE_RESULT_ERROR_INVALID_GROUP_SIZE_DIMENSION;
+        }
+    }
 
     auto simdSize = kernelDescriptor.kernelAttributes.simdSize;
     this->numThreadsPerThreadGroup = static_cast<uint32_t>((itemsInGroup + simdSize - 1u) / simdSize);
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 f97ea83df9..ec2bf535eb 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
@@ -106,6 +106,32 @@ HWTEST_F(KernelImpSetGroupSizeTest, givenLocalIdGenerationByRuntimeDisabledWhenS
     EXPECT_EQ(nullptr, mockKernel.perThreadDataForWholeThreadGroup);
 }
 
+HWTEST_F(KernelImpSetGroupSizeTest, givenIncorrectGroupSizeWhenSettingGroupSizeThenInvalidGroupSizeDimensionErrorIsReturned) {
+    Mock<Kernel> mockKernel;
+    Mock<Module> mockModule(this->device, nullptr);
+    for (auto i = 0u; i < 3u; i++) {
+        mockKernel.descriptor.kernelAttributes.requiredWorkgroupSize[i] = 2;
+    }
+    mockKernel.module = &mockModule;
+
+    uint32_t groupSize[3] = {1, 1, 1};
+    auto ret = mockKernel.setGroupSize(groupSize[0], groupSize[1], groupSize[2]);
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_GROUP_SIZE_DIMENSION, ret);
+}
+
+HWTEST_F(KernelImpSetGroupSizeTest, givenZeroGroupSizeWhenSettingGroupSizeThenInvalidArgumentErrorIsReturned) {
+    Mock<Kernel> mockKernel;
+    Mock<Module> mockModule(this->device, nullptr);
+    for (auto i = 0u; i < 3u; i++) {
+        mockKernel.descriptor.kernelAttributes.requiredWorkgroupSize[i] = 2;
+    }
+    mockKernel.module = &mockModule;
+
+    uint32_t groupSize[3] = {0, 0, 0};
+    auto ret = mockKernel.setGroupSize(groupSize[0], groupSize[1], groupSize[2]);
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, ret);
+}
+
 using SetKernelArg = Test<ModuleFixture>;
 using ImageSupport = IsWithinProducts<IGFX_SKYLAKE, IGFX_TIGERLAKE_LP>;