use single thread on cpu

Author: clearhanhui

gammagl/mpops/paddle_ext/setup.py

@@ -2,10 +2,10 @@
 from paddle.utils.cpp_extension import CppExtension, CUDAExtension, setup
 
 cuda_macro = ('COMPILE_WITH_CUDA', None) # Paddle offer `PADDLE_WITH_CUDA` macro
-omp_macro = ('COMPILE_WITH_OMP', None) # Note: OpenMP needs gcc>4.2.0
-compile_args = {
-    'cxx':['-fopenmp']
-}
+# omp_macro = ('COMPILE_WITH_OMP', None) # Note: OpenMP needs gcc>4.2.0
+# compile_args = {
+#     'cxx':['-fopenmp']
+# }
 
 def get_exts():
     if paddle.is_compiled_with_cuda():
@@ -16,17 +16,20 @@ def get_exts():
                 'cpu/segment_sum_cpu.cpp',
                 'cuda/segment_sum_cuda.cu',
             ],
-            define_macros=[omp_macro, cuda_macro],
-            extra_compile_args=compile_args
+            define_macros=[
+                cuda_macro, 
+                # omp_macro,
+                ],
+            # extra_compile_args=compile_args
         )
     else:
         return CppExtension(
             sources=[
                 'segment_sum.cpp',
                 'cpu/segment_sum_cpu.cpp',
             ],
-            define_macros=[omp_macro],
-            extra_compile_args=compile_args
+            # define_macros=[omp_macro],
+            # extra_compile_args=compile_args
         )
 
 setup(

gammagl/mpops/torch_ext/setup.py

@@ -4,10 +4,10 @@
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
 
 cuda_macro = ('COMPILE_WITH_CUDA', None)
-omp_macro = ('COMPILE_WITH_OMP', None) # Note: OpenMP needs gcc>4.2.0
-compile_args = {
-    'cxx':['-fopenmp']
-}
+# omp_macro = ('COMPILE_WITH_OMP', None) # Note: OpenMP needs gcc>4.2.0
+# compile_args = {
+#     'cxx':['-fopenmp']
+# }
 
 def get_exts():
     if torch.cuda.is_available():
@@ -19,8 +19,11 @@ def get_exts():
                     'cpu/segment_max_cpu.cpp',
                     'cuda/segment_max_cuda.cu'
                     ],
-                define_macros=[cuda_macro, omp_macro],
-                extra_compile_args=compile_args
+                define_macros=[
+                    cuda_macro,
+                    # omp_macro,
+                    ],
+                # extra_compile_args=compile_args
             ),
             CUDAExtension(
                 name='torch_gspmm', # Note: same with TORCH_LIBRARY (import)
@@ -29,8 +32,11 @@ def get_exts():
                     'cpu/spmm_sum_cpu.cpp',
                     'cuda/spmm_sum_cuda.cu'
                     ],
-                define_macros=[cuda_macro, omp_macro],
-                extra_compile_args=compile_args
+                define_macros=[
+                    cuda_macro, 
+                    # omp_macro,
+                    ],
+                # extra_compile_args=compile_args
             )
         ]
     else:
@@ -41,17 +47,17 @@ def get_exts():
                     'segment_max.cpp',
                     'cpu/segment_max_cpu.cpp'
                     ],
-                define_macros=[omp_macro],
-                extra_compile_args=compile_args
+                # define_macros=[omp_macro],
+                # extra_compile_args=compile_args
             ),
             CppExtension(
                 name='torch_gspmm', 
                 sources=[
                     'gspmm.cpp',
                     'cpu/spmm_sum_cpu.cpp'
                     ],
-                define_macros=[omp_macro],
-                extra_compile_args=compile_args
+                # define_macros=[omp_macro],
+                # extra_compile_args=compile_args
             )
         ]
 

profiler/mpops/paddle_ext_.py

@@ -2,7 +2,6 @@
 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 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)