chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/BUILD

@@ -0,0 +1,48 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["@stablehlo//:license"],
+    default_visibility = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "save_report",
+    srcs = ["save_report.cc"],
+    hdrs = ["save_report.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:report",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "save_report_test",
+    srcs = ["save_report_test.cc"],
+    deps = [
+        ":save_report",
+        "//tensorflow/compiler/mlir/quantization/common:test_base",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:passes",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:io",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tsl//tsl/platform:protobuf",
+    ],
+)

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.cc

@@ -0,0 +1,97 @@
+/* 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/instrumentations/save_report.h"
+
+#include <optional>
+#include <string>
+
+#include "absl/base/nullability.h"
+#include "absl/log/log.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/report.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+// Converts `std::optional<absl::string_view>` to `std::optional<std::string>`.
+// A `std::nullopt` is returned when `view` is `std::nullopt`.
+std::optional<std::string> OptionalStringViewToOptionalString(
+    std::optional<absl::string_view> view) {
+  if (view == std::nullopt) return std::nullopt;
+  return std::make_optional<std::string>(*view);
+}
+
+using quant::stablehlo::QuantizationReport;
+
+// Whether the pass is `QuantizeCompositeFunctionPass`.
+bool IsQuantizeCompositeFunctionPass(Pass* absl_nullable pass,
+                                     Operation* absl_nullable op) {
+  // It is known that `op` is `ModuleOp` when `pass` is
+  // `QuantizeCompositeFunctionPass`, but the check is still performed to be
+  // defensive.
+  return pass != nullptr &&
+         pass->getArgument() == "tf-stablehlo-quantize-composite-functions" &&
+         isa_and_nonnull<ModuleOp>(op);
+}
+
+// Report is saved only when:
+// * After running `QuantizeCompositeFunctionPass`.
+// * The pass is run on `ModuleOp`.
+// * `file_path` is not `nullopt`.
+bool ShouldSaveReport(Pass* absl_nullable pass, Operation* absl_nullable op,
+                      const std::optional<std::string>& file_path) {
+  return file_path != std::nullopt && IsQuantizeCompositeFunctionPass(pass, op);
+}
+
+void SaveReport(const QuantizationReport& report,
+                const absl::string_view file_path) {
+  if (const absl::Status save_status = report.Save(file_path);
+      save_status.ok()) {
+    LOG(INFO) << "Successfully saved quantization report to: " << file_path;
+  } else {
+    LOG(ERROR) << "Failed to save quantization report to: " << file_path
+               << " with status: " << save_status;
+  }
+}
+
+}  // namespace
+
+SaveQuantizationReportInstrumentation::SaveQuantizationReportInstrumentation(
+    std::optional<absl::string_view> file_path)
+    : file_path_(OptionalStringViewToOptionalString(file_path)) {}
+
+void SaveQuantizationReportInstrumentation::runAfterPass(Pass* pass,
+                                                         Operation* op) {
+  // Only run after `QuantizeCompositeFunctionPass`.
+  if (!IsQuantizeCompositeFunctionPass(pass, op)) return;
+
+  auto module_op = cast<ModuleOp>(op);
+  const QuantizationReport report(module_op);
+
+  // Print a human-readable report to stdout regardless of whether the report
+  // is saved to file.
+  report.Print();
+
+  // Exit early if the report should not be saved to file.
+  if (!ShouldSaveReport(pass, op, file_path_)) return;
+
+  SaveReport(report, *file_path_);
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.h

@@ -0,0 +1,52 @@
+/* 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_INSTRUMENTATIONS_SAVE_REPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_INSTRUMENTATIONS_SAVE_REPORT_H_
+
+#include <optional>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// A `PassInstrumentation` that saves quantization report to file after
+// `QuantizeCompositeFunctionsPass` is run. It inspects the `ModuleOp` after
+// quantization and analyzes the quantizable units and quantization methods
+// used. The report file will be saved at the `file_path`. The report file
+// contains textproto of `QuantizationResults`. `file_path`'s base directories
+// should exist (this pass instrumentation will not `mkdir` them).
+//
+// See `QuantizationReport` for further details on the quantization report.
+class SaveQuantizationReportInstrumentation : public PassInstrumentation {
+ public:
+  // `file_path` is the path to save the report file. The report file is in
+  // textproto format so a `.txtpb` extension is preferred but it doesn't result
+  // in error if other extension is used. This instrumentation will not be run
+  // if `file_path` is a `nullopt`.
+  explicit SaveQuantizationReportInstrumentation(
+      std::optional<absl::string_view> file_path);
+
+  void runAfterPass(Pass* pass, Operation* op) override;
+
+ private:
+  std::optional<std::string> file_path_;  // Path to file to save the report.
+};
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_INSTRUMENTATIONS_SAVE_REPORT_H_

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report_test.cc

@@ -0,0 +1,185 @@
+/* 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/instrumentations/save_report.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/io.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::stablehlo::quantization::QuantizationResults;
+using ::stablehlo::quantization::io::ReadFileToString;
+using ::testing::SizeIs;
+using ::testing::StrEq;
+using ::tsl::protobuf::TextFormat;
+
+using SaveQuantizationReportInstrumentationTest = QuantizationTestBase;
+
+TEST_F(SaveQuantizationReportInstrumentationTest, SaveReport) {
+  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
+    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
+      %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+      %0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
+      %2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+
+    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleWithCompositeDotGeneral);
+  ASSERT_TRUE(module_op);
+
+  // Create a pass manager with `SaveQuantizationReportInstrumentation` and
+  // `QuantizeCompositeFunctionsPass`. Run the passes against `module_op`.
+  PassManager pm(ctx_.get());
+
+  QuantizeCompositeFunctionsPassOptions options;
+  pm.addPass(createQuantizeCompositeFunctionsPass(options));
+
+  const std::string report_file_path =
+      absl::StrCat(testing::TempDir(), "/save_report.txtpb");
+  pm.addInstrumentation(std::make_unique<SaveQuantizationReportInstrumentation>(
+      report_file_path));
+
+  const LogicalResult run_result = pm.run(*module_op);
+  ASSERT_TRUE(succeeded(run_result));
+
+  // Check that the report file contains `QuantizationResults` textproto,
+  // reflecting the quantization results, in this case the
+  // `composite_dot_general_fn` with quantized with `static_range_ptq` method.
+  const absl::StatusOr<std::string> file_data =
+      ReadFileToString(report_file_path);
+  ASSERT_THAT(file_data, absl_testing::IsOk());
+
+  /*
+  results {
+    quantizable_unit {
+      name: "composite_dot_general_fn"
+    }
+    method { static_range_ptq { } }
+  }
+  */
+  QuantizationResults results{};
+  ASSERT_TRUE(TextFormat::ParseFromString(*file_data, &results));
+  ASSERT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).quantizable_unit().name(),
+              StrEq("composite_dot_general_fn"));
+  EXPECT_TRUE(results.results(0).method().has_static_range_ptq());
+}
+
+TEST_F(SaveQuantizationReportInstrumentationTest,
+       ReportNotSavedWhenNoQuantizeCompositeFunctionsPass) {
+  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
+    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
+      %cst = "stablehlo.constant"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+      %0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
+      %2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+
+    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleWithCompositeDotGeneral);
+  ASSERT_TRUE(module_op);
+
+  // Create a pass manager with `SaveQuantizationReportInstrumentation` a pass
+  // that is not `QuantizeCompositeFunctionsPass`. Run the passes against
+  // `module_op`.
+  PassManager pm(ctx_.get());
+
+  pm.addPass(createPrepareQuantizePass());
+
+  const std::string report_file_path = absl::StrCat(
+      testing::TempDir(),
+      "/report_not_saved_no_quantize_composite_functions_pass.txtpb");
+  pm.addInstrumentation(std::make_unique<SaveQuantizationReportInstrumentation>(
+      report_file_path));
+
+  const LogicalResult run_result = pm.run(*module_op);
+  ASSERT_TRUE(succeeded(run_result));
+
+  // The report file is not created because `QuantizeCompositeFunctionsPass` was
+  // not run.
+  EXPECT_THAT(ReadFileToString(report_file_path),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST_F(SaveQuantizationReportInstrumentationTest,
+       ReportNotSavedWhenReportFilePathIsNullopt) {
+  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
+    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
+      %cst = "stablehlo.constant"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+      %0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
+      %2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+
+    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleWithCompositeDotGeneral);
+  ASSERT_TRUE(module_op);
+
+  PassManager pm(ctx_.get());
+
+  QuantizeCompositeFunctionsPassOptions options;
+  pm.addPass(createQuantizeCompositeFunctionsPass(options));
+  pm.addInstrumentation(std::make_unique<SaveQuantizationReportInstrumentation>(
+      /*file_path=*/std::nullopt));
+
+  // The report file is not created and `SaveQuantizationReportInstrumentation`
+  // is not run, but the passes still run without errors.
+  const LogicalResult run_result = pm.run(*module_op);
+  ASSERT_TRUE(succeeded(run_result));
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo