Update on "[Executorch][llm] Compile custom op with -O2"

This enable some optimizations for inlining vectorized lib functions

Differential Revision: [D64910576](https://our.internmc.facebook.com/intern/diff/D64910576/)

[ghstack-poisoned]

Author: kimishpatel

extension/llm/custom_ops/targets.bzl

@@ -1,4 +1,9 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load(
+    "@fbsource//xplat/executorch/kernels/portable:op_registration_util.bzl",
+    "get_compiler_optimization_flags",
+)
+
 
 def define_common_targets():
     """Defines targets that should be shared between fbcode and xplat.
@@ -34,21 +39,7 @@ def define_common_targets():
                 "//executorch/kernels/portable/cpu/util:reduce_util",
                 "//executorch/extension/llm/custom_ops/spinquant:fast_hadamard_transform",
             ],
-            compiler_flags = ["-Wno-missing-prototypes", "-Wno-global-constructors"] + select({
-              "DEFAULT": [],
-              "ovr_config//os:android-arm64": [
-                    "-O2",
-              ] if not runtime.is_oss else [],
-              "ovr_config//os:iphoneos": [
-                  "-O2",
-              ] if not runtime.is_oss else [],
-              "ovr_config//os:macos-arm64": [
-                  "-O2",
-              ] if not runtime.is_oss else [],
-              "ovr_config//os:macos-x86_64": [
-                  "-O2",
-              ] if not runtime.is_oss else [],
-            }),
+            compiler_flags = ["-Wno-missing-prototypes", "-Wno-global-constructors"] + get_compiler_optimization_flags(),
             visibility = [
                 "//executorch/...",
                 "//executorch/extension/llm/custom_ops/...",

kernels/optimized/lib_defs.bzl

@@ -39,6 +39,23 @@ def get_vec_fbcode_preprocessor_flags():
     ]
     return preprocessor_flags
 
+def get_compiler_optimization_flags():
+    if not runtime.is_oss:
+      compiler_flags = select({
+        "DEFAULT": [],
+        "ovr_config//os:android-arm64": [
+              "-O2",
+        ],
+        "ovr_config//os:iphoneos": [
+            "-O2",
+        ],
+        "ovr_config//os:macos-arm64": [
+            "-O2",
+        ],
+      })
+      return compiler_flags
+    return []
+
 # Currently, having a dependency on fbsource//third-party/sleef:sleef may cause
 # duplicate symbol errors when linking fbcode targets in opt mode that also
 # depend on ATen. This is because ATen accesses sleef via the third-party folder
@@ -121,21 +138,7 @@ def define_libs():
             exported_headers = native.glob([
                 "blas/**/*.h",
             ]),
-            compiler_flags = select({
-              "DEFAULT": [],
-              "ovr_config//os:android-arm64": [
-                    "-O2",
-              ] if not runtime.is_oss else [],
-              "ovr_config//os:iphoneos": [
-                  "-O2",
-              ] if not runtime.is_oss else [],
-              "ovr_config//os:macos-arm64": [
-                  "-O2",
-              ] if not runtime.is_oss else [],
-              "ovr_config//os:macos-x86_64": [
-                  "-O2",
-              ] if not runtime.is_oss else [],
-            }),
+            compiler_flags = get_compiler_optimization_flags(),
             header_namespace = "executorch/kernels/optimized",
             visibility = [
                 "//executorch/...",

kernels/optimized/op_registration_util.bzl

@@ -3,6 +3,7 @@ load("@fbsource//xplat/executorch/build:selects.bzl", "selects")
 load(
     "@fbsource//xplat/executorch/kernels/optimized:lib_defs.bzl",
     "get_vec_android_preprocessor_flags",
+    "get_compiler_optimization_flags",
 )
 
 def op_target(name, deps = []):
@@ -87,21 +88,7 @@ def define_op_library(name, deps):
         ],
         # kernels often have helpers with no prototypes just disabling the warning here as the headers
         # are codegend and linked in later
-        compiler_flags = ["-Wno-missing-prototypes"] + select({
-          "DEFAULT": [],
-          "ovr_config//os:android": [
-                "-O2",
-          ] if not runtime.is_oss else [],
-          "ovr_config//os:iphoneos": [
-              "-O2",
-          ] if not runtime.is_oss else [],
-          "ovr_config//os:macos-arm64": [
-              "-O2",
-          ] if not runtime.is_oss else [],
-          "ovr_config//os:macos-x86_64": [
-              "-O2",
-          ] if not runtime.is_oss else [],
-        }),
+        compiler_flags = ["-Wno-missing-prototypes"] + get_compiler_optimization_flags(),
         deps = [
             "//executorch/runtime/kernel:kernel_includes",
         ] + augmented_deps,