Merge branch 'main' of github.com:NVIDIA/numba-cuda into ajm/vendor-imputils

Author: ashermancinelli

numba_cuda/numba/cuda/_internal/cuda_fp16.py

@@ -124,7 +124,7 @@ def __init__(self):
         self.bitwidth = 2 * 8
 
     def can_convert_from(self, typingctx, other):
-        from numba.core.typeconv import Conversion
+        from numba.cuda.typeconv import Conversion
 
         if other in []:
             return Conversion.safe
@@ -174,7 +174,7 @@ def __init__(self):
         self.bitwidth = 4 * 8
 
     def can_convert_from(self, typingctx, other):
-        from numba.core.typeconv import Conversion
+        from numba.cuda.typeconv import Conversion
 
         if other in []:
             return Conversion.safe
@@ -7903,9 +7903,9 @@ def generic(self, args, kws):
                 #  - Conversion.safe
 
                 if (
-                    (convertible == numba.core.typeconv.Conversion.exact)
-                    or (convertible == numba.core.typeconv.Conversion.promote)
-                    or (convertible == numba.core.typeconv.Conversion.safe)
+                    (convertible == numba.cuda.typeconv.Conversion.exact)
+                    or (convertible == numba.cuda.typeconv.Conversion.promote)
+                    or (convertible == numba.cuda.typeconv.Conversion.safe)
                 ):
                     return signature(retty, types.float16, types.float16)
 

numba_cuda/numba/cuda/cext/_helperlib.c

@@ -0,0 +1,71 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+/*
+ * Helper functions used by Numba CUDA at runtime.
+ * This C file is meant to be included after defining the
+ * NUMBA_EXPORT_FUNC() and NUMBA_EXPORT_DATA() macros.
+ */
+
+#include "_pymodule.h"
+#include <stddef.h>
+
+/*
+ * Unicode helpers
+ */
+
+/* Developer note:
+ *
+ * The hash value of unicode objects is obtained via:
+ * ((PyASCIIObject *)(obj))->hash;
+ * The use comes from this definition:
+ * https://github.com/python/cpython/blob/6d43f6f081023b680d9db4542d19b9e382149f0a/Objects/unicodeobject.c#L119-L120
+ * and it's used extensively throughout the `cpython/Object/unicodeobject.c`
+ * source, not least in `unicode_hash` itself:
+ * https://github.com/python/cpython/blob/6d43f6f081023b680d9db4542d19b9e382149f0a/Objects/unicodeobject.c#L11662-L11679
+ *
+ * The Unicode string struct layouts are described here:
+ * https://github.com/python/cpython/blob/6d43f6f081023b680d9db4542d19b9e382149f0a/Include/cpython/unicodeobject.h#L82-L161
+ * essentially, all the unicode string layouts start with a `PyASCIIObject` at
+ * offset 0 (as of commit 6d43f6f081023b680d9db4542d19b9e382149f0a, somewhere
+ * in the 3.8 development cycle).
+ *
+ * For safety against future CPython internal changes, the code checks that the
+ * _base members of the unicode structs are what is expected in 3.7, and that
+ * their offset is 0. It then walks the struct to the hash location to make sure
+ * the offset is indeed the same as PyASCIIObject->hash.
+ * Note: The large condition in the if should evaluate to a compile time
+ * constant.
+ */
+
+#define MEMBER_SIZE(structure, member) sizeof(((structure *)0)->member)
+
+NUMBA_EXPORT_FUNC(void *)
+numba_extract_unicode(PyObject *obj, Py_ssize_t *length, int *kind,
+                      unsigned int *ascii, Py_ssize_t *hash) {
+    if (!PyUnicode_READY(obj)) {
+        *length = PyUnicode_GET_LENGTH(obj);
+        *kind = PyUnicode_KIND(obj);
+        /* could also use PyUnicode_IS_ASCII but it is not publicly advertised in https://docs.python.org/3/c-api/unicode.html */
+        *ascii = (unsigned int)(PyUnicode_MAX_CHAR_VALUE(obj) == (0x7f));
+        /* this is here as a crude check for safe casting of all unicode string
+         * structs to a PyASCIIObject */
+        if (MEMBER_SIZE(PyCompactUnicodeObject, _base) == sizeof(PyASCIIObject)             &&
+            MEMBER_SIZE(PyUnicodeObject, _base) == sizeof(PyCompactUnicodeObject)           &&
+            offsetof(PyCompactUnicodeObject, _base) == 0                                    &&
+            offsetof(PyUnicodeObject, _base) == 0                                           &&
+            offsetof(PyCompactUnicodeObject, _base.hash) == offsetof(PyASCIIObject, hash)   &&
+            offsetof(PyUnicodeObject, _base._base.hash) == offsetof(PyASCIIObject, hash)
+           ) {
+            /* Grab the hash from the type object cache, do not compute it. */
+            *hash = ((PyASCIIObject *)(obj))->hash;
+        }
+        else {
+            /* cast is not safe, fail */
+            return NULL;
+        }
+        return PyUnicode_DATA(obj);
+    } else {
+        return NULL;
+    }
+}

