numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/cache_hints.py

@@ -0,0 +1,287 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from llvmlite import ir
+from numba import types
+from numba.cuda import cgutils
+from numba.cuda.extending import intrinsic, overload
+from numba.cuda.core.errors import NumbaTypeError
+from numba.cuda.api_util import normalize_indices
+
+# Docs references:
+# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-ld
+# https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#load-functions-using-cache-hints
+
+
+def ldca(array, i):
+    """Generate a `ld.global.ca` instruction for element `i` of an array."""
+
+
+def ldcg(array, i):
+    """Generate a `ld.global.cg` instruction for element `i` of an array."""
+
+
+def ldcs(array, i):
+    """Generate a `ld.global.cs` instruction for element `i` of an array."""
+
+
+def ldlu(array, i):
+    """Generate a `ld.global.lu` instruction for element `i` of an array."""
+
+
+def ldcv(array, i):
+    """Generate a `ld.global.cv` instruction for element `i` of an array."""
+
+
+def stcg(array, i, value):
+    """Generate a `st.global.cg` instruction for element `i` of an array."""
+
+
+def stcs(array, i, value):
+    """Generate a `st.global.cs` instruction for element `i` of an array."""
+
+
+def stwb(array, i, value):
+    """Generate a `st.global.wb` instruction for element `i` of an array."""
+
+
+def stwt(array, i, value):
+    """Generate a `st.global.wt` instruction for element `i` of an array."""
+
+
+# See
+# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#restricted-use-of-sub-word-sizes
+# for background on the choice of "r" for 8-bit operands - there is
+# no constraint for 8-bit operands, but the operand for loads and
+# stores is permitted to be greater than 8 bits.
+CONSTRAINT_MAP = {1: "b", 8: "r", 16: "h", 32: "r", 64: "l", 128: "q"}
+
+
+def _validate_arguments(instruction, array, index):
+    is_array = isinstance(array, types.Array)
+    is_pointer = isinstance(array, types.CPointer)
+    if not (is_array or is_pointer):
+        msg = f"{instruction} operates on arrays or pointers. Got type {array}"
+        raise NumbaTypeError(msg)
+
+    valid_index = False
+
+    if isinstance(index, types.Integer):
+        if is_array and array.ndim != 1:
+            # for pointers, any integer index is valid
+            msg = f"Expected {array.ndim} indices, got a scalar"
+            raise NumbaTypeError(msg)
+        valid_index = True
+
+    if isinstance(index, types.UniTuple):
+        if is_pointer:
+            msg = f"Pointers only support scalar indexing, got tuple of {index.count}"
+            raise NumbaTypeError(msg)
+
+        if index.count != array.ndim:
+            msg = f"Expected {array.ndim} indices, got {index.count}"
+            raise NumbaTypeError(msg)
+
+        if isinstance(index.dtype, types.Integer):
+            valid_index = True
+
+    if not valid_index:
+        raise NumbaTypeError(f"{index} is not a valid index")
+
+
+def _validate_bitwidth(instruction, array):
+    dtype = array.dtype
+
+    if not isinstance(dtype, (types.Integer, types.Float)):
+        msg = (
+            f"{instruction} requires array of integer or float type, "
+            f"got {dtype}"
+        )
+        raise NumbaTypeError(msg)
+
+    bitwidth = dtype.bitwidth
+    if bitwidth not in CONSTRAINT_MAP:
+        valid_widths = sorted(CONSTRAINT_MAP.keys())
+        msg = (
+            f"{instruction} requires array dtype with bitwidth "
+            f"in {valid_widths}, got bitwidth {bitwidth}"
+        )
+        raise NumbaTypeError(msg)
+
+
+def _get_element_pointer(
+    context, builder, index_type, indices, array_type, array
+):
+    if isinstance(array_type, types.CPointer):
+        return builder.gep(array, [indices])
+    else:
+        index_type, indices = normalize_indices(
+            context,
+            builder,
+            index_type,
+            indices,
+            array_type,
+            array_type.dtype,
+        )
+        array_struct = context.make_array(array_type)(
+            context, builder, value=array
+        )
+        return cgutils.get_item_pointer(
+            context,
+            builder,
+            array_type,
+            array_struct,
+            indices,
+            wraparound=True,
+        )
+
+
+def ld_cache_operator(operator):
+    @intrinsic
+    def impl(typingctx, array, index):
+        _validate_arguments(f"ld{operator}", array, index)
+        _validate_bitwidth(f"ld{operator}", array)
+
+        signature = array.dtype(array, index)
+
+        def codegen(context, builder, sig, args):
+            array_type, index_type = sig.args
+            loaded_type = context.get_value_type(array_type.dtype)
+            ptr_type = loaded_type.as_pointer()
+            ldcs_type = ir.FunctionType(loaded_type, [ptr_type])
+
+            array, indices = args
+
+            ptr = _get_element_pointer(
+                context, builder, index_type, indices, array_type, array
+            )
+
+            bitwidth = array_type.dtype.bitwidth
+            inst = f"ld.global.{operator}.b{bitwidth}"
+            constraints = f"={CONSTRAINT_MAP[bitwidth]},l"
+            ldcs = ir.InlineAsm(ldcs_type, f"{inst} $0, [$1];", constraints)
+            return builder.call(ldcs, [ptr])
+
+        return signature, codegen
+
+    return impl
+
+
+ldca_intrinsic = ld_cache_operator("ca")
+ldcg_intrinsic = ld_cache_operator("cg")
+ldcs_intrinsic = ld_cache_operator("cs")
+ldlu_intrinsic = ld_cache_operator("lu")
+ldcv_intrinsic = ld_cache_operator("cv")
+
+
+def st_cache_operator(operator):
+    @intrinsic
+    def impl(typingctx, array, index, value):
+        _validate_arguments(f"st{operator}", array, index)
+        _validate_bitwidth(f"st{operator}", array)
+
+        signature = types.void(array, index, value)
+
+        def codegen(context, builder, sig, args):
+            array_type, index_type, value_type = sig.args
+            stored_type = context.get_value_type(array_type.dtype)
+            ptr_type = stored_type.as_pointer()
+            stcs_type = ir.FunctionType(ir.VoidType(), [ptr_type, stored_type])
+
+            array, indices, value = args
+
+            ptr = _get_element_pointer(
+                context, builder, index_type, indices, array_type, array
+            )
+
+            casted_value = context.cast(
+                builder, value, value_type, array_type.dtype
+            )
+
+            bitwidth = array_type.dtype.bitwidth
+            inst = f"st.global.{operator}.b{bitwidth}"
+            constraints = f"l,{CONSTRAINT_MAP[bitwidth]},~{{memory}}"
+            stcs = ir.InlineAsm(stcs_type, f"{inst} [$0], $1;", constraints)
+            builder.call(stcs, [ptr, casted_value])
+
+        return signature, codegen
+
+    return impl
+
+
+stcg_intrinsic = st_cache_operator("cg")
+stcs_intrinsic = st_cache_operator("cs")
+stwb_intrinsic = st_cache_operator("wb")
+stwt_intrinsic = st_cache_operator("wt")
+
+
+@overload(ldca, target="cuda")
+def ol_ldca(array, i):
+    def impl(array, i):
+        return ldca_intrinsic(array, i)
+
+    return impl
+
+
+@overload(ldcg, target="cuda")
+def ol_ldcg(array, i):
+    def impl(array, i):
+        return ldcg_intrinsic(array, i)
+
+    return impl
+
+
+@overload(ldcs, target="cuda")
+def ol_ldcs(array, i):
+    def impl(array, i):
+        return ldcs_intrinsic(array, i)
+
+    return impl
+
+
+@overload(ldlu, target="cuda")
+def ol_ldlu(array, i):
+    def impl(array, i):
+        return ldlu_intrinsic(array, i)
+
+    return impl
+
+
+@overload(ldcv, target="cuda")
+def ol_ldcv(array, i):
+    def impl(array, i):
+        return ldcv_intrinsic(array, i)
+
+    return impl
+
+
+@overload(stcg, target="cuda")
+def ol_stcg(array, i, value):
+    def impl(array, i, value):
+        return stcg_intrinsic(array, i, value)
+
+    return impl
+
+
+@overload(stcs, target="cuda")
+def ol_stcs(array, i, value):
+    def impl(array, i, value):
+        return stcs_intrinsic(array, i, value)
+
+    return impl
+
+
+@overload(stwb, target="cuda")
+def ol_stwb(array, i, value):
+    def impl(array, i, value):
+        return stwb_intrinsic(array, i, value)
+
+    return impl
+
+
+@overload(stwt, target="cuda")
+def ol_stwt(array, i, value):
+    def impl(array, i, value):
+        return stwt_intrinsic(array, i, value)
+
+    return impl

