Ship the quantized kernels as their own library

.ci/scripts/wheel/test_cpp_sdk.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -58,6 +58,25 @@
 with torch.no_grad():
     expected = model(*example)
 
+if mode == "quantized":
+    # Quantize with the same flow the documentation shows, so the exported program
+    # references the quantized operator set rather than the plain one.
+    # Importing this loads the ahead-of-time library, which is what registers the out
+    # variants of the quantized operators with torch. Without it the export fails with
+    # "Missing out variants: quantized_decomposed::quantize_per_tensor", because the
+    # lowering step has no out variant to select.
+    import executorch.kernels.quantized  # noqa: F401
+    from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
+        get_symmetric_quantization_config,
+        XNNPACKQuantizer,
+    )
+    from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+
+    quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
+    prepared = prepare_pt2e(torch.export.export(model, example).module(), quantizer)
+    prepared(*example)
+    model = convert_pt2e(prepared)
+
 partitioners = []
 if mode == "delegate":
     from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
@@ -86,6 +105,11 @@
             "expected": expected.flatten().tolist(),
             "delegated": mode == "delegate",
             "has_xnnpack": b"XnnpackBackend" in bytes(buffer),
+            # Whether the program actually carries quantized operators. The numeric comparison alone
+            # cannot tell: an unquantized export of the same model produces a closer match than the
+            # tolerance a quantized one needs, so it would pass while proving nothing about the
+            # quantized kernels.
+            "has_quantized": b"quantized_decomposed" in bytes(buffer),
         }
     )
 )
@@ -189,8 +213,14 @@
   for (size_t i = 0; i < expected.size(); ++i) {
     worst = std::fmax(worst, std::fabs((double)actual[i] - (double)expected[i]));
   }
-  if (worst > 1e-4) {
-    std::printf("output differs from eager PyTorch by %g\n", worst);
+  // Passed in rather than fixed, because the acceptable difference depends on the
+  // model. A float32 model should match to within rounding, while an int8 quantized one
+  // legitimately differs by about one quantization step, and using the looser number
+  // for both would stop the float path catching a real regression.
+  const double tolerance = argc > 7 ? std::atof(argv[7]) : 1e-4;
+  if (worst > tolerance) {
+    std::printf(
+        "output differs from eager PyTorch by %g, tolerance %g\n", worst, tolerance);
     return 1;
   }
 
@@ -328,8 +358,16 @@
     return consumer
 
 
-def _run_consumer(consumer: Path, model: Path, reference, work_dir: Path) -> str:
-    """Run the application and require it to match eager PyTorch."""
+def _run_consumer(
+    consumer: Path, model: Path, reference, work_dir: Path, tolerance: float = 1e-4
+) -> str:
+    """Run the application and require it to match eager PyTorch within `tolerance`.
+
+    The tolerance is a parameter because the acceptable difference depends on the model.
+    A float32 model should match to within rounding, while an int8 quantized one
+    legitimately differs by about one quantization step, and using the looser number for
+    both would stop the float path catching a real regression.
+    """
     inputs = reference["inputs"]
     shape_a, data_a = _write_tensor(work_dir, "a", inputs[0])
     shape_b, data_b = _write_tensor(work_dir, "b", inputs[1])
@@ -351,6 +389,7 @@
             str(shape_b),
             str(data_b),
             str(expected),
+            repr(tolerance),
         ],
         capture_output=True,
         text=True,
@@ -823,12 +862,53 @@
     print("✓ the C++ example in the documentation compiles against the wheel")
 
 
+def test_quantized_kernels_component_runs_a_model(work_dir: Path) -> None:
+    """A C++ application can run a quantized model using the shipped quantized kernels.
+
+    Before the quantized kernels became their own library they existed only inside the
+    ahead-of-time extension beside the Python bindings, so a C++ application loading a
+    quantized program had nothing to link and failed at run time with the operators
+    reported missing.
+
+    Skipped rather than failed when the wheel ships no such library, because building
+    without the quantized kernels is a supported configuration.
+    """
+    package_dir = _installed_package_dir()
+    shipped = package_dir / "lib" / "libexecutorch_kernels_quantized.so"
+    if not shipped.is_file():
+        print("- this wheel ships no quantized kernels, skipping")
+        return
+
+    model, reference = _export(work_dir, "quantized")
+    # The export has to have produced a quantized program, or the rest of this proves nothing about the
+    # quantized kernels. The numeric comparison cannot tell the difference: an unquantized export of the
+    # same model lands well inside the tolerance a quantized one needs, so it would pass while linking a
+    # library it never exercised.
+    assert reference["has_quantized"], (
+        "the quantized export produced a program with no quantized operators, so this check would "
+        "prove nothing about the quantized kernels"
+    )
+    consumer = _build_consumer(
+        work_dir,
+        "with-quantized",
+        ["runtime", "kernels_optimized", "kernels_quantized"],
+    )
+    # One int8 quantization step over this model's output range is about 5e-3, so a
+    # float32 tolerance cannot be met by a correct quantized run.
+    output = _run_consumer(consumer, model, reference, work_dir, tolerance=2e-2)
+    print(
+        f"✓ a C++ app linking executorch::kernels_quantized runs a quantized model "
+        f"({output})"
+    )
+
+
 def run_tests(work_dir: Path) -> None:
     test_find_package_honours_a_version_request(work_dir)
     test_every_shipped_header_compiles(work_dir)
     test_documented_example_compiles(work_dir)
     test_runtime_alone_links_but_has_no_kernels(work_dir)
     test_kernels_component_runs_a_model(work_dir)
