From 59da1d97e84e4709d8838be042ad5d49b4dcf8c8 Mon Sep 17 00:00:00 2001
From: "akmkhale@ansatnuc04" <akm.khaled.talukder@intel.com>
Date: Fri, 10 Feb 2023 15:33:05 -0600
Subject: [PATCH] Test with CFD and usm_ndarray

---
 .../core/kernel_interface/dispatcher.py       |  2 -
 .../tests/kernel_tests/test_arg_types.py      | 68 +++++++++++++++----
 2 files changed, 54 insertions(+), 16 deletions(-)

diff --git a/numba_dpex/core/kernel_interface/dispatcher.py b/numba_dpex/core/kernel_interface/dispatcher.py
index 39ca09cfc0..d96bcc139d 100644
--- a/numba_dpex/core/kernel_interface/dispatcher.py
+++ b/numba_dpex/core/kernel_interface/dispatcher.py
@@ -636,8 +636,6 @@ def __call__(self, *args):
             tuple(argtypes),
             self.pyfunc,
             dpex_kernel_target.target_context.codegen(),
-            # backend=backend,
-            # device_type=exec_queue.sycl_device,
             exec_queue=exec_queue,
         )
 
diff --git a/numba_dpex/tests/kernel_tests/test_arg_types.py b/numba_dpex/tests/kernel_tests/test_arg_types.py
index 5e1e1bcc9e..18fd4469ae 100644
--- a/numba_dpex/tests/kernel_tests/test_arg_types.py
+++ b/numba_dpex/tests/kernel_tests/test_arg_types.py
@@ -5,6 +5,7 @@
 import sys
 
 import dpctl
+import dpctl.tensor as dpt
 import numpy as np
 import pytest
 
@@ -39,13 +40,39 @@ def input_arrays(request):
 
 @pytest.mark.parametrize("filter_str", filter_strings)
 def test_kernel_arg_types(filter_str, input_arrays):
-    kernel = dpex.kernel(mul_kernel)
-    a, actual, c = input_arrays
+    usm_type = "device"
+
+    a, b, c = input_arrays
     expected = a * c
+
     device = dpctl.SyclDevice(filter_str)
-    with dpctl.device_context(device):
-        kernel[global_size, local_size](a, actual, c)
-    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=0)
+    queue = dpctl.SyclQueue(device)
+
+    da = dpt.usm_ndarray(
+        a.shape,
+        dtype=a.dtype,
+        buffer=usm_type,
+        buffer_ctor_kwargs={"queue": queue},
+    )
+    da.usm_data.copy_from_host(a.reshape((-1)).view("|u1"))
+
+    db = dpt.usm_ndarray(
+        b.shape,
+        dtype=b.dtype,
+        buffer=usm_type,
+        buffer_ctor_kwargs={"queue": queue},
+    )
+    db.usm_data.copy_from_host(b.reshape((-1)).view("|u1"))
+
+    kernel = dpex.kernel(mul_kernel)
+    kernel[dpex.NdRange(dpex.Range(global_size), dpex.Range(local_size))](
+        da, db, c
+    )
+
+    result = np.zeros_like(b)
+    db.usm_data.copy_to_host(result.reshape((-1)).view("|u1"))
+
+    np.testing.assert_allclose(result, expected, rtol=1e-5, atol=0)
 
 
 def check_bool_kernel(A, test):
@@ -55,17 +82,30 @@ def check_bool_kernel(A, test):
         A[0] = 222
 
 
-@pytest.mark.skipif(
-    sys.platform.startswith("win"), reason="failing on teamcity windows CI"
-)
 @pytest.mark.parametrize("filter_str", filter_strings)
 def test_bool_type(filter_str):
-    kernel = dpex.kernel(check_bool_kernel)
+    usm_type = "device"
     a = np.array([2], np.int64)
 
     device = dpctl.SyclDevice(filter_str)
-    with dpctl.device_context(device):
-        kernel[a.size, dpex.DEFAULT_LOCAL_SIZE](a, True)
-        assert a[0] == 111
-        kernel[a.size, dpex.DEFAULT_LOCAL_SIZE](a, False)
-        assert a[0] == 222
+    queue = dpctl.SyclQueue(device)
+
+    da = dpt.usm_ndarray(
+        a.shape,
+        dtype=a.dtype,
+        buffer=usm_type,
+        buffer_ctor_kwargs={"queue": queue},
+    )
+    da.usm_data.copy_from_host(a.reshape((-1)).view("|u1"))
+
+    kernel = dpex.kernel(check_bool_kernel)
+
+    kernel[dpex.Range(a.size)](da, True)
+    result = np.zeros_like(a)
+    da.usm_data.copy_to_host(result.reshape((-1)).view("|u1"))
+    assert result[0] == 111
+
+    kernel[dpex.Range(a.size)](da, False)
+    result = np.zeros_like(a)
+    da.usm_data.copy_to_host(result.reshape((-1)).view("|u1"))
+    assert result[0] == 222