Small changes

utils/CMakeLists.txt

@@ -1,22 +0,0 @@
-project(PointNet2)
-cmake_minimum_required(VERSION 3.5)
-
-find_package(CUDA REQUIRED)
-
-include_directories("${CMAKE_SOURCE_DIR}/cinclude")
-cuda_include_directories("${CMAKE_SOURCE_DIR}/cinclude")
-file(GLOB cuda_kernels_src "csrc/*.cu")
-cuda_compile(cuda_kernels SHARED ${cuda_kernels_src} OPTIONS -O3)
-
-file(GLOB wrapper_headers "cinclude/*wrapper.h")
-file(GLOB wrapper_sources "csrs/*.c")
-add_custom_command(OUTPUT "${CMAKE_SOURCE_DIR}/_ext/__ext.so"
-		   WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
-		   COMMAND python "${CMAKE_SOURCE_DIR}/build_ffi.py" ${cuda_kernels}
-		   DEPENDS ${cuda_kernels}
-		   DEPENDS ${wrapper_headers}
-		   DEPENDS ${wrapper_sources}
-		   VERBATIM)
-
-add_custom_target(ext ALL
-	          DEPENDS "${CMAKE_SOURCE_DIR}/_ext/__ext.so")

utils/build_ffi.py

@@ -1,24 +1,50 @@
 import glob
 import torch
-from os import path
+import os.path as osp
 from torch.utils.ffi import create_extension
-import sys
+import sys, argparse, shutil
 
-base_dir = path.dirname(path.abspath(__file__))
-extra_objects = sys.argv[1:]
-extra_objects += [a for a in glob.glob('/usr/local/cuda/lib64/*.a')]
+base_dir = osp.dirname(osp.abspath(__file__))
 
-ffi = create_extension(
-    '_ext',
-    headers=[a for a in glob.glob("cinclude/*_wrapper.h")],
-    sources=[a for a in glob.glob("csrc/*.c")],
-    define_macros=[('WITH_CUDA', None)],
-    relative_to=__file__,
-    with_cuda=True,
-    extra_objects=extra_objects,
-    include_dirs=[path.join(base_dir, 'cinclude')],
-    verbose=False)
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Arguments for building pointnet2 ffi extension")
+    parser.add_argument("--objs", nargs="*")
+    clean_arg = parser.add_mutually_exclusive_group()
+    clean_arg.add_argument("--build", dest='build', action="store_true")
+    clean_arg.add_argument("--clean", dest='clean', action="store_true")
+    parser.set_defaults(build=False, clean=False)
+
+    args = parser.parse_args()
+    assert args.build or args.clean
+
+    return args
+
+
+def build(args):
+    extra_objects = args.objs
+    extra_objects += [a for a in glob.glob('/usr/local/cuda/lib64/*.a')]
+
+    ffi = create_extension(
+        '_ext',
+        headers=[a for a in glob.glob("cinclude/*_wrapper.h")],
+        sources=[a for a in glob.glob("csrc/*.c")],
+        define_macros=[('WITH_CUDA', None)],
+        relative_to=__file__,
+        with_cuda=True,
+        extra_objects=extra_objects,
+        include_dirs=[osp.join(base_dir, 'cinclude')],
+        verbose=False,
+        package=False)
+    ffi.build()
+
+def clean(args):
+    shutil.rmtree(osp.join(base_dir, "_ext"))
 
 if __name__ == "__main__":
-    assert torch.cuda.is_available(), "Needs CUDA!"
-    ffi.build()
+    args = parse_args()
+    if args.clean:
+        clean(args)
+    else:
+        build(args)