fix bug

Author: clearhanhui

gammagl/mpops/paddle.py

@@ -18,7 +18,7 @@ def unsorted_segment_sum(x, segment_ids, num_segments=None):
     else:
         num_segments = pd.max(segment_ids)+1
     if use_ext:
-        return paddle_segment.segment_sum(x, segment_ids, num_segments)
+        return paddle_ext.segment_sum(x, segment_ids, num_segments)
     idx_ = pd.argsort(segment_ids)
     x = pd.gather(x, idx_)
     segment_ids = pd.gather(segment_ids, idx_)

gammagl/mpops/paddle_ext/cuda/segment_sum_cuda.cu

@@ -12,10 +12,8 @@ __global__ void segment_sum_cuda_forward_kernel(const data_t *x_data, const int6
   int64_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
   int64_t e = (thread_idx / K) % E;
   int64_t k = thread_idx % K;
-  printf(" thread_idx = %d \n", thread_idx);
   if (thread_idx < numel)  {
     int64_t idx = index_data[e];
-    printf("%f \n", *out_data);
     atomicAdd(out_data + idx * K + k,
               x_data[thread_idx]);
   }

gammagl/mpops/paddle_ext/readme.md

@@ -5,3 +5,5 @@ Compile Steps:  (CMake, not work)
 TODO: support cmake
 
 > In `paddle/utils/cpp_extension/extension_utils.py:341L`, flags `"-ccbin"` & `"cc"` may cause error, since Paddle needs `nvcc` compiling with higher c++ standard while these flags may cause flag like `-std=c++14` does not work, when you are using a lower version gcc. Besides, we usually recommand to set gcc path to `CC` rather than `cc` in Linux, it will also cause error. Just annotate them.
+
+> Please keep the version of nvcc and the paddle-cuda consist, it may occur `the provided PTX was compiled with an unsupported toolchain.`

gammagl/mpops/torch_ext/cpu/segment_max_cpu.cpp

@@ -50,15 +50,13 @@ std::tuple<torch::Tensor, torch::Tensor> segment_max_cpu_forward(torch::Tensor&
   for (auto e = 0; e < E; ++e) {
     idx = index_data[e];
     for (auto k = 0; k < K; ++k) {
-      if (out_data[idx * K + k] < x_data[e * K + k]) {
 #ifdef COMPILE_WITH_OMP
-#pragma omp atomic
+#pragma omp critical
 #endif
-      {
+      if (out_data[idx * K + k] < x_data[e * K + k]) {
         out_data[idx * K + k] = x_data[e * K + k];
         arg_out_data[idx * K + k] = e;
       }
-      }
     }
   }
   out.masked_fill_(out == std::numeric_limits<int64_t>::lowest(), (scalar_t)0);

profiler/mpops/paddle_ext_.py

@@ -2,7 +2,7 @@
 from paddle_ext import unsorted_segment_sum
 
 src = paddle.to_tensor([[1, 1], [2, 2], [3, 3]], dtype='float32', stop_gradient=False)
-## TODO: it still successfully run, but it will get wrong answer. 
+## TODO: it still successfully run, but it will get wrong answer on GPU. 
 # src = paddle.to_tensor([1, 2, 3, 4, 5, 6], dtype='float32').reshape((2, 3)) 
 index = paddle.to_tensor([0, 1, 0], dtype=paddle.int64)
 out = unsorted_segment_sum(src, index, 3)