numba_cuda/numba/cuda/cext/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause

numba_cuda/numba/cuda/cext/_devicearray.cpp

@@ -0,0 +1,159 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+/* This file contains the base class implementation for all device arrays. The
+ * base class is implemented in C so that computing typecodes for device arrays
+ * can be implemented efficiently. */
+
+#include "_pymodule.h"
+
+
+/* Include _devicearray., but make sure we don't get the definitions intended
+ * for consumers of the Device Array API.
+ */
+#define NUMBA_IN_DEVICEARRAY_CPP_
+#include "_devicearray.h"
+
+/* DeviceArray PyObject implementation. Note that adding more members here is
+ * presently prohibited because mapped and managed arrays derive from both
+ * DeviceArray and NumPy's ndarray, which is also a C extension class - the
+ * layout of the object cannot be resolved if this class also has members beyond
+ * PyObject_HEAD. */
+class DeviceArray {
+    PyObject_HEAD
+};
+
+/* Trivial traversal - DeviceArray instances own nothing. */
+static int
+DeviceArray_traverse(DeviceArray *self, visitproc visit, void *arg)
+{
+    return 0;
+}
+
+/* Trivial clear of all references - DeviceArray instances own nothing. */
+static int
+DeviceArray_clear(DeviceArray *self)
+{
+    return 0;
+}
+
+/* The _devicearray.DeviceArray type */
+PyTypeObject DeviceArrayType = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "_devicearray.DeviceArray",                  /* tp_name */
+    sizeof(DeviceArray),                         /* tp_basicsize */
+    0,                                           /* tp_itemsize */
+    0,                                           /* tp_dealloc */
+    0,                                           /* tp_vectorcall_offset */
+    0,                                           /* tp_getattr */
+    0,                                           /* tp_setattr */
+    0,                                           /* tp_as_async */
+    0,                                           /* tp_repr */
+    0,                                           /* tp_as_number */
+    0,                                           /* tp_as_sequence */
+    0,                                           /* tp_as_mapping */
+    0,                                           /* tp_hash */
+    0,                                           /* tp_call*/
+    0,                                           /* tp_str*/
+    0,                                           /* tp_getattro*/
+    0,                                           /* tp_setattro*/
+    0,                                           /* tp_as_buffer*/
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
+                                                 /* tp_flags*/
+    "DeviceArray object",                        /* tp_doc */
+    (traverseproc) DeviceArray_traverse,         /* tp_traverse */
+    (inquiry) DeviceArray_clear,                 /* tp_clear */
+    0,                                           /* tp_richcompare */
+    0,                                           /* tp_weaklistoffset */
+    0,                                           /* tp_iter */
+    0,                                           /* tp_iternext */
+    0,                                           /* tp_methods */
+    0,                                           /* tp_members */
+    0,                                           /* tp_getset */
+    0,                                           /* tp_base */
+    0,                                           /* tp_dict */
+    0,                                           /* tp_descr_get */
+    0,                                           /* tp_descr_set */
+    0,                                           /* tp_dictoffset */
+    0,                                           /* tp_init */
+    0,                                           /* tp_alloc */
+    0,                                           /* tp_new */
+    0,                                           /* tp_free */
+    0,                                           /* tp_is_gc */
+    0,                                           /* tp_bases */
+    0,                                           /* tp_mro */
+    0,                                           /* tp_cache */
+    0,                                           /* tp_subclasses */
+    0,                                           /* tp_weaklist */
+    0,                                           /* tp_del */
+    0,                                           /* tp_version_tag */
+    0,                                           /* tp_finalize */
+    0,                                           /* tp_vectorcall */
+#if (PY_MAJOR_VERSION == 3) && (PY_MINOR_VERSION == 12)
+/* This was introduced first in 3.12
+ * https://github.com/python/cpython/issues/91051
+ */
+    0,                                           /* tp_watched */
+#endif
+
+/* WARNING: Do not remove this, only modify it! It is a version guard to
+ * act as a reminder to update this struct on Python version update! */
+#if (PY_MAJOR_VERSION == 3)
+#if ! (NB_SUPPORTED_PYTHON_MINOR)
+#error "Python minor version is not supported."
+#endif
+#else
+#error "Python major version is not supported."
+#endif
+/* END WARNING*/
+};
+
+/* CUDA device array C API */
+static void *_DeviceArray_API[1] = {
+    (void*)&DeviceArrayType
+};
+
+MOD_INIT(_devicearray) {
+    PyObject *m = nullptr;
+    PyObject *d = nullptr;
+    PyObject *c_api = nullptr;
+    int error = 0;
+
+    MOD_DEF(m, "_devicearray", "No docs", NULL)
+    if (m == NULL)
+        goto error_occurred;
+
+    c_api = PyCapsule_New((void *)_DeviceArray_API, NUMBA_DEVICEARRAY_IMPORT_NAME "._DEVICEARRAY_API", NULL);
+    if (c_api == NULL)
+        goto error_occurred;
+
+    DeviceArrayType.tp_new = PyType_GenericNew;
+    if (PyType_Ready(&DeviceArrayType) < 0)
+        goto error_occurred;
+
+    Py_INCREF(&DeviceArrayType);
+    error = PyModule_AddObject(m, "DeviceArray", (PyObject*)(&DeviceArrayType));
+    if (error)
+        goto error_occurred;
+
+    d = PyModule_GetDict(m);
+    if (d == NULL)
+        goto error_occurred;
+
+    error = PyDict_SetItemString(d, "_DEVICEARRAY_API", c_api);
+    /* Decref and set c_api to NULL, Py_XDECREF in error_occurred will have no
+     * effect. */
+    Py_CLEAR(c_api);
+
+    if (error)
+        goto error_occurred;
+
+    return MOD_SUCCESS_VAL(m);
+
+error_occurred:
+    Py_XDECREF(m);
+    Py_XDECREF(c_api);
+    Py_XDECREF((PyObject*)&DeviceArrayType);
+
+    return MOD_ERROR_VAL;
+}

numba_cuda/numba/cuda/cext/_devicearray.h

@@ -0,0 +1,29 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+#ifndef NUMBA_DEVICEARRAY_H_
+#define NUMBA_DEVICEARRAY_H_
+
+#ifdef __cplusplus
+    extern "C" {
+#endif
+
+#define NUMBA_DEVICEARRAY_IMPORT_NAME "numba.cuda.cext._devicearray"
+/* These definitions should only be used by consumers of the Device Array API.
+ * Consumers access the API through the opaque pointer stored in
+ * _devicearray._DEVICEARRAY_API.  We don't want these definitions in
+ * _devicearray.cpp itself because they would conflict with the actual
+ * implementations there.
+ */
+#ifndef NUMBA_IN_DEVICEARRAY_CPP_
+
+    extern void **DeviceArray_API;
+    #define DeviceArrayType (*(PyTypeObject*)DeviceArray_API[0])
+
+#endif /* ndef NUMBA_IN_DEVICEARRAY_CPP */
+
+#ifdef __cplusplus
+    }
+#endif
+
+#endif  /* NUMBA_DEVICEARRAY_H_ */