Error when non stable/headeronly/shim headers are included by stable extension (pytorch#167855)

Address Nikita's offline comment on pytorch#167496

Pull Request resolved: pytorch#167855
Approved by: https://github.com/janeyx99
ghstack dependencies: pytorch#167496

Author: mikaylagawarecki

.ci/pytorch/smoke_test/check_binary_symbols.py

@@ -145,14 +145,17 @@ def _compile_and_extract_symbols(
 
 def check_stable_only_symbols(install_root: Path) -> None:
     """
-    Test TORCH_STABLE_ONLY and TORCH_TARGET_VERSION by compiling test code and comparing symbol counts.
+    Test TORCH_STABLE_ONLY and TORCH_TARGET_VERSION by compiling test code.
 
     This approach tests:
-    1. WITHOUT macros -> many torch symbols exposed
-    2. WITH TORCH_STABLE_ONLY -> zero torch symbols (all hidden)
-    3. WITH TORCH_TARGET_VERSION -> zero torch symbols (all hidden)
-    4. WITH both macros -> zero torch symbols (all hidden)
+    1. WITHOUT macros -> many torch symbols exposed (compilation succeeds)
+    2. WITH TORCH_STABLE_ONLY -> compilation fails with #error directive
+    3. WITH TORCH_TARGET_VERSION -> compilation fails with #error directive
+    4. WITH both macros -> compilation fails with #error directive
     """
+    import subprocess
+    import tempfile
+
     include_dir = install_root / "include"
     assert include_dir.exists(), f"Expected {include_dir} to be present"
 
@@ -182,7 +185,7 @@ def check_stable_only_symbols(install_root: Path) -> None:
         "-c",  # Compile only, don't link
     ]
 
-    # Compile WITHOUT any macros
+    # Compile WITHOUT any macros - should succeed
     symbols_without = _compile_and_extract_symbols(
         cpp_content=test_cpp_content,
         compile_flags=base_compile_flags,
@@ -196,49 +199,56 @@ def check_stable_only_symbols(install_root: Path) -> None:
         "Expected a non-zero number of symbols without any macros"
     )
 
-    # Compile WITH TORCH_STABLE_ONLY (expect 0 symbols)
-    compile_flags_with_stable_only = base_compile_flags + ["-DTORCH_STABLE_ONLY"]
+    # Helper to verify compilation fails with expected error
+    def _expect_compilation_failure(compile_flags: list[str], macro_name: str) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmppath = Path(tmpdir)
+            cpp_file = tmppath / "test.cpp"
+            obj_file = tmppath / "test.o"
+
+            cpp_file.write_text(test_cpp_content)
+
+            result = subprocess.run(
+                compile_flags + [str(cpp_file), "-o", str(obj_file)],
+                capture_output=True,
+                text=True,
+                timeout=60,
+            )
+
+            if result.returncode == 0:
+                raise RuntimeError(
+                    f"Expected compilation to fail with {macro_name} defined, but it succeeded"
+                )
+
+            stderr = result.stderr
+            expected_error_msg = (
+                "This file should not be included when either TORCH_STABLE_ONLY "
+                "or TORCH_TARGET_VERSION is defined."
+            )
+
+            if expected_error_msg not in stderr:
+                raise RuntimeError(
+                    f"Expected error message to contain:\n  '{expected_error_msg}'\n"
+                    f"but got:\n{stderr[:1000]}"
+                )
+
+            print(f"Compilation correctly failed with {macro_name} defined")
 
-    symbols_with_stable_only = _compile_and_extract_symbols(
-        cpp_content=test_cpp_content,
-        compile_flags=compile_flags_with_stable_only,
-    )
-
-    num_symbols_with_stable_only = len(symbols_with_stable_only)
-    assert num_symbols_with_stable_only == 0, (
-        f"Expected no symbols with TORCH_STABLE_ONLY macro, but found {num_symbols_with_stable_only}"
-    )
+    compile_flags_with_stable_only = base_compile_flags + ["-DTORCH_STABLE_ONLY"]
+    _expect_compilation_failure(compile_flags_with_stable_only, "TORCH_STABLE_ONLY")
 
-    # Compile WITH TORCH_TARGET_VERSION (expect 0 symbols)
     compile_flags_with_target_version = base_compile_flags + [
         "-DTORCH_TARGET_VERSION=1"
     ]
-
-    symbols_with_target_version = _compile_and_extract_symbols(
-        cpp_content=test_cpp_content,
-        compile_flags=compile_flags_with_target_version,
+    _expect_compilation_failure(
+        compile_flags_with_target_version, "TORCH_TARGET_VERSION"
     )
 
-    num_symbols_with_target_version = len(symbols_with_target_version)
-    assert num_symbols_with_target_version == 0, (
-        f"Expected no symbols with TORCH_TARGET_VERSION macro, but found {num_symbols_with_target_version}"
-    )
-
-    # Compile WITH both macros (expect 0 symbols)
     compile_flags_with_both = base_compile_flags + [
         "-DTORCH_STABLE_ONLY",
         "-DTORCH_TARGET_VERSION=1",
     ]
-
-    symbols_with_both = _compile_and_extract_symbols(
-        cpp_content=test_cpp_content,
-        compile_flags=compile_flags_with_both,
-    )
-
-    num_symbols_with_both = len(symbols_with_both)
-    assert num_symbols_with_both == 0, (
-        f"Expected no symbols with both macros, but found {num_symbols_with_both}"
-    )
+    _expect_compilation_failure(compile_flags_with_both, "both macros")
 
 
 def check_stable_api_symbols(install_root: Path) -> None:

setup.py

@@ -1089,6 +1089,60 @@ def check_pydep(importname: str, module: str) -> None:
 
 
 class build_ext(setuptools.command.build_ext.build_ext):
+    def _wrap_headers_with_macro(self, include_dir: Path) -> None:
+        """Wrap all header files with #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION).
+
+        Excludes:
+        - torch/headeronly/*
+        - torch/csrc/stable/*
+        - torch/csrc/inductor/aoti_torch/c/ (only shim headers)
+        - torch/csrc/inductor/aoti_torch/generated/
+
+        This method is idempotent - it will not wrap headers that are already wrapped.
+        """
+        header_extensions = (".h", ".hpp", ".cuh")
+        header_files = [
+            f for ext in header_extensions for f in include_dir.rglob(f"*{ext}")
+        ]
+
+        # Paths to exclude from wrapping (relative to include_dir)
+        exclude_dir_patterns = [
+            "torch/headeronly/",
+            "torch/csrc/stable/",
+            "torch/csrc/inductor/aoti_torch/c/",
+            "torch/csrc/inductor/aoti_torch/generated/",
+        ]
+
+        # Marker to detect if a header is already wrapped
+        wrap_start_marker = (
+            "#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)\n"
+        )
+
+        for header_file in header_files:
+            rel_path = header_file.relative_to(include_dir).as_posix()
+
+            if any(rel_path.startswith(pattern) for pattern in exclude_dir_patterns):
+                report(f"Skipping header: {rel_path}")
+                continue
+
+            original_content = header_file.read_text(encoding="utf-8")
+
+            # Check if already wrapped (idempotency check)
+            if original_content.startswith(wrap_start_marker):
+                report(f"Already wrapped, skipping: {rel_path}")
+                continue
+
+            wrapped_content = (
+                wrap_start_marker
+                + f"{original_content}"
+                + "\n#else\n"
+                + '#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."\n'
+                + "#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)\n"
+            )
+
+            header_file.write_text(wrapped_content, encoding="utf-8")
+            report(f"Wrapped header: {rel_path}")
+
     def _embed_libomp(self) -> None:
         # Copy libiomp5.dylib/libomp.dylib inside the wheel package on MacOS
         build_lib = Path(self.build_lib)
@@ -1256,6 +1310,15 @@ def run(self) -> None:
 
         super().run()
 
+        # Wrap headers with TORCH_STABLE_ONLY and TORCH_TARGET_VERSION guards
+        build_lib = Path(self.build_lib)
+        build_torch_include_dir = build_lib / "torch" / "include"
+        if build_torch_include_dir.exists():
+            report(
+                "-- Wrapping header files with if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)"
+            )
+            self._wrap_headers_with_macro(build_torch_include_dir)
+
         if IS_DARWIN:
             self._embed_libomp()
 
@@ -1358,45 +1421,6 @@ def __exit__(self, *exc_info: object) -> None:
 
 # Need to create the proper LICENSE.txt for the wheel
 class bdist_wheel(setuptools.command.bdist_wheel.bdist_wheel):
-    def _wrap_headers_with_macro(self, bdist_dir: Path) -> None:
-        """Wrap all header files with #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION).
-
-        Excludes:
-        - torch/include/torch/headeronly/*
-        - torch/include/torch/csrc/stable/*
-        - torch/include/torch/csrc/inductor/aoti_torch/c/ (only shim headers)
-        - torch/include/torch/csrc/inductor/aoti_torch/generated/
-        """
-        header_extensions = (".h", ".hpp", ".cuh")
-        header_files = [
-            f for ext in header_extensions for f in bdist_dir.rglob(f"*{ext}")
-        ]
-
-        # Paths to exclude from wrapping
-        exclude_dir_patterns = [
-            "torch/include/torch/headeronly/",
-            "torch/include/torch/csrc/stable/",
-            "torch/include/torch/csrc/inductor/aoti_torch/c/",
-            "torch/include/torch/csrc/inductor/aoti_torch/generated/",
-        ]
-
-        for header_file in header_files:
-            rel_path = header_file.relative_to(bdist_dir).as_posix()
-
-            if any(rel_path.startswith(pattern) for pattern in exclude_dir_patterns):
-                report(f"Skipping header: {rel_path}")
-                continue
-
-            original_content = header_file.read_text(encoding="utf-8")
-            wrapped_content = (
-                "#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)\n"
-                f"{original_content}"
-                "\n#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)\n"
-            )
-
-            header_file.write_text(wrapped_content, encoding="utf-8")
-            report(f"Wrapped header: {rel_path}")
-
     def run(self) -> None:
         with concat_license_files(include_files=True):
             super().run()
@@ -1419,14 +1443,6 @@ def write_wheelfile(self, *args: Any, **kwargs: Any) -> None:
             # need an __init__.py file otherwise we wouldn't have a package
             (bdist_dir / "torch" / "__init__.py").touch()
 
-        # Wrap all header files with TORCH_STABLE_ONLY macro
-        assert self.bdist_dir is not None, "bdist_dir should be set during wheel build"
-        bdist_dir = Path(self.bdist_dir)
-        report(
-            "-- Wrapping header files with if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)"
-        )
-        self._wrap_headers_with_macro(bdist_dir)
-
 
 class clean(Command):
     user_options: ClassVar[list[tuple[str, str | None, str]]] = []