+    test_quantized_kernels_component_runs_a_model(work_dir)
     test_delegated_model_needs_the_delegate_component(work_dir)
     test_consumer_is_relocatable(work_dir)
     test_one_registry_in_the_cpp_process(work_dir)

.ci/scripts/wheel/test_shared_libraries.py

@@ -52,6 +52,14 @@
 # the operators are registered twice, which aborts at startup.
 _KERNEL_SYMBOLS = ("torch::executor::native::abs_out",)
 
+# The quantized kernels, whose own library the wheel ships when they are built.
+# A separate group because they have a separate owner, and because a wheel built
+# without them ships neither the library nor these symbols.
+_QUANTIZED_KERNEL_SYMBOLS = (
+    "torch::executor::native::quantize_per_tensor_out",
+    "torch::executor::native::dequantize_per_tensor_out",
+)
+
 # The registry entry points, kept separate from the kernel implementations above.
 # A library that carries its own copy of these has its own registration code, which
 # is what this split is meant to prevent: one owner of the operator table. Checking
@@ -238,6 +246,33 @@ def _defines_symbol(library: Path, symbol: str) -> bool:
     return False
 
 
+def _is_export_only(library: Path) -> bool:
+    """Whether a library exists to export a model rather than to run one.
+
+    The ahead-of-time operator libraries register kernels into torch so a model can be
+    exported, and they link torch to do it. They deliberately carry their own copy of
+    the kernels, because export happens in a Python process that never loads the
+    runtime libraries a C++ application links.
+
+    Excluded from the single-owner checks for that reason. Counting them would report a
+    duplicate for something that is not one, and the alternative, making them resolve
+    the kernels from the shipped library, would mean an export-time library depending on
+    a runtime layout it never uses.
+
+    Matched by what the file links rather than by its name, so a library renamed later
+    is still recognised.
+    """
+    if _tool("readelf") is None:
+        return library.name.endswith("_aot_lib.so")
+    dynamic = subprocess.run(
+        [_tool("readelf"), "-d", str(library)],
+        capture_output=True,
+        text=True,
+        check=False,
+    ).stdout
+    return "libtorch.so" in dynamic and "_aot_lib" in library.name
+
+
 def _assert_single_definer(symbols, what: str, owner: str | None = None) -> None:
     """At most one shipped library may define each of `symbols`.
 
@@ -253,7 +288,11 @@ def _assert_single_definer(symbols, what: str, owner: str | None = None) -> None
     assert _tool("nm") is not None, "nm is required to inspect the wheel"
 
     package_dir = _installed_package_dir()
-    libraries = _shipped_shared_objects(package_dir)
+    libraries = [
+        library
+        for library in _shipped_shared_objects(package_dir)
+        if not _is_export_only(library)
+    ]
     assert libraries, f"no shared libraries found under {package_dir}"
 
     # Every symbol is resolved before anything is reported, so a component that is only
@@ -326,6 +365,12 @@ def _assert_single_definer(symbols, what: str, owner: str | None = None) -> None
         "libexecutorch_kernels_optimized.so",
         False,
     ),
+    (
+        "set of quantized kernels",
+        _QUANTIZED_KERNEL_SYMBOLS,
+        "libexecutorch_kernels_quantized.so",
+        False,
+    ),
     # The third-party code these libraries bundle, checked separately from the
     # wrappers above. A wrapper can have a single owner while the implementation
     # underneath it is bundled into two of these, which is two real thread pools or
@@ -1135,8 +1180,24 @@ def test_extension_contains_no_component() -> None:
         ).stdout.splitlines()
         if "NEEDED" in line
     }
-    shipped = {path.name for path in _shipped_runtime_libraries(package_dir)}
-    unused = sorted(shipped - needed)
+    # Only the libraries whose code the extension used to contain. Those are the ones
+    # this split moved out of it, so the extension must now resolve them from outside or
+    # a retention option silently failed.
+    #
+    # Not every shipped library serves Python. The quantized kernels and the CUDA
+    # delegate exist for a C++ application: Python registers quantized operators through
+    # the torch-linked ahead-of-time library at export time, and never loads the CUDA
+    # delegate from this extension at all. Requiring a dependency on those would demand
+    # the extension link code it has no use for.
+    expected = {
+        name
+        for name in (path.name for path in _shipped_runtime_libraries(package_dir))
+        if not any(
+            marker in name
+            for marker in ("kernels_quantized", "backend_cuda", "extension_cuda")
+        )
+    }
+    unused = sorted(expected - needed)
     assert not unused, (
         f"the wheel ships {unused} but {extension.name} does not depend on them, so "
         "either they are dead weight or a retention option did not hold"
@@ -1167,7 +1228,7 @@ def test_extension_contains_no_component() -> None:
     print(
         f"✓ {extension.name} ({extension.stat().st_size // 1024} KiB) contains no "
         f"component, imports the runtime symbols it uses, and depends on all "
-        f"{len(shipped)} shipped libraries"
+        f"{len(expected)} shipped libraries it used to contain"
     )
 
 
@@ -1211,6 +1272,7 @@ def test_shipped_library_names_are_expected() -> None:
     known = (
         "libexecutorch",
         "libexecutorch_kernels_optimized",
+        "libexecutorch_kernels_quantized",
         "libexecutorch_backend_xnnpack",
         "libexecutorch_threadpool",
         "libexecutorch_etdump",

kernels/quantized/CMakeLists.txt

@@ -192,4 +192,15 @@ if(EXECUTORCH_BUILD_SHARED)
     executorch_quantized_ops quantized_ops_lib quantized_kernels
     executorch_shared
   )
+  # Named after what the library provides rather than after the target that
+  # produces it, matching the optimized kernels next to it, so the shipped file
+  # reads as libexecutorch_kernels_quantized.so. The target name stays as it is
+  # because a source build already refers to it.
+  set_target_properties(
+    executorch_quantized_ops PROPERTIES OUTPUT_NAME
+                                        executorch_kernels_quantized
+  )
+  # Ships beside libexecutorch.so in the wheel's lib/ directory, so it resolves
+  # the runtime from there rather than from wherever it was built.
+  executorch_target_shipped_runtime_path(executorch_quantized_ops)
 endif()

setup.py

@@ -1196,6 +1196,16 @@ def run(self):  # noqa C901
             if cmake_cache.is_enabled("EXECUTORCH_BUILD_MLX"):
                 cmake_build_args += ["--target", "mlxdelegate"]
 
+            # Named explicitly because nothing else links it. The other shipped
+            # libraries are built as dependencies of the Python extension, but a C++
+            # application is the only consumer of this one, so without naming it the
+            # target is generated and never built, and packaging then looks for a file
+            # that does not exist.
+            if cmake_cache.is_enabled("EXECUTORCH_BUILD_SHARED") and (
+                cmake_cache.is_enabled("EXECUTORCH_BUILD_KERNELS_QUANTIZED")
+            ):
+                cmake_build_args += ["--target", "executorch_quantized_ops"]
+
             if cmake_cache.is_enabled("EXECUTORCH_BUILD_KERNELS_LLM_AOT"):
                 cmake_build_args += ["--target", "custom_ops_aot_lib"]
                 cmake_build_args += ["--target", "quantized_ops_aot_lib"]
@@ -1313,6 +1323,19 @@ def run(self):  # noqa C901
                         "EXECUTORCH_BUILD_KERNELS_OPTIMIZED",
                     ],
                 ),
+                # The quantized kernels, as their own library rather than code
+                # fused into the AOT-only extension beside the Python bindings.
+                # A C++ application running a quantized model could not link
+                # them before.
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/kernels/quantized/",
+                    src_name="libexecutorch_kernels_quantized.so",
+                    dst="executorch/lib/libexecutorch_kernels_quantized.so",
+                    dependent_cmake_flags=[
+                        "EXECUTORCH_BUILD_SHARED",
+                        "EXECUTORCH_BUILD_KERNELS_QUANTIZED",
+                    ],
+                ),
                 # Install the XNNPACK delegate beside them, so a process has one
                 # copy of it instead of one per component that uses it.
                 BuiltFile(

tools/cmake/executorch-wheel-config.cmake

@@ -378,6 +378,9 @@ executorch_define_component(threadpool executorch_threadpool)
 # checks, so it has to be defined here or a consumer following the documentation
 # gets a bare name that CMake hands to the linker as a literal flag.
 executorch_define_component(kernels_optimized executorch_kernels_optimized)
+# The quantized kernels, optional in the same way: a wheel built without them
+# simply has no such library and the component is not defined.
+executorch_define_component(kernels_quantized executorch_kernels_quantized)
 # The profiler. A C++ application could not record timing data from an installed
 # package before, because the implementation shipped only inside the Python
 # extension.