Fix wheel build and smoke test (pytorch#4429)

Summary: Pull Request resolved: pytorch#4429

Test Plan:
1. Label `ciflow/binaries` to build wheel
2. CI and OSS CI. This makes sure that smoke test can pass, and a portable model can export and run E2E. Python 3.9 is not supported for now.
3. Run these helps validate unittest failure for quantized ops

```
pip uninstall executorch -y
rm dist/*.whl
source build/packaging/env_var_script_m1.sh # or Linux
python setup.py bdist_wheel
pip install dist/*.whl
```

Then in Python, try
```
from executorch.extension.pybindings import portable_lib
from executorch.exir.dialects._ops import ops
import executorch.kernels.quantized
ops.edge.quantized_decomposed.add.default.to_out_variant().name()
```

Reviewed By: larryliu0820

Differential Revision: D60694247

Pulled By: kirklandsign

fbshipit-source-id: 4da77493c0e7a868f86b583c9c33d54d08c3e48b

CMakeLists.txt

@@ -680,7 +680,6 @@ if(EXECUTORCH_BUILD_PYBIND)
       executorch
       extension_data_loader
       portable_ops_lib
-      quantized_ops_aot_lib
       util
       torch
   )
@@ -701,7 +700,6 @@ if(EXECUTORCH_BUILD_PYBIND)
 
   if(EXECUTORCH_BUILD_KERNELS_QUANTIZED)
     target_link_options_shared_lib(quantized_ops_lib)
-    list(APPEND _dep_libs quantized_kernels quantized_ops_lib)
   endif()
 
   # compile options for pybind
@@ -755,6 +753,25 @@ if(EXECUTORCH_BUILD_PYBIND)
                  # the torch libs are in `site-packages/torch/lib`.
                  BUILD_RPATH "@loader_path/../../../torch/lib"
                  INSTALL_RPATH "@loader_path/../../../torch/lib"
+                 # Assume <executorch> is the root `site-packages/executorch`
+                 # Need to add <executorch>/extension/llm/custom_ops for
+                 # libcustom_ops_aot_lib.dylib
+                 BUILD_RPATH "@loader_path/../../extension/llm/custom_ops"
+                 INSTALL_RPATH "@loader_path/../../extension/llm/custom_ops"
+                 # Need to add <executorch>/kernels/quantized for
+                 # libquantized_ops_aot_lib.dylib
+                 BUILD_RPATH "@loader_path/../../kernels/quantized"
+                 INSTALL_RPATH "@loader_path/../../kernels/quantized"
+    )
+  else()
+    set_target_properties(
+      portable_lib
+      PROPERTIES # Assume <executorch> is the root `site-packages/executorch`
+                 # Need to add <executorch>/extension/llm/custom_ops for
+                 # libcustom_ops_aot_lib
+                 # Need to add <executorch>/kernels/quantized for
+                 # libquantized_ops_aot_lib
+                 BUILD_RPATH "$ORIGIN:$ORIGIN/../../extension/llm/custom_ops:$ORIGIN/../../kernels/quantized"
     )
   endif()
 

build/Codegen.cmake

@@ -151,7 +151,7 @@ function(gen_custom_ops_aot_lib)
 
   target_link_options_shared_lib(${GEN_LIB_NAME})
   if(EXECUTORCH_BUILD_PYBIND AND APPLE)
-    target_link_libraries(${GEN_LIB_NAME} PRIVATE executorch_no_prim_ops_shared)
+    target_link_libraries(${GEN_LIB_NAME} PRIVATE executorch_no_prim_ops)
     target_link_options(${GEN_LIB_NAME} PRIVATE -undefined dynamic_lookup)
   else()
     target_link_libraries(${GEN_LIB_NAME} PRIVATE executorch_no_prim_ops)

build/packaging/post_build_script.sh

@@ -7,4 +7,8 @@
 
 set -eux
 
-echo "This script is run after building ExecuTorch binaries"
+# This script is run after building ExecuTorch binaries
+
+# Rename pip-out directory, to avoid using shared libraries in pip-out during
+# smoke test.
+mv pip-out BACKUP-pip-out

kernels/quantized/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+try:
+    from pathlib import Path
+
+    libs = list(Path(__file__).parent.resolve().glob("**/libquantized_ops_aot_lib.*"))
+    del Path
+    assert len(libs) == 1, f"Expected 1 library but got {len(libs)}"
+    import torch as _torch
+
+    _torch.ops.load_library(libs[0])
+    del _torch
+except:
+    import logging
+
+    logging.info("libquantized_ops_aot_lib is not loaded")
+    del logging

kernels/quantized/targets.bzl

@@ -88,3 +88,15 @@ def define_common_targets():
             ],
             define_static_targets = True,
         )
+
+    runtime.python_library(
+        name = "quantized_ops_lib",
+        srcs = ["__init__.py"],
+        deps = [
+            "//caffe2:torch",
+        ],
+        visibility = [
+            "//executorch/kernels/quantized/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+    )

kernels/quantized/test/TARGETS

@@ -15,5 +15,6 @@ python_unittest(
         "//caffe2:torch",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/passes:quant_fusion_pass",
+        "//executorch/kernels/quantized:quantized_ops_lib",
     ],
 )

kernels/quantized/test/test_out_variants.py

@@ -7,6 +7,8 @@
 
 import unittest
 
+import executorch.kernels.quantized  # noqa[F401] 'executorch.kernels.quantized' imported but unused
+
 import torch
 import torch.ao.quantization.fx._decomposed  # noqa[F401] 'torch.ao.quantization.fx._decomposed' imported but unused
 from executorch.exir.dialects._ops import ops

setup.py

@@ -489,8 +489,11 @@ def run(self):
             cmake_args += [
                 "-DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON",  # add llama sdpa ops to pybindings.
                 "-DEXECUTORCH_BUILD_KERNELS_CUSTOM_AOT=ON",
+                "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON",  # add quantized ops to pybindings.
+                "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON",
             ]
             build_args += ["--target", "custom_ops_aot_lib"]
+            build_args += ["--target", "quantized_ops_aot_lib"]
         # Allow adding extra cmake args through the environment. Used by some
         # tests and demos to expand the set of targets included in the pip
         # package.
@@ -570,7 +573,14 @@ def get_ext_modules() -> List[Extension]:
             # Install the prebuilt library for custom ops used in llama.
             BuiltFile(
                 "extension/llm/custom_ops/libcustom_ops_aot_lib.*",
-                "executorch/extension/llm/custom_ops",
+                "executorch/extension/llm/custom_ops/",
+            )
+        )
+        ext_modules.append(
+            # Install the prebuilt library for quantized ops required by custom ops.
+            BuiltFile(
+                "kernels/quantized/libquantized_ops_aot_lib.*",
+                "executorch/kernels/quantized/",
             )
         )
 
@@ -594,6 +604,7 @@ def get_ext_modules() -> List[Extension]:
         "executorch/examples/models": "examples/models",
         "executorch/exir": "exir",
         "executorch/extension": "extension",
+        "executorch/kernels/quantized": "kernels/quantized",
         "executorch/schema": "schema",
         "executorch/sdk": "sdk",
         "executorch/sdk/bundled_program": "sdk/bundled_program",