1// RUN: tf-tfrt-opt -tf-to-tfrt=enable-native-ops=false %s | FileCheck %s --dump-input=fail --dump-input-filter=all 2// RUN: tf-tfrt-opt -pass-pipeline='tf-to-tfrt{enable-native-ops=false target-tpurt=true tpu-use-core-selector=false}' %s | FileCheck %s --dump-input=fail --dump-input-filter=all 3 4// CHECK-LABEL: func @_tfrt_fallback_init 5// CHECK-SAME: {{.*}} !tfrt.chain 6// CHECK: tfrt_fallback_async.createop(%arg0) key(0) device("/device:CPU:0") "tf.ParseExampleV2"() 7// CHECK-SAME: Tdense = [f32, f32], dense_shapes = [#corert.shape<>, #corert.shape<>] 8// CHECK-SAME: num_sparse = 2 : i64, ragged_split_types = [], ragged_value_types = [] 9// CHECK-SAME: sparse_types = [!corert.string, i64]} 10// CHECK-SAME: num_args(7) 11 12// CHECK: tfrt_fallback_async.createop(%0) key(1) device("/device:CPU:0") "tf.ReadVariableOp"() {dtype = f32} num_args(1) 13 14// CHECK: tfrt_fallback_async.createop(%1) key(2) device("/device:CPU:0") "tf.MatMul"() {T = f32, transpose_a = false, transpose_b = false} num_args(2) 15 16// CHECK-LABEL: func @main 17// CHECK-SAME: {{.*}} !tfrt.chain 18// CHECK-SAME: [[serialized:%.*]]: !corert.tensorhandle 19func.func @main(%serialized: tensor<32x!tf_type.string>) -> (tensor<?x2xi64>) attributes {tf.entry_function = {inputs = "input0", outputs = "ParseExample/ParseExampleV2"}} { 20 %dense_default_0 = "tf.Const"() {device = "/device:CPU:0", dtype = f32, value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32> 21 %dense_default_1 = "tf.Const"() {device = "/device:CPU:0", dtype = f32, value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32> 22 %dense_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string> 23 %names = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string> 24 %ragged_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string> 25 %sparse_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string> 26 27 // CHECK: [[fallback_serialized:%.*]] = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[serialized]] 28 // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0" 29 // CHECK: [[outputs:%.*]]:8 = tfrt_fallback_async.executeop key(0) cost({{.*}}) device("/device:CPU:0") "tf.ParseExampleV2" 30 // CHECK-SAME: ([[fallback_serialized]] 31 // CHECK-NOT: device 32 // CHECK-SAME: Tdense = [f32, f32] 33 // CHECK-SAME: dense_shapes = [#corert.shape<>, #corert.shape<>] 34 // CHECK-SAME: num_sparse = 2 : i64 35 // CHECK-SAME: ragged_split_types = [] 36 // CHECK-SAME: ragged_value_types = [] 37 // CHECK-SAME: sparse_types = [!corert.string, i64] 38 %outputs:8 = "tf.ParseExampleV2"(%serialized, %names, %sparse_keys, %dense_keys, %ragged_keys, %dense_default_0, %dense_default_1) 39 { 40 Tdense = [f32, f32], dense_shapes = [#tf_type.shape<>, #tf_type.shape<>], 41 device = "/device:CPU:0", num_sparse = 2 : i64, ragged_split_types = [], ragged_value_types = [], 42 result_segment_sizes = array<i32: 2, 2, 2, 2, 0, 0>, 43 sparse_types = [!tf_type.string, !tf_type.string] 44 } : (tensor<32x!tf_type.string>, tensor<0x!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>, tensor<0x!tf_type.string>, tensor<0xf32>, tensor<0xf32>) 45 -> (tensor<?x2xi64>, tensor<?x2xi64>, tensor<?x!tf_type.string>, tensor<?xi64>, tensor<2xi64>, tensor<2xi64>, tensor<32xf32>, tensor<32xf32>) 46 47 // CHECK: [[result:%.*]] = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle [[outputs]]#0 48 // CHECK-SAME: device = "/device:CPU:0" 49 // CHECK: tfrt.return {{.*}}, [[result]] 50 func.return %outputs#0 : tensor<?x2xi64> 51} 52 53// CHECK-LABEL: func @no_native 54func.func @no_native(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> { 55 // CHECK-NOT: corert.executeop 56 // CHECK: tfrt_fallback_async.executeop.seq({{.*}}) key(1) cost({{.*}}) device("/device:CPU:0") "tf.ReadVariableOp" 57 // CHECK: tfrt_fallback_async.executeop key(2) cost({{.*}}) device("/device:CPU:0") "tf.MatMul" 58 %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32> 59 %1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> 60 func.return %1 : tensor<3x3xf32> 61} 62 63// CHECK-LABEL: func @gpu_device 64func.func @gpu_device(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> { 65 // CHECK-NOT: corert.executeop 66 // CHECK: tfrt_fallback_async.executeop.seq({{.*}}) key({{.*}}) cost({{.*}}) device("/device:GPU:0") "tf.ReadVariableOp" 67 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:GPU:0") "tf.MatMul" 68 %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:GPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32> 69 %1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> 70 func.return %1 : tensor<3x3xf32> 71} 72 73// CHECK-LABEL: func @tpu_device 74func.func @tpu_device(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> { 75 // CHECK-NOT: corert.executeop 76 // CHECK: tfrt_fallback_async.executeop.seq({{.*}}) key({{.*}}) cost({{.*}}) device("/device:TPU:0") "tf.ReadVariableOp" 77 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.MatMul" 78 %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:TPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32> 79 %1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> 80 func.return %1 : tensor<3x3xf32> 81} 82 83// CHECK-LABEL: func @tfrt_set_resource 84// CHECK-SAME: ([[in_ch:%.*]]: !tfrt.chain, 85func.func @tfrt_set_resource(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) { 86 // CHECK: [[ch0:%.*]] = tfrt_fallback_async.set_resource [[in_ch]], {{.*}} {device = "/device:CPU:0", index = 0 : i64} 87 // CHECK: [[ch1:%.*]], [[result:%.*]] = tfrt_fallback_async.executeop.seq([[ch0]]) key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.ReadVariableOp" 88 // CHECK: [[ch2:%.*]] = tfrt_fallback_async.set_resource [[ch1]], [[result]] {device = "/device:CPU:0", index = 1 : i64} 89 90 "tf._TfrtSetResource"(%arg0) {device = "/device:CPU:0", index = 0} : (tensor<3x1xf32>) -> () 91 %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32> 92 "tf._TfrtSetResource"(%0) {device = "/device:CPU:0", index = 1} : (tensor<1x3xf32>) -> () 93 func.return 94} 95 96// CHECK-LABEL: func @tfrt_get_resource 97func.func @tfrt_get_resource() -> tensor<3x3xf32> { 98 // CHECK: [[ready_ch:%.*]] = tfrt.new.chain 99 // CHECK: [[ch3:%.*]], [[results:%.*]]:2 = tfrt_fallback_async.get_resource [[ready_ch]] {device = "/device:CPU:0", indices = [0, 1]} 100 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/device:CPU:0") "tf.MatMul"([[results]]#0, [[results]]#1) 101 %a, %b = "tf._TfrtGetResource"() {device = "/device:CPU:0", indices = [0, 1], shared_name = ["", ""], container = ["", ""]} : () -> (tensor<3x1xf32>, tensor<1x3xf32>) 102 %1 = "tf.MatMul"(%a, %b) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> 103 func.return %1 : tensor<3x3xf32> 104} 105 106// CHECK-LABEL: func @tensor_array 107func.func @tensor_array() -> (tensor<1x1x512xf32>) { 108 %value = "tf.Const"() {device = "/device:CPU:0", value = dense<0.1> : tensor<1x512xf32>} : () -> (tensor<1x512xf32>) 109 %index = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> (tensor<i32>) 110 %size = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> (tensor<i32>) 111 %indices = "tf.Const"() {device = "/device:CPU:0", value = dense<[0]> : tensor<1xi32>} : () -> (tensor<1xi32>) 112 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayV3" 113 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayWriteV3" 114 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayGatherV3" 115 %handle, %flow = "tf.TensorArrayV3"(%size) {clear_after_read = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<?x512>, identical_element_shapes = true, tensor_array_name = "output"} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<1x512xf32>>>, tensor<f32>) 116 %flow_1 = "tf.TensorArrayWriteV3"(%handle, %index, %value, %flow) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<1x512xf32>>>, tensor<i32>, tensor<1x512xf32>, tensor<f32>) -> tensor<f32> 117 %result = "tf.TensorArrayGatherV3"(%handle, %indices, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", element_shape = #tf_type.shape<1x512>} : (tensor<2x!tf_type.resource<tensor<1x512xf32>>>, tensor<1xi32>, tensor<f32>) -> tensor<1x1x512xf32> 118 func.return %result : tensor<1x1x512xf32> 119} 120 121// CHECK-LABEL: func @gpu_device_cost 122func.func @gpu_device_cost(%arg0: tensor<3x1xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<3x3xf32> { 123 // CHECK: tfrt_fallback_async.executeop.seq({{.*}}) key({{.*}}) cost({{1}}) device({{.*}}) "tf.ReadVariableOp" 124 // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{1}}) device({{.*}}) "tf.MatMul" 125 %0 = "tf.ReadVariableOp"(%arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32> 126 %1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/job:localhost/replica:0/task:0/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> 127 func.return %1 : tensor<3x3xf32> 128} 129