remodify

paddle/operators/gather_func.h

@@ -14,9 +14,9 @@ limitations under the License. */
 
 #pragma once
 #include <cstring>
+#include "paddle/framework/ddim.h"
 #include "paddle/framework/tensor.h"
 #include "paddle/platform/place.h"
-#include "paddle/framework/ddim.h"
 
 /**
  * Return a new tensor from source tensor, gathered according to index
@@ -27,7 +27,7 @@ limitations under the License. */
 template <typename Place, typename T>
 Tensor* Gather(Tensor* src, Tensor* index) {
   // check index of shape 1-D
-  PADDLE_ENFORCE(index->dims().size()==1);
+  PADDLE_ENFORCE(index->dims().size() == 1);
   int index_size = index->dims()[0];
 
   // Source shape
@@ -41,60 +41,66 @@ Tensor* Gather(Tensor* src, Tensor* index) {
 
   /* slice size */
   int slice_size = 1;
-  for(unsigned int i = 0; i < src_dims.size(); ++i)
-	slice_size *= src_dims[i];
+  for (size_t i = 0; i < src_dims.size(); ++i) slice_size *= src_dims[i];
 
   /* Gathering */
   if (place == CPUPlace()) {
     // init for CPU
     output = New_tensor.mutable_data<T>(output_dims, CPUPlace());
-	CPUGather(src->data(), index->data(), slice_size, new_tensor->mutable_data());
+    CPUGather(
+        src->data(), index->data(), slice_size, new_tensor->mutable_data());
   } else {  // GPU
     // init for GPU
     output = New_tensor.mutable_data<T>(output_dims, GPUPlace());
     /* how to specialize device??*/
-	GPUGather(d, src->data(), index->data(), slice_size, new_tensor->mutable_data());
+    GPUGather(
+        d, src->data(), index->data(), slice_size, new_tensor->mutable_data());
   }
   return New_tensor;
 }
 
 /* Implementation of CPU copy */
-template<typename T>
-void CPUGather(const T* params, const int* indices, 
-			   const int slice_size, const int index_size,
+template <typename T>
+void CPUGather(const T* params,
+               const int* indices,
+               const int slice_size,
+               const int index_size,
                T* output) {
   const size_t slice_bytes = slice_size * sizeof(T);
 
-  for(int i = 0; i < index_size; ++i)
+  for (size_t i = 0; i < index_size; ++i) {
     int index_ = indices[i];
     /* copy src[index_] to output[i] */
-  	memcpy(output + i * slice_bytes,
-  		params + index_ * slice_bytes,
-  		slice_bytes);
+    memcpy(
+        output + i * slice_bytes, params + index_ * slice_bytes, slice_bytes);
+  }
 }
 
 /* Implementation of GPU copy:
    I suppose the GPUDevice& d, contains gpu_id and thread_id
    d = cuda_stream(gpu_id_, stream_id_);
 */
-template<typename T>
+template <typename T>
 void GPUGather(const GPUDevice& d,
-			   const T* src, const int* index, 
-	           const int slice_size, const int index_size,
+               const T* src,
+               const int* index,
+               const int slice_size,
+               const int index_size,
                T* output) {
   int block_count = slice_size * index_size;
   int thread_per_block = 1024;
 
-  GatherOpKernel<T>
-          <<<block_count, thread_per_block, 0, d.stream()>>>(
-              src, index, output, slice_size,
-              indices_size, slice_size, out_size);
+  GatherOpKernel<T><<<block_count, thread_per_block, 0, d.stream()>>>(
+      src, index, output, slice_size, indices_size, slice_size, out_size);
 }
 
 template <typename T>
-__global__ void GatherOpKernel(const T* params, const int* indices, T* out,
+__global__ void GatherOpKernel(const T* params,
+                               const int* indices,
+                               T* out,
                                int64 indices_size,
-                               int64 slice_size, int64 out_size) {
+                               int64 slice_size,
+                               int64 out_size) {
   /* I suppose we have the following macro,
      which I strongly suggest that we should put in cuda:
   #define CUDA_1D_KERNEL_LOOP(i, n)                            \

paddle/operators/scatter_func.h

@@ -14,95 +14,92 @@ limitations under the License. */
 
 #pragma once
 #include <cstring>
+#include "paddle/framework/ddim.h"
 #include "paddle/framework/tensor.h"
 #include "paddle/platform/place.h"
-#include "paddle/framework/ddim.h"
 
 /**
  * Return a updated tensor from source tensor, scattered according to index:
  * dst[i] += src[index[i]]
  * input[src]: type-T source Tensor
- * input[Index]: type-int index Tensor (1-D)
+ * input[index]: type-int index Tensor (1-D)
  * return: output tensor
  */
-template <typename place, typename T>
-void ScatterUpdate_func(Tensor* Src, Tensor* Dst, Tensor* Index) {
-	// assert index is an int-type tensor
-	assert(Index->istype(int));
-
+template <typename Place, typename T>
+void ScatterUpdate(Tensor* src, Tensor* dst, Tensor* index) {
   // Source shape
-	auto src_dims = Src->dims();
-	auto dst_dims = Dst->dims();
+  auto src_dims = src->dims();
+  auto dst_dims = dst->dims();
   DDim output_dims(dims_src);
 
-	// check Src shape and Dst shape should match
-	for(int i = 1; i < src_dims.size(); i++)
-		assert(src_dims[i]==dst_dims[i]);
+  // check src shape and dst shape should match
+  for (size_t i = 1; i < src_dims.size(); i++)
+    PADDLE_ENFORCE(src_dims[i] == dst_dims[i]);
 
-	int index_size = Index->dims()[0];
+  int index_size = index->dims()[0];
 
   /* slice size */
   int slice_size = 1;
-	for(unsigned int i = 0; i < src_dims.size(); ++i)
-		slice_size *= src_dims[i];
+  for (size_t i = 0; i < src_dims.size(); ++i) slice_size *= src_dims[i];
 
   if (place == CPUPlace()) {
     // init
     output = new_tensor.mutable_data<T>(output_dims, CPUPlace());
-		CPUScatterUpdate(src->data(), index->data(), slice_size, new_tensor->mutable_data());
+    CPUScatterUpdate(
+        src->data(), index->data(), slice_size, new_tensor->mutable_data());
 
   } else {  // GPU
     // init
     output = new_tensor.mutable_data<T>(output_dims, GPUPlace());
     /* how to specialize device??*/
-		GPUScatterUpdate(d, src->data(), index->data(), slice_size, new_tensor->mutable_data());
+    GPUScatterUpdate(
+        d, src->data(), index->data(), slice_size, new_tensor->mutable_data());
   }
 }
 
 /* Implementation of CPU copy */
-template<typename T>
-void CPUScatterUpdate(const T* src, const int* Index, 
-			   const int slice_size, const int index_size,
+template <typename T>
+void CPUScatterUpdate(const T* src,
+                      const int* index,
+                      const int slice_size,
+                      const int index_size,
                       T* output) {
-  //const size_t slice_bytes = slice_size * sizeof(T);
+  // const size_t slice_bytes = slice_size * sizeof(T);
 
-  for(int i = 0; i < index_size; ++i)
+  for (size_t i = 0; i < index_size; ++i) {
     int index_ = index[i];
-  	/* dst[index_] += src[index_]
-  	   add operation size: slice_size
-  	 */
-    math::vAdd<T>(slice_size, src + index_ * slice_bytes,
+    math::vAdd<T>(slice_size,
+                  src + index_ * slice_bytes,
                   output + i * slice_bytes,
                   output + i * slice_bytes);
-  	/* Scatter update, not just assign
-  	memcpy(output + i * slice_bytes,
-  		src + index_ * slice_bytes,
-  		slice_bytes);
-  	*/
+  }
 }
 
 /* Implementation of GPU scatter:
    I suppose the GPUDevice& d, contains gpu_id and thread_id
    d = cuda_stream(gpu_id_, stream_id_);
 */
-template<typename T>
+template <typename T>
 void GPUScatterUpdate(const GPUDevice& d,
-			   const T* src, const int* Index, 
-	           const int slice_size, const int index_size,
+                      const T* src,
+                      const int* index,
+                      const int slice_size,
+                      const int index_size,
                       T* output) {
   int block_count = slice_size * index_size;
   int thread_per_block = 1024;
 
-	ScatterOpKernel<T>
-          <<<block_count, thread_per_block, 0, d.stream()>>>(
-              src, Index, output, slice_size,
-              indices_size, slice_size, out_size);
+  ScatterOpKernel<T><<<block_count, thread_per_block, 0, d.stream()>>>(
+      src, index, output, slice_size, indices_size, slice_size, out_size);
 }
 
 template <typename T>
-__global__ void ScatterOpKernel(const T* params, const int* indices, T* out,
+__global__ void ScatterOpKernel(const T* params,
+                                const int* indices,
+                                T* out,
                                 int64 indices_size,
-                               int64 slice_size, int64 out_size) {
+                                int64 slice_size,
+                                int64 out_size) {
   /* I suppose we have the following macro,
      which I strongly suggest that we should put in cuda:
   #define CUDA_1D_KERNEL_LOOP(i, n)                            \