Skip to content

Commit

Permalink
Unit test for IntEnumLiteral compilation
Browse files Browse the repository at this point in the history
  • Loading branch information
Diptorup Deb committed Nov 25, 2023
1 parent a5409d7 commit 72177b8
Showing 1 changed file with 1 addition and 44 deletions.
45 changes: 1 addition & 44 deletions numba_dpex/tests/experimental/IntEnumLiteral/test_compilation.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,53 +29,10 @@ def update_with_flag(a):

def test_compilation_of_flag_enum():
"""Tests if a FlagEnum subclass can be used inside a kernel function."""
a = dpnp.ones(10)
a = dpnp.ones(10, dpnp.int64)
exp_dpex.call_kernel(update_with_flag, Range(10), a)

assert a[0] == MockFlags.FLAG1
assert a[1] == MockFlags.FLAG2
for idx in range(2, 9):
assert a[idx] == 1


def test_compilation_of_flag_enum_as_a_literal():
"""Tests if a FlagEnum subclass can be used inside a kernel function."""

# pylint: disable=W0613
def _inspect_kernel(kernel_fn, index_space, *kernel_args):
pass

@intrinsic
def intrin_inspect_kernel(
typingctx, kernel_fn, index_space, kernel_args # pylint: disable=W0613
):
kernel_args_list = list(kernel_args)
# signature of the kernel_fn
kernel_sig = types.void(*kernel_args_list)
kernel_fn.dispatcher.compile(kernel_sig)
kcres = kernel_fn.dispatcher.get_overload_kcres(kernel_sig)
breakpoint()

@overload(_inspect_kernel, target="cpu")
def _ol_inspect_kernel(kernel_fn, index_space, *kernel_args):
def impl(kernel_fn, index_space, *kernel_args):
intrin_inspect_kernel( # pylint: disable=E1120
kernel_fn, index_space, kernel_args
)

return impl

@dpjit
def inspect_kernel(kernel_fn, index_space, *kernel_args):
"""An overload that is used to compile a kernel and inspect the
generated LLVM IR.
"""
_inspect_kernel(kernel_fn, index_space, *kernel_args)

a = dpnp.ones(10, dtype=dpnp.int64)
inspect_kernel(update_with_flag, Range(10), a)

assert a[0] == MockFlags.FLAG1
assert a[1] == MockFlags.FLAG2
for idx in range(2, 9):
assert a[idx] == 1

0 comments on commit 72177b8

Please sign in to comment.