wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 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/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
}