[cpp_extension][inductor] Fix sleef windows depends. (#128770) (#128811)

xuhancn · web-flow · commit 04e98d3d0e82 · 2024-06-26T10:24:17.000-04:00
# Issue: During I'm working on enable inductor on PyTorch Windows, I found the sleef lib dependency issue. <img width="1011" alt="image" src="https://github.com/pytorch/pytorch/assets/8433590/423bd854-3c5f-468f-9a64-a392d9b514e3"> # Analysis: After we enabled SIMD on PyTorch Windows(#118980 ), the sleef functions are called from VEC headers. It bring the sleef to the dependency. Here is a different between Windows and Linux OS. ## Linux : Linux is default export its functions, so libtorch_cpu.so static link to sleef.a, and then It also export sleef's functions. <img width="647" alt="image" src="https://github.com/pytorch/pytorch/assets/8433590/00ac536c-33fc-4943-a435-25590508840d"> ## Windows: Windows is by default not export its functions, and have many limitation to export functions, reference: #80604 We can't package sleef functions via torch_cpu.dll like Linux. # Solution: Acturally, we also packaged sleef static lib as a part of release. We just need to help user link to sleef.lib, it should be fine. 1. Add sleef to cpp_builder for inductor. 2. Add sleef to cpp_extension for C++ extesion. Pull Request resolved: #128770 Approved by: https://github.com/jgong5, https://github.com/jansel
diff --git a/torch/_inductor/cpp_builder.py b/torch/_inductor/cpp_builder.py
@@ -571,6 +571,9 @@ def _get_torch_related_args(include_pytorch: bool, aot_mode: bool):
         if not aot_mode:
             libraries.append("torch_python")
 
+    if _IS_WINDOWS:
+        libraries.append("sleef")
+
     # Unconditionally import c10 for non-abi-compatible mode to use TORCH_CHECK - See PyTorch #108690
     if not config.abi_compatible:
         libraries.append("c10")
diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py
@@ -961,6 +961,9 @@ def CppExtension(name, sources, *args, **kwargs):
     libraries.append('torch')
     libraries.append('torch_cpu')
     libraries.append('torch_python')
+    if IS_WINDOWS:
+        libraries.append("sleef")
+
     kwargs['libraries'] = libraries
 
     kwargs['language'] = 'c++'
