gather test completed

source/tensor/test/TGather.cpp

@@ -20,6 +20,8 @@
  */
 
 #include "TGather.h"
+#include "../core/movement/Spread.h"
+#include "../core/getandset/ConvertDataType.h"
 
 namespace nts { // namespace nts(NiuTrans.Tensor)
 
@@ -362,6 +364,203 @@ bool TestGather3()
 #endif // USE_CUDA
 }
 
+
+/*
+case 4: gather indexed sub-tensors backward
+In this case, (3, 3) -> (3, 3), dim = 0,
+srcIndex = [0, 1, 2]
+*/
+bool TestGather4()
+{
+    /* a input tensor of size (3, 3) */
+    int sOrder = 2;
+    int * sDimSize = new int[sOrder];
+    sDimSize[0] = 3;
+    sDimSize[1] = 3;
+
+    int sUnitNum = 1;
+    for (int i = 0; i < sOrder; i++)
+        sUnitNum *= sDimSize[i];
+
+    /* a output tensor of size (3, 3) */
+    int tOrder = 2;
+    int * tDimSize = new int[tOrder];
+    tDimSize[0] = 3;
+    tDimSize[1] = 3;
+
+    int tUnitNum = 1;
+    for (int i = 0; i < tOrder; i++)
+        tUnitNum *= tDimSize[i];
+
+    /* a index tensor of size (2) */
+    int indexOrder = 1;
+    int * indexDimSize = new int[indexOrder];
+    indexDimSize[0] = 3;
+
+    int indexUnitNum = 1;
+    for (int i = 0; i < indexOrder; i++)
+        indexUnitNum *= indexDimSize[i];
+
+    DTYPE sData[3][3] = { {0.0F, -1.0F, 2.0F},
+                          {2.0F, 1.0F, 3.0F},
+                          {1.0F, 2.0F, 4.0F} };
+    DTYPE tData[3][3] = { {1.0F, 1.0F, 1.0F},
+                          {1.0F, 1.0F, 1.0F},
+                          {1.0F, 1.0F, 1.0F} };
+    DTYPE answer[3][3] = { {1.0F, 0.0F, 3.0F},
+                           {3.0F, 2.0F, 4.0F},
+                           {2.0F, 3.0F, 5.0F} };
+
+    int dim = 0;
+    int indexSize = 3;
+    int srcIndex[3] = {0, 1, 2};
+
+    /* CPU test */
+    bool cpuTest = true;
+
+#ifdef USE_CUDA
+    /* GPU test */
+    bool gpuTest = true;
+
+    /* create tensors */
+    XTensor * sGPU = NewTensor(sOrder, sDimSize, X_FLOAT, 1.0F, 0);
+    XTensor * tGPU = NewTensor(sOrder, tDimSize, X_FLOAT, 1.0F, 0);
+    XTensor * indexGPU = NewTensor(indexOrder, indexDimSize, X_INT, 1.0F, 0);
+
+    /* initialize variables */
+    sGPU->SetData(sData, sUnitNum);
+    tGPU->SetData(tData, tUnitNum);
+    indexGPU->SetData(srcIndex, indexSize);
+
+    /* call gather function */
+    _SpreadForGather(sGPU, tGPU, indexGPU);
+
+    /* check results */
+    gpuTest = sGPU->CheckData(answer, tUnitNum);
+
+    /* destroy variables */
+    delete sGPU;
+    delete tGPU;
+    delete indexGPU;
+    delete[] sDimSize;
+    delete[] tDimSize;
+    delete[] indexDimSize;
+
+    return cpuTest && gpuTest;
+#else
+    /* destroy variables */;
+    delete[] sDimSize;
+    delete[] tDimSize;
+    delete[] indexDimSize;
+
+    return cpuTest;
+#endif // USE_CUDA
+}
+
+/*
+case 5: float16 gather indexed sub-tensors backward
+In this case, (3, 3) -> (3, 3), dim = 0,
+srcIndex = [0, 1, 2]
+*/
+bool TestGather5()
+{
+    /* a input tensor of size (3, 3) */
+    int sOrder = 2;
+    int * sDimSize = new int[sOrder];
+    sDimSize[0] = 3;
+    sDimSize[1] = 3;
+
+    int sUnitNum = 1;
+    for (int i = 0; i < sOrder; i++)
+        sUnitNum *= sDimSize[i];
+
+    /* a output tensor of size (3, 3) */
+    int tOrder = 2;
+    int * tDimSize = new int[tOrder];
+    tDimSize[0] = 3;
+    tDimSize[1] = 3;
+
+    int tUnitNum = 1;
+    for (int i = 0; i < tOrder; i++)
+        tUnitNum *= tDimSize[i];
+
+    /* a index tensor of size (2) */
+    int indexOrder = 1;
+    int * indexDimSize = new int[indexOrder];
+    indexDimSize[0] = 3;
+
+    int indexUnitNum = 1;
+    for (int i = 0; i < indexOrder; i++)
+        indexUnitNum *= indexDimSize[i];
+
+    DTYPE sData[3][3] = { {0.0F, -1.0F, 2.0F},
+                          {2.0F, 1.0F, 3.0F},
+                          {1.0F, 2.0F, 4.0F} };
+    DTYPE tData[3][3] = { {1.0F, 1.0F, 1.0F},
+                          {1.0F, 1.0F, 1.0F},
+                          {1.0F, 1.0F, 1.0F} };
+    DTYPE answer[3][3] = { {1.0F, 0.0F, 3.0F},
+                           {3.0F, 2.0F, 4.0F},
+                           {2.0F, 3.0F, 5.0F} };
+
+    int dim = 0;
+    int indexSize = 3;
+    int srcIndex[3] = {0, 1, 2};
+
+    /* CPU test */
+    bool cpuTest = true;
+
+#ifdef USE_CUDA
+    /* GPU test */
+    bool gpuTest = true;
+
+    /* create tensors */
+    XTensor * sGPU = NewTensor(sOrder, sDimSize, X_FLOAT, 1.0F, 0);
+    XTensor * tGPU = NewTensor(sOrder, tDimSize, X_FLOAT, 1.0F, 0);
+    XTensor * indexGPU = NewTensor(indexOrder, indexDimSize, X_INT, 1.0F, 0);
+
+    /* create float16 tensors */
+    XTensor sHalfGPU;
+    XTensor tHalfGPU;
+
+    /* initialize variables */
+    sGPU->SetData(sData, sUnitNum);
+    tGPU->SetData(tData, tUnitNum);
+    indexGPU->SetData(srcIndex, indexSize);
+
+    /* convert data type from float to float16 */
+    sHalfGPU = ConvertDataType(*sGPU, X_FLOAT16);
+    tHalfGPU = ConvertDataType(*tGPU, X_FLOAT16);
+
+    /* call gather function */
+    _SpreadForGather(&sHalfGPU, &tHalfGPU, indexGPU);
+
+    /* convert data type from float16 to float */
+    _ConvertDataType(&sHalfGPU, sGPU);
+
+    /* check results */
+    gpuTest = sGPU->CheckData(answer, tUnitNum);
+
+    /* destroy variables */
+    delete sGPU;
+    delete tGPU;
+    delete indexGPU;
+    delete[] sDimSize;
+    delete[] tDimSize;
+    delete[] indexDimSize;
+
+    return cpuTest && gpuTest;
+#else
+    /* destroy variables */;
+    delete[] sDimSize;
+    delete[] tDimSize;
+    delete[] indexDimSize;
+
+    return cpuTest;
+#endif // USE_CUDA
+}
+
+
 /* other cases */
 /*
 TODO!!
@@ -398,8 +597,26 @@ bool TestGather()
         XPRINT(0, stdout, ">> case 3 failed!\n");
     }
     else
+        XPRINT(0, stdout, ">> case 3 passed!\n");
+
+    /* case 4 test */
+    caseFlag = TestGather4();
+    if (!caseFlag) {
+        returnFlag = false;
+        XPRINT(0, stdout, ">> case 4 failed!\n");
+    }
+    else
         XPRINT(0, stdout, ">> case 4 passed!\n");
 
+    /* case 5 test */
+    caseFlag = TestGather5();
+    if (!caseFlag) {
+        returnFlag = false;
+        XPRINT(0, stdout, ">> case 5 failed!\n");
+    }
+    else
+        XPRINT(0, stdout, ">> case 5 passed!\n");
+
     /* other cases test */
     /*
     TODO!!

source/tensor/test/Test.cpp

@@ -41,7 +41,7 @@ bool Test()
     //wrong = !TestDiv() || wrong;
     //wrong = !TestDivDim() || wrong;
     //wrong = !TestExp() || wrong;
-    //wrong = !TestGather() || wrong;
+    wrong = !TestGather() || wrong;
     //wrong = !TestLog() || wrong;
     //wrong = !TestMatrixMul() || wrong;
     //wrong = !TestMatrixMul2D() || wrong;