[setup.py] update extension module (#219)

* update extension module

* update

---------

Co-authored-by: gyzhou2000 <[email protected]>

gammagl/mpops/torch_ext/cpu/bspmm_sum_cpu.cpp

@@ -13,7 +13,7 @@ torch::Tensor bspmm_sum_cpu_forward(torch::Tensor &index, torch::Tensor &weight,
         index = index.contiguous();
     }
 
-    int num_nodes = x.size(0);
+    // int num_nodes = x.size(0);
     int heads = x.size(1);
     int out_channels = x.size(2);
 
@@ -62,7 +62,7 @@ std::tuple<torch::Tensor, torch::Tensor> bspmm_sum_cpu_backward(torch::Tensor &i
         index = index.contiguous();
     }
 
-    int num_nodes = grad.size(0);
+    // int num_nodes = grad.size(0);
     int heads = grad.size(1);
     int out_channels = grad.size(2);
 

gammagl/ops/segment/cpu/segment_csr.cpp

@@ -105,7 +105,7 @@ py::list segment_csr_cpu(
 
   //    auto indptr_info = getTensorInfo<int64_t>(indptr);
   //    auto stride = indptr_info.strides[indptr_info.dims - 1];
-  auto stride = 1;
+  // auto stride = 1;
   std::vector<int64_t> args(K);
 
   auto src_data = src.unchecked();

profiler/mpops/complete_test/mp_cpu/bspmm_sum_cpu.py

@@ -0,0 +1,67 @@
+import os
+
+os.environ['TL_BACKEND'] = 'torch'
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+import numpy as np
+import tensorlayerx as tlx
+from gammagl.mpops import *
+import time
+
+relative_path = '/home/zgy/review/zgy/GammaGL/profiler/mpops/edge_index/'
+file_name = ['cora.npy', 'pubmed.npy', 'ogbn-arxiv.npy']
+embedding = [16, 64, 256]
+heads = [8, 16, 32, 64]
+iter = 100
+
+
+with open('test_results.txt', 'w') as result_file:
+    for name in file_name:
+        path = relative_path + name
+        info = f"Loading data from {path}\n"
+        result_file.write(info)
+        print(info)
+
+        edge_index = np.load(path)
+
+        num_nodes = np.max(edge_index) + 1
+        src = tlx.convert_to_tensor(edge_index[0, :], tlx.int64)
+        dst = tlx.convert_to_tensor(edge_index[1, :], tlx.int64)
+        edge_index = tlx.convert_to_tensor(edge_index)
+
+        for head in heads:
+
+            weight = torch.ones((edge_index.shape[1], head), dtype=tlx.float32)
+
+            for embedding_dim in embedding:
+                info = f"**********embedding_dim={embedding_dim}  head={head}**********\n"
+                result_file.write(info)
+                print(info)
+                x = tlx.convert_to_tensor(np.random.randn(num_nodes, head, embedding_dim), dtype=tlx.float32)
+
+                start = time.time()
+                for j in range(iter):
+                    bspmm(edge_index, weight=weight, x=x, reduce='sum')
+                end = time.time()
+                info = "bspmm_sum:{:.3f}\n".format(end-start)
+                result_file.write(info)
+                print(info)
+
+                start = time.time()
+                for j in range(iter):
+                    msg = tlx.gather(x, src)
+                    edge_weight = tlx.expand_dims(weight, -1)
+                    msg = msg * edge_weight
+                    unsorted_segment_sum(msg, dst, num_nodes)
+                end = time.time()
+                info = "segment_sum:{:.3f}\n".format(end-start)
+                result_file.write(info)
+                print(info)
+
+
+                info = f"**********embedding_dim={embedding_dim}  head={head}**********\n"
+                result_file.write(info)
+                print(info)
+
+        info = f"Data tensors are on device: {x.device}\n"
+        result_file.write(info)
+        print(info)

setup.py

@@ -3,8 +3,9 @@
 import os
 import os.path as osp
 from setuptools import setup, find_packages
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
-from ggl_build_extension import PyCudaExtension, PyCPUExtension
+# from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
+# from ggl_build_extension import PyCudaExtension, PyCPUExtension
+from tensorlayerx.utils import PyCppExtension, PyCUDAExtension, PyBuildExtension
 
 # TODO will depend on different host
 WITH_CUDA = False
@@ -45,20 +46,22 @@ def load_mpops_extensions():
             file_list.extend([osp.join(src_dir, f) for f in src_files])
 
         if not WITH_CUDA:
-            extensions.append(CppExtension(
+            extensions.append(PyCppExtension(
                 name=osp.join(mpops_dir, f'_{mpops_prefix}').replace(osp.sep, "."),
                 sources=[f for f in file_list],
-                extra_compile_args=compile_args
+                extra_compile_args=compile_args,
+                use_torch=True
             ))
         else:
-            extensions.append(CUDAExtension(
+            extensions.append(PyCUDAExtension(
                 name=osp.join(mpops_dir, f'_{mpops_prefix}').replace(osp.sep, "."),
                 sources=[f for f in file_list],
                 define_macros=[
                     cuda_macro,
                     omp_macro
                 ],
-                extra_compile_args=compile_args
+                extra_compile_args=compile_args,
+                use_torch=True
             ))
 
     return extensions
@@ -91,14 +94,14 @@ def load_ops_extensions():
             if not src_files:
                 continue
             if not is_cuda_ext:
-                extensions.append(PyCPUExtension(
+                extensions.append(PyCppExtension(
                     name=osp.join(ops_dir, f'_{ops_prefix}').replace(osp.sep, "."),
                     sources=[osp.join(src_dir, f) for f in src_files],
                     include_dirs=[osp.abspath(osp.join('third_party', d)) for d in ops_third_party_deps[i]],
                     extra_compile_args=['-std=c++17']
                 ))
             else:
-                extensions.append(PyCudaExtension(
+                extensions.append(PyCUDAExtension(
                     name=osp.join(ops_dir, f'_{ops_prefix}_cuda').replace(osp.sep, "."),
                     sources=[osp.join(src_dir, f) for f in src_files],
                     include_dirs=[osp.abspath(osp.join('third_party', d)) for d in ops_third_party_deps[i]],
@@ -136,7 +139,7 @@ def readme():
     author_email="[email protected]",
     maintainer="Tianyu Zhao",
     license="Apache-2.0 License",
-    cmdclass={'build_ext': BuildExtension},
+    cmdclass={'build_ext': PyBuildExtension},
     ext_modules=load_extensions(),
     description=" ",
     long_description=readme(),