chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tf2xla/internal/inference/BUILD

@@ -0,0 +1,62 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//visibility:private",
+    ],
+)
+
+gentbl_cc_library(
+    name = "inference_passes_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"inference_passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=TF2XLA",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "inference_passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "inference_metrics_pass",
+    srcs = ["inference_metrics_pass.cc"],
+    hdrs = ["inference_passes.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":inference_passes_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Dialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "inference_metrics_pass_test",
+    srcs = ["inference_metrics_pass_test.cc"],
+    deps = [
+        "inference_metrics_pass",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
+        "//tensorflow/core/lib/monitoring:cell_reader",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)

tensorflow/compiler/mlir/tf2xla/internal/inference/inference_metrics_pass.cc

@@ -0,0 +1,73 @@
+/* Copyright 2023 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 <memory>
+#include <string>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#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 "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+
+namespace mlir {
+namespace tf2xla {
+namespace internal {
+
+auto* has_tpu_partitioned_call_streamz =
+    tensorflow::monitoring::Counter<1>::New(
+        "/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call",
+        "Whether the model has TPUPartitionedCallOp.",
+        "has_tpu_partitioned_call");
+
+namespace {
+
+#define GEN_PASS_DEF_INFERENCEMETRICSPASS
+#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"
+
+class InferenceMetricsPass
+    : public impl::InferenceMetricsPassBase<InferenceMetricsPass> {
+ public:
+  void runOnOperation() override;
+};
+
+void InferenceMetricsPass::runOnOperation() {
+  bool has_tpu_partitioned_call = false;
+  ModuleOp module = getOperation();
+
+  for (auto func_op : module.getOps<func::FuncOp>()) {
+    func_op->walk(
+        [&](TF::TPUPartitionedCallOp op) { has_tpu_partitioned_call = true; });
+
+    if (has_tpu_partitioned_call) break;
+  }
+
+  std::string has_tpu_partitioned_call_str =
+      has_tpu_partitioned_call ? "true" : "false";
+  has_tpu_partitioned_call_streamz->GetCell(has_tpu_partitioned_call_str)
+      ->IncrementBy(1);
+}
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateInferenceMetricsPass() {
+  return std::make_unique<InferenceMetricsPass>();
+}
+
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace mlir

tensorflow/compiler/mlir/tf2xla/internal/inference/inference_metrics_pass_test.cc

@@ -0,0 +1,105 @@
+/* Copyright 2023 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 <cstdint>
+#include <memory>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h"
+#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/lib/monitoring/cell_reader.h"
+
+namespace mlir {
+namespace tf2xla {
+namespace internal {
+namespace {
+
+using ::mlir::MLIRContext;
+using ::mlir::ModuleOp;
+using ::mlir::OwningOpRef;
+using ::mlir::hlo::test::GetMlirModuleFromString;
+using ::tensorflow::monitoring::testing::CellReader;
+
+static constexpr char kHasTpuPartitionedCallStreamzName[] =
+    "/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call";
+
+class InferenceMetricsPassTest : public ::testing::Test {
+ protected:
+  void CreateModule(const char* module_string) {
+    TF_ASSERT_OK_AND_ASSIGN(module_,
+                            GetMlirModuleFromString(module_string, &context_));
+
+    pm_ = std::make_unique<mlir::PassManager>(&context_);
+    pm_->addPass(CreateInferenceMetricsPass());
+  }
+  mlir::LogicalResult Run() { return pm_->run(module_.get()); }
+
+ private:
+  MLIRContext context_;
+  OwningOpRef<ModuleOp> module_;
+  std::unique_ptr<mlir::PassManager> pm_;
+};
+
+TEST_F(InferenceMetricsPassTest, RecordsTrueForTPUPartitionedCallOp) {
+  static constexpr char kMlirModuleStr[] = R"(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
+      func.return %arg0 : tensor<?xi32>
+    }
+
+    func.func @main(%arg0: tensor<20xi32>, %arg1: tensor<?xi32>) -> tensor<*xi32> {
+      %2 = "tf.TPUPartitionedCall"(%arg0, %arg1) {f = @tpu_partitioned_call_func} : (tensor<20xi32>, tensor<?xi32>) -> tensor<*xi32>
+      func.return %2 : tensor<*xi32>
+    }
+  })";
+
+  CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);
+  CreateModule(kMlirModuleStr);
+
+  auto result = Run();
+
+  EXPECT_TRUE(result.succeeded());
+
+  EXPECT_EQ(error.Delta("true"), 1);
+  EXPECT_EQ(error.Delta("false"), 0);
+}
+
+TEST_F(InferenceMetricsPassTest, RecordsFalseForNonTPUPartitionedCallOp) {
+  static constexpr char kMlirModuleStr[] = R"(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @main() -> tensor<1xi32> {
+      %0 = "tf.BadValue"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
+      func.return %0 : tensor<1xi32>
+    }
+  })";
+  CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);
+  CreateModule(kMlirModuleStr);
+
+  auto result = Run();
+
+  EXPECT_TRUE(result.succeeded());
+  EXPECT_EQ(error.Delta("false"), 1);
+  EXPECT_EQ(error.Delta("true"), 0);
+}
+
+}  // namespace
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace mlir

tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h

@@ -0,0 +1,39 @@
+/* Copyright 2023 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_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace tf2xla {
+namespace internal {
+
+std::unique_ptr<OperationPass<ModuleOp>> CreateInferenceMetricsPass();
+
+#define GEN_PASS_REGISTRATION
+#define GEN_PASS_DECL_INFERENCEMETRICSPASS
+#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"
+
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_

tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.td

@@ -0,0 +1,25 @@
+/* Copyright 2023 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 "mlir/Pass/PassBase.td"
+
+// Internal TF2XLA Bridge Passes used during interence only.
+def InferenceMetricsPass : Pass<"tf2xla-inference-metrics-pass", "ModuleOp"> {
+  let summary = "A pass to go over a Module and collect various metrics";
+  let description = [{
+    Collects metrics about a Module during inference, such as which ops are used.
+  }];
+  let constructor = "internal::CreateInferenceMetricsPass()";
+}