Add kernel export API via cuda.tile.compilation module

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

cext/cuda_loader.h

@@ -11,7 +11,7 @@
 
 #define FOREACH_CUDA_FUNCTION_TO_LOAD(X) \
     X(cuInit, 2000) \
-    X(cuLibraryLoadFromFile, 12000) \
+    X(cuLibraryLoadData, 12000) \
     X(cuLibraryUnload, 12000) \
     X(cuLibraryGetKernel, 12000) \
     X(cuGetErrorString, 6000) \

cext/hash_map.h

@@ -30,6 +30,17 @@ class HashMap {
     static constexpr size_t kInitialBuckets = 4;
 
 public:
+    struct Item {
+        const K key;
+        V value;
+
+        Item(K&& k, V&& v) : key(std::move(k)), value(std::move(v)) {}
+        Item(Item&&) = default;
+
+        Item(const Item&) = delete;
+        void operator=(const Item&) = delete;
+    };
+
     HashMap()
       : nbuckets_(kInitialBuckets),
         size_(0),
@@ -50,17 +61,17 @@ class HashMap {
     }
 
     template <typename Q>
-    V* find(Q&& key) {
+    Item* find(Q&& key) {
         uint64_t needle = compute_hash(key) | kOccupiedBit;
         auto [found, pos] = lookup(key, needle);
-        return found ? &items_[pos].value : nullptr;
+        return found ? &items_[pos] : nullptr;
     }
 
-    V* insert(K key, V value) {
+    Item* insert(K key, V value) {
         uint64_t needle = compute_hash(key) | kOccupiedBit;
         auto [found, pos] = lookup(key, needle);
         if (found) {
-            return &items_[pos].value;
+            return &items_[pos];
         } else {
             if (size_ >= nbuckets_ / 2) {
                 rehash();
@@ -70,18 +81,11 @@ class HashMap {
             }
             ++size_;
             hashes_[pos] = needle;
-            Item* new_item = new (&items_[pos]) Item(std::move(key), std::move(value));
-            return &new_item->value;
+            return new (&items_[pos]) Item(std::move(key), std::move(value));
         }
     }
 
 private:
-    struct Item {
-        K key;
-        V value;
-
-        Item(K&& k, V&& v) : key(std::move(k)), value(std::move(v)) {}
-    };
 
     size_t nbuckets_;  // power of two
     size_t size_;

cext/py.h

@@ -41,11 +41,28 @@ T pylong_as(PyObject* obj) {
         return PyLong_AsLong(obj);
     } else if constexpr (std::is_same_v<T, long long>) {
         return PyLong_AsLongLong(obj);
+    } else if constexpr (std::is_same_v<T, unsigned long>) {
+        return PyLong_AsUnsignedLong(obj);
+    } else if constexpr (std::is_same_v<T, unsigned long long>) {
+        return PyLong_AsUnsignedLongLong(obj);
     } else {
         static_assert(!sizeof(T*), "pylong_as<T> not implemented for given T");
     }
 }
 
+template <typename T>
+T pylong_as_overflow_and(PyObject* obj, int* overflow) {
+    if constexpr (std::is_same_v<T, int>) {
+        return pylong_as_int(obj);
+    } else if constexpr (std::is_same_v<T, long>) {
+        return PyLong_AsLongAndOverflow(obj, overflow);
+    } else if constexpr (std::is_same_v<T, long long>) {
+        return PyLong_AsLongLongAndOverflow(obj, overflow);
+    } else {
+        static_assert(!sizeof(T*), "pylong_as_overflow_and<T> not implemented for given T");
+    }
+}
+
 template <typename T>
 struct PythonWrapper {
     PyObject_HEAD
@@ -75,6 +92,21 @@ void pywrapper_dealloc(PyObject* self) {
     Py_TYPE(self)->tp_free(self);
 }
 
+template <typename T>
+PyObject* pywrapper_richcompare_via_operator_equals(PyObject* self, PyObject* other, int op) {
+    if (!PyObject_TypeCheck(self, &T::pytype) || !PyObject_TypeCheck(other, &T::pytype))
+        return Py_NewRef(Py_NotImplemented);
+
+    T& a = py_unwrap<T>(self);
+    T& b = py_unwrap<T>(other);
+
+    switch (op) {
+    case Py_EQ: return Py_NewRef(a == b ? Py_True : Py_False);
+    case Py_NE: return Py_NewRef(a == b ? Py_False : Py_True);
+    default: return Py_NewRef(Py_NotImplemented);
+    }
+}
+
 struct OK_t{};
 struct ErrorRaised_t{};
 

cext/test/test_hash_map.cpp

@@ -13,18 +13,18 @@ int main() {
 
     hm.insert(0, 0);
     CHECK(hm.find(0));
-    CHECK(*hm.find(0) == 0);
+    CHECK(hm.find(0)->value == 0);
 
     // Insert doesn't overwrite existing values
     hm.insert(0, 20);
-    CHECK(*hm.find(0) == 0);
+    CHECK(hm.find(0)->value == 0);
 
     for (int i = 1; i < 1000; ++i) {
         hm.insert(i * 16, i * 10);
         for (int j = 0; j <= i; ++j) {
-            int* v = hm.find(j * 16);
-            CHECK(v);
-            CHECK(*v == j * 10);
+            auto* item = hm.find(j * 16);
+            CHECK(item);
+            CHECK(item->value == j * 10);
         }
     }