numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

numba_cuda/numba/cuda/cext/typeconv.cpp

@@ -0,0 +1,212 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+#include <cstring>
+#include <cstdio>
+#include <algorithm>
+#include <limits.h>
+
+#include "typeconv.hpp"
+
+
+// ------ TypeManager ------
+
+TCCMap::TCCMap()
+    : nb_records(0)
+{
+}
+
+size_t TCCMap::hash(const TypePair &key) const {
+    return std::hash<size_t>()(std::hash<Type>()(key.first)) ^
+           std::hash<Type>()(key.second);
+}
+
+void TCCMap::insert(const TypePair &key, TypeCompatibleCode val) {
+    size_t i = hash(key) & (TCCMAP_SIZE - 1);
+    TCCMapBin &bin = records[i];
+    for (unsigned int j = 0; j < bin.size(); ++j) {
+        if (bin[j].key == key) {
+            bin[j].val = val;
+            return;
+        }
+    }
+    bin.push_back({key, val});
+    nb_records++;
+}
+
+TypeCompatibleCode TCCMap::find(const TypePair &key) const {
+    size_t i = hash(key) & (TCCMAP_SIZE - 1);
+    const TCCMapBin &bin = records[i];
+    for (unsigned int j = 0; j < bin.size(); ++j) {
+        if (bin[j].key == key) {
+            return bin[j].val;
+        }
+    }
+    return TCC_FALSE;
+}
+
+// ----- Ratings -----
+Rating::Rating() : promote(0), safe_convert(0), unsafe_convert(0) { }
+
+inline bool Rating::operator < (const Rating &other) const {
+    if (unsafe_convert < other.unsafe_convert)
+        return true;
+    else if (unsafe_convert > other.unsafe_convert)
+        return false;
+    if (safe_convert < other.safe_convert)
+        return true;
+    else if (safe_convert > other.safe_convert)
+        return false;
+    return (promote < other.promote);
+}
+
+inline bool Rating::operator == (const Rating &other) const {
+    return promote == other.promote && safe_convert == other.safe_convert &&
+           unsafe_convert == other.unsafe_convert;
+}
+
+// ------ TypeManager ------
+
+bool TypeManager::canPromote(Type from, Type to) const {
+    return isCompatible(from, to) == TCC_PROMOTE;
+}
+
+bool TypeManager::canSafeConvert(Type from, Type to) const {
+    return isCompatible(from, to) == TCC_CONVERT_SAFE;
+}
+
+bool TypeManager::canUnsafeConvert(Type from, Type to) const {
+    return isCompatible(from, to) == TCC_CONVERT_UNSAFE;
+}
+
+void TypeManager::addPromotion(Type from, Type to) {
+    return addCompatibility(from, to, TCC_PROMOTE);
+}
+
+void TypeManager::addUnsafeConversion(Type from, Type to) {
+    return addCompatibility(from, to, TCC_CONVERT_UNSAFE);
+}
+
+void TypeManager::addSafeConversion(Type from, Type to) {
+    return addCompatibility(from, to, TCC_CONVERT_SAFE);
+}
+
+void TypeManager::addCompatibility(Type from, Type to, TypeCompatibleCode tcc) {
+    TypePair pair(from, to);
+    tccmap.insert(pair, tcc);
+}
+
+TypeCompatibleCode TypeManager::isCompatible(Type from, Type to) const {
+    if (from == to)
+        return TCC_EXACT;
+    TypePair pair(from, to);
+    return tccmap.find(pair);
+}
+
+
+int TypeManager::selectOverload(const Type sig[], const Type ovsigs[],
+                                int &selected,
+                                int sigsz, int ovct, bool allow_unsafe,
+                                bool exact_match_required
+                               ) const {
+    int count;
+    if (ovct <= 16) {
+        Rating ratings[16];
+        int candidates[16];
+        count = _selectOverload(sig, ovsigs, selected, sigsz, ovct,
+                                allow_unsafe, exact_match_required, ratings,
+                                candidates);
+    }
+    else {
+        Rating *ratings = new Rating[ovct];
+        int *candidates = new int[ovct];
+        count = _selectOverload(sig, ovsigs, selected, sigsz, ovct,
+                                allow_unsafe, exact_match_required, ratings,
+                                candidates);
+        delete [] ratings;
+        delete [] candidates;
+    }
+    return count;
+}
+
+int TypeManager::_selectOverload(const Type sig[], const Type ovsigs[],
+                                 int &selected, int sigsz, int ovct,
+                                 bool allow_unsafe, bool exact_match_required,
+                                 Rating ratings[], int candidates[]) const {
+    // Generate rating table
+    // Use a penalize scheme.
+    int nb_candidates = 0;
+
+    for (int i = 0; i < ovct; ++i) {
+        const Type *entry = &ovsigs[i * sigsz];
+        Rating rate;
+
+        for (int j = 0; j < sigsz; ++j) {
+            TypeCompatibleCode tcc = isCompatible(sig[j], entry[j]);
+            if (tcc == TCC_FALSE ||
+                (tcc == TCC_CONVERT_UNSAFE && !allow_unsafe) ||
+                (tcc != TCC_EXACT && exact_match_required)) {
+                // stop the loop early
+                goto _incompatible;
+            }
+            switch(tcc) {
+            case TCC_PROMOTE:
+                rate.promote += 1;
+                break;
+            case TCC_CONVERT_SAFE:
+                rate.safe_convert += 1;
+                break;
+            case TCC_CONVERT_UNSAFE:
+                rate.unsafe_convert += 1;
+                break;
+            default:
+                break;
+            }
+        }
+        ratings[nb_candidates] = rate;
+        candidates[nb_candidates] = i;
+        nb_candidates++;
+    _incompatible:
+        ;
+    }
+
+    // Bail if no match
+    if (nb_candidates == 0)
+        return 0;
+
+    // Find lowest rating
+    Rating best = ratings[0];
+    selected = candidates[0];
+
+    int matchcount = 1;
+    for (int i = 1; i < nb_candidates; ++i) {
+        if (ratings[i] < best) {
+            best = ratings[i];
+            selected = candidates[i];
+            matchcount = 1;
+        }
+        else if (ratings[i] == best) {
+            matchcount += 1;
+        }
+    }
+    return matchcount;
+}
+
+// ----- utils -----
+
+const char* TCCString(TypeCompatibleCode tcc) {
+    switch(tcc) {
+    case TCC_EXACT:
+        return "exact";
+    case TCC_SUBTYPE:
+        return "subtype";
+    case TCC_PROMOTE:
+        return "promote";
+    case TCC_CONVERT_SAFE:
+        return "safe_convert";
+    case TCC_CONVERT_UNSAFE:
+        return "unsafe_convert";
+    default:
+        return "false";
+    }
+}

