Adds a new literal type to store IntEnum as Literal types.

    - Adds a new IntEnumLiteral type with corresponding
      data model into the DpexExpKernelTargetContext.
      The type is used to pass in or define an IntEnum flag as
      an Integer literal inside a kernel function.

Author: Diptorup Deb

numba_dpex/core/exceptions.py

@@ -386,3 +386,15 @@ def __init__(self, extra_msg=None) -> None:
         if extra_msg:
             self.message += " due to " + extra_msg
         super().__init__(self.message)
+
+
+class IllegalIntEnumLiteralValueError(Exception):
+    """Exception raised when an IntEnumLiteral is attempted to be created from
+    a non FlagEnum attribute.
+    """
+
+    def __init__(self) -> None:
+        self.message = (
+            "An IntEnumLiteral can only be initialized from a FlagEnum member"
+        )
+        super().__init__(self.message)

numba_dpex/experimental/__init__.py

@@ -11,6 +11,7 @@
 from .decorators import kernel
 from .kernel_dispatcher import KernelDispatcher
 from .launcher import call_kernel, call_kernel_async
+from .literal_intenum_type import IntEnumLiteral
 from .models import *
 from .types import KernelDispatcherType
 
@@ -26,4 +27,10 @@ def dpex_dispatcher_const(context):
     return context.get_dummy_value()
 
 
-__all__ = ["kernel", "KernelDispatcher", "call_kernel", "call_kernel_async"]
+__all__ = [
+    "kernel",
+    "call_kernel",
+    "call_kernel_async",
+    "IntEnumLiteral",
+    "KernelDispatcher",
+]

numba_dpex/experimental/flag_enum.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Provides a FlagEnum class to help distinguish IntEnum types that numba-dpex
+intends to use as Integer literal types inside the compiler type inferring
+infrastructure.
+"""
+from enum import IntEnum
+
+
+class FlagEnum(IntEnum):
+    """Helper class to distinguish IntEnum types that numba-dpex should consider
+    as Numba Literal types.
+    """
+
+    @classmethod
+    def basetype(cls) -> int:
+        """Returns an dummy int object that helps numba-dpex infer the type of
+        an instance of a FlagEnum class.
+
+        Returns:
+            int: Dummy int value
+        """
+        return int(0)

numba_dpex/experimental/literal_intenum_type.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Definition of a new Literal type in numba-dpex that allows treating IntEnum
+members as integer literals inside a JIT compiled function.
+"""
+from enum import IntEnum
+
+from numba.core.pythonapi import box
+from numba.core.typeconv import Conversion
+from numba.core.types import Integer, Literal
+from numba.core.typing.typeof import typeof
+
+from numba_dpex.core.exceptions import IllegalIntEnumLiteralValueError
+from numba_dpex.experimental.flag_enum import FlagEnum
+
+
+class IntEnumLiteral(Literal, Integer):
+    """A Literal type for IntEnum objects. The type contains the original Python
+    value of the IntEnum class in it.
+    """
+
+    #  pylint: disable=W0231
+    def __init__(self, value):
+        self._literal_init(value)
+        self.name = f"Literal[IntEnum]({value})"
+        if issubclass(value, FlagEnum):
+            basetype = typeof(value.basetype())
+            Integer.__init__(
+                self,
+                name=self.name,
+                bitwidth=basetype.bitwidth,
+                signed=basetype.signed,
+            )
+        else:
+            raise IllegalIntEnumLiteralValueError
+
+    def can_convert_to(self, typingctx, other) -> bool:
+        conv = typingctx.can_convert(self.literal_type, other)
+        if conv is not None:
+            return max(conv, Conversion.promote)
+        return False
+
+
+Literal.ctor_map[IntEnum] = IntEnumLiteral
+
+
+@box(IntEnumLiteral)
+def box_literal_integer(typ, val, c):
+    """Defines how a Numba representation for an IntEnumLiteral object should
+    be converted to a PyObject* object and returned back to Python.
+    """
+    val = c.context.cast(c.builder, val, typ, typ.literal_type)
+    return c.box(typ.literal_type, val)

numba_dpex/experimental/models.py

