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

[xuhancn]

Issue:

During I'm working on enable inductor on PyTorch Windows, I found the sleef lib dependency issue.

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.

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

cc @peterjc123 @mszhanyi @skyline75489 @nbcsm @vladimir-aubrecht @iremyux @Blackhex @cristianPanaite @malfet @zou3519 @jbschlosser @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128811

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 8 New Failures, 1 Unrelated Failure

As of commit f0189bd with merge base b66e3f0 ():

NEW FAILURES - The following jobs have failed:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• pull / linux-focal-py3_8-clang9-xla / test (xla, 1, 1, linux.12xlarge) (gh)
  test_repeat_truncated
• pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, linux.2xlarge) (gh)
  onnx/test_fx_op_consistency.py::TestOnnxModelOutputConsistency_opset_version_18_model_type_TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAMCPU::test_output_match_nn_functional_huber_loss_cpu_float16

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.24xl.spr-metal, unstable) (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[xuhancn]

@pytorchbot rebase -b release/2.4

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/release/2.4. Check the current status here

[pytorchmergebot]

Tried to rebase and push PR #128811, but it was already up to date. Try rebasing against main by issuing:
@pytorchbot rebase -b main

[malfet]

Do we have any testing for the feature?

[xuhancn]

@pytorchmergebot rebase -b release/2.4

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/release/2.4. Check the current status here

[pytorchmergebot]

Successfully rebased xu_fix_windows_sleef_depends_24 onto refs/remotes/origin/release/2.4, please pull locally before adding more changes (for example, via git checkout xu_fix_windows_sleef_depends_24 && git pull --rebase)

[xuhancn]

Do we have any testing for the feature?

1. the origin PR is passed all PreCI with cpp-extension and binaries_libtorch tag, I believe it passed the libtorch and cpp_extension UTs.
2. I double comfirmed with cmake config file,

   pytorch/cmake/TorchConfig.cmake.in

   Line 116 in b5d5416

   append_torchlib_if_found(sleef asmjit)

   add sleef as dependency is existing.

[xuhancn]

Hi @malfet , @atalman

I write a small UT for this PR:

vec.cpp

#include <torch/extension.h>
#if defined(CPU_CAPABILITY_AVX512) || defined(CPU_CAPABILITY_AVX2) || defined(CPU_CAPABILITY_ZVECTOR) || defined(CPU_CAPABILITY_NEON)
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
#endif

#if defined(__GNUC__)
#define __at_align__ __attribute__((aligned(64)))
#elif defined(_WIN32)
#define __at_align__ __declspec(align(64))
#else
#define __at_align__
#endif
__at_align__ float in_out_ptr0[16] = {0.0};

extern "C" void __avx_chk_kernel() {
    auto tmp0 = at::vec::Vectorized<float>(1);
    auto tmp1 = tmp0.exp();
    tmp1.store(in_out_ptr0);
}

PYBIND11_MODULE(vec, m) {
  m.def("__avx_chk_kernel", &__avx_chk_kernel, "");
}

setup.py

from setuptools import setup

from torch.utils.cpp_extension import (
    CppExtension,
    BuildExtension,
)

setup(
    name="vec",
    ext_modules=[
        CppExtension(
            "vec",
            [
                "vec.cpp",
            ],
            extra_compile_args=['/arch:AVX2', '/DCPU_CAPABILITY_AVX2'],
        )
    ],
    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
)

reproduce steps:

python setup.py install

result on release/2.4 branch:

result on xuhancn:xu_fix_windows_sleef_depends_24 this branch:

[xuhancn]

@pytorchmergebot rebase -b release/2.4

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/release/2.4. Check the current status here

[pytorchmergebot]

Successfully rebased xu_fix_windows_sleef_depends_24 onto refs/remotes/origin/release/2.4, please pull locally before adding more changes (for example, via git checkout xu_fix_windows_sleef_depends_24 && git pull --rebase)

[xuhancn]

@atalman Any comments?