merge/resolve

Author: brandon-b-miller

docs/source/user/cudapysupported.rst

@@ -214,6 +214,7 @@ The following functions from the :mod:`math` module are supported:
 * :func:`math.erf`
 * :func:`math.erfc`
 * :func:`math.exp`
+* :func:`math.exp2`
 * :func:`math.expm1`
 * :func:`math.fabs`
 * :func:`math.frexp`

numba_cuda/numba/cuda/__init__.py

@@ -7,6 +7,12 @@
 import warnings
 import sys
 
+# Re-export types itself
+import numba.cuda.types as types
+
+# Re-export all type names
+from numba.cuda.types import *
+
 
 # Require NVIDIA CUDA bindings at import time
 if not (

numba_cuda/numba/cuda/_internal/cuda_bf16.py

@@ -18,8 +18,8 @@
 
 import numba
 from llvmlite import ir
-from numba import types
-from numba.core.datamodel import PrimitiveModel, StructModel
+from numba.cuda import types
+from numba.cuda.datamodel import PrimitiveModel, StructModel
 from numba.cuda.extending import (
     lower_cast,
     make_attribute_wrapper,
@@ -41,7 +41,7 @@
 from numba.cuda import CUSource, declare_device
 from numba.cuda.vector_types import vector_types
 from numba.cuda.extending import as_numba_type
-from numba.types import (
+from numba.cuda.types import (
     CPointer,
     Function,
     Number,
@@ -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/_internal/cuda_fp16.py

@@ -18,9 +18,9 @@
 
 import numba
 from llvmlite import ir
-from numba import types
+from numba.cuda import types
 from numba.cuda.cudadrv.driver import _have_nvjitlink
-from numba.core.datamodel import PrimitiveModel, StructModel
+from numba.cuda.datamodel import PrimitiveModel, StructModel
 from numba.core.errors import NumbaPerformanceWarning
 from numba.cuda.extending import (
     lower_cast,
@@ -40,7 +40,7 @@
 from numba.cuda.typing.templates import Registry as TypingRegistry
 from numba.cuda.vector_types import vector_types
 from numba.cuda.extending import as_numba_type
-from numba.types import (
+from numba.cuda.types import (
     CPointer,
     Function,
     Number,
@@ -221,7 +221,7 @@ class _ctor_template_unnamed1362180(ConcreteTemplate):
 
 register_global(unnamed1362180, Function(_ctor_template_unnamed1362180))
 
-__half = _type___half = numba.core.types.float16
+__half = _type___half = numba.cuda.types.float16
 setattr(__half, "alignof_", 2)
 setattr(__half, "align", 2)
 

numba_cuda/numba/cuda/api.py

@@ -39,10 +39,9 @@ def from_cuda_array_interface(desc, owner=None, sync=True):
 
     shape = desc["shape"]
     strides = desc.get("strides")
-    dtype = np.dtype(desc["typestr"])
 
     shape, strides, dtype = prepare_shape_strides_dtype(
-        shape, strides, dtype, order="C"
+        shape, strides, desc["typestr"], order="C"
     )
     size = driver.memory_size_from_info(shape, strides, dtype.itemsize)
 

numba_cuda/numba/cuda/api_util.py

@@ -3,6 +3,8 @@
 
 import numpy as np
 
+import functools
+
 
 def prepare_shape_strides_dtype(shape, strides, dtype, order):
     dtype = np.dtype(dtype)
@@ -14,25 +16,33 @@ def prepare_shape_strides_dtype(shape, strides, dtype, order):
         raise TypeError("shape must be an integer or tuple of integers")
     if isinstance(shape, int):
         shape = (shape,)
+    else:
+        shape = tuple(shape)
     if isinstance(strides, int):
         strides = (strides,)
     else:
-        strides = strides or _fill_stride_by_order(shape, dtype, order)
+        if not strides:
+            strides = _fill_stride_by_order(shape, dtype, order)
+        else:
+            strides = tuple(strides)
     return shape, strides, dtype
 
 
+@functools.cache
 def _fill_stride_by_order(shape, dtype, order):
-    nd = len(shape)
-    if nd == 0:
+    ndims = len(shape)
+    if not ndims:
         return ()
-    strides = [0] * nd
+    strides = [0] * ndims
     if order == "C":
         strides[-1] = dtype.itemsize
-        for d in reversed(range(nd - 1)):
+        # -2 because we subtract one for zero-based indexing and another one
+        # for skipping the already-filled-in last element
+        for d in range(ndims - 2, -1, -1):
             strides[d] = strides[d + 1] * shape[d + 1]
     elif order == "F":
         strides[0] = dtype.itemsize
-        for d in range(1, nd):
+        for d in range(1, ndims):
             strides[d] = strides[d - 1] * shape[d - 1]
     else:
         raise ValueError("must be either C/F order")

numba_cuda/numba/cuda/bf16.py

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
+import sys
 
 from numba.cuda._internal.cuda_bf16 import (
     typing_registry,
@@ -191,14 +192,12 @@ def exp_ol(a):
     return _make_unary(a, hexp)
 
 
-try:
-    from math import exp2
+if sys.version_info >= (3, 11):
 
-    @overload(exp2, target="cuda")
+    @overload(math.exp2, target="cuda")
     def exp2_ol(a):
         return _make_unary(a, hexp2)
-except ImportError:
-    pass
+
 
 ## Public aliases using Numba/Numpy-style type names
 # Floating-point

numba_cuda/numba/cuda/cg.py

@@ -1,12 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
-from numba.core import types
+from numba.cuda import types
 from numba.cuda.extending import overload, overload_method
 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/cgutils.py

@@ -11,9 +11,9 @@
 
 from llvmlite import ir
 
-from numba.core import types
+from numba.cuda import types
 from numba.cuda import config, utils, debuginfo
-import numba.core.datamodel
+import numba.cuda.datamodel
 
 
 bool_t = ir.IntType(1)
@@ -104,7 +104,7 @@ class _StructProxy(object):
     def __init__(self, context, builder, value=None, ref=None):
         self._context = context
         self._datamodel = self._context.data_model_manager[self._fe_type]
-        if not isinstance(self._datamodel, numba.core.datamodel.StructModel):
+        if not isinstance(self._datamodel, numba.cuda.datamodel.StructModel):
             raise TypeError(
                 "Not a structure model: {0}".format(self._datamodel)
             )

numba_cuda/numba/cuda/compiler.py

@@ -7,10 +7,8 @@
 import copy
 
 from numba.core import ir as numba_ir
-from numba.core import (
-    types,
-    bytecode,
-)
+from numba.core import bytecode
+from numba.cuda import types
 from numba.cuda.core.options import ParallelOptions
 from numba.core.compiler_lock import global_compiler_lock
 from numba.core.errors import NumbaWarning, NumbaInvalidConfigWarning