From 5e77b4dafef58aec4f186359788d70d35ac54dda Mon Sep 17 00:00:00 2001
From: sunwen <sunwen@realcan.com>
Date: Mon, 4 Dec 2023 11:27:17 +0800
Subject: [PATCH] Add cuda auroraUnion and unittest.

---
 src/Function1D.cu             | 21 +++++++++++++++++++++
 src/Function1D.cuh            |  2 ++
 test/Function1D_Cuda_Test.cpp | 15 +++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/src/Function1D.cu b/src/Function1D.cu
index df54c33..3d5e121 100644
--- a/src/Function1D.cu
+++ b/src/Function1D.cu
@@ -5,6 +5,7 @@
 #include "Matrix.h"
 
 #include <cmath>
+#include <cstddef>
 #include <cstdlib>
 #include <thrust/device_vector.h>
 #include <thrust/transform.h>
@@ -1232,3 +1233,23 @@ CudaMatrix Aurora::linspaceCuda(float aStart, float aEnd, int aNum)
     cudaDeviceSynchronize();
     return Aurora::CudaMatrix::fromRawData(data,aNum);
 }
+
+CudaMatrix Aurora::auroraUnion(const CudaMatrix& aMatrix1, const CudaMatrix& aMatrix2)
+{
+    if(aMatrix1.isNull() || aMatrix2.isNull() || aMatrix1.isComplex() || aMatrix2.isComplex())
+    {
+        std::cerr<<"auroraUnion not support complex cudamatrix"<<std::endl;
+        return CudaMatrix();
+    }
+
+    size_t size1= aMatrix1.getDataSize();
+    size_t size2= aMatrix2.getDataSize();
+    float* data = nullptr;
+    cudaMalloc((void**)&data, sizeof(float) * (size1 + size2));
+    cudaMemcpy(data, aMatrix1.getData(), sizeof(float) * size1, cudaMemcpyDeviceToDevice);
+    cudaMemcpy(data + size1, aMatrix2.getData(), sizeof(float) * size2, cudaMemcpyDeviceToDevice);
+    thrust::sort(thrust::device, data, data+size1+size2);
+    float* endPointer = thrust::unique(thrust::device, data, data+size1+size2);
+
+    return CudaMatrix::fromRawData(data, endPointer - data);    
+}
diff --git a/src/Function1D.cuh b/src/Function1D.cuh
index 491471c..adaefcc 100644
--- a/src/Function1D.cuh
+++ b/src/Function1D.cuh
@@ -73,6 +73,8 @@ namespace Aurora
 
     CudaMatrix linspaceCuda(float aStart, float aEnd, int aNum);
 
+    CudaMatrix auroraUnion(const CudaMatrix& aMatrix1, const CudaMatrix& aMatrix2);
+
     // ------compareSet----------------------------------------------------
 
 
diff --git a/test/Function1D_Cuda_Test.cpp b/test/Function1D_Cuda_Test.cpp
index 07824b8..f0f45a9 100644
--- a/test/Function1D_Cuda_Test.cpp
+++ b/test/Function1D_Cuda_Test.cpp
@@ -950,3 +950,18 @@ TEST_F(Function1D_Cuda_Test, linspace) {
         EXPECT_FLOAT_AE(result1[i], result2[i]);
     }
 }
+
+TEST_F(Function1D_Cuda_Test, auroraUnion) {
+    float* data1 = new float[9]{3,3,2,2,2,1,4,4,7};
+    auto matrix1 = Aurora::Matrix::fromRawData(data1, 9,1,1);
+    float* data2 = new float[8]{6,6,7,7,8,1,2};
+    auto matrix2 = Aurora::Matrix::fromRawData(data2, 7,1,1);
+
+    auto result1 = Aurora::auroraUnion(matrix1, matrix2);
+    auto result2 = Aurora::auroraUnion(matrix1.toDeviceMatrix(), matrix2.toDeviceMatrix()).toHostMatrix();
+    EXPECT_FLOAT_AE(result1.getDataSize(), result2.getDataSize());
+    for(int i=0;i<result1.getDataSize();++i)
+    {
+        EXPECT_FLOAT_AE(result1[i], result2[i]);
+    }
+}