...tnn

4 files changed, 7 insertions, 1 deletions

diff --git a/nerv/lib/matrix/cukernel.h b/nerv/lib/matrix/cukernel.h
index 4654d29..c84200e 100644
--- a/nerv/lib/matrix/cukernel.h
+++ b/nerv/lib/matrix/cukernel.h
@@ -34,4 +34,6 @@ void cudak_(cuda_copy_rows_by_colidx)(const Matrix *a, Matrix *b,
                                                     double alpha, double beta);
 #endif
 
+void cudak_(cuda_prefixsum_row)(const Matrix *a, Matrix *b);
+
 #endif
diff --git a/nerv/lib/matrix/cumatrix.c b/nerv/lib/matrix/cumatrix.c
index 58bdfe7..d998871 100644
--- a/nerv/lib/matrix/cumatrix.c
+++ b/nerv/lib/matrix/cumatrix.c
@@ -1,5 +1,5 @@
 #define NERV_GENERIC_CUMATRIX
-#include "../common.h"
+#include "cumatrix.h"
 #include "cuda_helper.h"
 #include <string.h>
 #include <time.h>
diff --git a/nerv/lib/matrix/cumatrix.h b/nerv/lib/matrix/cumatrix.h
index e6def66..b47e14b 100644
--- a/nerv/lib/matrix/cumatrix.h
+++ b/nerv/lib/matrix/cumatrix.h
@@ -1,7 +1,9 @@
 #ifndef NERV_CUMATRIX_H
 #define NERV_CUMATRIX_H
 #include "matrix.h"
+#include "../common.h"
 void nerv_cumatrix_print_profile();
 void nerv_cumatrix_clear_profile();
 void nerv_cumatrix_init();
+void nerv_cumatrix_select_gpu(int dev, Status *status);
 #endif
diff --git a/nerv/lib/matrix/generic/cumatrix.h b/nerv/lib/matrix/generic/cumatrix.h
index 5b8076f..f3c2df8 100644
--- a/nerv/lib/matrix/generic/cumatrix.h
+++ b/nerv/lib/matrix/generic/cumatrix.h
@@ -73,3 +73,5 @@ void nerv_matrix_(update_select_rows_by_rowidx)(Matrix *c, const Matrix *a, cons
 void nerv_matrix_(update_select_rows_by_colidx)(Matrix *c, const Matrix *a, const Matrix *idx,
                                                 double alpha, double beta, Status *status);
 #endif
+
+void nerv_matrix_(prefixsum_row)(Matrix *a, const Matrix *b, Status *status);