feat: add support for math.exp2 (#541)

Adding support for `math.exp2`. Follow up of
numba/numba#10276.

---------

Co-authored-by: Graham Markall <[email protected]>

Author: kaeun97

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/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/cudamath.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import math
 from numba.cuda import types
 from numba.cuda.typing.templates import ConcreteTemplate, signature, Registry
@@ -58,6 +59,10 @@ class Math_unary_with_fp16(ConcreteTemplate):
     ]
 
 
+if sys.version_info >= (3, 11):
+    Math_unary_with_fp16 = infer_global(math.exp2)(Math_unary_with_fp16)
+
+
 @infer_global(math.atan2)
 class Math_atan2(ConcreteTemplate):
     key = math.atan2

numba_cuda/numba/cuda/fp16.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import numba.cuda.types as types
 from numba.cuda._internal.cuda_fp16 import (
     typing_registry,
@@ -190,6 +191,13 @@ def exp_ol(a):
     return _make_unary(a, hexp)
 
 
+if sys.version_info >= (3, 11):
+
+    @overload(math.exp2, target="cuda")
+    def exp2_ol(a):
+        return _make_unary(a, hexp2)
+
+
 @overload(math.tanh, target="cuda")
 def tanh_ol(a):
     return _make_unary(a, htanh)

numba_cuda/numba/cuda/mathimpl.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import math
 import operator
 from llvmlite import ir
@@ -25,6 +26,8 @@
 unarys += [("floor", "floorf", math.floor)]
 unarys += [("fabs", "fabsf", math.fabs)]
 unarys += [("exp", "expf", math.exp)]
+if sys.version_info >= (3, 11):
+    unarys += [("exp2", "exp2f", math.exp2)]
 unarys += [("expm1", "expm1f", math.expm1)]
 unarys += [("erf", "erff", math.erf)]
 unarys += [("erfc", "erfcf", math.erfc)]
@@ -330,6 +333,7 @@ def tanhf_impl_fastmath():
 impl_unary_int(math.tanh, int64, libdevice.tanh)
 impl_unary_int(math.tanh, uint64, libdevice.tanh)
 
+
 # Complex power implementations - translations of _Py_c_pow from CPython
 # https://github.com/python/cpython/blob/a755410e054e1e2390de5830befc08fe80706c66/Objects/complexobject.c#L123-L151
 #

numba_cuda/numba/cuda/tests/cudapy/test_bfloat16.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import numpy as np
 from ml_dtypes import bfloat16 as mldtypes_bf16
 from numba import cuda
@@ -134,12 +135,8 @@ def test_math_bindings(self):
         self.skip_unsupported()
 
         exp_functions = [math.exp]
-        try:
-            from math import exp2
-
-            exp_functions += [exp2]
-        except ImportError:
-            pass
+        if sys.version_info >= (3, 11):
+            exp_functions += [math.exp2]
 
         functions = [
             math.trunc,

numba_cuda/numba/cuda/tests/cudapy/test_fastmath.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 from typing import List
 from dataclasses import dataclass, field
 from numba import cuda
@@ -142,6 +143,19 @@ def test_expf(self):
             ),
         )
 
+    @unittest.skipUnless(sys.version_info >= (3, 11), "Python 3.11+ required")
+    def test_exp2f(self):
+        from math import exp2
+
+        self._test_fast_math_unary(
+            exp2,
+            FastMathCriterion(
+                fast_expected=["ex2.approx.ftz.f32 "],
+                prec_expected=["ex2.approx.f32 "],
+                prec_unexpected=["ex2.approx.ftz.f32 "],
+            ),
+        )
+
     def test_logf(self):
         # Look for constant used to convert from log base 2 to log base e
         self._test_fast_math_unary(

numba_cuda/numba/cuda/tests/cudapy/test_math.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import numpy as np
 from numba.cuda.testing import (
     skip_unless_cc_53,
@@ -84,6 +85,11 @@ def math_exp(A, B):
     B[i] = math.exp(A[i])
 
 
+def math_exp2(A, B):
+    i = cuda.grid(1)
+    B[i] = math.exp2(A[i])
+
+
 def math_erf(A, B):
     i = cuda.grid(1)
     B[i] = math.erf(A[i])
@@ -401,6 +407,8 @@ def test_math_fp16(self):
         self.unary_template_float16(math_sqrt, np.sqrt)
         self.unary_template_float16(math_ceil, np.ceil)
         self.unary_template_float16(math_floor, np.floor)
+        if sys.version_info >= (3, 11):
+            self.unary_template_float16(math_exp2, np.exp2)
 
     @skip_on_cudasim("numpy does not support trunc for float16")
     @skip_unless_cc_53
@@ -496,6 +504,16 @@ def test_math_exp(self):
         self.unary_template_int64(math_exp, np.exp)
         self.unary_template_uint64(math_exp, np.exp)
 
+    # ---------------------------------------------------------------------------
+    # test_math_exp2
+
+    @unittest.skipUnless(sys.version_info >= (3, 11), "Python 3.11+ required")
+    def test_math_exp2(self):
+        self.unary_template_float32(math_exp2, np.exp2)
+        self.unary_template_float64(math_exp2, np.exp2)
+        self.unary_template_int64(math_exp2, np.exp2)
+        self.unary_template_uint64(math_exp2, np.exp2)
+
     # ---------------------------------------------------------------------------
     # test_math_expm1
 

numba_cuda/numba/cuda/typing/mathdecl.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import math
 from numba.cuda import types
 from numba.cuda.typing.templates import ConcreteTemplate, signature, Registry
@@ -44,6 +45,10 @@ class Math_unary(ConcreteTemplate):
     ]
 
 
+if sys.version_info >= (3, 11):
+    Math_unary = infer_global(math.exp2)(Math_unary)
+
+
 @infer_global(math.atan2)
 class Math_atan2(ConcreteTemplate):
     cases = [