[Bazel] Build ATen_CPU_AVX2 lib with AVX2 arch flags enabled (#37381)

Summary:
Make sleef dependency public so that `ATen_CPU_{capability}` libs can depend on it
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37381

Test Plan: CI

Differential Revision: D21273443

Pulled By: malfet

fbshipit-source-id: 7f756c7f3c605e51cf0c27ea37f687913cd48708

BUILD.bazel

@@ -637,6 +637,7 @@ intern_build_aten_ops(
     copts = ATEN_COPTS,
     deps = [
         ":aten_headers",
+        "@sleef",
         "@fbgemm",
     ],
 )

aten.bzl

@@ -1,6 +1,12 @@
 load("@rules_cc//cc:defs.bzl", "cc_library")
 
 CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX", "AVX2"]
+CAPABILITY_COMPILER_FLAGS = {
+    "AVX2": ["-mavx2", "-mfma"],
+    "AVX": ["-mavx"],
+    "DEFAULT": [],
+}
+
 PREFIX = "aten/src/ATen/native/"
 
 def intern_build_aten_ops(copts, deps):
@@ -27,11 +33,12 @@ def intern_build_aten_ops(copts, deps):
             copts = copts + [
                 "-DCPU_CAPABILITY=" + cpu_capability,
                 "-DCPU_CAPABILITY_" + cpu_capability,
-            ],
+            ] + CAPABILITY_COMPILER_FLAGS[cpu_capability],
             deps = deps,
+            linkstatic = 1,
         )
     cc_library(
         name = "ATen_CPU",
-        srcs = ["ATen_CPU_" + cpu_capability for cpu_capability in CPU_CAPABILITY_NAMES],
+        deps = [":ATen_CPU_" + cpu_capability for cpu_capability in CPU_CAPABILITY_NAMES],
         linkstatic = 1,
     )

third_party/sleef.BUILD

@@ -51,7 +51,7 @@ SLEEF_PUBLIC_INCLUDES = [
 ]
 
 SLEEF_VISIBILITY = [
-    "@pytorch//:__subpackages__",
+    "//visibility:public",
 ]
 
 cc_binary(