wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/mlir/tfrt/tests/BUILD
+29
View File
@@ -0,0 +1,29 @@
load("//tensorflow:tensorflow.bzl", "if_oss")
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",
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:tf-mlir-translate",
"//tensorflow/compiler/mlir:tf-opt",
"//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
"@llvm-project//mlir:run_lit.sh",
],
)
tensorflow/compiler/mlir/tfrt/tests/analysis/BUILD
+50
View File
@@ -0,0 +1,50 @@
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:tf-mlir-translate",
"//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 = "update_op_cost_in_tfrt_mlir_test",
srcs = ["update_op_cost_in_tfrt_mlir_test.cc"],
data = [
"testdata/test.mlir",
],
deps = [
"//tensorflow/compiler/mlir/tfrt:transforms/update_op_cost_in_tfrt_mlir",
"//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
"//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_sync_opdefs",
"//tensorflow/core:test",
"//tensorflow/core/platform:resource_loader",
"//tensorflow/core/tfrt/fallback:cost_recorder",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@tf_runtime//:init_tfrt_dialects",
],
)
tensorflow/compiler/mlir/tfrt/tests/analysis/cost_analysis.mlir
+165
View File
@@ -0,0 +1,165 @@
// 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 -tfrt-test-cost-analysis -verify-diagnostics %s | FileCheck %s
// CHECK-LABEL: test_cheap_ops_0
func.func @test_cheap_ops_0(%arg: tensor<?x!tf_type.string>) -> (tensor<?x8xf32>) {
// expected-remark@+1 {{Cost: 1}}
%0 = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64>
// expected-remark@+1 {{Cost: 1}}
%1 = "tf.Const"() {value = dense<"has_login_page_feature"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%2 = "tf.Const"() {value = dense<"num_terms_inside_postform"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%3 = "tf.Const"() {value = dense<"num_terms_outside_postform"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%4 = "tf.Const"() {value = dense<"num_terms_outside_postform_without_bp"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%5 = "tf.Const"() {value = dense<"password_not_in_bp_area"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%6 = "tf.Const"() {value = dense<"query_params_contains_url"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%7 = "tf.Const"() {value = dense<"title_with_login_phase"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%8 = "tf.Const"() {value = dense<"url_contains_login_terms"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%9 = "tf.Const"() {value = dense<> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string>
// expected-remark@+1 {{Cost: 1}}
%10 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
// expected-remark@+1 {{Cost: 1}}
%11 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
// expected-remark@+1 {{Cost: 19}}
%dense_values:8 = "tf.ParseExample"(%arg, %9, %1, %2, %3, %4, %5, %6, %7, %8, %0, %0, %0, %0, %0, %0, %0, %0) {dense_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>], device = "/job:localhost/replica:0/task:0/device:CPU:0", operandSegmentSizes = array<i32: 1, 1, 0, 8, 8>, resultSegmentSizes = array<i32: 0, 0, 0, 8>} : (tensor<?x!tf_type.string>, tensor<0x!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>) -> (tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>)
// expected-remark@+1 {{Cost: 2}}
%28 = "tf.Cast"(%dense_values#0) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%29 = "tf.Cast"(%dense_values#1) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%30 = "tf.Cast"(%dense_values#2) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%31 = "tf.Cast"(%dense_values#3) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%32 = "tf.Cast"(%dense_values#4) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%33 = "tf.Cast"(%dense_values#5) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%34 = "tf.Cast"(%dense_values#6) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 2}}
%35 = "tf.Cast"(%dense_values#7) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
// expected-remark@+1 {{Cost: 1}}
%36 = "tf.ExpandDims"(%28, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%37 = "tf.ExpandDims"(%29, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%38 = "tf.ExpandDims"(%30, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%39 = "tf.ExpandDims"(%31, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%40 = "tf.ExpandDims"(%32, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%41 = "tf.ExpandDims"(%33, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%42 = "tf.ExpandDims"(%34, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 1}}
%43 = "tf.ExpandDims"(%35, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
// expected-remark@+1 {{Cost: 10}}
%44 = "tf.ConcatV2"(%36, %37, %38, %39, %40, %41, %42, %43, %10) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<i32>) -> tensor<?x8xf32>
// expected-remark@+1 {{Cost: 1}}
func.return %44 : tensor<?x8xf32>
}
// CHECK-LABEL: test_cheap_ops_1
func.func @test_cheap_ops_1(%arg: tensor<?x8x?x?xf32>) -> (tensor<4xi32>, tensor<?x8x?x?xf32>) {
// expected-remark@+1 {{Cost: 1}}
%0 = "tf.Const"() {value = dense<8> : tensor<i32>} : () -> tensor<i32>
// expected-remark@+1 {{Cost: 1}}
%1 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32>
// expected-remark@+1 {{Cost: 1}}
%2 = "tf.Const"() {value = dense<64> : tensor<i32>} : () -> tensor<i32>
// expected-remark@+1 {{Cost: 1}}
%3 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
// expected-remark@+1 {{Cost: 1}}
%4 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
// expected-remark@+1 {{Cost: 1}}
%5 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
// expected-remark@+1 {{Cost: 1}}
%6 = "tf.Const"() {value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
// expected-remark@+1 {{Cost: 9}}
%7 = "tf.Softmax"(%arg) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>) -> tensor<?x8x?x?xf32>
// expected-remark@+1 {{Cost: 1}}
%8 = "tf.Shape"(%7) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>) -> tensor<4xi32>
// expected-remark@+1 {{Cost: 1}}
%9 = "tf.StridedSlice"(%8, %5, %4, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// expected-remark@+1 {{Cost: 1}}
%10 = "tf.StridedSlice"(%8, %3, %6, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// expected-remark@+1 {{Cost: 5}}
%11 = "tf.Pack"(%9, %0, %10, %2) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4xi32>
// expected-remark@+1 {{Cost: 1}}
%12 = "tf.StridedSlice"(%8, %6, %1, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// expected-remark@+1 {{Cost: 5}}
%13 = "tf.Pack"(%9, %0, %10, %12) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4xi32>
// expected-remark@+1 {{Cost: 1}}
%14 = "tf.Reshape"(%7, %13) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>, tensor<4xi32>) -> tensor<?x8x?x?xf32>
// expected-remark@+1 {{Cost: 8}}
func.return %11, %14 : tensor<4xi32>, tensor<?x8x?x?xf32>
}
// CHECK-LABEL: test_expensive_ops
func.func @test_expensive_ops(%arg: tensor<?x512xf32>) -> tensor<?x512xf32> {
// expected-remark@+1 {{Cost: 1}}
%0 = "tf.VarHandleOp"() {allowed_devices = [], container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shared_name = "var"} : () -> tensor<!tf_type.resource<tensor<512x512xf32>>>
// expected-remark@+1 {{Cost: 2}}
%1 = "tf.ReadVariableOp"(%0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource<tensor<512x512xf32>>>) -> tensor<512x512xf32>
// 262657 = 1 + 512 + 512 * 512
// expected-remark@+1 {{Cost: 262657}}
%2 = "tf.MatMul"(%arg, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<?x512xf32>, tensor<512x512xf32>) -> tensor<?x512xf32>
// expected-remark@+1 {{Cost: 512}}
func.return %2 : tensor<?x512xf32>
}
// CHECK-LABEL: test_dynamic_shape
func.func @test_dynamic_shape(%key: tensor<?x!tf_type.string>, %value: tensor<8xi64>) -> tensor<*xi1> {
// expected-remark@+1 {{Cost: 1}}
%default = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i64>} : () -> tensor<i64>
// expected-remark@+1 {{Cost: 1}}
%0 = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
// expected-remark@+1 {{Cost: 1024}}
%1 = "tf.LookupTableFindV2"(%0, %key, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
// 17 = 1 + 8 + 8
// expected-remark@+1 {{Cost: 17}}
%2 = "tf.NotEqual"(%1, %value) {device = "/job:localhost/replica:0/task:0/device:CPU:0", incompatible_shape_error = true} : (tensor<*xi64>, tensor<8xi64>) -> tensor<*xi1>
// expected-remark@+1 {{Cost: 8}}
func.return %2 : tensor<*xi1>
}
// CHECK-LABEL: test_gather
func.func @test_gather(%arg0 : tensor<1x2x20xf32>, %arg1 : tensor<3x5xi32>) -> (tensor<1x3x5x20xf32>){
// expected-remark@+1 {{Cost: 1}}
%0 = "tf.Const"() { value = dense<[1]> : tensor<1xi32> } : () -> tensor<1xi32>
// expected-remark@+1 {{Cost: 300}}
%1 = "tf.GatherV2"(%arg0, %arg1, %0) : (tensor<1x2x20xf32>, tensor<3x5xi32>, tensor<1xi32>) -> tensor<1x3x5x20xf32>
// expected-remark@+1 {{Cost: 40}}
func.return %1 : tensor<1x3x5x20xf32>
}
// CHECK-LABEL: test_sparse_segment_sum
func.func @test_sparse_segment_sum(%indices: tensor<3xi64>, %segment_ids: tensor<3xi64>) -> (tensor<?x28xf32>){
// expected-remark@+1 {{Cost: 1}}
%data = "tf.Const"() { value = dense<0.1> : tensor<476x28xf32> } : () -> tensor<476x28xf32>
// expected-remark@+1 {{Cost: 28}}
%1 = "tf.SparseSegmentSum"(%data, %indices, %segment_ids) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<476x28xf32>, tensor<3xi64>, tensor<3xi64>) -> tensor<?x28xf32>
// expected-remark@+1 {{Cost: 3}}
func.return %1 : tensor<?x28xf32>
}
tensorflow/compiler/mlir/tfrt/tests/analysis/tensor_array_side_effect_analysis.mlir
+27
View File
@@ -0,0 +1,27 @@
// 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 -tfrt-test-tensor-array-effect -verify-diagnostics %s | FileCheck %s
// CHECK-LABEL: @test_tensor_array_effect
// expected-remark@+1 {{HasAtMostTensorArrayEffect: 1}}
func.func @test_tensor_array_effect(%index: tensor<i32>, %size: tensor<i32>, %flow_0: tensor<f32>, %flow_1: tensor<f32>, %handle_0: tensor<2x!tf_type.resource<tensor<?x100xf32>>>, %handle_1: tensor<2x!tf_type.resource<tensor<?x512xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>) {
%cst = "tf.Const"() {value = dense<1.1> : tensor<100x512xf32>} : () -> tensor<100x512xf32>
%one = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%x = "tf.TensorArrayReadV3"(%handle_0, %index, %flow_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<i32>, tensor<f32>) -> tensor<?x100xf32>
%y = "tf.MatMul"(%x, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x100xf32>, tensor<100x512xf32>) -> (tensor<?x512xf32>)
%flow_1_out = "tf.TensorArrayWriteV3"(%handle_1, %index, %y, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<i32>, tensor<?x512xf32>, tensor<f32>) -> tensor<f32>
%next_index = "tf.AddV2"(%index, %one) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %next_index, %size, %flow_0, %flow_1_out, %handle_0, %handle_1 : tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>
}
tensorflow/compiler/mlir/tfrt/tests/analysis/testdata/test.mlir
Vendored
+22
View File
@@ -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/analysis/update_op_cost_in_tfrt_mlir_test.cc
+89
View File
@@ -0,0 +1,89 @@
/* 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/tfrt/transforms/update_op_cost_in_tfrt_mlir.h"
#include <cstdint>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_map.h"
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#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/IR/Operation.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 "tensorflow/core/tfrt/fallback/cost_recorder.h"
#include "tfrt/init_tfrt_dialects.h" // from @tf_runtime
namespace tensorflow {
namespace {
constexpr char kCostAttrName[] = "_tfrt_cost";
constexpr char kOpKeyAttrName[] = "op_key";
absl::flat_hash_map<int64_t, uint64_t> GetOpCostMap(mlir::ModuleOp op) {
absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
op.walk([&](mlir::Operation* op) {
const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
if (!cost_attr) return;
const auto op_key_attr =
op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
if (!op_key_attr) return;
op_cost_map[op_key_attr.getInt()] = cost_attr.getInt();
});
return op_cost_map;
}
TEST(CostUpdateTest, Basic) {
std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/analysis/testdata/test.mlir");
mlir::DialectRegistry registry;
tfrt::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);
// Create a cost recorder with fake cost records.
auto expected_op_cost_map = GetOpCostMap(module.get());
EXPECT_EQ(expected_op_cost_map.size(), 1);
unsigned int seed = 23579;
for (auto& [op_key, cost] : expected_op_cost_map) {
cost = rand_r(&seed) % 1000;
}
tensorflow::tfrt_stub::CostRecorder cost_recorder;
for (const auto& [op_key, cost] : expected_op_cost_map) {
cost_recorder.RecordCost(op_key, cost);
}
// Update the TFRT MLIR with the cost recorder.
tfrt_compiler::UpdateOpCostInTfrtMlir(module.get(), cost_recorder);
// Check the updated costs.
const auto got_op_cost_map = GetOpCostMap(module.get());
EXPECT_THAT(got_op_cost_map, ::testing::ContainerEq(expected_op_cost_map));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tfrt/tests/batch_function_fallback_resource_variable_as_captured_tensor.mlir
+40
View File
@@ -0,0 +1,40 @@
// 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 -pass-pipeline='builtin.module(tf-executor-to-tfrt-pipeline{target-tpurt=true})' %s | FileCheck %s
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func @main
func.func @main_func() -> (tensor<*xf32> {tf_saved_model.index_path = ["a"]}) attributes {tf_saved_model.exported_names = ["main_func"]} {
%0 = tf_executor.graph {
%outputs_0, %control_0 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<501000x128xf32>>>
%outputs_1, %control_1 = tf_executor.island wraps "tf.Cast"(%outputs_0) {Truncate = false} : (tensor<!tf_type.resource<tensor<501000x128xf32>>>) -> tensor<*x!tf_type.resource>
// CHECK: tfrt_fallback_async.batch_function device([[DEVICE:.*]]) @batched_func ([[BATCHED_FUNC_ARG:%.*]])
// CHECK-SAME: Tcaptured = [!corert.resource]
// CHECK-SAME: Tin = []
// CHECK-SAME: Tout = [f32]
%outputs_2, %control_2 = tf_executor.island wraps "tf.BatchFunction"(%outputs_1) {batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 0, 1>, shared_name = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
tf_executor.fetch %outputs_2 : tensor<*xf32>
}
func.return %0 : tensor<*xf32>
}
func.func private @batched_func(%arg0: tensor<*x!tf_type.resource>) -> tensor<?xf32> {
%0 = tf_executor.graph {
%outputs_0, %control_0 = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource>) -> tensor<?xf32>
tf_executor.fetch %outputs_0 : tensor<?xf32>
}
func.return %0 : tensor<?xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/batch_function_lowering.mlir
+41
View File
@@ -0,0 +1,41 @@
// 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 -tf-executor-to-tfrt-pipeline %s | FileCheck %s --dump-input=always
func.func private @batched_function(%arg0: tensor<1x3xf32> {tf._user_specified_name = "0"}, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32> attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
%0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<1x3xf32>
%1 = "tf.AddV2"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "batch/BatchFunction:0"}} {
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK: tfrt_fallback_async.batch_function device("/device:CPU:0") @batched_function
// CHECK-SAME: Tin = [f32]
// CHECK-SAME: Tout = [f32]
// CHECK-SAME: allowed_batch_sizes = [6]
// CHECK-SAME: batch_timeout_micros = 100000 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: max_batch_size = 6 : i64
// CHECK-SAME: max_enqueued_batches = 10 : i64
// CHECK-SAME: num_batch_threads = 1 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
tensorflow/compiler/mlir/tfrt/tests/convert_ref_variables.mlir
+120
View File
@@ -0,0 +1,120 @@
// 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 -tfrt-convert-ref-variables -split-input-file -verify-diagnostics %s | FileCheck %s
// Test the basic cases where all uses of a ref variable can be converted.
// CHECK-LABEL: @init
func.func @init() {
// CHECK-NOT: tf.VariableV2
// CHECK-NOT: tf.Assign
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"
// CHECK-SAME: shared_name = "x"
// CHECK: "tf.AssignVariableOp"([[handle]], {{%.*}})
%0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
%1 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Assign"(%0, %1) {T = i32, device = "", use_locking = true, validate_shape = true} : (tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<!tf_type.int32ref>
func.return
}
// CHECK-LABEL: @inference
func.func @inference() -> tensor<i32> {
// CHECK-NOT: tf.VariableV2
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"
// CHECK-SAME: shared_name = "x"
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
%1 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// -----
// Test the cases when there are both reads and writes, the order of the reads and writes are preserved.
// CHECK-LABEL: @init
func.func @init() -> tensor<i32> {
// CHECK-NOT: tf.VariableV2
// CHECK: [[zero:%.*]] = "tf.Const"
// CHECK-SAME: dense<0>
%0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"
// CHECK-SAME: shared_name = "x"
// CHECK-NEXT: "tf.AssignVariableOp"([[handle]], [[zero]])
// CHECK-NEXT: "tf.ReadVariableOp"([[handle]])
%1 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
%2 = "tf.Assign"(%1, %0) {T = i32, device = "", use_locking = true, validate_shape = true} : (tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<!tf_type.int32ref>
%3 = "tf.Identity"(%1) : (tensor<!tf_type.int32ref>) -> tensor<i32>
// CHECK: [[one:%.*]] = "tf.Const"
// CHECK-SAME: dense<1>
// CHECK-NEXT: "tf.AssignVariableOp"([[handle]], [[one]])
// CHECK-NEXT: "tf.ReadVariableOp"([[handle]])
%4 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%5 = "tf.Assign"(%1, %4) {T = i32, device = "", use_locking = true, validate_shape = true} : (tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<!tf_type.int32ref>
%6 = "tf.Identity"(%1) : (tensor<!tf_type.int32ref>) -> tensor<i32>
func.return %6 : tensor<i32>
}
// CHECK-LABEL: @inference
func.func @inference() -> (tensor<i32>, tensor<i32>, tensor<i32>) {
// CHECK-NOT: tf.VariableV2
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"
// CHECK-SAME: shared_name = "x"
// CHECK: "tf.ReadVariableOp"([[handle]])
// CHECK: "tf.ReadVariableOp"([[handle]])
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
%1 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
%2 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
%3 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
func.return %1, %2, %3 : tensor<i32>, tensor<i32>, tensor<i32>
}
// -----
// Test report error when the shared_name of the tf.VariableV2 op is empty.
// CHECK-LABEL: @inference
func.func @inference() -> tensor<i32> {
// expected-error @+1 {{unable to convert reference variables with empty shared_names.}}
%0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = ""} : () -> tensor<!tf_type.int32ref>
%1 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// -----
// Test conversion when the user is a side-effect-free op.
// CHECK-LABEL: @side_effect_free_user
func.func @side_effect_free_user() -> tensor<2xi32> {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"
// CHECK-SAME: shared_name = "x"
// CHECK: [[value0:%.*]] = "tf.ReadVariableOp"([[handle]])
// CHECK: [[value1:%.*]] = "tf.ReadVariableOp"([[handle]])
// CHECK: "tf.ConcatV2"([[value1]], [[value0]]
// CHECK-SAME: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2xi32>
%axis = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
%1 = "tf.ConcatV2"(%0, %0, %axis) : (tensor<!tf_type.int32ref>, tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<2xi32>
func.return %1 : tensor<2xi32>
}
tensorflow/compiler/mlir/tfrt/tests/cross_device_transfer.mlir
+45
View File
@@ -0,0 +1,45 @@
// 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 -tfrt-cross-device-transfer %s | FileCheck %s
// CHECK-LABEL: test_transfer_op_result
func.func @test_transfer_op_result(%arg0: !tfrt.chain) -> () {
// CHECK-NEXT: %[[RESULT_0:.*]] = corert.get_op_handler %[[ARG_0:.*]] "cpu"
%0 = corert.get_op_handler %arg0 "cpu"
// CHECK-NEXT: %[[RESULT_1:.*]] = corert.get_op_handler %[[ARG_0]] "gpu"
%1 = corert.get_op_handler %arg0 "gpu"
// CHECK-NEXT: %[[RESULT_2:.*]] = corert.create_dense_tensor.i32 {shape = [0], value = []}
%2 = corert.create_dense_tensor.i32 {shape = [0], value = []}
// CHECK-NEXT: %[[RESULT_3:.*]] = corert.executeop(%[[RESULT_0]]) "tf.AddV2"(%[[RESULT_2]], %[[RESULT_2]])
%3 = corert.executeop(%0) "tf.AddV2"(%2, %2) {T = f32, device = "/device:CPU:0"} : 1
// CHECK-NEXT: %[[RESULT_4:.*]] = tfrt.get_device %[[ARG_0]] {device_name = "/device:GPU:0"}
// CHECK-NEXT: %[[RESULT_5:.*]] = corert.get_dst_tensor_type %[[RESULT_3]], %[[RESULT_4]]
// CHECK-NEXT: %[[RESULT_6:.*]] = corert.transfer %[[RESULT_3]], %[[RESULT_4]], %[[RESULT_5]]
// CHECK-NEXT: %[[RESULT_7:.*]] = corert.executeop(%[[RESULT_1]]) "tf.AddV2"(%[[RESULT_6]], %[[RESULT_6]])
%4 = corert.executeop(%1) "tf.AddV2"(%3, %3) {T = f32, device = "/device:GPU:0"} : 1
tfrt.return
}
// CHECK: func @test_transfer_func_arg(%[[ARG_0:.*]]: !tfrt.chain, %[[ARG_1:.*]]: !corert.tensorhandle
func.func @test_transfer_func_arg(%arg0: !tfrt.chain, %arg1: !corert.tensorhandle {tfrt.device = "/device:CPU:0"}) -> () {
// CHECK-NEXT: %[[RESULT_0:.*]] = corert.get_op_handler %[[ARG_0]] "gpu"
%0 = corert.get_op_handler %arg0 "gpu"
// CHECK-NEXT: %[[RESULT_1:.*]] = tfrt.get_device %[[ARG_0]] {device_name = "/device:GPU:0"}
// CHECK-NEXT: %[[RESULT_2:.*]] = corert.get_dst_tensor_type %[[ARG_1]], %[[RESULT_1]]
// CHECK-NEXT: %[[RESULT_3:.*]] = corert.transfer %[[ARG_1]], %[[RESULT_1]], %[[RESULT_2]]
// CHECK-NEXT: %[[RESULT_4:.*]] = corert.executeop(%[[RESULT_0]]) "tf.AddV2"(%[[RESULT_3]], %[[RESULT_3]])
%1 = corert.executeop(%0) "tf.AddV2"(%arg1, %arg1) {T = f32, device = "/device:GPU:0"} : 1
tfrt.return
}
tensorflow/compiler/mlir/tfrt/tests/deduplicate_if_results.mlir
+61
View File
@@ -0,0 +1,61 @@
// 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 -split-input-file -tfrt-deduplicate-if-result %s | FileCheck %s -dump-input=fail
func.func private @then(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
func.return %x, %x : tensor<i32>, tensor<i32>
}
func.func private @else(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
func.return %y, %y : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: then/tfrt_dedup_results
// CHECK: return {{%.*}} : tensor<i32>
// CHECK-LABEL: else/tfrt_dedup_results
// CHECK: return {{%.*}} : tensor<i32>
// CHECK-LABEL: @basic
func.func @basic(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
// CHECK-NEXT: [[r:%.*]] = "tf.If"
// CHECK-NEXT: return [[r]], [[r]] : tensor<i32>, tensor<i32>
%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @else, then_branch = @then, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
return %0, %1 : tensor<i32>, tensor<i32>
}
// -----
func.func private @unmatched_then(%x: tensor<*xi32>, %y: tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) {
func.return %x, %x : tensor<*xi32>, tensor<*xi32>
}
func.func private @unmatched_else(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
func.return %y, %y : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: unmatched_then/tfrt_dedup_results
// CHECK: return {{%.*}} : tensor<*xi32>
// CHECK-LABEL: unmatched_else/tfrt_dedup_results
// CHECK: return {{%.*}} : tensor<i32>
// CHECK-LABEL: @unmatched_then_else_type
func.func @unmatched_then_else_type(%cond: tensor<i1>, %x: tensor<*xi32>, %y: tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) {
// CHECK-NEXT: [[r:%.*]] = "tf.If"
// CHECK-NEXT: return [[r]], [[r]] : tensor<*xi32>, tensor<*xi32>
%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @unmatched_else, then_branch = @unmatched_then, is_stateless = true} : (tensor<i1>, tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>)
return %0, %1 : tensor<*xi32>, tensor<*xi32>
}
tensorflow/compiler/mlir/tfrt/tests/fuse_tpu_compile_and_execute_ops.mlir
+190
View File
@@ -0,0 +1,190 @@
// 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 -verify-diagnostics -split-input-file -tfrt-fuse-tpu-compile-and-execute-ops -canonicalize %s | FileCheck %s --dump-input=fail --dump-input-filter=all
module attributes {tf_saved_model.semantics} {
// Test fusing _TPUCompileMlirOp and TPUExecuteOp into TPUCompileMlirAndExecuteOp.
// CHECK-LABEL: func private @test_fuse_tpu_ops
func.func private @test_fuse_tpu_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
// CHECK-NOT: tf._TPUCompileMlirOp
// CHECK-NOT: tf.TPUCompileSucceededAssert
// CHECK-NOT: tf.TPUExecuteOp
// CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1)
// CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
// CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
%0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
%1 = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
%2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
%compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%3 = "tf.TPUExecute"(%arg0, %0, %program) {device = "/TPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> tensor<*xi32>
func.return %3 : tensor<*xi32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test models using Outside Compilation
// CHECK-LABEL: func private @test_outside_compilation
func.func private @test_outside_compilation(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
// CHECK-NOT: tf._TPUCompileMlirOp
// CHECK-NOT: tf.TPUCompileSucceededAssert
// CHECK-NOT: tf.TPUExecuteOp
// CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1)
// CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
// CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) <{device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"}> {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> ()
// CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
%0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
%1 = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
%2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
%compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
"tf._XlaSendFromHost"(%arg0, %0, %program) {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> ()
%3 = "tf.TPUExecute"(%arg0, %0, %program) {device = "/TPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> tensor<*xi32>
func.return %3 : tensor<*xi32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test models with dynamic bounds ops.
// CHECK-LABEL: func private @test_fuse_dynamic_dimension_ops
func.func private @test_fuse_dynamic_dimension_ops(%arg0: tensor<?x?xi32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<2xi64>, %arg3: tensor<?xi64>, %arg4: tensor<?xi64>) -> tensor<*xi32> {
// CHECK-NOT: tf._TPUCompileMlirOp
// CHECK-NOT: tf.TPUCompileSucceededAssert
// CHECK-NOT: tf.TPUExecuteOp
// CHECK-NOT: tf.SetStaticDimensionBounds
// CHECK: [[read_result:%.*]] = "tf.ReadVariableOp"(%arg1)
// CHECK: [[shape_result_1:%.*]] = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<?x?xi32>) -> tensor<?xi64>
// CHECK: [[shape_result_2:%.*]] = "tf.Shape"([[read_result]]) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
// CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, [[shape_result_2]], %0, %0, %arg2, %arg4, %arg3) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 4, 3>, operands_with_static_shape = [0 : i32, 1 : i32, 3 : i32], producer_name = "default"}> : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<*xi32>, tensor<2xi64>, tensor<?xi64>, tensor<?xi64>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
// CHECK: [[key_1:%.*]], [[exec_result_1:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %2, %0, %1) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 4, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
// CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
%0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
%dyn_arg0 = "tf.SetStaticDimensionBounds" (%arg0, %arg2) :(tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%dyn_0 = "tf.SetStaticDimensionBounds" (%0, %arg3) :(tensor<*xi32>, tensor<?xi64>) -> tensor<?xi64>
%1 = "tf.Shape"(%dyn_arg0) {device = "/CPU:0"} : (tensor<?x?xi32>) -> tensor<?xi64>
%2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
%dyn_2 = "tf.SetStaticDimensionBounds" (%2, %arg4) :(tensor<?xi64>, tensor<?xi64>) -> tensor<?xi64>
%compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%3 = "tf.TPUExecute"(%dyn_arg0, %dyn_2, %0, %dyn_0, %program) {device = "/TPU:0"} : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>, tensor<3x!tf_type.string>) -> tensor<*xi32>
%compilation_status_2, %program_2 = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%4 = "tf.TPUExecute"(%arg0, %2, %0, %1, %program_2) {device = "/TPU:0"} : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>, tensor<3x!tf_type.string>) -> tensor<*xi32>
func.return %3 : tensor<*xi32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func private @reorder_execute_arg_defining_ops
// CHECK: tf.VarHandleOp
// CHECK-NEXT: tf.ReadVariableOp
// CHECK-NEXT: tf.TPUCompileMlirAndExecute
func.func private @reorder_execute_arg_defining_ops(%arg0: tensor<1x3xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
%compilation_status, %program = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%2 = "tf.TPUExecute"(%arg0, %1, %program) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x3xf32>, tensor<3x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
return %2 : tensor<1x1xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func private @spmd_fuse_mulitple_execute_ops
// CHECK-NEXT: %0 = "tf.VarHandleOp"()
// CHECK-NEXT: %1 = "tf.ReadVariableOp"(%0)
// CHECK-NEXT: %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %1)
func.func private @spmd_fuse_mulitple_execute_ops(%arg0: tensor<1x4xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
%cst = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%compilation_status, %program:2 = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<2x1xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<2x1xf32>>>) -> tensor<2x1xf32>
%2:2 = "tf.Split"(%cst, %arg0) {device = "/CPU:0"} : (tensor<i32>, tensor<1x4xf32>) -> (tensor<1x2xf32>, tensor<1x2xf32>)
%3 = "tf.TPUExecute"(%2#0, %1, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x2xf32>, tensor<2x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
%4 = "tf.TPUExecute"(%2#1, %1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x2xf32>, tensor<2x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
return %3 : tensor<1x1xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func private @spmd_fuse_mulitple_execute_ops_2
// CHECK-NEXT: %0 = "tf.VarHandleOp"()
// CHECK-NEXT: %1 = "tf.ReadVariableOp"(%0)
// CHECK-NEXT: %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %1)
func.func private @spmd_fuse_mulitple_execute_ops_2(%arg0: tensor<1x1xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
%cst = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%compilation_status, %program:2 = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<2x1xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<2x1xf32>>>) -> tensor<2x1xf32>
%2:2 = "tf.Split"(%cst, %1) {device = "/CPU:0"} : (tensor<i32>, tensor<2x1xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)
%3 = "tf.TPUExecute"(%arg0, %2#0, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
%4 = "tf.TPUExecute"(%arg0, %2#1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
return %3 : tensor<1x1xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func private @spmd_fuse_split_nd_ops
// CHECK-NEXT: %0 = "tf.VarHandleOp"()
// CHECK-NEXT: %1 = "tf.ReadVariableOp"(%0)
// CHECK-NEXT: %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %1)
func.func private @spmd_fuse_split_nd_ops(%arg0: tensor<1x4xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
%cst = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%compilation_status, %program:4 = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>)
"tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
%0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<1x1xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<1x1xf32>>>) -> tensor<1x1xf32>
%2:2 = "tf.Split"(%cst, %arg0) {device = "/CPU:0"} : (tensor<i32>, tensor<1x4xf32>) -> (tensor<1x2xf32>, tensor<1x2xf32>)
%3:2 = "tf.Split"(%cst, %2#0) {device = "/CPU:0"} : (tensor<i32>, tensor<1x2xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)
%4:2 = "tf.Split"(%cst, %2#1) {device = "/CPU:0"} : (tensor<i32>, tensor<1x2xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)
%5 = "tf.TPUExecute"(%3#0, %1, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
%6 = "tf.TPUExecute"(%3#1, %1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
%7 = "tf.TPUExecute"(%4#0, %1, %program#2) {_producer_name = "UNKNOWN", device = "/TPU:2"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
%8 = "tf.TPUExecute"(%4#1, %1, %program#3) {_producer_name = "UNKNOWN", device = "/TPU:3"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
return %5 : tensor<1x1xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops.mlir
+329
View File
@@ -0,0 +1,329 @@
// 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 -split-input-file -tfrt-lower-tf-savedmodel=hoist-invariant-ops=true %s | FileCheck %s --dump-input=fail --dump-input-filter=all
module attributes {tf_saved_model.semantics} {
// Test hoisting varhandle op.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: [[x:%.*]] = "tf.ReadVariableOp"([[handle]]) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
// CHECK-LABEL: func @test_hoist_varhandleop
func.func @hoist_varhandleop(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_hoist_varhandleop"]} {
// CHECK-NOT: tf.VarHandleOp
// CHECK-NOT: tf.ReadVariableOp
// CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK: [[r:%.*]] = "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: return [[r]]
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
%x = "tf.ReadVariableOp"(%handle) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test hoisting hash table op.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[handle:%.*]] = "tf.HashTableV2"()
// CHECK-SAME: shared_name = "x"
// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]])
// CHECK: "tf._TfrtSetResource"([[x]]) <{index = [[size_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i64>) -> ()
// CHECK: func @test_hoist_hash_table
func.func @hoist_hash_table(%arg: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["input"]}, %default: tensor<i64> {tf_saved_model.index_path = ["default"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]})
attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} {
// CHECK-NOT: tf.HashTableV2
// CHECK-NOT: tf.LookupTableSizeV2
// CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() <{container = ["", ""], indices = [0, 1], shared_name = [{{.*}}, {{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK: [[r:%.*]] = "tf.LookupTableFindV2"([[v]]#[[handle_idx]]
// CHECK: return [[v]]#[[size_idx]], [[r]]
%0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "x", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
%1 = "tf.LookupTableSizeV2"(%0) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
%2 = "tf.LookupTableFindV2"(%0, %arg, %default) {device = "/CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
func.return %1, %2 : tensor<i64>, tensor<*xi64>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test hoisting const op.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK: [[x:%.*]] = "tf.AddV2"([[const]], [[const]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
// CHECK: [[const_1:%.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK: "tf._TfrtSetResource"([[const_1]]) <{index = 1 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
// CHECK-LABEL: func @test_hoist_const
func.func @hoist_const(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_hoist_const"]} {
// CHECK-NOT: tf.Const
// CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK-NEXT: "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: return
%const = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%x = "tf.AddV2"(%const, %const) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
// CHECK-LABEL: func @test_hoist_const_return
func.func @hoist_const_return(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_hoist_const_return"]} {
// CHECK-NOT: tf.Const
// CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [1], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK-NEXT: return [[v]]
%const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
func.return %const : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test not hoisting `tf.BatchFunction`.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {device = "/CPU:0"} : () -> tensor<1xi32>
// CHECK: "tf._TfrtSetResource"([[const]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<1xi32>) -> ()
// CHECK-LABEL: func.func private @func_with_batch_function
func.func private @func_with_batch_function() -> tensor<*xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "StatefulPartitionedCall:0"}} {
// CHECK: "tf._TfrtGetResource"()
%cst = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {device = "/CPU:0"} : () -> tensor<1xi32>
// CHECK: "tf.BatchFunction"
%0 = "tf.BatchFunction"(%cst) <{allowed_batch_sizes = [1], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", enable_large_batch_splitting = true, f = @_batched, low_priority_allowed_batch_sizes = [], low_priority_batch_timeout_micros = 0 : i64, low_priority_max_batch_size = 0 : i64, low_priority_max_enqueued_batches = 0 : i64, max_batch_size = 1 : i64, max_enqueued_batches = 1 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch_function___inference_signature_wrapper_fn_with_defaults_36"}> {device = "/CPU:0"} : (tensor<1xi32>) -> tensor<*xi32>
return %0 : tensor<*xi32>
}
func.func private @_batched(%arg0: tensor<1xi32>) -> tensor<1xi32> {
return %arg0 : tensor<1xi32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test hoisting write side-effect ops.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
// CHECK: "tf._TfrtSetResource"([[const]]) <{index = [[const_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> ()
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> ()
// CHECK: func @test_hoist_var_read_write
func.func @hoist_var_read_write() -> (tensor<i32> {tf_saved_model.index_path = ["x"]}, tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_hoist_var_read_write"]} {
// CHECK-NOT: tf.Const
// CHECK-NOT: tf.VarHandleOp
// CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() <{container = ["", ""], indices = [0, 1], shared_name = [{{.*}}, {{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ({{.*}})
// CHECK: [[x:%.*]] = "tf.ReadVariableOp"([[v]]#[[handle_idx]]) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
// CHECK-NEXT: "tf.AssignVariable"([[v]]#[[handle_idx]], [[v]]#[[const_idx]]) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
// CHECK-NEXT: [[r:%.*]] = "tf.ReadVariableOp"([[v]]#[[handle_idx]]) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
// CHECK-NEXT: return [[x]], [[r]]
%const = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
%x = "tf.ReadVariableOp"(%handle) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
"tf.AssignVariable"(%handle, %const) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
%r = "tf.ReadVariableOp"(%handle) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %x, %r : tensor<i32>, tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test not hoisting read variable op that used by control flow ops if var handle op and read variable op are separated, but still hoists const ops and var handle ops.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf._TfrtSetResource"([[handle]])
// CHECK-SAME: index = [[handle_index:.*]]
// CHECK: [[handle1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf._TfrtSetResource"([[handle1]])
// CHECK-SAME: index = [[handle1_index:.*]]
// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = "/CPU:0"} : () -> tensor<i1>
// CHECK: "tf._TfrtSetResource"([[const]])
// CHECK-SAME: index = [[const_index:.*]]
func.func private @some_func(
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: tf.ReadVariableOp
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @test_not_hoist_stateful_call
func.func @not_hoist_stateful_call(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_not_hoist_stateful_call"]} {
// CHECK-NOT: tf.VarHandleOp
// CHECK: "tf._TfrtGetResource"()
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: tf.StatefulPartitionedCall
%x = "tf.StatefulPartitionedCall"(%handle) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
// CHECK-LABEL: func @test_not_hoist_if
func.func @not_hoist_if(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_not_hoist_if"]} {
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK-NOT: tf.Const
// CHECK: "tf._TfrtGetResource"()
%cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
// CHECK: tf.If
%x = "tf.If"(%cond, %handle) {then_branch = @some_func, else_branch = @some_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test hoist var handle op and read variable op in the batch function.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK-NOT: tf.VarHandleOp
// CHECK-NOT: tf.ReadVariableOp
// CHECK: "tf._TfrtGetResource"()
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%3 = "tf.Identity"(%2) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK-NOT: tf.VarHandleOp
// CHECK: "tf._TfrtGetResource"()
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK: "tf.BatchFunction"(%arg0, %0)
// CHECK: operandSegmentSizes = array<i32: 1, 1>
%1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test not hoisting callees in init functions.
"tf_saved_model.session_initializer"() {initializers = [@init]} : () -> ()
func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
%var0 = "tf.VarHandleOp"() {container = "", shared_name = "var0"} : () -> tensor<!tf_type.resource<tensor<i1>>>
%cond = "tf.ReadVariableOp"(%var0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i1>>>) -> tensor<i1>
%x = "tf.StatefulPartitionedCall"(%cond) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<i1>) -> (tensor<i32>)
%var1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<i32>>>
"tf.AssignVariable"(%var1, %x) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
func.return
}
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK-NEXT: return
// CHECK-LABEL: func private @some_func
func.func private @some_func(%arg: tensor<i1>) -> tensor<i32> {
// CHECK-NOT: tf._TfrtGetResource
%const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32> } : () -> tensor<i32>
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
%0 = "tf.ReadVariableOp"(%handle) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%r = "tf.SelectV2"(%arg, %const, %0) {device = "/CPU:0"} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test not hoisting callees in xla launch functions.
// CHECK-LABEL: func private @xla_func
func.func private @xla_func(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK-NOT: tf._TfrtGetResource
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%3 = "tf.Identity"(%2) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.XlaLaunch"(%arg0, %0) {device = "/device:GPU:0", function = @xla_func, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test not hoisting in TPU functions.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK-NEXT: return
// CHECK-LABEL: func private @func2
func.func private @func2(%arg: tensor<i1>) -> tensor<i32> {
// CHECK-NOT: tf._TfrtGetResource
"tf.TPUReplicateMetadata"() {_tpu_replicate = "0", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 4 : i64, num_replicas = 1 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", tpu_compile_options_proto = "", use_spmd_for_xla_partitioning = true, use_tpu = true} : () -> ()
%const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32> } : () -> tensor<i32>
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
%0 = "tf.ReadVariableOp"(%handle) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%r = "tf.SelectV2"(%arg, %const, %0) {device = "/CPU:0"} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
}
tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops_mlrt.mlir
+45
View File
@@ -0,0 +1,45 @@
// 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 -split-input-file -tfrt-lower-tf-savedmodel="hoist-invariant-ops=true fuse-get-resource-ops=false" %s | FileCheck %s --dump-input=fail --dump-input-filter=all
module attributes {tf_saved_model.semantics} {
// Test hoisting hash table op.
// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[handle:%.*]] = "tf.HashTableV2"()
// CHECK-SAME: shared_name = "x"
// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]])
// CHECK: "tf._TfrtSetResource"([[x]]) <{index = [[size_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i64>) -> ()
// CHECK: func @test_hoist_hash_table
func.func @hoist_hash_table(%arg: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["input"]}, %default: tensor<i64> {tf_saved_model.index_path = ["default"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]})
attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} {
// CHECK-NOT: tf.HashTableV2
// CHECK-NOT: tf.LookupTableSizeV2
// CHECK-DAG: [[v0:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [[[handle_id]]], shared_name = [{{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK-DAG: [[v1:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [[[size_id]]], shared_name = [{{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK-DAG: [[r:%.*]] = "tf.LookupTableFindV2"([[v0]]
// CHECK-DAG: return [[v1]], [[r]]
%0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "x", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
%1 = "tf.LookupTableSizeV2"(%0) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
%2 = "tf.LookupTableFindV2"(%0, %arg, %default) {device = "/CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
func.return %1, %2 : tensor<i64>, tensor<*xi64>
}
}
// -----
tensorflow/compiler/mlir/tfrt/tests/ifrt/BUILD
+21
View File
@@ -0,0 +1,21 @@
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "//tensorflow/compiler/mlir:run_lit.sh",
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//mlir:run_lit.sh",
],
)
tensorflow/compiler/mlir/tfrt/tests/ifrt/lower_to_ifrt_restore_variable.mlir
+322
View File
@@ -0,0 +1,322 @@
// 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 -split-input-file -verify-diagnostics -lower-to-ifrt-restore-variable %s | FileCheck %s
// -----
// single variable
// CHECK-LABEL: func.func @restore_single() {
// CHECK-NEXT: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
// CHECK-SAME: {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_single() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
}
// -----
// single variable: VarHandleOp is before RestoreV2
// CHECK-LABEL: func.func @varhandle_before_restore() {
// CHECK-NEXT: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
// CHECK-SAME: {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @varhandle_before_restore() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
"tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
}
// -----
// multiple variables
// CHECK-LABEL: func.func @restore_multiple() {
// CHECK-NEXT: [[PREFIX:%.*]] = "tf.Const"()
// CHECK-NEXT: [[SLICE:%.*]] = "tf.Const"()
// CHECK-NEXT: [[NAME:%.*]] = "tf.Const"()
// CHECK-NEXT: [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "x", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK-NEXT: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]], [[HANDLEZ]])
// CHECK-SAME: {restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_multiple() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<["", ""]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<["y", "z"]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
%0:2 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<3x1xf32>, tensor<1x3xf32>)
%1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
%2 = "tf.VarHandleOp"() <{container = "x", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
"tf.AssignVariableOp"(%2, %0#1) : (tensor<!tf_type.resource<tensor<1x3xf32>>>, tensor<1x3xf32>) -> ()
return
}
}
// -----
// Restored variable is not assigned with a name is an error.
module {
func.func @unassigned_restore_return_error() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<["", ""]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<["y", "z"]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
//expected-error@below {{'tf.RestoreV2' op expects 2 valid users, but got 1}}
%0:2 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<3x1xf32>, tensor<1x3xf32>)
%1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
}
// -----
// Restored tensor is consumed by an op other than AssignVariableOp, it is returned as an output by IfrtRestoreVariableOp.
// CHECK-LABEL: func.func @restore_with_consumer() {
// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
// CHECK-NEXT: %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: %1 = "tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %0) <{restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: false>}> : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
// CHECK-NEXT: %2 = "tf.ReluOp"(%1) : (tensor<3x1xf32>) -> tensor<3x1xf32>
// CHECK-NEXT: %3 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: "tf.AssignVariableOp"(%3, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_with_consumer() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%2 = "tf.ReluOp"(%0) : (tensor<3x1xf32>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
}
// -----
// variable with cast
// CHECK-LABEL: func.func @restore_with_cast() {
// CHECK-NEXT: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
// CHECK-NEXT: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
// CHECK-SAME: {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_with_cast() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.Cast"(%0) <{Truncate = false}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
%2 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
"tf.AssignVariableOp"(%2, %1) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
return
}
}
// -----
// variable and table lookup
// CHECK-LABEL: func.func @restore_var_and_table()
// CHECK-NEXT: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"model/foo"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[NAMES:%.*]] = "tf.Const"() <{value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
// CHECK-NEXT: [[VAR1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "var1"}> : () -> tensor<!tf_type.resource<tensor<*xi32>>>
// CHECK-NEXT: [[TABLE1_KEYS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_keys"}> : () -> tensor<!tf_type.resource<tensor<*xi64>>>
// CHECK-NEXT: [[TABLE1_VALS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_vals"}> : () -> tensor<!tf_type.resource<tensor<*xf32>>>
// CHECK-NEXT: [[RETURNED_VALUES:%.*]]:2 = "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAMES]], [[SLICE]], [[VAR1]], [[TABLE1_KEYS]], [[TABLE1_VALS]])
// CHECK-SAME: <{restored_dtypes = [i32, i64, f32], returned_tensor_names = ["table1_keys", "table1_vals"], truncate_in_cast = array<i1: false, false, false>}> : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<!tf_type.resource<tensor<*xi32>>>, tensor<!tf_type.resource<tensor<*xi64>>>, tensor<!tf_type.resource<tensor<*xf32>>>) -> (tensor<*xi64>, tensor<*xf32>)
// CHECK-NEXT: [[TABLE1:%.*]] = "tf.HashTableV2"() <{container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32}> : () -> tensor<!tf_type.resource>
// CHECK-NEXT: "tf.LookupTableImportV2"([[TABLE1]], [[RETURNED_VALUES]]#0, [[RETURNED_VALUES]]#1) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_var_and_table() {
%cst = "tf.Const"() {value = dense<"model/foo"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() {value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
%cst_1 = "tf.Const"() {value = dense<["", "", ""]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
%0:3 = "tf.RestoreV2"(%cst, %cst_0, %cst_1) {dtypes = [i32, i64, f32]} : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi64>, tensor<*xf32>)
%1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<*xi32>>>
"tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<*xi32>>>, tensor<*xi32>) -> ()
%2 = "tf.HashTableV2"() {container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32} : () -> tensor<!tf_type.resource>
"tf.LookupTableImportV2"(%2, %0#1, %0#2) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
return
}
}
// -----
// variable and dense table lookup
// CHECK-LABEL: func.func @restore_var_and_dense_table()
// CHECK-NEXT: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"model/foo"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[NAMES:%.*]] = "tf.Const"() <{value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
// CHECK-NEXT: [[VAR1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "var1"}> : () -> tensor<!tf_type.resource<tensor<*xi32>>>
// CHECK-NEXT: [[TABLE1_KEYS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_keys"}> : () -> tensor<!tf_type.resource<tensor<*xi64>>>
// CHECK-NEXT: [[TABLE1_VALS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_vals"}> : () -> tensor<!tf_type.resource<tensor<*xf32>>>
// CHECK-NEXT: [[RETURNED_VALUES:%.*]]:2 = "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAMES]], [[SLICE]], [[VAR1]], [[TABLE1_KEYS]], [[TABLE1_VALS]])
// CHECK-SAME: <{restored_dtypes = [i32, i64, f32], returned_tensor_names = ["table1_keys", "table1_vals"], truncate_in_cast = array<i1: false, false, false>}> : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<!tf_type.resource<tensor<*xi32>>>, tensor<!tf_type.resource<tensor<*xi64>>>, tensor<!tf_type.resource<tensor<*xf32>>>) -> (tensor<*xi64>, tensor<*xf32>)
// CHECK-NEXT: [[EMPTY_KEY:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: [[DEFAULT_VALUE:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: [[TABLE1:%.*]] = "tf.MutableDenseHashTableV2"(%cst_2, %cst_3) <{container = "", shared_name = "table1", value_dtype = f32}> {key_dtype = i64} : (tensor<i64>, tensor<f32>) -> tensor<!tf_type.resource>
// CHECK-NEXT: "tf.LookupTableImportV2"([[TABLE1]], [[RETURNED_VALUES]]#0, [[RETURNED_VALUES]]#1) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_var_and_dense_table() {
%cst = "tf.Const"() {value = dense<"model/foo"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() {value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
%cst_1 = "tf.Const"() {value = dense<["", "", ""]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
%0:3 = "tf.RestoreV2"(%cst, %cst_0, %cst_1) {dtypes = [i32, i64, f32]} : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi64>, tensor<*xf32>)
%1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<*xi32>>>
"tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<*xi32>>>, tensor<*xi32>) -> ()
%empty_key = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
%default_value = "tf.Const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
%2 = "tf.MutableDenseHashTableV2"(%empty_key, %default_value) {container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32} : (tensor<i64>, tensor<f32>) -> tensor<!tf_type.resource>
"tf.LookupTableImportV2"(%2, %0#1, %0#2) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
return
}
}
// -----
// restored variable assigned to one var handle and then read and assigned to another
// CHECK-LABEL: func.func @restore_and_copy() {
// CHECK-DAG: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}>
// CHECK-DAG: [[VAR_Y:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-DAG: [[VAR_Z:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-DAG: [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
// CHECK-DAG: [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
// CHECK: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[VAR_Y]], [[VAR_Z]])
// CHECK-SAME: <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}>
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NOT: "tf.ReadVariableOp"
// CHECK-NEXT: return
module {
func.func @restore_and_copy() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
%read = "tf.ReadVariableOp"(%1) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%2 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%2, %read) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
}
// -----
// restored variable assigned to one var handle and then read and assigned to another in a different block/function
// CHECK-LABEL: func.func @restore_derived() {
// CHECK-DAG: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_variables"> : tensor<!tf_type.string>}>
// CHECK-DAG: [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-DAG: [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-DAG: [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
// CHECK-DAG: [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
// CHECK: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[HANDLEY]], [[HANDLEZ]])
// CHECK-SAME: <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}>
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NEXT: return
module {
func.func @restore_derived() {
%cst = "tf.Const"() <{value = dense<"restore_variables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
func.func @derived_init() {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%2 = "tf.Identity"(%1) : (tensor<3x1xf32>) -> tensor<3x1xf32>
%3 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%3, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
}
// -----
// restored variable assigned to one var handle and then read, cast, and assigned to another
// CHECK-LABEL: func.func @restore_and_cast_copy() {
// CHECK-DAG: [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}>
// CHECK-DAG: [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-DAG: [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
// CHECK-DAG: [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
// CHECK-DAG: [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
// CHECK: "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[HANDLEY]], [[HANDLEZ]])
// CHECK-SAME: <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, true>}>
// CHECK-NOT: "tf.RestoreV2"
// CHECK-NOT: "tf.ReadVariableOp"
// CHECK-NEXT: return
module {
func.func @restore_and_cast_copy() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
%read = "tf.ReadVariableOp"(%1) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%cast = "tf.Cast"(%read) <{Truncate = true}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
%2 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
"tf.AssignVariableOp"(%2, %cast) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
return
}
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/propagate_static_shapes.mlir
+104
View File
@@ -0,0 +1,104 @@
// 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 -split-input-file -propagate-static-shapes %s | FileCheck %s
// -----
// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 123 : i64}
// CHECK: return %arg0
// CHECK-LABEL: func.func @main
// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 123 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
// CHECK-NEXT: return %[[C1]]
module {
func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 123 : i64} {
func.return %arg0 : tensor<?x?xi32>
}
func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
%0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
%1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%2 = "tf.IfrtCall"(%1) {program_id = 123 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
func.return %2 : tensor<?x?xi32>
}
}
// -----
// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 3 : i32}, %arg1: tensor<?xi32> {tf._static_shape_arg_idx = 4 : i32}, %arg2: tensor<?x?xf32>, %arg3: tensor<2xi64>, %arg4: tensor<1xi64>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) attributes {tfrt_ifrt_serving.program_id = 456 : i64}
// CHECK: return %arg0, %arg1, %arg2
// CHECK-LABEL: func.func @main
// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
// CHECK-NEXT: %[[C1:.*]] = "tf.Const"
// CHECK-NEXT: %[[R:.*]]:3 = "tf.IfrtCall"(%arg0, %arg1, %arg2, %[[C0]], %[[C1]]) <{operandSegmentSizes = array<i32: 3, 2>, program_id = 456 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>, tensor<2xi64>, tensor<1xi64>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>)
// CHECK-NEXT: return %[[R]]#0, %[[R]]#1, %[[R]]#2
module {
func.func @callee(%arg0: tensor<?x?xi32>, %arg1: tensor<?xi32>, %arg2: tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) attributes {tfrt_ifrt_serving.program_id = 456 : i64} {
func.return %arg0, %arg1, %arg2 : tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>
}
func.func @main(%arg0: tensor<?x?xi32>, %arg1: tensor<?xi32>, %arg2: tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) {
%c0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
%c1 = "tf.Const"() {value = dense<4> : tensor<1xi64>} : () -> tensor<1xi64>
%0 = "tf.SetStaticDimensionBounds"(%arg0, %c0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%1 = "tf.SetStaticDimensionBounds"(%arg1, %c1) : (tensor<?xi32>, tensor<1xi64>) -> tensor<?xi32>
%2:3 = "tf.IfrtCall"(%0, %1, %arg2) {program_id = 456 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 3, 0>} : (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>)
func.return %2#0, %2#1, %2#2 : tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>
}
}
// -----
// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 789 : i64}
// CHECK: return %arg0
// CHECK-LABEL: func.func @main
// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 789 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
// CHECK-NEXT: %[[C2:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 789 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
// CHECK-NEXT: return %[[C2]]
module {
func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 789 : i64} {
func.return %arg0 : tensor<?x?xi32>
}
func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
%0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
%1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%2 = "tf.IfrtCall"(%1) {program_id = 789 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
%3 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%4 = "tf.IfrtCall"(%3) {program_id = 789 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
func.return %4 : tensor<?x?xi32>
}
}
// -----
// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 999 : i64}
// CHECK: return %arg0
// CHECK-LABEL: func.func @main
// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 999 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
// CHECK-NEXT: %[[C2:.*]] = "tf.AsyncIfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 999 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
// CHECK-NEXT: return %[[C2]]
module {
func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 999 : i64} {
func.return %arg0 : tensor<?x?xi32>
}
func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
%0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
%1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%2 = "tf.IfrtCall"(%1) {program_id = 999 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
%3 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
%4 = "tf.AsyncIfrtCall"(%3) {program_id = 999 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
func.return %4 : tensor<?x?xi32>
}
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/rewrite_cluster_to_ifrt_call.mlir
+140
View File
@@ -0,0 +1,140 @@
// 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 -split-input-file -rewrite-cluster-to-ifrt-call %s | FileCheck %s
// TODO(b/316226111): the printer may not guarantee the same order of fields. Rewrite the checks to be less sensitive to proto serialization formats.
// -----
// Non-SPMD: one input and one output
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
// CHECK-NEXT: "tf.IfrtCall"(%arg0)
// CHECK-SAME: {operandSegmentSizes = array<i32: 1, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
// CHECK-SAME: (tensor<1x3xf32>) -> tensor<1x3xf32>
// CHECK: return
//
// CHECK: func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>)
// CHECK-SAME: __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 1 } dim { size: 3 } } kind: PARAMETER sharding { } is_bounded_dynamic_dim: false } retvals { sharding { } } num_replicas: 1 num_cores_per_replica: 1 "
// CHECK-SAME: device_assignment = []
// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
// CHECK: return
module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
%outputs = "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
%0 = "tf_device.cluster_func"(%arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = [""], num_cores_per_replica = 1 : i64, device_assignment = [], topology = "", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = false, use_tpu = true} : (tensor<1x3xf32>) -> (tensor<1x3xf32>)
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
return %arg0 : tensor<1x3xf32>
}
}
// -----
// SPMD: one input and no return
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) {
// CHECK-NEXT: "tf.IfrtCall"(%arg0)
// CHECK-SAME: {operandSegmentSizes = array<i32: 1, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
// CHECK-SAME: (tensor<1x3xf32>) -> ()
// CHECK: return
//
// CHECK: func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>)
// CHECK-SAME: __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 1 } dim { size: 3 } } kind: PARAMETER sharding { type: OTHER tile_assignment_dimensions: 2 tile_assignment_dimensions: 1 tile_assignment_devices: 0 tile_assignment_devices: 1 } is_bounded_dynamic_dim: false } num_replicas: 1 num_cores_per_replica: 2 use_spmd_for_xla_partitioning: true "
// CHECK-SAME: device_assignment = [0, 0, 0, 0, 0, 0, 0, 1]
// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
// CHECK: return
module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
func.func @serving_default(%arg0: tensor<1x3xf32>) -> () {
%outputs = "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
"tf_device.cluster_func"(%arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}"], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>) -> ()
return
}
// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>) -> () {
return
}
}
// -----
// Multiple ifrt calls and have two sharded arguments
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
// CHECK-NEXT: %0 = "tf.IfrtCall"(%arg1, %arg0)
// CHECK-SAME: {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
// CHECK-SAME: (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT: %1 = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
// CHECK-NEXT: %2 = "tf.IfrtCall"(%1, %arg0)
// CHECK-SAME: {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID]] : i64, variable_arg_indices = []
// CHECK-SAME: (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT: %3 = "tf.add"(%0, %2) : (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
// CHECK: return
//
// CHECK: func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-SAME: device_assignment = [0, 0, 0, 0, 0, 0, 0, 1]
// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
// CHECK-NEXT: %0 = "tf.MatMul"(%arg0, %arg1)
// CHECK: return
module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
%outputs = "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
%outputs_0 = "tf_device.cluster_func"(%arg1, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%duplicate_arg = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
%outputs_1 = "tf_device.cluster_func"(%duplicate_arg, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%outputs_2 = "tf.add"(%outputs_0, %outputs_1): (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
return %outputs_2 : tensor<1x1xf32>
}
// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {
%outputs_0 = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
return %outputs_0 : tensor<1x1xf32>
}
}
// -----
// Missing topology and device assignment attribute in spmd is ok
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
// CHECK-NEXT: %0 = "tf.IfrtCall"(%arg1, %arg0)
// CHECK-SAME: {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
// CHECK-SAME: (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK: return
//
// CHECK: func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-SAME: device_assignment = []
// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
// CHECK-NEXT: %0 = "tf.MatMul"(%arg0, %arg1)
// CHECK: return
module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
%outputs = "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
%outputs_0 = "tf_device.cluster_func"(%arg1, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
return %outputs_0 : tensor<1x1xf32>
}
// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {
%outputs_0 = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
return %outputs_0 : tensor<1x1xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/sink_variable_as_named_array.mlir
+144
View File
@@ -0,0 +1,144 @@
// 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 -split-input-file -tf-device-decompose-resource-ops -sink-variable-as-named-array %s | FileCheck %s
// -----
// Basic test: all variables tensors are for devices and sinked as named ifrt arrays
//
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
// CHECK-NEXT: [[HANDLE2:%.*]] = "tf.VarHandleOp"
// CHECK-NEXT: [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
// CHECK-SAME: used_by_host = false
// CHECK-NEXT: [[RES:%.*]] = "tf.IfrtCall"([[KEY]], %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
// CHECK-SAME: : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT: return [[RES]] : tensor<1x1xf32>
//
module {
func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%result = "tf.IfrtCall"(%2, %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<3x1xf32>, tensor<1x3xf32>) -> (tensor<1x1xf32>)
return %result : tensor<1x1xf32>
}
}
// -----
// Variable tensor for host can still be used.
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
// CHECK: "tf.VarHandleOp"
// CHECK-NOT: [[VARIABLE:%.*]] = "tf.ReadVariableOp"
// CHECK-NEXT: [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
// CHECK-SAME: used_by_host = true
// CHECK-NEXT: [[MATRES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
// CHECK-NEXT: [[RES:%.*]] = "tf.IfrtCall"(%arg0, [[KEY]]) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}>
// CHECK-NEXT: return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>
//
module {
func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%3 = "tf.MatMul"(%arg0, %2) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%result = "tf.IfrtCall"(%arg0, %2) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x3xf32>, tensor<3x1xf32>) -> (tensor<1x1xf32>)
return %result, %3 : tensor<1x1xf32>, tensor<1x1xf32>
}
}
// -----
// Variable tensor is only for host
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
// CHECK: "tf.VarHandleOp"
// CHECK-NOT: tf.ReadVariableOp
// CHECK-NEXT: [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
// CHECK-SAME: used_by_host = true
// CHECK-NEXT: [[RES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
// CHECK-NEXT: return [[RES]] : tensor<1x1xf32>
//
module {
func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%3 = "tf.MatMul"(%arg0, %2) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
return %3: tensor<1x1xf32>
}
}
// -----
// Resources that are created in the same module are not sinked.
//
// CHECK-LABEL: func.func @serving_default
// CHECK-NOT: IfrtLoadVariable
// CHECK: "tf.VarHandleOp"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: "tf.ReadVariableOp"
// CHECK-NEXT: "tf.StatefulPartitionedCall"
// CHECK-NEXT: return
//
module {
func.func @serving_default() -> tensor<*xi32> {
%cst = "tf.Const"() <{value = dense<"some_test.txt"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
"tf.AssignVariableOp"(%0, %cst) <{validate_shape = false}> : (tensor<!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()
%2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<!tf_type.string>>>) -> tensor<*x!tf_type.string>
%4 = "tf.StatefulPartitionedCall"(%2) <{config = "", config_proto = "", executor_type = "", f = @__initializer}> : (tensor<*x!tf_type.string>) -> tensor<*xi32>
return %4: tensor<*xi32>
}
func.func @__initializer(%arg0: tensor<*x!tf_type.string>) -> tensor<i32> {
%0 = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
return %0 : tensor<i32>
}
}
// -----
// Decomposable Resource Ops usage
//
// CHECK-LABEL: func.func @serving_default
// CHECK: "tf.VarHandleOp"
// CHECK-NEXT: "tf.IfrtLoadVariable"
// CHECK-NEXT: "tf.GatherV2"
// CHECK-NEXT: return
//
module {
func.func @serving_default() -> tensor<1x3xbf16> {
%cst = "tf.Const"() <{value = dense<[1]> : tensor<1xi32>}> : () -> tensor<1xi32>
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<2x3xbf16>>>
%1 = "tf.ResourceGather"(%0, %cst) <{batch_dims = 0 : i64, validate_indices = true}> : (tensor<!tf_type.resource<tensor<2x3xbf16>>>, tensor<1xi32>) -> tensor<1x3xbf16>
return %1: tensor<1x3xbf16>
}
}
// -----
// AsyncIfrtCall test: all variables tensors are for devices and sinked as named ifrt arrays
//
//
// CHECK-LABEL: func.func @serving_default_async(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
// CHECK-NEXT: [[HANDLE2:%.*]] = "tf.VarHandleOp"
// CHECK-NEXT: [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
// CHECK-SAME: used_by_host = false
// CHECK-NEXT: [[RES:%.*]] = "tf.AsyncIfrtCall"([[KEY]], %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
// CHECK-SAME: : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT: return [[RES]] : tensor<1x1xf32>
//
module {
func.func @serving_default_async(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%result = "tf.AsyncIfrtCall"(%2, %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<3x1xf32>, tensor<1x3xf32>) -> (tensor<1x1xf32>)
return %result : tensor<1x1xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_device_cleanup.mlir
+22
View File
@@ -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.
// ==============================================================================
// RUN: tf-tfrt-opt %s -tf-device-cleanup | FileCheck %s
// CHECK-LABEL: func @ops_with_device
func.func @ops_with_device() {
%0 = "tf.VarHandleOp"() {container = "", shared_name = "var", device = "/device/..."} : () -> tensor<!tf_type.resource<tensor<1xf32>>>
// CHECK-NOT: device = "/device/..."
func.return
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_identity_propagation.mlir
+52
View File
@@ -0,0 +1,52 @@
// 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-identity-propagation -canonicalize | FileCheck %s
// CHECK-LABEL: func @identity
// CHECK-SAME: (%[[ARG0:.*]]: tensor<i32>)
func.func @identity(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK-NOT: "tf.Identity"
%0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
// CHECK: return %[[ARG0]]
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @identity_terminator
// CHECK-SAME: (%[[ARG0:.*]]: tensor<i32>)
func.func @identity_terminator(%arg0: tensor<i32>) -> (tensor<*xi32>, tensor<i32>) {
// CHECK: %[[IDENTITY:.*]] = "tf.Identity"
%0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<*xi32>
// CHECK-NOT: "tf.Identity"
%1 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
// CHECK: return %[[IDENTITY]], %[[ARG0]]
func.return %0, %1 : tensor<*xi32>, tensor<i32>
}
// CHECK-LABEL: func @xla_sharding
func.func @xla_sharding(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: %[[OUTPUT:.*]] = "tf.Identity"
%0 = "tf.Identity"(%arg0) {_XlaSharding = ""} : (tensor<i32>) -> tensor<i32>
// CHECK: return %[[OUTPUT]]
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @identity_n
// CHECK-SAME: (%[[ARG0:.*]]: tensor<i32>, %[[ARG1:.*]]: tensor<f32>)
func.func @identity_n(%arg0: tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>, tensor<f32>) {
// CHECK-NOT: "tf.IdentityN"
%0:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>)
// CHECK: return %[[ARG0]], %[[ARG1]]
func.return %0#0, %0#1 : tensor<i32>, tensor<f32>
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_merging.mlir
+35
View File
@@ -0,0 +1,35 @@
// 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-restore-merging | FileCheck %s
// CHECK-LABEL: func @single_restore_group
// CHECK-SAME: (%[[ARG0:.*]]: {{.*}})
func.func @single_restore_group(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
%0 = "tf.Const"() {value = dense<"foo"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%1 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%2 = "tf.RestoreV2"(%arg0, %0, %1) : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<*xf32>
%3 = "tf.Const"() {value = dense<"bar"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%4 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%5 = "tf.RestoreV2"(%arg0, %3, %4) : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<*xi32>
// CHECK: %[[NAMES:.*]] = "tf.Const"() <{value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>}>
// CHECK-NEXT: %[[SHAPES:.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
// CHECK-NEXT: %[[TENSORS:.*]]:2 = "tf.RestoreV2"(%[[ARG0]], %[[NAMES]], %[[SHAPES]])
// CHECK-SAME: -> (tensor<*xf32>, tensor<*xi32>)
// CHECK: return %[[TENSORS]]#0, %[[TENSORS]]#1
func.return %2, %5 : tensor<*xf32>, tensor<*xi32>
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_pruning.mlir
+39
View File
@@ -0,0 +1,39 @@
// 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 -tf-restore-pruning %s | FileCheck %s
// CHECK-LABEL: func.func @prune_unused_restore
func.func @prune_unused_restore() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK-NOT: tf.RestoreV2
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
return
}
// CHECK-LABEL: func.func @used_restore_remains
func.func @used_restore_remains() {
%cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
%cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
%cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
// CHECK: tf.RestoreV2
%0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
%1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
"tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
return
}
tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_splitting.mlir
+32
View File
@@ -0,0 +1,32 @@
// 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-restore-splitting | FileCheck %s
// CHECK-LABEL: func @single_restore
// CHECK-SAME: (%[[ARG0:.*]]: {{.*}})
func.func @single_restore(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
%0 = "tf.Const"() {value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%1 = "tf.Const"() {value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%2:2 = "tf.RestoreV2"(%arg0, %0, %1) : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>)
// CHECK: %[[FOO_NAME:.*]] = "tf.Const"() <{value = dense<"foo"> : tensor<1x!tf_type.string>}>
// CHECK: %[[FOO:.*]] = "tf.RestoreV2"(%[[ARG0]], %[[FOO_NAME]], {{.*}})
// CHECK: %[[BAR_NAME:.*]] = "tf.Const"() <{value = dense<"bar"> : tensor<1x!tf_type.string>}>
// CHECK: %[[BAR:.*]] = "tf.RestoreV2"(%[[ARG0]], %[[BAR_NAME]], {{.*}})
// CHECK: return %[[FOO]], %[[BAR]]
func.return %2#0, %2#1 : tensor<*xf32>, tensor<*xi32>
}
tensorflow/compiler/mlir/tfrt/tests/ifrt_set_tpu_host_allocator.mlir
+70
View File
@@ -0,0 +1,70 @@
// 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 -tf-mlrt-ifrt-set-tpu-host-allocator %s | FileCheck %s --dump-input=fail --dump-input-filter=all
// All arguments are non-variables
// CHECK-LABEL: func.func @serving_default
// CHECK-NEXT: "tf.MatMul"
// CHECK-SAME: tf_mlrt.custom_device = "tpu_host_device"
// CHECK-NEXT: "tf.MatMul"
// CHECK-SAME: tf_mlrt.custom_device = "tpu_host_device"
// CHECK-NEXT: "tf.IfrtCall"
func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
%producer_0= "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%producer_1= "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%result = "tf.IfrtCall"(%producer_0, %producer_1) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<1x1xf32>)
return %result: tensor<1x1xf32>
}
// Arguments to the IfrtCall are a mix of non-variables and variables
// CHECK-LABEL: func.func @serving_default1
// CHECK-NEXT: "tf.VarHandleOp"
// CHECK-NOT: tf_mlrt.custom_device
// CHECK-NEXT: "tf.ReadVariableOp"
// CHECK-NOT: tf_mlrt.custom_device
// CHECK-NEXT: "tf.MatMul"
// CHECK-SAME: tf_mlrt.custom_device = "tpu_host_device"
// CHECK-NEXT: "tf.MatMul"
// CHECK-NOT: tf_mlrt.custom_device
// CHECK-NEXT: "tf.IfrtCall"
func.func @serving_default1(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%2 = "tf.MatMul"(%arg1, %arg0) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%3 = "tf.MatMul"(%arg1, %arg0) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%result = "tf.IfrtCall"(%1, %2) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}> : (tensor<3x1xf32>, tensor<1x1xf32>) -> (tensor<1x1xf32>)
return %result: tensor<1x1xf32>
}
// -----
// Async test: All arguments are non-variables
//
// CHECK-LABEL: func.func @serving_default_async
// CHECK-NEXT: "tf.MatMul"
// CHECK-SAME: tf_mlrt.custom_device = "tpu_host_device"
// CHECK-NEXT: "tf.MatMul"
// CHECK-SAME: tf_mlrt.custom_device = "tpu_host_device"
// CHECK-NEXT: "tf.AsyncIfrtCall"
func.func @serving_default_async(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
%producer_0= "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%producer_1= "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%result = "tf.AsyncIfrtCall"(%producer_0, %producer_1) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<1x1xf32>)
return %result: tensor<1x1xf32>
}
tensorflow/compiler/mlir/tfrt/tests/ir/BUILD
+47
View File
@@ -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
+98
View File
@@ -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
Vendored
+22
View File
@@ -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
+63
View File
@@ -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
tensorflow/compiler/mlir/tfrt/tests/mlrt/BUILD
+21
View File
@@ -0,0 +1,21 @@
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "//tensorflow/compiler/mlir:run_lit.sh",
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//mlir:run_lit.sh",
],
)
tensorflow/compiler/mlir/tfrt/tests/mlrt/assign_op_key.mlir
+63
View File
@@ -0,0 +1,63 @@
// 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 -split-input-file -tf-mlrt-assign-op-key %s | FileCheck %s
// CHECK-LABEL: func @main
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 0 : i32}
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 1 : i32}
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 2 : i32}
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 3 : i32}
// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 4 : i32}
// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 5 : i32}
// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 6 : i32}
// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 7 : i32}
// CHECK: [[x:%.*]] = "tf.AddV2"
// CHECK-SAME: {__op_key = 8 : i32}
// CHECK: return [[x]]
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
%a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b0 = "tf.Sub"(%b, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %c : tensor<i32>
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/async_while.mlir
+244
View File
@@ -0,0 +1,244 @@
// 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 -split-input-file -tf-mlrt-async-while %s | FileCheck %s
// This is a simple case that should be pipelined.
// CHECK-LABEL: func.func private @"map/while_cond"
func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
%0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
return %0 : tensor<i1>
}
// CHECK-LABEL: func.func private @"map/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
// CHECK-NEXT: %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: return %0 : tensor<i1>
// CHECK-LABEL: func.func private @"map/while_body"
func.func private @"map/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) {
%cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
%2 = "tf.MatMul"(%1, %matrix) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
return %0, %max_iterations, %handle, %flow_in, %2: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>
}
// CHECK-LABEL: func.func private @"map/while_body/TfMlrtAsyncWhileBody"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor<i32>, %arg6: tensor<?x!tf_type.resource>, %arg7: tensor<*xf32>) {
// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>
// CHECK-NEXT: %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor<i1>) -> ()
// CHECK-NEXT: %3 = "tf.TensorArrayReadV3"(%arg6, %0, %arg7) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: %5 = "tf.MatMul"(%3, %4) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg4, %5) : (!mlrt.promise, tensor<3x3xf32>) -> ()
// CHECK-NEXT: return
//CHECK-LABEL: func.func @serving_default
func.func @serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<*xf32>) {
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
%1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
// CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#5) : (!mlrt.future) -> tensor<*xf32>
// CHECK-NEXT: %3 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: return %3, %2 : tensor<3x3xf32>, tensor<*xf32>
return %1#4, %1#3 : tensor<3x3xf32>, tensor<*xf32>
}
//CHECK-LABEL: func.func @multi_while_test
func.func @multi_while_test(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>, %array_handle_2: tensor<?x!tf_type.resource>, %array_flow_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<3x3xf32>) {
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
%1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
// CHECK: %2 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %3:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%2, %cst, %arg6, %arg0, %arg4, %arg5) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
%2:5 = "tf.While"(%cst_0, %max_iterations, %array_handle_2, %array_flow_2, %matrix_2) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
// CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: %5 = "tf_mlrt.tf_await"(%3#2) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: return %4, %5 : tensor<3x3xf32>, tensor<3x3xf32>
return %1#4, %2#4 : tensor<3x3xf32>, tensor<3x3xf32>
}
// -----
// Test a case in which predicate is updated after mutables and shall not be converted to AsyncWhile.
// CHECK-LABEL: func.func private @"map/while_cond"
func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
%0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
return %0 : tensor<i1>
}
// CHECK-LABEL: func.func private @"map/while_body"
func.func private @"map/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) {
%cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
%2 = "tf.MatMul"(%1, %matrix) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// Predicate is update at the last stage.
%0 = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
return %0, %max_iterations, %handle, %flow_in, %2: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>
}
//CHECK-LABEL: func.func @serving_default
func.func @serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<3x3xf32> {
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: tf.While
// CHECK-NOT: AsyncWhile
%1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
return %1#4 : tensor<3x3xf32>
}
// -----
// The newly create function name may have conflict with existing functions (very rare).
// CHECK-LABEL: func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
return %0: tensor<i32>
}
// CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
return %0: tensor<i32>
}
// CHECK-LABEL: func.func private @"random/while_cond"
func.func private @"random/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
%0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
return %0 : tensor<i1>
}
// CHECK-LABEL: func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate_0"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
// CHECK-NEXT: %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: return %0 : tensor<i1>
// CHECK-LABEL: func.func private @"random/while_body"
func.func private @"random/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) {
%cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
%2 = "tf.MatMul"(%1, %matrix) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
return %0, %max_iterations, %handle, %flow_in, %2: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>
}
// CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody_1"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor<i32>, %arg6: tensor<?x!tf_type.resource>, %arg7: tensor<*xf32>) {
// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>
// CHECK-NEXT: %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor<i1>) -> ()
// CHECK-NEXT: %3 = "tf.TensorArrayReadV3"(%arg6, %0, %arg7) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: %5 = "tf.MatMul"(%3, %4) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg4, %5) : (!mlrt.promise, tensor<3x3xf32>) -> ()
// CHECK-NEXT: return
//CHECK-LABEL: func.func @random_serving_default
func.func @random_serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<*xf32>) {
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"random/while_body/TfMlrtAsyncWhileBody_1"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
%1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"random/while_body", cond = @"random/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
// CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#5) : (!mlrt.future) -> tensor<*xf32>
// CHECK-NEXT: %3 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: return %3, %2 : tensor<3x3xf32>, tensor<*xf32>
return %1#4, %1#3 : tensor<3x3xf32>, tensor<*xf32>
}
// -----
// This case test the re-ordering of the while body function to maximize pipelining between iterations.
// CHECK-LABEL: func.func private @"sort_map/while_cond"
func.func private @"sort_map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>, %handle_2: tensor<?x!tf_type.resource>, %flow_in_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>, %bound: tensor<i32>) -> tensor<i1> {
%0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
return %0 : tensor<i1>
}
// CHECK-LABEL: func.func private @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
// CHECK-NEXT: %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: return %0 : tensor<i1>
// CHECK-LABEL: func.func private @"sort_map/while_body"
// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %0 = "tf.AddV2"(%arg0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: %1 = "tf.TensorArrayReadV3"(%arg2, %arg0, %arg3) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %2 = "tf.TensorArrayReadV3"(%arg5, %arg0, %arg6) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %3 = "tf.MatMul"(%1, %arg4) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %4 = "tf.GreaterEqual"(%0, %arg8) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %5 = "tf.MatMul"(%3, %arg7) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %6 = "tf.Select"(%4, %3, %arg4) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %7 = "tf.Identity"(%4) : (tensor<i1>) -> tensor<i1>
// CHECK-NEXT: %8 = "tf.MatMul"(%5, %arg7) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %9 = "tf.Select"(%7, %8, %arg7) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: return %0, %arg1, %arg2, %arg3, %6, %arg5, %arg6, %9, %arg8
func.func private @"sort_map/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>, %handle_2: tensor<?x!tf_type.resource>, %flow_in_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>, %bound: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<i32>) {
%cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%updated_loop_count = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%in_matrix = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
%out_matrix = "tf.MatMul"(%in_matrix, %matrix) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%in_matrix1 = "tf.TensorArrayReadV3"(%handle_2, %loop_count, %flow_in_2) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
%out_matrix1 = "tf.MatMul"(%out_matrix, %matrix_2) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%out_matrix2 = "tf.MatMul"(%out_matrix1, %matrix_2): (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%cond = "tf.GreaterEqual"(%updated_loop_count, %bound) : (tensor<i32>, tensor<i32>) -> tensor<i1>
%result = "tf.Select"(%cond, %out_matrix, %matrix) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%cond2 = "tf.Identity"(%cond) : (tensor<i1>) -> tensor<i1>
%result2 = "tf.Select"(%cond2, %out_matrix2, %matrix_2) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
return %updated_loop_count, %max_iterations, %handle, %flow_in, %result, %handle_2, %flow_in_2, %result2, %bound: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<i32>
}
// CHECK-LABEL: func.func private @"sort_map/while_body/TfMlrtAsyncWhileBody"
// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>
// CHECK-NEXT: %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg7) <{config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor<i1>) -> ()
// CHECK-NEXT: %3 = "tf.TensorArrayReadV3"(%arg8, %0, %arg9) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %4 = "tf.TensorArrayReadV3"(%arg10, %0, %arg11) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %5 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: %6 = "tf.MatMul"(%3, %5) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %7 = "tf.GreaterEqual"(%1, %arg12) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %8 = "tf_mlrt.tf_await"(%arg5) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: %9 = "tf.MatMul"(%6, %8) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %10 = "tf.Select"(%7, %6, %5) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg4, %10) : (!mlrt.promise, tensor<3x3xf32>) -> ()
// CHECK-NEXT: %11 = "tf.Identity"(%7) : (tensor<i1>) -> tensor<i1>
// CHECK-NEXT: %12 = "tf.MatMul"(%9, %8) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: %13 = "tf.Select"(%11, %12, %8) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg6, %13) : (!mlrt.promise, tensor<3x3xf32>) -> ()
// CHECK-NEXT: return
//CHECK-LABEL: func.func @sort_serving_default
func.func @sort_serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>, %bound: tensor<i32>) -> (tensor<3x3xf32>, tensor<3x3xf32>, tensor<*xf32>) {
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT: %1:10 = tf_mlrt.tf_async_while @"sort_map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg3, %arg0, %arg1, %arg2, %arg1, %arg2, %arg4) {invariant_size = 6 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<i32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
%1:9 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix , %array_handle, %array_flow, %matrix, %bound) {body= @"sort_map/while_body", cond = @"sort_map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>,tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<i32>)
// CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#6) : (!mlrt.future) -> tensor<*xf32>
// CHECK-NEXT: %3 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%1#3) : (!mlrt.future) -> tensor<3x3xf32>
// CHECK-NEXT: return %3, %4, %2 :
return %1#4, %1#7, %1#3 : tensor<3x3xf32>, tensor<3x3xf32>, tensor<*xf32>
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/fuse_mlrt_ops.mlir
+92
View File
@@ -0,0 +1,92 @@
// 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 -split-input-file -tf-mlrt-fuse %s | FileCheck %s
// CHECK-LABEL: @main
// CHECK-SAME: ([[f0:%.*]]: !mlrt.future, [[f1:%.*]]: !mlrt.future, [[f2:%.*]]: !mlrt.future)
func.func @main(%f0: !mlrt.future, %f1: !mlrt.future, %f2: !mlrt.future) -> (!tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor) {
// CHECK-NEXT: [[t:%.*]]:3 = tf_mlrt.await_all [[f0]], [[f1]], [[f2]]
// CHECK-NOT: tf_mlrt.await
// CHECK-NEXT: return [[t]]#0, [[t]]#1, [[t]]#2
%t0 = tf_mlrt.await %f0
%t1 = tf_mlrt.await %f1
%t2 = tf_mlrt.await %f2
func.return %t0, %t1, %t2 : !tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor
}
// -----
// CHECK-LABEL: @main
// CHECK-SAME: ([[f0:%.*]]: !mlrt.future, [[f1:%.*]]: !mlrt.future, [[f2:%.*]]: !mlrt.future)
func.func @main(%f0: !mlrt.future, %f1: !mlrt.future, %f2: !mlrt.future) -> (!tf_mlrt.tensor, !tf_mlrt.tensor) {
// CHECK-NEXT: [[t:%.*]]:2 = tf_mlrt.await_all [[f0]], [[f1]]
// CHECK-NOT: tf_mlrt.await
// CHECK-NEXT: [[t2:%.*]] = tf_mlrt.executeop([[t]]#0, [[t]]#1)
// CHECK-NEXT: [[t3:%.*]] = tf_mlrt.await [[f2]]
// CHECK-NEXT: return [[t2]], [[t3]]
%t0 = tf_mlrt.await %f0
%t1 = tf_mlrt.await %f1
%t2 = tf_mlrt.executeop(%t0, %t1) {node_def = "AddV2", op_key = 0 : i32} : (!tf_mlrt.tensor, !tf_mlrt.tensor) -> (!tf_mlrt.tensor)
%t3 = tf_mlrt.await %f2
func.return %t2, %t3 : !tf_mlrt.tensor, !tf_mlrt.tensor
}
// -----
// CHECK-LABEL: @main
// CHECK-SAME: ([[f0:%.*]]: !mlrt.async_handle, [[f1:%.*]]: !mlrt.async_handle, [[f2:%.*]]: !mlrt.async_handle)
func.func @main(%f0: !mlrt.async_handle, %f1: !mlrt.async_handle, %f2: !mlrt.async_handle) -> () {
// CHECK-NEXT: mlrt.await_all_handle [[f0]], [[f1]], [[f2]]
// CHECK-NOT: mlrt.await_handle
// CHECK-NEXT: return
mlrt.await_handle %f0
mlrt.await_handle %f1
mlrt.await_handle %f2
func.return
}
// -----
// CHECK-LABEL: @main
func.func @main() -> (!tf_mlrt.tensor, !tf_mlrt.tensor) {
// CHECK-NEXT: [[r:%.*]]:3 = tf_mlrt.get_resource {indices = [2, 0, 1]}
// CHECK-NEXT: [[v:%.*]] = tf_mlrt.executeop([[r]]#0, [[r]]#1)
// CHECK-NEXT: return [[v]], [[r]]#2
%0 = tf_mlrt.get_resource {indices = [2]} : !tf_mlrt.tensor
%1 = tf_mlrt.get_resource {indices = [0]} : !tf_mlrt.tensor
%r = tf_mlrt.executeop(%0, %1) {node_def = "AddV2", op_key = 0 : i32} : (!tf_mlrt.tensor, !tf_mlrt.tensor) -> (!tf_mlrt.tensor)
%2 = tf_mlrt.get_resource {indices = [1]} : !tf_mlrt.tensor
func.return %r, %2 : !tf_mlrt.tensor, !tf_mlrt.tensor
}
// -----
// CHECK-LABEL: @fuse_promise_return
// CHECK-SAME: ([[p:%.*]]: !mlrt.promise, [[v:%.*]]: !tf_mlrt.tensor)
func.func @fuse_promise_return(%p: !mlrt.promise, %v: !tf_mlrt.tensor) -> () {
// CHECK: tf_mlrt.promise_return [[p]], [[v]]
tf_mlrt.promise %p, %v
func.return
}
// -----
// CHECK-LABEL: @not_fuse_promise_return
// CHECK-SAME: ([[p:%.*]]: !mlrt.promise, [[v:%.*]]: !tf_mlrt.tensor)
func.func @not_fuse_promise_return(%p: !mlrt.promise, %v: !tf_mlrt.tensor) -> (!tf_mlrt.tensor) {
// CHECK-NOT: tf_mlrt.promise_return
tf_mlrt.promise %p, %v
func.return %v : !tf_mlrt.tensor
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/inline.mlir
+64
View File
@@ -0,0 +1,64 @@
// 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 -split-input-file -pass-pipeline='builtin.module(tf-to-mlrt, inline)' %s | FileCheck %s -dump-input=fail
// Test generated tf_mlrt while body and predicate is inlined.
func.func @then(%x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> tensor<i1> {
return %x: tensor<i1>
}
func.func @else(%x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> tensor<i1> {
return %y: tensor<i1>
}
// CHECK-LABEL: func @while_cond_if
// CHECK: [[cond:%.*]] = tf_mlrt.predicate
// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else
// CHECK: return [[z]]
func.func @while_cond_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>) {
%r = "tf.If"(%cond, %x, %y, %z) {then_branch = @then, else_branch = @else, is_stateless = true} : (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) -> tensor<i1>
return %r : tensor<i1>
}
// CHECK-LABEL: func @while_body_if
func.func @while_body_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) {
%0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%z, %0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %cond, %x, %y, %1 : tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>
}
// CHECK-LABEL: func @while_test_if
// CHECK-SAME: -> !tf_mlrt.tensor
func.func @while_test_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>) -> (tensor<i32>) {
// CHECK: [[CONST:%.*]] = tf_mlrt.constop {tensor_proto = "\08\03\12\00"}
%cst = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// Predicate should be inlined.
// CHECK-NEXT: tf_mlrt.predicate
// CHECK-NEXT: mlrt.cond
// CHECK-NEXT: tf_mlrt.predicate
// CHECK-NEXT: mlrt.while
%0:4 = "tf.While"(%cond, %x, %y, %cst) { cond = @while_cond_if, body = @while_body_if, is_stateless = false, parallel_iterations = 1} : (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>)
// CHECK: return
// CHECK-SAME: !tf_mlrt.tensor
func.return %0#3 : tensor<i32>
}
// CHECK-LABEL: func @"while_body_if/tf_mlrt_body"
// CHECK-NOT: call
// CHECK-LABEL: func @"while_cond_if/tf_mlrt_predicate"
// CHECK-NOT: call
tensorflow/compiler/mlir/tfrt/tests/mlrt/parallelization.mlir
+392
View File
@@ -0,0 +1,392 @@
// 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 -split-input-file -tf-mlrt-parallelization %s | FileCheck %s --dump-input=fail --dump-input-filter=all
// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
// CHECK-SAME: ({{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: [[RES:%.*]] = "tf.Sub"
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
// CHECK: return
// CHECK-LABEL: func @main
// CHECK: [[PROMISE:%.*]], [[FUTURE:%.*]] = "tf_mlrt.allocate_futures"
// CHECK: [[HANDLE:%.*]] = mlrt.async({{%.*}}, [[PROMISE]])
// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: [[x:%.*]] = "tf.AddV2"
// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
// CHECK: mlrt.await_handle [[HANDLE]]
// CHECK: return [[RES]]
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
%a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b0 = "tf.Sub"(%b, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %c : tensor<i32>
}
// -----
// Test merging child streams
// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
// CHECK-SAME: ({{%.*}}: tensor<i32>, {{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: [[RES:%.*]] = "tf.Sub"
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
// CHECK: return
// CHECK-LABEL: func @main
// CHECK: [[PROMISE:%.*]], [[FUTURE:%.*]] = "tf_mlrt.allocate_futures"
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: [[VALUE:%.*]] = "tf.AddV2"
// CHECK: [[HANDLE:%.*]] = mlrt.async([[VALUE]], {{%.*}}, [[PROMISE]])
// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: [[x:%.*]] = "tf.AddV2"
// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
// CHECK: mlrt.await_handle [[HANDLE]]
// CHECK: return [[RES]]
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
%a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a4 = "tf.AddV2"(%a3, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a5 = "tf.AddV2"(%a4, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a6 = "tf.AddV2"(%a5, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a7 = "tf.AddV2"(%a6, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b0 = "tf.Sub"(%a3, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c = "tf.AddV2"(%a7, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %c : tensor<i32>
}
// -----
// Test side-effecting ops
// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
// CHECK-SAME: ([[ARG:%.*]]: tensor<i32>, [[FUTURE:%.*]]: !mlrt.future, [[CONTROL_PROMISE:%.*]]: !mlrt.promise)
// CHECK: [[HANDLE:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
// CHECK: "tf.AssignVariableOp"([[HANDLE]], [[ARG]])
// CHECK-NEXT: mlrt.promise_control [[CONTROL_PROMISE]]
// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
// CHECK-SAME: ({{%.*}}: tensor<i32>, {{%.*}}: tensor<i32>, [[FUTURE:%.*]]: !mlrt.future, [[PROMISE:%.*]]: !mlrt.promise)
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: [[V:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
// CHECK-NEXT: [[RES:%.*]] = "tf.Sub"({{%.*}}, [[V]])
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
// CHECK: return
// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
// CHECK-SAME: ([[CONTROL_FUTURE:%.*]]: !mlrt.future, [[PROMISE:%.*]]: !mlrt.promise, [[PROMISE_HANDLE:%.*]]: !mlrt.promise)
// CHECK: [[HANDLE:%.*]] = "tf.VarHandleOp"
// CHECK-NEXT: "tf_mlrt.tf_promise"([[PROMISE_HANDLE]], [[HANDLE]])
// CHECK: mlrt.await_control [[CONTROL_FUTURE]]
// CHECK-NEXT: [[V:%.*]] = "tf.ReadVariableOp"([[HANDLE]])
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[V]])
// CHECK-LABEL: func @main
// CHECK: [[PROMISE:%.*]]:3, [[FUTURE:%.*]]:3 = "tf_mlrt.allocate_futures"
// CHECK: [[CONTROL_PROMISE:%.*]], [[CONTROL_FUTURE:%.*]] = "mlrt.allocate_control_futures"
// CHECK: [[ASYNC_HANDLE_0:%.*]] = mlrt.async([[CONTROL_FUTURE]], [[PROMISE]]#0, [[PROMISE]]#1)
// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
// CHECK: [[ASYNC_HANDLE_1:%.*]] = mlrt.async({{%.*}}, {{%.*}}, [[FUTURE]]#0, [[PROMISE]]#2)
// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: [[x:%.*]] = "tf.AddV2"
// CHECK: [[ASYNC_HANDLE_2:%.*]] = mlrt.async([[x]], [[FUTURE]]#1, [[CONTROL_PROMISE]])
// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]]#2)
// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
// CHECK: mlrt.await_handle [[ASYNC_HANDLE_0]]
// CHECK-NEXT: mlrt.await_handle [[ASYNC_HANDLE_1]]
// CHECK-NEXT: mlrt.await_handle [[ASYNC_HANDLE_2]]
// CHECK-NEXT: return [[RES]]
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
%handle = "tf.VarHandleOp"() {container = "", shared_name = "var"} : () -> tensor<!tf_type.resource_handle<tensor<i32>>>
%a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"tf.AssignVariableOp"(%handle, %a3) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
%b0 = "tf.Sub"(%a, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%var = "tf.ReadVariableOp"(%handle) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%b3 = "tf.Sub"(%b2, %var) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %c : tensor<i32>
}
// -----
// Test multiple promises and futures
// CHECK-LABEL: func private @main_stream_1
// CHECK: mlrt.await_control
// CHECK: "tf.DummySideEffecting"() {id = 4
// CHECK: return
// CHECK-LABEL: func private @main_stream_2
// CHECK: mlrt.await_control
// CHECK: "tf.DummySideEffecting"() {id = 3
// CHECK: mlrt.promise_control
// CHECK: return
// CHECK-LABEL: func private @main_stream_3
// CHECK: mlrt.await_control
// CHECK: "tf.DummySideEffecting"() {id = 2
// CHECK: mlrt.promise_control
// CHECK: return
// CHECK-LABEL: func private @main_stream_4
// CHECK: "tf.DummySideEffecting"() {id = 1
// CHECK: mlrt.promise_control
// CHECK: return
// CHECK-LABEL: func @main()
// CHECK: [[PROMISES:%.*]]:3, [[FUTURES:%.*]]:3 = "mlrt.allocate_control_futures"
// CHECK: mlrt.async([[PROMISES]]#2) {callee = @main_stream_4
// CHECK: mlrt.async([[FUTURES]]#2, [[PROMISES]]#1) {callee = @main_stream_3
// CHECK: mlrt.async([[FUTURES]]#1, [[PROMISES]]#0) {callee = @main_stream_2
// CHECK: mlrt.async([[FUTURES]]#0) {callee = @main_stream_1
// CHECK: mlrt.await_handle
// CHECK: mlrt.await_handle
// CHECK: mlrt.await_handle
// CHECK: mlrt.await_handle
func.func @main() {
"tf.DummySideEffecting"() {id = 1} : () -> ()
"tf.DummySideEffecting"() {id = 2} : () -> ()
"tf.DummySideEffecting"() {id = 3} : () -> ()
"tf.DummySideEffecting"() {id = 4} : () -> ()
func.return
}
// -----
// Test correctness when there are both data and control promises in a stream function.
// CHECK-LABEL: func private @main_stream_1
// CHECK-SAME: ([[PROMISE:%.*]]: !mlrt.promise, [[CONTROL_PROMISE:%.*]]: !mlrt.promise)
// CHECK: tf.DummySideEffecting
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]]
// CHECK: mlrt.promise_control [[CONTROL_PROMISE]]
func.func @main() -> tensor<i32> {
%v = "tf.DummySideEffecting"() {id = 1} : () -> tensor<i32>
%w = "tf.DummySideEffecting"() {id = 2} : () -> tensor<i32>
%r = "tf.AddV2"(%w, %v) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
// -----
// Test inputs to the child streams are merged to the parent streams
// CHECK-LABEL: func private @main_stream_1
// CHECK-SAME: ([[INPUT0:%.*]]: tensor<i32>, [[INPUT1:%.*]]: tensor<i32>
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: mlrt.async({{%.*}}, [[INPUT1]]
// CHECK-LABEL: func @main
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
%a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b0 = "tf.Sub"(%b, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c = "tf.AddV2"(%b1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%d = "tf.AddN"(%a3, %b3, %c) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %d : tensor<i32>
}
// -----
// Test that constants are copied instead of using promise/await.
// CHECK-LABEL: func private @main_stream_1
// CHECK-SAME: ({{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
// CHECK: tf._TfrtGetResource
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: [[RES:%.*]] = "tf.Sub"
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
// CHECK: return
// CHECK-NOT: func private @main_stream
// CHECK-LABEL: func @main
// CHECK: [[PROMISE:%.*]], [[FUTURE:%.*]] = "tf_mlrt.allocate_futures"
// CHECK-NEXT: [[HANDLE:%.*]] = mlrt.async({{%.*}}, [[PROMISE]])
// CHECK-SAME: callee = @main_stream_1
// CHECK: tf._TfrtGetResource
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: tf.AddV2
// CHECK: [[x:%.*]] = "tf.AddV2"
// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
// CHECK: mlrt.await_handle [[HANDLE]]
// CHECK: return [[RES]]
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
%c0 = "tf._TfrtGetResource"() {indices = [0], shared_name = [""], container = [""]} : () -> (tensor<i32>)
%a0 = "tf.AddV2"(%a, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a1 = "tf.AddV2"(%a0, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a2 = "tf.AddV2"(%a1, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%a3 = "tf.AddV2"(%a2, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b0 = "tf.Sub"(%b, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b1 = "tf.Sub"(%b0, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b2 = "tf.Sub"(%b1, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%b3 = "tf.Sub"(%b2, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %c : tensor<i32>
}
// -----
// Test that constants private to a stream are still handled properly when we are copying shared constants.
// CHECK-LABEL: func private @main_stream_1
// CHECK: [[r:%.*]] = "tf._TfrtGetResource"
// CHECK-SAME: indices = [1]
// CHECK: "tf.DummySideEffecting"([[r]])
// CHECK-LABEL: func private @main_stream_2
// CHECK: [[r:%.*]] = "tf._TfrtGetResource"
// CHECK-SAME: indices = [0]
// CHECK: "tf.DummySideEffecting"([[r]])
// CHECK-LABEL: func @main
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> () {
%c0 = "tf._TfrtGetResource"() {indices = [0], shared_name = [""], container = [""]} : () -> (tensor<i32>)
"tf.DummySideEffecting"(%c0) : (tensor<i32>) -> ()
%c1 = "tf._TfrtGetResource"() {indices = [1], shared_name = [""], container = [""]} : () -> (tensor<i32>)
"tf.DummySideEffecting"(%c1) : (tensor<i32>) -> ()
func.return
}
// -----
// Test that streams with no args but side-effecting ops are still created properly
// CHECK-LABEL: func private @main_stream_1()
// CHECK: [[r:%.*]] = "tf._TfrtGetResource"
// CHECK-SAME: indices = [0]
// CHECK: "tf.DummySideEffecting"([[r]])
// CHECK-LABEL: func @main
func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> () {
%c0 = "tf._TfrtGetResource"() {indices = [0], shared_name = [""], container = [""]} : () -> (tensor<i32>)
"tf.DummySideEffecting"(%c0) : (tensor<i32>) -> ()
func.return
}
// -----
// Test control deps of tf.Assert is skipped.
// CHECK-LABEL: func.func private @skip_assert_stream_3(
// CHECK-NOT: mlrt.await_control
// CHECK: tf.Assert
// CHECK-NOT: mlrt.promise_control
// CHECK: return
// CHECK-LABEL: func.func private @skip_assert_stream_2(
// CHECK-NOT: mlrt.await_control
// CHECK: tf.Assert
// CHECK-NOT: mlrt.promise_control
// CHECK: return
func.func @skip_assert(%key: tensor<!tf_type.string>) -> (tensor<i64>, tensor<i64>) {
%error_message = "tf.Const"() {value = dense<"error"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%default = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
%handle = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
%keys = "tf.Const"() {value = dense<["a", "b", "c", "d"]> : tensor<4x!tf_type.string>} : () -> tensor<4x!tf_type.string>
%values = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi64>} : () -> tensor<4xi64>
"tf.LookupTableImportV2"(%handle, %keys, %values) {device = ""} : (tensor<!tf_type.resource>, tensor<4x!tf_type.string>, tensor<4xi64>) -> ()
%value0 = "tf.LookupTableFindV2"(%handle, %key, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
%cond = "tf.Equal"(%value0, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0", incompatible_shape_error = true} : (tensor<i64>, tensor<i64>) -> tensor<i1>
"tf.Assert"(%cond, %error_message) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
"tf.Assert"(%cond, %error_message) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
%value1 = "tf.LookupTableFindV2"(%handle, %key, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
func.return %value0, %value1 : tensor<i64>, tensor<i64>
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/rewrite_ifrt_load_variable.mlir
+57
View File
@@ -0,0 +1,57 @@
// 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 -split-input-file -tf-mlrt-rewrite-ifrt-load-variable %s | FileCheck %s
// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT: [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT: [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])
// CHECK-SAME: <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
// CHECK-NEXT: [[TENSOR:%.*]] = "tf_mlrt.tf_await"([[FURTURE]]) : (!mlrt.future) -> tensor<3x1xf32>
// CHECK-NEXT: "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT: "tf.IfrtCall"(%arg0, [[ARRAYKEY]]) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
// CHECK-NEXT: return
//
func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%array_key, %tensor = "tf.IfrtLoadVariable"(%0) <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, tensor<3x1xf32>)
%1 = "tf.MatMul"(%arg0, %tensor) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
%2 = "tf.IfrtCall"(%arg0, %array_key) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
return %2 : tensor<1x1xf32>
}
// -----
// Variable is used by two CPU ops
//
// CHECK-LABEL: func @serving_default
// CHECK-NEXT: [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT: [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])
// CHECK-SAME: <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
// CHECK: [[TENSOR:%.*]] = "tf_mlrt.tf_await"([[FURTURE]]) : (!mlrt.future) -> tensor<3x1xf32>
// CHECK-NEXT: "tf.AddV2"([[TENSOR]], %cst) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
// CHECK-NEXT: "tf.Sub"([[TENSOR]], %cst) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
// CHECK-NEXT: return
//
func.func @serving_default() {
%0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%array_key, %tensor = "tf.IfrtLoadVariable"(%0) <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, tensor<3x1xf32>)
%cst_24 = "tf.Const"() <{value = dense<[[0.0], [1.0], [2.0]]> : tensor<3x1xf32>}> : () -> tensor<3x1xf32>
%1 = "tf.AddV2"(%tensor, %cst_24) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
%2 = "tf.Sub"(%tensor, %cst_24) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
return
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/tf_to_mlrt.mlir
+616
View File
@@ -0,0 +1,616 @@
// 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 -split-input-file -tf-to-mlrt %s | FileCheck %s
// CHECK-LABEL: @main_stream_0
// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[promise_b:%.*]]: !mlrt.promise)
func.func @main_stream_0(%input0: tensor<i32>, %promise_b: !mlrt.promise) {
%const = "tf.Const"() {__op_key = 0 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[a:%.*]] = tf_mlrt.executeop([[input0]],
// CHECK-SAME: AddV2
%a = "tf.AddV2"(%input0, %const) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: [[b:%.*]] = tf_mlrt.executeop([[a]])
// CHECK-SAME: Abs
%b = "tf.Abs"(%a) {__op_key = 2 : i32}: (tensor<i32>) -> tensor<i32>
// CHECK: tf_mlrt.promise [[promise_b]], [[b]]
"tf_mlrt.tf_promise"(%promise_b, %b) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK: return
return
}
// CHECK-LABEL: @main_stream_1
// CHECK-SAME: ([[input1:%.*]]: !tf_mlrt.tensor, [[promise_c:%.*]]: !mlrt.promise, [[promise_d:%.*]]: !mlrt.promise)
func.func @main_stream_1(%input1: tensor<i32>, %promise_c: !mlrt.promise, %promise_d: !mlrt.promise) {
%const = "tf.Const"() {__op_key = 3 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[c:%.*]] = tf_mlrt.executeop([[input1]],
// CHECK-SAME: Sub
%c = "tf.Sub"(%input1, %const) {__op_key = 4: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: tf_mlrt.promise [[promise_c]], [[c]]
"tf_mlrt.tf_promise"(%promise_c, %c) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK: [[d:%.*]] = tf_mlrt.executeop([[c]])
// CHECK-SAME: Abs
%d = "tf.Abs"(%c) {__op_key = 5: i32}: (tensor<i32>) -> tensor<i32>
// CHECK: tf_mlrt.promise [[promise_d]], [[d]]
"tf_mlrt.tf_promise"(%promise_d, %d) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK: return
return
}
// CHECK-LABEL: @main
// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[input1:%.*]]: !tf_mlrt.tensor)
func.func @main(%input0: tensor<i32>, %input1: tensor<i32>) -> tensor<i32> {
// CHECK: [[promises:%.*]]:3, [[futures:%.*]]:3 = "tf_mlrt.allocate_futures"
// CHECK-SAME: num_futures = 3
%promise_b, %promise_c, %promise_d, %future_b, %future_c, %future_d =
"tf_mlrt.allocate_futures"()
{num_futures = 3 : i32, resultSegmentSizes = array<i32: 3, 3>} : () ->
(!mlrt.promise, !mlrt.promise, !mlrt.promise,
!mlrt.future, !mlrt.future, !mlrt.future)
// CHECK: [[handle_0:%.*]] = mlrt.async([[input0]], [[promises]]#0)
// CHECK-SAME: callee = @main_stream_0
%handle_0 = mlrt.async(%input0, %promise_b)
{callee = @main_stream_0} :
(tensor<i32>, !mlrt.promise) -> !mlrt.async_handle
// CHECK: [[handle_1:%.*]] = mlrt.async([[input1]], [[promises]]#1, [[promises]]#2)
// CHECK-SAME: callee = @main_stream_1
%handle_1 = mlrt.async(%input1, %promise_c, %promise_d)
{callee = @main_stream_1} :
(tensor<i32>, !mlrt.promise, !mlrt.promise) -> !mlrt.async_handle
%const = "tf.Const"() {__op_key = 6: i32, value = dense<2> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[e:%.*]] = tf_mlrt.executeop([[input1]],
// CHECK-SAME: Mul
%e = "tf.Mul"(%input1, %const) {__op_key = 7: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: [[c:%.*]] = tf_mlrt.await [[futures]]#1
%c = "tf_mlrt.tf_await"(%future_c) : (!mlrt.future) ->tensor<i32>
// CHECK: [[f:%.*]] = tf_mlrt.executeop([[e]], [[c]])
// CHECK-SAME: Div
%f = "tf.Div"(%e, %c) {__op_key = 8: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: [[b:%.*]] = tf_mlrt.await [[futures]]#0
%b = "tf_mlrt.tf_await"(%future_b) : (!mlrt.future) ->tensor<i32>
// CHECK: [[d:%.*]] = tf_mlrt.await [[futures]]#2
%d = "tf_mlrt.tf_await"(%future_d) : (!mlrt.future) ->tensor<i32>
// CHECK: [[result:%.*]] = tf_mlrt.executeop([[b]], [[d]], [[f]])
// CHECK-SAME: AddN
%result = "tf.AddN"(%b, %d, %f) {__op_key = 9: i32}: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: mlrt.await_handle [[handle_0]]
// CHECK: mlrt.await_handle [[handle_1]]
mlrt.await_handle %handle_0
mlrt.await_handle %handle_1
// CHECK: return [[result]]
return %result : tensor<i32>
}
// -----
// Test lowering tf.If
func.func @then(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
return %x: tensor<i32>
}
func.func @else(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
return %y: tensor<i32>
}
// CHECK-LABEL: func @main
// CHECK-SAME: ([[cond_tensor:%.*]]: !tf_mlrt.tensor, [[x:%.*]]: !tf_mlrt.tensor, [[y:%.*]]: !tf_mlrt.tensor)
// CHECK: [[cond:%.*]] = tf_mlrt.predicate [[cond_tensor]]
// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else([[x]], [[y]])
// CHECK: return [[z]]
func.func @main(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
%z = "tf.If"(%cond, %x, %y) {then_branch = @then, else_branch = @else, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
return %z: tensor<i32>
}
// -----
// Test lowering AsyncOpKernel
// CHECK-LABEL: func @main
func.func @main(%x: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
// CHECK: [[y_future:%.*]] = tf_mlrt.async_executeop
%y = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
// CHECK: [[z:%.*]] = tf_mlrt.executeop
%z = "tf.Identity"(%x) {__op_key = 1: i32}: (tensor<i32>) -> tensor<i32>
// CHECK: [[y:%.*]] = tf_mlrt.await [[y_future]]
// CHECK-NEXT: tf_mlrt.executeop([[y]]
%w = "tf.AddV2"(%y, %z) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: tf_mlrt.executeop([[y]]
%u = "tf.AddV2"(%y, %z) {__op_key = 3: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: tf_mlrt.executeop([[y]]
%v = "tf.AddV2"(%y, %z) {__op_key = 4: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
return %w, %u, %v : tensor<i32>, tensor<i32>, tensor<i32>
}
// -----
// Test lowering BatchFunction op.
func.func @batched_function(%x: tensor<?xi32>) -> tensor<?xi32> {
return %x : tensor<?xi32>
}
// CHECK-LABEL: func @main
func.func @main(%x: tensor<1xi32>) -> (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) {
// CHECK: [[y_future:%.*]] = tf_mlrt.batch_function
// CHECK-SAME: f = @batched_function
// CHECK-SAME: \22batch_function\22
%y = "tf.BatchFunction"(%x) {
allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64,
batching_queue = "", container = "", device = "/device:CPU:0",
enable_large_batch_splitting = false, f = @batched_function,
max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64,
num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>,
shared_name = "batch_function"
} : (tensor<1xi32>) -> tensor<1xi32>
// CHECK: [[z:%.*]] = tf_mlrt.executeop
%z = "tf.Identity"(%x) {__op_key = 0: i32} : (tensor<1xi32>) -> tensor<1xi32>
// CHECK: [[y:%.*]] = tf_mlrt.await [[y_future]]
// CHECK-NEXT: tf_mlrt.executeop([[y]]
%w = "tf.AddV2"(%y, %z) {__op_key = 1: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
// CHECK-NEXT: tf_mlrt.executeop([[y]]
%u = "tf.AddV2"(%y, %z) {__op_key = 2: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
// CHECK-NEXT: tf_mlrt.executeop([[y]]
%v = "tf.AddV2"(%y, %z) {__op_key = 3: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
return %w, %u, %v : tensor<1xi32>, tensor<1xi32>, tensor<1xi32>
}
// -----
// Test node names are preserved.
// CHECK-LABEL: func @main
func.func @main(%x: tensor<i32>) -> tensor<i32> {
// CHECK: tf_mlrt.executeop
// CHECK-SAME: name: \22name_loc/AddV2_0\22
%y = "tf.AddV2"(%x, %x) {__op_key = 0: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32> loc("name_loc:AddV2")
// CHECK: tf_mlrt.executeop
// CHECK-SAME: name: \22fused_loc/AddV2_1\22
%z = "tf.AddV2"(%y, %x) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32> loc(fused["fused_loc:", "AddV2"])
// CHECK: tf_mlrt.executeop
// CHECK-SAME: name: \22AddV2_2\22
%w = "tf.AddV2"(%z, %x) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
return %z : tensor<i32>
}
// -----
// Test function name canonicalization
// CHECK-LABEL: func @__inference_pruned_35
func.func @__inference_pruned_35() -> tensor<!tf_type.variant> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "flatmapdataset__4_RetVal"}} {
%0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Const"() {__op_key = 1: i32, device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.RangeDataset"(%0, %1, %2) {__op_key = 3: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tf_mlrt.executeop{{.*}}op: \22FlatMapDataset\22
// CHECK-SAME: \22__inference_Dataset_flat_map_lambda_19\22
%4 = "tf.FlatMapDataset"(%3) {__op_key = 4: i32, Targuments = [], device = "/device:CPU:0", f = @__inference_Dataset_flat_map_lambda_190, output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return %4 : tensor<!tf_type.variant>
}
// CHECK-LABEL: __inference_Dataset_flat_map_lambda_190
func.func private @__inference_Dataset_flat_map_lambda_190(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<!tf_type.variant> attributes {tf._original_func_name = "__inference_Dataset_flat_map_lambda_19", tf._tf_data_function = true, tf.signature.is_stateful} {
%0 = "tf.Const"() {__op_key = 5: i32, device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Const"() {__op_key = 6: i32,device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.Const"() {__op_key = 7: i32,device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.RangeDataset"(%0, %2, %1) {__op_key = 8: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tf_mlrt.executeop{{.*}}op: \22MapDataset\22
// CHECK-SAME: \22__inference_Dataset_map_lambda_16\22
%4 = "tf.MapDataset"(%3) {__op_key = 9: i32, device = "/device:CPU:0", f = @__inference_Dataset_map_lambda_160, f._tf_data_function = true, output_shapes = [#tf_type.shape<>], output_types = [i64], preserve_cardinality = true, use_inter_op_parallelism = true, metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
%5 = "tf.Identity"(%4) {__op_key = 10: i32, device = "/device:CPU:0"} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return %5 : tensor<!tf_type.variant>
}
// CHECK-LABEL: __inference_Dataset_map_lambda_160
func.func private @__inference_Dataset_map_lambda_160(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<i64> attributes {tf._tf_data_function = true} {
%0 = "tf.Const"() {__op_key = 11: i32, device = "/device:CPU:0", value = dense<2> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Mul"(%arg0, %0) {__op_key = 12: i32, device = "/device:CPU:0"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
%2 = "tf.Identity"(%1) {__op_key = 13: i32, device = "/device:CPU:0"} : (tensor<i64>) -> tensor<i64>
func.return %2 : tensor<i64>
}
// -----
// Test while conversion
// CHECK-LABEL: func @while_cond_lt9
// CHECK-SAME: ([[arg0:%.*]]: !tf_mlrt.tensor) -> !tf_mlrt.tensor
func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
// CHECK-LABEL: func @while_body_add2
// CHECK-SAME: ([[arg0:%.*]]: !tf_mlrt.tensor) -> !tf_mlrt.tensor
func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%arg0, %0) {__op_key = 3: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @while_test()
// CHECK-SAME: -> !tf_mlrt.tensor
func.func @while_test() -> (tensor<i32>) {
// CHECK: [[CONST:%.*]] = tf_mlrt.constop
%0 = "tf.Const"() {__op_key = 4: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[pred_res:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"([[CONST]]) : (!tf_mlrt.tensor) -> i1
// CHECK: [[while_res:%.*]]:2 = mlrt.while
// CHECK-SAME: @"while_body_add2/tf_mlrt_body"([[CONST]])
// CHECK-SAME: (!tf_mlrt.tensor) -> (!tf_mlrt.tensor, i1)
%1 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
// CHECK: return [[while_res]]#0 : !tf_mlrt.tensor
func.return %1 : tensor<i32>
}
// CHECK: func @"while_body_add2/tf_mlrt_body"([[arg:%.*]]: !tf_mlrt.tensor) -> (!tf_mlrt.tensor, i1)
// CHECK: [[body_res:%.*]] = call @while_body_add2([[arg]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
// CHECK: [[pred_res:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"([[body_res]]) : (!tf_mlrt.tensor) -> i1
// CHECK: return [[body_res]], [[pred_res]] : !tf_mlrt.tensor, i1
// CHECK: func @"while_cond_lt9/tf_mlrt_predicate"([[arg:%.*]]: !tf_mlrt.tensor) -> i1
// CHECK: [[cond_res:%.*]] = call @while_cond_lt9([[arg]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
// CHECK: [[bool_res:%.*]] = tf_mlrt.predicate [[cond_res]]
// CHECK: return [[bool_res]] : i1
// CHECK-LABEL: func @multi_while_test
func.func @multi_while_test() -> (tensor<i32>, tensor<i32>) {
%0 = "tf.Const"() {__op_key = 5: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Const"() {__op_key = 6: i32, device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[pred_0:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"
// CHECK: mlrt.while [[pred_0]] @"while_body_add2/tf_mlrt_body"
// CHECK: [[pred_1:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"
// CHECK: mlrt.while [[pred_1]] @"while_body_add2/tf_mlrt_body"
%2 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
%3 = "tf.While"(%1) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
func.return %2, %3 : tensor<i32>, tensor<i32>
}
// -----
// Test async output to function is converted
// CHECK-LABEL: @serving_default_stream_1
// CHECK-SAME: !mlrt.future
func.func private @serving_default_stream_1(%arg0: tensor<i32>) {
// CHECK: [[tensor:%.*]] = tf_mlrt.await
// CHECK: tf_mlrt.executeop([[tensor]])
%0 = "tf.StringFormat"(%arg0) {__op_key = 0: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", placeholder = "{}", strtemplate = "%s", summarize = 3 : i64, template = "Outside compiled {}"} : (tensor<i32>) -> tensor<!tf_type.string>
"tf.PrintV2"(%0) {__op_key = 1: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", end = "\0A", output_stream = "stderr"} : (tensor<!tf_type.string>) -> ()
return
}
func.func @callee(%arg: tensor<i32>) -> (tensor<i32>) {
func.return %arg: tensor<i32>
}
// CHECK-LABEL: @executeop_input
func.func @executeop_input(%arg0: tensor<i32>) -> (tensor<i32>) {
// CHECK: [[async_out:%.*]] = tf_mlrt.batch_function
%2 = "tf.BatchFunction"(%arg0) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = ""} : (tensor<i32>) -> tensor<i32>
// CHECK-NEXT: mlrt.async([[async_out]]) {{.*}} : (!mlrt.future)
%3 = mlrt.async(%2) {callee = @serving_default_stream_1} : (tensor<i32>) -> !mlrt.async_handle
// CHECK: mlrt.await_handle
mlrt.await_handle %3
// CHECK: return
// CHECK-SAME: !tf_mlrt.tensor
func.return %2 : tensor<i32>
}
// -----
// Support pre-assigned op_key
// CHECK-LABEL: @main
// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[promise_b:%.*]]: !mlrt.promise)
func.func @main(%input0: tensor<i32>, %promise_b: !mlrt.promise) {
%const = "tf.Const"() {__op_key = 0 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[a:%.*]] = tf_mlrt.executeop([[input0]],
// CHECK-SAME: AddV2
// CHECK-SAME: op_key = 1
// CHECK-NOT: __op_key
%a = "tf.AddV2"(%input0, %const) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: [[b:%.*]] = tf_mlrt.executeop([[a]])
// CHECK-SAME: Abs
// CHECK-SAME: op_key = 2
// CHECK-NOT: __op_key
%b = "tf.Abs"(%a) {__op_key = 2: i32 }: (tensor<i32>) -> tensor<i32>
// CHECK: tf_mlrt.promise [[promise_b]], [[b]]
"tf_mlrt.tf_promise"(%promise_b, %b) : (!mlrt.promise, tensor<i32>) -> ()
// CHECK: return
return
}
// -----
// Test future as input to promise
// CHECK-LABEL: func @main_stream_0
func.func @main_stream_0(%x: tensor<i32>, %p: !mlrt.promise) -> () {
// CHECK: [[y_future:%.*]] = tf_mlrt.async_executeop
%y = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
// CHECK: tf_mlrt.promise_future
// CHECK-SAME: [[y_future]]
"tf_mlrt.tf_promise"(%p, %y): (!mlrt.promise, tensor<i32>) -> ()
return
}
// CHECK-LABEL: @main
// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor)
func.func @main(%input0: tensor<i32>) -> tensor<i32> {
// CHECK: [[promises:%.*]], [[futures:%.*]] = "tf_mlrt.allocate_futures"
// CHECK-SAME: num_futures = 1
%promise_b, %future_b = "tf_mlrt.allocate_futures"()
{num_futures = 1 : i32, resultSegmentSizes = array<i32: 1, 1>} : () ->
(!mlrt.promise, !mlrt.future)
// CHECK: [[handle_0:%.*]] = mlrt.async([[input0]], [[promises]])
// CHECK-SAME: callee = @main_stream_0
%handle_0 = mlrt.async(%input0, %promise_b)
{callee = @main_stream_0} :
(tensor<i32>, !mlrt.promise) -> !mlrt.async_handle
// CHECK: [[const:%.*]] = tf_mlrt.const
%const = "tf.Const"() {__op_key = 1: i32, value = dense<2> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[b:%.*]] = tf_mlrt.await [[futures]]
%b = "tf_mlrt.tf_await"(%future_b) : (!mlrt.future) ->tensor<i32>
// CHECK: [[result:%.*]] = tf_mlrt.executeop([[b]], [[const]])
// CHECK-SAME: AddV2
%result = "tf.AddV2"(%b, %const) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: mlrt.await_handle [[handle_0]]
mlrt.await_handle %handle_0
// CHECK: return [[result]]
return %result : tensor<i32>
}
// -----
// Test lowering of tf call ops
// CHECK-LABEL: @callee
func.func @callee(%arg0: tensor<i32>) -> (tensor<i32>) {
func.return %arg0: tensor<i32>
}
// CHECK-LABEL: func @call_test
func.func @call_test(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
%0 = "tf.Add"(%arg0, %arg0) {__op_key = 0, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: [[results_0:%.*]] = call @callee(
// CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
%1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
// CHECK-NEXT: [[results_1:%.*]] = call @callee(
// CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
%2 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
// CHECK-NEXT: [[results_2:%.*]] = call @callee(
// CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
%3 = "tf.LegacyCall"(%0) {f = @callee} : (tensor<i32>) -> (tensor<i32>)
// CHECK: [[results_0]], [[results_1]], [[results_2]]
func.return %1, %2, %3 : tensor<i32>, tensor<i32>, tensor<i32>
}
// CHECK-LABEL: @branch0
func.func @branch0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
%0 = "tf.Add" (%arg0, %arg1) {__op_key = 1, device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: @branch1
func.func @branch1(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
%0 = "tf.Add" (%arg0, %arg1) {__op_key = 2, device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
%1 = "tf.Add" (%arg0, %0) {__op_key = 3, device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %1 : tensor<f32>
}
// CHECK-LABEL: func @case_test
// CHECK-SAME: ([[tf_idx:%.*]]: !tf_mlrt.tensor, [[branch_arg0:%.*]]: !tf_mlrt.tensor, [[branch_arg1:%.*]]: !tf_mlrt.tensor)
func.func @case_test(%arg0: tensor<i32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// CHECK: [[idx:%.*]] = tf_mlrt.tensor_to_int32 [[tf_idx]]
// CHECK-NEXT: [[out:%.*]] = mlrt.case [[idx]] [@branch0, @branch1]([[branch_arg0]], [[branch_arg1]])
%0 = "tf.Case"(%arg0, %arg1, %arg2) {_lower_using_switch_merge = true, branches = [@branch0, @branch1], is_stateless = true} : (tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
// Test await is added for unused futures
// CHECK-LABEL: func @unused_future_arg
// CHECK-SAME: ({{%.*}}: !tf_mlrt.tensor, [[unused:%.*]]: !mlrt.future)
func.func @unused_future_arg(%x: tensor<i32>, %unused: !mlrt.future) -> tensor<i32> {
// CHECK: mlrt.await_all_control [[unused]]
return %x : tensor<i32>
}
// CHECK-LABEL: func @unused_future
func.func @unused_future(%x: tensor<i32>) -> tensor<i32> {
// CHECK: [[unused:%.*]] = tf_mlrt.async_executeop
%unused = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
// CHECK: mlrt.await_all_control [[unused]]
return %x : tensor<i32>
}
// -----
// Test for XlaLaunch
func.func private @xla_func_0(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._XlaMustCompile = true, tf._noinline = true, tf._original_func_name = "should_not_be_used"} {
%1 = "tf.AddV2"(%arg0, %arg1) {__op_key = 0: i32} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %1 : tensor<1x3xf32>
}
// CHECK-LABEL: func @xla_func
func.func @xla_func(%arg0: tensor<1x3xf32>) -> tensor<*xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "output:0"}} {
%0 = "tf.VarHandleOp"() {__op_key = 1: i32, device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.ReadVariableOp"(%0) {__op_key = 2: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
// CHECK: tf_mlrt.executeop
// CHECK: tf_mlrt.async_executeop{{.*}}op: \22XlaLaunch\22\0A
// CHECK: tf_mlrt.await
// CHECK: return
// CHECK-SAME: !tf_mlrt.tensor
%2 = "tf.XlaLaunch"(%arg0, %1) {__op_key = 3: i32, _noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<*xf32>
func.return %2 : tensor<*xf32>
}
// -----
// Test lowering of IfrtLoadVariableOp
// CHECK-LABEL: func @ifrt_load_variable_test
func.func @ifrt_load_variable_test() -> () {
// CHECK: [[HANDLE:%.*]] = tf_mlrt.executeop()
// CHECK-SAME: VarHandleOp
%0 = "tf.VarHandleOp"() {__op_key = 1: i32, device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK-NEXT: "tf_mlrt.ifrt_load_variable"([[HANDLE]])
// CHECK-SAME: used_by_host = true
%1, %2 = "tf_mlrt.tf_ifrt_load_variable"(%0) {used_by_host = true, __op_key = 2: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
// CHECK-NEXT: mlrt.await_all_control
// CHECK-NEXT: return
func.return
}
// -----
// Test lowering of IfrtRestoreVariableOp
// CHECK-LABEL: func @ifrt_restore_variable_test
func.func @ifrt_restore_variable_test() -> () {
// CHECK-NEXT: [[PREFIX:%.*]] = tf_mlrt.constop
%cst = "tf.Const"() {__op_key = 0: i32, value = dense<"restore_ariables"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = tf_mlrt.constop
%cst_0 = "tf.Const"() {__op_key = 1: i32, value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[NAME:%.*]] = tf_mlrt.constop
%cst_1 = "tf.Const"() {__op_key = 2: i32, value = dense<["y"]> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[HANDLE:%.*]] = tf_mlrt.executeop
%handle = "tf.VarHandleOp"() {__op_key = 3: i32, container = "x", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: "tf_mlrt.ifrt_restore_variable"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLE]]) <{restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: true>}>
"tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %handle) {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: true>} : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> ()
// CHECK-NEXT: return
func.return
}
// -----
// Test lowering of tf.IfrtResourceDeserializeOp to tf_mlrt.ifrt_resource_deserialize
// CHECK-LABEL: func @ifrt_resource_deserialize_test
func.func @ifrt_resource_deserialize_test(%arg0: tensor<!tf_type.resource<tensor<f32>>>) {
%input_dir = "tf.Const"() { value = dense<"some/path"> : tensor<!tf_type.string> } : () -> tensor<!tf_type.string>
// CHECK: "tf_mlrt.ifrt_resource_deserialize"(%arg0, %{{.*}}) <{tensor_name = "my_tensor"}>
"tf.IfrtResourceDeserialize"(%arg0, %input_dir) {require_matching_crc = false, tensor_name = "my_tensor"} : (tensor<!tf_type.resource<tensor<f32>>>, tensor<!tf_type.string>) -> ()
func.return
}
// -----
// Test lowering of tf.IfrtRestoreVariableOp with outputs to tf_mlrt.ifrt_restore_variable
// CHECK-LABEL: func @ifrt_restore_variable_with_output_test
func.func @ifrt_restore_variable_with_output_test() -> (tensor<3x1xf32>) {
// CHECK-NEXT: [[PREFIX:%.*]] = tf_mlrt.constop
%cst = "tf.Const"() {__op_key = 0: i32, value = dense<"restore_ariables"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
// CHECK-NEXT: [[SLICE:%.*]] = tf_mlrt.constop
%cst_0 = "tf.Const"() {__op_key = 1: i32, value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[NAME:%.*]] = tf_mlrt.constop
%cst_1 = "tf.Const"() {__op_key = 2: i32, value = dense<["y"]> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
// CHECK-NEXT: [[HANDLE:%.*]] = tf_mlrt.executeop
%handle = "tf.VarHandleOp"() {__op_key = 3: i32, container = "x", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT: [[RESFUTURE:%.*]] = "tf_mlrt.ifrt_restore_variable"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLE]]) <{restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: true>}> : (!tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor) -> !mlrt.future
%result = "tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %handle) {restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: true>} : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
// CHECK-NEXT: [[RESULT:%.*]] = tf_mlrt.await [[RESFUTURE]]
// CHECK-NEXT: return [[RESULT]] : !tf_mlrt.tensor
func.return %result : tensor<3x1xf32>
}
// -----
// Test lowering of tf.IfrtCall (should use fallback)
// CHECK-LABEL: func @ifrt_call_fallback_test
func.func @ifrt_call_fallback_test(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: tf_mlrt.executeop
// CHECK-SAME: IfrtCall
%0 = "tf.IfrtCall"(%arg0) {program_id = 123 : i64, variable_arg_indices = [], __op_key = 0 : i32, operandSegmentSizes = array<i32: 1, 0>} : (tensor<i32>) -> tensor<i32>
return %0 : tensor<i32>
}
// -----
// Test lowering of tf.AsyncIfrtCall (should split into async call and await)
// CHECK-LABEL: func @async_ifrt_call_test
func.func @async_ifrt_call_test(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: [[FUTURE:%.*]] = tf_mlrt.async_ifrt_call
// CHECK: tf_mlrt.await [[FUTURE]]
%0 = "tf.AsyncIfrtCall"(%arg0) {program_id = 123 : i64, variable_arg_indices = [], __op_key = 0 : i32, operandSegmentSizes = array<i32: 1, 0>} : (tensor<i32>) -> tensor<i32>
return %0 : tensor<i32>
}
// -----
// Test lowering of tf.PartitionedCall and tf.StatefulPartitionedCall to func.call
// CHECK-LABEL: func private @callee
// CHECK-SAME: (%arg0: !tf_mlrt.tensor) -> !tf_mlrt.tensor
func.func private @callee(%arg0: tensor<i32>) -> tensor<i32> {
return %arg0 : tensor<i32>
}
// CHECK-LABEL: func @test_call_ops
// CHECK-SAME: ([[ARG:%.*]]: !tf_mlrt.tensor) -> (!tf_mlrt.tensor, !tf_mlrt.tensor)
func.func @test_call_ops(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
// CHECK: [[RES1:%.*]] = call @callee([[ARG]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
%0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> tensor<i32>
// CHECK: [[RES2:%.*]] = call @callee([[ARG]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
%1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> tensor<i32>
// CHECK: return [[RES1]], [[RES2]] : !tf_mlrt.tensor, !tf_mlrt.tensor
return %0, %1 : tensor<i32>, tensor<i32>
}
// -----
// Test lowering of tf.PartitionedCall with a future input (should insert await)
// CHECK-LABEL: func private @callee_future
// CHECK-SAME: (%arg0: !tf_mlrt.tensor) -> !tf_mlrt.tensor
func.func private @callee_future(%arg0: tensor<3x1xf32>) -> tensor<3x1xf32> {
return %arg0 : tensor<3x1xf32>
}
// CHECK-LABEL: func @test_call_with_future
// CHECK-SAME: ([[ARG0:%.*]]: !tf_mlrt.tensor, [[ARG1:%.*]]: !tf_mlrt.tensor, [[ARG2:%.*]]: !tf_mlrt.tensor, [[ARG3:%.*]]: !tf_mlrt.tensor) -> !tf_mlrt.tensor
func.func @test_call_with_future(%arg0: tensor<!tf_type.string>, %arg1: tensor<1x!tf_type.string>, %arg2: tensor<1x!tf_type.string>, %arg3: tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32> {
// CHECK: [[FUTURE:%.*]] = "tf_mlrt.ifrt_restore_variable"([[ARG0]], [[ARG1]], [[ARG2]], [[ARG3]])
%result = "tf.IfrtRestoreVariableOp"(%arg0, %arg1, %arg2, %arg3) {restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: true>} : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
// CHECK-NEXT: [[AWAITED:%.*]] = tf_mlrt.await [[FUTURE]]
// CHECK-NEXT: [[RESULT:%.*]] = call @callee_future([[AWAITED]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
%0 = "tf.PartitionedCall"(%result) {config = "", config_proto = "", executor_type = "", f = @callee_future} : (tensor<3x1xf32>) -> tensor<3x1xf32>
// CHECK: return [[RESULT]] : !tf_mlrt.tensor
return %0 : tensor<3x1xf32>
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/tpu_conversions.mlir
+243
View File
@@ -0,0 +1,243 @@
// 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 --split-input-file -pass-pipeline='builtin.module(pre-parallel-tf-to-mlrt{use-tpu-host-allocator-for-inputs=true},tf-mlrt-parallelization{tfrt-cost-threshold=4},tf-to-mlrt)' %s | FileCheck %s --dump-input=fail --dump-input-filter=all
func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
func.return %arg0: tensor<i32>
}
// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>) -> (tensor<i32>) {
// CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function
// CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
// CHECK-NEXT: [[rendezvous_key_base:%.*]] = tf_mlrt_tpu.compile_and_execute([[batch_result]])
// CHECK-NEXT: return [[rendezvous_key_base]]
%0 = "tf.BatchFunction"(%arg0, %arg0) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// -----
func.func @executeop_input(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
// CHECK-NOT: tf_mlrt.executeop(
// CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
// CHECK: [[cast:%.*]] = tf_mlrt.executeop.device([[device]]){{.*}}op: \22Cast\22
// CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute([[cast]])
// CHECK: tf_mlrt.await [[result_future]]
%0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
%1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
func.return %1, %2 : tensor<i32>, tensor<i32>
}
// -----
func.func @executeop_side_effecting_input(%arg0: tensor<!tf_type.resource<tensor<4xf32>>>, %indices: tensor<i32>) -> (tensor<i32>) {
// CHECK-NOT: tf_mlrt.executeop(
// CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
// CHECK: [[var:%.*]] = tf_mlrt.executeop.device([[device]]){{.*}}op: \22ResourceGather\22
// CHECK: [[rendezvous_key_base:%.*]] = tf_mlrt_tpu.compile_and_execute([[var]])
%0 = "tf.ResourceGather"(%arg0, %indices) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<4xf32>>>, tensor<i32>) -> tensor<f32>
%1 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// -----
func.func @executeop_input_same_execute_op(%arg0: tensor<i32>, %arg1: tensor<2xf32>) -> (tensor<i32>) {
// CHECK-NOT: tf_mlrt.executeop(
// CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
// CHECK: [[split:%.*]]:2 = tf_mlrt.executeop.device([[device]])
// CHECK: tf_mlrt_tpu.compile_and_execute([[split]]#0, [[split]]#1)
%0, %1 = "tf.Split"(%arg0, %arg1) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<2xf32>) -> (tensor<f32>, tensor<f32>)
%2 = "tf.TPUCompileMlirAndExecute"(%0, %1) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, producer_name = "producer_name"} : (tensor<f32>, tensor<f32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// -----
// Test that inputs are lowered correctly when they form a DAG.
// CHECK-LABEL: executeop_dag
func.func @executeop_dag(%arg0: tensor<i32>) -> (tensor<i32>) {
// CHECK-NEXT: tf_mlrt_tpu.get_tpu_host_device
// CHECK-NEXT: tf_mlrt.executeop.device{{.*}}op: \22Cast\22
// CHECK-NEXT: tf_mlrt_tpu.get_tpu_host_device
// CHECK-NEXT: tf_mlrt.executeop.device{{.*}}op: \22Relu\22
// CHECK-NEXT: tf_mlrt_tpu.compile_and_execute
%0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
%1 = "tf.Relu"(%0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<f32>) -> (tensor<f32>)
%2 = "tf.TPUCompileMlirAndExecute"(%1, %0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, producer_name = "producer_name"} : (tensor<f32>, tensor<f32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// -----
func.func @test_fuse_dynamic_dimension_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<*xi32>, %arg3: tensor<*xi32>, %arg4: tensor<*xi32>, %arg5: tensor<?xi64>, %arg6: tensor<?xi64>, %arg7: tensor<?xi64>) -> tensor<*xi32> {
%0 = "tf.ReadVariableOp"(%arg1) {__op_key = 0: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
%1 = "tf.Shape"(%arg0) {__op_key = 1: i32, device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
%2 = "tf.Shape"(%0) {__op_key = 2: i32, device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
// CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute
// CHECK-SAME: constant_operand_indices = array<i32: 2>
// CHECK-SAME: num_operands = 4
// CHECK-SAME: operands_with_static_shape = array<i32: 0, 1, 3>
%rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %2, %0, %1, %arg5, %arg6, %arg7) {operands_with_static_shape = [0 : i32, 1 : i32, 3 : i32], metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 4, 3>, producer_name = "producer_name"} : (tensor<*xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
func.return %results : tensor<*xi32>
}
// -----
// Test async output of tf.TPUCompileMlirAndExecute to function is converted
// CHECK-LABEL: @executeop_input_stream_1
// CHECK-SAME: ([[future:%.*]]: !mlrt.future
// CHECK: [[tensor:%.*]] = tf_mlrt.await [[future]]
// CHECK: tf_mlrt.executeop([[tensor]])
// CHECK-SAME: StringFormat
// CHECK-LABEL: @executeop_input
func.func @executeop_input(%arg0: tensor<i32>) -> (tensor<i32>) {
// CHECK: tf_mlrt.executeop
%0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
// CHECK: [[rendezvous_key_base:%.*]], [[result:%.*]] = tf_mlrt_tpu.compile_and_execute
%1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
%3 = "tf.StringFormat"(%2) {__op_key = 1: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", placeholder = "{}", strtemplate = "%s", summarize = 3 : i64, template = "Outside compiled {}"} : (tensor<i32>) -> tensor<!tf_type.string>
"tf.PrintV2"(%3) {__op_key = 2: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", end = "\0A", output_stream = "stderr"} : (tensor<!tf_type.string>) -> ()
// CHECK: [[handle:%.*]] = mlrt.async([[result]])
// CHECK-SAME: (!mlrt.future)
// CHECK: mlrt.await_handle [[handle]]
// CHECK: return [[rendezvous_key_base]]
// CHECK-SAME: !tf_mlrt.tensor
func.return %1 : tensor<i32>
}
// -----
// Test constant arguments to tf.TPUCompileMlirAndExecute are preserved during parallelization.
// CHECK-LABEL: @preserve_constant_args(
func.func @preserve_constant_args(%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<*x!tf_type.resource>, %arg3: tensor<*x!tf_type.resource>) -> (tensor<i32>) {
// CHECK-NOT: ReadVariableOp
// CHECK: mlrt.async(
%v0 = "tf.ReadVariableOp"(%arg1) {__op_key = 0: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
%v1 = "tf.ReadVariableOp"(%arg2) {__op_key = 1: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
// CHECK: [[cast:%.*]] = tf_mlrt.executeop(
// CHECK-SAME: ReadVariableOp
%v2 = "tf.ReadVariableOp"(%arg3) {__op_key = 2: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
// CHECK: [[cast:%.*]] = tf_mlrt.executeop.device
// CHECK-SAME: Cast
%0 = "tf.Cast"(%arg0) {__op_key = 3: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
// CHECK: tf_mlrt_tpu.compile_and_execute({{%.*}}, [[cast]]
// CHECK-SAME: constant_operand_indices = array<i32: 1, 3, 4>
%1, %2 = "tf.TPUCompileMlirAndExecute"(%0, %v1, %0, %v2, %v0, %arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 6, 0>, producer_name = "producer_name"} : (tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.return %2 : tensor<i32>
}
// -----
func.func @executeop_input_async() -> (tensor<i32>, tensor<i32>) {
// CHECK-NOT: tf_mlrt.executeop(
// CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
// CHECK: [[recv_future:%.*]] = tf_mlrt.async_executeop.device([[device]]){{.*}}op: \22Recv\22
// CHECK: [[recv:%.*]] = tf_mlrt.await [[recv_future]]
// CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute([[recv]])
// CHECK: tf_mlrt.await [[result_future]]
%0 = "tf.Recv"() {__op_key = 0: i32, device = "/device:CPU:0", tensor_name = "tensor", send_device = "/device:CPU:0", send_device_incarnation = 0, recv_device = "/device:CPU:0"} : () -> tensor<f32>
%1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
func.return %1, %2 : tensor<i32>, tensor<i32>
}
// -----
// Test the output from TPU op is properly awaited before its use by map_fn.
// CHECK-LABEL: @main
// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[input1:%.*]]: !tf_mlrt.tensor)
func.func @main(%input0: tensor<i32>, %input1: tensor<i32>, %input2: tensor<!tf_type.variant<tensor<*xf32>>> ) -> tensor<i32> {
%0 = "tf.Cast"(%input0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
// CHECK: tf_mlrt_tpu.compile_and_execute
%1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
%max_iter = "tf.Const"() {__op_key = 1, value = dense<2> : tensor<i32>} : () -> tensor<i32>
// CHECK: tf_mlrt.map_fn
%result = "tf_mlrt.tf_map_fn"(%max_iter, %input2, %2) { operandSegmentSizes = array<i32: 1, 1, 1>, body_fn = @NopMapFnBody, num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> tensor<i32>
return %result : tensor<i32>
}
// CHECK-LABEL: @NopMapFnBody
func.func private @NopMapFnBody(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>) -> () {
%const = "tf.Const"() {__op_key = 2 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
%a = "tf.AddV2"(%arg2, %const) {__op_key = 3: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
return
}
// -----
// If not all in_tensors of the batched op are used by TPUCompileMlirAndExecute,
// should not lower batch_function to tf_mlrt.batch_function.device.
func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
%1 = "tf.TPUCompileMlirAndExecute"(%arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
%const = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%a = "tf.AddV2"(%arg1, %const) {__op_key = 3: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1: tensor<i32>
}
// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
// CHECK: tf_mlrt.batch_function(%arg0, %arg1)
// CHECK-NOT: batch_function.device
%0 = "tf.BatchFunction"(%arg0, %arg1) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 2, 0>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// captured tensors not used by TPUCompileMlirAndExecute is ok.
func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
%1 = "tf.TPUCompileMlirAndExecute"(%arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
%const = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%a = "tf.AddV2"(%arg1, %const) {__op_key = 3: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1: tensor<i32>
}
// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
// CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
// CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function.device([[device]]) (%arg0, %arg1)
// CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
// CHECK: return [[batch_result]]
%0 = "tf.BatchFunction"(%arg0, %arg1) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
func.func @batched_func(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
%0 = "tf.TPUCompileMlirAndExecute"(%arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
%2 = "tf.TPUCompileMlirAndExecute"(%arg1) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
func.return %2: tensor<i32>
}
// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
// CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
// CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function.device([[device]]) (%arg0, %arg1)
// CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
// CHECK: return [[batch_result]]
%0 = "tf.BatchFunction"(%arg0, %arg1) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 2, 0>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
tensorflow/compiler/mlir/tfrt/tests/mlrt/while_to_map_fn.mlir
+829
View File
@@ -0,0 +1,829 @@
// 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 -split-input-file -tf-mlrt-while-to-map-fn %s | FileCheck %s
// Test a while to map_fn conversion in which the max iteration is hard coded inside the predicate body.
// CHECK-LABEL: map/while_cond
func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
return %2 : tensor<i1>
}
// CHECK-LABEL: map/while_body
func.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00, 9.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_3 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
%cst_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.AddV2"(%arg0, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.Mul"(%arg3, %cst_3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<9xf32>) -> tensor<9xf32>
%2 = "tf.Reshape"(%1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
%3 = "tf.AddV2"(%arg1, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%4 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%5 = "tf.Cast"(%4) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
%6 = "tf.Mul"(%5, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<9xf32>) -> tensor<9xf32>
%7 = "tf.Reshape"(%6, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
%8 = "tf.MatMul"(%2, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%9 = "tf.MatrixDeterminant"(%8) {T = f32, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<f32>
%10 = "tf.TensorListSetItem"(%arg2, %arg1, %9) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %0, %3, %10, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
}
// CHECK-LABEL: map/while_body/MapFnBody
// CHECK-SAME: (%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<?xf32>)
// CHECK: [[det:%.*]] = "tf.MatrixDeterminant"
// CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0) : (!mlrt.future) -> tensor<!tf_type.variant<tensor<*xf32>>>
// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[det]]) <{
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[ta_1]]) : (!mlrt.promise, tensor<!tf_type.variant<tensor<*xf32>>>) -> ()
// CHECK-NEXT: return
//CHECK-LABEL: @serving_default
func.func @serving_default(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> tensor<3xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input:0", outputs = "PartitionedCall:0"}} {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[tensor_list:%.*]] = "tf.TensorListReserve"([[shape:%.*]], [[reserve_size:%.*]]) {
%0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
// CHECK: [[map_fn_result:%.*]] = tf_mlrt.tf_map_fn([[reserve_size]], [[tensor_list]], %arg0)
// CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
// CHECK-NOT: tf.While
%1:4 = "tf.While"(%cst, %cst, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
// CHECK-NEXT: "tf.TensorListStack"([[map_fn_result]], %cst_0) <{
%2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 3 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<0xi32>) -> tensor<3xf32>
return %2 : tensor<3xf32>
}
// -----
// Test a while to map_fn conversion in which max_iterations are passed
// into the predicate function.
// CHECK-LABEL: @"map/while_cond"
func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>, %arg4: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg5: tensor<?x3xf32>, %arg6: tensor<?x4xf32>) -> tensor<i1> {
%outputs = "tf.Less"(%arg0, %arg3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%outputs_0 = "tf.Less"(%arg1, %arg3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%outputs_2 = "tf.LogicalAnd"(%outputs_0, %outputs) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
return %outputs_2 : tensor<i1>
}
// CHECK-LABEL: @"map/while_body"
func.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>, %arg4: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg5: tensor<?x3xf32>, %arg6: tensor<?x4xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>) {
%outputs = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_2 = "tf.AddV2"(%arg0, %outputs_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%outputs_4 = "tf.ReadVariableOp"(%arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
%outputs_6 = "tf.Identity"(%outputs_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
%outputs_8 = "tf.MatMul"(%arg5, %outputs_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<?x3xf32>, tensor<3x1xf32>) -> tensor<?x1xf32>
%outputs_10 = "tf.AddV2"(%arg1, %outputs_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%outputs_12 = "tf.Identity"(%outputs_10) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
%outputs_14 = "tf.GatherV2"(%arg6, %arg1, %outputs) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x4xf32>, tensor<i32>, tensor<i32>) -> tensor<4xf32>
%outputs_16 = "tf.AddV2"(%outputs_8, %outputs_14) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x1xf32>, tensor<4xf32>) -> tensor<?x4xf32>
%outputs_18 = "tf.TensorListSetItem"(%arg2, %arg1, %outputs_16) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<?x4xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %outputs_6, %outputs_12, %outputs_18, %arg3, %arg4, %arg5, %arg6 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>
}
// CHECK-LABEL: @"map/while_body/MapFnBody"
// CHECK-SAME (%arg0: !mlrt.Future, %arg1: !mlrt.Promise, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg6: tensor<?x3xf32>, %arg7: tensor<?x4xf32>)
// CHECK-NEXT: [[cst_0:%.*]] = "tf.Const"
// CHECK-NEXT: [[cst_1:%.*]] = "tf.Const"
// CHECK-NEXT: [[loop_counter:%.*]] = "tf.AddV2"(%arg2, [[cst_1]])
// CHECK-NEXT: [[weight:%.*]] = "tf.ReadVariableOp"(%arg5)
// CHECK-NEXT: [[mpy:%.*]] = "tf.MatMul"(%arg6, [[weight]])
// CHECK-NEXT: [[element_index:%.*]] = "tf.AddV2"(%arg3, [[cst_1]])
// CHECK-NEXT: [[bias:%.*]] = "tf.GatherV2"(%arg7, %arg3, [[cst_0]])
// CHECK-NEXT: [[res:%.*]] = "tf.AddV2"([[mpy]], [[bias]])
// CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0)
// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[res]])
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[ta_1]])
// CHECK-NEXT: return
// CHECK-LABEL: func @main_while
func.func @main_while(%arg0: tensor<?x3xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x?x4xf32> {
%outputs = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[-1, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%outputs_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_6 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_8 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[elems:%.*]] = "tf.VarHandleOp"
%outputs_10 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shared_name = "w"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
%outputs_12 = "tf.Shape"(%arg1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x4xf32>) -> tensor<2xi32>
// CHECK: [[max_iter:%.*]] = "tf.StridedSlice"
%outputs_14 = "tf.StridedSlice"(%outputs_12, %outputs_2, %outputs_4, %outputs_4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: [[tensor_list:%.*]] = "tf.TensorListReserve"
%outputs_16 = "tf.TensorListReserve"(%outputs_0, %outputs_14) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
// CHECK: tf_mlrt.tf_map_fn
// CHECK-SAME: ([[max_iter]], [[tensor_list]], [[max_iter]], [[elems]], %arg0, %arg1)
// CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
// CHECK-NOT: tf.while
%outputs_18:7 = "tf.While"(%outputs_6, %outputs_6, %outputs_16, %outputs_14, %outputs_10, %arg0, %arg1) {_lower_using_switch_merge = true, _num_original_outputs = 8 : i64, _read_only_resource_inputs = [6], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>)
%outputs_20 = "tf.TensorListStack"(%outputs_18#2, %outputs) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = -1 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<2xi32>) -> tensor<?x?x4xf32>
return %outputs_20 : tensor<?x?x4xf32>
}
// -----
// Test a while to map_fn conversion in which the passed in max_iterations
// is not in typical location of %arg3 and there are identify chains in function bodies.
// CHECK-LABEL: @map_while_cond_170
func.func private @map_while_cond_170(%arg0: tensor<i32> {tf._user_specified_name = "map/while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "map/while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>, %arg4: tensor<*x!tf_type.variant>, %arg5: tensor<*xf32>) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime", tf._original_func_name = "map_while_cond_17"} {
%outputs = "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
%outputs_0 = "tf.Less"(%arg0, %arg1) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%outputs_2 = "tf.Less"(%arg2, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%outputs_4 = "tf.LogicalAnd"(%outputs_0, %outputs_2) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
%outputs_6 = "tf.Identity"(%outputs_4) {device = ""} : (tensor<*xi1>) -> tensor<*xi1>
return %outputs_6 : tensor<*xi1>
}
// Original input argument list (loop_counter, max_iterations, element_index, tensor_list, read_only_tensor_list, scale)
// CHECK-LABEL: @map_while_body_180
func.func private @map_while_body_180(%arg0: tensor<i32> {tf._user_specified_name = "map/while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "map/while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>, %arg4: tensor<!tf_type.variant> {tf._user_specified_name = "map/TensorArrayUnstack/TensorListFromTensor"}, %arg5: tensor<?xf32> {tf._user_specified_name = "input"}) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>) attributes {tf._construction_context = "kEagerRuntime", tf._original_func_name = "map_while_body_18"} {
%outputs = "tf.Const"() {device = "", value = dense<16> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_0 = "tf.Const"() {device = "", value = dense<16> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_2 = "tf.Const"() {device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_4 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_6 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_8 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_10 = "tf.Const"() {device = "", value = dense<256> : tensor<i32>} : () -> tensor<i32>
%outputs_12 = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_14 = "tf.Range"(%outputs_12, %outputs_10, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_16 = "tf.Cast"(%outputs_14) {Truncate = false, device = ""} : (tensor<*xi32>) -> tensor<*xf32>
%outputs_18 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_20 = "tf.Const"() {device = "", value = dense<257> : tensor<i32>} : () -> tensor<i32>
%outputs_22 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_24 = "tf.Range"(%outputs_22, %outputs_20, %outputs_18) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_26 = "tf.Cast"(%outputs_24) {Truncate = false, device = ""} : (tensor<*xi32>) -> tensor<*xf32>
%outputs_28 = "tf.Const"() {device = "", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_30 = "tf.Transpose"(%outputs_26, %outputs_28) {device = ""} : (tensor<*xf32>, tensor<1xi32>) -> tensor<*xf32>
%outputs_32 = "tf.AddV2"(%arg0, %outputs_6) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_34 = "tf.Identity"(%outputs_32) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_36 = "tf.Identity"(%arg1) {device = ""} : (tensor<i32>) -> tensor<*xi32>
%outputs_38 = "tf.Mul"(%outputs_16, %arg5) {device = ""} : (tensor<*xf32>, tensor<?xf32>) -> tensor<*xf32>
%outputs_40 = "tf.Reshape"(%outputs_38, %outputs) {device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
%outputs_42 = "tf.AddV2"(%arg2, %outputs_4) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_44 = "tf.Identity"(%outputs_42) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_46 = "tf.TensorListGetItem"(%arg4, %arg2, %outputs_2) {device = ""} : (tensor<!tf_type.variant>, tensor<i32>, tensor<0xi32>) -> tensor<*xi32>
%outputs_48 = "tf.Cast"(%outputs_46) {Truncate = false, device = ""} : (tensor<*xi32>) -> tensor<*xf32>
%outputs_50 = "tf.Mul"(%outputs_30, %outputs_48) {device = ""} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%outputs_52 = "tf.Reshape"(%outputs_50, %outputs_0) {device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
%outputs_54 = "tf.MatMul"(%outputs_40, %outputs_52) {device = "", transpose_a = false, transpose_b = false} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%outputs_56 = "tf.MatrixDeterminant"(%outputs_54) {T = f32, device = ""} : (tensor<*xf32>) -> tensor<*xf32>
%outputs_58 = "tf.TensorListSetItem"(%arg3, %arg2, %outputs_56) {device = "", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant>, tensor<i32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
%outputs_60 = "tf.Identity"(%outputs_58) {device = ""} : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>
return %outputs_34, %outputs_36, %outputs_44, %outputs_60, %arg4, %arg5 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>
}
// Converted input argument list (loop_counter, element_index, max_iterations, tensor_list, read_only_tensor_list, scale)
// CHECK-LABEL: @"map_while_body_180/MapFnBody"
// CHECK-SAME: (%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor<i32> {tf._user_specified_name = "map/while/loop_counter"}, %arg3: tensor<i32>, %arg4: tensor<i32> {tf._user_specified_name = "map/while/maximum_iterations"}, %arg5: tensor<!tf_type.variant> {tf._user_specified_name = "map/TensorArrayUnstack/TensorListFromTensor"}, %arg6: tensor<?xf32> {tf._user_specified_name = "input"})
// CHECK: [[res:%.*]] = "tf.MatrixDeterminant"
// CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0)
// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[res]])
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[ta_1]])
// CHECK-NEXT: return
// CHECK-LABEL: __inference_while_from_map_fn_810
// CHECK-SAME: ([[scale:%.*]]: tensor<?xf32>
func.func private @__inference_while_from_map_fn_810(%arg0: tensor<?xf32> {tf._user_specified_name = "input"}) -> tensor<*xf32> attributes {tf._construction_context = "kEagerRuntime", tf._original_func_name = "__inference_while_from_map_fn_81"} {
// CHECK: [[element_index:%.*]] = "tf.Const"
%outputs = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_0 = "tf.Const"() {device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_2= "tf.Const"() {device = "", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%outputs_4 = "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
// CHECK: tf.TensorListReserve
%outputs_6 = "tf.TensorListReserve"(%outputs_2, %outputs_4) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>>
%outputs_8 = "tf.Const"() {device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_10 = "tf.Const"() {device = "", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%outputs_12 = "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[tensor_list:%.*]] = "tf.TensorListReserve"([[shape:%.*]], [[reserve_size:%.*]]) {
%outputs_14 = "tf.TensorListReserve"(%outputs_10, %outputs_12) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
// CHECK-NEXT: [[loop_counter:%.*]] = "tf.Const"
%outputs_16 = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK-NEXT: [[max_iterations:%.*]] = "tf.Const"
%outputs_18 = "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
%outputs_20 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_22 = "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
%outputs_24 = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_26 = "tf.Range"(%outputs_24, %outputs_22, %outputs_20) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
// CHECK: [[read_only_tensor_list:%.*]] = "tf.TensorListFromTensor"
%outputs_28 = "tf.TensorListFromTensor"(%outputs_26, %outputs_0) {device = ""} : (tensor<*xi32>, tensor<0xi32>) -> tensor<*x!tf_type.variant>
// CHECK: [[map_fn_out:%.*]] = tf_mlrt.tf_map_fn
// CHECK-SAME: ([[reserve_size]], [[tensor_list]], [[max_iterations]], [[read_only_tensor_list]], [[scale]])
// CHECK-SAME: {body_fn = @"map_while_body_180/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
// CHECK-NOT: tf.While
%outputs_30:6 = "tf.While"(%outputs_16, %outputs_18, %outputs, %outputs_14, %outputs_28, %arg0) {T = [i32, i32, i32, !tf_type.variant, !tf_type.variant, f32], _lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], body = @map_while_body_180, cond = @map_while_cond_170, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<?>], parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<*x!tf_type.variant>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>)
// CHECK-NEXT: "tf.TensorListStack"
// CHECK-SAME: ([[map_fn_out]],
%outputs_32 = "tf.TensorListStack"(%outputs_30#3, %outputs_8) {device = "", num_elements = 16 : i64} : (tensor<!tf_type.variant>, tensor<0xi32>) -> tensor<*xf32>
%outputs_34 = "tf.Identity"(%outputs_32) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
return %outputs_34 : tensor<*xf32>
}
// -----
// Test a while to map_fn conversion in which tensor array is used instead of
// tensor list.
// CHECK-LABEL: map/while/LoopCond_cond
func.func private @"map/while/LoopCond_cond"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> tensor<i1> {
%outputs = "tf.Less"(%arg0, %arg3) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%outputs_0 = "tf.Less"(%arg1, %arg3) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%outputs_2 = "tf.LogicalAnd"(%outputs, %outputs_0) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
%outputs_4 = "tf.ToBool"(%outputs_2) : (tensor<*xi1>) -> tensor<i1>
return %outputs_4 : tensor<i1>
}
// CHECK-LABEL: map/while/LoopCond_body
func.func private @"map/while/LoopCond_body"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) {
%outputs = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_4 = "tf.Identity"(%arg0) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_6 = "tf.AddV2"(%outputs_4, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%outputs_8 = "tf.Identity"(%arg1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_10 = "tf.AddV2"(%outputs_8, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%outputs_12 = "tf.Identity"(%arg2) {device = ""} : (tensor<f32>) -> tensor<f32>
%outputs_14 = "tf.TensorArrayReadV3"(%arg4, %outputs_8, %arg5) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<*xi32>, tensor<f32>) -> tensor<*x!tf_type.string>
%outputs_16 = "tf.DecodeJpeg"(%outputs_14) {acceptable_fraction = 1.000000e+00 : f32, channels = 3 : i64, dct_method = "INTEGER_FAST", device = "", fancy_upscaling = true, ratio = 1 : i64, try_recover_truncated = false} : (tensor<*x!tf_type.string>) -> tensor<?x?x3xui8>
%outputs_18 = "tf.ExpandDims"(%outputs_16, %outputs_0) {device = ""} : (tensor<?x?x3xui8>, tensor<i32>) -> tensor<1x?x?x3xui8>
%outputs_20 = "tf.ResizeBilinear"(%outputs_18, %outputs) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x?x?x3xui8>, tensor<2xi32>) -> tensor<1x224x224x3xf32>
%outputs_22 = "tf.Squeeze"(%outputs_20) {device = "", squeeze_dims = [0]} : (tensor<1x224x224x3xf32>) -> tensor<224x224x3xf32>
%outputs_24 = "tf.Cast"(%outputs_22) {Truncate = false, device = ""} : (tensor<224x224x3xf32>) -> tensor<224x224x3xui8>
%outputs_26 = "tf.TensorArrayWriteV3"(%arg6, %outputs_8, %outputs_24, %outputs_12) {device = ""} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<*xi32>, tensor<224x224x3xui8>, tensor<f32>) -> tensor<f32>
return %outputs_6, %outputs_10, %outputs_26, %arg3, %arg4, %arg5, %arg6: tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>
}
// CHECK-LABEL: @"map/while/LoopCond_body/MapFnBody"
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.TensorArrayReadV3
// CHECK-NEXT: tf.DecodeJpeg
// CHECK-NEXT: tf.ExpandDims
// CHECK-NEXT: tf.ResizeBilinear
// CHECK-NEXT: tf.Squeeze
// CHECK-NEXT: tf.Cast
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorArrayWriteV3
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: return
//CHECK-LABEL: map_while_test
func.func @map_while_test(%arg0: tensor<?x!tf_type.string>) -> tensor<?x224x224x3xui8> {
%outputs = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<1xi32>
%outputs_0 = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_4 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_6 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_8 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_10 = "tf.Shape"(%arg0) {device = ""} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
// CHECK: [[max_iter:%.*]] = "tf.StridedSlice"
%outputs_12 = "tf.StridedSlice"(%outputs_10, %outputs_6, %outputs_4, %outputs_4) {begin_mask = 0 : i64, device = "", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: tf.Range
%outputs_14 = "tf.Range"(%outputs_2, %outputs_12, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
// CHECK-NEXT: [[handle_1:%.*]], [[flow_in_1:%.*]] = "tf.TensorArrayV3"
%outputs_16:2 = "tf.TensorArrayV3"(%outputs_12) {clear_after_read = true, device = "", dtype = !tf_type.string, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>)
// CHECK-NEXT: [[handle_2:%.*]] = "tf.TensorArrayScatterV3"
%outputs_18 = "tf.TensorArrayScatterV3"(%outputs_16#0, %outputs_14, %arg0, %outputs_16#1) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<?xi32>, tensor<?x!tf_type.string>, tensor<f32>) -> tensor<f32>
// CHECK-NEXT: tf.Range
%outputs_20 = "tf.Range"(%outputs_2, %outputs_12, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
// CHECK-NEXT: [[tensor_array:%.*]], [[flow_in:%.*]] = "tf.TensorArrayV3"
%outputs_22:2 = "tf.TensorArrayV3"(%outputs_12) {clear_after_read = true, device = "", dtype = ui8, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<f32>)
// CHECK-NEXT: tf_mlrt.tf_map_fn
// CHECK-SAME: ([[max_iter]], [[flow_in]], [[max_iter]], [[handle_1]], [[handle_2]], [[tensor_array]])
// CHECK-SAME: {body_fn = @"map/while/LoopCond_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
// CHECK-NOT: tf.While
%outputs_24:7 = "tf.While"(%outputs, %outputs, %outputs_22#1, %outputs_12, %outputs_16#0, %outputs_18, %outputs_22#0) {_xla_propagate_compile_time_consts = true, body = @"map/while/LoopCond_body", cond = @"map/while/LoopCond_cond", device = "", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>)
// CHECK-NEXT: tf.TensorArrayGatherV3
%outputs_26 = "tf.TensorArrayGatherV3"(%outputs_22#0, %outputs_20, %outputs_24#2) {device = "", element_shape = #tf_type.shape<224x224x3>} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<?xi32>, tensor<f32>) -> tensor<?x224x224x3xui8>
return %outputs_26 : tensor<?x224x224x3xui8>
}
// -----
// Test non-applicable while is NOT converted to map_fn.
// CHECK-LABEL: func @while_cond_lt9
func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
// CHECK-LABEL: func @while_body_add2
func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @while_test()
func.func @while_test() -> (tensor<i32>) {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: tf.While
%1 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
func.return %1 : tensor<i32>
}
// -----
// Test a case that the while body has multiple tensor lists.
// CHECK-LABEL: tf.MultiListWhileRegion_body
func.func private @tf.MultiListWhileRegion_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>, %arg4: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00], [8.000000e+00, 9.000000e+00, 1.000000e+01, 1.100000e+01, 1.200000e+01, 1.300000e+01, 1.400000e+01, 1.500000e+01]]> : tensor<2x8xf32>} : () -> tensor<2x8xf32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[[1.600000e+01, 1.700000e+01, 1.800000e+01, 1.900000e+01, 2.000000e+01, 2.100000e+01, 2.200000e+01, 2.300000e+01], [2.400000e+01, 2.500000e+01, 2.600000e+01, 2.700000e+01, 2.800000e+01, 2.900000e+01, 3.000000e+01, 3.100000e+01]]> : tensor<2x8xf32>} : () -> tensor<2x8xf32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.GatherV2"(%arg4, %cst_2, %cst_2) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>, tensor<i32>) -> tensor<f32>
%1 = "tf.AddV2"(%arg0, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.AddV2"(%arg1, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%3 = "tf.GatherV2"(%cst_0, %arg1, %cst_2) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x8xf32>, tensor<i32>, tensor<i32>) -> tensor<8xf32>
%4 = "tf.Mul"(%0, %3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<8xf32>) -> tensor<8xf32>
%5 = "tf.TensorListSetItem"(%arg2, %arg1, %4) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<8xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
%6 = "tf.GatherV2"(%cst, %arg1, %cst_2) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x8xf32>, tensor<i32>, tensor<i32>) -> tensor<8xf32>
%7 = "tf.Mul"(%0, %6) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<8xf32>) -> tensor<8xf32>
%8 = "tf.TensorListSetItem"(%arg3, %arg1, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<8xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %1, %2, %5, %8, %arg4 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
}
// CHECK-LABEL: tf.MultiListWhileRegion_body/MapFnBody
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: return
// CHECK-LABEL: tf.MultiListWhileRegion_cond
func.func private @tf.MultiListWhileRegion_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>, %arg4: tensor<?xf32>) -> tensor<i1> {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
return %2 : tensor<i1>
}
// CHECK-LABEL: multilist_serving
func.func private @multilist_serving(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<2x8xf32>, tensor<2x8xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<8> : tensor<1xi32>} : () -> tensor<1xi32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
// CHECK: TensorListReserve
%0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
// CHECK-NEXT: tf_mlrt.tf_map_fn
%1:5 = "tf.While"(%cst, %cst, %0, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 8 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @tf.MultiListWhileRegion_body, cond = @tf.MultiListWhileRegion_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
// CHECK-NEXT: TensorListStack
%2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>
%3 = "tf.TensorListStack"(%1#3, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>
return %3, %2 : tensor<2x8xf32>, tensor<2x8xf32>
}
// -----
// Convert a while with multiple tensor array to map_fn
// CHECK-LABEL: tf.WhileRegion1_body(
func.func private @tf.WhileRegion1_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<i32>, %arg5: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg6: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg7: tensor<*xi32>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.AddV2"(%arg0, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.AddV2"(%arg1, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%3 = "tf.RaggedTensorToVariant"(%arg7) {RAGGED_RANK = 0 : i64, Tsplits = i64, Tvalues = i32, batched_input = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>) -> tensor<!tf_type.variant>
%4 = "tf.TensorArrayWriteV3"(%arg5, %arg1, %3, %arg2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<!tf_type.variant>, tensor<f32>) -> tensor<f32>
%5 = "tf.RaggedTensorToVariant"(%arg7) {RAGGED_RANK = 0 : i64, Tsplits = i64, Tvalues = f32, batched_input = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>) -> tensor<!tf_type.variant>
%6 = "tf.TensorArrayWriteV3"(%arg6, %arg1, %5, %arg3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<!tf_type.variant>, tensor<f32>) -> tensor<f32>
return %0, %1, %4, %6, %arg4, %arg5, %arg6, %arg7 : tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>
}
// CHECK-LABEL: func.func private @"tf.WhileRegion1_body/MapFnBody"(%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: !mlrt.future, %arg3: !mlrt.promise, %arg4: tensor<i32>, %arg5: tensor<i32>, %arg6: tensor<i32>, %arg7: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg8: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg9: tensor<*xi32>) attributes {tfrt.cost_threshold = 4294967295 : i64}
// CHECK: [[result_0:%.*]] = "tf.RaggedTensorToVariant"
// CHECK: [[result_1:%.*]] = "tf.RaggedTensorToVariant"
// CHECK-NEXT: [[flow_in_0:%.*]] = "tf_mlrt.tf_await"(%arg0) : (!mlrt.future) -> tensor<f32>
// CHECK-NEXT: [[flow_in_1:%.*]] = "tf_mlrt.tf_await"(%arg2) : (!mlrt.future) -> tensor<f32>
// CHECK-NEXT: [[flow_out_0:%.*]] = "tf.TensorArrayWriteV3"(%arg7, %arg5, [[result_0]], [[flow_in_0]])
// CHECK-NEXT: [[flow_out_1:%.*]] = "tf.TensorArrayWriteV3"(%arg8, %arg5, [[result_1]], [[flow_in_1]])
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[flow_out_0]]) : (!mlrt.promise, tensor<f32>) -> ()
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg3, [[flow_out_1]]) : (!mlrt.promise, tensor<f32>) -> ()
// CHECK-NEXT: return
// CHECK-LABEL: tf.WhileRegion1_cond
func.func private @tf.WhileRegion1_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<i32>, %arg5: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg6: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg7: tensor<*xi32>) -> (tensor<i1>) {
%0 = "tf.Less"(%arg0, %arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%1 = "tf.Less"(%arg1, %arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
%3 = "tf.ToBool"(%2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi1>) -> tensor<i1>
return %3 : tensor<i1>
}
// CHECK-LABEL: func.func private @tf.WhileRegion2_body(
func.func private @tf.WhileRegion2_body(%arg0: tensor<*xi32>) -> (tensor<?x!tf_type.variant>, tensor<?x!tf_type.variant>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[-1, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%max_iter = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<4> : tensor<i32>} : () -> tensor<i32>
// CHECK: "tf.TensorArrayV3"
%handle_12, %flow_13 = "tf.TensorArrayV3"(%max_iter) {device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = !tf_type.variant, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<f32>)
// CHECK: "tf.TensorArrayV3"
%handle_14, %flow_15 = "tf.TensorArrayV3"(%max_iter) {device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = !tf_type.variant, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<f32>)
// CHECK: tf_mlrt.tf_map_fn
// CHECK-SAME: {body_fn = @"tf.WhileRegion1_body/MapFnBody", num_tensor_list_or_flow_in = 2 : i32}
%4:8 = "tf.While"(%cst_0, %cst_0, %flow_13, %flow_15, %max_iter, %handle_12, %handle_14, %arg0) {body = @tf.WhileRegion1_body, cond = @tf.WhileRegion1_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>)
// CHECK: TensorArrayGatherV3
%5 = "tf.TensorArrayGatherV3"(%handle_12, %1, %4#2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<f32>) -> tensor<?x!tf_type.variant>
// CHECK: TensorArrayGatherV3
%6 = "tf.TensorArrayGatherV3"(%handle_14, %2, %4#3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<f32>) -> tensor<?x!tf_type.variant>
return %5, %6 : tensor<?x!tf_type.variant>, tensor<?x!tf_type.variant>
}
// -----
// Test a while to map_fn conversion in which tensor array is used instead of
// tensor list and the tensor array size and the number of iterations are bounded
// by separate constants of the same value.
// CHECK-LABEL: map2/while/LoopCond_body
func.func private @"map2/while/LoopCond_body"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) {
%outputs = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_4 = "tf.Identity"(%arg0) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_6 = "tf.AddV2"(%outputs_4, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%outputs_8 = "tf.Identity"(%arg1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_10 = "tf.AddV2"(%outputs_8, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%outputs_12 = "tf.Identity"(%arg2) {device = ""} : (tensor<f32>) -> tensor<f32>
%outputs_14 = "tf.TensorArrayReadV3"(%arg4, %outputs_8, %arg5) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<*xi32>, tensor<f32>) -> tensor<*x!tf_type.string>
%outputs_16 = "tf.DecodeJpeg"(%outputs_14) {acceptable_fraction = 1.000000e+00 : f32, channels = 3 : i64, dct_method = "INTEGER_FAST", device = "", fancy_upscaling = true, ratio = 1 : i64, try_recover_truncated = false} : (tensor<*x!tf_type.string>) -> tensor<?x?x3xui8>
%outputs_18 = "tf.ExpandDims"(%outputs_16, %outputs_0) {device = ""} : (tensor<?x?x3xui8>, tensor<i32>) -> tensor<1x?x?x3xui8>
%outputs_20 = "tf.ResizeBilinear"(%outputs_18, %outputs) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x?x?x3xui8>, tensor<2xi32>) -> tensor<1x224x224x3xf32>
%outputs_22 = "tf.Squeeze"(%outputs_20) {device = "", squeeze_dims = [0]} : (tensor<1x224x224x3xf32>) -> tensor<224x224x3xf32>
%outputs_24 = "tf.Cast"(%outputs_22) {Truncate = false, device = ""} : (tensor<224x224x3xf32>) -> tensor<224x224x3xui8>
%outputs_26 = "tf.TensorArrayWriteV3"(%arg6, %outputs_8, %outputs_24, %outputs_12) {device = ""} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<*xi32>, tensor<224x224x3xui8>, tensor<f32>) -> tensor<f32>
return %outputs_6, %outputs_10, %outputs_26, %arg3, %arg4, %arg5, %arg6: tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>
}
// CHECK-LABEL: @"map2/while/LoopCond_body/MapFnBody"
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.TensorArrayReadV3
// CHECK-NEXT: tf.DecodeJpeg
// CHECK-NEXT: tf.ExpandDims
// CHECK-NEXT: tf.ResizeBilinear
// CHECK-NEXT: tf.Squeeze
// CHECK-NEXT: tf.Cast
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorArrayWriteV3
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: return
// CHECK-LABEL: map2/while/LoopCond_cond
func.func private @"map2/while/LoopCond_cond"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> tensor<i1> {
%cst = "tf.Const"() {value = dense<224> : tensor<i32>} : () -> tensor<i32>
%outputs = "tf.Less"(%arg0, %cst) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%outputs_0 = "tf.Less"(%arg1, %cst) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%outputs_2 = "tf.LogicalAnd"(%outputs, %outputs_0) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
%outputs_4 = "tf.ToBool"(%outputs_2) : (tensor<*xi1>) -> tensor<i1>
return %outputs_4 : tensor<i1>
}
//CHECK-LABEL: map2_while_test
func.func private @map2_while_test(%arg0: tensor<?x!tf_type.string>) -> tensor<?x224x224x3xui8> {
// CHECK-NEXT: tf.Const
%outputs = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<1xi32>
// CHECK-NEXT: [[max_iter:%.*]] = "tf.Const"
%cst_0 = "tf.Const"() {value = dense<224> : tensor<i32>} : () -> tensor<i32>
%cst_1 = "tf.Const"() {value = dense<256> : tensor<i32>} : () -> tensor<i32>
%outputs_2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_4 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_6 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_8 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_10 = "tf.Shape"(%arg0) {device = ""} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
// CHECK: tf.Range
%outputs_14 = "tf.Range"(%outputs_2, %cst_0, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
// CHECK-NEXT: tf.TensorArrayV3
%outputs_16:2 = "tf.TensorArrayV3"(%cst_0) {clear_after_read = true, device = "", dtype = !tf_type.string, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>)
// CHECK-NEXT: tf.TensorArrayScatterV3
%outputs_18 = "tf.TensorArrayScatterV3"(%outputs_16#0, %outputs_14, %arg0, %outputs_16#1) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<?xi32>, tensor<?x!tf_type.string>, tensor<f32>) -> tensor<f32>
// CHECK-NEXT: tf.Range
%outputs_20 = "tf.Range"(%outputs_2, %cst_0, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
// CHECK-NEXT: [[tensor_array:%.*]], [[flow_in:%.*]] = "tf.TensorArrayV3"
%outputs_22:2 = "tf.TensorArrayV3"(%cst_0) {clear_after_read = true, device = "", dtype = ui8, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<f32>)
// CHECK-NEXT: tf_mlrt.tf_map_fn
// CHECK-SAME: ([[max_iter]], [[flow_in]], %cst_1
// CHECK-SAME: {body_fn = @"map2/while/LoopCond_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
// CHECK-NOT: tf.While
%outputs_24:7 = "tf.While"(%outputs, %outputs, %outputs_22#1, %cst_1, %outputs_16#0, %outputs_18, %outputs_22#0) {_xla_propagate_compile_time_consts = true, body = @"map2/while/LoopCond_body", cond = @"map2/while/LoopCond_cond", device = "", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>)
// CHECK-NEXT: tf.TensorArrayGatherV3
%outputs_26 = "tf.TensorArrayGatherV3"(%outputs_22#0, %outputs_20, %outputs_24#2) {device = "", element_shape = #tf_type.shape<224x224x3>} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<?xi32>, tensor<f32>) -> tensor<?x224x224x3xui8>
return %outputs_26 : tensor<?x224x224x3xui8>
}
// -----
// Test a nest while in which the while body is after the usage.
// CHECK-LABEL: nested_while
func.func @nested_while(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<16x16x?xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[16, -1]> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
// CHECK: tf.TensorListReserve
%0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
// CHECK-NEXT: tf_mlrt.tf_map_fn
%1:4 = "tf.While"(%cst, %cst, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @tf.NestedWhileRegion1_body, cond = @tf.NestedWhileRegion1_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
%2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 16 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<2xi32>) -> tensor<16x16x?xf32>
return %2 : tensor<16x16x?xf32>
}
// CHECK-LABEL: tf.NestedWhileRegion1_body
func.func private @tf.NestedWhileRegion1_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi32>} : () -> tensor<16xi32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%cst_3 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
%cst_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.TensorListReserve"(%cst_4, %cst_3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
%1 = "tf.AddV2"(%arg0, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.AddV2"(%arg1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%3 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<16xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%4 = "tf.Cast"(%3) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
%5 = "tf.Mul"(%arg3, %4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
%6:4 = "tf.While"(%cst_0, %cst_0, %0, %5) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @tf.NestedWhileRegion_body, cond = @tf.NestedWhileRegion_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
%7 = "tf.TensorListStack"(%6#2, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 16 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<16x?xf32>
%8 = "tf.TensorListSetItem"(%arg2, %arg1, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<16x?xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %1, %2, %8, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
}
//CHECK-LABEL: @"tf.NestedWhileRegion1_body/MapFnBody"(%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<?xf32>)
// CHECK: tf.TensorListReserve
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Cast
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf_mlrt.tf_map_fn
// CHECK-NEXT: tf.TensorListStack
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: tf_mlrt.tf_promise
// CHECK-NEXT: return
// CHECK-LABEL: tf.NestedWhileRegion1_cond
func.func private @tf.NestedWhileRegion1_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
return %2 : tensor<i1>
}
// CHECK-LABEL: tf.NestedWhileRegion_body
func.func private @tf.NestedWhileRegion_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi32>} : () -> tensor<16xi32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.AddV2"(%arg0, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.AddV2"(%arg1, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.GatherV2"(%cst_0, %arg1, %cst) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<16xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%3 = "tf.Cast"(%2) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
%4 = "tf.Mul"(%arg3, %3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
%5 = "tf.TensorListSetItem"(%arg2, %arg1, %4) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<?xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %0, %1, %5, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
}
// CHECK-LABEL: tf.NestedWhileRegion_body/MapFnBody
// CHECK: tf.AddV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Cast
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: "tf_mlrt.tf_promise
// CHECK-NEXT: return
// CHECK-LABEL: tf.NestedWhileRegion_cond
func.func private @tf.NestedWhileRegion_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
return %2 : tensor<i1>
}
// -----
// Test a while to map_fn conversion is skipped if the tensor list cannot be found in the current function body.
// CHECK-LABEL: map/while_cond
func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
return %2 : tensor<i1>
}
// CHECK-LABEL: map/while_body
func.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00, 9.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_3 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
%cst_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%0 = "tf.AddV2"(%arg0, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.Mul"(%arg3, %cst_3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<9xf32>) -> tensor<9xf32>
%2 = "tf.Reshape"(%1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
%3 = "tf.AddV2"(%arg1, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%4 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%5 = "tf.Cast"(%4) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
%6 = "tf.Mul"(%5, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<9xf32>) -> tensor<9xf32>
%7 = "tf.Reshape"(%6, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
%8 = "tf.MatMul"(%2, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%9 = "tf.MatrixDeterminant"(%8) {T = f32, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<f32>
%10 = "tf.TensorListSetItem"(%arg2, %arg1, %9) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %0, %3, %10, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
}
//CHECK-LABEL: @func
func.func @func(%arg0: tensor<?xf32>, %arg1: tensor<!tf_type.variant<tensor<*xf32>>>) -> tensor<3xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input:0", outputs = "PartitionedCall:0"}} {
%cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
// CHECK-NOT: tf_map_fn
%1:4 = "tf.While"(%cst, %cst, %arg1, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
%2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 3 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<0xi32>) -> tensor<3xf32>
return %2 : tensor<3xf32>
}
// -----
// Test a while to map_fn conversion in which a tf.StopGradient is inserted to consume the while result.
// CHECK-LABEL: @while_map_while_body_884030
func.func private @while_map_while_body_884030(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, %arg4: tensor<!tf_type.variant<tensor<?x?x1xui8>>> {tf._user_specified_name = "while/map/TensorArrayUnstack/TensorListFromTensor"}) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>) {
%cst = "tf.Const"() <{value = dense<[-1, -1, 1]> : tensor<3xi32>}> : () -> tensor<3xi32>
%cst_0 = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
%0 = "tf.AddV2"(%arg2, %cst_0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%1 = "tf.Identity"(%0) : (tensor<i32>) -> tensor<i32>
%2 = "tf.TensorListGetItem"(%arg4, %arg2, %cst) : (tensor<!tf_type.variant<tensor<?x?x1xui8>>>, tensor<i32>, tensor<3xi32>) -> tensor<?x?x1xui8>
%3 = "tf.EncodePng"(%2) <{compression = -1 : i64}> : (tensor<?x?x1xui8>) -> tensor<!tf_type.string>
%4 = "tf.TensorListSetItem"(%arg3, %arg2, %3) <{resize_if_index_out_of_bounds = false}> : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<i32>, tensor<!tf_type.string>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
%5 = "tf.Identity"(%4) : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
%6 = "tf.AddV2"(%arg0, %cst_0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%7 = "tf.Identity"(%6) : (tensor<i32>) -> tensor<i32>
%8 = "tf.Identity"(%arg1) : (tensor<i32>) -> tensor<i32>
return %7, %8, %1, %5, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>
}
// CHECK-LABEL: while_map_while_body_884030/MapFnBody
// CHECK: tf.AddV2
// CHECK-NEXT: tf.TensorListGetItem
// CHECK-NEXT: tf.EncodePng
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf_await
// CHECK-NEXT: tf.TensorListSetItem
// CHECK-NEXT: tf_promise
// CHECK-LABEL: @while_map_while_cond_884020
func.func private @while_map_while_cond_884020(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, %arg4: tensor<!tf_type.variant<tensor<?x?x1xui8>>>) -> tensor<i1> {
%cst = "tf.Const"() <{value = dense<11> : tensor<i32>}> : () -> tensor<i32>
%0 = "tf.Less"(%arg2, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i1>
%1 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
%2 = "tf.LogicalAnd"(%1, %0) : (tensor<i1>, tensor<i1>) -> tensor<i1>
%3 = "tf.Identity"(%2) : (tensor<i1>) -> tensor<i1>
return %3 : tensor<i1>
}
// CHECK-LABEL: @main
// CHECK: tf.Cast
// CHECK-NEXT: tf.TensorListReserve
// CHECK-NEXT: tf.Transpose
// CHECK-NEXT: tf.TensorListFromTensor
// CHECK-NEXT: tf_mlrt.tf_map_fn
// CHECK-SAME: {body_fn = @"while_map_while_body_884030/MapFnBody", num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<i32>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
// CHECK-NEXT: tf.StopGradient
// CHECK-NEXT: tf.TensorListStack
func.func @main(%arg0: tensor<1x?x?x11xf32>) -> tensor<11x!tf_type.string> {
%cst_0 = "tf.Const"() <{value = dense<[3, 1, 2, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
%cst_10 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%cst_11 = "tf.Const"() <{value = dense<2> : tensor<i32>}> : () -> tensor<i32>
%cst_12 = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
%cst_13 = "tf.Const"() <{value = dense<[-1, -1, 1]> : tensor<3xi32>}> : () -> tensor<3xi32>
%cst_14 = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32>
%cst_15 = "tf.Const"() <{value = dense<-1> : tensor<i32>}> : () -> tensor<i32>
%cst_16 = "tf.Const"() <{value = dense<11> : tensor<i32>}> : () -> tensor<i32>
%92 = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x?x?x11xf32>) -> tensor<1x?x?x11xui8>
%0 = "tf.TensorListReserve"(%cst_15, %cst_16) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
%93 = "tf.Transpose"(%92, %cst_0) : (tensor<1x?x?x11xui8>, tensor<4xi32>) -> tensor<11x?x?x1xui8>
%94 = "tf.TensorListFromTensor"(%93, %cst_13) : (tensor<11x?x?x1xui8>, tensor<3xi32>) -> tensor<!tf_type.variant<tensor<?x?x1xui8>>>
%95:5 = "tf.While"(%cst_10, %cst_16, %cst_10, %0, %94) <{body = @while_map_while_body_884030, cond = @while_map_while_cond_884020, is_stateless = true, parallel_iterations = 16 : i64, shape_invariant}> {T = [i32, i32, i32, !tf_type.variant, !tf_type.variant], _lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], device = "", output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>]} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>)
%96 = "tf.StopGradient"(%95#3) : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
%97 = "tf.TensorListStack"(%96, %cst_14) <{num_elements = 11 : i64}> : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<0xi32>) -> tensor<11x!tf_type.string>
return %97 : tensor<11x!tf_type.string>
}
// -----
// Test a while to map_fn conversion where constants are hoisted to resources.
// CHECK-LABEL: hoisted_constants_map_while_cond
func.func private @hoisted_constants_map_while_cond(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>) -> tensor<i1> {
%0 = "tf.Less"(%arg0, %arg3) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%1 = "tf.Less"(%arg1, %arg3) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%2 = "tf.LogicalAnd"(%0, %1) : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
%3 = "tf.ToBool"(%2) : (tensor<*xi1>) -> tensor<i1>
return %3 : tensor<i1>
}
// CHECK-LABEL: hoisted_constants_map_while_body
func.func private @hoisted_constants_map_while_body(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) {
// Resource 1 is constant 1
%c1 = "tf._TfrtGetResource"() <{container = [""], indices = [1 : i64], shared_name = [""]}> : () -> tensor<i32>
%0 = "tf.AddV2"(%arg0, %c1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%1 = "tf.AddV2"(%arg1, %c1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
// Dummy work
%dummy = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
%2 = "tf.TensorListSetItem"(%arg2, %arg1, %dummy) {resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<*xi32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
return %0, %1, %2, %arg3 : tensor<*xi32>, tensor<*xi32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>
}
// CHECK-LABEL: @hoisted_constants_main
// CHECK: tf.TensorListReserve
// CHECK-NEXT: tf_mlrt.tf_map_fn
// CHECK-SAME: {body_fn = @"hoisted_constants_map_while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.TensorListStack
// CHECK-NOT: tf.While
func.func @hoisted_constants_main() -> tensor<3xf32> {
// Resource 0 is constant 0
%c0 = "tf._TfrtGetResource"() <{container = [""], indices = [0 : i64], shared_name = [""]}> : () -> tensor<i32>
// Resource 2 is constant 3 (max iterations)
%c3 = "tf._TfrtGetResource"() <{container = [""], indices = [2 : i64], shared_name = [""]}> : () -> tensor<i32>
%c_neg1 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%tensor_list = "tf.TensorListReserve"(%c_neg1, %c3) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
%0:4 = "tf.While"(%c0, %c0, %tensor_list, %c3) {body = @hoisted_constants_map_while_body, cond = @hoisted_constants_map_while_cond, is_stateless = true, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>)
%cst_shape = "tf.Const"() {value = dense<[3]> : tensor<1xi32>} : () -> tensor<1xi32>
%stacked = "tf.TensorListStack"(%0#2, %cst_shape) {num_elements = 3 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<3xf32>
return %stacked : tensor<3xf32>
}
// This function defines the resources.
func.func private @_tfrt_resource_init() {
%c0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
"tf._TfrtSetResource"(%c0) {index = 0 : i64} : (tensor<i32>) -> ()
%c1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
"tf._TfrtSetResource"(%c1) {index = 1 : i64} : (tensor<i32>) -> ()
%c3 = "tf.Const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
"tf._TfrtSetResource"(%c3) {index = 2 : i64} : (tensor<i32>) -> ()
return
}
tensorflow/compiler/mlir/tfrt/tests/optimize.mlir
+63
View File
@@ -0,0 +1,63 @@
// 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 -optimize-tf-for-tfrt -split-input-file -verify-diagnostics %s | FileCheck %s
// CHECK-LABEL: @fold_device_index
func.func @fold_device_index() -> tensor<i32> {
// CHECK-NOT: tf.DeviceIndex
// CHECK: tf.Const
// CHECK-SAME: value = dense<1> : tensor<i32>
%0 = "tf.DeviceIndex"() {device = "/device:CPU:0", device_names = ["GPU", "CPU"]} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: @not_fold_device_index
func.func @not_fold_device_index() -> tensor<i32> {
// CHECK-NOT: tf.Const
// CHECK: tf.DeviceIndex
%0 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: @eliminate_multinomial
func.func @eliminate_multinomial(%0: tensor<*xf32>, %1: tensor<*xi32>) -> (tensor<*xi64>, tensor<*xi64>) {
// CHECK-NEXT: tf.Multinomial
// CHECK-NEXT: return
%2 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
%3 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
func.return %2, %3 : tensor<*xi64>, tensor<*xi64>
}
// -----
// CHECK-LABEL: @not_eliminate_multinomial
func.func @not_eliminate_multinomial(%0: tensor<*xf32>, %1: tensor<*xi32>) -> (tensor<*xi64>, tensor<*xi64>) {
// CHECK-NEXT: tf.Multinomial
// CHECK-SAME: seed = 0
// CHECK-NEXT: tf.Multinomial
// CHECK-SAME: seed = 1
// CHECK-NEXT: tf.Multinomial
// CHECK-SAME: seed = 0
// CHECK-NEXT: return
%2 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
%3 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 1 : i64, seed2 = 1 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
%4 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
%5 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
func.return %2, %3 : tensor<*xi64>, tensor<*xi64>
}
tensorflow/compiler/mlir/tfrt/tests/pack_inputs.mlir
+34
View File
@@ -0,0 +1,34 @@
// 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 -pack-inputs="slices=1,0,24,2,24,24" | FileCheck %s
// Tests basic contiguous slice packing:
// - Input 0 is preserved (not in slice list).
// - Inputs 1 and 2 are packed into a single contiguous i8 buffer with explicit slices.
module {
// CHECK-LABEL: func.func @main
// CHECK-SAME: (%arg0: tensor<10x10xf32>, %arg1: tensor<48xi8>) -> (tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>)
func.func @main(%arg0: tensor<10x10xf32>, %arg1: tensor<3x2xf32>, %arg2: tensor<3x2xf32>) -> (tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>) attributes {ifrt.function} {
// CHECK: %[[SLICE0:.*]] = stablehlo.slice %arg1 [0:24] : (tensor<48xi8>) -> tensor<24xi8>
// CHECK: %[[RESHAPE0:.*]] = stablehlo.reshape %[[SLICE0]] : (tensor<24xi8>) -> tensor<3x2x4xi8>
// CHECK: %[[BITCAST0:.*]] = stablehlo.bitcast_convert %[[RESHAPE0]] : (tensor<3x2x4xi8>) -> tensor<3x2xf32>
// CHECK: %[[SLICE1:.*]] = stablehlo.slice %arg1 [24:48] : (tensor<48xi8>) -> tensor<24xi8>
// CHECK: %[[RESHAPE1:.*]] = stablehlo.reshape %[[SLICE1]] : (tensor<24xi8>) -> tensor<3x2x4xi8>
// CHECK: %[[BITCAST1:.*]] = stablehlo.bitcast_convert %[[RESHAPE1]] : (tensor<3x2x4xi8>) -> tensor<3x2xf32>
// CHECK: return %arg0, %[[BITCAST0]], %[[BITCAST1]] : tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>
return %arg0, %arg1, %arg2 : tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/pack_inputs_exclude.mlir
+35
View File
@@ -0,0 +1,35 @@
// 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 -pack-inputs="slices=1,0,24,2,28,16" | FileCheck %s
// Tests advanced slice packing with gaps/padding and unmerged variables:
// - Input 0 is preserved.
// - Inputs 1 and 2 are packed into a single i8 buffer with a 4-byte gap (start 0 and start 28).
// - Inputs 3 and 4 are skipped and remain unchanged.
module {
// CHECK-LABEL: func.func @main
// CHECK-SAME: (%arg0: tensor<10x10xf32>, %arg1: tensor<5x2xf32>, %arg2: tensor<6x1xf32>, %arg3: tensor<44xi8>) -> (tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>)
func.func @main(%arg0: tensor<10x10xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<4x1xf32>, %arg3: tensor<5x2xf32>, %arg4: tensor<6x1xf32>) -> (tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>) attributes {ifrt.function} {
// CHECK: %[[SLICE0:.*]] = stablehlo.slice %arg3 [0:24] : (tensor<44xi8>) -> tensor<24xi8>
// CHECK: %[[RESHAPE0:.*]] = stablehlo.reshape %[[SLICE0]] : (tensor<24xi8>) -> tensor<2x3x4xi8>
// CHECK: %[[BITCAST0:.*]] = stablehlo.bitcast_convert %[[RESHAPE0]] : (tensor<2x3x4xi8>) -> tensor<2x3xf32>
// CHECK: %[[SLICE1:.*]] = stablehlo.slice %arg3 [28:44] : (tensor<44xi8>) -> tensor<16xi8>
// CHECK: %[[RESHAPE1:.*]] = stablehlo.reshape %[[SLICE1]] : (tensor<16xi8>) -> tensor<4x1x4xi8>
// CHECK: %[[BITCAST1:.*]] = stablehlo.bitcast_convert %[[RESHAPE1]] : (tensor<4x1x4xi8>) -> tensor<4x1xf32>
// CHECK: return %arg0, %[[BITCAST0]], %[[BITCAST1]], %arg1, %arg2 : tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>
return %arg0, %arg1, %arg2, %arg3, %arg4 : tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op.mlir
+123
View File
@@ -0,0 +1,123 @@
// 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 -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=2 tfrt-min-max-enqueued-batches=3 tfrt-batch-padding-policy=PAD_UP" %s | FileCheck %s --dump-input=always
// -----
// The num_batch_threads is lowered bound to 2 from the original attribute of 1
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [6]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 100000 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: max_batch_size = 6 : i64
// CHECK-SAME: max_enqueued_batches = 10 : i64
// CHECK-SAME: num_batch_threads = 2 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
// -----
// The num_batch_threads remains 3 (the same as the original attribute)
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [6]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 100000 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: max_batch_size = 6 : i64
// CHECK-SAME: max_enqueued_batches = 10 : i64
// CHECK-SAME: num_batch_threads = 3 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 3 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
// -----
// The max_enqueued_batches is changed to 3 from the original attribute of 2
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [6]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 100000 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: max_batch_size = 6 : i64
// CHECK-SAME: max_enqueued_batches = 3 : i64
// CHECK-SAME: num_batch_threads = 2 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
// -----
// The max_enqueued_batches remains 10 (the same as the original attribute)
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [6]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 100000 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: max_batch_size = 6 : i64
// CHECK-SAME: max_enqueued_batches = 10 : i64
// CHECK-SAME: num_batch_threads = 3 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 3 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op_global_prioritization.mlir
+81
View File
@@ -0,0 +1,81 @@
// 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 -split-input-file -tfrt-reconfig-batch-op="tfrt-batch-queue-global-prioritization-num-threads=4" %s | FileCheck %s --dump-input=always
// -----
// The num_batch_threads is updated to 4 from the original attribute of 2,
// a mixed_priority_policy is set along with all low priority batching params
// being copied from the high priority batching params since no low priority
// settings are provided.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [1, 2]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 2 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: enable_priority_aware_batch_scheduler = true
// CHECK-SAME: max_batch_size = 2 : i64
// CHECK-SAME: max_enqueued_batches = 2 : i64
// CHECK-SAME: num_batch_threads = 4 : i64
// CHECK-SAME: num_warmup_batch_threads = {{[0-9]+}} : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [1, 2], batch_padding_policy = "PAD_UP", batch_timeout_micros = 2 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 2 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
// -----
// Same as first test, but low_priority_* parameters are already provided
// so not overriden.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [1, 2]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 2 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: enable_priority_aware_batch_scheduler = true
// CHECK-SAME: low_priority_allowed_batch_sizes = [1, 10]
// CHECK-SAME: low_priority_batch_timeout_micros = 7 : i64
// CHECK-SAME: low_priority_max_batch_size = 8 : i64
// CHECK-SAME: low_priority_max_enqueued_batches = 9 : i64
// CHECK-SAME: max_batch_size = 2 : i64
// CHECK-SAME: max_enqueued_batches = 2 : i64
// CHECK-SAME: num_batch_threads = 4 : i64
// CHECK-SAME: num_warmup_batch_threads = {{[0-9]+}} : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [1, 2], batch_padding_policy = "PAD_UP", batch_timeout_micros = 2 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 2 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/", low_priority_batch_timeout_micros = 7 : i64, low_priority_max_batch_size = 8 : i64, low_priority_max_enqueued_batches = 9 : i64, low_priority_allowed_batch_sizes = [1, 10]} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op_params.mlir
+64
View File
@@ -0,0 +1,64 @@
// 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 -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=4 tfrt-min-max-enqueued-batches=4 tfrt-num-batch-threads=3 tfrt-max-batch-size=3 tfrt-batch-timeout-micros=3 tfrt-allowed-batch-sizes=3,4 tfrt-max-enqueued-batches=3 tfrt-enable-large-batch-splitting=true tfrt-mixed-priority-batching-policy=priority_merge tfrt-low-priority-max-batch-size=5 tfrt-low-priority-batch-timeout-micros=5 tfrt-low-priority-allowed-batch-sizes=5,6 tfrt-low-priority-max-enqueued-batches=5 tfrt-num-warmup-batch-threads=2" %s | FileCheck %s --dump-input=always
// -----
// The num_batch_threads is updated to 3 from the original attribute of 2,
// overriding the min_num_batch_threads of 4.
// The max_batch_size is updated to 3 from the original attribute of 2.
// The batch_timeout_micros is updated to 3 from the original attribute of 2.
// The allowed_batch_sizes is updated to [3, 4] from the original attribute of
// [1, 2].
// The max_enqueued_batches is updated to 3 from the original attribute of 2,
// overriding the min_max_enqueued_batches of 4.
// The enable_large_batch_splitting is updated to true from the original
// attribute of false.
// The mixed_priority_policy is updated to "priority_merge" from the original
// attribute of "".
// The low_priority_max_batch_size is updated to 5 from the original attribute of 2.
// The low_priority_batch_timeout_micros is updated to 5 from the original attribute of 2.
// The low_priority_allowed_batch_sizes is updated to [5, 6] from the original attribute of [1, 2].
// The low_priority_max_enqueued_batches is updated to 5 from the original attribute of 2.
// The num_warmup_batch_threads is updated to 2 from the original attribute of 0.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [3, 4]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 3 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = true
// CHECK-SAME: low_priority_allowed_batch_sizes = [5, 6]
// CHECK-SAME: low_priority_batch_timeout_micros = 5 : i64
// CHECK-SAME: low_priority_max_batch_size = 5 : i64
// CHECK-SAME: low_priority_max_enqueued_batches = 5 : i64
// CHECK-SAME: max_batch_size = 3 : i64
// CHECK-SAME: max_enqueued_batches = 3 : i64
// CHECK-SAME: mixed_priority_policy = "priority_merge"
// CHECK-SAME: num_batch_threads = 3 : i64
// CHECK-SAME: num_warmup_batch_threads = 2 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [1, 2], batch_padding_policy = "PAD_UP", batch_timeout_micros = 2 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, low_priority_allowed_batch_sizes = [1, 2], low_priority_batch_timeout_micros = 2 : i64, low_priority_max_batch_size = 2 : i64, low_priority_max_enqueued_batches = 2 : i64, max_batch_size = 2 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op_skip_pass.mlir
+43
View File
@@ -0,0 +1,43 @@
// 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 -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=0 tfrt-min-max-enqueued-batches=0" %s | FileCheck %s --dump-input=always
// -----
// Confirm no batch options change when the tfrt-reconfig-batch-op flag is
// empty.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
// CHECK: "tf.BatchFunction"
// CHECK-SAME: allowed_batch_sizes = [6]
// CHECK-SAME: batch_padding_policy = "PAD_UP"
// CHECK-SAME: batch_timeout_micros = 100000 : i64
// CHECK-SAME: batching_queue = ""
// CHECK-SAME: container = ""
// CHECK-SAME: enable_large_batch_splitting = false
// CHECK-SAME: max_batch_size = 6 : i64
// CHECK-SAME: max_enqueued_batches = 10 : i64
// CHECK-SAME: num_batch_threads = 1 : i64
// CHECK-SAME: shared_name = "batch/"
%1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
tensorflow/compiler/mlir/tfrt/tests/remove_device_attribute.mlir
+22
View File
@@ -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.
// ==============================================================================
// RUN: tf-tfrt-opt -tfrt-remove-device-attribute %s | FileCheck %s
func.func @test(%arg0: !tfrt.chain, %arg1: !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle) {
%0 = corert.get_op_handler %arg0 "cpu"
// CHECK: %[[RESULT:.*]] = corert.executeop(%[[ARG_0:.*]]) "tf.MatMul"(%[[ARG_1:.*]], %[[ARG_1]]) {T = f32, transpose_a = false, transpose_b = false} : 1
%1 = corert.executeop(%0) "tf.MatMul"(%arg1, %arg1) {T = f32, device = "cpu", transpose_a = false, transpose_b = false} : 1
tfrt.return %arg0, %1 : !tfrt.chain, !corert.tensorhandle
}
tensorflow/compiler/mlir/tfrt/tests/runtime_lowering_gpu.mlir
+29
View File
@@ -0,0 +1,29 @@
// 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 -tfrt-lower-cluster-to-runtime-ops-non-tpu -split-input-file -verify-diagnostics %s | FileCheck %s
module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:GPU:0"]} {
// CHECK-LABEL: @converts_cluster
func.func @converts_cluster() {
// CHECK: "tf.XlaLaunch"()
"tf_device.cluster_func"() {_xla_compile_device_type = "GPU", _replication_info = "cluster0", func = @empty_func, num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = [], input_sharding_configuration = [], output_sharding_configuration = [], use_spmd_for_xla_partitioning = false} : () -> ()
func.return
}
func.func @empty_func() {
func.return
}
}
tensorflow/compiler/mlir/tfrt/tests/runtime_lowering_tpu.mlir
+43
View File
@@ -0,0 +1,43 @@
// 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 -tfrt-lower-cluster-to-runtime-ops-tpu -split-input-file -verify-diagnostics %s | FileCheck %s
module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
// CHECK-LABEL: @converts_cluster
func.func @converts_cluster() {
// CHECK: %0:2 = "tf_device.launch"() <{{.*}}> ({
// CHECK: %compilation_status, %program = "tf._TPUCompileMlir"()
"tf_device.cluster_func"() {_xla_compile_device_type = "TPU", _replication_info = "cluster0", func = @empty_func, num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = [], input_sharding_configuration = [], output_sharding_configuration = [], use_spmd_for_xla_partitioning = false} : () -> ()
func.return
}
func.func @empty_func() {
func.return
}
}
// -----
module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
func.func @missing_num_cores_per_replica() {
// expected-error@+1 {{requires attribute 'num_cores_per_replica'}}
"tf_device.cluster_func"() {_xla_compile_device_type = "TPU", _replication_info = "cluster0", func = @empty_func, step_marker_location = "STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP", topology = "", device_assignment = [], input_sharding_configuration = [], output_sharding_configuration = [], use_spmd_for_xla_partitioning = false} : () -> ()
func.return
}
func.func @empty_func() {
func.return
}
}
tensorflow/compiler/mlir/tfrt/tests/saved_model/BUILD
+47
View File
@@ -0,0 +1,47 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
tf_cc_test(
name = "saved_model_test",
srcs = ["saved_model_test.cc"],
data = [
"testdata/test.mlir",
"testdata/xla_launch.mlir",
"testdata/xla_launch_xla_reduce_window.mlir",
],
tags = ["no_oss"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tfrt:import_model",
"//tensorflow/compiler/mlir/tfrt:saved_model",
"//tensorflow/compiler/mlir/tfrt:tfrt_compile_options",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:resource_loader",
"//tensorflow/core/tfrt/fallback:fallback_state",
"//tensorflow/core/tfrt/graph_executor:graph_execution_options",
"//tensorflow/core/tfrt/runtime",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@tf_runtime//:bef",
"@tf_runtime//:hostcontext",
"@xla//xla/tsl/platform:statusor",
],
)
filegroup(
name = "testdata",
srcs = glob(
["testdata/**"],
),
visibility = ["//visibility:public"],
)
tensorflow/compiler/mlir/tfrt/tests/saved_model/saved_model_test.cc
+248
View File
@@ -0,0 +1,248 @@
/* 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/saved_model/saved_model.h"
#include <iterator>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tfrt/translate/import_model.h"
#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/tfrt/fallback/fallback_state.h"
#include "tensorflow/core/tfrt/graph_executor/graph_execution_options.h"
#include "tensorflow/core/tfrt/runtime/runtime.h"
#include "tfrt/bef/bef_buffer.h" // from @tf_runtime
#include "tfrt/host_context/resource_context.h" // from @tf_runtime
namespace tensorflow {
namespace {
TEST(SavedModelTest, MapSignatures) {
std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/test.mlir");
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::MLIRContext context(registry);
auto module =
mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
ASSERT_TRUE(module);
std::vector<std::string> inputs;
std::vector<std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
in_specs;
std::vector<std::string> outputs;
std::vector<std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
out_specs;
std::vector<mlir::Operation*> bound_inputs;
TF_ASSERT_OK(MapFunctionSignaturesFromTFSavedModelMLIR(
module.get(), [&](const TFRTSavedModelSignatureInfo& sig_info) {
// Only check the signature of "serving_default".
if (sig_info.func_name != "serving_default") return;
transform(sig_info.input_names, std::back_inserter(inputs),
[](llvm::StringRef x) { return x.str(); });
in_specs.assign(sig_info.input_specs.begin(),
sig_info.input_specs.end());
transform(sig_info.output_names, std::back_inserter(outputs),
[](llvm::StringRef x) { return x.str(); });
out_specs.assign(sig_info.output_specs.begin(),
sig_info.output_specs.end());
bound_inputs.assign(sig_info.bound_inputs.begin(),
sig_info.bound_inputs.end());
}));
ASSERT_EQ(inputs.size(), 1);
EXPECT_EQ(inputs[0], "x");
ASSERT_EQ(outputs.size(), 1);
EXPECT_EQ(outputs[0], "r");
ASSERT_EQ(in_specs.size(), 1);
ASSERT_EQ(in_specs[0].first, tensorflow::DT_INT32);
ASSERT_TRUE(in_specs[0].second.IsIdenticalTo(PartialTensorShape({1, 3})));
ASSERT_EQ(out_specs.size(), 1);
ASSERT_EQ(out_specs[0].first, tensorflow::DT_INT32);
ASSERT_TRUE(out_specs[0].second.IsIdenticalTo(PartialTensorShape({1, 1})));
ASSERT_EQ(bound_inputs.size(), 2);
auto global_tensor =
llvm::cast<mlir::tf_saved_model::GlobalTensorOp>(bound_inputs[0]);
auto asset = llvm::cast<mlir::tf_saved_model::AssetOp>(bound_inputs[1]);
EXPECT_EQ(global_tensor.getSymName(), "y");
EXPECT_EQ(asset.getSymName(), "z");
}
TEST(SavedModelTest, CompileToBEF) {
std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/test.mlir");
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::MLIRContext context(registry);
auto module =
mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
ASSERT_TRUE(module);
tfrt::BefBuffer bef_buffer;
auto runtime =
tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/1);
tfrt_stub::GraphExecutionOptions options(runtime.get());
tfrt::ResourceContext resource_context;
tfrt_stub::ModelRuntimeContext model_context(
&options, options.compile_options.saved_model_dir, &resource_context);
TF_ASSERT_OK(ConvertTfMlirToBef(options.compile_options, module.get(),
&bef_buffer, model_context));
}
TEST(SavedModelTest, ConvertTfMlirToBefWithXlaFuncExport) {
std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/"
"xla_launch.mlir");
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::MLIRContext context(registry);
auto module =
mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
ASSERT_TRUE(module);
tfrt::BefBuffer bef_buffer;
auto runtime =
tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/1);
tfrt_stub::GraphExecutionOptions options(runtime.get());
options.compile_options.device_target = TfrtDeviceInfraTarget::kGpu;
TF_ASSERT_OK_AND_ASSIGN(
std::unique_ptr<tfrt_stub::FallbackState> fallback_state,
tfrt_stub::FallbackState::Create(SessionOptions(), FunctionDefLibrary()));
tfrt::ResourceContext resource_context;
tfrt_stub::ModelRuntimeContext model_context(
&options, options.compile_options.saved_model_dir, &resource_context);
TF_ASSERT_OK(ConvertTfMlirToBef(options.compile_options, module.get(),
&bef_buffer, model_context,
fallback_state.get()));
// The module contains an XLA function, as well as a while body and a while
// condition within the XLA function.
EXPECT_EQ(fallback_state->process_function_library_runtime()
.GetFunctionLibraryDefinition()
->num_functions(),
3);
}
TEST(SavedModelTest, ConvertTfMlirToBefExportingXlaReduceWindow) {
std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/"
"xla_launch_xla_reduce_window.mlir");
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::MLIRContext context(registry);
auto module =
mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
ASSERT_TRUE(module);
tfrt::BefBuffer bef_buffer;
auto runtime =
tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/1);
tfrt_stub::GraphExecutionOptions options(runtime.get());
options.compile_options.device_target = TfrtDeviceInfraTarget::kGpu;
TF_ASSERT_OK_AND_ASSIGN(
std::unique_ptr<tfrt_stub::FallbackState> fallback_state,
tfrt_stub::FallbackState::Create(SessionOptions(), FunctionDefLibrary()));
tfrt::ResourceContext resource_context;
tfrt_stub::ModelRuntimeContext model_context(
&options, options.compile_options.saved_model_dir, &resource_context);
TF_ASSERT_OK(ConvertTfMlirToBef(options.compile_options, module.get(),
&bef_buffer, model_context,
fallback_state.get()));
// The module contains an XLA function, as well as a sum_reducer function
// referenced by an XlaReduceWindow op.
EXPECT_EQ(fallback_state->process_function_library_runtime()
.GetFunctionLibraryDefinition()
->num_functions(),
2);
}
TEST(SavedModelTest, AddXlaFunctionsOutputFunctionNames) {
std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/"
"xla_launch_xla_reduce_window.mlir");
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::MLIRContext context(registry);
auto module =
mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
ASSERT_TRUE(module);
tfrt::BefBuffer bef_buffer;
auto runtime =
tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/1);
tfrt_stub::GraphExecutionOptions options(runtime.get());
options.compile_options.device_target = TfrtDeviceInfraTarget::kGpu;
TF_ASSERT_OK_AND_ASSIGN(
std::unique_ptr<tfrt_stub::FallbackState> fallback_state,
tfrt_stub::FallbackState::Create(SessionOptions(), FunctionDefLibrary()));
tfrt::ResourceContext resource_context;
tfrt_stub::ModelRuntimeContext model_context(
&options, options.compile_options.saved_model_dir, &resource_context);
std::vector<std::string> function_names;
TF_ASSERT_OK(ConvertTfMlirToBef(options.compile_options, module.get(),
&bef_buffer, model_context,
fallback_state.get(), &function_names));
EXPECT_THAT(function_names, ::testing::SizeIs(1));
}
// TODO(b/162442824): Add a SavedModel test that covers the error pass.
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/test.mlir
Vendored
+43
View File
@@ -0,0 +1,43 @@
// 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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 398 : i32}, tf_saved_model.semantics} {
"tf_saved_model.global_tensor"() {is_mutable, sym_name = "y", type = tensor<3x1xi32>, value = dense<[[1], [2], [3]]> : tensor<3x1xi32>} : () -> ()
"tf_saved_model.asset"() {sym_name = "z", filename = "file"} : () -> ()
func.func @serving_default(
%arg0: tensor<1x3xi32> {tf_saved_model.index_path = ["x"]},
%arg1: tensor<!tf_type.resource<tensor<3x1xi32>>> {tf_saved_model.bound_input = @y},
%arg2: tensor<!tf_type.string> {tf_saved_model.bound_input = @z}
) -> (tensor<1x1xi32> {tf_saved_model.index_path = ["r"]})
attributes {
tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "result:0"},
tf_saved_model.exported_names = ["serving_default"]
}
{
%0 = "tf.ReadVariableOp"(%arg1) {device = ""} : (tensor<!tf_type.resource<tensor<3x1xi32>>>) -> tensor<3x1xi32>
%1 = "tf.MatMul"(%arg0, %0) {device = "", transpose_a = false, transpose_b = false} : (tensor<1x3xi32>, tensor<3x1xi32>) -> tensor<1x1xi32>
func.return %1 : tensor<1x1xi32>
}
func.func @predict(
) -> (tensor<0x!tf_type.string> {tf_saved_model.index_path = ["r"]})
attributes {
tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "result:0"},
tf_saved_model.exported_names = ["predict"]
}
{
%0 = "tf.Const"() {dtype = !tf_type.string, value = dense<[]> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string>
func.return %0 : tensor<0x!tf_type.string>
}
}
tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/xla_launch.mlir
Vendored
+39
View File
@@ -0,0 +1,39 @@
// 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 @while_cond(%arg0: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func @while_body(%arg0: tensor<i32>) -> tensor<i32> {
%1 = "tf.AddV2"(%arg0, %arg0) {} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
func.func private @xla_func_0(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._XlaMustCompile = true, tf._noinline = true, tf._original_func_name = "should_not_be_used"} {
%1 = "tf.AddV2"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.While"(%2) { cond = @while_cond, body = @while_body, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
func.return %1 : tensor<1x3xf32>
}
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "output:0"}} {
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/xla_launch_xla_reduce_window.mlir
Vendored
+36
View File
@@ -0,0 +1,36 @@
// 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 private @sum_reducer(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
%0 = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
func.func @xla_func_0(%arg0: tensor<7xf32>, %arg1: tensor<f32>) -> tensor<10xf32> {
%cst = "tf.Const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
%cst_0 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%cst_1 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
%cst_2 = "tf.Const"() {value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%cst_3 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32>
%0 = "tf.XlaReduceWindow"(%arg0, %arg1, %cst_0, %cst_1, %cst_2, %cst_3, %cst) {computation = @sum_reducer} : (tensor<7xf32>, tensor<f32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1x2xi32>) -> tensor<10xf32>
func.return %0 : tensor<10xf32>
}
func.func @main(%arg0: tensor<7xf32>) -> tensor<10xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "output:0"}} {
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<f32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<f32>>>) -> tensor<f32>
%2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<7xf32>, tensor<f32>) -> tensor<10xf32>
func.return %2 : tensor<10xf32>
}
tensorflow/compiler/mlir/tfrt/tests/sink_in_invariant_ops.mlir
+502
View File
@@ -0,0 +1,502 @@
// 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 -split-input-file -tfrt-sink-in-invariant-ops %s | FileCheck %s --dump-input=fail --dump-input-filter=all
module attributes {tf_saved_model.semantics} {
// Test sinks in var handle op to batch function.
// CHECK-LABEL: func private @batched_function
// CHECK: arg1
func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<1x3xf32>
%1 = "tf.AddV2"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK: tf.VarHandleOp
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK: "tf.BatchFunction"(%arg0, %0)
// CHECK: operandSegmentSizes = array<i32: 1, 1>
%1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %1 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks in const op to batch function.
// CHECK-LABEL: func private @batched_function
// CHECK: arg1
func.func private @batched_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK: tf.Const
%1 = "tf.AddV2"(%arg0, %arg1) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK: [[handle:%.*]] = "tf.Const"()
%0 = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: "tf.BatchFunction"(%arg0, [[handle]])
// CHECK-SAME: operandSegmentSizes = array<i32: 1, 1>
%1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks in HashTableV2 op to batch function.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<1xi32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
%default = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
// CHECK: tf.HashTableV2
%0 = "tf.LookupTableFindV2"(%arg1, %arg0, %default) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32> {tf_saved_model.index_path = ["input"]}) -> (tensor<1xi32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
%0 = "tf.HashTableV2"() {device = "/device:CPU:0", container = "", shared_name = "variable", key_dtype = i32, value_dtype = i32} : () -> tensor<*x!tf_type.resource>
%1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1xi32>, tensor<*x!tf_type.resource>) -> tensor<1xi32>
func.return %1 : tensor<1xi32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sink in multiple invariant ops.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<!tf_type.resource<tensor<1x3xf32>>>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK-DAG: [[handle1:%.*]] = "tf.VarHandleOp"() <{{{.*}}, shared_name = "variable1"}>
// CHECK-DAG: [[handle2:%.*]] = "tf.VarHandleOp"() <{{{.*}}, shared_name = "variable2"}>
// CHECK: "tf.ReadVariableOp"([[handle1]])
// CHECK: "tf.ReadVariableOp"([[handle2]])
%0 = "tf.ReadVariableOp"(%arg0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%1 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%2 = "tf.AddV2"(%0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%3 = "tf.Identity"(%2) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK: tf.VarHandleOp
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable1"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable2"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK: "tf.BatchFunction"(%0, %1)
%2 = "tf.BatchFunction"(%0, %1) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %2 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks in var handle op that used by control flow ops.
// CHECK-LABEL: func private @some_func
func.func private @some_func(
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @some_other_func
func.func private @some_other_func(
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @sink_in_stateful_call
func.func @sink_in_stateful_call(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_stateful_call"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf.StatefulPartitionedCall"([[handle]])
%x = "tf.StatefulPartitionedCall"(%handle) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
// CHECK-LABEL: func @sink_in_if
func.func @sink_in_if(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_if"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: [[cond:%.*]] = "tf.Const"()
%cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
// CHECK: "tf.If"([[cond]], [[handle]])
%x = "tf.If"(%cond, %handle) {then_branch = @some_other_func, else_branch = @some_other_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test doesn't sink in to the callee that invoked by multiple callers.
// CHECK: func private @some_func([[arg0:.+]]: tensor<!tf_type.resource<tensor<i32>>>)
func.func private @some_func(%arg0: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK-NOT: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
%0 = "tf.ReadVariableOp"(%arg0) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @sink_in_stateful_call
func.func @sink_in_stateful_call(%arg0: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_stateful_call"]} {
// CHECK: tf.VarHandleOp
%0 = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf.StatefulPartitionedCall"(%0)
%1 = "tf.StatefulPartitionedCall"(%0) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)
%2 = "tf.AddV2"(%arg0, %1) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// Test VarHandleOp getting sinked when it is used by the called function and returned by the called function.
// CHECK: func private @func_use_and_return_varhandle([[arg0:.+]]: tensor<!tf_type.resource<tensor<i32>>>)
func.func private @func_use_and_return_varhandle(%arg0: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) {
// CHECK: tf.VarHandleOp
// CHECK-NEXT: tf.ReadVariableOp
%0 = "tf.ReadVariableOp"(%arg0) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0, %arg0 : tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>
}
// CHECK-LABEL: func @sink_in_stateful_call_varhandle_return
func.func @sink_in_stateful_call_varhandle_return(%arg0: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_stateful_call_varhandle_return"]} {
// CHECK: tf.VarHandleOp
%0 = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf.StatefulPartitionedCall"(%0)
%1:2 = "tf.StatefulPartitionedCall"(%0) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @func_use_and_return_varhandle} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
%2 = "tf.AddV2"(%arg0, %1#0) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// CHECK-LABEL: func @sink_in_if
func.func @sink_in_if(%arg0: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_if"]} {
// CHECK: tf.VarHandleOp
%0 = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
%cst = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
// CHECK: "tf.If"(%cst, %0)
%1 = "tf.If"(%cst, %0) {then_branch = @some_func, else_branch = @some_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%2 = "tf.AddV2"(%arg0, %1) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test doesn't sink in var handle op + read variable op. Consider implement when we see it from production.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK-NOT: tf.VarHandleOp
// CHECK-NOT: tf.ReadVariableOp
%1 = "tf.AddV2"(%arg0, %arg1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
// CHECK: "tf.BatchFunction"(%arg0, %1)
%2 = "tf.BatchFunction"(%arg0, %1) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<*xf32>
func.return %2 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks in var handle op if it's used by one callee, and also by read only ops in the current funciton.
// CHECK-LABEL: func private @batched_function
// CHECK: arg1
func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
%1 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%3 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK: "tf.ReadVariableOp"([[handle]])
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
// CHECK: "tf.BatchFunction"(%arg0, [[handle]])
// CHECK-SAME: operandSegmentSizes = array<i32: 1, 1>
%2 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %2 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks in var handle op crossing nested tf.BatchFunction, while the var handle op is only copied at the target.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
%1 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%3 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
func.return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func private @nested_batched_function
func.func private @nested_batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK-NEXT: tf.BatchFunction
%0 = "tf.BatchFunction"(%arg0, %arg1) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
// CHECK: "tf.ReadVariableOp"([[handle]])
%1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
// CHECK: "tf.BatchFunction"(%arg0, [[handle]])
// CHECK-SAME: operandSegmentSizes = array<i32: 1, 1>
%2 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @nested_batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
func.return %2 : tensor<*xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks crossing nested tf.If, while the sinkable ops are only copied at the target.
// CHECK-LABEL: func private @then_func
func.func private @then_func(
%cond: tensor<i1>,
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @else_func
func.func private @else_func(
%cond: tensor<i1>,
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @nested_else_func
func.func private @nested_else_func(
%cond: tensor<i1>,
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @nested_then_func(
// CHECK-SAME: [[arg0:%.*]]: tensor<i1>,
// CHECK-SAME: [[arg1:%.*]]: tensor<!tf_type.resource<tensor<i32>>>)
func.func private @nested_then_func(
%cond: tensor<i1>,
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK-NOT: tf.VarHandleOp
// CHECK: [[const:%.*]] = "tf.Const"
// CHECK: "tf.If"([[const]], [[arg0]], [[arg1]])
%0 = "tf.If"(%cond, %cond, %arg) {then_branch = @then_func, else_branch = @else_func, is_stateless = false} : (tensor<i1>, tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @nested_sink_in_if
func.func @nested_sink_in_if(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_if"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: [[cond:%.*]] = "tf.Const"()
%cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
// CHECK: "tf.If"([[cond]], [[cond]], [[handle]])
%x = "tf.If"(%cond, %cond, %handle) {then_branch = @nested_then_func, else_branch = @nested_else_func, is_stateless = false} : (tensor<i1>, tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %r : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks crossing nested tf.While and BatchFunction, while the sinkable ops are only copied at the target.
// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
// CHECK: tf.VarHandleOp
// CHECK-NEXT: tf.ReadVariableOp
%1 = "tf.ReadVariableOp"(%arg0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// CHECK-LABEL: func private @while_cond_func
func.func private @while_cond_func(
%arg0: tensor<i32>,
%arg1: tensor<i32>,
%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[handle]])
%0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @while_body_func
func.func private @while_body_func(
%arg0: tensor<i32>,
%arg1: tensor<i32>,
%arg2: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) {
// CHECK: "tf.BatchFunction"(%arg2)
%0 = "tf.BatchFunction"(%arg2) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %0, %arg0, %arg2 : tensor<i32>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>
}
// CHECK-LABEL: func @nested_sink_in_while_and_batch_functions
func.func @nested_sink_in_while_and_batch_functions(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["test_sink_in_while_and_batch_functions"]} {
// CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: [[cond:%.*]] = "tf.Const"()
%cond = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: "tf.While"([[cond]], [[cond]], [[handle]])
%x:3 = "tf.While"(%cond, %cond, %handle) {body = @while_body_func, cond = @while_cond_func, is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
func.return %x#0 : tensor<i32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// Test sinks in var handle op to WhileOp with unranked loop signature.
// This verifies that the pass does not corrupt the return type of the loop body
// when the resource is also returned (pass-through).
// CHECK-LABEL: func private @while_cond
func.func private @while_cond(%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<i1> {
%cst = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
func.return %cst : tensor<i1>
}
// CHECK-LABEL: func private @while_body
// CHECK: (%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource>) -> (tensor<i32>, tensor<*x!tf_type.resource>)
func.func private @while_body(%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource>) -> (tensor<i32>, tensor<*x!tf_type.resource>) {
// CHECK: [[sunk_handle:%.*]] = "tf.VarHandleOp"()
// CHECK: "tf.ReadVariableOp"([[sunk_handle]])
%0 = "tf.ReadVariableOp"(%arg1) {device = "cpu"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
%1 = "tf.AddV2"(%arg0, %0) {device = "cpu"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: return {{.*}}, %arg1
func.return %1, %arg1 : tensor<i32>, tensor<*x!tf_type.resource>
}
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
attributes {tf_saved_model.exported_names = ["main"]} {
%handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: "tf.While"
%x:2 = "tf.While"(%arg0, %handle) {body = @while_body, cond = @while_cond, is_stateless = false, parallel_iterations = 10 : i64} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
func.return %x#0 : tensor<i32>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/BUILD
+38
View File
@@ -0,0 +1,38 @@
load("//tensorflow:tensorflow.bzl", "if_oss")
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"],
# Custom driver is unsupported in OSS. Fails if one is provided.
# copybara:uncomment driver = "//tensorflow/compiler/mlir:run_lit.sh",
# copybara:comment_begin(JitRt/Auto fusion depreciated)
exclude = [
"auto-fusion.mlir",
"tf_to_corert_pipeline_cpurt.mlir",
"outline-cpurt-cluster.mlir",
],
# copybara:comment_end
features = if_oss(["--path=org_tensorflow/tensorflow/compiler/mlir/tfrt"]),
size_override = {
"fallback.mlir": "medium",
},
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",
],
)
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/attributes.mlir
+97
View File
@@ -0,0 +1,97 @@
// 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 -tf-to-tfrt %s | FileCheck %s --dump-input=fail
// _output_shapes and f.* attributes are removed during tf-to-tfrt lowering.
// CHECK-LABEL: func @remove_unused_attr
func.func @remove_unused_attr() {
// CHECK: %out_op_chain = tfrt_fallback_async.executeop.seq(%arg0) key(0) cost({{.*}}) device("/device:CPU:0") "tf.SomeOp2"()
"tf.SomeOp2"() {device = "/device:CPU:0", _output_shapes = ["tfshape$"], f.Tin = [f32], f._read_only_resource_inputs = []} : () -> ()
func.return
}
// CHECK-LABEL: func @basic
func.func @basic(
%arg0: tensor<3x1xf32>,
%arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> (tensor<3x3xf32>) {
%1 = "tf.ReadVariableOp"(%arg1) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
// CHECK: {{%.*}} = tfrt_fallback_async.executeop {{.*}} device("/device:CPU:0") "tf.MatMul"
// CHECK-SAME: {T = f32, transpose_a = false, transpose_b = false}
%2 = "tf.MatMul"(%arg0, %1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
func.return %2 : tensor<3x3xf32>
}
// CHECK-LABEL: func @string_type
func.func @string_type(%arg: tensor<1x2x!tf_type.string>) -> tensor<?x6x!tf_type.string> {
%multiples = "tf.Const"() { device = "/device:CPU:0", value = dense<[7,3]> : tensor<2xi32> } : () -> tensor<2xi32>
// CHECK: T = !corert.string
%output = "tf.Tile"(%arg, %multiples) { device = "/device:CPU:0" } : (tensor<1x2x!tf_type.string>, tensor<2xi32>) -> tensor<?x6x!tf_type.string>
func.return %output : tensor<?x6x!tf_type.string>
}
// CHECK-LABEL: func @shape
func.func @shape() {
%size = "tf.Const"() {value = dense<5> : tensor<i32>} : () -> tensor<i32>
// CHECK: tf.TensorArrayV3
// CHECK-SAME: element_shape = #corert.shape<*>
%ta:2 = "tf.TensorArrayV3"(%size) {device = "/device:CPU:0", dtype = f32, element_shape = #tf_type.shape<*>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor<i32>) -> (tensor<!tf_type.resource>, tensor<f32>)
func.return
}
// CHECK-LABEL: func @resource
func.func @resource() {
// CHECK: tf.SomeOp
// CHECK-SAME: dtype = !corert.resource
"tf.SomeOp"() {device = "/device:CPU:0", dtype = !tf_type.resource} : () -> ()
func.return
}
// CHECK-LABEL: func @variant
func.func @variant(%arg: tensor<!tf_type.variant>) {
// CHECK: tf.ZerosLike
// CHECK-SAME: T = !corert.variant
%0 = "tf.ZerosLike"(%arg) {device = "/device:CPU:0", T = !tf_type.variant} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return
}
// Checks that TF quantized attrs are lowered to the corert types
// CHECK-LABEL: func @quantized_types
func.func @quantized_types(%arg0: tensor<!tf_type.resource<tensor<1x3x!tf_type.quint8>>>,
%arg1: tensor<!tf_type.resource<tensor<1x3x!tf_type.quint16>>>,
%arg2: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint8>>>,
%arg3: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint16>>>,
%arg4: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint32>>>) {
// CHECK: tf.ReadVariableOp
// CHECK-SAME: dtype = !corert.quint8
%0 = "tf.ReadVariableOp"(%arg0) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = !tf_type.quint8} : (tensor<!tf_type.resource<tensor<1x3x!tf_type.quint8>>>) -> tensor<1x3x!tf_type.quint8>
// CHECK: tf.ReadVariableOp
// CHECK-SAME: dtype = !corert.quint16
%1 = "tf.ReadVariableOp"(%arg1) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = !tf_type.quint16} : (tensor<!tf_type.resource<tensor<1x3x!tf_type.quint16>>>) -> tensor<1x3x!tf_type.quint16>
// CHECK: tf.ReadVariableOp
// CHECK-SAME: dtype = !corert.qint8
%2 = "tf.ReadVariableOp"(%arg2) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = !tf_type.qint8} : (tensor<!tf_type.resource<tensor<1x3x!tf_type.qint8>>>) -> tensor<1x3x!tf_type.qint8>
// CHECK: tf.ReadVariableOp
// CHECK-SAME: dtype = !corert.qint16
%3 = "tf.ReadVariableOp"(%arg3) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = !tf_type.qint16} : (tensor<!tf_type.resource<tensor<1x3x!tf_type.qint16>>>) -> tensor<1x3x!tf_type.qint16>
// CHECK: tf.ReadVariableOp
// CHECK-SAME: dtype = !corert.qint32
%4 = "tf.ReadVariableOp"(%arg4) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = !tf_type.qint32} : (tensor<!tf_type.resource<tensor<1x3x!tf_type.qint32>>>) -> tensor<1x3x!tf_type.qint32>
func.return
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/basic.mlir
+70
View File
@@ -0,0 +1,70 @@
// 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 -pass-pipeline='builtin.module(func.func(tf-tensor-device-copy),tfrt-lower-tf-savedmodel{hoist-invariant-ops=true},tf-to-tfrt{tfrt-cost-threshold=1024 tfrt-merge-inter-dependent-streams=true})' %s | FileCheck %s --dump-input-filter=all
// CHECK-NOT: tf_saved_model.semantics
// CHECK: tfrt.cost_threshold = 1024
// CHECK-SAME: tfrt.merge_inter_dependent_streams = true
module attributes {tf_saved_model.semantics} {
// CHECK-NOT: "tf_saved_model.session_initializer"
"tf_saved_model.session_initializer"() { initializers = [@func_init] } : () -> ()
// CHECK-LABEL: _tfrt_resource_init
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
// CHECK: tfrt_fallback_async.set_resource
// CHECK-SAME: {device = "/device:CPU:0", index = 0 : i64}
// CHECK-LABEL: func @init
// CHECK-SAME: {tfrt.cost_threshold = 1 : i64}
func.func @func_init() attributes {tf_saved_model.exported_names = ["init"]} {
func.return
}
// CHECK-LABEL: func @basic
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
// CHECK-SAME: [[arg0:%.*]]: !tfrt_fallback.tf_tensor,
// CHECK-SAME: [[arg1:%.*]]: !tfrt_fallback.tf_tensor)
// CHECK-SAME: -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @func_basic(
%arg0: tensor<3x1xf32> {tf_saved_model.index_path = [0]},
%arg1: tensor<!tf_type.resource<tensor<1x3xf32>>> {tf_saved_model.index_path = [1]})
-> (tensor<3x3xf32> {tf_saved_model.index_path = []})
attributes {tf_saved_model.exported_names = ["basic"]} {
%handle = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<3xf32>>>
%0 = "tf.ReadVariableOp"(%handle) {_output_shapes = ["tfshape$dim { size: 3 }"], device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<3xf32>>>) -> tensor<3xf32>
%1 = "tf.ReadVariableOp"(%arg1) {_output_shapes = ["tfshape$dim { size: 1 } dim { size: 3 }"], device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
// CHECK-NEXT: [[ready_ch:%.*]] = tfrt.new.chain
// CHECK-NEXT: [[ch:%.*]], [[result:%.*]] = tfrt_fallback_async.get_resource [[ready_ch]] {device = "/device:CPU:0", indices = [0]} : (!tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
// CHECK-NEXT: [[ch1:%.*]], [[var:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) {{.*}} "tf.ReadVariableOp"([[arg1]])
// CHECK-NEXT: [[r0:%.*]] = tfrt_fallback_async.executeop {{.*}} "tf.MatMul"([[arg0]], [[var]])
%2 = "tf.MatMul"(%arg0, %1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: [[r1:%.*]] = tfrt_fallback_async.executeop {{.*}} "tf.BiasAdd"([[r0]], [[result]])
%3 = "tf.BiasAdd"(%2, %0) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], data_format = "NHWC", device = "/device:CPU:0"} : (tensor<3x3xf32>, tensor<3xf32>) -> tensor<3x3xf32>
// CHECK-NEXT: [[r2:%.*]] = tfrt_fallback_async.executeop {{.*}} "tf.Tanh"([[r1]]) {T = f32}
%4 = "tf.Tanh"(%3) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NOT: tf.Identity
%5 = "tf.Identity"(%4) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<3x3xf32>
// CHECK-NOT: tf.IdentityN
%6:2 = "tf.IdentityN"(%5, %4) {T = [f32, f32], _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }", "tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0"} : (tensor<3x3xf32>, tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<3x3xf32>)
// CHECK-NEXT: [[out_ch:%.*]] = tfrt.merge.chains [[ch]], [[ch1]] : !tfrt.chain, !tfrt.chain
// CHECK-NEXT: tfrt.return [[out_ch]], [[r2]] : !tfrt.chain, !tfrt_fallback.tf_tensor
func.return %6#0 : tensor<3x3xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/batch_function_deduplicate.mlir
+60
View File
@@ -0,0 +1,60 @@
// 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 -tfrt-deduplicate-functions-invoked-by-batch-function %s | FileCheck %s
// This test verifies the function `compute_1` will be removed to deduplicate
// the functions invoked by BatchFunction with the same shared_name and the
// function `compute_2` will not be removed as the shared_name is different.
// CHECK-LABEL: func private @batch_0
// CHECK: f = @compute_0
func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {
%0:2= "tf.BatchFunction"(%arg0, %arg0) {_xla_inferred_shapes = [#tf_type.shape<*>, #tf_type.shape<*>], allowed_batch_sizes = [64, 128, 256], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", enable_large_batch_splitting = true, f = @compute_0, max_batch_size = 256 : i64, max_enqueued_batches = 10000 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "computation"} : (tensor<?x?xi32>, tensor<?x?xi32>) -> (tensor<*xi32>, tensor<*xi32>)
func.return %0#0 : tensor<*xi32>
}
// CHECK: func private @compute_0
func.func private @compute_0(%arg0: tensor<?x?xi32> {tf._user_specified_name = "0"}, %arg1: tensor<?x?xi32>) -> (tensor<?x?xi32>, tensor<?x?xi32>) {
func.return %arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>
}
// Batch function in batch_1 uses the same shared_name as the one in batch_0,
// so compute_1 is deduped, and compute_0 will be used here.
// CHECK-LABEL: func private @batch_1
// CHECK: f = @compute_0
// CHECK-NOT: f = @compute_1
func.func private @batch_1(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {
%0:2 = "tf.BatchFunction"(%arg0, %arg0) {_xla_inferred_shapes = [#tf_type.shape<*>, #tf_type.shape<*>], allowed_batch_sizes = [64, 128, 256], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", enable_large_batch_splitting = true, f = @compute_1, max_batch_size = 256 : i64, max_enqueued_batches = 10000 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "computation"} : (tensor<?x?xi32>, tensor<?x?xi32>) -> (tensor<*xi32>, tensor<*xi32>)
func.return %0#0 : tensor<*xi32>
}
// CHECK-NOT: func private @compute_1
func.func private @compute_1(%arg0: tensor<?x?xi32> {tf._user_specified_name = "0"}, %arg1: tensor<?x?xi32>) -> (tensor<?x?xi32>, tensor<?x?xi32>) {
func.return %arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>
}
// Batch function in batch_2 uses a different shared_name from the one in
// batch_0, so it should be kept.
// CHECK-LABEL: func private @batch_2
// CHECK: f = @compute_2
func.func private @batch_2(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {
%0:2 = "tf.BatchFunction"(%arg0, %arg0) {_xla_inferred_shapes = [#tf_type.shape<*>, #tf_type.shape<*>], allowed_batch_sizes = [64, 128, 256], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", enable_large_batch_splitting = true, f = @compute_2, max_batch_size = 256 : i64, max_enqueued_batches = 10000 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "computation_unique_name"} : (tensor<?x?xi32>, tensor<?x?xi32>) -> (tensor<*xi32>, tensor<*xi32>)
func.return %0#0 : tensor<*xi32>
}
// CHECK: func private @compute_2
func.func private @compute_2(%arg0: tensor<?x?xi32> {tf._user_specified_name = "0"}, %arg1: tensor<?x?xi32>) -> (tensor<?x?xi32>, tensor<?x?xi32>) {
func.return %arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/batch_function_deduplicate_failed.mlir
+38
View File
@@ -0,0 +1,38 @@
// 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: not tf-tfrt-opt -tfrt-deduplicate-functions-invoked-by-batch-function %s 2>&1 | FileCheck %s
// This test verifies the error when two functions are different but invoked by
// the batch functions with same shared_name.
func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {
%0:2 = "tf.BatchFunction"(%arg0, %arg0) {_xla_inferred_shapes = [#tf_type.shape<*>, #tf_type.shape<*>], allowed_batch_sizes = [64, 128, 256], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", enable_large_batch_splitting = true, f = @compute_0, max_batch_size = 256 : i64, max_enqueued_batches = 10000 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "computation"} : (tensor<?x?xi32>, tensor<?x?xi32>) -> (tensor<*xi32>, tensor<*xi32>)
func.return %0#0 : tensor<*xi32>
}
func.func private @compute_0(%arg0: tensor<?x?xi32> {tf._user_specified_name = "0"}, %arg1: tensor<?x?xi32>) -> (tensor<?x?xi32>, tensor<?x?xi32>) {
func.return %arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>
}
func.func private @batch_3(%arg0: tensor<?x1xi32>) -> tensor<*xi32> {
%0:2 = "tf.BatchFunction"(%arg0, %arg0) {_xla_inferred_shapes = [#tf_type.shape<*>, #tf_type.shape<*>], allowed_batch_sizes = [64, 128, 256], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", enable_large_batch_splitting = true, f = @compute_3, max_batch_size = 256 : i64, max_enqueued_batches = 10000 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "computation"} : (tensor<?x1xi32>, tensor<?x1xi32>) -> (tensor<*xi32>, tensor<*xi32>)
func.return %0#0 : tensor<*xi32>
}
// compute_3 has different argument types from compute_1, thus error is reported.
// CHECK: error: func_ops for BatchFunctionOp with the same shared name are different
func.func private @compute_3(%arg0: tensor<?x1xi32> {tf._user_specified_name = "0"}, %arg1: tensor<?x1xi32>) -> (tensor<?x1xi32>, tensor<?x1xi32>) {
func.return %arg0, %arg1 : tensor<?x1xi32>, tensor<?x1xi32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/const_tensor.mlir
+44
View File
@@ -0,0 +1,44 @@
// 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 -tf-to-tfrt %s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @string_tensor
func.func @string_tensor() -> (tensor<0x!tf_type.string>, tensor<7x!tf_type.string>) {
// CHECK: {shape = [0], value = []}
%0 = "tf.Const"() {value = dense<[]> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string>
// CHECK: {shape = [7], value = ["has_login_page_feature", "num_terms_inside_postform", "num_terms_outside_postform", "num_terms_outside_postform_without_bp", "query_params_contains_url", "title_with_login_phase", "url_contains_login_terms"]}
%1 = "tf.Const"() {value = dense<["has_login_page_feature", "num_terms_inside_postform", "num_terms_outside_postform", "num_terms_outside_postform_without_bp", "query_params_contains_url", "title_with_login_phase", "url_contains_login_terms"]> : tensor<7x!tf_type.string>} : () -> tensor<7x!tf_type.string>
func.return %0, %1 : tensor<0x!tf_type.string>, tensor<7x!tf_type.string>
}
// Convert tf.Const to tfrt_fallback_async.const_dense_tensor only on cpu device
// CHECK-LABEL: func @dense_tensor
func.func @dense_tensor() -> tensor<4xui64> {
// CHECK: tfrt_fallback_async.const_dense_tensor dense<[1, 2, 3, 4]> : tensor<4xui64>
%0 = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : () -> tensor<4xui64>
// CHECK: tfrt_fallback_async.const_dense_tensor dense<1.000000e+00> : tensor<1xbf16>
%1 = "tf.Const"() {device = "/device:CPU:0", value = dense<[1.0]> : tensor<1xbf16>} : () -> tensor<4xbf16>
// CHECK: corert.executeop({{.*}}) "tf.Const"() {dtype = ui64, value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : 1
%2 = "tf.Const"() {device = "/device:GPU:0", value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : () -> tensor<4xui64>
func.return %0 : tensor<4xui64>
}
// CHECK-LABEL: func @tensor_proto
func.func @tensor_proto() -> tensor<!tf_type.quint8> {
// tfrt_fallback_async.const_tensor_proto accepts a serialized tensor proto.
// CHECK: tfrt_fallback_async.const_tensor_proto "\08\0C\12\00\22\01@"
%0 = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F5155494E54382074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20224022"> : tensor<!tf_type.quint8>} : () -> tensor<!tf_type.quint8>
func.return %0 : tensor<!tf_type.quint8>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/control_flow.mlir
+215
View File
@@ -0,0 +1,215 @@
// 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 -tf-to-tfrt="enable-while-parallel-iterations=true" %s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @cond_false(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @cond_false(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @cond_true(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @cond_true(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @cond(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor, %arg2: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @cond(%arg0: tensor<i1>, %arg1: tensor<i32>) -> tensor<i32> {
// CHECK: [[cond:%.*]] = tfrt_fallback_async.predicate
// CHECK: [[cond_res:%.*]]:2 = tfrt.cond [[cond]]
// CHECK-SAME: @cond_true @cond_false(%arg0, %arg2) : (!tfrt.chain, !tfrt_fallback.tf_tensor)
%2 = "tf.If"(%arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = true} : (tensor<i1>, tensor<i32>) -> tensor<i32>
// CHECK: [[out_ch:%.*]] = tfrt.merge.chains [[cond_res]]#0, %arg0 : !tfrt.chain, !tfrt.chain
// CHECK: tfrt.return [[out_ch]], [[cond_res]]#1 : !tfrt.chain, !tfrt_fallback.tf_tensor
func.return %2 : tensor<i32>
}
// CHECK-LABEL: func @cond_stateful(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @cond_stateful(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK: [[cond_res:%.*]]:2 = tfrt.cond
// CHECK-SAME: @cond_true @cond_false(%arg0, %arg1) : (!tfrt.chain, !tfrt_fallback.tf_tensor)
%2 = "tf.If"(%1, %arg0) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = false} : (tensor<i1>, tensor<i32>) -> tensor<i32>
// Note: returns %out_op_chain.
// CHECK: tfrt.return [[cond_res]]#0, [[cond_res]]#1 : !tfrt.chain, !tfrt_fallback.tf_tensor
func.return %2 : tensor<i32>
}
// CHECK-LABEL: func @while_cond_lt9
// CHECK-SAME: ({{%.+}}: !tfrt.chain, {{%.+}}: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
// CHECK-LABEL: func @while_body_add2
// CHECK-SAME: ({{%.+}}: !tfrt.chain, {{%.+}}: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @while_test
// CHECK-SAME: ([[ARG0:%.+]]: !tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @while_test() -> (tensor<i32>) {
// CHECK: [[CONST:%.*]] = tfrt_fallback_async.const_dense_tensor dense<0> : tensor<i32>
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[pred_res:%.*]]:2 = tfrt.call @"while_cond_lt9/tfrt_predicate"([[ARG0]], [[CONST]]) : (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, i1)
// CHECK: [[while_res:%.]]:2 = tfrt.while [[pred_res]]#1 @"while_body_add2/tfrt_body_1"([[pred_res]]#0, [[CONST]])
// CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
%1 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
// CHECK: [[out_chain:%.*]] = tfrt.merge.chains [[while_res]]#0, [[ARG0]]
// CHECK: tfrt.return [[out_chain]], [[while_res]]#1 : !tfrt.chain, !tfrt_fallback.tf_tensor
func.return %1 : tensor<i32>
}
// CHECK: func @"while_body_add2/tfrt_body_1"([[ch:%.*]]: !tfrt.chain, [[arg:%.*]]: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor, i1)
// CHECK: [[body_res:%.*]]:2 = tfrt.call @while_body_add2([[ch]], [[arg]]) : (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
// CHECK: [[pred_res:%.*]]:2 = tfrt.call @"while_cond_lt9/tfrt_predicate"([[body_res]]#0, [[body_res]]#1) : (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, i1)
// CHECK: tfrt.return [[pred_res]]#0, [[body_res]]#1, [[pred_res]]#1 : !tfrt.chain, !tfrt_fallback.tf_tensor, i1
// CHECK: func @"while_cond_lt9/tfrt_predicate"([[ch:%.*]]: !tfrt.chain, [[arg:%.*]]: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, i1)
// CHECK: [[cond_res:%.*]]:2 = tfrt.call @while_cond_lt9([[ch]], [[arg]]) : (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
// CHECK: [[bool_cond:%.*]] = tfrt_fallback_async.predicate [[cond_res]]#1
// CHECK: tfrt.return [[cond_res]]#0, [[bool_cond]] : !tfrt.chain, i1
// CHECK-LABEL: func @multi_while_test
func.func @multi_while_test() -> (tensor<i32>, tensor<i32>) {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
// CHECK: [[pred_0:%.*]]:2 = tfrt.call @"while_cond_lt9/tfrt_predicate"
// CHECK: tfrt.while [[pred_0]]#1 @"while_body_add2/tfrt_body_10"
// CHECK-SAME: parallel_iterations(10)
// CHECK: [[pred_1:%.*]]:2 = tfrt.call @"while_cond_lt9/tfrt_predicate"
// CHECK: tfrt.while [[pred_1]]#1 @"while_body_add2/tfrt_body_1"
// CHECK-SAME: parallel_iterations(1)
%2 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 10} : (tensor<i32>) -> (tensor<i32>)
%3 = "tf.While"(%1) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
func.return %2, %3 : tensor<i32>, tensor<i32>
}
func.func @side_effect_while_cond_lt9(%arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i1> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.ReadVariableOp"(%arg) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%2 = "tf.Less"(%1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %2 : tensor<i1>
}
func.func @side_effect_while_body_add2(%arg: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<!tf_type.resource<tensor<i32>>>) {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.ReadVariableOp"(%arg) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%2 = "tf.Add"(%1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
"tf.AssignVariableOp"(%arg, %2) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
func.return %arg : tensor<!tf_type.resource<tensor<i32>>>
}
// CHECK-LABEL: func @side_effect_while_test
func.func @side_effect_while_test() -> (tensor<i32>) {
%0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "c", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<i32>>>
// CHECK: [[while_res:%.]]:2 = tfrt.while {{%.*}} @"side_effect_while_body_add2/tfrt_body_1"
// CHECK: [[out_ch:%.*]], [[res:%.*]] = tfrt_fallback_async.executeop.seq([[while_res]]#0) {{.*}} "tf.ReadVariableOp"
%1 = "tf.While"(%0) { cond = @side_effect_while_cond_lt9, body = @side_effect_while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<!tf_type.resource<tensor<i32>>>)
%2 = "tf.ReadVariableOp"(%1) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
func.return %2 : tensor<i32>
}
func.func @tensor_array_while_cond(%index: tensor<i32>, %size: tensor<i32>, %flow_0: tensor<f32>, %flow_1: tensor<f32>, %handle_0: tensor<2x!tf_type.resource<tensor<?x100xf32>>>, %handle_1: tensor<2x!tf_type.resource<tensor<?x512xf32>>>) -> (tensor<i1>) {
%0 = "tf.Less"(%index, %size) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %0 : tensor<i1>
}
func.func @tensor_array_while_body(%index: tensor<i32>, %size: tensor<i32>, %flow_0: tensor<f32>, %flow_1: tensor<f32>, %handle_0: tensor<2x!tf_type.resource<tensor<?x100xf32>>>, %handle_1: tensor<2x!tf_type.resource<tensor<?x512xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>) {
%cst = "tf.Const"() {value = dense<1.1> : tensor<100x512xf32>} : () -> tensor<100x512xf32>
%one = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%x = "tf.TensorArrayReadV3"(%handle_0, %index, %flow_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<i32>, tensor<f32>) -> tensor<?x100xf32>
%y = "tf.MatMul"(%x, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x100xf32>, tensor<100x512xf32>) -> (tensor<?x512xf32>)
%flow_1_out = "tf.TensorArrayWriteV3"(%handle_1, %index, %y, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<i32>, tensor<?x512xf32>, tensor<f32>) -> tensor<f32>
%next_index = "tf.AddV2"(%index, %one) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %next_index, %size, %flow_0, %flow_1_out, %handle_0, %handle_1 : tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>
}
// CHECK-LABEL: func @tensor_array_while_test
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
func.func @tensor_array_while_test(%indices: tensor<?xi32>, %input_0: tensor<?x?x?xf32>, %input_1: tensor<?x?x?xf32>) -> (tensor<?x?x512xf32>, tensor<?x?x512xf32>) {
%index = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> (tensor<i32>)
%size = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> (tensor<i32>)
%handle_0, %flow_0 = "tf.TensorArrayV3"(%size) {clear_after_read = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<?x100>, identical_element_shapes = true, tensor_array_name = "processed_embeddings/bidirectional_rnn/bw/bw/dynamic_rnn/input_0"} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<f32>)
%handle_1, %flow_1 = "tf.TensorArrayV3"(%size) {clear_after_read = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<?x512>, identical_element_shapes = true, tensor_array_name = "processed_embeddings/bidirectional_rnn/bw/bw/dynamic_rnn/output_0"} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<f32>)
%flow_01 = "tf.TensorArrayScatterV3"(%handle_0, %indices, %input_0, %flow_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<?xi32>, tensor<?x?x?xf32>, tensor<f32>) -> tensor<f32>
// CHECK: [[pred_0:%.*]]:2 = tfrt.call @"tensor_array_while_cond/tfrt_predicate"([[in_chain]]
// CHECK: [[while_res_0:%.*]]:7 = tfrt.while {{%.*}} @"tensor_array_while_body/tfrt_body_10"([[pred_0]]#0
// CHECK-SAME: parallel_iterations(10)
%res_0:6 = "tf.While"(%index, %size, %flow_01, %flow_1, %handle_0, %handle_1) {body = @tensor_array_while_body, cond = @tensor_array_while_cond, device = "", is_stateless = false, parallel_iterations = 10 : i64} : (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>)
%output_0 = "tf.TensorArrayGatherV3"(%handle_1, %indices, %res_0#3) {device = "/job:localhost/replica:0/task:0/device:CPU:0", element_shape = #tf_type.shape<?x512>} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<?xi32>, tensor<f32>) -> tensor<?x?x512xf32>
%handle_2, %flow_2 = "tf.TensorArrayV3"(%size) {clear_after_read = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<?x100>, identical_element_shapes = true, tensor_array_name = "processed_embeddings/bidirectional_rnn/bw/bw/dynamic_rnn/input_0"} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<f32>)
%handle_3, %flow_3 = "tf.TensorArrayV3"(%size) {clear_after_read = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<?x512>, identical_element_shapes = true, tensor_array_name = "processed_embeddings/bidirectional_rnn/bw/bw/dynamic_rnn/output_0"} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<f32>)
%flow_21 = "tf.TensorArrayScatterV3"(%handle_2, %indices, %input_1, %flow_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<?xi32>, tensor<?x?x?xf32>, tensor<f32>) -> tensor<f32>
// CHECK: [[pred_1:%.*]]:2 = tfrt.call @"tensor_array_while_cond/tfrt_predicate"([[in_chain]]
// CHECK: [[while_res_1:%.*]]:7 = tfrt.while {{%.*}} @"tensor_array_while_body/tfrt_body_10"([[pred_1]]#0
// CHECK-SAME: parallel_iterations(10)
%res_1:6 = "tf.While"(%index, %size, %flow_21, %flow_3, %handle_2, %handle_3) {body = @tensor_array_while_body, cond = @tensor_array_while_cond, device = "", is_stateless = false, parallel_iterations = 10 : i64} : (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>)
%output_1 = "tf.TensorArrayGatherV3"(%handle_3, %indices, %res_1#3) {device = "/job:localhost/replica:0/task:0/device:CPU:0", element_shape = #tf_type.shape<?x512>} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<?xi32>, tensor<f32>) -> tensor<?x?x512xf32>
func.return %output_0, %output_1 : tensor<?x?x512xf32>, tensor<?x?x512xf32>
}
// CHECK: func @"tensor_array_while_body/tfrt_body_10"
func.func @callee(%arg0: tensor<i32>) -> (tensor<i32>) {
func.return %arg0: tensor<i32>
}
// CHECK-LABEL: func @call_test
// CHECK-SAME: ([[chain:%.*]]: !tfrt.chain,
func.func @call_test(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
%0 = "tf.Add"(%arg0, %arg0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: [[results_0:%.*]]:2 = tfrt.call @callee([[chain]]
// CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
%1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
// CHECK-NEXT: [[results_1:%.*]]:2 = tfrt.call @callee([[chain]]
// CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
%2 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
// CHECK-NEXT: [[results_2:%.*]]:2 = tfrt.call @callee([[chain]]
// CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
%3 = "tf.LegacyCall"(%0) {f = @callee} : (tensor<i32>) -> (tensor<i32>)
// CHECK: [[results_0]]#1, [[results_1]]#1, [[results_2]]#1
func.return %1, %2, %3 : tensor<i32>, tensor<i32>, tensor<i32>
}
func.func @branch0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
%0 = "tf.Add" (%arg0, %arg1) {device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
func.func @branch1(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
%0 = "tf.Add" (%arg0, %arg1) {device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
%1 = "tf.Add" (%arg0, %0) {device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %1 : tensor<f32>
}
// CHECK-LABEL: func @case_test
// CHECK-SAME: ([[chain:%.*]]: !tfrt.chain, [[tf_idx:%.*]]: !tfrt_fallback.tf_tensor, [[branch_arg0:%.*]]: !tfrt_fallback.tf_tensor, [[branch_arg1:%.*]]: !tfrt_fallback.tf_tensor)
func.func @case_test(%arg0: tensor<i32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// CHECK: [[th_idx:%.*]] = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle [[tf_idx]]
// CHECK-NEXT: [[idx:%.*]] = corert.tensorhandle_to_int32 [[th_idx]]
// CHECK-NEXT: [[out:%.*]] = tfrt.case [[idx]] [@branch0, @branch1]([[chain]], [[branch_arg0]], [[branch_arg1]])
%0 = "tf.Case"(%arg0, %arg1, %arg2) {_lower_using_switch_merge = true, branches = [@branch0, @branch1], is_stateless = true} : (tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/decompose_resource_op.mlir
+32
View File
@@ -0,0 +1,32 @@
// 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 -tf-executor-to-tfrt-pipeline=decompose-resource-ops=true %s | FileCheck %s --dump-input=fail
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 293 : i32}} {
// CHECK-LABEL: func @gather
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
// CHECK-SAME: [[arg0:%.*]]: !tfrt_fallback.tf_tensor, [[arg1:%.*]]: !tfrt_fallback.tf_tensor)
// CHECK: [[const:%.*]] = tfrt_fallback_async.const_dense_tensor
// CHECK-NEXT: [[out_chain:%.*]], [[value:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"({{.*}})
// CHECK-NEXT: [[res:%.*]] = tfrt_fallback_async.executeop key(1) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.GatherV2"([[value]], {{.*}}, [[const]])
// CHECK-NEXT: tfrt.return [[out_chain]], [[res]] : !tfrt.chain, !tfrt_fallback.tf_tensor
func.func @gather(%indices: tensor<?xi32>,
%resource: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
%0 = "tf.ResourceGather"(%resource, %indices) {batch_dims = 0 : i64, device = "/device:CPU:0", validate_indices = true}: (tensor<*x!tf_type.resource>, tensor<?xi32>) -> (tensor<*xi32>)
func.return %0 : tensor<*xi32>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/derived_attrs.mlir
+35
View File
@@ -0,0 +1,35 @@
// 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 -tf-to-tfrt %s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @derived_attrs
func.func @derived_attrs(
%serialized: tensor<?x!tf_type.string>,
%names: tensor<0x!tf_type.string>,
%sparse_keys: tensor<0x!tf_type.string>,
%dense_keys: tensor<1x!tf_type.string>,
%ragged_keys: tensor<0x!tf_type.string>,
%dense_default: tensor<0xi64>) -> tensor<?xi64> {
%dense_value =
"tf.ParseExampleV2"(%serialized, %names, %sparse_keys, %dense_keys, %ragged_keys, %dense_default)
// CHECK: Tdense = [i64]
// CHECK-SAME: dense_shapes = [#corert.shape<>]
{ device = "/device:CPU:0", num_sparse = 0 : i64, dense_shapes = [#tf_type.shape<>], resultSegmentSizes = array<i32: 0, 0, 0, 1, 0, 0>}
: (tensor<?x!tf_type.string>, tensor<0x!tf_type.string>, tensor<0x!tf_type.string>, tensor<1x!tf_type.string>, tensor<0x!tf_type.string>, tensor<0xi64>)
-> tensor<?xi64>
func.return %dense_value : tensor<?xi64>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/device_conversion.mlir
+25
View File
@@ -0,0 +1,25 @@
// 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 -tf-to-tfrt %s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @device_test
func.func @device_test(
%arg0: tensor<3x1xf32> {tf_saved_model.index_path = [0]},
%arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
-> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
// CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"
%2 = "tf.MatMul"(%arg0, %arg1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
func.return %2 : tensor<3x3xf32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/errors.mlir
+21
View File
@@ -0,0 +1,21 @@
// 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 -tf-to-tfrt %s -split-input-file -verify-diagnostics
// expected-error @+1 {{failed to legalize operation 'func.func' that was explicitly marked illegal}}
func.func @test_identity_wrong_type(%arg0: tensor<4x2x!tf_type.string>) -> tensor<4x2x!tf_type.stringref> {
%0 = "tf.SomeOp"(%arg0) : (tensor<4x2x!tf_type.string>) -> tensor<4x2x!tf_type.stringref>
func.return %0 : tensor<4x2x!tf_type.stringref>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/fallback.mlir
+129
View File
@@ -0,0 +1,129 @@
// 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 -tf-to-tfrt %s | FileCheck %s --dump-input=fail --dump-input-filter=all
// RUN: tf-tfrt-opt -pass-pipeline='builtin.module(tf-to-tfrt{target-tpurt=true tpu-use-core-selector=false})' %s | FileCheck %s --dump-input=fail --dump-input-filter=all
// CHECK-LABEL: func @_tfrt_fallback_init
// CHECK-SAME: {{.*}} !tfrt.chain
// CHECK: tfrt_fallback_async.createop(%arg0) key(0) device("/device:CPU:0") "tf.ParseExampleV2"()
// CHECK-SAME: Tdense = [f32, f32], dense_shapes = [#corert.shape<>, #corert.shape<>]
// CHECK-SAME: num_sparse = 2 : i64, ragged_split_types = [], ragged_value_types = []
// CHECK-SAME: sparse_types = [!corert.string, i64]}
// CHECK-SAME: num_args(7)
// CHECK: tfrt_fallback_async.createop(%0) key(1) device("/device:CPU:0") "tf.ReadVariableOp"() {dtype = f32} num_args(1)
// CHECK: tfrt_fallback_async.createop(%1) key(2) device("/device:CPU:0") "tf.MatMul"() {T = f32, transpose_a = false, transpose_b = false} num_args(2)
// CHECK-LABEL: func @main
// CHECK-SAME: {{.*}} !tfrt.chain
// CHECK-SAME: [[serialized:%.*]]: !tfrt_fallback.tf_tensor
func.func @main(%serialized: tensor<32x!tf_type.string>) -> (tensor<?x2xi64>) attributes {tf.entry_function = {inputs = "input0", outputs = "ParseExample/ParseExampleV2"}} {
%dense_default_0 = "tf.Const"() {device = "/device:CPU:0", dtype = f32, value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32>
%dense_default_1 = "tf.Const"() {device = "/device:CPU:0", dtype = f32, value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32>
%dense_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%names = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string>
%ragged_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string>
%sparse_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
// CHECK: [[outputs:%.*]]:8 = tfrt_fallback_async.executeop key(0) cost({{.*}}) device("/device:CPU:0") "tf.ParseExampleV2"
// CHECK-SAME: ([[serialized]]
// CHECK-NOT: device
// CHECK-SAME: Tdense = [f32, f32]
// CHECK-SAME: dense_shapes = [#corert.shape<>, #corert.shape<>]
// CHECK-SAME: num_sparse = 2 : i64
// CHECK-SAME: ragged_split_types = []
// CHECK-SAME: ragged_value_types = []
// CHECK-SAME: sparse_types = [!corert.string, i64]
%outputs:8 = "tf.ParseExampleV2"(%serialized, %names, %sparse_keys, %dense_keys, %ragged_keys, %dense_default_0, %dense_default_1)
{
Tdense = [f32, f32], dense_shapes = [#tf_type.shape<>, #tf_type.shape<>],
device = "/device:CPU:0", num_sparse = 2 : i64, ragged_split_types = [], ragged_value_types = [],
resultSegmentSizes = array<i32: 2, 2, 2, 2, 0, 0>,
sparse_types = [!tf_type.string, !tf_type.string]
} : (tensor<32x!tf_type.string>, tensor<0x!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>, tensor<0x!tf_type.string>, tensor<0xf32>, tensor<0xf32>)
-> (tensor<?x2xi64>, tensor<?x2xi64>, tensor<?x!tf_type.string>, tensor<?xi64>, tensor<2xi64>, tensor<2xi64>, tensor<32xf32>, tensor<32xf32>)
// CHECK: tfrt.return {{.*}}, [[outputs]]#0
func.return %outputs#0 : tensor<?x2xi64>
}
// CHECK-LABEL: func @no_native
func.func @no_native(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> {
// CHECK-NOT: corert.executeop
// CHECK: tfrt_fallback_async.executeop.seq({{.*}}) key(1) cost({{.*}}) device("/device:CPU:0") "tf.ReadVariableOp"
// CHECK: tfrt_fallback_async.executeop key(2) cost({{.*}}) device("/device:CPU:0") "tf.MatMul"
%0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
func.return %1 : tensor<3x3xf32>
}
// CHECK-LABEL: func @gpu_device
func.func @gpu_device(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> {
// CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"
// CHECK: {{.*}} = corert.executeop.seq({{.*}}) "tf.ReadVariableOp"({{.*}}) {dtype = f32} : 1
// CHECK: {{.*}} = corert.executeop({{.*}}) "tf.MatMul"({{.*}}) {T = f32, transpose_a = false, transpose_b = false} : 1
%0 = "tf.ReadVariableOp"(%arg1) {device = "/device:GPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
func.return %1 : tensor<3x3xf32>
}
// CHECK-LABEL: func @tpu_device
func.func @tpu_device(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> {
// CHECK-NOT: corert.executeop
// CHECK: tfrt_fallback_async.executeop.seq({{.*}}) key({{.*}}) cost({{.*}}) device("/device:TPU:0") "tf.ReadVariableOp"
// CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.MatMul"
%0 = "tf.ReadVariableOp"(%arg1) {device = "/device:TPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
%1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
func.return %1 : tensor<3x3xf32>
}
// CHECK-LABEL: func @tfrt_set_resource
// CHECK-SAME: ([[in_ch:%.*]]: !tfrt.chain,
func.func @tfrt_set_resource(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) {
// CHECK: [[ch0:%.*]] = tfrt_fallback_async.set_resource [[in_ch]], {{.*}} {device = "/device:CPU:0", index = 0 : i64}
// CHECK: [[ch1:%.*]], [[result:%.*]] = tfrt_fallback_async.executeop.seq([[ch0]]) key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.ReadVariableOp"
// CHECK: [[ch2:%.*]] = tfrt_fallback_async.set_resource [[ch1]], [[result]] {device = "/device:CPU:0", index = 1 : i64}
"tf._TfrtSetResource"(%arg0) {device = "/device:CPU:0", index = 0} : (tensor<3x1xf32>) -> ()
%0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
"tf._TfrtSetResource"(%0) {device = "/device:CPU:0", index = 1} : (tensor<1x3xf32>) -> ()
func.return
}
// CHECK-LABEL: func @tfrt_get_resource
func.func @tfrt_get_resource() -> tensor<3x3xf32> {
// CHECK: [[ready_ch:%.*]] = tfrt.new.chain
// CHECK: [[ch3:%.*]], [[results:%.*]]:2 = tfrt_fallback_async.get_resource [[ready_ch]] {device = "/device:CPU:0", indices = [0, 1]}
// CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.MatMul"([[results]]#0, [[results]]#1)
%a, %b = "tf._TfrtGetResource"() {device = "/device:CPU:0", indices = [0, 1], shared_name = ["", ""], container = ["", ""]} : () -> (tensor<3x1xf32>, tensor<1x3xf32>)
%1 = "tf.MatMul"(%a, %b) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
func.return %1 : tensor<3x3xf32>
}
// CHECK-LABEL: func @tensor_array
func.func @tensor_array() -> (tensor<1x1x512xf32>) {
%value = "tf.Const"() {device = "/device:CPU:0", value = dense<0.1> : tensor<1x512xf32>} : () -> (tensor<1x512xf32>)
%index = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> (tensor<i32>)
%size = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> (tensor<i32>)
%indices = "tf.Const"() {device = "/device:CPU:0", value = dense<[0]> : tensor<1xi32>} : () -> (tensor<1xi32>)
// CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayV3"
// CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayWriteV3"
// CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayGatherV3"
%handle, %flow = "tf.TensorArrayV3"(%size) {clear_after_read = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<?x512>, identical_element_shapes = true, tensor_array_name = "output"} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<1x512xf32>>>, tensor<f32>)
%flow_1 = "tf.TensorArrayWriteV3"(%handle, %index, %value, %flow) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<1x512xf32>>>, tensor<i32>, tensor<1x512xf32>, tensor<f32>) -> tensor<f32>
%result = "tf.TensorArrayGatherV3"(%handle, %indices, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", element_shape = #tf_type.shape<1x512>} : (tensor<2x!tf_type.resource<tensor<1x512xf32>>>, tensor<1xi32>, tensor<f32>) -> tensor<1x1x512xf32>
func.return %result : tensor<1x1x512xf32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/fallback_canonicalization.mlir
+84
View File
@@ -0,0 +1,84 @@
// 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 -pass-pipeline='builtin.module(func.func(canonicalize))' | FileCheck %s
// CHECK-LABEL: func @test_const_tensor_canonicalization_single_denst_tensor_operand
func.func @test_const_tensor_canonicalization_single_denst_tensor_operand() -> !tfrt_fallback.tf_tensor {
// CHECK: tfrt_fallback_async.const_dense_tensor
%a = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
%ra = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %a {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
tfrt.return %ra : !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_const_tensor_canonicalization_single_string_operand
func.func @test_const_tensor_canonicalization_single_string_operand() -> !tfrt_fallback.tf_tensor {
// CHECK: tfrt_fallback_async.const_string_tensor
%a = corert.const_string_tensor {shape = [2], value = ["string", "tensor"]}
%ra = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %a {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
tfrt.return %ra : !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_const_tensor_canonicalization_multiple_operands
func.func @test_const_tensor_canonicalization_multiple_operands() -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
// CHECK: tfrt_fallback_async.const_dense_tensor
// CHECK-NEXT: tfrt_fallback_async.const_string_tensor
%a = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
%b = corert.const_string_tensor {shape = [2], value = ["string", "tensor"]}
%ra, %rb = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %a, %b {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
tfrt.return %ra, %rb : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// Tests the case where the conversion op is partially canonicalizable.
// CHECK-LABEL: func @test_const_tensor_canonicalization_mixed_operands
// CHECK-SAME: ([[arg0:%.*]]: !corert.tensorhandle, [[arg1:%.*]]: !corert.tensorhandle)
func.func @test_const_tensor_canonicalization_mixed_operands(%arg0: !corert.tensorhandle, %arg1: !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
%a = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
%b = corert.const_dense_tensor dense<[false, true]> : tensor<2xi1>
// CHECK: [[b:%.*]] = tfrt_fallback_async.const_dense_tensor dense<[false, true]> : tensor<2xi1>
// CHECK-NEXT: [[a:%.*]] = tfrt_fallback_async.const_dense_tensor dense<[true, false]> : tensor<2xi1>
// CHECK-NEXT: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[arg0]]
// CHECK-NEXT: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[arg1]]
%ra, %rarg0, %rb, %rarg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %a, %arg0, %b, %arg1 {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
tfrt.return %ra, %rarg0, %rb, %rarg1 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// Tests that if the conversion op is partially canonicalizable, the non-canonicalizable operands are always separated into individual conversion ops.
// CHECK-LABEL: func @test_const_tensor_canonicalization_mixed_operands_no_consolidation
// CHECK-SAME: ([[arg0:%.*]]: !corert.tensorhandle, [[arg1:%.*]]: !corert.tensorhandle)
func.func @test_const_tensor_canonicalization_mixed_operands_no_consolidation(%arg0: !corert.tensorhandle, %arg1: !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
// CHECK-NEXT: tfrt_fallback_async.const_dense_tensor dense<[true, false]> : tensor<2xi1>
// CHECK-NEXT: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[arg0]]
// CHECK-NEXT: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[arg1]]
%a = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
%rarg0, %rarg1, %ra = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg0, %arg1, %a {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
tfrt.return %rarg0, %rarg1, %ra : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_remove_double_conversion
// CHECK-SAME: ([[arg:%.*]]: !tfrt_fallback.tf_tensor
func.func @test_remove_double_conversion(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor) {
// CHECK-NOT: fallback_tensor_to_corert_tensorhandle
// CHECK-NOT: corert_tensorhandle_to_fallback_tensor
%0 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %arg {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
%1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %0 {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
// CHECK: tfrt.return [[arg]]
tfrt.return %1 : !tfrt_fallback.tf_tensor
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/fallback_inline.mlir
+37
View File
@@ -0,0 +1,37 @@
// 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 -inline | FileCheck %s
func.func @_tfrt_fallback_init(%arg0: !tfrt.chain) -> !tfrt.chain {
%0 = tfrt_fallback_async.createop(%arg0) key(0) device("/device:CPU:0") "tf.Less"() {T = i32} num_args(2)
tfrt.return %0 : !tfrt.chain
}
func.func @callee(%ch: !tfrt.chain, %arg: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
%const = tfrt_fallback_async.const_dense_tensor dense<9> : tensor<i32> {_tfrt_cost = 1 : i64}
%result = tfrt_fallback_async.executeop key(0) cost(3) device("/device:CPU:0") "tf.Less"(%arg, %const) {T = i32} : 1
tfrt.return %ch, %result : !tfrt.chain, !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_inline_fallback_ops
// CHECK-SAME: ([[ch:%.*]]: !tfrt.chain, [[arg:%.*]]: !tfrt_fallback.tf_tensor
func.func @test_inline_fallback_ops(%ch: !tfrt.chain, %arg: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
// CHECK-NOT: tfrt.call
// CHECK: [[const:%.*]] = tfrt_fallback_async.const_dense_tensor dense<9> : tensor<i32>
// CHECK-NEXT: [[result:%.*]] = tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.Less"([[arg]], [[const]]) {T = i32} : 1
// CHECK-NEXT: tfrt.return [[ch]], [[result]] : !tfrt.chain, !tfrt_fallback.tf_tensor
%out_ch, %result = tfrt.call @callee(%ch, %arg) : (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
tfrt.return %out_ch, %result : !tfrt.chain, !tfrt_fallback.tf_tensor
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/func_attributes.mlir
+46
View File
@@ -0,0 +1,46 @@
// 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 -tf-executor-to-tfrt-pipeline %s | FileCheck %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 567 : i32}} {
// CHECK-LABEL: func @__inference_pruned_35
func.func @__inference_pruned_35() -> tensor<!tf_type.variant> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "flatmapdataset__4_RetVal"}} {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Const"() {device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.RangeDataset"(%0, %1, %2) {device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "__inference_Dataset_flat_map_lambda_19"} : 1
%4 = "tf.FlatMapDataset"(%3) {Targuments = [], device = "/device:CPU:0", f = @__inference_Dataset_flat_map_lambda_190, output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return %4 : tensor<!tf_type.variant>
}
// CHECK-LABEL: __inference_Dataset_flat_map_lambda_190
func.func private @__inference_Dataset_flat_map_lambda_190(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<!tf_type.variant> attributes {tf._original_func_name = "__inference_Dataset_flat_map_lambda_19", tf._tf_data_function = true, tf.signature.is_stateful} {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.Const"() {device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.RangeDataset"(%0, %2, %1) {device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.MapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64], preserve_cardinality = true, use_inter_op_parallelism = true} {f = "__inference_Dataset_map_lambda_16"} : 1
%4 = "tf.MapDataset"(%3) {device = "/device:CPU:0", f = @__inference_Dataset_map_lambda_160, f._tf_data_function = true, output_shapes = [#tf_type.shape<>], output_types = [i64], preserve_cardinality = true, use_inter_op_parallelism = true, metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
%5 = "tf.Identity"(%4) {device = "/device:CPU:0"} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return %5 : tensor<!tf_type.variant>
}
// CHECK-LABEL: __inference_Dataset_map_lambda_160
func.func private @__inference_Dataset_map_lambda_160(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<i64> attributes {tf._tf_data_function = true} {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Mul"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<i64>) -> tensor<i64>
func.return %2 : tensor<i64>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/func_attributes_multiple_callers.mlir
+54
View File
@@ -0,0 +1,54 @@
// 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 -tf-executor-to-tfrt-pipeline %s | FileCheck %s
// Checks that the ops' function attribute references the original function name
// `funcB` for `funcB_renamed` after the module is lowered to TFRT. Note that,
// `funcB_renamed` are called twice, so `CreateGuaranteeAllFuncsOneUsePass` will
// make a replicaion of `funcB_renamed` with a different name.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 567 : i32}} {
// CHECK-LABEL: @funcA
func.func @funcA() -> (tensor<!tf_type.variant>, tensor<!tf_type.variant>) attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "flatmapdataset__4_RetVal"}} {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Const"() {device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.RangeDataset"(%0, %1, %2) {device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "funcB"} : 1
%4 = "tf.FlatMapDataset"(%3) {Targuments = [], device = "/device:CPU:0", f = @funcB_renamed, output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
%5 = "tf.RangeDataset"(%1, %2, %0) {device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "funcB"} : 1
%6 = "tf.FlatMapDataset"(%5) {Targuments = [], device = "/device:CPU:0", f = @funcB_renamed, output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return %4, %6 : tensor<!tf_type.variant>, tensor<!tf_type.variant>
}
// CHECK-LABEL: @funcB_renamed
func.func private @funcB_renamed(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<!tf_type.variant> attributes {tf._original_func_name = "funcB", tf._tf_data_function = true, tf.signature.is_stateful} {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.Const"() {device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.RangeDataset"(%0, %2, %1) {device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
// CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.MapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64], preserve_cardinality = true, use_inter_op_parallelism = true} {f = "funcC"} : 1
%4 = "tf.MapDataset"(%3) {device = "/device:CPU:0", f = @funcC_renamed, f._tf_data_function = true, output_shapes = [#tf_type.shape<>], output_types = [i64], preserve_cardinality = true, use_inter_op_parallelism = true, metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
%5 = "tf.Identity"(%4) {device = "/device:CPU:0"} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
func.return %5 : tensor<!tf_type.variant>
}
// CHECK-LABEL: @funcC_renamed
func.func private @funcC_renamed(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<i64> attributes {tf._tf_data_function = true, tf._original_func_name = "funcC"} {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i64>} : () -> tensor<i64>
%1 = "tf.Mul"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
%2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<i64>) -> tensor<i64>
func.return %2 : tensor<i64>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/insert_fallback_tensor_copy.mlir
+69
View File
@@ -0,0 +1,69 @@
// 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 -tfrt-insert-fallback-tensor-copy | FileCheck %s
// CHECK-LABEL: func @test_insert_copy
// CHECK-SAME: ([[arg:%.*]]: !tfrt_fallback.tf_tensor
func.func @test_insert_copy(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) attributes {tfrt.cost_threshold = 1024} {
// CHECK: [[value:%.*]] = tfrt_fallback_async.executeop key({{.*}}) {{.*}} "tf.AddV2"([[arg]], [[arg]])
%0 = tfrt_fallback_async.executeop key(0) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
// CHECK: [[copy:%.*]] = tfrt_fallback_async.copy_if_small [[value]]
// CHECK: tfrt_fallback_async.executeop key({{.*}}) {{.*}} "tf.AddV2"([[copy]], [[copy]])
%1 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%2 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%3 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%4 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%5 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%6 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%7 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%8 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%9 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
%10 = tfrt_fallback_async.executeop key(1) cost(512) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%0, %0) {T = f32} : 1
tfrt.return %1, %2, %3 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_insert_copy_for_arg
// CHECK-SAME: ([[arg:%.*]]: !tfrt_fallback.tf_tensor
func.func @test_insert_copy_for_arg(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) attributes {tfrt.cost_threshold = 1024} {
// CHECK: [[copy:%.*]] = tfrt_fallback_async.copy_if_small [[arg]]
// CHECK: tfrt_fallback_async.executeop key({{.*}}) {{.*}} "tf.AddV2"([[copy]], [[copy]])
%0 = tfrt_fallback_async.executeop key(0) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%1 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%2 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%3 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%4 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%5 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%6 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%7 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%8 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%9 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
%10 = tfrt_fallback_async.executeop key(1) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
tfrt.return %0, %1, %2 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_no_copy_for_return
func.func @test_no_copy_for_return(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) attributes {tfrt.cost_threshold = 1024} {
// CHECK-NOT: tfrt_fallback_async.copy_if_small
%0 = tfrt_fallback_async.executeop key(0) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AddV2"(%arg, %arg) {T = f32} : 1
tfrt.return %arg, %0 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: func @test_no_copy_for_few_uses
func.func @test_no_copy_for_few_uses(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) attributes {tfrt.cost_threshold = 1024} {
// CHECK-NOT: tfrt_fallback_async.copy_if_small
%0 = tfrt_fallback_async.executeop key(0) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.Relu"(%arg) {T = f32} : 1
%1 = tfrt_fallback_async.executeop key(0) cost(1024) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.Relu"(%arg) {T = f32} : 1
tfrt.return %0, %1 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/merge_tf_if_ops.mlir
+143
View File
@@ -0,0 +1,143 @@
// 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 -tfrt-merge-tf-if-ops %s | FileCheck %s -dump-input=fail
func.func @no_side_effect_then_0(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0 = "tf.AddV2"(%x, %y) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %0 : tensor<i32>, tensor<i32>
}
func.func @no_side_effect_else_0(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0 = "tf.Const"() {value = dense<1> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.AddV2"(%x, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.AddV2"(%y, %1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1, %2 : tensor<i32>, tensor<i32>
}
func.func @no_side_effect_then_1(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
%0 = "tf.AddV2"(%x, %y) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @no_side_effect_else_1(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {value = dense<2> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.AddV2"(%x, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.AddV2"(%y, %1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
func.func @nested_if_op_then_0(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0 = "tf.AddV2"(%x, %y) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %0 : tensor<i32>, tensor<i32>
}
func.func @nested_if_op_else_0(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.return %0, %1 : tensor<i32>, tensor<i32>
}
func.func @nested_if_op_then_1(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
%0 = "tf.AddV2"(%x, %y) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @nested_if_op_else_1(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
%0 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_1, then_branch = @no_side_effect_then_1, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @merge_stateless_merged_if_0_0_then
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK: [[r0:%.*]] = "tf.AddV2"([[x]], [[y]])
// CHECK: return [[r0]], [[r0]], [[r0]]
// CHECK-LABEL: func private @merge_stateless_merged_if_0_0_else
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK-DAG: [[cst:%.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}>
// CHECK-DAG: [[cst_0:%.*]] = "tf.Const"() <{value = dense<2> : tensor<i32>}>
// CHECK: [[r0:%.*]] = "tf.AddV2"([[x]], [[cst]])
// CHECK: [[r1:%.*]] = "tf.AddV2"([[y]], [[r0]])
// CHECK: [[r2:%.*]] = "tf.AddV2"([[x]], [[cst_0]])
// CHECK: [[r3:%.*]] = "tf.AddV2"([[y]], [[r2]])
// CHECK: return [[r0]], [[r1]], [[r3]]
// CHECK-LABEL: func @merge_stateless
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>, [[cond:%.*]]: tensor<i1>)
func.func @merge_stateless(%x: tensor<i32>, %y: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
// CHECK-NEXT: [[res:%.*]]:3 = "tf.If"([[cond]], [[x]], [[y]])
// CHECK-SAME: <{else_branch = @merge_stateless_merged_if_0_0_else, is_stateless = true, then_branch = @merge_stateless_merged_if_0_0_then}>
// CHECK-SAME: (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>)
// CHECK-NEXT: return [[res]]#0, [[res]]#1, [[res]]#2
%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
%2 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_1, then_branch = @no_side_effect_then_1, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %1, %2 : tensor<i32>, tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func private @merge_nested_if_op_merged_if_0_0_then
// CHECK-SAME: ([[cond:%.*]]: tensor<i1>, [[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK-NEXT: [[r0:%.*]] = "tf.AddV2"([[x]], [[y]]) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: return [[r0]], [[r0]], [[r0]]
// CHECK-LABEL: func private @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_then
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK-NEXT: [[r0:%.*]] = "tf.AddV2"([[x]], [[y]]) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-NEXT: return [[r0]], [[r0]], [[r0]]
// CHECK-LABEL: func private @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_else
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK-NEXT: [[cst:%.*]] = "tf.Const"() <{value = dense<2> : tensor<i32>}>
// CHECK-NEXT: [[cst_0:%.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}>
// CHECK-NEXT: [[r0:%.*]] = "tf.AddV2"([[x]], [[cst_0]])
// CHECK-NEXT: [[r1:%.*]] = "tf.AddV2"([[y]], [[r0]])
// CHECK-NEXT: [[r2:%.*]] = "tf.AddV2"([[x]], [[cst]])
// CHECK-NEXT: [[r3:%.*]] = "tf.AddV2"([[y]], [[r2]])
// CHECK-NEXT: return [[r0]], [[r1]], [[r3]]
// CHECK-LABEL: func private @merge_nested_if_op_merged_if_0_0_else
// CHECK-SAME: ([[cond:%.*]]: tensor<i1>, [[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK-NEXT: [[r0:%.*]]:3 = "tf.If"(%arg0, %arg1, %arg2) <{else_branch = @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_else, is_stateless = true, then_branch = @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_then}> : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>)
// CHECK-NEXT: return [[r0]]#0, [[r0]]#1, [[r0]]#2
// CHECK-LABEL: func @merge_nested_if_op
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>, [[cond:%.*]]: tensor<i1>, [[nested_cond:%.*]]: tensor<i1>)
func.func @merge_nested_if_op(%x: tensor<i32>, %y: tensor<i32>, %cond: tensor<i1>, %nested_cond: tensor<i1>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
// CHECK-NEXT: [[res:%.*]]:3 = "tf.If"([[cond]], [[nested_cond]], [[x]], [[y]])
// CHECK-SAME: <{else_branch = @merge_nested_if_op_merged_if_0_0_else, is_stateless = true, then_branch = @merge_nested_if_op_merged_if_0_0_then}>
// CHECK-SAME: (tensor<i1>, tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>)
// CHECK-NEXT: return [[res]]#0, [[res]]#1, [[res]]#2
%0, %1 = "tf.If"(%cond, %nested_cond, %x, %y) {else_branch = @nested_if_op_else_0, then_branch = @nested_if_op_then_0, is_stateless = true} : (tensor<i1>, tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
%2 = "tf.If"(%cond, %nested_cond, %x, %y) {else_branch = @nested_if_op_else_1, then_branch = @nested_if_op_then_1, is_stateless = true} : (tensor<i1>, tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %1, %2 : tensor<i32>, tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @merge_side_effect
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>, [[cond:%.*]]: tensor<i1>)
func.func @merge_side_effect(%x: tensor<i32>, %y: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
// CHECK-NEXT: tf.If
// CHECK-SAME: is_stateless = false
// CHECK-NEXT: return
%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
%2 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_1, then_branch = @no_side_effect_then_1, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %1, %2 : tensor<i32>, tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @multiple_uses
func.func @multiple_uses(%x: tensor<i32>, %y: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
// CHECK-NEXT: tf.If
// CHECK-SAME: <{else_branch = @multiple_uses_merged_if_0_0_else, is_stateless = true, then_branch = @multiple_uses_merged_if_0_0_then}>
%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
%2 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_1, then_branch = @no_side_effect_then_1, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %1, %2 : tensor<i32>, tensor<i32>, tensor<i32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/optimize_tf_control_flow_side_effect.mlir
+190
View File
@@ -0,0 +1,190 @@
// 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 -split-input-file -tfrt-optimize-tf-control-flow-side-effect %s | FileCheck %s
func.func @no_side_effect_cond(%arg: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.Less"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func @no_side_effect_body(%arg: tensor<i32>) -> tensor<i32> {
%0 = "tf.AddV2"(%arg, %arg) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @no_side_effect_body2(%arg: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.AddV2"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @set_stateless
func.func @set_stateless(%arg: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>) {
// CHECK: tf.While
// CHECK-SAME: is_stateless = true
%0 = "tf.While"(%arg) { cond = @no_side_effect_cond, body = @no_side_effect_body, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
// CHECK: tf.If
// CHECK-SAME: is_stateless = true
%1 = "tf.If"(%cond, %arg) {else_branch = @no_side_effect_body, then_branch = @no_side_effect_body2, is_stateless = false} : (tensor<i1>, tensor<i32>) -> tensor<i32>
func.return %0, %1 : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @nested_set_stateless
func.func @nested_set_stateless(%arg: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>) {
// CHECK: tf.If
// CHECK-SAME: is_stateless = true
%1 = "tf.If"(%cond, %arg) {else_branch = @no_side_effect_body, then_branch = @no_side_effect_nested_body, is_stateless = false} : (tensor<i1>, tensor<i32>) -> tensor<i32>
// CHECK: tf.While
// CHECK-SAME: is_stateless = true
%0 = "tf.While"(%arg) { cond = @no_side_effect_cond, body = @no_side_effect_nested_body, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
func.return %0, %1 : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @no_side_effect_nested_body
func.func @no_side_effect_nested_body(%arg: tensor<i32>) -> tensor<i32> {
%const = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%cond = "tf.Less"(%arg, %const) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK: tf.If
// CHECK-SAME: is_stateless = true
%0 = "tf.If"(%cond, %arg) {else_branch = @no_side_effect_body, then_branch = @no_side_effect_body2, is_stateless = false} : (tensor<i1>, tensor<i32>) -> tensor<i32>
// CHECK: tf.While
// CHECK-SAME: is_stateless = true
%1 = "tf.While"(%0) { cond = @no_side_effect_cond, body = @no_side_effect_body, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
func.return %1 : tensor<i32>
}
// -----
func.func @no_side_effect_cond(%arg: tensor<i32>, %handle: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i1> {
%0 = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.Less"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func @side_effect_body(%arg: tensor<i32>, %handle: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) {
%0 = "tf.ReadVariableOp"(%handle) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
%1 = "tf.AddV2"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1, %handle : tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>
}
func.func @no_side_effect_body(%arg: tensor<i32>, %handle: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) {
%0 = "tf.AddV2"(%arg, %arg) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %handle : tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>
}
// CHECK-LABEL: func @no_set_stateless
func.func @no_set_stateless(%arg: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>) {
%handle = "tf.VarHandleOp"() { container = "", shared_name = "var" } : () -> tensor<!tf_type.resource<tensor<i32>>>
"tf.AssignVariableOp"(%handle, %arg) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
// CHECK: tf.While
// CHECK-SAME: is_stateless = false
%0, %h1 = "tf.While"(%arg, %handle) { cond = @no_side_effect_cond, body = @side_effect_body, is_stateless = false} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
// CHECK: tf.If
// CHECK-SAME: is_stateless = false
%1, %h2 = "tf.If"(%cond, %arg, %handle) {else_branch = @no_side_effect_body, then_branch = @side_effect_body, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
func.return %0, %1 : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @side_effect_nested_body
func.func @side_effect_nested_body(%arg: tensor<i32>, %handle: tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) {
%const = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
// CHECK: tf.While
// CHECK-SAME: is_stateless = false
%0, %h1 = "tf.While"(%arg, %handle) { cond = @no_side_effect_cond, body = @side_effect_body, is_stateless = false} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
%cond = "tf.Less"(%0, %const) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK: tf.If
// CHECK-SAME: is_stateless = false
%1, %h2 = "tf.If"(%cond, %0, %handle) {else_branch = @no_side_effect_body, then_branch = @side_effect_body, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
func.return %1, %handle : tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>
}
// CHECK-LABEL: func @nested_no_set_stateless
func.func @nested_no_set_stateless(%arg: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>, tensor<i32>) {
%handle = "tf.VarHandleOp"() { container = "", shared_name = "var" } : () -> tensor<!tf_type.resource<tensor<i32>>>
"tf.AssignVariableOp"(%handle, %arg) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
// CHECK: tf.While
// CHECK-SAME: is_stateless = false
%0, %h1 = "tf.While"(%arg, %handle) { cond = @no_side_effect_cond, body = @side_effect_nested_body, is_stateless = false} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
// CHECK: tf.If
// CHECK-SAME: is_stateless = false
%1, %h2 = "tf.If"(%cond, %arg, %handle) {else_branch = @no_side_effect_body, then_branch = @side_effect_nested_body, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>)
func.return %0, %1 : tensor<i32>, tensor<i32>
}
// -----
// Set stateless if the body contains only read-only side-effecting ops.
func.func private @cond(%arg: tensor<i32>, %handle: tensor<!tf_type.resource>) -> tensor<i1> {
%0 = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.Less"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func private @body(%arg: tensor<i32>, %handle: tensor<!tf_type.resource>) -> (tensor<i32>, tensor<!tf_type.resource>) {
%default = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%0 = "tf.LookupTableFindV2"(%handle, %arg, %default) : (tensor<!tf_type.resource>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %handle : tensor<i32>, tensor<!tf_type.resource>
}
// CHECK-LABEL: func @set_readonly_stateless
func.func @set_readonly_stateless(%arg: tensor<i32>, %cond: tensor<i1>) -> tensor<i32> {
%handle = "tf.HashTableV2"() {container = "", key_dtype = i32, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i32} : () -> tensor<!tf_type.resource>
// CHECK: tf.While
// CHECK-SAME: is_stateless = true
%0, %handle_2 = "tf.While"(%arg, %handle) { cond = @cond, body = @body, is_stateless = false} : (tensor<i32>, tensor<!tf_type.resource>) -> (tensor<i32>, tensor<!tf_type.resource>)
func.return %0: tensor<i32>
}
// -----
// Set stateless if write side-effecting ops are inside the initializer.
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init]} : () -> ()
func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
%keys = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
%values = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
%handle = "tf.HashTableV2"() {container = "", key_dtype = i32, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i32} : () -> tensor<!tf_type.resource>
"tf.LookupTableImportV2"(%handle, %keys, %values) : (tensor<!tf_type.resource>, tensor<4xi32>, tensor<4xi32>) -> ()
func.return
}
func.func private @cond(%arg: tensor<i32>, %handle: tensor<!tf_type.resource>) -> tensor<i1> {
%0 = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.Less"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func private @body(%arg: tensor<i32>, %handle: tensor<!tf_type.resource>) -> (tensor<i32>, tensor<!tf_type.resource>) {
%default = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
%0 = "tf.LookupTableFindV2"(%handle, %arg, %default) : (tensor<!tf_type.resource>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0, %handle : tensor<i32>, tensor<!tf_type.resource>
}
// CHECK-LABEL: func @set_readonly_stateless
func.func @set_readonly_stateless(%arg: tensor<i32> {tf_saved_model.index_path = ["arg"]}, %cond: tensor<i1> {tf_saved_model.index_path = ["cond"]}) -> (tensor<i32> {tf_saved_model.index_path = ["output"]}) attributes {tf_saved_model.exported_names = ["main"]} {
%handle = "tf.HashTableV2"() {container = "", key_dtype = i32, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i32} : () -> tensor<!tf_type.resource>
// CHECK: tf.While
// CHECK-SAME: is_stateless = true
%0, %handle_2 = "tf.While"(%arg, %handle) { cond = @cond, body = @body, is_stateless = false} : (tensor<i32>, tensor<!tf_type.resource>) -> (tensor<i32>, tensor<!tf_type.resource>)
func.return %0: tensor<i32>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/remove_tf_if_const_args.mlir
+63
View File
@@ -0,0 +1,63 @@
// 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 -tfrt-remove-tf-if-const-args %s | FileCheck %s -dump-input-filter=all
func.func @then(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>) {
%0 = "tf.AddV2"(%x, %y) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @else(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>) {
%0 = "tf.Const"() {value = dense<1> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.AddV2"(%x, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.AddV2"(%y, %1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %2 : tensor<i32>
}
// CHECK-LABEL: func private @then_removed_const_args_0
// CHECK-SAME: ([[x:%.*]]: tensor<i32>)
// CHECK: [[const:%.*]] = "tf.Const"
// CHECK-SAME: value = dense<10> : tensor<i32>
// CHECK: [[r:%.*]] = "tf.StatefulPartitionedCall"([[x]], [[const]])
// CHECK-SAME: f = @then}
// CHECK: return [[r]]
// CHECK-LABEL: func private @else_removed_const_args_0
// CHECK-SAME: ([[x:%.*]]: tensor<i32>)
// CHECK: [[const:%.*]] = "tf.Const"
// CHECK-SAME: value = dense<10> : tensor<i32>
// CHECK: [[r:%.*]] = "tf.StatefulPartitionedCall"([[x]], [[const]])
// CHECK-SAME: f = @else}
// CHECK: return [[r]]
// CHECK-LABEL: func @remove_const_args
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[cond:%.*]]: tensor<i1>)
func.func @remove_const_args(%x: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>) {
%0 = "tf.Const"() {value = dense<10> : tensor<i32> } : () -> tensor<i32>
// CHECK: [[res:%.*]] = "tf.If"([[cond]], [[x]])
// CHECK-SAME: {else_branch = @else_removed_const_args_0, is_stateless = false, then_branch = @then_removed_const_args_0}
// CHECK-NEXT: return [[res]]
%1 = "tf.If"(%cond, %x, %0) {else_branch = @else, then_branch = @then, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @multiple_uses
func.func @multiple_uses(%x: tensor<i32>, %cond: tensor<i1>) -> (tensor<i32>) {
%0 = "tf.Const"() {value = dense<10> : tensor<i32> } : () -> tensor<i32>
// CHECK: [[res:%.*]] = "tf.If"
// CHECK-SAME: {else_branch = @else_removed_const_args_1, is_stateless = false, then_branch = @then_removed_const_args_1}
%1 = "tf.If"(%cond, %0, %x) {else_branch = @else, then_branch = @then, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/reorder_assert.mlir
+61
View File
@@ -0,0 +1,61 @@
// 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 -tfrt-reorder-tf-assert %s | FileCheck %s
// CHECK-LABEL: @reorder_assert
func.func @reorder_assert(%key0: tensor<!tf_type.string>, %key1: tensor<!tf_type.string>) -> (tensor<i64>, tensor<i64>) {
%error_message = "tf.Const"() {value = dense<"error"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%default = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
%handle = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
// CHECK: tf.LookupTableFindV2
// CHECK-NOT: tf.Assert
// CHECK: tf.LookupTableFindV2
// CHECK: tf.Assert
%value0 = "tf.LookupTableFindV2"(%handle, %key0, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
%cond = "tf.Equal"(%value0, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0", incompatible_shape_error = true} : (tensor<i64>, tensor<i64>) -> tensor<i1>
"tf.Assert"(%cond, %error_message) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
%value1 = "tf.LookupTableFindV2"(%handle, %key1, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
func.return %value0, %value1 : tensor<i64>, tensor<i64>
}
func.func private @else_branch(%arg0: tensor<i1>) -> tensor<i1> {
%cst = "tf.Const"() {value = dense<"Empty SparseTensor with shape"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
"tf.Assert"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
func.return %arg0 : tensor<i1>
}
func.func private @then_branch(%arg0: tensor<i1>) -> tensor<i1> {
func.return %arg0 : tensor<i1>
}
// CHECK-LABEL: @reorder_assert_only_if
func.func @reorder_assert_only_if(%key0: tensor<!tf_type.string>, %key1: tensor<!tf_type.string>) -> (tensor<i64>, tensor<i64>) {
%error_message = "tf.Const"() {value = dense<"error"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%default = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
%handle = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
// CHECK: tf.LookupTableFindV2
// CHECK-NOT: tf.If
// CHECK: tf.LookupTableFindV2
// CHECK: tf.If
%value0 = "tf.LookupTableFindV2"(%handle, %key0, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
%cond = "tf.Equal"(%value0, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0", incompatible_shape_error = true} : (tensor<i64>, tensor<i64>) -> tensor<i1>
%unused = "tf.If"(%cond, %cond) {device = "/job:localhost/replica:0/task:0/device:CPU:0", else_branch = @else_branch, is_stateless = false, then_branch = @then_branch} : (tensor<i1>, tensor<i1>) -> tensor<i1>
%value1 = "tf.LookupTableFindV2"(%handle, %key1, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
func.return %value0, %value1 : tensor<i64>, tensor<i64>
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/side_effects.mlir
+29
View File
@@ -0,0 +1,29 @@
// 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 -tf-to-tfrt %s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @assign_variable
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain) -> !tfrt.chain
func.func @assign_variable() {
// CHECK: [[ch1:%.*]], %results = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/device:CPU:0") "tf.VarHandleOp"
// CHECK-NEXT: [[ch2:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(1) cost({{.*}}) device("/device:CPU:0") "tf.AssignVariableOp"
// CHECK-NEXT: [[out_ch:%.*]] = tfrt.merge.chains [[ch1]], [[ch2]]
// CHECK-NEXT: tfrt.return [[out_ch]]
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<4.200000e+01> : tensor<f32>} : () -> tensor<f32>
%1 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<f32>>>
"tf.AssignVariableOp"(%1, %0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<f32>>>, tensor<f32>) -> ()
func.return
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/tf_to_corert_pipeline.mlir
+87
View File
@@ -0,0 +1,87 @@
// 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 -tf-executor-to-tfrt-pipeline="enable-optimizer=true tfrt-cost-threshold=1024" %s | FileCheck %s --dump-input=fail
// CHECK: tfrt.cost_threshold = 1024 : i64
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 462 : i32}} {
// CHECK-LABEL: func @__forward_call_369
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
// CHECK-SAME: [[arg1:%.*]]: !tfrt_fallback.tf_tensor {tf._user_specified_name = "inputs"},
// CHECK-SAME: [[arg2:%.*]]: !tfrt_fallback.tf_tensor, [[arg3:%.*]]: !tfrt_fallback.tf_tensor, [[arg4:%.*]]: !tfrt_fallback.tf_tensor, [[arg5:%.*]]: !tfrt_fallback.tf_tensor)
// CHECK-SAME: -> (!tfrt.chain
// CHECK: [[o1:%.*]] = tfrt_fallback_async.const_dense_tensor
// CHECK: [[o2_chain:%.*]], [[o2:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"([[arg3]])
// CHECK-NEXT: [[o3_chain:%.*]], [[o3:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"([[arg2]])
// CHECK-NEXT: [[o4_chain:%.*]], [[o4:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"([[arg5]])
// CHECK-NEXT: [[o5_chain:%.*]], [[o5:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"([[arg4]])
// CHECK-NEXT: [[o6:%.*]] = tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf._FusedConv2D"([[arg1]], [[o3]], [[o2]])
// CHECK-NEXT: [[o7:%.*]] = tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.AvgPool"([[o6]])
// CHECK-NEXT: [[o8:%.*]] = tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.Reshape"([[o7]], [[o1]])
// CHECK-NEXT: [[o9:%.*]] = tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf._FusedMatMul"([[o8]], [[o5]], [[o4]])
// CHECK-NEXT: [[out_chain:%.*]] = tfrt.merge.chains [[o2_chain]], [[o3_chain]], [[o4_chain]], [[o5_chain]]
// CHECK-NEXT: tfrt.return [[out_chain]], [[o9]], [[o5]], [[o8]], [[o6]], [[arg1]], [[o3]] : !tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
func.func @__forward_call_369(%arg0: tensor<16x224x224x3xf32> {tf._user_specified_name = "inputs"}, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<*x!tf_type.resource>, %arg3: tensor<*x!tf_type.resource>, %arg4: tensor<*x!tf_type.resource>) -> (tensor<?x?xf32>, tensor<*xf32>, tensor<?x16384xf32>, tensor<16x112x112x?xf32>, tensor<16x224x224x3xf32>, tensor<*xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "inputs_0,conv1_conv2d_readvariableop_resource,conv1_biasadd_readvariableop_resource,fc1000_matmul_readvariableop_resource,fc1000_biasadd_readvariableop_resource", outputs = "identity_RetVal,fc1000_matmul_readvariableop_RetVal,flatten_reshape_RetVal,relu_RetVal,inputs_RetVal,conv1_conv2d_readvariableop_RetVal"}} {
%0:6 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) {device = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.ReadVariableOp"(%arg4) {device = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
%outputs_4, %control_5 = tf_executor.island wraps "tf.ReadVariableOp"(%arg3) {device = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
%outputs_6, %control_7 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<[-1, 16384]> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Conv2D"(%arg0, %outputs_0) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<16x224x224x3xf32>, tensor<*xf32>) -> tensor<16x112x112x?xf32>
%outputs_10, %control_11 = tf_executor.island wraps "tf.BiasAdd"(%outputs_8, %outputs) {data_format = "NHWC", device = ""} : (tensor<16x112x112x?xf32>, tensor<*xf32>) -> tensor<16x112x112x?xf32>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Relu"(%outputs_10) {device = ""} : (tensor<16x112x112x?xf32>) -> tensor<16x112x112x?xf32>
%outputs_14, %control_15 = tf_executor.island wraps "tf.AvgPool"(%outputs_12) {data_format = "NHWC", device = "", ksize = [1, 7, 7, 1], padding = "VALID", strides = [1, 7, 7, 1]} : (tensor<16x112x112x?xf32>) -> tensor<16x16x16x?xf32>
%outputs_16, %control_17 = tf_executor.island wraps "tf.Reshape"(%outputs_14, %outputs_6) {device = ""} : (tensor<16x16x16x?xf32>, tensor<2xi32>) -> tensor<?x16384xf32>
%outputs_18, %control_19 = tf_executor.island wraps "tf.MatMul"(%outputs_16, %outputs_4) {device = "", transpose_a = false, transpose_b = false} : (tensor<?x16384xf32>, tensor<*xf32>) -> tensor<?x?xf32>
%outputs_20, %control_21 = tf_executor.island wraps "tf.BiasAdd"(%outputs_18, %outputs_2) {data_format = "NHWC", device = ""} : (tensor<?x?xf32>, tensor<*xf32>) -> tensor<?x?xf32>
%outputs_22, %control_23 = tf_executor.island wraps "tf.Identity"(%outputs_20) {device = ""} : (tensor<?x?xf32>) -> tensor<?x?xf32>
tf_executor.fetch %outputs_22, %outputs_4, %outputs_16, %outputs_12, %arg0, %outputs_0 : tensor<?x?xf32>, tensor<*xf32>, tensor<?x16384xf32>, tensor<16x112x112x?xf32>, tensor<16x224x224x3xf32>, tensor<*xf32>
}
func.return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 : tensor<?x?xf32>, tensor<*xf32>, tensor<?x16384xf32>, tensor<16x112x112x?xf32>, tensor<16x224x224x3xf32>, tensor<*xf32>
}
func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func @while_test
// CHECK-SAME: ([[ARG0:%.+]]: !tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
func.func @while_test() -> (tensor<i32>) {
// The predicate function should be inlined.
// CHECK-DAG: tfrt_fallback_async.const_dense_tensor dense<9> : tensor<i32>
// CHECK-DAG: tfrt_fallback_async.const_dense_tensor dense<0> : tensor<i32>
// CHECK-NEXT: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.Less"
// CHECK-NEXT: [[pred:%.*]] = tfrt_fallback_async.predicate
// CHECK-NEXT: tfrt.while [[pred]] @"[[while_func_prefix:.*]]/tfrt_body_1"
// CHECK-NEXT: tfrt.merge.chains
// CHECK-NEXT: tfrt.return
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
func.return %1 : tensor<i32>
}
// CHECK: func @"[[while_func_prefix]]/tfrt_body_1"
// CHECK-NOT: tfrt.call
// CHECK: func @"[[while_cond_prefix:.*]]/tfrt_predicate"
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/tf_to_corert_pipeline_refvar.mlir
+32
View File
@@ -0,0 +1,32 @@
// 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 -tf-executor-to-tfrt-pipeline %s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @__inference_pruned_131
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
// CHECK-NEXT: [[o_chain:%.*]], [[o:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.VarHandleOp"()
// CHECK-NEXT: [[o_chain_0:%.*]], [[o1:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(1) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"([[o]]) {dtype = f32} : 1
// CHECK-NEXT: [[out_ch:%.*]] = tfrt.merge.chains [[o_chain]], [[o_chain_0]]
// CHECK-NEXT: tfrt.return [[out_ch]], [[o1]] : !tfrt.chain, !tfrt_fallback.tf_tensor
module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0"], tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 679 : i32}} {
func.func @__inference_pruned_131() -> tensor<*xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "variable", outputs = "identity_retval_RetVal"}} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.VariableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shape = #tf_type.shape<>, shared_name = "v_load_44"} : () -> tensor<!tf_type.f32ref>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Identity"(%outputs) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.f32ref>) -> tensor<*xf32>
tf_executor.fetch %outputs_0 : tensor<*xf32>
}
func.return %0 : tensor<*xf32>
}
}
tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/whileop.mlir
+47
View File
@@ -0,0 +1,47 @@
// 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 -tf-executor-to-tfrt-pipeline="enable-optimizer=true tfrt-cost-threshold=1024" %s | FileCheck %s --dump-input=fail
// Check that unused While op results and the associated ops are removed.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 462 : i32}} {
func.func @while_cond_lt9(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
func.return %1 : tensor<i1>
}
func.func @while_body_add2(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.AddV2"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%2 = "tf.Div"(%arg1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1, %2 : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @while_test_remove_unused_results
// CHECK: [[pred:%.*]] = tfrt_fallback_async.predicate
// CHECK-NEXT: tfrt.while [[pred]] @"[[while_func_prefix:.*]]/tfrt_body_1"
// CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
// CHECK-NOT: func.call
func.func @while_test_remove_unused_results(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = "tf.While"(%arg0, %arg1) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
%1:2 = func.call @while_body_add2(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.return %0#0, %1#0 : tensor<i32>, tensor<i32>
}
// CHECK: func @"[[while_func_prefix]]/tfrt_body_1"
// CHECK: "tf.AddV2"
// CHECK-NOT: "tf.Div"
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/BUILD
+142
View File
@@ -0,0 +1,142 @@
load("@tf_runtime//tools:mlir_to_bef.bzl", "glob_tfrt_lit_tests", "mlir_to_bef")
load("//tensorflow:tensorflow.bzl", "tf_cc_shared_test")
# copybara:uncomment load("//third_party/tf_runtime_google/cpp_tests:gen_tests.bzl", "tfrt_cc_test_and_strict_benchmark")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
srcs = [
"//tensorflow/compiler/mlir/tfrt:tfrt_fallback_translate",
"//tensorflow/core/runtime_fallback:tf_bef_executor",
"//tensorflow/core/runtime_fallback/util:fallback_test_util",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
"@llvm-project//mlir:run_lit.sh",
"@tf_runtime//tools:tfrt_translate",
],
)
# copybara:uncomment_begin(TFRT lit issue b/290857552)
# glob_tfrt_lit_tests(
# data = [":test_utilities"],
# # Custom driver is unsupported in OSS. Fails if one is provided.
# # copybara:uncomment driver = "//tensorflow/compiler/mlir:run_lit.sh",
# exclude = [
# "compile.benchmark.large.mlir",
# "batch_function_fallback.mlir",
# "create_op.mlir",
# "custom_thread_pool.mlir",
# ],
# # copybara:uncomment flaky = ["compile.error.mlir"],
# size_override = {
# "compile.benchmark.small.mlir": "medium",
# "batching_fallback.mlir": "medium",
# },
# tags_override = {
# "async_op_thread.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "batching_fallback.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "compile.benchmark.small.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "convert_tensorhandle_to_fallback_tensor.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "fallback.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "fallback_tensor_conversion_host.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "kernel_fallback_op_handler.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "mnist.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "runtime_fallback_op_handler.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "tf_delegate.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "tf_ops.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "tf_ops_error.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# "tfrt_forwarding.mlir": ["nomsan"], # Can't instrument code in precompiled lib (cuDNN)
# },
# tfrt_translate = "//tensorflow/compiler/mlir/tfrt:tfrt_fallback_translate",
# )
# copybara:uncomment_end
mlir_to_bef(
name = "batch_function_fallback.mlir",
tfrt_translate = "//tensorflow/compiler/mlir/tfrt:tfrt_fallback_translate",
)
mlir_to_bef(
name = "create_op.mlir",
tfrt_translate = "//tensorflow/compiler/mlir/tfrt:tfrt_fallback_translate",
)
mlir_to_bef(
name = "custom_thread_pool.mlir",
tfrt_translate = "//tensorflow/compiler/mlir/tfrt:tfrt_fallback_translate",
)
# copybara:uncomment_begin(internal benchmarking)
# # C++ benchmarks for batch function runtime fallback.
# tfrt_cc_test_and_strict_benchmark(
# name = "batch_function_fallback_benchmark_test",
# srcs = ["batch_function_fallback_benchmark_test.cc"],
# data = ["batch_function_fallback.mlir.bef"],
# enable_xprof = True,
# includes = ["third_party/tf_runtime/include"],
# owners = ["tf-runtime-testing"],
# tags = [
# "need_main",
# "no_gpu",
# ],
# deps = [
# "//devtools/build/runtime:get_runfiles_dir",
# "@com_google_absl//absl/log:check",
# "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
# "//tensorflow/core/platform:env",
# "//tensorflow/core/platform:resource_loader",
# "//tensorflow/core/platform:status",
# "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_op_handler",
# "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_tensor",
# "//tensorflow/core/runtime_fallback/runtime:runtime_fallback_alwayslink",
# "//tensorflow/core/runtime_fallback/util:fallback_test_util",
# "//tensorflow/core/runtime_fallback/util:tensor_util",
# "//tensorflow/core/tfrt/utils:fallback_tensor",
# "@eigen_archive//:eigen3",
# "@tf_runtime//:bef",
# "@tf_runtime//:befexecutor",
# "@tf_runtime//:core_runtime_alwayslink",
# "@tf_runtime//:hostcontext_alwayslink",
# "@tf_runtime//:mlirtobef",
# "@tf_runtime//:support",
# "@tf_runtime//:tensor",
# "@tf_runtime//backends/cpu:core_runtime_alwayslink",
# "@tf_runtime//backends/cpu:test_ops_alwayslink",
# ],
# )
# copybara:uncomment_end
tf_cc_shared_test(
name = "kernel_fallback_compat_test",
srcs = ["kernel_fallback_compat_test.cc"],
data = [
"create_op.mlir.bef",
"custom_thread_pool.mlir.bef",
],
tags = ["no_oss"],
deps = [
"//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
"//tensorflow/core:all_kernels",
"//tensorflow/core:lib",
"//tensorflow/core/platform:resource_loader",
"//tensorflow/core/runtime_fallback/kernel:kernel_fallback_compat_request_state",
"//tensorflow/core/runtime_fallback/runtime:runtime_fallback_alwayslink",
"//tensorflow/core/runtime_fallback/util:fallback_test_util",
"//tensorflow/core/tfrt/fallback:op_kernel_runner",
"//tensorflow/core/tfrt/runtime",
"//tensorflow/core/tfrt/utils:thread_pool",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@tf_runtime//:bef",
"@tf_runtime//:befexecutor",
"@tf_runtime//:core_runtime",
"@tf_runtime//:hostcontext",
"@tf_runtime//:init_tfrt_dialects",
"@tf_runtime//:support",
"@tf_runtime//:tracing",
],
)
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/async_op_thread.mlir
+44
View File
@@ -0,0 +1,44 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd 2>&1 | FileCheck %s
func.func @test_async_op_kernel_thread() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("/CPU:0") "tf.Const"()
{ dtype = i32, value = dense<[2]> : tensor<1xi32> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("/CPU:0")
"tf.TestAsyncTfrtAsyncThread"() {T = i32} num_args(1)
%ch3 = tfrt_fallback_async.createop(%ch2) key(2) device("/CPU:0")
"tf.TestPrintThreadName"() num_args(0)
%ch4, %0 = tfrt_fallback_async.executeop.seq(%ch3) key(0) cost(100)
device("/CPU:0") "tf.Const"()
{ dtype = i32, value = dense<[2]> : tensor<1xi32> } : 1
// Given TestAsyncTfrtAsyncThread will invoke done callback in thread with
// name "test_thread_in_compute_async",
// CHECK: TestAsyncTfrtAsyncThread thread name: test_thread_in_compute_async
%ch5, %1 = tfrt_fallback_async.executeop.seq(%ch4) key(1) cost(100)
device("/CPU:0") "tf.TestAsyncTfrtAsyncThread"(%0) {T = i32} : 1
// ... when TestPrintThreadName is part of the next op in sequence of
// TestAsyncTfrtAsyncThread op,
// ... then TestPrintThreadName should not run in the thread in
// TestAsyncTfrtAsyncThread.
// CHECK-NOT: TestPrintThreadName thread name: test_thread_in_compute_async
%ch6 = tfrt_fallback_async.executeop.seq(%ch5) key(2) cost(100)
device("/CPU:0") "tf.TestPrintThreadName"() : 0
tfrt.return %ch6: !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/batch_function_fallback.mlir
+175
View File
@@ -0,0 +1,175 @@
// 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 @f(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor, %arg2: !tfrt_fallback.tf_tensor, %arg3: !tfrt_fallback.tf_tensor, %arg4: !tfrt_fallback.tf_tensor,
%arg5: !tfrt_fallback.tf_tensor, %arg6: !tfrt_fallback.tf_tensor, %arg7: !tfrt_fallback.tf_tensor, %arg8: !tfrt_fallback.tf_tensor, %arg9: !tfrt_fallback.tf_tensor,
%arg10: !tfrt_fallback.tf_tensor, %arg11: !tfrt_fallback.tf_tensor, %arg12: !tfrt_fallback.tf_tensor, %arg13: !tfrt_fallback.tf_tensor, %arg14: !tfrt_fallback.tf_tensor,
%arg15: !tfrt_fallback.tf_tensor, %arg16: !tfrt_fallback.tf_tensor, %arg17: !tfrt_fallback.tf_tensor, %arg18: !tfrt_fallback.tf_tensor, %arg19: !tfrt_fallback.tf_tensor,
%arg20: !tfrt_fallback.tf_tensor, %arg21: !tfrt_fallback.tf_tensor, %arg22: !tfrt_fallback.tf_tensor, %arg23: !tfrt_fallback.tf_tensor, %arg24: !tfrt_fallback.tf_tensor,
%arg25: !tfrt_fallback.tf_tensor, %arg26: !tfrt_fallback.tf_tensor, %arg27: !tfrt_fallback.tf_tensor, %arg28: !tfrt_fallback.tf_tensor, %arg29: !tfrt_fallback.tf_tensor,
%arg30: !tfrt_fallback.tf_tensor, %arg31: !tfrt_fallback.tf_tensor, %arg32: !tfrt_fallback.tf_tensor, %arg33: !tfrt_fallback.tf_tensor, %arg34: !tfrt_fallback.tf_tensor,
%arg35: !tfrt_fallback.tf_tensor, %arg36: !tfrt_fallback.tf_tensor, %arg37: !tfrt_fallback.tf_tensor, %arg38: !tfrt_fallback.tf_tensor, %arg39: !tfrt_fallback.tf_tensor,
%arg40: !tfrt_fallback.tf_tensor, %arg41: !tfrt_fallback.tf_tensor, %arg42: !tfrt_fallback.tf_tensor, %arg43: !tfrt_fallback.tf_tensor, %arg44: !tfrt_fallback.tf_tensor,
%arg45: !tfrt_fallback.tf_tensor, %arg46: !tfrt_fallback.tf_tensor, %arg47: !tfrt_fallback.tf_tensor, %arg48: !tfrt_fallback.tf_tensor, %arg49: !tfrt_fallback.tf_tensor,
%arg50: !tfrt_fallback.tf_tensor, %arg51: !tfrt_fallback.tf_tensor, %arg52: !tfrt_fallback.tf_tensor, %arg53: !tfrt_fallback.tf_tensor, %arg54: !tfrt_fallback.tf_tensor,
%arg55: !tfrt_fallback.tf_tensor, %arg56: !tfrt_fallback.tf_tensor, %arg57: !tfrt_fallback.tf_tensor, %arg58: !tfrt_fallback.tf_tensor, %arg59: !tfrt_fallback.tf_tensor,
%arg60: !tfrt_fallback.tf_tensor, %arg61: !tfrt_fallback.tf_tensor, %arg62: !tfrt_fallback.tf_tensor, %arg63: !tfrt_fallback.tf_tensor, %arg64: !tfrt_fallback.tf_tensor,
%arg65: !tfrt_fallback.tf_tensor, %arg66: !tfrt_fallback.tf_tensor, %arg67: !tfrt_fallback.tf_tensor, %arg68: !tfrt_fallback.tf_tensor, %arg69: !tfrt_fallback.tf_tensor,
%arg70: !tfrt_fallback.tf_tensor, %arg71: !tfrt_fallback.tf_tensor, %arg72: !tfrt_fallback.tf_tensor, %arg73: !tfrt_fallback.tf_tensor, %arg74: !tfrt_fallback.tf_tensor,
%arg75: !tfrt_fallback.tf_tensor, %arg76: !tfrt_fallback.tf_tensor, %arg77: !tfrt_fallback.tf_tensor, %arg78: !tfrt_fallback.tf_tensor, %arg79: !tfrt_fallback.tf_tensor,
%arg80: !tfrt_fallback.tf_tensor, %arg81: !tfrt_fallback.tf_tensor, %arg82: !tfrt_fallback.tf_tensor, %arg83: !tfrt_fallback.tf_tensor, %arg84: !tfrt_fallback.tf_tensor,
%arg85: !tfrt_fallback.tf_tensor, %arg86: !tfrt_fallback.tf_tensor, %arg87: !tfrt_fallback.tf_tensor, %arg88: !tfrt_fallback.tf_tensor, %arg89: !tfrt_fallback.tf_tensor,
%arg90: !tfrt_fallback.tf_tensor, %arg91: !tfrt_fallback.tf_tensor, %arg92: !tfrt_fallback.tf_tensor, %arg93: !tfrt_fallback.tf_tensor, %arg94: !tfrt_fallback.tf_tensor,
%arg95: !tfrt_fallback.tf_tensor, %arg96: !tfrt_fallback.tf_tensor, %arg97: !tfrt_fallback.tf_tensor, %arg98: !tfrt_fallback.tf_tensor, %arg99: !tfrt_fallback.tf_tensor,
%arg100: !tfrt_fallback.tf_tensor, %arg101: !tfrt_fallback.tf_tensor, %arg102: !tfrt_fallback.tf_tensor, %arg103: !tfrt_fallback.tf_tensor, %arg104: !tfrt_fallback.tf_tensor,
%arg105: !tfrt_fallback.tf_tensor, %arg106: !tfrt_fallback.tf_tensor, %arg107: !tfrt_fallback.tf_tensor, %arg108: !tfrt_fallback.tf_tensor, %arg109: !tfrt_fallback.tf_tensor,
%arg110: !tfrt_fallback.tf_tensor, %arg111: !tfrt_fallback.tf_tensor, %arg112: !tfrt_fallback.tf_tensor)
-> (!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
tfrt.return %arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10,%arg11, %arg12, %arg13, %arg14,
%arg15, %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23, %arg24, %arg25, %arg26, %arg27, %arg28, %arg29,
%arg30, %arg31, %arg32, %arg33, %arg34, %arg35, %arg36, %arg37, %arg38, %arg39, %arg40, %arg41, %arg42, %arg43, %arg44,
%arg45, %arg46, %arg47, %arg48, %arg49, %arg50, %arg51, %arg52, %arg53, %arg54, %arg55, %arg56, %arg57, %arg58, %arg59,
%arg60, %arg61, %arg62, %arg63, %arg64, %arg65, %arg66, %arg67, %arg68, %arg69, %arg70, %arg71, %arg72, %arg73, %arg74,
%arg75, %arg76, %arg77, %arg78, %arg79, %arg80, %arg81, %arg82, %arg83, %arg84, %arg85, %arg86, %arg87, %arg88, %arg89,
%arg90, %arg91, %arg92, %arg93, %arg94, %arg95, %arg96, %arg97, %arg98, %arg99, %arg100, %arg101, %arg102, %arg103, %arg104,
%arg105, %arg106, %arg107, %arg108, %arg109, %arg110, %arg111, %arg112
:!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
func.func @main(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor, %arg2: !tfrt_fallback.tf_tensor, %arg3: !tfrt_fallback.tf_tensor, %arg4: !tfrt_fallback.tf_tensor,
%arg5: !tfrt_fallback.tf_tensor, %arg6: !tfrt_fallback.tf_tensor, %arg7: !tfrt_fallback.tf_tensor, %arg8: !tfrt_fallback.tf_tensor, %arg9: !tfrt_fallback.tf_tensor,
%arg10: !tfrt_fallback.tf_tensor, %arg11: !tfrt_fallback.tf_tensor, %arg12: !tfrt_fallback.tf_tensor, %arg13: !tfrt_fallback.tf_tensor, %arg14: !tfrt_fallback.tf_tensor,
%arg15: !tfrt_fallback.tf_tensor, %arg16: !tfrt_fallback.tf_tensor, %arg17: !tfrt_fallback.tf_tensor, %arg18: !tfrt_fallback.tf_tensor, %arg19: !tfrt_fallback.tf_tensor,
%arg20: !tfrt_fallback.tf_tensor, %arg21: !tfrt_fallback.tf_tensor, %arg22: !tfrt_fallback.tf_tensor, %arg23: !tfrt_fallback.tf_tensor, %arg24: !tfrt_fallback.tf_tensor,
%arg25: !tfrt_fallback.tf_tensor, %arg26: !tfrt_fallback.tf_tensor, %arg27: !tfrt_fallback.tf_tensor, %arg28: !tfrt_fallback.tf_tensor, %arg29: !tfrt_fallback.tf_tensor,
%arg30: !tfrt_fallback.tf_tensor, %arg31: !tfrt_fallback.tf_tensor, %arg32: !tfrt_fallback.tf_tensor, %arg33: !tfrt_fallback.tf_tensor, %arg34: !tfrt_fallback.tf_tensor,
%arg35: !tfrt_fallback.tf_tensor, %arg36: !tfrt_fallback.tf_tensor, %arg37: !tfrt_fallback.tf_tensor, %arg38: !tfrt_fallback.tf_tensor, %arg39: !tfrt_fallback.tf_tensor,
%arg40: !tfrt_fallback.tf_tensor, %arg41: !tfrt_fallback.tf_tensor, %arg42: !tfrt_fallback.tf_tensor, %arg43: !tfrt_fallback.tf_tensor, %arg44: !tfrt_fallback.tf_tensor,
%arg45: !tfrt_fallback.tf_tensor, %arg46: !tfrt_fallback.tf_tensor, %arg47: !tfrt_fallback.tf_tensor, %arg48: !tfrt_fallback.tf_tensor, %arg49: !tfrt_fallback.tf_tensor,
%arg50: !tfrt_fallback.tf_tensor, %arg51: !tfrt_fallback.tf_tensor, %arg52: !tfrt_fallback.tf_tensor, %arg53: !tfrt_fallback.tf_tensor, %arg54: !tfrt_fallback.tf_tensor,
%arg55: !tfrt_fallback.tf_tensor, %arg56: !tfrt_fallback.tf_tensor, %arg57: !tfrt_fallback.tf_tensor, %arg58: !tfrt_fallback.tf_tensor, %arg59: !tfrt_fallback.tf_tensor,
%arg60: !tfrt_fallback.tf_tensor, %arg61: !tfrt_fallback.tf_tensor, %arg62: !tfrt_fallback.tf_tensor, %arg63: !tfrt_fallback.tf_tensor, %arg64: !tfrt_fallback.tf_tensor,
%arg65: !tfrt_fallback.tf_tensor, %arg66: !tfrt_fallback.tf_tensor, %arg67: !tfrt_fallback.tf_tensor, %arg68: !tfrt_fallback.tf_tensor, %arg69: !tfrt_fallback.tf_tensor,
%arg70: !tfrt_fallback.tf_tensor, %arg71: !tfrt_fallback.tf_tensor, %arg72: !tfrt_fallback.tf_tensor, %arg73: !tfrt_fallback.tf_tensor, %arg74: !tfrt_fallback.tf_tensor,
%arg75: !tfrt_fallback.tf_tensor, %arg76: !tfrt_fallback.tf_tensor, %arg77: !tfrt_fallback.tf_tensor, %arg78: !tfrt_fallback.tf_tensor, %arg79: !tfrt_fallback.tf_tensor,
%arg80: !tfrt_fallback.tf_tensor, %arg81: !tfrt_fallback.tf_tensor, %arg82: !tfrt_fallback.tf_tensor, %arg83: !tfrt_fallback.tf_tensor, %arg84: !tfrt_fallback.tf_tensor,
%arg85: !tfrt_fallback.tf_tensor, %arg86: !tfrt_fallback.tf_tensor, %arg87: !tfrt_fallback.tf_tensor, %arg88: !tfrt_fallback.tf_tensor, %arg89: !tfrt_fallback.tf_tensor,
%arg90: !tfrt_fallback.tf_tensor, %arg91: !tfrt_fallback.tf_tensor, %arg92: !tfrt_fallback.tf_tensor, %arg93: !tfrt_fallback.tf_tensor, %arg94: !tfrt_fallback.tf_tensor,
%arg95: !tfrt_fallback.tf_tensor, %arg96: !tfrt_fallback.tf_tensor, %arg97: !tfrt_fallback.tf_tensor, %arg98: !tfrt_fallback.tf_tensor, %arg99: !tfrt_fallback.tf_tensor,
%arg100: !tfrt_fallback.tf_tensor, %arg101: !tfrt_fallback.tf_tensor, %arg102: !tfrt_fallback.tf_tensor, %arg103: !tfrt_fallback.tf_tensor, %arg104: !tfrt_fallback.tf_tensor,
%arg105: !tfrt_fallback.tf_tensor, %arg106: !tfrt_fallback.tf_tensor, %arg107: !tfrt_fallback.tf_tensor, %arg108: !tfrt_fallback.tf_tensor, %arg109: !tfrt_fallback.tf_tensor,
%arg110: !tfrt_fallback.tf_tensor, %arg111: !tfrt_fallback.tf_tensor, %arg112: !tfrt_fallback.tf_tensor)
-> (!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
%res:112 = tfrt_fallback_async.batch_function device("/device:CPU:0") @f (%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10,%arg11, %arg12, %arg13, %arg14,
%arg15, %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23, %arg24, %arg25, %arg26, %arg27, %arg28, %arg29,
%arg30, %arg31, %arg32, %arg33, %arg34, %arg35, %arg36, %arg37, %arg38, %arg39, %arg40, %arg41, %arg42, %arg43, %arg44,
%arg45, %arg46, %arg47, %arg48, %arg49, %arg50, %arg51, %arg52, %arg53, %arg54, %arg55, %arg56, %arg57, %arg58, %arg59,
%arg60, %arg61, %arg62, %arg63, %arg64, %arg65, %arg66, %arg67, %arg68, %arg69, %arg70, %arg71, %arg72, %arg73, %arg74,
%arg75, %arg76, %arg77, %arg78, %arg79, %arg80, %arg81, %arg82, %arg83, %arg84, %arg85, %arg86, %arg87, %arg88, %arg89,
%arg90, %arg91, %arg92, %arg93, %arg94, %arg95, %arg96, %arg97, %arg98, %arg99, %arg100, %arg101, %arg102, %arg103, %arg104,
%arg105, %arg106, %arg107, %arg108, %arg109, %arg110, %arg111, %arg112) {
num_batch_threads = 16,
max_batch_size = 1,
allowed_batch_sizes = [1],
batch_timeout_micros = 0,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32, i32, i32, i32, i32],
Tcaptured = [i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32],
Tout = [i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32,
i32, i32, i32, i32, i32, i32, i32]} : 112
tfrt.return %arg0, %res#0, %res#1, %res#2, %res#3, %res#4, %res#5, %res#6, %res#7, %res#8, %res#9, %res#10, %res#11, %res#12, %res#13, %res#14,
%res#15, %res#16, %res#17, %res#18, %res#19, %res#20, %res#21, %res#22, %res#23, %res#24, %res#25, %res#26, %res#27, %res#28, %res#29,
%res#30, %res#31, %res#32, %res#33, %res#34, %res#35, %res#36, %res#37, %res#38, %res#39, %res#40, %res#41, %res#42, %res#43, %res#44,
%res#45, %res#46, %res#47, %res#48, %res#49, %res#50, %res#51, %res#52, %res#53, %res#54, %res#55, %res#56, %res#57, %res#58, %res#59,
%res#60, %res#61, %res#62, %res#63, %res#64, %res#65, %res#66, %res#67, %res#68, %res#69, %res#70, %res#71, %res#72, %res#73, %res#74,
%res#75, %res#76, %res#77, %res#78, %res#79, %res#80, %res#81, %res#82, %res#83, %res#84, %res#85, %res#86, %res#87, %res#88, %res#89,
%res#90, %res#91, %res#92, %res#93, %res#94, %res#95, %res#96, %res#97, %res#98, %res#99, %res#100, %res#101, %res#102, %res#103, %res#104,
%res#105, %res#106, %res#107, %res#108, %res#109, %res#110, %res#111
:!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor,
!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/batch_function_fallback_benchmark_test.cc
+157
View File
@@ -0,0 +1,157 @@
/* 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 <string>
#include <utility>
#include <vector>
#include "base/logging.h"
#include "testing/base/public/benchmark.h"
#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_op_handler.h"
#include "tensorflow/core/runtime_fallback/util/fallback_test_util.h"
#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
#include "tfrt/bef/bef_buffer.h" // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h" // from @tf_runtime
#include "tfrt/core_runtime/core_runtime.h" // from @tf_runtime
#include "tfrt/host_context/async_value.h" // from @tf_runtime
#include "tfrt/host_context/async_value_ref.h" // from @tf_runtime
#include "tfrt/host_context/chain.h" // from @tf_runtime
#include "tfrt/host_context/concurrent_work_queue.h" // from @tf_runtime
#include "tfrt/host_context/execution_context.h" // from @tf_runtime
#include "tfrt/host_context/function.h" // from @tf_runtime
#include "tfrt/host_context/host_allocator.h" // from @tf_runtime
#include "tfrt/host_context/host_context.h" // from @tf_runtime
#include "tfrt/host_context/resource_context.h" // from @tf_runtime
#include "tfrt/support/forward_decls.h" // from @tf_runtime
#include "tfrt/support/rc_array.h" // from @tf_runtime
#include "tfrt/tensor/dense_host_tensor.h" // from @tf_runtime
#include "tfrt/tensor/tensor.h" // from @tf_runtime
namespace tensorflow {
namespace {
// Creates a BEF file with a program that runs
// tfrt_fallback_async.batch_function with a empty function forwarding inputs or
// outputs.
std::pair<tfrt::BefBuffer, tfrt::RCReference<tfrt::BEFFile>> CreateBefFile(
tfrt::HostContext* host) {
std::string file_path = GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/"
"batch_function_fallback.mlir.bef");
std::string data;
CHECK_OK(ReadFileToString(Env::Default(), file_path, &data));
tfrt::BefBuffer bef_buffer(data.begin(), data.end());
auto bef_file = tfrt::BEFFile::Open(bef_buffer, host->GetKernelRegistry(),
host->diag_handler(), host->allocator());
CHECK(bef_file);
return std::make_pair(std::move(bef_buffer), std::move(bef_file));
}
std::unique_ptr<tfrt::CoreRuntime> CreateTestCoreRuntime() {
auto corert = tfrt::CoreRuntime::Create(
/*diag_handler=*/[](const tfrt::DecodedDiagnostic&
diag) { LOG(ERROR) << diag.message(); },
tfrt::CreateMallocAllocator(),
tfrt::CreateMultiThreadedWorkQueue(16, 16));
CHECK(corert);
auto fallback_op_handler = tensorflow::tfd::CreateKernelFallbackOpHandler(
corert->get(), corert->get()->GetHostContext()->GetHostDeviceRef());
CHECK(fallback_op_handler);
corert.get()->RegisterOpHandler("tfkernel", fallback_op_handler.get());
return std::move(corert.get());
}
tfrt::RCArray<tfrt::AsyncValue> CreateTestArguments(const tfrt::Function* func,
tfrt::HostContext* host) {
Tensor tensor(DataType::DT_INT32, TensorShape({1}));
std::vector<tfrt::RCReference<tfrt::AsyncValue>> arguments;
arguments.reserve(func->argument_types().size());
arguments.push_back(tfrt::GetReadyChain());
for (int i = 1, e = func->argument_types().size(); i < e; ++i) {
arguments.push_back(
tfrt::MakeAvailableAsyncValueRef<tfrt_stub::FallbackTensor>(tensor));
}
return tfrt::RCArray<tfrt::AsyncValue>(arguments);
}
TEST(BatchFunctionTest, Basic) {
auto corert = CreateTestCoreRuntime();
auto* host = corert->GetHostContext();
auto [bef_buffer, bef_file] = CreateBefFile(host);
auto* func = bef_file->GetFunction("main");
CHECK(func);
CHECK_EQ(func->result_types().size(), 113);
CHECK_EQ(func->argument_types().size(), 113);
auto arguments = CreateTestArguments(func, host);
tfrt::ResourceContext resource_ctx;
auto exec_ctx = tfd::CreateFallbackTestExecutionContext(host, &resource_ctx);
std::vector<tfrt::RCReference<tfrt::AsyncValue>> results;
results.resize(func->result_types().size());
std::vector<tfrt::RCReference<tfrt::AsyncValue>> result_tensors;
result_tensors.resize(func->result_types().size() - 1);
func->Execute(exec_ctx, arguments.values(), results);
host->Await(results);
for (auto& result : results) {
EXPECT_FALSE(result->IsError());
}
}
// Runs a BEF function that batches a function that does nothing just to measure
// the runtime overhead. The BEF function signature is adapted from a real model
// and is useful for benchmarking ops with large attributes and many
// input/output.
void BM_BatchFunctionFallbackWithLargeAttributesAndManyInputsOutputs(
benchmark::State& state) {
auto corert = CreateTestCoreRuntime();
auto* host = corert->GetHostContext();
auto [bef_buffer, bef_file] = CreateBefFile(host);
auto* func = bef_file->GetFunction("main");
CHECK(func);
CHECK_EQ(func->result_types().size(), 113);
CHECK_EQ(func->argument_types().size(), 113);
auto arguments = CreateTestArguments(func, host);
tfrt::ResourceContext resource_ctx;
auto exec_ctx = tfd::CreateFallbackTestExecutionContext(host, &resource_ctx);
std::vector<tfrt::RCReference<tfrt::AsyncValue>> results;
results.resize(func->result_types().size());
for (auto _ : state) {
func->Execute(exec_ctx, arguments.values(), results);
host->Await(results);
results.clear();
results.resize(func->result_types().size());
}
}
BENCHMARK(BM_BatchFunctionFallbackWithLargeAttributesAndManyInputsOutputs);
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/batching_fallback.mlir
+224
View File
@@ -0,0 +1,224 @@
// 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: tfrt_fallback_translate --mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd:1,1 | FileCheck %s
// RUN: tfrt_fallback_translate --mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd:8 | FileCheck %s
func.func @matmul_cpu(%ch: !tfrt.chain, %a: !tfrt_fallback.tf_tensor, %b: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
// Enqueue a sleep onto blocking work queue, %ch0 is fulfilled when sleeping is done.
%us = tfrt.constant.i32 1000
%ch0 = "tfrt_test.blocking.usleep"(%us) : (i32) -> (!tfrt.chain)
%ch1 = tfrt.merge.chains %ch, %ch0 : !tfrt.chain, !tfrt.chain
%ch2 = tfrt_fallback_async.createop(%ch1) key(0) device("/CPU:0") "tf.MatMul"() {T = i32} num_args(2)
%ch3, %result = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("/CPU:0") "tf.MatMul"(%a, %b) {T = i32} : 1
%s = "tfrt_test.get_string"() { value = "Running @matmul_cpu" } : () -> !tfrt.string
%ch4 = "tfrt_test.print_string"(%s, %ch3) : (!tfrt.string, !tfrt.chain) -> (!tfrt.chain)
%ch5 = "tfrt_fallback_async.print_tensor"(%result, %ch4) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch5, %result : !tfrt.chain, !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: --- Running 'batch_function_fallback_concat_test'
func.func @batch_function_fallback_concat_test() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a1 = tfrt_fallback_async.const_dense_tensor dense<[[2, 2], [2, 2]]> : tensor<2x2xi32>
%a2 = tfrt_fallback_async.const_dense_tensor dense<[[3, 3], [3, 3]]> : tensor<2x2xi32>
%b = tfrt_fallback_async.const_dense_tensor dense<[[1, 1], [1, 1]]> : tensor<2x2xi32>
// Two batch_size=2 batches get concatenated.
%result_1 = tfrt_fallback_async.batch_function device("/device:CPU:0") @matmul_cpu (%a1, %b) {
num_batch_threads = 1,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 1000000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32],
Tcaptured = [i32],
Tout = [i32]} : 1
%result_2 = tfrt_fallback_async.batch_function device("/device:CPU:0") @matmul_cpu (%a2, %b) {
num_batch_threads = 1,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 1000000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32],
Tcaptured = [i32],
Tout = [i32]} : 1
// Since batch function kernel scheduling is async, the above 2 batches can arrive in any order.
// CHECK: Running @matmul_cpu
// CHECK-NEXT: Tensor<type: int32 shape: [4,2] values: [[value_output:.*]]>
// CHECK: Tensor<type: int32 shape: [2,2] values: [4 4]
%ch1 = "tfrt_fallback_async.print_tensor"(%result_1, %ch0) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: int32 shape: [2,2] values: [6 6]
%ch2 = "tfrt_fallback_async.print_tensor"(%result_2, %ch1) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch2 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'batch_function_fallback_timeout_test'
func.func @batch_function_fallback_timeout_test() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_dense_tensor dense<[[4, 4], [4, 4]]> : tensor<2x2xi32>
%b = tfrt_fallback_async.const_dense_tensor dense<[[1, 1], [1, 1]]> : tensor<2x2xi32>
// One batch_size=2 batches get padded and processed after timeout.
%result = tfrt_fallback_async.batch_function device("/device:CPU:0") @matmul_cpu (%a, %b) {
num_batch_threads = 1,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 1000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32],
Tcaptured = [i32],
Tout = [i32]} : 1
// CHECK: Running @matmul_cpu
// CHECK-NEXT: Tensor<type: int32 shape: [4,2] values: [[value_output:.*]]>
// CHECK: Tensor<type: int32 shape: [2,2] values: [8 8]
%ch1 = "tfrt_fallback_async.print_tensor"(%result, %ch0) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'batch_function_fallback_no_padding_test'
func.func @batch_function_fallback_no_padding_test() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_dense_tensor dense<[[4, 4], [4, 4]]> : tensor<2x2xi32>
%b = tfrt_fallback_async.const_dense_tensor dense<[[1, 1], [1, 1]]> : tensor<2x2xi32>
// One batch_size=2 batches get processed after timeout.
%result = tfrt_fallback_async.batch_function device("/device:CPU:0") @matmul_cpu (%a, %b) {
num_batch_threads = 1,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 1000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
disable_padding = true,
Tin = [i32],
Tcaptured = [i32],
Tout = [i32]} : 1
// CHECK: Running @matmul_cpu
// As no padding is appended, the tensor shape printed inside the batch function
// is [2,2] instead of [4,2]
// CHECK-NEXT: Tensor<type: int32 shape: [2,2] values: [[value_output:.*]]>
// CHECK: Tensor<type: int32 shape: [2,2] values: [8 8]
%ch1 = "tfrt_fallback_async.print_tensor"(%result, %ch0) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// This test is for testing support for Adaptive Batch Scheduler, and it is
// triggered by num_batch_threads<=0
// CHECK-LABEL: --- Running 'batch_function_fallback_zero_batch_thread_test'
func.func @batch_function_fallback_zero_batch_thread_test() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_dense_tensor dense<[[2, 2], [2, 2]]> : tensor<2x2xi32>
%b = tfrt_fallback_async.const_dense_tensor dense<[[2, 2], [2, 2]]> : tensor<2x2xi32>
// One batch_size=2 batches get padded and processed after timeout.
%result = tfrt_fallback_async.batch_function device("/device:CPU:0") @matmul_cpu (%a, %b) {
num_batch_threads = 0,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 1000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32],
Tcaptured = [i32],
Tout = [i32]} : 1
// CHECK: Running @matmul_cpu
// CHECK-NEXT: Tensor<type: int32 shape: [4,2] values: [[value_output:.*]]>
// CHECK: Tensor<type: int32 shape: [2,2] values: [8 8]
%ch1 = "tfrt_fallback_async.print_tensor"(%result, %ch0) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// Test to verify the fix for b/190394141:
// Given a function that returns multiple values referenced to the same value,
func.func @returns_multiple_refs(%ch: !tfrt.chain, %a: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
tfrt.return %ch, %a, %a : !tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
}
// ... when the function returns_multiple_refs is invoked by the batch_function,
// ... then the code should not crash.
// CHECK-LABEL: Running 'test_batch_returns_multiple_refs'
func.func @test_batch_returns_multiple_refs() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%0 = tfrt_fallback_async.const_dense_tensor dense<[[1, 1], [1, 1]]> : tensor<2x2xi32>
%1, %2 = tfrt_fallback_async.batch_function device("/device:CPU:0") @returns_multiple_refs (%0) {
num_batch_threads = 1,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 10000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32],
Tcaptured = [],
Tout = [i32, i32]} : 2
%3, %4 = tfrt_fallback_async.batch_function device("/device:CPU:0") @returns_multiple_refs (%0) {
num_batch_threads = 1,
max_batch_size = 4,
allowed_batch_sizes = [4],
batch_timeout_micros = 10000,
container = "container",
shared_name = "shared_name",
batching_queue = "batching_queue",
enable_large_batch_splitting = false,
Tin = [i32],
Tcaptured = [],
Tout = [i32, i32]} : 2
%ch1 = "tfrt_fallback_async.print_tensor"(%1, %ch0) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch2 = "tfrt_fallback_async.print_tensor"(%3, %ch1) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch2 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/compile.benchmark.large.mlir
+62
View File
@@ -0,0 +1,62 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s \
// RUN: | tf_bef_executor --test_init_function=register_op_handlers_kernel \
// RUN: --host_allocator_type=malloc \
// RUN: --work_queue_type=mstd:72 \
// RUN: | FileCheck %s --dump-input=always
// A set of benchmarks for large tensor inputs. These benchmarks measure
// codegen/runtime efficiency for executing parallel/concurrent code. Fallback
// kernels rely on Eigen for parallelizing compute operation.
module @kernels attributes { tfrt.compiled } {
func.func @log1p(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = "tf.Log1p"(%arg0): (tensor<?x?xf32>) -> tensor<?x?xf32>
func.return %0 : tensor<?x?xf32>
}
}
// CHECK: --- Running 'BM_fallback_log1p_f32'
func.func @BM_fallback_log1p_f32() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(1) device("/CPU:0")
"tf.Const"() {
dtype = f32, value = dense<1.0> : tensor<1024x1024xf32>
} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(2) device("/CPU:0")
"tf.Log1p"() { T = f32 } num_args(1)
%ch3, %const = tfrt_fallback_async.executeop.seq(%ch2) key(1) cost(100)
device("/CPU:0")
"tf.Const"() {
dtype = f32, value = dense<1.0> : tensor<1024x1024xf32>
} : 1
%ch4 = tfrt_test.benchmark "BM_fallback_log1p_f32"(
%const : !tfrt_fallback.tf_tensor,
%ch3 : !tfrt.chain
)
duration_secs = 1, max_count = 10000, num_warmup_runs = 10
{
%result = tfrt_fallback_async.executeop key(2) cost(100) device("/CPU:0")
"tf.Log1p"(%const) { T = f32 } : 1
tfrt.return %result : !tfrt_fallback.tf_tensor
}
tfrt.return %ch4 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/compile.benchmark.small.mlir
+105
View File
@@ -0,0 +1,105 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s \
// RUN: | tf_bef_executor --test_init_function=register_op_handlers_kernel \
// RUN: | FileCheck %s --dump-input=always
// A set of benchmarks measuring runtime/compiler overheads for running
// kernels with small inputs. These benchmarks do not launch any concurrent
// tasks and do not use concurrent work queue or intra-op thread pool.
module @kernels attributes { tfrt.compiled } {
func.func @rsqrt(%arg0: tensor<?xf32>) -> tensor<?xf32> {
%0 = "tf.Rsqrt"(%arg0): (tensor<?xf32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
func.func @rsqrt_tanh(%arg0: tensor<?xf32>) -> tensor<?xf32> {
%0 = "tf.Rsqrt"(%arg0): (tensor<?xf32>) -> tensor<?xf32>
%1 = "tf.Tanh"(%0): (tensor<?xf32>) -> tensor<?xf32>
func.return %1 : tensor<?xf32>
}
}
// CHECK: --- Running 'BM_fallback_rsqrt_f32'
func.func @BM_fallback_rsqrt_f32() {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(1) device("/CPU:0")
"tf.Const"() {
dtype = f32, value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(2) device("/CPU:0")
"tf.Rsqrt"() { T = f32 } num_args(1)
%ch3, %const = tfrt_fallback_async.executeop.seq(%ch1) key(1) cost(100)
device("/CPU:0")
"tf.Const"() {
dtype = f32, value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
} : 1
tfrt_test.benchmark "BM_fallback_rsqrt_f32"(
%const : !tfrt_fallback.tf_tensor,
%ch2 : !tfrt.chain,
%ch3 : !tfrt.chain
)
duration_secs = 3, max_count = 100000, num_warmup_runs = 10
{
%result = tfrt_fallback_async.executeop key(2) cost(100) device("/CPU:0")
"tf.Rsqrt"(%const) { T = f32 } : 1
tfrt.return %result : !tfrt_fallback.tf_tensor
}
tfrt.return
}
// CHECK: --- Running 'BM_fallback_rsqrt_tanh_f32'
func.func @BM_fallback_rsqrt_tanh_f32() {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(3) device("/CPU:0")
"tf.Const"() {
dtype = f32, value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(4) device("/CPU:0")
"tf.Rsqrt"() { T = f32 } num_args(1)
%ch3 = tfrt_fallback_async.createop(%ch2) key(5) device("/CPU:0")
"tf.Tanh"() { T = f32 } num_args(1)
%ch4, %const = tfrt_fallback_async.executeop.seq(%ch3) key(3) cost(100)
device("/CPU:0")
"tf.Const"() {
dtype = f32, value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
} : 1
tfrt_test.benchmark "BM_fallback_rsqrt_tanh_f32"(
%const : !tfrt_fallback.tf_tensor,
%ch4 : !tfrt.chain
)
duration_secs = 3, max_count = 100000, num_warmup_runs = 10
{
%result0 = tfrt_fallback_async.executeop key(4) cost(100) device("/CPU:0")
"tf.Rsqrt"(%const) { T = f32 } : 1
%result1 = tfrt_fallback_async.executeop key(5) cost(100) device("/CPU:0")
"tf.Tanh"(%result0) { T = f32 } : 1
tfrt.return %result1 : !tfrt_fallback.tf_tensor
}
tfrt.return
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/convert_tensorhandle_to_fallback_tensor.mlir
+41
View File
@@ -0,0 +1,41 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd 2>&1 | FileCheck %s
// This test is to verify the fix for the bug in converting tensorhandle to
// fallback tensor. The TensorHandle should not be moved to the output,
// otherwise, the input will be set to null and cannot be referenced again.
// This bug happens if the AsyncTensor is unavailable during the conversion.
func.func @test_convert_tensorhandle_to_fallback_tensor() {
%us = tfrt.constant.i32 1000
// Given an available TensorHandle with an unavailable AsyncTensor, which will
// become available after 1000 ms.
// CHECK: Created TensorHandle (test string)
%1 = "tfrt_fallback_test.create_tensorhandle_with_delayed_async_tensor"(%us) : (i32) -> !corert.tensorhandle
%2 = "tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor"(%1) {_tfrt_cost = 1 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (!corert.tensorhandle) -> !tfrt_fallback.tf_tensor
// ... when `tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor` is
// invoked twice with the same input,
// ... then the test should not crash.
%3 = "tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor"(%1) {_tfrt_cost = 1 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:1"} : (!corert.tensorhandle) -> !tfrt_fallback.tf_tensor
// CHECK: Slept for 1000 microseconds
// CHECK: Marked AsyncTensor available.
tfrt.return
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/create_op.mlir
+38
View File
@@ -0,0 +1,38 @@
// 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 @init(%in_ch: !tfrt.chain) -> !tfrt.chain {
%ch0 = "tfrt_fallback_async.createop"(%in_ch)
{
device = "cpu",
num_args = 2 : i64,
op_attrs = [["T", i32]],
op_func_attrs = [],
op_key = 0 : i64,
op_name = "tf.AddV2"
} : (!tfrt.chain) -> (!tfrt.chain)
%ch1 = "tfrt_fallback_async.createop"(%ch0)
{
device = "cpu",
num_args = 1 : i64,
op_attrs = [["Targuments", []],
["output_shapes", [#corert.shape<>]],
["output_types", [i64]]],
op_func_attrs = [["f", "dummy_fn"]],
op_key = 1 : i64,
op_name = "tf.FlatMapDataset"
} : (!tfrt.chain) -> (!tfrt.chain)
tfrt.return %ch1 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/custom_thread_pool.mlir
+34
View File
@@ -0,0 +1,34 @@
// 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 @init(%in_ch: !tfrt.chain) -> !tfrt.chain {
%ch0 = "tfrt_fallback_async.createop"(%in_ch)
{
device = "cpu",
num_args = 2 : i64,
op_attrs = [["T", i32]],
op_func_attrs = [],
op_key = 0 : i64,
op_name = "tf.AddV2"
} : (!tfrt.chain) -> (!tfrt.chain)
tfrt.return %ch0 : !tfrt.chain
}
func.func @run(%in_ch: !tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor){
%x_th= corert.const_dense_tensor dense<1> : tensor<i32>
%x = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %x_th {_tfrt_cost = 1 : i64, device = "cpu"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
%res = tfrt_fallback_async.executeop key(0) cost(100) device("cpu") "tf.AddV2"(%x, %x) { T = i32 } : 1
tfrt.return %in_ch, %res : !tfrt.chain, !tfrt_fallback.tf_tensor
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/fallback.mlir
+468
View File
@@ -0,0 +1,468 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor --test_init_function=register_op_handlers_kernel | FileCheck %s
func.func @register_op_handlers_kernel() {
%fallback = "corert.create_kernel_fallback_op_handler"() : () -> !corert.ophandler
%cpu = "corert.create_cpu_op_handler"(%fallback) : (!corert.ophandler) -> !corert.ophandler
%cpu_ordinal = tfrt.constant.i32 1
%cpu1 = "corert.create_cpu_op_handler_with_ordinal"(%fallback, %cpu_ordinal) : (!corert.ophandler, i32) -> !corert.ophandler
corert.register_op_handler %cpu "cpu"
corert.register_op_handler %cpu1 "cpu"
corert.register_op_handler %fallback "tfkernel"
tfrt.return
}
// CHECK-LABEL: --- Running 'test_tf_random_uniform'
func.func @test_tf_random_uniform() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("CPU:0") "tf.Const"()
{ dtype = i64, value = dense<[2,2]> : tensor<2xi64> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("CPU:0") "tf.RandomUniform"() {dtype = f32, T = i64} num_args(1)
// Test tf.RandomUniform.
%ch, %shape_th = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("CPU:0") "tf.Const"()
{ dtype = i64, value = dense<[2,2]> : tensor<2xi64> } : 1
%random_uniform = tfrt_fallback_async.executeop key(1) cost(100) device("CPU:0") "tf.RandomUniform"(%shape_th) {dtype = f32, T = i64} : 1
// CHECK: Tensor<type: float shape: [2,2] values:
%ch3 = "tfrt_fallback_async.print_tensor"(%random_uniform, %ch2)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch3 : !tfrt.chain
}
func.func @test_tf_random_uniform_on_cpu1() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("CPU:1") "tf.Const"()
{ dtype = i64, value = dense<[2,2]> : tensor<2xi64> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("CPU:1") "tf.RandomUniform"() {dtype = f32, T = i64} num_args(1)
// Test tf.RandomUniform.
%ch, %shape_th = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("CPU:1") "tf.Const"()
{ dtype = i64, value = dense<[2,2]> : tensor<2xi64> } : 1
%random_uniform = tfrt_fallback_async.executeop key(1) cost(100) device("CPU:1") "tf.RandomUniform"(%shape_th) {dtype = f32, T = i64} : 1
%res_th = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %random_uniform {_tfrt_cost = 1 : i64, device = "CPU:1"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
// CHECK: float shape: [2,2] values:
%ch3 = "corert.print_tensorhandle"(%res_th, %ch) : (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch3 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_parse_single_example'
func.func @test_parse_single_example() -> !tfrt.chain {
%ch0 = tfrt.new.chain
// %serialized is a tensorflow::Example proto whose content is
// features {
// feature {
// key: "key_0"
// value {
// int64_list {
// value: 100
// }
// }
// }
// feature {
// key: "key_1"
// value {
// int64_list {
// value: 200
// }
// }
// }
// }
//
// Note that in MLIR, non-printable bytes are printed as escaped hex numbers.
%serialized_th = corert.const_string_tensor {shape = [], value = ["\0A!\0A\0E\0A\05key_0\12\05\1A\03\0A\01d\0A\0F\0A\05key_1\12\06\1A\04\0A\02\C8\01"]}
%serialized = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %serialized_th {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
%ch1 = tfrt_fallback_async.createop(%ch0) key(2) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<1xi64>} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(3) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<1> : tensor<1xi64>} num_args(0)
%ch3 = tfrt_fallback_async.createop(%ch2) key(4) device("CPU:0") "tf.ParseSingleExample"()
{Tdense = [i64, i64], dense_keys = ["key_0", "key_1"], num_sparse = 0 : i64, sparse_types = [], sparse_keys = [], dense_shapes = [#corert.shape<?>, #corert.shape<?>]} num_args(3)
%ch4, %zero = tfrt_fallback_async.executeop.seq(%ch3) key(2) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<1xi64>} : 1
%ch5, %one = tfrt_fallback_async.executeop.seq(%ch3) key(3) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<1> : tensor<1xi64>} : 1
%dense_values:2 = tfrt_fallback_async.executeop key(4) cost(100) device("CPU:0") "tf.ParseSingleExample"(%serialized, %zero, %one)
{Tdense = [i64, i64], dense_keys = ["key_0", "key_1"], num_sparse = 0 : i64, sparse_types = [], sparse_keys = [], dense_shapes = [#corert.shape<?>, #corert.shape<?>]} : 2
// CHECK: Tensor<type: int64 shape: [1] values: 100>
%ch6 = "tfrt_fallback_async.print_tensor"(%dense_values#0, %ch5)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: int64 shape: [1] values: 200>
%ch7 = "tfrt_fallback_async.print_tensor"(%dense_values#1, %ch6)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch7 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_const_tensor_proto'
func.func @test_const_tensor_proto() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_tensor_proto "\08\0C\12\00\22\01@"
// CHECK: Tensor<type: quint8 shape: [] values: 64>
%ch1 = "tfrt_fallback_async.print_tensor"(%a, %ch0)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_const_dense_tensor'
func.func @test_const_dense_tensor() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_dense_tensor dense<[true, false]> : tensor<2xi1> {_tfrt_cost = 1 : i64}
// CHECK: Tensor<type: bool shape: [2] values: 1 0>
%ch1 = "tfrt_fallback_async.print_tensor"(%a, %ch0)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_const_string_tensor'
func.func @test_const_string_tensor() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_string_tensor {shape = [2], value = ["const", "string"], _tfrt_cost = 1 : i64}
// CHECK: Tensor<type: string shape: [2] values: const string>
%ch1 = "tfrt_fallback_async.print_tensor"(%a, %ch0)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_const_string_tensor_same_value'
func.func @test_const_string_tensor_same_value() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a = tfrt_fallback_async.const_string_tensor {shape = [2], value = ["string"], _tfrt_cost = 1 : i64}
// CHECK: Tensor<type: string shape: [2] values: string string>
%ch1 = "tfrt_fallback_async.print_tensor"(%a, %ch0)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_predicate'
func.func @test_predicate() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
%a = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
%ra = tfrt_fallback_async.predicate %a {_tfrt_cost = 1 : i64}
// CHECK: false
%ch1 = "tfrt_test.print_bool"(%ch0, %ra) : (!tfrt.chain, i1) -> !tfrt.chain
%b = tfrt_fallback_async.const_dense_tensor dense<1> : tensor<i32> {_tfrt_cost = 1 : i64}
%rb = tfrt_fallback_async.predicate %b {_tfrt_cost = 1 : i64}
// CHECK: true
%ch2 = "tfrt_test.print_bool"(%ch1, %rb) : (!tfrt.chain, i1) -> !tfrt.chain
%c = tfrt_fallback_async.const_string_tensor {shape = [], value = [""], _tfrt_cost = 1 : i64}
%rc = tfrt_fallback_async.predicate %c {_tfrt_cost = 1 : i64}
// CHECK: false
%ch3 = "tfrt_test.print_bool"(%ch2, %rc) : (!tfrt.chain, i1) -> !tfrt.chain
%d = tfrt_fallback_async.const_string_tensor {shape = [], value = ["string"], _tfrt_cost = 1 : i64}
%rd = tfrt_fallback_async.predicate %d {_tfrt_cost = 1 : i64}
// CHECK: true
%ch4 = "tfrt_test.print_bool"(%ch3, %rd) : (!tfrt.chain, i1) -> !tfrt.chain
%e = tfrt_fallback_async.const_dense_tensor dense<[]> : tensor<0xi32> {_tfrt_cost = 1 : i64}
%re = tfrt_fallback_async.predicate %e {_tfrt_cost = 1 : i64}
// CHECK: false
%ch5 = "tfrt_test.print_bool"(%ch4, %re) : (!tfrt.chain, i1) -> !tfrt.chain
%f = tfrt_fallback_async.const_dense_tensor dense<[0]> : tensor<1xi32> {_tfrt_cost = 1 : i64}
%rf = tfrt_fallback_async.predicate %f {_tfrt_cost = 1 : i64}
// CHECK: true
%ch6 = "tfrt_test.print_bool"(%ch5, %rf) : (!tfrt.chain, i1) -> !tfrt.chain
tfrt.return %ch6 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_fallback_resource'
func.func @test_fallback_resource() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a_th = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
%b_th = corert.const_dense_tensor dense<[false, true]> : tensor<2xi1>
%ra, %rb = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %a_th, %b_th {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%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, %b, %a = tfrt_fallback_async.get_resource %ch2 {_tfrt_cost = 1 : i64, device = "cpu", indices = [1 : i64, 0 : i64]} : (!tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
// CHECK: Tensor<type: bool shape: [2] values: 1 0>
%ch4 = "tfrt_fallback_async.print_tensor"(%a, %ch3)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: bool shape: [2] values: 0 1>
%ch5 = "tfrt_fallback_async.print_tensor"(%b, %ch4)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch5 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_i1_dtype'
func.func @test_i1_dtype() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%a_th = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
%b_th = corert.const_dense_tensor dense<[false, true]> : tensor<2xi1>
%a, %b = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %a_th, %b_th {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%ch1 = tfrt_fallback_async.createop(%ch0) key(5) device("CPU:0") "tf.LogicalAnd"() num_args(2)
%ch2, %c = tfrt_fallback_async.executeop.seq(%ch1) key(5) cost(100) device("CPU:0") "tf.LogicalAnd"(%a, %b) : 1
// CHECK: Tensor<type: bool shape: [2] values: 0 0>
%ch3 = "tfrt_fallback_async.print_tensor"(%c, %ch2)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch3 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_step_container'
func.func @test_step_container() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%size_th = corert.const_dense_tensor dense<5> : tensor<i32>
%index_th = corert.const_dense_tensor dense<1> : tensor<i32>
%value_th = corert.const_dense_tensor dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
%size, %index, %value = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %size_th, %index_th, %value_th {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%ch1 = tfrt_fallback_async.createop(%ch0) key(6) device("CPU:0") "tf.TensorArrayV3"() {dtype = f32, element_shape = #corert.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} num_args(1)
%ch2 = tfrt_fallback_async.createop(%ch1) key(7) device("CPU:0") "tf.TensorArrayWriteV3"() {T = f32} num_args(4)
%ch3 = tfrt_fallback_async.createop(%ch2) key(8) device("CPU:0") "tf.TensorArrayReadV3"() {dtype = f32} num_args(3)
%ch4 = tfrt_fallback_async.createop(%ch3) key(9) device("CPU:0") "tf.TensorArrayCloseV3"() num_args(1)
%ch5, %ta:2 = tfrt_fallback_async.executeop.seq(%ch4) key(6) cost(100) device("CPU:0") "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #corert.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : 2
%ch6, %write = tfrt_fallback_async.executeop.seq(%ch5) key(7) cost(100) device("CPU:0") "tf.TensorArrayWriteV3"(%ta#0, %index, %value, %ta#1) {T = f32} : 1
%ch7, %read = tfrt_fallback_async.executeop.seq(%ch6) key(8) cost(100) device("CPU:0") "tf.TensorArrayReadV3"(%ta#0, %index, %write) {dtype = f32} : 1
%ch8 = tfrt_fallback_async.executeop.seq(%ch7) key(9) cost(100) device("CPU:0") "tf.TensorArrayCloseV3"(%ta#0)
// CHECK: Tensor<type: float shape: [3] values: 1 2 3>
%ch9 = "tfrt_fallback_async.print_tensor"(%read, %ch8)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch9 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_select'
func.func @test_select() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%arg_th = corert.const_dense_tensor dense<[1, 2, 3, 4]> : tensor<4xi64>
%cond_th = corert.const_dense_tensor dense<[true, false, true, false]> : tensor<4xi1>
%arg, %cond = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg_th, %cond_th {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%ch1 = tfrt_fallback_async.createop(%ch0) key(10) device("CPU:0") "tf.ZerosLike"() {T = i64} num_args(1)
%ch2 = tfrt_fallback_async.createop(%ch1) key(11) device("CPU:0") "tf.Select"() {T = i64} num_args(3)
%ch3, %zeros = tfrt_fallback_async.executeop.seq(%ch2) key(10) cost(100) device("CPU:0") "tf.ZerosLike"(%arg) {T = i64} : 1
%0 = tfrt_fallback_async.executeop key(11) cost(100) device("CPU:0") "tf.Select"(%cond, %arg, %zeros) {T = i64} : 1
// CHECK: Tensor<type: int64 shape: [4] values: 1 0 3...>
%ch4 = "tfrt_fallback_async.print_tensor"(%0, %ch3)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
func.func @while_cond_false(%in: !tfrt.chain, %x: !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle) {
%ch0 = tfrt_fallback_async.createop(%in) key(12) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<i64>} num_args(0)
%ch1, %zero = tfrt_fallback_async.executeop.seq(%ch0) key(12) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<i64>} : 1
%zero_th = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %zero {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
tfrt.return %ch1, %zero_th : !tfrt.chain, !corert.tensorhandle
}
func.func @while_body_print(%in: !tfrt.chain, %x: !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle) {
%ch0 = tfrt.new.chain
%ch1 = "corert.print_tensorhandle"(%x, %ch0) : (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %in, %x : !tfrt.chain, !corert.tensorhandle
}
// While loop test
// CHECK-LABEL: --- Running 'while_fallback_condition'
func.func @while_fallback_condition() {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
%one = corert.executeop(%cpu)
"tfrt_test.create_dense_tensor"() { shape = [1, 1], values = [1 : i32] } : 1
// CHECK-NOT: DenseHostTensor dtype = i32, shape = [1, 1], values = [1]
%ch1, %result = corert.while @while_cond_false @while_body_print (%ch0, %one) : (!corert.tensorhandle) -> (!corert.tensorhandle)
tfrt.return
}
// CHECK-LABEL: --- Running 'test_copy_if_small'
func.func @test_copy_if_small() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%small = tfrt_fallback_async.const_dense_tensor dense<[1, 2, 3, 4]> : tensor<4xi64> {_tfrt_cost = 1 : i64}
%small_copies:2 = tfrt_fallback_async.copy_if_small %small {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
// CHECK: Tensor<type: int64 shape: [4] values: 1 2 3...>
// CHECK: Tensor<type: int64 shape: [4] values: 1 2 3...>
%ch1 = "tfrt_fallback_async.print_tensor"(%small_copies#0, %ch0) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch2 = "tfrt_fallback_async.print_tensor"(%small_copies#1, %ch1) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%large = tfrt_fallback_async.const_dense_tensor dense<100> : tensor<128xi64> {_tfrt_cost = 1 : i64}
%large_copies:2 = tfrt_fallback_async.copy_if_small %large {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
// CHECK: Tensor<type: int64 shape: [128] values: 100 100 100...>
// CHECK: Tensor<type: int64 shape: [128] values: 100 100 100...>
%ch3 = "tfrt_fallback_async.print_tensor"(%large_copies#0, %ch2) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch4 = "tfrt_fallback_async.print_tensor"(%large_copies#1, %ch3) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%string = tfrt_fallback_async.const_string_tensor {shape = [1], value = ["string"], _tfrt_cost = 1 : i64}
%string_copies:2 = tfrt_fallback_async.copy_if_small %string {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
// CHECK: Tensor<type: string shape: [1] values: string>
// CHECK: Tensor<type: string shape: [1] values: string>
%ch5 = "tfrt_fallback_async.print_tensor"(%string_copies#0, %ch4) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch6 = "tfrt_fallback_async.print_tensor"(%string_copies#1, %ch5) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch4 : !tfrt.chain
}
func.func @branch0(%ch0: !tfrt.chain, %arg0: !corert.tensorhandle, %arg1: !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle) {
%cpu = corert.get_op_handler %ch0 "cpu"
%res = corert.executeop(%cpu) "tfrt_test.add"(%arg0, %arg1) : 1
tfrt.return %ch0, %res : !tfrt.chain, !corert.tensorhandle
}
func.func @branch1(%ch0: !tfrt.chain, %arg0: !corert.tensorhandle, %arg1: !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle) {
%cpu = corert.get_op_handler %ch0 "cpu"
%th = corert.executeop(%cpu)
"tfrt_test.create_dense_tensor"() { shape = [1], values = [4 : i32] } : 1
%add0 = corert.executeop(%cpu) "tfrt_test.add"(%arg0, %arg1) : 1
%res = corert.executeop(%cpu) "tfrt_test.add"(%add0, %th) : 1
tfrt.return %ch0, %res : !tfrt.chain, !corert.tensorhandle
}
// CHECK-LABEL: --- Running 'test_case_fallback_index'
func.func @test_case_fallback_index() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
%ch1 = tfrt_fallback_async.createop(%ch0) key(13) device("CPU:0") "tf.Const"() {dtype = i32, value = dense<0> : tensor<i32>} num_args(0)
%ch2, %fallback_idx0 = tfrt_fallback_async.executeop.seq(%ch1) key(13) cost(100) device("CPU:0") "tf.Const"() {dtype = i32, value = dense<0> : tensor<i32>} : 1
%fallback_idx_th0 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %fallback_idx0 {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
%idx0 = corert.tensorhandle_to_int32 %fallback_idx_th0
%ch5 = tfrt_fallback_async.createop(%ch0) key(14) device("CPU:0") "tf.Const"() {dtype = i32, value = dense<1> : tensor<i32>} num_args(0)
%ch6, %fallback_idx1 = tfrt_fallback_async.executeop.seq(%ch5) key(14) cost(100) device("CPU:0") "tf.Const"() {dtype = i32, value = dense<1> : tensor<i32>} : 1
%fallback_idx_th1 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %fallback_idx1 {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
%idx1 = corert.tensorhandle_to_int32 %fallback_idx_th1
%arg0 = corert.executeop(%cpu)
"tfrt_test.create_dense_tensor"() { shape = [1], values = [2 : i32] } : 1
%arg1 = corert.executeop(%cpu)
"tfrt_test.create_dense_tensor"() { shape = [1], values = [4 : i32] } : 1
%ch3, %res0 = tfrt.case %idx0 [@branch0, @branch1] (%ch0, %arg0, %arg1) : (!tfrt.chain, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle)
// CHECK: DenseHostTensor dtype = i32, shape = [1], values = [6]
%ch4 = "corert.print_tensorhandle"(%res0, %ch3) : (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
%ch7, %res1 = tfrt.case %idx1 [@branch0, @branch1] (%ch4, %arg0, %arg1) : (!tfrt.chain, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle)
// CHECK: DenseHostTensor dtype = i32, shape = [1], values = [10]
%ch8 = "corert.print_tensorhandle"(%res1, %ch7) : (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch8 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_fallback_to_host'
func.func @test_fallback_to_host() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(2) device("/CPU:0") "tf.Const" () { dtype = f32, value = dense<1.0> : tensor<2x2xf32> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(3) device("/CPU:0") "tf.VarHandleOp" () { dtype = f32, shape = #corert.shape<2x2>, container = "c", shared_name = "v0"} num_args(0)
%ch3 = tfrt_fallback_async.createop(%ch2) key(4) device("/CPU:0") "tf.AssignVariableOp" () { dtype = f32 } num_args(2)
%ch4 = tfrt_fallback_async.createop(%ch3) key(5) device("/TPU:0") "tf.ReadVariableOp" () { dtype = f32 } num_args(1)
%ch5, %val = tfrt_fallback_async.executeop.seq(%ch4) key(2) cost(100) device("/CPU:0") "tf.Const" () { dtype = f32, value = dense<1.0> : tensor<2x2xf32> } : 1
%ch6, %varh = tfrt_fallback_async.executeop.seq(%ch4) key(3) cost(100) device("/CPU:0") "tf.VarHandleOp" () { dtype = f32, shape = #corert.shape<2x2>, container = "c", shared_name = "v0"} : 1
%ch7 = tfrt_fallback_async.executeop.seq(%ch6) key(4) cost(100) device("/CPU:0") "tf.AssignVariableOp" (%varh, %val) { dtype = f32 } : 0
%ch8, %val2 = tfrt_fallback_async.executeop.seq(%ch7) key(5) cost(100) device("/TPU:0") "tf.ReadVariableOp" (%varh) { dtype = f32 } : 1
// CHECK: Tensor<type: float shape: [2,2]
%ch9 = "tfrt_fallback_async.print_tensor"(%val2, %ch8)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch9 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_custom_allocator'
func.func @test_custom_allocator() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%allocator = "tfrt_fallback_async.get_test_allocator"() : () -> (!tfrt_fallback.tf_allocator)
%ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("/CPU:0") "tf.Cast"() { SrcT = i32, DstT = f32 } num_args(1)
%ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("/CPU:0") "tf.Cast"() { SrcT = f32, DstT = i32 } num_args(1)
%ch3 = tfrt_fallback_async.createop(%ch2) key(2) device("/CPU:0") "tf.Const"() { dtype = i32, value = dense<123> : tensor<i32> } num_args(0)
%ch4, %const = tfrt_fallback_async.executeop.seq(%ch3) key(2) cost(100) device("/CPU:0") "tf.Const"() { dtype = i32, value = dense<123> : tensor<i32> } : 1
// CHECK: Using TestAllocator
%val0 = tfrt_fallback_async.executeop.allocator(%allocator) key(0) cost(100) device("/CPU:0") "tf.Cast"(%const) { SrcT = i32, DstT = f32 } : 1
// CHECK: Using TestAllocator
%ch5, %val1 = tfrt_fallback_async.executeop.seq.allocator(%ch4, %allocator) key(1) cost(100) device("/CPU:0") "tf.Cast"(%val0) { SrcT = f32, DstT = i32 } : 1
// CHECK: Tensor<type: float shape: [] values: 123>
%ch6 = "tfrt_fallback_async.print_tensor"(%val0, %ch5)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: int32 shape: [] values: 123>
%ch7 = "tfrt_fallback_async.print_tensor"(%val1, %ch6)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch7 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_custom_allocator_async_opkernel'
func.func @test_custom_allocator_async_opkernel() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%allocator = "tfrt_fallback_async.get_test_allocator"() : () -> (!tfrt_fallback.tf_allocator)
%ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("/CPU:0") "tf.Const"() { dtype = i32, value = dense<123> : tensor<i32> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("/CPU:0") "tf.TestAsyncIdentity"() { T = i32 } num_args(1)
%ch3, %const = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("/CPU:0") "tf.Const"() { dtype = i32, value = dense<123> : tensor<i32> } : 1
%val = tfrt_fallback_async.executeop.allocator(%allocator) key(1) cost(100) device("/CPU:0") "tf.TestAsyncIdentity"(%const) { SrcT = i32, DstT = f32 } : 1
// CHECK: Tensor<type: int32 shape: [] values: 123>
%ch4 = "tfrt_fallback_async.print_tensor"(%val, %ch3)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch4 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/fallback_tensor_conversion_host.mlir
+168
View File
@@ -0,0 +1,168 @@
// 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: tfrt_translate -mlir-to-bef %s | tf_bef_executor --test_init_function=register_op_handlers | FileCheck %s --dump-input=fail
func.func @register_op_handlers() {
%fallback = "corert.create_runtime_fallback_op_handler"() {tf_device_name="/device:CPU:0"} : () -> !corert.ophandler
%cpu = "corert.create_cpu_op_handler"(%fallback) : (!corert.ophandler) -> !corert.ophandler
corert.register_op_handler %fallback "tf"
corert.register_op_handler %cpu "cpu"
tfrt.return
}
// CHECK: --- Running 'transfer_to_host'
func.func @transfer_to_host() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%rtfb_handler = "corert.create_runtime_fallback_op_handler"() {tf_device_name="/device:CPU:0"} : () -> !corert.ophandler
%cpu_device = "tfrt.get_device"(%ch0) {device_name="CPU:0"} : (!tfrt.chain) -> !tfrt.device
%th0 = corert.executeop(%rtfb_handler) "tf.Const"()
{ dtype = f32, value = dense<1.0> : tensor<1xf32> } : 1
%th1 = corert.executeop(%rtfb_handler) "tf.Exp" (%th0) { T = f32 } : 1
// TFRuntimeFallbackT->DHT
%th2_tensor_type = corert.get_dst_tensor_type %th1, %cpu_device
%th2 = corert.transfer %th1, %cpu_device, %th2_tensor_type
// CHECK: DenseHostTensor dtype = f32, shape = [1], values = [2.718282e+00]
%ch1 = "corert.print_tensorhandle"(%th2, %ch0)
: (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK: --- Running 'scalar_to_runtime_fallback'
func.func @scalar_to_runtime_fallback() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
%cpu_device = "tfrt.get_device"(%ch0) {device_name="CPU:0"} : (!tfrt.chain) -> !tfrt.device
%th1 = corert.executeop(%cpu) "tfrt_test.create_from_scalar"()
{shape = [2: i64, 2: i64], value = 1: i32} : 1
// ScalarHT->TFRuntimeFallbackT
%th2_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "RuntimeFallback" } : () -> !tfrt.tensor_type
%th2 = corert.transfer %th1, %cpu_device, %th2_tensor_type
// TFRuntimeFallbackT->DHT
%th3_tensor_type = corert.get_dst_tensor_type %th2, %cpu_device
%th3 = corert.transfer %th2, %cpu_device, %th3_tensor_type
// CHECK: DenseHostTensor dtype = i32, shape = [2, 2], values = [1, 1, 1, 1]
%ch1 = "corert.print_tensorhandle"(%th3, %ch0)
: (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}
// CHECK: --- Running 'dht_to_runtime_fallback'
func.func @dht_to_runtime_fallback() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
%cpu_device = "tfrt.get_device"(%ch0) {device_name="CPU:0"} : (!tfrt.chain) -> !tfrt.device
%th1 = corert.executeop(%cpu) "tfrt_test.create_dense_tensor"()
{ shape = [2: i64, 2: i64], values = [42 : i32] } : 1
// DHT->TFRuntimeFallbackT
%th2_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "RuntimeFallback" } : () -> !tfrt.tensor_type
%th2 = corert.transfer %th1, %cpu_device, %th2_tensor_type
// TFRuntimeFallbackT->DHT
%th3_tensor_type = corert.get_dst_tensor_type %th2, %cpu_device
%th3 = corert.transfer %th2, %cpu_device, %th3_tensor_type
// CHECK: DenseHostTensor dtype = i32, shape = [2, 2], values = [42, 42, 42, 42]
%ch1 = "corert.print_tensorhandle"(%th3, %ch0)
: (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
// DHT->TFKernelFallbackT
%th4_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "KernelFallback" } : () -> !tfrt.tensor_type
%th4 = corert.transfer %th1, %cpu_device, %th4_tensor_type
// TFKernelFallbackT->TFRuntimeFallbackT
%th5_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "RuntimeFallback" } : () -> !tfrt.tensor_type
%th5 = corert.transfer %th4, %cpu_device, %th5_tensor_type
// TFRuntimeFallbackT->TFKernelFallbackT
%th6_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "KernelFallback" } : () -> !tfrt.tensor_type
%th6 = corert.transfer %th5, %cpu_device, %th6_tensor_type
// TFKernelFallbackT->DHT
%th7_tensor_type = corert.get_dst_tensor_type %th6, %cpu_device
%th7 = corert.transfer %th6, %cpu_device, %th7_tensor_type
// CHECK: DenseHostTensor dtype = i32, shape = [2, 2], values = [42, 42, 42, 42]
%ch2 = "corert.print_tensorhandle"(%th7, %ch1)
: (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch2 : !tfrt.chain
}
// CHECK: --- Running 'sht_to_fallback'
func.func @sht_to_fallback() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
%cpu_device = "tfrt.get_device"(%ch0) {device_name="CPU:0"} : (!tfrt.chain) -> !tfrt.device
%th1 = corert.const_string_tensor {shape = [2], value = ["uiuc", "berkeley"]}
// SHT->TFRuntimeFallbackT
%th2_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "RuntimeFallback" } : () -> !tfrt.tensor_type
%th2 = corert.transfer %th1, %cpu_device, %th2_tensor_type
// TFRuntimeFallbackT->SHT
%th3_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "StringHost" } : () -> !tfrt.tensor_type
%th3 = corert.transfer %th2, %cpu_device, %th3_tensor_type
// CHECK: StringHostTensor shape = [2], values = ["uiuc", "berkeley"]
%ch1 = "corert.print_tensorhandle"(%th3, %ch0)
: (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
// SHT->TFKernelFallbackT
%th4_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "KernelFallback" } : () -> !tfrt.tensor_type
%th4 = corert.transfer %th1, %cpu_device, %th4_tensor_type
// TFKernelFallbackT->TFRuntimeFallbackT
%th5_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "RuntimeFallback" } : () -> !tfrt.tensor_type
%th5 = corert.transfer %th4, %cpu_device, %th5_tensor_type
// TFRuntimeFallbackT->TFKernelFallbackT
%th6_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "KernelFallback" } : () -> !tfrt.tensor_type
%th6 = corert.transfer %th5, %cpu_device, %th6_tensor_type
// TFKernelFallbackT->SHT
%th7_tensor_type = "tfrt_test.get_static_tensor_type"()
{ tensor_type = "StringHost" } : () -> !tfrt.tensor_type
%th7 = corert.transfer %th6, %cpu_device, %th7_tensor_type
// CHECK: StringHostTensor shape = [2], values = ["uiuc", "berkeley"]
%ch2 = "corert.print_tensorhandle"(%th7, %ch1)
: (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch2 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/kernel_fallback_compat_test.cc
+150
View File
@@ -0,0 +1,150 @@
/* 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 <string>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
#include "tensorflow/core/runtime_fallback/util/fallback_test_util.h"
#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
#include "tensorflow/core/tfrt/runtime/runtime.h"
#include "tensorflow/core/tfrt/utils/thread_pool.h"
#include "tfrt/bef/bef_buffer.h" // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h" // from @tf_runtime
#include "tfrt/core_runtime/core_runtime.h" // from @tf_runtime
#include "tfrt/host_context/async_value.h" // from @tf_runtime
#include "tfrt/host_context/chain.h" // from @tf_runtime
#include "tfrt/host_context/function.h" // from @tf_runtime
#include "tfrt/host_context/host_context.h" // from @tf_runtime
#include "tfrt/host_context/resource_context.h" // from @tf_runtime
#include "tfrt/support/ref_count.h" // from @tf_runtime
#include "tfrt/tracing/tracing.h" // from @tf_runtime
namespace tensorflow {
namespace {
// Creates a BEF file with a program that runs tfrt_fallback.batch_function with
// a empty function forwarding inputs or outputs.
//
// TODO(b/175648326): Move the function below to the common test utilities for
// BEF.
std::pair<tfrt::BefBuffer, tfrt::RCReference<tfrt::BEFFile>> CreateBefFile(
absl::string_view file_name, tfrt::HostContext* host) {
std::string file_path = GetDataDependencyFilepath(absl::StrCat(
"tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/", file_name));
std::string data;
CHECK_OK(ReadFileToString(Env::Default(), file_path, &data));
tfrt::BefBuffer bef_buffer(data.begin(), data.end());
auto bef_file = tfrt::BEFFile::Open(bef_buffer, host->GetKernelRegistry(),
host->diag_handler(), host->allocator());
CHECK(bef_file);
return std::make_pair(std::move(bef_buffer), std::move(bef_file));
}
TEST(KernelFallbackCompatTest, CreateOp) {
tfrt::tracing::SetTracingLevel(tfrt::tracing::TracingLevel::Debug);
auto runtime =
tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/4);
auto* host = runtime->core_runtime()->GetHostContext();
auto pair = CreateBefFile("create_op.mlir.bef", host);
auto& bef_file = pair.second;
tfrt::ResourceContext resource_ctx;
auto exec_ctx = tfd::CreateFallbackTestExecutionContext(host, &resource_ctx);
auto chain = tfrt::GetReadyChain();
auto* func = bef_file->GetFunction("init");
ASSERT_TRUE(func != nullptr);
std::vector<tfrt::RCReference<tfrt::AsyncValue>> results;
results.resize(1);
func->Execute(exec_ctx, {chain.GetAsyncValue()}, results);
host->Await(results);
ASSERT_FALSE(results[0]->IsError()) << results[0]->GetError().message();
auto* fallback_request_state =
exec_ctx.request_ctx()
->GetDataIfExists<tfd::KernelFallbackCompatRequestState>();
ASSERT_TRUE(fallback_request_state != nullptr);
auto* rendezvous = fallback_request_state->rendezvous();
ASSERT_TRUE(rendezvous != nullptr);
auto* runner_table = fallback_request_state->runner_table();
// TODO(tfrt-devs): Create a special key type for OpKernelRunnerCache instead
// of using tfrt::Location. The key should be generated by higher-level
// applications such as compiler, and it is higher-level applications'
// responsibility to make sure the key is unique in an OpKernelRunnerCache
// instance.
auto* kernel_runner_add = runner_table->Get(0);
ASSERT_TRUE(kernel_runner_add);
auto* kernel_runner_flatmap = runner_table->Get(1);
ASSERT_TRUE(kernel_runner_flatmap);
}
TEST(KernelFallbackCompatTest, CustomThreadPool) {
auto runtime =
tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/4);
auto* host = runtime->core_runtime()->GetHostContext();
auto pair = CreateBefFile("custom_thread_pool.mlir.bef", host);
auto& bef_file = pair.second;
tfrt::ResourceContext resource_ctx;
tensorflow::tfrt_stub::TfThreadPool thread_pool(/*name=*/"test",
/*num_threads=*/1);
auto exec_ctx = tfd::CreateFallbackTestExecutionContext(host, &resource_ctx,
&thread_pool);
auto chain = tfrt::GetReadyChain();
auto* init_func = bef_file->GetFunction("init");
ASSERT_TRUE(init_func != nullptr);
std::vector<tfrt::RCReference<tfrt::AsyncValue>> results;
results.resize(1);
init_func->Execute(exec_ctx, {chain.GetAsyncValue()}, results);
host->Await(results);
ASSERT_FALSE(results[0]->IsError()) << results[0]->GetError().message();
auto* run_func = bef_file->GetFunction("run");
ASSERT_TRUE(run_func != nullptr);
results.clear();
results.resize(2);
run_func->Execute(exec_ctx, {chain.GetAsyncValue()}, results);
host->Await(results);
ASSERT_FALSE(results[0]->IsError()) << results[0]->GetError().message();
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/kernel_fallback_op_handler.mlir
+45
View File
@@ -0,0 +1,45 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor | FileCheck %s
// CHECK-LABEL: --- Not running 'register_op_handler' because it has arguments.
func.func @register_op_handler(%ch0: !tfrt.chain) -> !tfrt.chain {
%op_handler = "corert.create_kernel_fallback_op_handler"() : () -> !corert.ophandler
%ch = corert.register_op_handler %op_handler "tfkernel0"
tfrt.return %ch : !tfrt.chain
}
// CHECK-LABEL: --- Not running 'get_op_handler' because it has arguments.
func.func @get_op_handler(%ch0: !tfrt.chain) -> !tfrt.chain {
%tfkernel0 = corert.get_op_handler %ch0 "tfkernel0"
tfrt.return %ch0 : !tfrt.chain
}
// CHECK-LABEL: --- Not running 'failed_get_op_handler' because it has arguments.
func.func @failed_get_op_handler(%ch0: !tfrt.chain) -> !tfrt.chain {
// expected-error @+1 {{runtime error: op_handler not found}}
%tf0 = corert.get_op_handler %ch0 "tfkernel0"
tfrt.return %ch0 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_op_handler_kernels'
func.func @test_op_handler_kernels() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt.call @failed_get_op_handler(%ch0) : (!tfrt.chain) -> !tfrt.chain
%ch2 = tfrt.call @register_op_handler(%ch1) : (!tfrt.chain) -> !tfrt.chain
%ch3 = tfrt.call @get_op_handler(%ch2) : (!tfrt.chain) -> !tfrt.chain
tfrt.return %ch3 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/mnist.mlir
+210
View File
@@ -0,0 +1,210 @@
// 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: tfrt_translate -mlir-to-bef %s | tf_bef_executor | FileCheck %s
// Initialize all the variables used in the model by creating variable handles
// and assigning the variables with the correct values. This function takes an
// input chain and returns an output chain and a list of variable handles that
// have been value-initialized. The returned tfd.tf_tensors are TF tensors that
// contain tensorflow::ResourceHandle.
func.func @mnist_init_variables(%start_c : !tfrt.chain) -> (
!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor, !tfd.tf_tensor, !tfd.tf_tensor) {
// Init w1.
%init_w1_c, %w1_h = "tfd.delegate_kernel"(%start_c) {
_name = "VarHandleOp", attr0_name = "container", attr0_value = "string$",
attr1_name = "shared_name", attr1_value = "string$w1",
attr2_name = "dtype", attr2_value = "tfdtype$DT_INT32",
attr3_name = "shape", attr3_value = "tfshape$[2,2]"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%const_w1_c, %w1_const_tensor = "tfd.delegate_kernel"(%init_w1_c) {
_name = "Const", attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32",
attr1_name = "value",
attr1_value = "tftensor$dtype:DT_INT32 tensor_shape { dim { size: 2 } dim { size: 2 }} int_val: 1 int_val: 1 int_val: 1 int_val: 1"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%assignv_w1_c = "tfd.delegate_kernel"(%const_w1_c, %w1_h, %w1_const_tensor) {
_name = "AssignVariableOp",
attr0_name = "dtype",attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain)
// Init b1.
%init_b1_c, %b1_h = "tfd.delegate_kernel"(%assignv_w1_c) {
_name = "VarHandleOp", attr0_name = "container", attr0_value = "string$",
attr1_name = "shared_name", attr1_value = "string$b1",
attr2_name = "dtype", attr2_value = "tfdtype$DT_INT32",
attr3_name = "shape", attr3_value = "tfshape$[2,2]"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%const_b1_c, %b1_const_tensor = "tfd.delegate_kernel"(%init_b1_c) {
_name = "Const", attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32",
attr1_name = "value",
attr1_value = "tftensor$dtype:DT_INT32 tensor_shape { dim { size: 2 } dim { size: 2 }} int_val: 1 int_val: 1 int_val: 1 int_val: 1"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%assignv_b1_c = "tfd.delegate_kernel"(%const_b1_c, %b1_h, %b1_const_tensor) {
_name = "AssignVariableOp",
attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain)
// Init w2.
%init_w2_c, %w2_h = "tfd.delegate_kernel"(%assignv_b1_c) {
_name = "VarHandleOp", attr0_name = "container", attr0_value = "string$",
attr1_name = "shared_name", attr1_value = "string$w2",
attr2_name = "dtype", attr2_value = "tfdtype$DT_INT32",
attr3_name = "shape", attr3_value = "tfshape$[2,2]"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%const_w2_c, %w2_const_tensor = "tfd.delegate_kernel"(%init_w2_c) {
_name = "Const", attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32",
attr1_name = "value",
attr1_value = "tftensor$dtype:DT_INT32 tensor_shape { dim { size: 2 } dim { size: 2 }} int_val: 1 int_val: 1 int_val: 1 int_val: 1"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%assignv_w2_c = "tfd.delegate_kernel"(%const_w2_c, %w2_h, %w2_const_tensor) {
_name = "AssignVariableOp",
attr0_name = "dtype",attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain)
// Init b2.
%init_b2_c, %b2_h = "tfd.delegate_kernel"(%assignv_w2_c) {
_name = "VarHandleOp", attr0_name = "container", attr0_value = "string$",
attr1_name = "shared_name", attr1_value = "string$b2",
attr2_name = "dtype", attr2_value = "tfdtype$DT_INT32",
attr3_name = "shape", attr3_value = "tfshape$[2,2]"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%const_b2_c, %b2_const_tensor = "tfd.delegate_kernel"(%init_b2_c) {
_name = "Const", attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32",
attr1_name = "value",
attr1_value = "tftensor$dtype:DT_INT32 tensor_shape { dim { size: 2 } dim { size: 2 }} int_val: 1 int_val: 1 int_val: 1 int_val: 1"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
%assignv_b2_c = "tfd.delegate_kernel"(%const_b2_c, %b2_h, %b2_const_tensor) {
_name = "AssignVariableOp",
attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain)
tfrt.return %assignv_b2_c, %w1_h, %b1_h, %w2_h, %b2_h : !tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor, !tfd.tf_tensor, !tfd.tf_tensor
}
// This function represents one forward pass on the model. It takes an input
// chain and the variable handles for all model variables, and returns an output
// chain and the result tensor value.
func.func @inference_call(
%start_c : !tfrt.chain, %inputx_tensor : !tfd.tf_tensor,
%w1_h : !tfd.tf_tensor, %b1_h : !tfd.tf_tensor, %w2_h : !tfd.tf_tensor,
%b2_h : !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor) {
// Read all variables into tensors.
%readv_c1, %w1_tensor = "tfd.delegate_kernel"(%start_c, %w1_h) {
_name = "ReadVariableOp",
attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
%readv_c2, %b1_tensor = "tfd.delegate_kernel"(%readv_c1, %b1_h) {
_name = "ReadVariableOp",
attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
%readv_c3, %w2_tensor = "tfd.delegate_kernel"(%readv_c2, %w2_h) {
_name = "ReadVariableOp",
attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
%readv_c4, %b2_tensor = "tfd.delegate_kernel"(%readv_c3, %b2_h) {
_name = "ReadVariableOp",
attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32"
} : (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
// Compute the model forward pass.
%matmul0_out_c, %matmul0_tensor = "tfd.delegate_kernel"(
%readv_c4, %inputx_tensor, %w1_tensor) {
_name = "MatMul",
attr1_name = "transpose_a", attr1_value = "bool$false",
attr2_name = "transpose_b", attr2_value = "bool$false"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (
!tfrt.chain, !tfd.tf_tensor)
%add0_out_c, %add0_tensor = "tfd.delegate_kernel"(
%matmul0_out_c, %matmul0_tensor, %b1_tensor) {
_name = "AddV2"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (
!tfrt.chain, !tfd.tf_tensor)
%relu0_out_c, %relu0_tensor = "tfd.delegate_kernel"(
%add0_out_c, %add0_tensor) {
_name = "Relu"
} : (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
%matmul1_out_c, %matmul1_tensor = "tfd.delegate_kernel"(
%relu0_out_c, %relu0_tensor, %w2_tensor) {
_name = "MatMul",
attr1_name = "transpose_a", attr1_value = "bool$false",
attr2_name = "transpose_b", attr2_value = "bool$false"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (
!tfrt.chain, !tfd.tf_tensor)
%add1_out_c, %add1_tensor = "tfd.delegate_kernel"(
%matmul1_out_c, %matmul1_tensor, %b2_tensor) {
_name = "AddV2"
} : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (
!tfrt.chain, !tfd.tf_tensor)
%identity_out_c, %identity_tensor = "tfd.delegate_kernel"(
%add1_out_c, %add1_tensor) {
_name = "Identity"
} : (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
tfrt.return %identity_out_c, %identity_tensor : !tfrt.chain, !tfd.tf_tensor
}
// This is the main driver function. It creates the EagerContext, initializes
// all variables and runs one forward pass on the model.
// CHECK: --- Running 'mnist_delegate_test'
func.func @mnist_delegate_test() {
%start_c = tfrt.new.chain
// Init eager context.
%context_init_c = "tfd.init_eager_context"(%start_c)
: (!tfrt.chain) -> !tfrt.chain
// Initialize all variables.
%init_var_c, %w1_h, %b1_h, %w2_h, %b2_h
= tfrt.call @mnist_init_variables(%context_init_c) : (!tfrt.chain) -> (
!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor, !tfd.tf_tensor,
!tfd.tf_tensor)
// Get input data.
%get_input_c, %inputx_tensor = "tfd.delegate_kernel"(%init_var_c) {
_name = "Const", attr0_name = "dtype", attr0_value = "tfdtype$DT_INT32",
attr1_name = "value",
attr1_value = "tftensor$dtype:DT_INT32 tensor_shape { dim { size: 2 } dim { size: 2 }} int_val: 1 int_val: 1 int_val: 1 int_val: 1"
} : (!tfrt.chain) -> (!tfrt.chain, !tfd.tf_tensor)
// Make one inference call.
%inference_call_out_c, %inference_call_out_tensor
= tfrt.call @inference_call(
%get_input_c, %inputx_tensor, %w1_h, %b1_h, %w2_h, %b2_h) : (
!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor, !tfd.tf_tensor,
!tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
// CHECK: shape = [2, 2], values = [7, 7, 7, 7]
%print_c = "tfd.print_tft"(%inference_call_out_tensor, %inference_call_out_c)
: (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/runtime_fallback_op_handler.mlir
+51
View File
@@ -0,0 +1,51 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor | FileCheck %s
// CHECK-LABEL: --- Not running 'register_runtime_fallback_op_handler_chain' because it has arguments.
func.func @register_runtime_fallback_op_handler_chain(%ch0: !tfrt.chain) -> !tfrt.chain {
%runtime_fallback = "corert.create_runtime_fallback_op_handler"() {tf_device_name="/device:CPU:0"} : () -> !corert.ophandler
%ch = corert.register_op_handler %runtime_fallback "tf0"
tfrt.return %ch : !tfrt.chain
}
// CHECK-LABEL: --- Not running 'get_runtime_fallback_op_handler' because it has arguments.
func.func @get_runtime_fallback_op_handler(%ch0: !tfrt.chain) -> !tfrt.chain {
%tf0 = corert.get_op_handler %ch0 "tf0"
%th = corert.executeop(%tf0) "tf.Const"()
{ dtype = f32, value = dense<1.0> : tensor<2x2xf32> } : 1
%ch2 = "corert.print_tensorhandle" (%th, %ch0) : (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch2 : !tfrt.chain
}
// CHECK-LABEL: --- Not running 'failed_runtime_fallback_get_op_handler' because it has arguments.
func.func @failed_runtime_fallback_get_op_handler(%ch0: !tfrt.chain) -> !tfrt.chain {
// expected-error @+1 {{runtime error: op_handler not found}}
%tf0 = corert.get_op_handler %ch0 "tf0"
tfrt.return %ch0 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_runtime_fallback_op_handler_chain_kernels'
func.func @test_runtime_fallback_op_handler_chain_kernels() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt.call @failed_runtime_fallback_get_op_handler(%ch0) : (!tfrt.chain) -> !tfrt.chain
%ch2 = tfrt.call @register_runtime_fallback_op_handler_chain(%ch1) : (!tfrt.chain) -> !tfrt.chain
// CHECK: RuntimeFallbackTensor dtype = float, shape = [2, 2], values = [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00]
%ch3 = tfrt.call @get_runtime_fallback_op_handler(%ch2) : (!tfrt.chain) -> !tfrt.chain
tfrt.return %ch3 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/tf_delegate.mlir
+139
View File
@@ -0,0 +1,139 @@
// 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: tfrt_translate -mlir-to-bef %s | tf_bef_executor | FileCheck %s
// CHECK: --- Running 'matmul_delegate_test'
func.func @matmul_delegate_test() {
%c0 = tfrt.new.chain
// Create 2x2 dht<i32, 2> with value 1
%dht_a = tfrt_dht.create_uninitialized_tensor.i32.2 [2 : i64, 2 : i64]
%c1 = tfrt_dht.fill_tensor_with_constant.i32 %dht_a, %c0 1 : i32
// Create 2x2 dht<i32, 2> with value 2
%dht_b = tfrt_dht.create_uninitialized_tensor.i32.2 [2 : i64, 2 : i64]
%c2 = tfrt_dht.fill_tensor_with_constant.i32 %dht_b, %c0 2 : i32
// Convert dht to tf tensor
%tft_a, %c3 = "tfd.move_dht_to_tft"(%dht_a, %c1)
: (!tfrt_tensor.tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
%tft_b, %c4 = "tfd.move_dht_to_tft"(%dht_b, %c2)
: (!tfrt_tensor.tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
// Print legacy TF tensors
// CHECK: shape = [2, 2], values = [1, 1, 1, 1]
%cc0 = "tfd.print_tft"(%tft_a, %c0) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: shape = [2, 2], values = [2, 2, 2, 2]
%cc1 = "tfd.print_tft"(%tft_b, %cc0) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// Create TF eager context
%c6 = "tfd.init_eager_context"(%c0): (!tfrt.chain) -> !tfrt.chain
// Delegate to tf.matmul
%c7, %tft_x = "tfd.delegate_kernel"(%c6, %tft_a, %tft_b) {op_name = "MatMul"}
: (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
// CHECK: shape = [2, 2], values = [4, 4, 4, 4]
%cc2 = "tfd.print_tft"(%tft_x, %cc1) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// Delegate to tf.matmul by using %tft_x from the previous delegation as input
%c8, %tft_y = "tfd.delegate_kernel"(%c6, %tft_x, %tft_b) {op_name = "MatMul"}
: (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
// Print legacy TF tensors
// CHECK: shape = [2, 2], values = [16, 16, 16, 16]
%cc3 = "tfd.print_tft"(%tft_y, %cc2) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// Convert tf tensor back to dht
%dht_c, %c9 = "tfd.convert_tft_to_dht"(%tft_x, %cc2)
: (!tfd.tf_tensor, !tfrt.chain) -> (!tfrt_tensor.tensor, !tfrt.chain)
%dht_d, %c10 = "tfd.convert_tft_to_dht"(%tft_y, %cc3)
: (!tfd.tf_tensor, !tfrt.chain) -> (!tfrt_tensor.tensor, !tfrt.chain)
// Print the result dht
// CHECK: shape = [2, 2], values = [4, 4, 4, 4]
%cc4 = tfrt_dht.print_tensor %dht_c, %cc3
// CHECK: shape = [2, 2], values = [16, 16, 16, 16]
%cc5 = tfrt_dht.print_tensor %dht_d, %cc4
tfrt.return
}
// CHECK: --- Running 'bad_op_name_test'
func.func @bad_op_name_test() {
%c0 = tfrt.new.chain
// Create 2x2 dht<i32, 2> with value 2
%dht_a = tfrt_dht.create_uninitialized_tensor.i32.2 [2 : i64, 2 : i64]
%c1 = tfrt_dht.fill_tensor_with_constant.i32 %dht_a, %c0 2 : i32
// Convert dht to tf tensor
%tft_a, %c2 = "tfd.move_dht_to_tft"(%dht_a, %c1)
: (!tfrt_tensor.tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
// Create TF eager context
%c3 = "tfd.init_eager_context"(%c0): (!tfrt.chain) -> !tfrt.chain
// expected-error @+1 {{runtime error: 'BadOp'}}
%c4, %tft_x = "tfd.delegate_kernel"(%c3, %tft_a) {op_name = "BadOp"}
: (!tfrt.chain, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
tfrt.return
}
// CHECK: --- Running 'addn_delegate_test'
func.func @addn_delegate_test() {
%c0 = tfrt.new.chain
// Create scalar dht<i32, 0> with value 1
%dht_a = tfrt_dht.create_uninitialized_tensor.i32.0 []
%c1 = tfrt_dht.fill_tensor_with_constant.i32 %dht_a, %c0 1 : i32
// Create scalar dht<i32, 0> with value 2
%dht_b = tfrt_dht.create_uninitialized_tensor.i32.0 []
%c2 = tfrt_dht.fill_tensor_with_constant.i32 %dht_b, %c0 2 : i32
// Convert dht to tf tensor
%tft_a, %c3 = "tfd.move_dht_to_tft"(%dht_a, %c1)
: (!tfrt_tensor.tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
%tft_b, %c4 = "tfd.move_dht_to_tft"(%dht_b, %c2)
: (!tfrt_tensor.tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
// Print legacy TF tensors
// CHECK: shape = [], values = [1]
%cc0 = "tfd.print_tft"(%tft_a, %c0) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: shape = [], values = [2]
%cc1 = "tfd.print_tft"(%tft_b, %cc0) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// Create TF eager context
%c6 = "tfd.init_eager_context"(%c0): (!tfrt.chain) -> !tfrt.chain
// Delegate to AddN
%c7, %tft_x = "tfd.delegate_kernel"(%c6, %tft_a, %tft_b) {_op_name = "AddN", attr0_name = "N", attr0_value = "int$2", attr1_name = "T", attr1_value = "tfdtype$DT_INT32"}
: (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
// CHECK: shape = [], values = [3]
%cc2 = "tfd.print_tft"(%tft_x, %cc1) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
// Delegate to AddN by using %tft_x from the previous delegation as input
%c8, %tft_y = "tfd.delegate_kernel"(%c6, %tft_x, %tft_b) {_op_name = "AddN", attr0_name = "N", attr0_value = "int$2", attr1_name = "T", attr1_value = "tfdtype$DT_INT32"}
: (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (!tfrt.chain, !tfd.tf_tensor)
// Print legacy TF tensors
// CHECK: shape = [], values = [5]
%cc3 = "tfd.print_tft"(%tft_y, %cc2) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/tf_ops.mlir
+195
View File
@@ -0,0 +1,195 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd --test_init_function=register_op_handlers 2>&1 | FileCheck %s
func.func @register_op_handlers() -> (!tfrt.chain, !tfrt.chain, !tfrt.chain) {
%fallback = "corert.create_kernel_fallback_op_handler"() : () -> !corert.ophandler
%cpu = "corert.create_cpu_op_handler"(%fallback) : (!corert.ophandler) -> !corert.ophandler
%ch0 = corert.register_op_handler %fallback "tfkernel"
%cpu_identity = "tfrt_test.identity"(%ch0, %cpu) : (!tfrt.chain, !corert.ophandler) -> !corert.ophandler
%ch1 = corert.register_op_handler %cpu_identity "cpu"
%ch2 = "corert.add_kernel_fallback_implicit_conversions" (%cpu_identity) : (!corert.ophandler) -> !tfrt.chain
tfrt.return %ch0, %ch1, %ch2 : !tfrt.chain, !tfrt.chain, !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_tf_random_uniform'
func.func @test_tf_random_uniform() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("CPU:0") "tf.Const"()
{ dtype = i64, value = dense<[2,2]> : tensor<2xi64> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("CPU:0") "tf.RandomUniform"() {T = i64, dtype = f32} num_args(1)
%ch3, %shape = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("CPU:0") "tf.Const"()
{ dtype = i64, value = dense<[2,2]> : tensor<2xi64> } : 1
%random_uniform = tfrt_fallback_async.executeop key(1) cost(100) device("CPU:0") "tf.RandomUniform"(%shape) {T = i64, dtype = f32} : 1
// CHECK: Tensor<type: float shape: [2,2]
%ch4 = "tfrt_fallback_async.print_tensor"(%random_uniform, %ch3)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch4 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_tf_vars'
func.func @test_tf_vars() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(2) device("CPU:0") "tf.Const" () { dtype = f32, value = dense<1.0> : tensor<2x2xf32> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(3) device("CPU:0") "tf.VarHandleOp" () { dtype = f32, shape = #corert.shape<2x2>, container = "c", shared_name = "v0"} num_args(0)
%ch3 = tfrt_fallback_async.createop(%ch2) key(4) device("CPU:0") "tf.AssignVariableOp" () { dtype = f32 } num_args(2)
%ch4 = tfrt_fallback_async.createop(%ch3) key(5) device("CPU:0") "tf.ReadVariableOp" () { dtype = f32 } num_args(1)
%ch5, %val = tfrt_fallback_async.executeop.seq(%ch4) key(2) cost(100) device("CPU:0") "tf.Const" () { dtype = f32, value = dense<1.0> : tensor<2x2xf32> } : 1
%ch6, %varh = tfrt_fallback_async.executeop.seq(%ch4) key(3) cost(100) device("CPU:0") "tf.VarHandleOp" () { dtype = f32, shape = #corert.shape<2x2>, container = "c", shared_name = "v0"} : 1
%ch7 = tfrt_fallback_async.executeop.seq(%ch6) key(4) cost(100) device("CPU:0") "tf.AssignVariableOp" (%varh, %val) { dtype = f32 } : 0
%ch8, %val2 = tfrt_fallback_async.executeop.seq(%ch7) key(5) cost(100) device("CPU:0") "tf.ReadVariableOp" (%varh) { dtype = f32 } : 1
// CHECK: Tensor<type: float shape: [2,2]
%ch9 = "tfrt_fallback_async.print_tensor"(%val2, %ch8)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch9 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_parse_example_v2'
func.func @test_parse_example_v2() -> !tfrt.chain {
%ch = tfrt.new.chain
%serialized_th = corert.const_string_tensor {shape = [1], value = [""]}
%names_th = corert.const_string_tensor {shape = [0], value = []}
%dense_keys_th = corert.const_string_tensor {shape = [7], value = ["has_login_page_feature", "num_terms_inside_postform", "num_terms_outside_postform", "num_terms_outside_postform_without_bp", "query_params_contains_url", "title_with_login_phase", "url_contains_login_terms"]}
%ch0 = tfrt_fallback_async.createop(%ch) key(6) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<i64>} num_args(0)
%ch1 = tfrt_fallback_async.createop(%ch0) key(7) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<1> : tensor<i64>} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(8) device("CPU:0") "tf.ParseExampleV2"()
{Tdense = [i64, i64, i64, i64, i64, i64, i64], dense_shapes = [#corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>], num_sparse = 0 : i64, ragged_split_types = [], ragged_value_types = [], sparse_types = []}
num_args(12)
%ch3, %zero = tfrt_fallback_async.executeop.seq(%ch2) key(6) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<i64>} : 1
%ch4, %one = tfrt_fallback_async.executeop.seq(%ch2) key(7) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<1> : tensor<i64>} : 1
%serialized, %names, %dense_keys = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %serialized_th, %names_th, %dense_keys_th
{_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%dense_values:7 = tfrt_fallback_async.executeop key(8) cost(100) device("CPU:0")
"tf.ParseExampleV2"(%serialized, %names, %names, %dense_keys, %names, %zero, %one, %zero, %one, %zero, %one, %zero)
{Tdense = [i64, i64, i64, i64, i64, i64, i64], dense_shapes = [#corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>], num_sparse = 0 : i64, ragged_split_types = [], ragged_value_types = [], sparse_types = []} : 7
// CHECK: Tensor<type: int64 shape: [1] values: 0>
// CHECK: Tensor<type: int64 shape: [1] values: 1>
// CHECK: Tensor<type: int64 shape: [1] values: 0>
// CHECK: Tensor<type: int64 shape: [1] values: 1>
// CHECK: Tensor<type: int64 shape: [1] values: 0>
// CHECK: Tensor<type: int64 shape: [1] values: 1>
// CHECK: Tensor<type: int64 shape: [1] values: 0>
%ch5 = "tfrt_fallback_async.print_tensor"(%dense_values#0, %ch4)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch6 = "tfrt_fallback_async.print_tensor"(%dense_values#1, %ch5)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch7 = "tfrt_fallback_async.print_tensor"(%dense_values#2, %ch6)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch8 = "tfrt_fallback_async.print_tensor"(%dense_values#3, %ch7)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch9 = "tfrt_fallback_async.print_tensor"(%dense_values#4, %ch8)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch10 = "tfrt_fallback_async.print_tensor"(%dense_values#5, %ch9)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
%ch11 = "tfrt_fallback_async.print_tensor"(%dense_values#6, %ch10)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch11 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_tile'
func.func @test_tile() -> !tfrt.chain {
%ch = tfrt.new.chain
%input_th = corert.const_string_tensor {shape = [1, 2], value = ["", ""]}
%multiples_th = corert.const_dense_tensor dense<[7,3]> : tensor<2xi32>
%input, %multiples = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %input_th, %multiples_th {_tfrt_cost = 1 : i64, device = "CPU:0"}
: (!corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%ch0 = tfrt_fallback_async.createop(%ch) key(9) device("CPU:0") "tf.Tile"() {T = !corert.string, Tmultiples = i32} num_args(2)
%ch1, %output = tfrt_fallback_async.executeop.seq(%ch0) key(9) cost(100) device("CPU:0") "tf.Tile"(%input, %multiples) {T = !corert.string, Tmultiples = i32} : 1
// CHECK: Tensor<type: string shape: [7,6]
%ch2 = "tfrt_fallback_async.print_tensor"(%output, %ch1)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch2 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'tfrt_test_benchmark_lifetime_regression_test'
func.func @tfrt_test_benchmark_lifetime_regression_test() {
%ch_0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch_0 "cpu"
%tensor_0 = corert.executeop(%cpu) "tfrt_test.create_dense_tensor"() { shape = [1 : i64], values = [1.0 : f32] } : 1
%tensor_1 = corert.executeop(%cpu) "tf.VarHandleOp" (){ dtype = f32, shape = #corert.shape<1>, container = "c", shared_name = "v_1" } : 1
%ch_1 = corert.executeop.seq(%cpu, %ch_0) "tf.AssignVariableOp" (%tensor_1, %tensor_0) { dtype = f32 } : 0
tfrt_test.benchmark "tfrt_test_benchmark_lifetime_regression_test"(
%ch_0 : !tfrt.chain,
%ch_1 : !tfrt.chain,
%tensor_1 : !corert.tensorhandle,
%cpu : !corert.ophandler
)
duration_secs = 1, max_count = 1, num_warmup_runs = 0
{
%ch_3_1, %tensor_2 = corert.executeop.seq(%cpu, %ch_1) "tf.ReadVariableOp" (%tensor_1) { dtype = f32 } : 1
tfrt.return %ch_0 : !tfrt.chain
}
tfrt.return
}
// CHECK-LABEL: --- Running 'test_quantized'
func.func @test_quantized() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%ch1 = tfrt_fallback_async.createop(%ch0) key(2) device("CPU:0") "tf.Const" () { dtype = f32, value = dense<1.0> : tensor<1x1xf32> } num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(3) device("CPU:0") "tf.QuantizeV2" () { T = !corert.quint8} num_args(3)
%ch3 = tfrt_fallback_async.createop(%ch2) key(4) device("CPU:0") "tf.QuantizeV2" () { T = !corert.quint16} num_args(3)
%ch4 = tfrt_fallback_async.createop(%ch3) key(5) device("CPU:0") "tf.QuantizeV2" () { T = !corert.qint8} num_args(3)
%ch5 = tfrt_fallback_async.createop(%ch4) key(6) device("CPU:0") "tf.QuantizeV2" () { T = !corert.qint16} num_args(3)
%ch6 = tfrt_fallback_async.createop(%ch5) key(7) device("CPU:0") "tf.QuantizeV2" () { T = !corert.qint32} num_args(3)
%ch7, %val = tfrt_fallback_async.executeop.seq(%ch6) key(2) cost(100) device("CPU:0") "tf.Const" () { dtype = f32, value = dense<1.0> : tensor<1x1xf32> } : 1
%ch8, %val1:3 = tfrt_fallback_async.executeop.seq(%ch7) key(3) cost(100) device("CPU:0") "tf.QuantizeV2" (%val, %val, %val) { T = !corert.quint8} : 3
%ch9, %val2:3 = tfrt_fallback_async.executeop.seq(%ch8) key(4) cost(100) device("CPU:0") "tf.QuantizeV2" (%val, %val, %val) { T = !corert.quint16} : 3
%ch10, %val3:3 = tfrt_fallback_async.executeop.seq(%ch9) key(5) cost(100) device("CPU:0") "tf.QuantizeV2" (%val, %val, %val) { T = !corert.qint8} : 3
%ch11, %val4:3 = tfrt_fallback_async.executeop.seq(%ch10) key(6) cost(100) device("CPU:0") "tf.QuantizeV2" (%val, %val, %val) { T = !corert.qint16} : 3
%ch12, %val5:3 = tfrt_fallback_async.executeop.seq(%ch11) key(7) cost(100) device("CPU:0") "tf.QuantizeV2" (%val, %val, %val) { T = !corert.qint32} : 3
// CHECK: Tensor<type: quint8 shape: [1,1]
%ch13 = "tfrt_fallback_async.print_tensor"(%val1#0, %ch12)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: quint16 shape: [1,1]
%ch14 = "tfrt_fallback_async.print_tensor"(%val2#0, %ch13)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: qint8 shape: [1,1]
%ch15 = "tfrt_fallback_async.print_tensor"(%val3#0, %ch14)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: qint16 shape: [1,1]
%ch16 = "tfrt_fallback_async.print_tensor"(%val4#0, %ch15)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
// CHECK: Tensor<type: qint32 shape: [1,1]
%ch17 = "tfrt_fallback_async.print_tensor"(%val5#0, %ch16)
: (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch17 : !tfrt.chain
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/tf_ops_error.mlir
+74
View File
@@ -0,0 +1,74 @@
// 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: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd --test_init_function=register_op_handlers 2>&1 | FileCheck %s --dump-input=fail
func.func @register_op_handlers() {
%fallback = "corert.create_kernel_fallback_op_handler"() : () -> !corert.ophandler
corert.register_op_handler %fallback "tfkernel"
tfrt.return
}
// CHECK-LABEL: --- Running 'test_parse_example_v2_error'
func.func @test_parse_example_v2_error() -> !tfrt_fallback.tf_tensor {
%serialized_th = corert.const_string_tensor {shape = [2, 1], value = ["", ""]}
%names_th = corert.const_string_tensor {shape = [0], value = []}
%dense_keys_th = corert.const_string_tensor {shape = [7], value = ["has_login_page_feature", "num_terms_inside_postform", "num_terms_outside_postform", "num_terms_outside_postform_without_bp", "query_params_contains_url", "title_with_login_phase", "url_contains_login_terms"]}
%ch = tfrt.new.chain
%ch0 = tfrt_fallback_async.createop(%ch) key(0) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<i64>} num_args(0)
%ch1 = tfrt_fallback_async.createop(%ch0) key(1) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<1> : tensor<i64>} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(2) device("CPU:0") "tf.ParseExampleV2"()
{Tdense = [i64, i64, i64, i64, i64, i64, i64], dense_shapes = [#corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>], num_sparse = 0 : i64, ragged_split_types = [], ragged_value_types = [], sparse_types = []}
num_args(12)
%ch3, %zero = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<0> : tensor<i64>} : 1
%ch4, %one = tfrt_fallback_async.executeop.seq(%ch2) key(1) cost(100) device("CPU:0") "tf.Const"() {dtype = i64, value = dense<1> : tensor<i64>} : 1
%serialized, %names, %dense_keys = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %serialized_th, %names_th, %dense_keys_th {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
// expected-error @+1 {{Expected serialized to be a scalar or vector, got shape: [2,1]}}
%dense_values:7 = tfrt_fallback_async.executeop key(2) cost(100) device("CPU:0")
"tf.ParseExampleV2"(%serialized, %names, %names, %dense_keys, %names, %zero, %one, %zero, %one, %zero, %one, %zero)
{Tdense = [i64, i64, i64, i64, i64, i64, i64], dense_shapes = [#corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>, #corert.shape<>], num_sparse = 0 : i64, ragged_split_types = [], ragged_value_types = [], sparse_types = []} : 7
tfrt.return %dense_values#0: !tfrt_fallback.tf_tensor
}
// CHECK-LABEL: --- Running 'test_assign_variable_error'
func.func @test_assign_variable_error() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%th0 = corert.const_dense_tensor dense<1> : tensor<i32>
%th1 = corert.const_dense_tensor dense<1.0> : tensor<f32>
%0, %1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %th0, %th1 {_tfrt_cost = 1 : i64, device = "CPU:0"} : (!corert.tensorhandle, !corert.tensorhandle) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
%ch1 = tfrt_fallback_async.createop(%ch0) key(3) device("CPU:0") "tf.VarHandleOp"() {container = "", dtype = f32, shape = #corert.shape<>, shared_name = "x"} num_args(0)
%ch2 = tfrt_fallback_async.createop(%ch1) key(4) device("CPU:0") "tf.AssignVariableOp"() {dtype = f32} num_args(2)
%ch3 = tfrt_fallback_async.createop(%ch2) key(5) device("CPU:0") "tf.AssignVariableOp"() {dtype = i32} num_args(2)
%ch4, %2 = tfrt_fallback_async.executeop.seq(%ch3) key(3) cost(100) device("CPU:0") "tf.VarHandleOp"() {container = "", dtype = f32, shape = #corert.shape<>, shared_name = "x"} : 1
%ch5 = tfrt_fallback_async.executeop.seq(%ch4) key(4) cost(100) device("CPU:0") "tf.AssignVariableOp"(%2, %1) {dtype = f32}
// expected-error @+1 {{Trying to assign variable with wrong dtype. Expected float got int32}}
%ch6 = tfrt_fallback_async.executeop.seq(%ch5) key(5) cost(100) device("CPU:0") "tf.AssignVariableOp"(%2, %0) {dtype = i32}
tfrt.return %ch6 : !tfrt.chain
}
// CHECK-LABEL: --- Running 'test_ref_args'
func.func @test_ref_args() {
%ch0 = tfrt.new.chain
// expected-error @+1 {{Unsupported ref args in VariableV2}}
%ch1 = tfrt_fallback_async.createop(%ch0) key(6) device("CPU:0") "tf.VariableV2"() {container = "", dtype = f32, shape = #corert.shape<>, shared_name = "x"} num_args(0)
%ch2, %0 = tfrt_fallback_async.executeop.seq(%ch1) key(6) cost(100) device("CPU:0") "tf.VariableV2"() {container = "", dtype = f32, shape = #corert.shape<>, shared_name = "x"} : 1
tfrt.return
}
tensorflow/compiler/mlir/tfrt/tests/tfrt_fallback/tfrt_forwarding.mlir
+64
View File
@@ -0,0 +1,64 @@
// 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: tfrt_translate --mlir-to-bef %s | tf_bef_executor | FileCheck %s
// CHECK: --- Running 'scalar_add_tfrt_forwarding_test'
func.func @scalar_add_tfrt_forwarding_test() {
%one = tfrt.constant.i32 1
%two = tfrt.constant.i32 2
%tft_a = "tfd.constant_tensor"(%one) : (i32) -> !tfd.tensor
%tft_b = "tfd.constant_tensor"(%two) : (i32) -> !tfd.tensor
%tft_c = "tfd.forward_kernel"(%tft_a, %tft_b) {op_name = "ScalarAdd"}: (!tfd.tensor, !tfd.tensor) -> !tfd.tensor
// CHECK: tensor=Tensor<type: int32 shape: [] values: 3>
"tfd.print_tensor"(%tft_c) : (!tfd.tensor) -> ()
tfrt.return
}
// CHECK: --- Running 'failing_kernel_tfrt_forwarding_test'
func.func @failing_kernel_tfrt_forwarding_test() -> !tfd.tensor {
%tft = "tfd.forward_kernel"() {op_name = "FailingKernel"}: () -> !tfd.tensor // expected-error {{runtime error: OP_REQUIRES failed at filename:999 : Internal: TFRT forwarding error!}}
tfrt.return %tft : !tfd.tensor
}
// CHECK-NEXT: 'failing_kernel_tfrt_forwarding_test' returned <<error: OP_REQUIRES failed at filename:999 : Internal: TFRT forwarding error!>>
// CHECK: --- Running 'missing_kernel_tfrt_forwarding_test'
func.func @missing_kernel_tfrt_forwarding_test() -> !tfd.tensor {
%tft = "tfd.forward_kernel"() {op_name = "MissingKernel"}: () -> !tfd.tensor // expected-error {{runtime error: Not found: Could not find kernel MissingKernel in the registry.}}
tfrt.return %tft : !tfd.tensor
}
// CHECK-NEXT: 'missing_kernel_tfrt_forwarding_test' returned <<error: Not found: Could not find kernel MissingKernel in the registry.>>
// CHECK: --- Running 'bool_attr_tfrt_forwarding_test'
func.func @bool_attr_tfrt_forwarding_test() {
// Note that op_name is prefixed with underscore because op_name must be the
// first attributes when all attributes are sorted by name (b/140896071).
%tft_t = "tfd.forward_kernel"() {_op_name = "KernelWithBoolAttr", attr1_name = "testattr", attr1_value = "bool$true"}: () -> !tfd.tensor
%tft_f = "tfd.forward_kernel"() {_op_name = "KernelWithBoolAttr", attr1_name = "testattr", attr1_value = "bool$false"}: () -> !tfd.tensor
// CHECK: tensor=Tensor<type: bool shape: [] values: 1>
"tfd.print_tensor"(%tft_t) : (!tfd.tensor) -> ()
// CHECK: tensor=Tensor<type: bool shape: [] values: 0>
"tfd.print_tensor"(%tft_f) : (!tfd.tensor) -> ()
tfrt.return
}

Some files were not shown because too many files have changed in this diff Show More