From e605151eb1bd7d8d5c3b45b55fd440c761c692ac Mon Sep 17 00:00:00 2001
From: Jatin Chaudhary <JatinJaikishan.Chaudhary@amd.com>
Date: Mon, 16 Oct 2023 23:11:31 +0100
Subject: [PATCH] SWDEV-411064 - check if image is supported before running
 image specific tests

Also fix CU Mask test.

Change-Id: I0fbb93a1bcce131e0f98dee7cd5da9704d2ca8ac
---
 catch/unit/errorHandling/hipGetLastError.cc | 6 ++++++
 catch/unit/memory/hipMemcpy3DAsync.cc       | 2 ++
 catch/unit/stream/hipStreamGetCUMask.cc     | 9 +++++++--
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/catch/unit/errorHandling/hipGetLastError.cc b/catch/unit/errorHandling/hipGetLastError.cc
index 0ca215394..c6424688e 100644
--- a/catch/unit/errorHandling/hipGetLastError.cc
+++ b/catch/unit/errorHandling/hipGetLastError.cc
@@ -203,6 +203,8 @@ TEST_CASE("Unit_hipGetLastError_with_hipMemcpyDtoHAsync") {
 *  - HIP_VERSION >= 6.0
 */
 TEST_CASE("Unit_hipGetLastError_with_hipMemcpyParam2DAsync") {
+  CHECK_IMAGE_SUPPORT
+
   float* A_h{nullptr}, *B_h{nullptr}, *C_h{nullptr}, * A_d{nullptr};
   size_t pitch_A;
   size_t width{WIDTH * sizeof(float)};
@@ -283,6 +285,8 @@ TEST_CASE("Unit_hipGetLastError_with_hipDrvMemcpy3DAsync") {
 */
 
 TEST_CASE("Unit_hipGetLastError_with_hipMemcpy3DAsync") {
+  CHECK_IMAGE_SUPPORT
+
   constexpr int width{10}, height{10}, depth{10};
   auto size = width * height * depth * sizeof(int);
   hipArray_t devArray;
@@ -469,6 +473,8 @@ TEST_CASE("Unit_hipGetLastError_with_hipMemPrefetchAsync") {
 */
 
 TEST_CASE("Unit_hipGetLastError_with_hipMemcpy2DAsync") {
+  CHECK_IMAGE_SUPPORT
+
   int* A_h{nullptr}, *A_d{nullptr};
   size_t pitch_A;
   size_t width{WIDTH * sizeof(int)};
diff --git a/catch/unit/memory/hipMemcpy3DAsync.cc b/catch/unit/memory/hipMemcpy3DAsync.cc
index b1d864132..870672844 100644
--- a/catch/unit/memory/hipMemcpy3DAsync.cc
+++ b/catch/unit/memory/hipMemcpy3DAsync.cc
@@ -89,6 +89,8 @@ TEST_CASE("Unit_hipMemcpy3DAsync_Positive_Parameters") {
 }
 
 TEST_CASE("Unit_hipMemcpy3DAsync_Positive_Array") {
+  CHECK_IMAGE_SUPPORT
+
   constexpr bool async = true;
   SECTION("Array from/to Host") { Memcpy3DArrayHostShell<async>(Memcpy3DWrapper<async>); }
 #if HT_NVIDIA // Disabled on AMD due to defect - EXSWHTEC-238
diff --git a/catch/unit/stream/hipStreamGetCUMask.cc b/catch/unit/stream/hipStreamGetCUMask.cc
index 2e689b351..695388993 100644
--- a/catch/unit/stream/hipStreamGetCUMask.cc
+++ b/catch/unit/stream/hipStreamGetCUMask.cc
@@ -25,6 +25,8 @@ Testcase Scenarios :
 */
 
 #include <hip_test_common.hh>
+
+#include <cmath>
 #include <vector>
 
 
@@ -33,8 +35,7 @@ Testcase Scenarios :
  * Scenario to verify hipExtStreamGetCUMask api returns custom mask set.
  */
 TEST_CASE("Unit_hipExtStreamGetCUMask_verifyDefaultAndCustomMask") {
-  constexpr int maxNum = 10;
-  std::vector<uint32_t> cuMask(maxNum);
+  constexpr unsigned maxCUPerValue = 32;
   hipDeviceProp_t props;
   std::stringstream ss;
   std::string gCUMask;
@@ -53,6 +54,10 @@ TEST_CASE("Unit_hipExtStreamGetCUMask_verifyDefaultAndCustomMask") {
   INFO("info: running on bus " << "0x" << props.pciBusID << " " <<
          props.name << " with " << props.multiProcessorCount << " CUs");
 
+  const unsigned int maxNum =
+      static_cast<unsigned int>(std::ceil((props.multiProcessorCount * 1.0f) / maxCUPerValue));
+  std::vector<uint32_t> cuMask(maxNum);
+
   // Get global CU Mask if exists
   gCUMask = TestContext::getEnvVar("ROC_GLOBAL_CU_MASK");
   if (!gCUMask.empty()) {