chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/lite/sparsity/BUILD

@@ -0,0 +1,65 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        ":friends",
+        "//tensorflow:__pkg__",
+    ],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        "//learning/brain/experimental/mlir/...",
+        "//tensorflow/compiler/mlir/lite/...",
+        "//tensorflow/lite/...",
+    ],
+)
+
+cc_library(
+    name = "sparsify_model",
+    srcs = [
+        "sparsify_model.cc",
+    ],
+    hdrs = [
+        "sparsify_model.h",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/lite:flatbuffer_translate_lib",
+        "//tensorflow/compiler/mlir/lite:pass_registry_utils",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite_d2s",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs",
+        "//tensorflow/compiler/mlir/lite/tools/optimize:reduced_precision_metadata",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "sparsify_model_test",
+    srcs = ["sparsify_model_test.cc"],
+    data = [
+        ":testdata/sparse_tensor.bin",
+    ],
+    deps = [
+        ":sparsify_model",
+        "//tensorflow/compiler/mlir/lite/core:absl_error_model_builder",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs",
+        "//tensorflow/compiler/mlir/lite/tools/optimize:reduced_precision_metadata",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)

tensorflow/compiler/mlir/lite/sparsity/sparsify_model.cc

@@ -0,0 +1,110 @@
+/* Copyright 2020 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/lite/sparsity/sparsify_model.h"
+
+#include <cstdint>
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.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/lite/flatbuffer_export.h"
+#include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/compiler/mlir/lite/tools/optimize/reduced_precision_metadata.h"
+#include "tensorflow/compiler/mlir/lite/transforms/dense_to_sparse_pass.h"
+#include "tensorflow/compiler/mlir/lite/transforms/pass_registry_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace mlir {
+namespace lite {
+
+absl::Status SparsifyModel(const tflite::ModelT& input_model,
+                           flatbuffers::FlatBufferBuilder* builder) {
+  MLIRContext context;
+  StatusScopedDiagnosticHandler statusHandler(&context,
+                                              /*propagate=*/true);
+
+  // Import input_model to a MLIR module
+  flatbuffers::FlatBufferBuilder input_builder;
+  flatbuffers::Offset<tflite::Model> input_model_location =
+      tflite::Model::Pack(input_builder, &input_model);
+  tflite::FinishModelBuffer(input_builder, input_model_location);
+
+  std::string serialized_model(
+      reinterpret_cast<const char*>(input_builder.GetBufferPointer()),
+      input_builder.GetSize());
+
+  OwningOpRef<mlir::ModuleOp> module = tflite::FlatBufferToMlir(
+      serialized_model, &context, UnknownLoc::get(&context));
+  if (!module) {
+    LOG(ERROR) << "Couldn't import flatbuffer to MLIR.";
+    return absl::InternalError("Couldn't import flatbuffer to MLIR.");
+  }
+
+  PassManager pm((*module)->getName(), OpPassManager::Nesting::Implicit);
+  pm.addPass(TFL::Create<TFL::DenseToSparsePass>());
+
+  if (failed(pm.run(module.get()))) {
+    LOG(ERROR) << "Failed to sparsify: "
+               << statusHandler.ConsumeStatus().message();
+    return absl::InternalError(absl::StrCat(
+        "Failed to sparsify: ", statusHandler.ConsumeStatus().message()));
+  }
+
+  // Export the results to the builder
+  std::string result;
+  tflite::FlatbufferExportOptions options;
+  options.converter_flags.set_force_select_tf_ops(false);
+  options.converter_flags.set_enable_select_tf_ops(true);
+  options.converter_flags.set_allow_custom_ops(true);
+
+  // Copy metadata for Reduced Precision Support from input model if it exists
+  for (const auto& metadata : input_model.metadata) {
+    if (metadata->name != tflite::optimize::kTfLiteReducedPrecisionKey) {
+      continue;
+    }
+
+    const auto& data = input_model.buffers[metadata->buffer]->data;
+    options.metadata[metadata->name] = std::string(data.begin(), data.end());
+    break;
+  }
+
+  if (!tflite::MlirToFlatBufferTranslateFunction(module.get(), options,
+                                                 &result)) {
+    LOG(ERROR) << "Failed to export MLIR to flatbuffer.";
+    return absl::InternalError("Failed to export MLIR to flatbuffer.");
+  }
+  builder->PushFlatBuffer(reinterpret_cast<const uint8_t*>(result.data()),
+                          result.size());
+
+  return absl::OkStatus();
+}
+
+}  // namespace lite
+}  // namespace mlir