numba_cuda/numba/cuda/cext/_helpermod.c

@@ -0,0 +1,82 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+/*
+ * Expose all functions as pointers in a dedicated C extension.
+ */
+
+/* Import _pymodule.h first, for a recent _POSIX_C_SOURCE */
+#include "_pymodule.h"
+
+/* Visibility control macros */
+#if defined(_WIN32) || defined(_WIN64)
+    #define VISIBILITY_HIDDEN
+    #define VISIBILITY_GLOBAL __declspec(dllexport)
+#else
+    #define VISIBILITY_HIDDEN __attribute__((visibility("hidden")))
+    #define VISIBILITY_GLOBAL __attribute__((visibility("default")))
+#endif
+
+/* Define all runtime-required symbols in this C module, but do not
+   export them outside the shared library if possible. */
+#define NUMBA_EXPORT_FUNC(_rettype) VISIBILITY_HIDDEN _rettype
+#define NUMBA_EXPORT_DATA(_vartype) VISIBILITY_HIDDEN _vartype
+
+/* Numba CUDA C helpers */
+#include "_helperlib.c"
+
+static PyObject *
+build_c_helpers_dict(void)
+{
+    PyObject *dct = PyDict_New();
+    if (dct == NULL)
+        goto error;
+
+#define _declpointer(name, value) do {                 \
+    PyObject *o = PyLong_FromVoidPtr(value);           \
+    if (o == NULL) goto error;                         \
+    if (PyDict_SetItemString(dct, name, o)) {          \
+        Py_DECREF(o);                                  \
+        goto error;                                    \
+    }                                                  \
+    Py_DECREF(o);                                      \
+} while (0)
+
+#define declmethod(func) _declpointer(#func, &numba_##func)
+
+    /* Unicode string support */
+    declmethod(extract_unicode);
+
+#undef declmethod
+    return dct;
+error:
+    Py_XDECREF(dct);
+    return NULL;
+}
+
+static PyMethodDef ext_methods[] = {
+    { NULL },
+};
+
+MOD_INIT(_helperlib) {
+    PyObject *m;
+    MOD_DEF(m, "_helperlib", "No docs", ext_methods)
+    if (m == NULL)
+        return MOD_ERROR_VAL;
+
+    PyModule_AddObject(m, "c_helpers", build_c_helpers_dict());
+    PyModule_AddIntConstant(m, "long_min", LONG_MIN);
+    PyModule_AddIntConstant(m, "long_max", LONG_MAX);
+    PyModule_AddIntConstant(m, "py_buffer_size", sizeof(Py_buffer));
+    PyModule_AddIntConstant(m, "py_gil_state_size", sizeof(PyGILState_STATE));
+    PyModule_AddIntConstant(m, "py_unicode_1byte_kind", PyUnicode_1BYTE_KIND);
+    PyModule_AddIntConstant(m, "py_unicode_2byte_kind", PyUnicode_2BYTE_KIND);
+    PyModule_AddIntConstant(m, "py_unicode_4byte_kind", PyUnicode_4BYTE_KIND);
+#if (PY_MAJOR_VERSION == 3)
+#if ((PY_MINOR_VERSION == 10) || (PY_MINOR_VERSION == 11))
+    PyModule_AddIntConstant(m, "py_unicode_wchar_kind", PyUnicode_WCHAR_KIND);
+#endif
+#endif
+
+    return MOD_SUCCESS_VAL(m);
+}

numba_cuda/numba/cuda/cext/_typeconv.cpp

