cumo 0.1.0

data/ext/cumo/cuda/driver.c

@@ -0,0 +1,421 @@
+#include <ruby.h>
+#include <ruby/thread.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include "cumo/cuda/driver.h"
+
+VALUE cumo_cuda_eDriverError;
+VALUE cumo_cuda_mDriver;
+#define eDriverError cumo_cuda_eDriverError
+#define mDriver cumo_cuda_mDriver
+
+static void
+check_status(CUresult status)
+{
+    if (status != 0) {
+        const char *errname = NULL;
+        const char *errstring = NULL;
+        cuGetErrorName(status, &errname);
+        cuGetErrorString(status, &errstring);
+        rb_raise(cumo_cuda_eDriverError, "%s %s (error=%d)", errname, errstring, status);
+    }
+}
+
+///////////////////////////////////////////////
+// Context Management
+//////////////////////////////////////////////
+
+static VALUE
+rb_cuCtxCreate(VALUE self, VALUE flags, VALUE dev)
+{
+    unsigned int _flags = NUM2INT(flags);
+    CUdevice _dev = (CUdevice)NUM2INT(dev);
+    CUcontext _pctx;
+    CUresult status;
+
+    status = cuCtxCreate(&_pctx, _flags, _dev);
+
+    check_status(status);
+    return SIZET2NUM((size_t)_pctx);
+}
+
+static VALUE
+rb_cuCtxGetCurrent(VALUE self)
+{
+    CUcontext ctx;
+    CUresult status;
+
+    status = cuCtxGetCurrent(&ctx);
+    check_status(status);
+
+    return SIZET2NUM((size_t)ctx);
+}
+
+///////////////////////////////////////////////
+// Device Management
+//////////////////////////////////////////////
+
+static VALUE
+rb_cuDeviceGet(VALUE self, VALUE ordinal)
+{
+    int _ordinal = NUM2INT(ordinal);
+    CUdevice _device;
+    CUresult status;
+
+    status = cuDeviceGet(&_device, _ordinal);
+
+    check_status(status);
+    return INT2NUM(_device);
+}
+
+///////////////////////////////////////////////
+// Module Load and Kernel Execution
+//////////////////////////////////////////////
+
+struct cuLinkAddDataParam {
+    CUlinkState state;
+    CUjitInputType type;
+    void* data;
+    size_t size;
+    const char* name;
+    unsigned int numOptions;
+    CUjit_option* options;
+    void ** optionValues;
+};
+
+static void *
+cuLinkAddData_without_gvl_cb(void *param)
+{
+    struct cuLinkAddDataParam *p = param;
+    CUresult status;
+    status = cuLinkAddData(p->state, p->type, p->data, p->size, p->name, p->numOptions, p->options, p->optionValues);
+    return (void *)status;
+}
+
+// TODO(sonots): Support options.
+static VALUE
+rb_cuLinkAddData(VALUE self, VALUE state, VALUE type, VALUE data, VALUE name)
+{
+    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
+    CUjitInputType _type = (CUjitInputType)NUM2INT(type);
+    void* _data = (void *)RSTRING_PTR(data);
+    size_t _size = RSTRING_LEN(data);
+    const char* _name = RSTRING_PTR(data);
+    CUresult status;
+
+    struct cuLinkAddDataParam param = {_state, _type, _data, _size, _name, 0, (CUjit_option*)0, (void**)0};
+    status = (CUresult)rb_thread_call_without_gvl(cuLinkAddData_without_gvl_cb, &param, NULL, NULL);
+    //status = cuLinkAddData(_state, _type, _data, _size, _name, 0, (CUjit_option*)0, (void**)0);
+
+    check_status(status);
+    return Qnil;
+}
+
+struct cuLinkAddFileParam {
+    CUlinkState state;
+    CUjitInputType type;
+    const char* path;
+    unsigned int numOptions;
+    CUjit_option* options;
+    void ** optionValues;
+};
+
+static void *
+cuLinkAddFile_without_gvl_cb(void *param)
+{
+    struct cuLinkAddFileParam *p = param;
+    CUresult status;
+    status = cuLinkAddFile(p->state, p->type, p->path, p->numOptions, p->options, p->optionValues);
+    return (void *)status;
+}
+
+// TODO(sonots): Support options.
+static VALUE
+rb_cuLinkAddFile(VALUE self, VALUE state, VALUE type, VALUE path)
+{
+    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
+    CUjitInputType _type = (CUjitInputType)NUM2INT(type);
+    const char* _path = RSTRING_PTR(path);
+    CUresult status;
+
+    struct cuLinkAddFileParam param = {_state, _type, _path, 0, (CUjit_option*)0, (void **)0};
+    status = (CUresult)rb_thread_call_without_gvl(cuLinkAddFile_without_gvl_cb, &param, NULL, NULL);
+    //status = cuLinkAddFile(_state, _type, _path, 0, (CUjit_option*)0, (void **)0);
+
+    check_status(status);
+    return Qnil;
+}
+
+struct cuLinkCompleteParam {
+    CUlinkState state;
+    void** cubinOut;
+    size_t* sizeOut;
+};
+
+static void *
+cuLinkComplete_without_gvl_cb(void *param)
+{
+    struct cuLinkCompleteParam *p = param;
+    CUresult status;
+    status = cuLinkComplete(p->state, p->cubinOut, p->sizeOut);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuLinkComplete(VALUE self, VALUE state)
+{
+    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
+    void* _cubinOut;
+    size_t _sizeOut;
+    CUresult status;
+
+    struct cuLinkCompleteParam param = {_state, &_cubinOut, &_sizeOut};
+    status = (CUresult)rb_thread_call_without_gvl(cuLinkComplete_without_gvl_cb, &param, NULL, NULL);
+    //status = cuLinkComplete(_state, &_cubinOut, &_sizeOut);
+
+    check_status(status);
+    return rb_str_new((char *)_cubinOut, _sizeOut);
+}
+
+struct cuLinkCreateParam {
+    unsigned int numOptions;
+    CUjit_option* options;
+    void** optionValues;
+    CUlinkState* state;
+};
+
+static void *
+cuLinkCreate_without_gvl_cb(void *param)
+{
+    struct cuLinkCreateParam *p = param;
+    CUresult status;
+    status = cuLinkCreate(p->numOptions, p->options, p->optionValues, p->state);
+    return (void *)status;
+}
+
+// TODO(sonots): Support options.
+static VALUE
+rb_cuLinkCreate(VALUE self)
+{
+    CUlinkState state;
+    CUresult status;
+
+    struct cuLinkCreateParam param = {0, (CUjit_option*)0, (void**)0, &state};
+    status = (CUresult)rb_thread_call_without_gvl(cuLinkCreate_without_gvl_cb, &param, NULL, NULL);
+    //status = cuLinkCreate(0, (CUjit_option*)0, (void**)0, &state);
+
+    check_status(status);
+    return SIZET2NUM((size_t)state);
+}
+
+struct cuLinkDestroyParam {
+    CUlinkState state;
+};
+
+static void *
+cuLinkDestroy_without_gvl_cb(void *param)
+{
+    struct cuLinkDestroyParam *p = param;
+    CUresult status;
+    status = cuLinkDestroy(p->state);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuLinkDestroy(VALUE self, VALUE state)
+{
+    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
+    CUresult status;
+
+    struct cuLinkDestroyParam param = {_state};
+    status = (CUresult)rb_thread_call_without_gvl(cuLinkDestroy_without_gvl_cb, &param, NULL, NULL);
+    //status = cuLinkDestroy(_state);
+
+    check_status(status);
+    return Qnil;
+}
+
+struct cuModuleGetFunctionParam {
+    CUfunction* hfunc;
+    CUmodule hmod;
+    const char* name;
+};
+
+static void *
+cuModuleGetFunction_without_gvl_cb(void *param)
+{
+    struct cuModuleGetFunctionParam *p = param;
+    CUresult status;
+    status = cuModuleGetFunction(p->hfunc, p->hmod, p->name);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuModuleGetFunction(VALUE self, VALUE hmod, VALUE name)
+{
+    CUfunction _hfunc;
+    CUmodule _hmod = (CUmodule)NUM2SIZET(hmod);
+    const char* _name = RSTRING_PTR(name);
+    CUresult status;
+
+    struct cuModuleGetFunctionParam param = {&_hfunc, _hmod, _name};
+    status = (CUresult)rb_thread_call_without_gvl(cuModuleGetFunction_without_gvl_cb, &param, NULL, NULL);
+    //status = cuModuleGetFunction(&_hfunc, _hmod, _name);
+
+    check_status(status);
+    return SIZET2NUM((size_t)_hfunc);
+}
+
+struct cuModuleGetGlobalParam {
+    CUdeviceptr* dptr;
+    size_t* bytes;
+    CUmodule hmod;
+    const char* name;
+};
+
+static void *
+cuModuleGetGlobal_without_gvl_cb(void *param)
+{
+    struct cuModuleGetGlobalParam *p = param;
+    CUresult status;
+    status = cuModuleGetGlobal(p->dptr, p->bytes, p->hmod, p->name);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuModuleGetGlobal(VALUE self, VALUE hmod, VALUE name)
+{
+    CUdeviceptr _dptr;
+    size_t _bytes;
+    CUmodule _hmod = (CUmodule)NUM2SIZET(hmod);
+    const char* _name = RSTRING_PTR(name);
+    CUresult status;
+
+    struct cuModuleGetGlobalParam param = {&_dptr, &_bytes, _hmod, _name};
+    status = (CUresult)rb_thread_call_without_gvl(cuModuleGetGlobal_without_gvl_cb, &param, NULL, NULL);
+    //status = cuModuleGetGlobal(&_dptr, &_bytes, _hmod, _name);
+
+    check_status(status);
+    return rb_str_new((char *)_dptr, _bytes);
+}
+
+struct cuModuleLoadParam {
+    CUmodule* module;
+    const char* fname;
+};
+
+static void *
+cuModuleLoad_without_gvl_cb(void *param)
+{
+    struct cuModuleLoadParam *p = param;
+    CUresult status;
+    status = cuModuleLoad(p->module, p->fname);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuModuleLoad(VALUE self, VALUE fname)
+{
+    CUmodule _module;
+    const char* _fname = RSTRING_PTR(fname);
+    CUresult status;
+
+    struct cuModuleLoadParam param = {&_module, _fname};
+    status = (CUresult)rb_thread_call_without_gvl(cuModuleLoad_without_gvl_cb, &param, NULL, NULL);
+    //status = cuModuleLoad(&_module, _fname);
+
+    check_status(status);
+    return SIZET2NUM((size_t)_module);
+}
+
+struct cuModuleLoadDataParam {
+    CUmodule* module;
+    const void* image;
+};
+
+static void *
+cuModuleLoadData_without_gvl_cb(void *param)
+{
+    struct cuModuleLoadDataParam *p = param;
+    CUresult status;
+    status = cuModuleLoadData(p->module, p->image);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuModuleLoadData(VALUE self, VALUE image)
+{
+    CUmodule _module;
+    const void* _image = (void*)RSTRING_PTR(image);
+    CUresult status;
+
+    struct cuModuleLoadDataParam param = {&_module, _image};
+    status = (CUresult)rb_thread_call_without_gvl(cuModuleLoadData_without_gvl_cb, &param, NULL, NULL);
+    //status = cuModuleLoadData(&_module, _image);
+
+    check_status(status);
+    return SIZET2NUM((size_t)_module);
+}
+
+struct cuModuleUnloadParam {
+    CUmodule hmod;
+};
+
+static void *
+cuModuleUnload_without_gvl_cb(void *param)
+{
+    struct cuModuleUnloadParam *p = param;
+    CUresult status;
+    status = cuModuleUnload(p->hmod);
+    return (void *)status;
+}
+
+static VALUE
+rb_cuModuleUnload(VALUE self, VALUE hmod)
+{
+    CUmodule _hmod = (CUmodule)NUM2SIZET(hmod);
+    CUresult status;
+
+    struct cuModuleUnloadParam param = {_hmod};
+    status = (CUresult)rb_thread_call_without_gvl(cuModuleUnload_without_gvl_cb, &param, NULL, NULL);
+    //status = cuModuleUnload(_hmod);
+
+    check_status(status);
+    return Qnil;
+}
+
+void
+Init_cumo_cuda_driver()
+{
+    VALUE mCumo = rb_define_module("Cumo");
+    VALUE mCUDA = rb_define_module_under(mCumo, "CUDA");
+    mDriver = rb_define_module_under(mCUDA, "Driver");
+    eDriverError = rb_define_class_under(mCUDA, "DriverError", rb_eStandardError);
+
+    rb_define_singleton_method(mDriver, "cuCtxGetCurrent", rb_cuCtxGetCurrent, 0);
+    rb_define_singleton_method(mDriver, "cuLinkAddData",   rb_cuLinkAddData,   4);
+    rb_define_singleton_method(mDriver, "cuLinkAddFile", rb_cuLinkAddFile, 3);
+    rb_define_singleton_method(mDriver, "cuLinkComplete", rb_cuLinkComplete, 1);
+    rb_define_singleton_method(mDriver, "cuLinkCreate", rb_cuLinkCreate, 0);
+    rb_define_singleton_method(mDriver, "cuLinkDestroy", rb_cuLinkDestroy, 1);
+    rb_define_singleton_method(mDriver, "cuModuleGetFunction", rb_cuModuleGetFunction, 2);
+    rb_define_singleton_method(mDriver, "cuModuleGetGlobal", rb_cuModuleGetGlobal, 2);
+    rb_define_singleton_method(mDriver, "cuModuleLoad", rb_cuModuleLoad, 1);
+    rb_define_singleton_method(mDriver, "cuModuleLoadData", rb_cuModuleLoadData, 1);
+    rb_define_singleton_method(mDriver, "cuModuleUnload", rb_cuModuleUnload, 1);
+
+    rb_define_singleton_method(mDriver, "cuDeviceGet", rb_cuDeviceGet, 1);
+    rb_define_singleton_method(mDriver, "cuCtxCreate", rb_cuCtxCreate, 2);
+
+    rb_define_const(mDriver, "CU_JIT_INPUT_CUBIN", INT2NUM(CU_JIT_INPUT_CUBIN));
+    rb_define_const(mDriver, "CU_JIT_INPUT_FATBINARY", INT2NUM(CU_JIT_INPUT_FATBINARY));
+    rb_define_const(mDriver, "CU_JIT_INPUT_LIBRARY", INT2NUM(CU_JIT_INPUT_LIBRARY));
+    rb_define_const(mDriver, "CU_JIT_INPUT_OBJECT", INT2NUM(CU_JIT_INPUT_OBJECT));
+    rb_define_const(mDriver, "CU_JIT_INPUT_PTX", INT2NUM(CU_JIT_INPUT_PTX));
+
+    CUdevice cuDevice;
+    CUcontext context;
+    cuInit(0);
+    cuDeviceGet(&cuDevice, 0);
+    cuCtxCreate(&context, 0, cuDevice);
+}

data/ext/cumo/cuda/memory_pool.cpp

@@ -0,0 +1,185 @@
+#include <ruby.h>
+#include <cuda_runtime.h>
+#include "memory_pool_impl.hpp"
+#include "cumo/cuda/memory_pool.h"
+#include "cumo/cuda/runtime.h"
+
+#include <cstdlib>
+#include <string>
+
+#if defined(__cplusplus)
+extern "C" {
+#if 0
+} /* satisfy cc-mode */
+#endif
+#endif
+
+static cumo::internal::MemoryPool pool{};
+static bool memory_pool_enabled;
+
+VALUE cumo_cuda_eOutOfMemoryError;
+
+char*
+cumo_cuda_runtime_malloc(size_t size)
+{
+    if (memory_pool_enabled) {
+        try {
+            // TODO(sonots): Get current CUDA stream and pass it
+            return reinterpret_cast<char*>(pool.Malloc(size));
+        } catch (const cumo::internal::CUDARuntimeError& e) {
+            cumo_cuda_runtime_check_status(e.status());
+        } catch (const cumo::internal::OutOfMemoryError& e) {
+            rb_raise(cumo_cuda_eOutOfMemoryError, "%s", e.what());
+        }
+    } else {
+        void *ptr = 0;
+        cumo_cuda_runtime_check_status(cudaMallocManaged(&ptr, size, cudaMemAttachGlobal));
+        return reinterpret_cast<char*>(ptr);
+    }
+    return 0; // should not reach here
+}
+
+void
+cumo_cuda_runtime_free(char *ptr)
+{
+    if (memory_pool_enabled) {
+        try {
+            // TODO(sonots): Get current CUDA stream and pass it
+            pool.Free(reinterpret_cast<intptr_t>(ptr));
+        } catch (const cumo::internal::CUDARuntimeError& e) {
+            cumo_cuda_runtime_check_status(e.status());
+        }
+    } else {
+        cumo_cuda_runtime_check_status(cudaFree((void*)ptr));
+    }
+}
+
+/*
+  Enable memory pool.
+
+  @return [Boolean] Returns previous state (true if enabled)
+ */
+static VALUE
+rb_memory_pool_enable(VALUE self)
+{
+    VALUE ret = (memory_pool_enabled ? Qtrue : Qfalse);
+    memory_pool_enabled = true;
+    return ret;
+}
+
+/*
+  Disable memory pool.
+
+  @return [Boolean] Returns previous state (true if enabled)
+ */
+static VALUE
+rb_memory_pool_disable(VALUE self)
+{
+    VALUE ret = (memory_pool_enabled ? Qtrue : Qfalse);
+    memory_pool_enabled = false;
+    return ret;
+}
+
+/*
+  Returns whether memory pool is enabled or not.
+
+  @return [Boolean] Returns the state (true if enabled)
+ */
+static VALUE
+rb_memory_pool_enabled_p(VALUE self)
+{
+    return (memory_pool_enabled ? Qtrue : Qfalse);
+}
+
+/*
+  Free all **non-split** chunks in all arenas.
+ */
+static VALUE
+rb_memory_pool_free_all_blocks(int argc, VALUE* argv, VALUE self)
+{
+    try {
+        if (argc < 1) {
+            pool.FreeAllBlocks();
+        } else {
+            // TODO(sonots): FIX if we create a Stream object
+            cudaStream_t stream_ptr = (cudaStream_t)NUM2SIZET(argv[0]);
+            pool.FreeAllBlocks(stream_ptr);
+        }
+    } catch (const cumo::internal::CUDARuntimeError& e) {
+        cumo_cuda_runtime_check_status(e.status());
+    }
+    return Qnil;
+}
+
+/*
+  Count the total number of free blocks.
+
+  @return [Integer] The total number of free blocks.
+ */
+static VALUE
+rb_memory_pool_n_free_blocks(VALUE self)
+{
+    return SIZET2NUM(pool.GetNumFreeBlocks());
+}
+
+/*
+  Get the total number of bytes used.
+
+  @return [Integer] The total number of bytes used.
+ */
+static VALUE
+rb_memory_pool_used_bytes(VALUE self)
+{
+    return SIZET2NUM(pool.GetUsedBytes());
+}
+
+/*
+  Get the total number of bytes acquired but not used in the pool.
+
+  @return [Integer] The total number of bytes acquired but not used in the pool.
+ */
+static VALUE
+rb_memory_pool_free_bytes(VALUE self)
+{
+    return SIZET2NUM(pool.GetFreeBytes());
+}
+
+/*
+  Get the total number of bytes acquired in the pool.
+
+  @return [Integer] The total number of bytes acquired in the pool.
+ */
+static VALUE
+rb_memory_pool_total_bytes(VALUE self)
+{
+    return SIZET2NUM(pool.GetTotalBytes());
+}
+
+void
+Init_cumo_cuda_memory_pool()
+{
+    VALUE mCumo = rb_define_module("Cumo");
+    VALUE mCUDA = rb_define_module_under(mCumo, "CUDA");
+    VALUE mMemoryPool = rb_define_module_under(mCUDA, "MemoryPool");
+    cumo_cuda_eOutOfMemoryError = rb_define_class_under(mCUDA, "OutOfMemoryError", rb_eStandardError);
+    
+    rb_define_singleton_method(mMemoryPool, "enable", RUBY_METHOD_FUNC(rb_memory_pool_enable), 0);
+    rb_define_singleton_method(mMemoryPool, "disable", RUBY_METHOD_FUNC(rb_memory_pool_disable), 0);
+    rb_define_singleton_method(mMemoryPool, "enabled?", RUBY_METHOD_FUNC(rb_memory_pool_enabled_p), 0);
+    rb_define_singleton_method(mMemoryPool, "free_all_blocks", RUBY_METHOD_FUNC(rb_memory_pool_free_all_blocks), -1);
+    rb_define_singleton_method(mMemoryPool, "n_free_blocks", RUBY_METHOD_FUNC(rb_memory_pool_n_free_blocks), 0);
+    rb_define_singleton_method(mMemoryPool, "used_bytes", RUBY_METHOD_FUNC(rb_memory_pool_used_bytes), 0);
+    rb_define_singleton_method(mMemoryPool, "free_bytes", RUBY_METHOD_FUNC(rb_memory_pool_free_bytes), 0);
+    rb_define_singleton_method(mMemoryPool, "total_bytes", RUBY_METHOD_FUNC(rb_memory_pool_total_bytes), 0);
+
+    // default is true
+    const char* env = std::getenv("CUMO_MEMORY_POOL");
+    memory_pool_enabled = env == nullptr || (std::string(env) != "OFF" && std::string(env) != "0" && std::string(env) != "NO");
+}
+
+#if defined(__cplusplus)
+#if 0
+{ /* satisfy cc-mode */
+#endif
+}  /* extern "C" { */
+#endif