From d14b117b33e641f4dc13c6fda91ede1b1ea21bb4 Mon Sep 17 00:00:00 2001
From: Aaron Greig <aaron.greig@codeplay.com>
Date: Wed, 8 Nov 2023 14:44:12 +0000
Subject: [PATCH] Query support for 2D USM operations in CTS tests.

---
 .../enqueue/urEnqueueUSMFill2D.cpp            |  8 ++++++++
 .../enqueue/urEnqueueUSMMemcpy2D.cpp          | 20 +++++++++++++------
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/test/conformance/enqueue/urEnqueueUSMFill2D.cpp b/test/conformance/enqueue/urEnqueueUSMFill2D.cpp
index 37470d40e2..fcb244e94a 100644
--- a/test/conformance/enqueue/urEnqueueUSMFill2D.cpp
+++ b/test/conformance/enqueue/urEnqueueUSMFill2D.cpp
@@ -33,6 +33,14 @@ struct urEnqueueUSMFill2DTestWithParam
     void SetUp() override {
         UUR_RETURN_ON_FATAL_FAILURE(urQueueTestWithParam::SetUp());
 
+        bool memfill2d_support = false;
+        ASSERT_SUCCESS(urContextGetInfo(
+            context, UR_CONTEXT_INFO_USM_FILL2D_SUPPORT,
+            sizeof(memfill2d_support), &memfill2d_support, nullptr));
+        if (!memfill2d_support) {
+            GTEST_SKIP() << "2D USM mem fill is not supported";
+        }
+
         pitch = std::get<1>(GetParam()).pitch;
         width = std::get<1>(GetParam()).width;
         height = std::get<1>(GetParam()).height;
diff --git a/test/conformance/enqueue/urEnqueueUSMMemcpy2D.cpp b/test/conformance/enqueue/urEnqueueUSMMemcpy2D.cpp
index 8eaed4b743..d0e3dd9f72 100644
--- a/test/conformance/enqueue/urEnqueueUSMMemcpy2D.cpp
+++ b/test/conformance/enqueue/urEnqueueUSMMemcpy2D.cpp
@@ -17,6 +17,14 @@ struct urEnqueueUSMMemcpy2DTestWithParam
             GTEST_SKIP() << "Device USM is not supported";
         }
 
+        bool memcpy2d_support = false;
+        ASSERT_SUCCESS(urContextGetInfo(
+            context, UR_CONTEXT_INFO_USM_MEMCPY2D_SUPPORT,
+            sizeof(memcpy2d_support), &memcpy2d_support, nullptr));
+        if (!memcpy2d_support) {
+            GTEST_SKIP() << "2D USM memcpy is not supported";
+        }
+
         const auto [inPitch, inWidth, inHeight] = getParam();
         std::tie(pitch, width, height) =
             std::make_tuple(inPitch, inWidth, inHeight);
@@ -28,9 +36,9 @@ struct urEnqueueUSMMemcpy2DTestWithParam
                                         num_elements, &pDst));
         ur_event_handle_t memset_event = nullptr;
 
-        ASSERT_SUCCESS(urEnqueueUSMFill2D(
-            queue, pSrc, pitch, sizeof(memset_value), &memset_value, width,
-            height, 0, nullptr, &memset_event));
+        ASSERT_SUCCESS(urEnqueueUSMFill(queue, pSrc, sizeof(memset_value),
+                                        &memset_value, pitch * height, 0,
+                                        nullptr, &memset_event));
 
         ASSERT_SUCCESS(urQueueFlush(queue));
         ASSERT_SUCCESS(urEventWait(1, &memset_event));
@@ -171,9 +179,9 @@ TEST_P(urEnqueueUSMMemcpy2DNegativeTest, InvalidEventWaitList) {
     // enqueue something to get an event
     ur_event_handle_t event = nullptr;
     uint8_t fill_pattern = 14;
-    ASSERT_SUCCESS(urEnqueueUSMFill2D(queue, pDst, pitch, sizeof(fill_pattern),
-                                      &fill_pattern, width, height, 0, nullptr,
-                                      &event));
+    ASSERT_SUCCESS(urEnqueueUSMFill(queue, pDst, sizeof(fill_pattern),
+                                    &fill_pattern, pitch * height, 0, nullptr,
+                                    &event));
     ASSERT_NE(event, nullptr);
     ASSERT_SUCCESS(urQueueFinish(queue));