Reduce framework extensions' binary size and fix debug builds (#1769)

* Remove default debug info from distutils
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add assert
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

build_tools/jax.py

@@ -9,7 +9,7 @@ from pathlib import Path
 
 import setuptools
 
-from .utils import get_cuda_include_dirs, all_files_in_dir
+from .utils import get_cuda_include_dirs, all_files_in_dir, debug_build_enabled
 from typing import List
 
 
@@ -57,6 +57,11 @@ def setup_jax_extension(
 
     # Compile flags
     cxx_flags = ["-O3"]
+    if debug_build_enabled():
+        cxx_flags.append("-g")
+        cxx_flags.append("-UNDEBUG")
+    else:
+        cxx_flags.append("-g0")
 
     # Define TE/JAX as a Pybind11Extension
     from pybind11.setup_helpers import Pybind11Extension

build_tools/pytorch.py

@@ -8,7 +8,7 @@ from pathlib import Path
 
 import setuptools
 
-from .utils import all_files_in_dir, cuda_version, get_cuda_include_dirs
+from .utils import all_files_in_dir, cuda_version, get_cuda_include_dirs, debug_build_enabled
 
 
 def setup_pytorch_extension(
@@ -37,10 +37,12 @@ def setup_pytorch_extension(
     )
 
     # Compiler flags
-    cxx_flags = [
-        "-O3",
-        "-fvisibility=hidden",
-    ]
+    cxx_flags = ["-O3", "-fvisibility=hidden"]
+    if debug_build_enabled():
+        cxx_flags.append("-g")
+        cxx_flags.append("-UNDEBUG")
+    else:
+        cxx_flags.append("-g0")
 
     # Version-dependent CUDA options
     try: