diff --git a/numba_dpex/tests/experimental/codegen/test_inline_threshold_codegen.py b/numba_dpex/tests/experimental/codegen/test_inline_threshold_codegen.py index bafb430903..39a2442698 100644 --- a/numba_dpex/tests/experimental/codegen/test_inline_threshold_codegen.py +++ b/numba_dpex/tests/experimental/codegen/test_inline_threshold_codegen.py @@ -5,14 +5,15 @@ import dpctl from numba.core import types -import numba_dpex as dpex from numba_dpex import DpctlSyclQueue, DpnpNdArray from numba_dpex import experimental as dpex_exp from numba_dpex import int64 +from numba_dpex.core.types.kernel_api.index_space_ids import ItemType +from numba_dpex.kernel_api import Item -def kernel_func(a, b, c): - i = dpex.get_global_id(0) +def kernel_func(item: Item, a, b, c): + i = item.get_id(0) c[i] = a[i] + b[i] @@ -36,7 +37,7 @@ def test_codegen_with_max_inline_threshold(): queue_ty = DpctlSyclQueue(dpctl.SyclQueue()) i64arr_ty = DpnpNdArray(ndim=1, dtype=int64, layout="C", queue=queue_ty) - kernel_sig = types.void(i64arr_ty, i64arr_ty, i64arr_ty) + kernel_sig = types.void(ItemType(1), i64arr_ty, i64arr_ty, i64arr_ty) disp = dpex_exp.kernel(inline_threshold=3)(kernel_func) disp.compile(kernel_sig) @@ -57,7 +58,7 @@ def test_codegen_without_max_inline_threshold(): queue_ty = DpctlSyclQueue(dpctl.SyclQueue()) i64arr_ty = DpnpNdArray(ndim=1, dtype=int64, layout="C", queue=queue_ty) - kernel_sig = types.void(i64arr_ty, i64arr_ty, i64arr_ty) + kernel_sig = types.void(ItemType(1), i64arr_ty, i64arr_ty, i64arr_ty) disp = dpex_exp.kernel(kernel_func) disp.compile(kernel_sig) @@ -70,4 +71,4 @@ def test_codegen_without_max_inline_threshold(): if not f.is_declaration: count_of_non_declaration_type_functions += 1 - assert count_of_non_declaration_type_functions == 2 + assert count_of_non_declaration_type_functions == 3 diff --git a/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_fetch_phi.py b/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_fetch_phi.py index 30f268c87a..c239c8f91a 100644 --- a/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_fetch_phi.py +++ b/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_fetch_phi.py @@ -6,9 +6,8 @@ import pytest from numba.core.errors import TypingError -import numba_dpex as dpex import numba_dpex.experimental as dpex_exp -from numba_dpex.kernel_api import AtomicRef +from numba_dpex.kernel_api import AtomicRef, Item, Range from numba_dpex.tests._helper import get_all_dtypes list_of_supported_dtypes = get_all_dtypes( @@ -45,8 +44,8 @@ def test_fetch_phi_fn(input_arrays, ref_index, fetch_phi_fn): """A test for all fetch_phi atomic functions.""" @dpex_exp.kernel - def _kernel(a, b, ref_index): - i = dpex.get_global_id(0) + def _kernel(item: Item, a, b, ref_index): + i = item.get_id(0) v = AtomicRef(b, index=ref_index) getattr(v, fetch_phi_fn)(a[i]) @@ -60,9 +59,9 @@ def _kernel(a, b, ref_index): # fetch_and, fetch_or, fetch_xor accept only int arguments. # test for TypingError when float arguments are passed. with pytest.raises(TypingError): - dpex_exp.call_kernel(_kernel, dpex.Range(10), a, b, ref_index) + dpex_exp.call_kernel(_kernel, Range(10), a, b, ref_index) else: - dpex_exp.call_kernel(_kernel, dpex.Range(10), a, b, ref_index) + dpex_exp.call_kernel(_kernel, Range(10), a, b, ref_index) # Verify that `a` accumulated at b[ref_index] by kernel # matches the `a` accumulated at b[ref_index+1] using Python for i in range(a.size): @@ -76,8 +75,8 @@ def test_fetch_phi_retval(fetch_phi_fn): """A test for all fetch_phi atomic functions.""" @dpex_exp.kernel - def _kernel(a, b, c): - i = dpex.get_global_id(0) + def _kernel(item: Item, a, b, c): + i = item.get_id(0) v = AtomicRef(b, index=i) c[i] = getattr(v, fetch_phi_fn)(a[i]) @@ -89,7 +88,7 @@ def _kernel(a, b, c): b_copy = dpnp.copy(b) c_copy = dpnp.copy(c) - dpex_exp.call_kernel(_kernel, dpex.Range(10), a, b, c) + dpex_exp.call_kernel(_kernel, Range(10), a, b, c) # Verify if the value returned by fetch_phi kernel # stored into `c` is same as the value returned @@ -108,8 +107,8 @@ def test_fetch_phi_diff_types(fetch_phi_fn): """ @dpex_exp.kernel - def _kernel(a, b): - i = dpex.get_global_id(0) + def _kernel(item: Item, a, b): + i = item.get_id(0) v = AtomicRef(b, index=0) getattr(v, fetch_phi_fn)(a[i]) @@ -118,19 +117,19 @@ def _kernel(a, b): b = dpnp.zeros(N, dtype=dpnp.int32) with pytest.raises(TypingError): - dpex_exp.call_kernel(_kernel, dpex.Range(10), a, b) + dpex_exp.call_kernel(_kernel, Range(10), a, b) @dpex_exp.kernel -def atomic_ref_0(a): - i = dpex.get_global_id(0) +def atomic_ref_0(item: Item, a): + i = item.get_id(0) v = AtomicRef(a, index=0) v.fetch_add(a[i + 2]) @dpex_exp.kernel -def atomic_ref_1(a): - i = dpex.get_global_id(0) +def atomic_ref_1(item: Item, a): + i = item.get_id(0) v = AtomicRef(a, index=1) v.fetch_add(a[i + 2]) @@ -144,24 +143,24 @@ def test_spirv_compiler_flags_add(): N = 10 a = dpnp.ones(N, dtype=dpnp.float32) - dpex_exp.call_kernel(atomic_ref_0, dpex.Range(N - 2), a) - dpex_exp.call_kernel(atomic_ref_1, dpex.Range(N - 2), a) + dpex_exp.call_kernel(atomic_ref_0, Range(N - 2), a) + dpex_exp.call_kernel(atomic_ref_1, Range(N - 2), a) assert a[0] == N - 1 assert a[1] == N - 1 @dpex_exp.kernel -def atomic_max_0(a): - i = dpex.get_global_id(0) +def atomic_max_0(item: Item, a): + i = item.get_id(0) v = AtomicRef(a, index=0) if i != 0: v.fetch_max(a[i]) @dpex_exp.kernel -def atomic_max_1(a): - i = dpex.get_global_id(0) +def atomic_max_1(item: Item, a): + i = item.get_id(0) v = AtomicRef(a, index=0) if i != 0: v.fetch_max(a[i]) @@ -177,8 +176,8 @@ def test_spirv_compiler_flags_max(): a = dpnp.arange(N, dtype=dpnp.float32) b = dpnp.arange(N, dtype=dpnp.float32) - dpex_exp.call_kernel(atomic_max_0, dpex.Range(N), a) - dpex_exp.call_kernel(atomic_max_1, dpex.Range(N), b) + dpex_exp.call_kernel(atomic_max_0, Range(N), a) + dpex_exp.call_kernel(atomic_max_1, Range(N), b) assert a[0] == N - 1 assert b[0] == N - 1 diff --git a/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_ref.py b/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_ref.py index fb40f9311d..d70f1aa03a 100644 --- a/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_ref.py +++ b/numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_ref.py @@ -6,22 +6,21 @@ import pytest from numba.core.errors import TypingError -import numba_dpex as dpex import numba_dpex.experimental as dpex_exp -from numba_dpex.kernel_api import AddressSpace, AtomicRef +from numba_dpex.kernel_api import AddressSpace, AtomicRef, Item, Range def test_atomic_ref_compilation(): @dpex_exp.kernel - def atomic_ref_kernel(a, b): - i = dpex.get_global_id(0) + def atomic_ref_kernel(item: Item, a, b): + i = item.get_id(0) v = AtomicRef(b, index=0, address_space=AddressSpace.GLOBAL) v.fetch_add(a[i]) a = dpnp.ones(10) b = dpnp.zeros(10) try: - dpex_exp.call_kernel(atomic_ref_kernel, dpex.Range(10), a, b) + dpex_exp.call_kernel(atomic_ref_kernel, Range(10), a, b) except Exception: pytest.fail("Unexpected execution failure") @@ -33,8 +32,8 @@ def test_atomic_ref_compilation_failure(): """ @dpex_exp.kernel - def atomic_ref_kernel(a, b): - i = dpex.get_global_id(0) + def atomic_ref_kernel(item: Item, a, b): + i = item.get_id(0) v = AtomicRef(b, index=0, address_space=AddressSpace.LOCAL) v.fetch_add(a[i]) @@ -42,4 +41,4 @@ def atomic_ref_kernel(a, b): b = dpnp.zeros(10) with pytest.raises(TypingError): - dpex_exp.call_kernel(atomic_ref_kernel, dpex.Range(10), a, b) + dpex_exp.call_kernel(atomic_ref_kernel, Range(10), a, b) diff --git a/numba_dpex/tests/experimental/test_async_kernel.py b/numba_dpex/tests/experimental/test_async_kernel.py index 8e772e2891..58ec040e43 100644 --- a/numba_dpex/tests/experimental/test_async_kernel.py +++ b/numba_dpex/tests/experimental/test_async_kernel.py @@ -7,10 +7,9 @@ import pytest from numba.core.errors import TypingError -import numba_dpex as dpex import numba_dpex.experimental as exp_dpex -from numba_dpex import Range from numba_dpex.experimental import testing +from numba_dpex.kernel_api import Item, Range @exp_dpex.kernel( @@ -19,8 +18,8 @@ no_cpython_wrapper=True, no_cfunc_wrapper=True, ) -def add(a, b, c): - i = dpex.get_global_id(0) +def add(item: Item, a, b, c): + i = item.get_id(0) c[i] = b[i] + a[i] diff --git a/numba_dpex/tests/experimental/test_compiler_warnings.py b/numba_dpex/tests/experimental/test_compiler_warnings.py index 5cbc2bc9e6..e7ae1ad550 100644 --- a/numba_dpex/tests/experimental/test_compiler_warnings.py +++ b/numba_dpex/tests/experimental/test_compiler_warnings.py @@ -6,14 +6,15 @@ import pytest from numba.core import types -import numba_dpex as dpex from numba_dpex import DpctlSyclQueue, DpnpNdArray from numba_dpex import experimental as dpex_exp from numba_dpex import int64 +from numba_dpex.core.types.kernel_api.index_space_ids import ItemType +from numba_dpex.kernel_api import Item -def _kernel(a, b, c): - i = dpex.get_global_id(0) +def _kernel(item: Item, a, b, c): + i = item.get_id(0) c[i] = a[i] + b[i] @@ -30,5 +31,5 @@ def test_inline_threshold_level_warning(): with pytest.warns(UserWarning): queue_ty = DpctlSyclQueue(dpctl.SyclQueue()) i64arr_ty = DpnpNdArray(ndim=1, dtype=int64, layout="C", queue=queue_ty) - kernel_sig = types.void(i64arr_ty, i64arr_ty, i64arr_ty) + kernel_sig = types.void(ItemType(1), i64arr_ty, i64arr_ty, i64arr_ty) dpex_exp.kernel(inline_threshold=3)(_kernel).compile(kernel_sig) diff --git a/numba_dpex/tests/experimental/test_inline_threshold_config.py b/numba_dpex/tests/experimental/test_inline_threshold_config.py index c90843f20b..75c0430171 100644 --- a/numba_dpex/tests/experimental/test_inline_threshold_config.py +++ b/numba_dpex/tests/experimental/test_inline_threshold_config.py @@ -6,10 +6,11 @@ import numba_dpex as dpex from numba_dpex import experimental as dpex_exp +from numba_dpex.kernel_api import Item -def kernel_func(a, b, c): - i = dpex.get_global_id(0) +def kernel_func(item: Item, a, b, c): + i = item.get_id(0) c[i] = a[i] + b[i] diff --git a/numba_dpex/tests/experimental/test_kernel_specialization.py b/numba_dpex/tests/experimental/test_kernel_specialization.py index 45275f8f44..b52062eeae 100644 --- a/numba_dpex/tests/experimental/test_kernel_specialization.py +++ b/numba_dpex/tests/experimental/test_kernel_specialization.py @@ -7,26 +7,30 @@ import pytest from numba.core.errors import TypingError -import numba_dpex as dpex import numba_dpex.experimental as dpex_exp from numba_dpex import DpnpNdArray, float32, int64 from numba_dpex.core.exceptions import InvalidKernelSpecializationError -from numba_dpex.kernel_api import Range +from numba_dpex.core.types.kernel_api.index_space_ids import ItemType +from numba_dpex.kernel_api import Item, Range i64arrty = DpnpNdArray(ndim=1, dtype=int64, layout="C") f32arrty = DpnpNdArray(ndim=1, dtype=float32, layout="C") +item_ty = ItemType(ndim=1) -specialized_kernel1 = dpex_exp.kernel((i64arrty, i64arrty, i64arrty)) +specialized_kernel1 = dpex_exp.kernel((item_ty, i64arrty, i64arrty, i64arrty)) specialized_kernel2 = dpex_exp.kernel( - [(i64arrty, i64arrty, i64arrty), (f32arrty, f32arrty, f32arrty)] + [ + (item_ty, i64arrty, i64arrty, i64arrty), + (item_ty, f32arrty, f32arrty, f32arrty), + ] ) -def data_parallel_sum(a, b, c): +def data_parallel_sum(item: Item, a, b, c): """ Vector addition using the ``kernel`` decorator. """ - i = dpex.get_global_id(0) + i = item.get_id(0) c[i] = a[i] + b[i] @@ -46,7 +50,9 @@ def test_invalid_specialization_error(): """Test if an InvalidKernelSpecializationError is raised when attempting to specialize with NumPy arrays. """ - specialized_kernel3 = dpex_exp.kernel((int64[::1], int64[::1], int64[::1])) + specialized_kernel3 = dpex_exp.kernel( + (item_ty, int64[::1], int64[::1], int64[::1]) + ) with pytest.raises(InvalidKernelSpecializationError): specialized_kernel3(data_parallel_sum) @@ -90,11 +96,14 @@ def test_string_specialization(): """Test if NotImplementedError is raised when signature is a string""" with pytest.raises(NotImplementedError): - dpex_exp.kernel("(i64arrty, i64arrty, i64arrty)") + dpex_exp.kernel("(item_ty, i64arrty, i64arrty, i64arrty)") with pytest.raises(NotImplementedError): dpex_exp.kernel( - ["(i64arrty, i64arrty, i64arrty)", "(f32arrty, f32arrty, f32arrty)"] + [ + "(item_ty, i64arrty, i64arrty, i64arrty)", + "(item_ty, f32arrty, f32arrty, f32arrty)", + ] ) with pytest.raises(ValueError): diff --git a/numba_dpex/tests/experimental/test_strided_dpnp_array_in_kernel.py b/numba_dpex/tests/experimental/test_strided_dpnp_array_in_kernel.py index 8be4b8e45f..90e48356ba 100644 --- a/numba_dpex/tests/experimental/test_strided_dpnp_array_in_kernel.py +++ b/numba_dpex/tests/experimental/test_strided_dpnp_array_in_kernel.py @@ -8,9 +8,8 @@ import numpy as np import pytest -import numba_dpex import numba_dpex.experimental as exp_dpex -from numba_dpex import Range +from numba_dpex.kernel_api import Item, Range def get_order(a): @@ -36,14 +35,14 @@ def get_order(a): @exp_dpex.kernel -def change_values_1d(x, v): +def change_values_1d(item: Item, x, v): """Assign values in a 1d dpnp.ndarray Args: x (dpnp.ndarray): Input array. v (int): Value to be assigned. """ - i = numba_dpex.get_global_id(0) + i = item.get_id(0) x[i] = v @@ -59,15 +58,15 @@ def change_values_1d_func(a, p): @exp_dpex.kernel -def change_values_2d(x, v): +def change_values_2d(item: Item, x, v): """Assign values in a 2d dpnp.ndarray Args: x (dpnp.ndarray): Input array. v (int): Value to be assigned. """ - i = numba_dpex.get_global_id(0) - j = numba_dpex.get_global_id(1) + i = item.get_id(0) + j = item.get_id(1) x[i, j] = v @@ -84,16 +83,16 @@ def change_values_2d_func(a, p): @exp_dpex.kernel -def change_values_3d(x, v): +def change_values_3d(item: Item, x, v): """Assign values in a 3d dpnp.ndarray Args: x (dpnp.ndarray): Input array. v (int): Value to be assigned. """ - i = numba_dpex.get_global_id(0) - j = numba_dpex.get_global_id(1) - k = numba_dpex.get_global_id(2) + i = item.get_id(0) + j = item.get_id(1) + k = item.get_id(2) x[i, j, k] = v diff --git a/numba_dpex/tests/experimental/test_supported_array_types_as_kernel_args.py b/numba_dpex/tests/experimental/test_supported_array_types_as_kernel_args.py index 7520a37373..3c4c1be6d2 100644 --- a/numba_dpex/tests/experimental/test_supported_array_types_as_kernel_args.py +++ b/numba_dpex/tests/experimental/test_supported_array_types_as_kernel_args.py @@ -8,8 +8,8 @@ import dpnp import pytest -import numba_dpex as dpex import numba_dpex.experimental as dpex_exp +from numba_dpex.kernel_api import Item, Range from numba_dpex.tests._helper import get_all_dtypes list_of_dtypes = get_all_dtypes( @@ -28,12 +28,12 @@ def input_array(request): @dpex_exp.kernel -def set_ones(a): - i = dpex.get_global_id(0) +def set_ones(item: Item, a): + i = item.get_id(0) a[i] = 1 def test_fetch_add(input_array): - dpex_exp.call_kernel(set_ones, dpex.Range(_SIZE), input_array) + dpex_exp.call_kernel(set_ones, Range(_SIZE), input_array) assert input_array[0] == 1 diff --git a/numba_dpex/tests/experimental/test_target_specific_overload.py b/numba_dpex/tests/experimental/test_target_specific_overload.py index b2c773b136..3dddeadc4b 100644 --- a/numba_dpex/tests/experimental/test_target_specific_overload.py +++ b/numba_dpex/tests/experimental/test_target_specific_overload.py @@ -5,13 +5,13 @@ import dpnp from numba.core.extending import overload -import numba_dpex as dpex import numba_dpex.experimental as dpex_exp from numba_dpex.core.descriptor import dpex_kernel_target from numba_dpex.experimental.target import ( DPEX_KERNEL_EXP_TARGET_NAME, dpex_exp_kernel_target, ) +from numba_dpex.kernel_api import Item, Range def scalar_add(a, b): @@ -27,8 +27,8 @@ def ol_scalar_add_impl(a, b): @dpex_exp.kernel -def kernel_calling_overload(a, b, c): - i = dpex.get_global_id(0) +def kernel_calling_overload(item: Item, a, b, c): + i = item.get_id(0) c[i] = scalar_add(a[i], b[i]) @@ -36,7 +36,7 @@ def kernel_calling_overload(a, b, c): b = dpnp.ones(10, dtype=dpnp.int64) c = dpnp.zeros(10, dtype=dpnp.int64) -dpex_exp.call_kernel(kernel_calling_overload, dpex.Range(10), a, b, c) +dpex_exp.call_kernel(kernel_calling_overload, Range(10), a, b, c) def test_end_to_end_overload_execution():