Redirect to numba.types and numba.core.datamodel if they are available

Author: VijayKandiah

numba_cuda/numba/cuda/_internal/cuda_bf16.py

@@ -60,7 +60,7 @@
     uint64,
     void,
 )
-from numba.cuda.types import bfloat16
+from numba.cuda.ext_types import bfloat16
 
 float32x2 = vector_types["float32x2"]
 __half = float16

numba_cuda/numba/cuda/cg.py

@@ -6,7 +6,7 @@
 from numba.cuda.typing import signature
 from numba.cuda import nvvmutils
 from numba.cuda.extending import intrinsic
-from numba.cuda.types import grid_group, GridGroup as GridGroupClass
+from numba.cuda.ext_types import grid_group, GridGroup as GridGroupClass
 
 
 class GridGroup:

numba_cuda/numba/cuda/core/typeinfer.py

@@ -51,7 +51,6 @@
 )
 from numba.cuda.core.funcdesc import qualifying_prefix
 from numba.cuda.typeconv import Conversion
-from numba.cuda.core.sigutils import is_numba_type, convert_to_cuda_type
 
 _logger = logging.getLogger(__name__)
 
@@ -75,19 +74,8 @@ def __init__(self, context, var):
         # Qualifiers
         self.literal_value = NOTSET
 
-    def _ensure_cuda_type(self, tp):
-        """
-        Convert numba.core types to numba.cuda types if necessary.
-        This ensures cross-compatibility.
-        """
-
-        if is_numba_type(tp):
-            tp = convert_to_cuda_type(tp)
-        return tp
-
     def add_type(self, tp, loc):
-        tp = self._ensure_cuda_type(tp)
-        assert isinstance(tp, types.Type) or is_numba_type(tp), type(tp)
+        assert isinstance(tp, types.Type), type(tp)
         # Special case for _undef_var.
         # If the typevar is the _undef_var, use the incoming type directly.
         if self.type is types._undef_var:
@@ -121,8 +109,7 @@ def add_type(self, tp, loc):
         return self.type
 
     def lock(self, tp, loc, literal_value=NOTSET):
-        tp = self._ensure_cuda_type(tp)
-        assert isinstance(tp, types.Type) or is_numba_type(tp), type(tp)
+        assert isinstance(tp, types.Type), type(tp)
 
         if self.locked:
             msg = (

numba_cuda/numba/cuda/datamodel/__init__.py renamed to numba_cuda/numba/cuda/cuda_datamodel/__init__.py

@@ -165,100 +165,6 @@
 longlong = np_longlong = _make_signed(np.longlong)
 ulonglong = np_ulonglong = _make_unsigned(np.longlong)
 
-
-class Dim3(types.Type):
-    """
-    A 3-tuple (x, y, z) representing the position of a block or thread.
-    """
-
-    def __init__(self):
-        super().__init__(name="Dim3")
-
-
-class GridGroup(types.Type):
-    """
-    The grid of all threads in a cooperative kernel launch.
-    """
-
-    def __init__(self):
-        super().__init__(name="GridGroup")
-
-
-dim3 = Dim3()
-grid_group = GridGroup()
-
-
-class CUDADispatcher(types.Dispatcher):
-    """The type of CUDA dispatchers"""
-
-    # This type exists (instead of using types.Dispatcher as the type of CUDA
-    # dispatchers) so that we can have an alternative lowering for them to the
-    # lowering of CPU dispatchers - the CPU target lowers all dispatchers as a
-    # constant address, but we need to lower to a dummy value because it's not
-    # generally valid to use the address of CUDA kernels and functions.
-    #
-    # Notes: it may be a bug in the CPU target that it lowers all dispatchers to
-    # a constant address - it should perhaps only lower dispatchers acting as
-    # first-class functions to a constant address. Even if that bug is fixed, it
-    # is still probably a good idea to have a separate type for CUDA
-    # dispatchers, and this type might get other differentiation from the CPU
-    # dispatcher type in future.
-
-
-class Bfloat16(Number):
-    """
-    A bfloat16 type. Has 8 exponent bits and 7 significand bits.
-
-    Conversion rules:
-    Floats:
-    from:
-        fp32, fp64: UNSAFE
-        fp16: UNSAFE (loses precision)
-    to:
-        fp32, fp64: PROMOTE (same exponent, more mantissa)
-        fp16: UNSAFE (loses range)
-
-    Integers:
-    from:
-        int8: SAFE
-        other int: All UNSAFE (bf16 cannot represent all integers in range)
-    to: UNSAFE (loses precision, round to zeros)
-
-    All other conversions are not allowed.
-    """
-
-    def __init__(self):
-        super().__init__(name="__nv_bfloat16")
-
-        self.alignof_ = 2
-        self.bitwidth = 16
-
-    def can_convert_from(self, typingctx, other):
-        if isinstance(other, types.Float):
-            return Conversion.unsafe
-
-        elif isinstance(other, types.Integer):
-            if other.bitwidth == 8:
-                return Conversion.safe
-            else:
-                return Conversion.unsafe
-
-    def can_convert_to(self, typingctx, other):
-        if isinstance(other, types.Float):
-            if other.bitwidth >= 32:
-                return Conversion.safe
-            else:
-                return Conversion.unsafe
-        elif isinstance(other, types.Integer):
-            return Conversion.unsafe
-
-    def unify(self, typingctx, other):
-        if isinstance(other, (types.Float, types.Integer)):
-            return typingctx.unify_pairs(self, other)
-
-
-bfloat16 = Bfloat16()
-
 all_str = """
     int8
     int16
@@ -312,9 +218,6 @@ def unify(self, typingctx, other):
     ffi_forced_object
     ffi
     deferred_type
-    dim3
-    grid_group
-    bfloat16
 """