From 041d317816d8497a406eca0299123bbf21d6515f Mon Sep 17 00:00:00 2001
From: Bram Veenboer <veenboer@astron.nl>
Date: Tue, 7 Nov 2023 12:02:47 +0100
Subject: [PATCH] Update cudawrappers to support COBALT (#240)

* Make Function::getAttribute const
* Add Function::name
* Add HostMemory::size
* Add DeviceMemory::size
* Add Module constructor with CUjit_option map
* Update CHANGELOG
---
 CHANGELOG.md                |  7 ++++++
 include/cudawrappers/cu.hpp | 46 +++++++++++++++++++++++++++++++++----
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2836fd9f..39447532 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,8 +7,15 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 - Added `cu::Context::getDevice()`
+- Added `cu::Module` constructor with `CUjit_option` map argument
+- Added `DeviceMemory::size`
+- Added `HostMemory::size`
+- Added `Function::name`
+
 ### Changed
 - Fixed the `cu::Module(CUmodule&)` constructor
+- Added `Function::getAttribute` is now const
+
 ### Removed
 
 ## [0.6.0] - 2023-10-06
diff --git a/include/cudawrappers/cu.hpp b/include/cudawrappers/cu.hpp
index 06405782..04974ddd 100644
--- a/include/cudawrappers/cu.hpp
+++ b/include/cudawrappers/cu.hpp
@@ -4,6 +4,7 @@
 #include <array>
 #include <cstddef>
 #include <exception>
+#include <map>
 #include <memory>
 #include <stdexcept>
 #include <string>
@@ -250,7 +251,7 @@ class Context : public Wrapper<CUcontext> {
 
 class HostMemory : public Wrapper<void *> {
  public:
-  explicit HostMemory(size_t size, unsigned int flags = 0) {
+  explicit HostMemory(size_t size, unsigned int flags = 0) : _size(size) {
     checkCudaCall(cuMemHostAlloc(&_obj, size, flags));
     manager = std::shared_ptr<void *>(new (void *)(_obj), [](void **ptr) {
       cuMemFreeHost(*ptr);
@@ -258,7 +259,8 @@ class HostMemory : public Wrapper<void *> {
     });
   }
 
-  explicit HostMemory(void *ptr, size_t size, unsigned int flags = 0) {
+  explicit HostMemory(void *ptr, size_t size, unsigned int flags = 0)
+      : _size(size) {
     _obj = ptr;
     checkCudaCall(cuMemHostRegister(&_obj, size, flags));
     manager = std::shared_ptr<void *>(
@@ -269,6 +271,11 @@ class HostMemory : public Wrapper<void *> {
   operator T *() {
     return static_cast<T *>(_obj);
   }
+
+  size_t size() const { return _size; }
+
+ private:
+  size_t _size;
 };
 
 class Array : public Wrapper<CUarray> {
@@ -342,6 +349,24 @@ class Module : public Wrapper<CUmodule> {
     });
   }
 
+  typedef std::map<CUjit_option, void *> optionmap_t;
+  explicit Module(const void *image, Module::optionmap_t &options) {
+    std::vector<CUjit_option> keys;
+    std::vector<void *> values;
+
+    for (const std::pair<CUjit_option, void *> &i : options) {
+      keys.push_back(i.first);
+      values.push_back(i.second);
+    }
+
+    checkCudaCall(cuModuleLoadDataEx(&_obj, image, options.size(), keys.data(),
+                                     values.data()));
+
+    for (size_t i = 0; i < keys.size(); ++i) {
+      options[keys[i]] = values[i];
+    }
+  }
+
   explicit Module(CUmodule &module) : Wrapper(module) {}
 
   CUdeviceptr getGlobal(const char *name) const {
@@ -353,13 +378,13 @@ class Module : public Wrapper<CUmodule> {
 
 class Function : public Wrapper<CUfunction> {
  public:
-  Function(const Module &module, const char *name) {
+  Function(const Module &module, const char *name) : _name(name) {
     checkCudaCall(cuModuleGetFunction(&_obj, module, name));
   }
 
   explicit Function(CUfunction &function) : Wrapper(function) {}
 
-  int getAttribute(CUfunction_attribute attribute) {
+  int getAttribute(CUfunction_attribute attribute) const {
     int value{};
     checkCudaCall(cuFuncGetAttribute(&value, attribute, _obj));
     return value;
@@ -368,6 +393,11 @@ class Function : public Wrapper<CUfunction> {
   void setCacheConfig(CUfunc_cache config) {
     checkCudaCall(cuFuncSetCacheConfig(_obj, config));
   }
+
+  const char *name() const { return _name; }
+
+ private:
+  const char *_name;
 };
 
 class Event : public Wrapper<CUevent> {
@@ -402,7 +432,8 @@ class Event : public Wrapper<CUevent> {
 class DeviceMemory : public Wrapper<CUdeviceptr> {
  public:
   explicit DeviceMemory(size_t size, CUmemorytype type = CU_MEMORYTYPE_DEVICE,
-                        unsigned int flags = 0) {
+                        unsigned int flags = 0)
+      : _size(size) {
     if (type == CU_MEMORYTYPE_DEVICE and !flags) {
       checkCudaCall(cuMemAlloc(&_obj, size));
     } else if (type == CU_MEMORYTYPE_UNIFIED) {
@@ -443,6 +474,11 @@ class DeviceMemory : public Wrapper<CUdeviceptr> {
           "Cannot return memory of type CU_MEMORYTYPE_DEVICE as pointer.");
     }
   }
+
+  size_t size() const { return _size; }
+
+ private:
+  size_t _size;
 };
 
 class Stream : public Wrapper<CUstream> {