@@ -6,14 +6,28 @@
 numba_dpex.experimental module.
 """
 
+from llvmlite import ir as llvmir
 from numba.core.datamodel import DataModelManager, models
+from numba.core.datamodel.models import PrimitiveModel
 from numba.core.extending import register_model
 
 import numba_dpex.core.datamodel.models as dpex_core_models
 
+from .literal_intenum_type import IntEnumLiteral
 from .types import KernelDispatcherType
 
 
+class LiteralIntEnumModel(PrimitiveModel):
+    """Representation of an object of LiteralIntEnum type using Numba's
+    PrimitiveModel that can be represented natively in the target in all
+    usage contexts.
+    """
+
+    def __init__(self, dmm, fe_type):
+        be_type = llvmir.IntType(fe_type.bitwidth)
+        super().__init__(dmm, fe_type, be_type)
+
+
 def _init_exp_data_model_manager() -> DataModelManager:
     """Initializes a DpexExpKernelTarget-specific data model manager.
 
@@ -28,7 +42,7 @@ def _init_exp_data_model_manager() -> DataModelManager:
     dmm = dpex_core_models.dpex_data_model_manager.copy()
 
     # Register the types and data model in the DpexExpTargetContext
-    # Add here...
+    dmm.register(IntEnumLiteral, LiteralIntEnumModel)
 
     return dmm
 

numba_dpex/experimental/target.py

@@ -8,8 +8,11 @@
 
 from functools import cached_property
 
+from llvmlite import ir as llvmir
+from numba.core import types
 from numba.core.descriptors import TargetDescriptor
 from numba.core.target_extension import GPU, target_registry
+from numba.core.types.scalars import IntEnumClass
 
 from numba_dpex.core.descriptor import DpexTargetOptions
 from numba_dpex.core.targets.kernel_target import (
@@ -18,6 +21,9 @@
 )
 from numba_dpex.experimental.models import exp_dmm
 
+from .flag_enum import FlagEnum
+from .literal_intenum_type import IntEnumLiteral
+
 
 #  pylint: disable=R0903
 class SyclDeviceExp(GPU):
@@ -39,6 +45,37 @@ class DpexExpKernelTypingContext(DpexKernelTypingContext):
     are stable enough to be migrated to DpexKernelTypingContext.
     """
 
+    def resolve_value_type(self, val):
+        """
+        Return the numba type of a Python value that is being used
+        as a runtime constant.
+        ValueError is raised for unsupported types.
+        """
+
+        ty = super().resolve_value_type(val)
+
+        if isinstance(ty, IntEnumClass) and issubclass(val, FlagEnum):
+            ty = IntEnumLiteral(val)
+
+        return ty
+
+    def resolve_getattr(self, typ, attr):
+        """
+        Resolve getting the attribute *attr* (a string) on the Numba type.
+        The attribute's type is returned, or None if resolution failed.
+        """
+        ty = None
+
+        if isinstance(typ, IntEnumLiteral):
+            try:
+                attrval = getattr(typ.literal_value, attr).value
+                ty = types.IntegerLiteral(attrval)
+            except ValueError:
+                pass
+        else:
+            ty = super().resolve_getattr(typ, attr)
+        return ty
+
 
 #  pylint: disable=W0223
 # FIXME: Remove the pylint disablement once we add an override for
@@ -56,6 +93,22 @@ def __init__(self, typingctx, target=DPEX_KERNEL_EXP_TARGET_NAME):
         super().__init__(typingctx, target)
         self.data_model_manager = exp_dmm
 
+    def get_getattr(self, typ, attr):
+        """
+        Overrides the get_getattr function to provide an implementation for
+        getattr call on an IntegerEnumLiteral type.
+        """
+
+        if isinstance(typ, IntEnumLiteral):
+            #  pylint: disable=W0613
+            def enum_literal_getattr_imp(context, builder, typ, val, attr):
+                enum_attr_value = getattr(typ.literal_value, attr).value
+                return llvmir.Constant(llvmir.IntType(64), enum_attr_value)
+
+            return enum_literal_getattr_imp
+
+        return super().get_getattr(typ, attr)
+
 
 class DpexExpKernelTarget(TargetDescriptor):
     """