tensorflow/compiler/mlir/lite/sparsity/sparsify_model.h

@@ -0,0 +1,31 @@
+/* Copyright 2020 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_LITE_SPARSITY_SPARSIFY_MODEL_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_SPARSITY_SPARSIFY_MODEL_H_
+
+#include "absl/status/status.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+
+namespace mlir {
+namespace lite {
+
+// Sparsify the `input_model` and write the result to a flatbuffer `builder`.
+absl::Status SparsifyModel(const tflite::ModelT& input_model,
+                           flatbuffers::FlatBufferBuilder* builder);
+}  // namespace lite
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_SPARSITY_SPARSIFY_MODEL_H_

tensorflow/compiler/mlir/lite/sparsity/sparsify_model_test.cc

@@ -0,0 +1,81 @@
+/* Copyright 2022 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/lite/sparsity/sparsify_model.h"
+
+#include <stdint.h>
+
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/compiler/mlir/lite/core/absl_error_model_builder.h"
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/compiler/mlir/lite/tools/optimize/reduced_precision_metadata.h"
+
+namespace mlir {
+namespace lite {
+namespace {
+
+
+TEST(SparsifyModelTest, MetadataIsAddedToOutputModel) {
+  std::string expected_key = tflite::optimize::kTfLiteReducedPrecisionKey;
+  std::string expected_value = "test_data";
+
+  // Load input model
+  auto input_fbm = mlir::TFL::FlatBufferModelAbslError::BuildFromFile(
+      "tensorflow/compiler/mlir/lite/sparsity/testdata/"
+      "sparse_tensor.bin");
+  tflite::ModelT input_model;
+  input_fbm->GetModel()->UnPackTo(&input_model);
+
+  // Populate input metadata
+  auto model_metadata_buffer = std::make_unique<tflite::BufferT>();
+  model_metadata_buffer->data =
+      std::vector<uint8_t>(expected_value.begin(), expected_value.end());
+  input_model.buffers.push_back(std::move(model_metadata_buffer));
+  auto metadata_t = std::make_unique<tflite::MetadataT>();
+  metadata_t->name = tflite::optimize::kTfLiteReducedPrecisionKey;
+  metadata_t->buffer = input_model.buffers.size() - 1;
+  input_model.metadata.push_back(std::move(metadata_t));
+
+  // Sparsify and create output model
+  flatbuffers::FlatBufferBuilder output_builder;
+  ASSERT_TRUE(SparsifyModel(input_model, &output_builder).ok());
+  auto output_fbm = mlir::TFL::FlatBufferModelAbslError::BuildFromBuffer(
+      reinterpret_cast<const char*>(output_builder.GetCurrentBufferPointer()),
+      output_builder.GetSize());
+  tflite::ModelT output_model;
+  output_fbm->GetModel()->UnPackTo(&output_model);
+
+  // Extract output metadata
+  std::map<std::string, std::string> output_metadata;
+  for (const auto& metadata : output_model.metadata) {
+    const auto& data = output_model.buffers[metadata->buffer]->data;
+    output_metadata[metadata->name] = std::string(data.begin(), data.end());
+  }
+
+  EXPECT_THAT(output_metadata,
+              testing::Contains(testing::Pair(expected_key, expected_value)));
+}
+
+}  // namespace
+}  // namespace lite
+}  // namespace mlir