chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/tests/ir/BUILD

@@ -0,0 +1,47 @@
+load("//tensorflow:tensorflow.bzl", "if_oss", "tf_cc_test")
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    exclude = ["testdata/**"],
+    features = if_oss(["--path=org_tensorflow/tensorflow/compiler/mlir/tfrt"]),
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)
+
+tf_cc_test(
+    name = "tfrt_fallback_util_test",
+    srcs = ["tfrt_fallback_util_test.cc"],
+    data = [
+        "testdata/test.mlir",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_sync_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_util",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:resource_loader",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@tf_runtime//:init_tfrt_dialects",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/ir/fallback_opt.mlir

@@ -0,0 +1,98 @@
+// Copyright 2026 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s | tf-tfrt-opt | FileCheck %s --dump-input=fail
+
+// CHECK-LABEL: func @const_tensor_proto
+func.func @const_tensor_proto() -> !tfrt_fallback.tf_tensor {
+  // CHECK: tfrt_fallback_async.const_tensor_proto "fake serialized proto"
+  %0 = tfrt_fallback_async.const_tensor_proto "fake serialized proto"
+  tfrt.return %0 : !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @const_dense_tensor
+func.func @const_dense_tensor() -> !tfrt_fallback.tf_tensor {
+  // CHECK: tfrt_fallback_async.const_dense_tensor
+  %0 = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
+  tfrt.return %0 : !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @const_string_tensor
+func.func @const_string_tensor() -> !tfrt_fallback.tf_tensor {
+  // CHECK: tfrt_fallback_async.const_string_tensor
+  %0 = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["const", "string"], _tfrt_cost = 1 : i64}
+  tfrt.return %0 : !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @convert
+func.func @convert() -> !corert.tensorhandle {
+  %0 = corert.const_dense_tensor dense<0.0> : tensor<f32>
+  // CHECK: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor
+  %1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %0 {_tfrt_cost = 1 : i64, device = "cpu"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
+  // CHECK: tfrt_fallback_async.executeop key(0) cost(100) device("cpu") "tf.Relu"(%{{.*}}) {T = f32} : 1
+  %2 = tfrt_fallback_async.executeop key(0) cost(100) device("cpu") "tf.Relu"(%1) {T = f32} : 1
+  // CHECK: tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle
+  %3 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %2 {_tfrt_cost = 1 : i64, device="cpu"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
+  tfrt.return %3 : !corert.tensorhandle
+}
+
+// CHECK-LABEL: func @predicate
+func.func @predicate() -> i1 {
+  %0 = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
+  // CHECK: tfrt_fallback_async.predicate
+  %1 = tfrt_fallback_async.predicate %0 {_tfrt_cost = 1 : i64, device = "cpu"}
+  tfrt.return %1 : i1
+}
+
+// CHECK-LABEL: func @createop
+func.func @createop(%in_ch: !tfrt.chain) -> !tfrt.chain {
+
+  // CHECK: [[ch:%.*]] = tfrt_fallback_async.createop(%{{.*}}) key(100) device("cpu") "tf.AddV2"() {T = i32} num_args(2)
+  %out_ch = tfrt_fallback_async.createop(%in_ch) key(100) device("cpu") "tf.AddV2"() {T = i32} num_args(2)
+
+  // CHECK: tfrt.return [[ch]]
+  tfrt.return %out_ch: !tfrt.chain
+}
+
+// CHECK-LABEL: func @fallback_resource
+func.func @fallback_resource(%ch0: !tfrt.chain) -> !tfrt.chain {
+
+  %ra = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
+  %rb = tfrt_fallback_async.const_dense_tensor dense<0.5> : tensor<f32> {_tfrt_cost = 1 : i64}
+
+  // CHECK: tfrt_fallback_async.set_resource {{%.*}}, {{%.*}} {device = "cpu", index = 0 : i64}
+  // CHECK: tfrt_fallback_async.set_resource {{%.*}}, {{%.*}} {device = "cpu", index = 1 : i64}
+  // CHECK: tfrt_fallback_async.get_resource {{%.*}} {_tfrt_cost = 1 : i64, device = "cpu", indices = [0, 1]}
+  %ch1 = tfrt_fallback_async.set_resource %ch0, %ra {device = "cpu", index = 0 : i64}
+  %ch2 = tfrt_fallback_async.set_resource %ch1, %rb {device = "cpu", index = 1 : i64}
+  %ch3, %a, %b = tfrt_fallback_async.get_resource %ch2 {_tfrt_cost = 1 : i64, device = "cpu", indices = [0 : i64, 1 : i64]} : (!tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
+
+  tfrt.return %ch3: !tfrt.chain
+}
+
+// CHECK-LABEL: func @copy_if_small
+func.func @copy_if_small(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
+  // CHECK: tfrt_fallback_async.copy_if_small {{%.*}} {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
+  %0:2 = tfrt_fallback_async.copy_if_small %arg {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
+  tfrt.return %0#0, %0#1 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @custom_allocator
+func.func @custom_allocator(%ch: !tfrt.chain, %arg: !tfrt_fallback.tf_tensor, %allocator: !tfrt_fallback.tf_allocator) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
+  // CHECK: tfrt_fallback_async.executeop.allocator(%{{.*}}) key(200) cost(100) device("cpu") "tf.Cast"(%{{.*}}) {Truncate = false, T = f32} : 1
+  %0 = tfrt_fallback_async.executeop.allocator(%allocator) key(200) cost(100) device("cpu") "tf.Cast"(%arg) {Truncate = false, T = f32} : 1
+  // CHECK: tfrt_fallback_async.executeop.seq.allocator(%{{.*}}, %{{.*}}) key(201) cost(100) device("cpu") "tf.Cast"(%{{.*}}) {Truncate = false, T = i32} : 1
+  %out_ch, %1 = tfrt_fallback_async.executeop.seq.allocator(%ch, %allocator) key(201) cost(100) device("cpu") "tf.Cast"(%0) {Truncate = false, T = i32} : 1
+  tfrt.return %out_ch, %1 : !tfrt.chain, !tfrt_fallback.tf_tensor
+}

tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir

@@ -0,0 +1,22 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @test(%ch: !tfrt.chain, %arg0: !corert.tensorhandle, %arg1_th: !corert.tensorhandle) {
+  %cpu = corert.get_op_handler %ch "cpu"
+  %0 = corert.executeop(%cpu) "tf.Relu"(%arg0) { T = f32 } : 1
+  %arg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg1_th {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
+  %1 = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "tf.Relu"(%arg1) { T = f32 } : 1
+  tfrt.return
+}

tensorflow/compiler/mlir/tfrt/tests/ir/tfrt_fallback_util_test.cc

@@ -0,0 +1,63 @@
+/* Copyright 2021 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/tfrt/ir/tfrt_fallback_util.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/test.h"
+#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_async {
+namespace {
+
+TEST(SavedModelTest, MapFallbackArgs) {
+  std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
+      "tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir");
+
+  mlir::DialectRegistry registry;
+  RegisterTFRTDialects(registry);
+  registry.insert<tfrt::fallback_async::FallbackAsyncDialect>();
+  registry.insert<tfrt::fallback_sync::FallbackSyncDialect>();
+
+  mlir::MLIRContext context(registry);
+  auto module =
+      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
+  ASSERT_TRUE(module);
+
+  std::vector<std::pair<std::string, int>> func_and_index;
+  ForEachArgConsumedByFallback(
+      module.get(),
+      [&func_and_index](llvm::StringRef func_name, int arg_index) {
+        func_and_index.push_back({func_name.str(), arg_index});
+      });
+
+  ASSERT_EQ(func_and_index.size(), 1);
+  EXPECT_EQ(func_and_index[0].first, "test");
+  EXPECT_EQ(func_and_index[0].second, 2);
+}
+
+}  // namespace
+}  // namespace fallback_async
+}  // namespace tfrt