@@ -0,0 +1,206 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+#include "_pymodule.h"
+#include "capsulethunk.h"
+#include "typeconv.hpp"
+
+extern "C" {
+
+
+static PyObject*
+new_type_manager(PyObject* self, PyObject* args);
+
+static void
+del_type_manager(PyObject *);
+
+static PyObject*
+select_overload(PyObject* self, PyObject* args);
+
+static PyObject*
+check_compatible(PyObject* self, PyObject* args);
+
+static PyObject*
+set_compatible(PyObject* self, PyObject* args);
+
+static PyObject*
+get_pointer(PyObject* self, PyObject* args);
+
+
+static PyMethodDef ext_methods[] = {
+#define declmethod(func) { #func , ( PyCFunction )func , METH_VARARGS , NULL }
+    declmethod(new_type_manager),
+    declmethod(select_overload),
+    declmethod(check_compatible),
+    declmethod(set_compatible),
+    declmethod(get_pointer),
+    { NULL },
+#undef declmethod
+};
+
+
+MOD_INIT(_typeconv) {
+    PyObject *m;
+    MOD_DEF(m, "_typeconv", "No docs", ext_methods)
+    if (m == NULL)
+        return MOD_ERROR_VAL;
+
+    return MOD_SUCCESS_VAL(m);
+}
+
+} // end extern C
+
+///////////////////////////////////////////////////////////////////////////////
+
+const char PY_CAPSULE_TM_NAME[] = "*tm";
+#define BAD_TM_ARGUMENT PyErr_SetString(PyExc_TypeError,                    \
+                                        "1st argument not TypeManager")
+
+static
+TypeManager* unwrap_TypeManager(PyObject *tm) {
+    void* p = PyCapsule_GetPointer(tm, PY_CAPSULE_TM_NAME);
+    return reinterpret_cast<TypeManager*>(p);
+}
+
+PyObject*
+new_type_manager(PyObject* self, PyObject* args)
+{
+    TypeManager* tm = new TypeManager();
+    return PyCapsule_New(tm, PY_CAPSULE_TM_NAME, &del_type_manager);
+}
+
+void
+del_type_manager(PyObject *tm)
+{
+    delete unwrap_TypeManager(tm);
+}
+
+PyObject*
+select_overload(PyObject* self, PyObject* args)
+{
+    PyObject *tmcap, *sigtup, *ovsigstup;
+    int allow_unsafe;
+    int exact_match_required;
+
+    if (!PyArg_ParseTuple(args, "OOOii", &tmcap, &sigtup, &ovsigstup,
+                          &allow_unsafe, &exact_match_required)) {
+        return NULL;
+    }
+
+    TypeManager *tm = unwrap_TypeManager(tmcap);
+    if (!tm) {
+        BAD_TM_ARGUMENT;
+    }
+
+    Py_ssize_t sigsz = PySequence_Size(sigtup);
+    Py_ssize_t ovsz = PySequence_Size(ovsigstup);
+
+    Type *sig = new Type[sigsz];
+    Type *ovsigs = new Type[ovsz * sigsz];
+
+    for (int i = 0; i < sigsz; ++i) {
+        sig[i] = Type(PyNumber_AsSsize_t(PySequence_Fast_GET_ITEM(sigtup,
+                                                                  i), NULL));
+    }
+
+    for (int i = 0; i < ovsz; ++i) {
+        PyObject *cursig = PySequence_Fast_GET_ITEM(ovsigstup, i);
+        for (int j = 0; j < sigsz; ++j) {
+            long tid = PyNumber_AsSsize_t(PySequence_Fast_GET_ITEM(cursig,
+                                                                   j), NULL);
+            ovsigs[i * sigsz + j] = Type(tid);
+        }
+    }
+
+    int selected = -42;
+    int matches = tm->selectOverload(sig, ovsigs, selected, sigsz, ovsz,
+                                     (bool) allow_unsafe,
+                                     (bool) exact_match_required);
+
+    delete [] sig;
+    delete [] ovsigs;
+
+    if (matches > 1) {
+        PyErr_SetString(PyExc_TypeError, "Ambiguous overloading");
+        return NULL;
+    } else if (matches == 0) {
+        PyErr_SetString(PyExc_TypeError, "No compatible overload");
+        return NULL;
+    }
+
+    return PyLong_FromLong(selected);
+}
+
+PyObject*
+check_compatible(PyObject* self, PyObject* args)
+{
+    PyObject *tmcap;
+    int from, to;
+    if (!PyArg_ParseTuple(args, "Oii", &tmcap, &from, &to)) {
+        return NULL;
+    }
+
+    TypeManager *tm = unwrap_TypeManager(tmcap);
+    if(!tm) {
+        BAD_TM_ARGUMENT;
+        return NULL;
+    }
+
+    switch(tm->isCompatible(Type(from), Type(to))){
+    case TCC_EXACT:
+        return PyString_FromString("exact");
+    case TCC_PROMOTE:
+        return PyString_FromString("promote");
+    case TCC_CONVERT_SAFE:
+        return PyString_FromString("safe");
+        case TCC_CONVERT_UNSAFE:
+        return PyString_FromString("unsafe");
+    default:
+        Py_RETURN_NONE;
+    }
+}
+
+PyObject*
+set_compatible(PyObject* self, PyObject* args)
+{
+    PyObject *tmcap;
+    int from, to, by;
+    if (!PyArg_ParseTuple(args, "Oiii", &tmcap, &from, &to, &by)) {
+        return NULL;
+    }
+
+    TypeManager *tm = unwrap_TypeManager(tmcap);
+    if (!tm) {
+        BAD_TM_ARGUMENT;
+        return NULL;
+    }
+    TypeCompatibleCode tcc;
+    switch (by) {
+    case 'p': // promote
+        tcc = TCC_PROMOTE;
+        break;
+    case 's': // safe convert
+        tcc = TCC_CONVERT_SAFE;
+        break;
+    case 'u': // unsafe convert
+        tcc = TCC_CONVERT_UNSAFE;
+        break;
+    default:
+        PyErr_SetString(PyExc_ValueError, "Unknown TCC");
+        return NULL;
+    }
+
+    tm->addCompatibility(Type(from), Type(to), tcc);
+    Py_RETURN_NONE;
+}
+
+
+PyObject*
+get_pointer(PyObject* self, PyObject* args)
+{
+    PyObject *tmcap;
+    if (!PyArg_ParseTuple(args, "O", &tmcap)) {
+        return NULL;
+    }
+    return PyLong_FromVoidPtr(unwrap_TypeManager(tmcap));
+}