From 15c0654c5c13963f1a804c6444b5a3c9a6ff9ea6 Mon Sep 17 00:00:00 2001
From: sunwen <sunwen@realcan.com>
Date: Mon, 9 Oct 2023 17:13:31 +0800
Subject: [PATCH] Add NewCuda.

---
 src/Function.cpp |  6 ++++++
 src/Function.h   |  1 +
 src/Matrix.cpp   | 22 ++++++++++++++++++++++
 src/Matrix.h     |  5 +++++
 4 files changed, 34 insertions(+)

diff --git a/src/Function.cpp b/src/Function.cpp
index 6000ba4..2f2957e 100644
--- a/src/Function.cpp
+++ b/src/Function.cpp
@@ -13,6 +13,7 @@
 #include <Eigen/Core>
 #include <Eigen/Eigen>
 #include <Eigen/Dense>
+#include <cuda_runtime.h>
 
 namespace Aurora {
 
@@ -25,4 +26,9 @@ namespace Aurora {
     void free(void* ptr){
         mkl_free(ptr);
     }
+
+    void gpuFree(void* ptr)
+    {
+        cudaFree(ptr);
+    }
 }
\ No newline at end of file
diff --git a/src/Function.h b/src/Function.h
index 4be1195..7882dd6 100644
--- a/src/Function.h
+++ b/src/Function.h
@@ -10,6 +10,7 @@
 namespace Aurora{
     float* malloc(size_t size,bool complex = false);
     void free(void* ptr);
+    void gpuFree(void* ptr);
 };
 
 
diff --git a/src/Matrix.cpp b/src/Matrix.cpp
index 7fa19dd..adccf22 100644
--- a/src/Matrix.cpp
+++ b/src/Matrix.cpp
@@ -16,6 +16,8 @@
 #include "Eigen/src/Core/Matrix.h"
 #include "Function.h"
 
+#include <cuda_runtime.h>
+
 namespace Aurora{
     typedef void(*CalcFuncD)(const MKL_INT n, const float a[], const MKL_INT inca, const float b[],
                              const MKL_INT incb, float r[], const MKL_INT incr);
@@ -391,6 +393,26 @@ namespace Aurora {
         return ret;
     }
 
+    Matrix Matrix::NewCuda(float *data, int rows, int cols, int slices, ValueType type)
+    {
+        if (!data) return Matrix();
+        std::vector<int> vector(3);
+        vector[0]=rows;
+        vector[1] = (cols > 0?cols:1);
+        vector[2] = (slices > 0?slices:1);
+        Matrix ret({data, gpuFree}, vector);
+        if (type != Normal)ret.setValueType(type);
+        ret.mCuda_Allocated = true;
+        return ret;
+    }
+
+    Matrix Matrix::toHostMatrix() const
+    {
+        float* data = new float[getDataSize()];
+        cudaMemcpy(data, mData.get(), sizeof(float) * getDataSize() * (mValueType == Normal ? 1 : 2), cudaMemcpyDeviceToHost);
+        return fromRawData(data, mInfo[0], mInfo[1], mInfo[2], getValueType());
+    }
+
     Matrix Matrix::New(float *data, const Matrix &shapeMatrix) {
         return New(data,
                    shapeMatrix.getDimSize(0),
diff --git a/src/Matrix.h b/src/Matrix.h
index d671af9..e2116e6 100644
--- a/src/Matrix.h
+++ b/src/Matrix.h
@@ -87,6 +87,7 @@ namespace Aurora {
          */
         static Matrix New(float *data, const Matrix &shapeMatrix);
 
+        static Matrix NewCuda(float *data, int rows, int cols = 1, int slices = 1, ValueType type = Normal);
         /**
          * New a mkl calculate based Matrix
          * @attention  Memory are allocate by Aurora:malloc function
@@ -287,11 +288,15 @@ namespace Aurora {
 
         void forceReshape(int rows, int columns, int slices);
 
+        Matrix toHostMatrix() const;
+
+
     private:
         ValueType mValueType = Normal;
         std::shared_ptr<float> mData;
         std::vector<int> mInfo;
         bool mMKL_Allocated = false;
+        bool mCuda_Allocated = false;
     };
 }