Add pywrap_quantization_lib.h/cc to provide a middle layer for exposing symbols to shared libraries.

PiperOrigin-RevId: 610248507

Author: dansuh17

tensorflow/BUILD

@@ -1377,7 +1377,7 @@ tf_cc_shared_library(
         "//tensorflow/compiler/mlir/lite/quantization/lite:quantize_model",
         "//tensorflow/compiler/mlir/lite/quantization:quantization_config",
         "//tensorflow/compiler/mlir/lite/sparsity:sparsify_model",
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq_impl",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/python:pywrap_quantization_lib_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:custom_aggregator_op",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:quantize_model_cc_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow:passes",

tensorflow/compiler/mlir/lite/BUILD

@@ -1383,7 +1383,6 @@ cc_library(
         "//tensorflow/compiler/mlir/lite/stablehlo:transforms",
         "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
         "//tensorflow/compiler/mlir/quantization/stablehlo:quantize_passes",
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq_impl",  # buildcleaner: keep; prevents undefined reference
         "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
         "//tensorflow/compiler/mlir/quantization/tensorflow:quantize_passes",
         "//tensorflow/compiler/mlir/quantization/tensorflow:quantize_preprocess",

tensorflow/compiler/mlir/lite/quantization/stablehlo/BUILD

@@ -37,7 +37,6 @@ tf_cc_test(
         "//tensorflow/cc/saved_model:loader",
         "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc:io",
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq_impl",  # buildcleaner: keep; prevents undefined reference
         "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibrator_singleton_impl",  # buildcleaner: keep; prevents undefined reference
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

tensorflow/compiler/mlir/lite/stablehlo/BUILD

@@ -647,7 +647,6 @@ tf_cc_binary(
         "//tensorflow/compiler/mlir/lite:flatbuffer_export",
         "//tensorflow/compiler/mlir/lite:tf_to_tfl_flatbuffer",
         "//tensorflow/compiler/mlir/lite/stablehlo/serializer:flatbuffer_export",
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow:quantize_preprocess",
         "//tensorflow/compiler/mlir/quantization/tensorflow:tf_quant_ops",
         "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibrator_singleton_impl",

tensorflow/compiler/mlir/quantization/stablehlo/cc/BUILD

@@ -33,21 +33,8 @@ cc_library(
     ],
 )
 
-# OSS: This is a header-only target. Do NOT directly depend on `config_impl` unless it is necessary
-# (e.g. undefined symbol error), to avoid ODR violation.
 cc_library(
     name = "config",
-    hdrs = ["config.h"],
-    compatible_with = get_compatible_with_portable(),
-    deps = [
-        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
-    ],
-)
-
-# OSS: This is a impl target corresponding to `config`. Do NOT directly depend on `config_impl`
-# unless it is necessary (e.g. undefined symbol error), to avoid ODR violation.
-cc_library(
-    name = "config_impl",
     srcs = ["config.cc"],
     hdrs = ["config.h"],
     compatible_with = get_compatible_with_portable(),
@@ -303,44 +290,11 @@ cc_library(
     ],
 )
 
-# OSS: This is a header-only target. Do NOT directly depend on `static_range_ptq_impl` unless it is
-# necessary (e.g. undefined symbol error), to avoid ODR violation.
 cc_library(
     name = "static_range_ptq",
-    hdrs = ["static_range_ptq.h"],
-    compatible_with = get_compatible_with_portable(),
-    # Must be header-only or unique (i.e. the same target is not a dependency of
-    # tensorflow_framework, tensorflow_cc, or pywrap_tensorflow_internal) dependencies.
-    deps = [
-        ":component",
-        ":types",
-        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
-        "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
-        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
-        "//tensorflow/core:protos_all_cc",
-        "@com_google_absl//absl/base:nullability",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:string_view",
-        "@llvm-project//mlir:IR",
-    ],
-)
-
-cc_library(
-    name = "static_range_ptq_impl",
     srcs = ["static_range_ptq.cc"],
     hdrs = ["static_range_ptq.h"],
     compatible_with = get_compatible_with_portable(),
-    visibility = [
-        "//tensorflow:__pkg__",
-        "//tensorflow/compiler/mlir/lite:__pkg__",  # For tf_tfl_translate binary.
-        # For odml_to_stablehlo binary.
-        "//tensorflow/compiler/mlir/lite/stablehlo:__pkg__",
-        # For StableHLO Quantizer adapter functionalities within TFLite. Testonly.
-        "//tensorflow/compiler/mlir/lite/quantization/stablehlo:__pkg__",
-        "//tensorflow/python:__pkg__",
-    ],
     deps = [
         ":component",
         ":context",

tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD

@@ -1,6 +1,7 @@
 load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
 load(
     "//tensorflow:tensorflow.default.bzl",
+    "get_compatible_with_portable",
     "tf_py_strict_test",
     "tf_python_pybind_extension",
 )
@@ -91,13 +92,56 @@ tf_py_strict_test(
     ],
 )
 
+# This is a header-only target. The purpose of `pywrap_quantization_lib_*` targets is to expose only
+# the symbols that are required by `pywrap_quantization` that translates them to python functions.
+# The only intended use case of this library is by `pywrap_quantization`. Not letting
+# `pywrap_quantization` directly depend on sub-libraries like `static_range_srq` and instead haiving
+# a consolidated impl library `pywrap_quantization_lib_impl` allows the maintainers to avoid
+# declaring multiple impl libraries to `libtensorflow_cc` and `lib_pywrap_tensorflow_internal`,
+# which is required to avoid ODR violations.
+cc_library(
+    name = "pywrap_quantization_lib_header_only",
+    srcs = [],
+    hdrs = ["pywrap_quantization_lib.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:private"],  # ONLY for `pywrap_quantization`.
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+# See the comments for `pywrap_quantization_lib_header_only`.
+cc_library(
+    name = "pywrap_quantization_lib_impl",
+    srcs = ["pywrap_quantization_lib.cc"],
+    hdrs = ["pywrap_quantization_lib.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = [
+        "//tensorflow:__pkg__",  # For libtensorflow_cc.so.
+        "//tensorflow/python:__pkg__",  # For lib_pywrap_tensorflow_internal.so.
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:config",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
 tf_python_pybind_extension(
     name = "pywrap_quantization",
     srcs = ["pywrap_quantization.cc"],
     pytype_srcs = ["pywrap_quantization.pyi"],
+    # Each dependency MUST be either header-only or exclusive.
     deps = [
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:config_impl",
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq",
+        ":pywrap_quantization_lib_header_only",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:type_casters",
         "@pybind11",
         "@pybind11_abseil//pybind11_abseil:absl_casters",

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization.cc

@@ -20,16 +20,15 @@ limitations under the License.
 #include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil   // IWYU pragma: keep
 #include "pybind11_abseil/import_status_module.h"  // from @pybind11_abseil
 #include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil  // IWYU pragma: keep
-#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
-#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/static_range_ptq.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/python/type_casters.h"  // IWYU pragma: keep
 
 namespace py = pybind11;
 
 namespace {
 
-using ::mlir::quant::stablehlo::QuantizeStaticRangePtq;
-using ::stablehlo::quantization::PopulateDefaults;
+using ::stablehlo::quantization::pywrap::PywrapPopulateDefaults;
+using ::stablehlo::quantization::pywrap::PywrapQuantizeStaticRangePtq;
 
 }  // namespace
 
@@ -42,7 +41,7 @@ PYBIND11_MODULE(pywrap_quantization, m) {
   // If the function signature changes, likely its corresponding .pyi type
   // hinting should also change.
   // LINT.IfChange(static_range_ptq)
-  m.def("static_range_ptq", &QuantizeStaticRangePtq,
+  m.def("static_range_ptq", &PywrapQuantizeStaticRangePtq,
         R"pbdoc(
         Runs static-range post-training quantization (PTQ) on a SavedModel at
         `src_saved_model_path` and saves the resulting model to
@@ -68,7 +67,7 @@ PYBIND11_MODULE(pywrap_quantization, m) {
   // If the function signature changes, likely its corresponding .pyi type
   // hinting should also change.
   // LINT.IfChange(populate_default_configs)
-  m.def("populate_default_configs", &PopulateDefaults,
+  m.def("populate_default_configs", &PywrapPopulateDefaults,
         R"pbdoc(
         Populates `QuantizationConfig` with default values.
 

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.cc

@@ -0,0 +1,55 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/static_range_ptq.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+
+namespace stablehlo::quantization::pywrap {
+
+using ::mlir::quant::stablehlo::QuantizeStaticRangePtq;
+using ::tensorflow::SignatureDef;
+using ::tensorflow::quantization::PyFunctionLibrary;
+
+// Note for maintainers: the definitions should ONLY mirror existing functions
+// defined in different targets. Do not include any extra business logic that
+// causes divergence from the semantics of mirrored functions.
+
+absl::Status PywrapQuantizeStaticRangePtq(
+    absl::string_view src_saved_model_path,
+    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
+    const std::vector<std::string>& signature_keys,
+    const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
+    const absl::flat_hash_map<std::string, std::string>& function_aliases,
+    const PyFunctionLibrary& py_function_library) {
+  return QuantizeStaticRangePtq(src_saved_model_path, dst_saved_model_path,
+                                config, signature_keys, signature_def_map,
+                                function_aliases, py_function_library);
+}
+
+QuantizationConfig PywrapPopulateDefaults(
+    const QuantizationConfig& user_provided_config) {
+  return PopulateDefaults(user_provided_config);
+}
+
+}  // namespace stablehlo::quantization::pywrap

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.h

@@ -0,0 +1,51 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PYTHON_PYWRAP_QUANTIZATION_LIB_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PYTHON_PYWRAP_QUANTIZATION_LIB_H_
+
+// Contains mirror functions from StableHLO Quantizer to be exposed to python
+// via `pywrap_quantization`.
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+
+namespace stablehlo::quantization::pywrap {
+
+// Function used by the pywrap_quantization module to mirror
+// `::mlir::quant::stablehlo::QuantizeStaticRangePtq`.
+absl::Status PywrapQuantizeStaticRangePtq(
+    absl::string_view src_saved_model_path,
+    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
+    const std::vector<std::string>& signature_keys,
+    const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
+        signature_def_map,
+    const absl::flat_hash_map<std::string, std::string>& function_aliases,
+    const tensorflow::quantization::PyFunctionLibrary& py_function_library);
+
+// Function used by the pywrap_quantization module to mirror
+// `::stablehlo::quantization::PopulateDefaults`.
+QuantizationConfig PywrapPopulateDefaults(
+    const QuantizationConfig& user_provided_config);
+
+}  // namespace stablehlo::quantization::pywrap
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PYTHON_PYWRAP_QUANTIZATION_LIB_H_

tensorflow/python/BUILD

@@ -765,7 +765,7 @@ pywrap_tensorflow_macro(
         "//tensorflow/cc/saved_model:fingerprinting_impl",
         "//tensorflow/cc/saved_model:loader_lite_impl",
         "//tensorflow/cc/saved_model:metrics_impl",
-        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq_impl",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/python:pywrap_quantization_lib_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:quantize_model_cc_impl",
         "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
         "//tensorflow/compiler/tf2tensorrt:op_converter_registry_impl",

tensorflow/tf_exported_symbols.lds

@@ -15,3 +15,4 @@
 *tsl*
 *lite*
 *TFL*
+*quantization*

tensorflow/tf_version_script.lds

@@ -17,6 +17,7 @@ tensorflow {
     *tsl*;
     *lite*;
     *TFL*;
+    *quantization*;
   local:
     *;
 };