select gpu code from mfy

1 files changed, 9 insertions, 0 deletions

diff --git a/nerv/matrix/cumatrix.c b/nerv/matrix/cumatrix.c
index bf92f92..7f22d68 100644
--- a/nerv/matrix/cumatrix.c
+++ b/nerv/matrix/cumatrix.c
@@ -8,6 +8,14 @@ static cublasHandle_t cublas_handle;
 static cudaEvent_t profile_start, profile_stop;
 static HashMap *profile;
 
+static int select_gpu(lua_State *L) {
+    Status status;
+    int dev = luaL_checkinteger(L, 1);
+    nerv_cumatrix_select_gpu(dev, &status);
+    NERV_LUA_CHECK_STATUS(L, status);
+    return 0;
+}
+
 static int print_profile(lua_State *L) {
     nerv_cumatrix_print_profile();
     return 0;
@@ -21,6 +29,7 @@ static int clear_profile(lua_State *L) {
 static const luaL_Reg cumatrix_methods[] = {
     {"print_profile", print_profile},
     {"clear_profile", clear_profile},
+    {"select_gpu", select_gpu},
     {NULL, NULL}
 };