Improve JAX build tool (#942)

* adding option to select only .cpp files in a dir in the build tool

* change cmake build path

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

build_tools/build_ext.py

@@ -94,7 +94,7 @@ def get_build_ext(extension_cls: Type[setuptools.Extension]):
                 if isinstance(ext, CMakeExtension):
                     print(f"Building CMake extension {ext.name}")
                     # Set up incremental builds for CMake extensions
-                    setup_dir = Path(__file__).resolve().parent
+                    setup_dir = Path(__file__).resolve().parent.parent
                     build_dir = setup_dir / "build" / "cmake"
 
                     # Ensure the directory exists

build_tools/jax.py

@@ -23,7 +23,7 @@ def setup_jax_extension(
     extensions_dir = csrc_source_files / "extensions"
     sources = [
         csrc_source_files / "utils.cu",
-    ] + all_files_in_dir(extensions_dir)
+    ] + all_files_in_dir(extensions_dir, ".cpp")
 
     # Header files
     cuda_home, _ = cuda_path()

build_tools/utils.py

@@ -28,10 +28,12 @@ def debug_build_enabled() -> bool:
     return False
 
 
-def all_files_in_dir(path):
+def all_files_in_dir(path, name_extension=None):
     all_files = []
     for dirname, _, names in os.walk(path):
         for name in names:
+            if name_extension is not None and name_extension not in name:
+                continue
             all_files.append(Path(dirname, name))
     return all_files