numba_cuda/numba/cuda/cext/typeconv.hpp

@@ -0,0 +1,101 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-2-Clause
+
+#ifndef NUMBA_TYPECONV_HPP_
+#define NUMBA_TYPECONV_HPP_
+#include <string>
+#include <vector>
+
+
+typedef int Type;
+
+enum TypeCompatibleCode{
+    // No match
+    TCC_FALSE = 0,
+    // Exact match
+    TCC_EXACT,
+    // Subtype is UNUSED
+    TCC_SUBTYPE,
+    // Promotion with no precision loss
+    TCC_PROMOTE,
+    // Conversion with no precision loss
+    // e.g. int32 to double
+    TCC_CONVERT_SAFE,
+    // Conversion with precision loss
+    // e.g. int64 to double (53 bits precision)
+    TCC_CONVERT_UNSAFE,
+};
+
+typedef std::pair<Type, Type> TypePair;
+
+struct TCCRecord {
+    TypePair key;
+    TypeCompatibleCode val;
+};
+
+typedef std::vector<TCCRecord> TCCMapBin;
+
+class TCCMap {
+public:
+    TCCMap();
+
+    void insert(const TypePair &key, TypeCompatibleCode val);
+    TypeCompatibleCode find(const TypePair &key) const;
+private:
+    size_t hash(const TypePair &key) const;
+
+    /* Must be a power of two */
+    static const size_t TCCMAP_SIZE = 512;
+    TCCMapBin records[TCCMAP_SIZE];
+    int nb_records;
+};
+
+struct Rating {
+    unsigned int promote;
+    unsigned int safe_convert;
+    unsigned int unsafe_convert;
+
+    Rating();
+
+    bool operator < (const Rating &other) const;
+    bool operator == (const Rating &other) const;
+};
+
+
+class TypeManager{
+public:
+    bool canPromote(Type from, Type to) const;
+    bool canUnsafeConvert(Type from, Type to) const;
+    bool canSafeConvert(Type from, Type to) const;
+
+    void addPromotion(Type from, Type to);
+    void addUnsafeConversion(Type from, Type to);
+    void addSafeConversion(Type from, Type to);
+    void addCompatibility(Type from, Type to, TypeCompatibleCode by);
+
+    TypeCompatibleCode isCompatible(Type from, Type to) const;
+
+    /**
+    Output stored in selected.
+    Returns
+        Number of matches
+    */
+    int selectOverload(const Type sig[], const Type ovsigs[], int &selected,
+                       int sigsz, int ovct, bool allow_unsafe,
+                       bool exact_match_required
+                      ) const;
+
+private:
+    int _selectOverload(const Type sig[], const Type ovsigs[], int &selected,
+                        int sigsz, int ovct, bool allow_unsafe,
+                        bool exact_match_required,
+                        Rating ratings[], int candidates[]) const;
+
+    TCCMap tccmap;
+};
+
+
+const char* TCCString(TypeCompatibleCode tcc);
+
+
+#endif // NUMBA_TYPECONV_HPP_

numba_cuda/numba/cuda/cg.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+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.ext_types import grid_group, GridGroup as GridGroupClass
+
+
+class GridGroup:
+    """A cooperative group representing the entire grid"""
+
+    def sync() -> None:
+        """Synchronize this grid group"""
+
+
+def this_grid() -> GridGroup:
+    """Get the current grid group."""
+    return GridGroup()
+
+
+@intrinsic
+def _this_grid(typingctx):
+    sig = signature(grid_group)
+
+    def codegen(context, builder, sig, args):
+        context.active_code_library.use_cooperative = True
+        one = context.get_constant(types.int32, 1)
+        mod = builder.module
+        return builder.call(
+            nvvmutils.declare_cudaCGGetIntrinsicHandle(mod), (one,)
+        )
+
+    return sig, codegen
+
+
+@overload(this_grid, target="cuda")
+def _ol_this_grid():
+    def impl():
+        return _this_grid()
+
+    return impl
+
+
+@intrinsic
+def _grid_group_sync(typingctx, group):
+    sig = signature(types.int32, group)
+
+    def codegen(context, builder, sig, args):
+        context.active_code_library.use_cooperative = True
+        flags = context.get_constant(types.int32, 0)
+        mod = builder.module
+        return builder.call(
+            nvvmutils.declare_cudaCGSynchronize(mod), (*args, flags)
+        )
+
+    return sig, codegen
+
+
+@overload_method(GridGroupClass, "sync", target="cuda")
+def _ol_grid_group_sync(group):
+    def impl(group):
+        return _grid_group_sync(group)
+
+    return impl