Reduce size of packaged PyPI wheels

- Strip all symbols from executable targets, which cuts their size by
  ~35%
- Dedupe the binaries from data_files by providing a python entrypoint.

Author: j2kun

frontend/heir/entrypoints.py

@@ -0,0 +1,33 @@
+"""Entrypoints for shimming heir-opt and heir-translate in the Python path"""
+
+import os
+import sys
+from pathlib import Path
+
+
+def _run_binary(binary_name):
+  """Locate and execute a bundled binary."""
+  # This file is at frontend/heir/entrypoints.py
+  # The binaries are installed in the 'heir' package directory (frontend/heir/)
+  binary_path = Path(__file__).parent / binary_name
+  if not binary_path.exists():
+    # Fallback for some environments where binaries might be in the parent
+    binary_path = Path(__file__).parent.parent / binary_name
+
+  if not binary_path.exists():
+    print(
+        f"Error: {binary_name} not found in package at {binary_path}",
+        file=sys.stderr,
+    )
+    sys.exit(1)
+
+  # os.execv replaces the current process with the binary
+  os.execv(binary_path, [str(binary_path)] + sys.argv[1:])
+
+
+def run_heir_opt():
+  _run_binary("heir-opt")
+
+
+def run_heir_translate():
+  _run_binary("heir-translate")

lib/Pipelines/BooleanPipelineRegistration.cpp

@@ -13,6 +13,7 @@
 #include "lib/Dialect/Debug/Transforms/ValidateNames.h"
 #include "lib/Dialect/Secret/Conversions/SecretToCGGI/SecretToCGGI.h"
 #include "lib/Dialect/Secret/Transforms/DistributeGeneric.h"
+#include "lib/Pipelines/PipelineRegistration.h"
 #include "lib/Transforms/BooleanVectorizer/BooleanVectorizer.h"
 #include "lib/Transforms/FoldConstantTensors/FoldConstantTensors.h"
 #include "lib/Transforms/ForwardInsertToExtract/ForwardInsertToExtract.h"
@@ -22,6 +23,7 @@
 #include "lib/Transforms/MemrefToArith/MemrefToArith.h"
 #include "lib/Transforms/Secretize/Passes.h"
 #include "lib/Transforms/TensorLinalgToAffineLoops/TensorLinalgToAffineLoops.h"
+#include "lib/Transforms/UnusedMemRef/UnusedMemRef.h"
 #include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project
 #include "mlir/include/mlir/Conversion/TensorToLinalg/TensorToLinalgPass.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Affine/Transforms/Passes.h"  // from @llvm-project

pyproject.toml

@@ -59,6 +59,10 @@ dev = [
 Homepage = "https://heir.dev"
 Issues = "https://github.com/google/heir/issues"
 
+[project.scripts]
+heir-opt = "heir.entrypoints:run_heir_opt"
+heir-translate = "heir.entrypoints:run_heir_translate"
+
 [tool.setuptools]
 package-dir = { "" = "frontend" }
 zip-safe = false

setup.py

@@ -101,6 +101,7 @@ def __init__(
       target_file: str,
       is_binary: bool = False,
       copy_include_files: bool = False,
+      aggressive_strip: bool = False,
       **kwargs: Any,
   ):
     super().__init__(name=name, sources=[], **kwargs)
@@ -115,6 +116,11 @@ def __init__(
     self.is_binary = is_binary
     self.copy_include_files = copy_include_files
 
+    # Determines whether to `strip-all` when building a target for PyPI.
+    # Critically, this should only be used for cc_binary targets that are run as
+    # a subprocess, not cc_library targets that are loaded into Python.
+    self.aggressive_strip = aggressive_strip
+
 
 class BuildBazelExtension(build_ext.build_ext):
   """A command that runs Bazel to build a C/C++ extension."""
@@ -226,10 +232,14 @@ def bazel_build(self, ext: BazelExtension) -> None:  # noqa: C901
         "--ui_event_filters=ERROR",
         f"--symlink_prefix={temp_path / 'bazel-'}",
         "--compilation_mode=opt",
+        "--strip=always",
         f"--cxxopt={'/std:c++20' if IS_WINDOWS else '-std=c++20'}",
         f"--@rules_python//python/config_settings:python_version={python_version}",
     ]
 
+    if ext.aggressive_strip:
+      bazel_argv.append("--stripopt=--strip-all")
+
     if is_cibuildwheel() and IS_LINUX:
       # TODO(#2249): OpenMP is disabled for manylinux, as libomp is not on the allowed libraries list,
       # and statically bundling OpenMP is non-trivial and left as future work.
@@ -288,9 +298,10 @@ def bazel_build(self, ext: BazelExtension) -> None:  # noqa: C901
           | stat.S_IXOTH,
       )
 
-      # Also copy binaries to project root so they can be included in data_files
+      # Also copy binaries to project root so they are visible when running from
+      # Python as a subprocess.
       root_path = Path(ext.target_file)
-      print(f"Copying {srcdir_path} to {root_path} for data_files")
+      print(f"Copying {srcdir_path} to {root_path}")
       shutil.copyfile(srcdir_path, root_path)
       os.chmod(
           root_path,
@@ -313,27 +324,30 @@ def bazel_build(self, ext: BazelExtension) -> None:  # noqa: C901
     ext_modules=[
         BazelExtension(
             name="heir_py._heir_opt",
-            bazel_target="//tools:heir-opt",
-            generated_so_file=Path("tools") / "heir-opt",
+            bazel_target="//tools:heir-opt.stripped",
+            generated_so_file=Path("tools") / "heir-opt.stripped",
             target_file="heir-opt",
             py_limited_api=py_limited_api,
             is_binary=True,
+            aggressive_strip=is_cibuildwheel(),
         ),
         BazelExtension(
             name="heir_py._heir_translate",
-            bazel_target="//tools:heir-translate",
-            generated_so_file=Path("tools") / "heir-translate",
+            bazel_target="//tools:heir-translate.stripped",
+            generated_so_file=Path("tools") / "heir-translate.stripped",
             target_file="heir-translate",
             py_limited_api=py_limited_api,
             is_binary=True,
+            aggressive_strip=is_cibuildwheel(),
         ),
         BazelExtension(
             name="heir_py._abc",
-            bazel_target="@abc//:abc_bin",
-            generated_so_file=Path("external") / "abc+" / "abc_bin",
+            bazel_target="@abc//:abc_bin.stripped",
+            generated_so_file=Path("external") / "abc+" / "abc_bin.stripped",
             target_file="abc_bin",
             py_limited_api=py_limited_api,
             is_binary=True,
+            aggressive_strip=is_cibuildwheel(),
         ),
         BazelExtension(
             name="heir_py._libopenfhe",
@@ -342,8 +356,8 @@ def bazel_build(self, ext: BazelExtension) -> None:  # noqa: C901
             target_file="libopenfhe.so",
             py_limited_api=py_limited_api,
             copy_include_files=True,
+            aggressive_strip=False,
         ),
     ],
-    data_files=[("bin", ["heir-opt", "heir-translate"])],
     options=options,
 )