// 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: litert-opt %s -canonicalize | env FILECHECK_OPTS="" FileCheck %s // CHECK-LABEL: @elementwise_unary_ops func.func @elementwise_unary_ops() -> (tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor) { %0 = arith.constant dense<-1.0> : tensor %1 = arith.constant dense<1.0> : tensor %2 = arith.constant dense<1.0> : tensor %3 = arith.constant dense<1.0> : tensor %4 = arith.constant dense<4.0> : tensor %5 = arith.constant dense<4.0> : tensor %6 = arith.constant dense<2.0> : tensor %one = arith.constant dense<1.0> : tensor // CHECK-DAG: [[cst0:%.*]] = arith.constant dense<1.000000e+00> : tensor // CHECK-DAG: [[cst1:%.*]] = arith.constant dense<0.841470957> : tensor // CHECK-DAG: [[cst2:%.*]] = arith.constant dense<0.540302277> : tensor // CHECK-DAG: [[cst3:%.*]] = arith.constant dense<0.000000e+00> : tensor // CHECK-DAG: [[cst4:%.*]] = arith.constant dense<2.000000e+00> : tensor // CHECK-DAG: [[cst5:%.*]] = arith.constant dense<5.000000e-01> : tensor // CHECK-DAG: [[cst6:%.*]] = arith.constant dense<4.000000e+00> : tensor // CHECK-DAG: [[cst7:%.*]] = arith.constant dense<0.761594176> : tensor // CHECK-DAG: [[cst8:%.*]] = arith.constant dense<0.73105859{{[78]}}> : tensor // CHECK: return [[cst0]], [[cst1]], [[cst2]], [[cst3]], [[cst4]], [[cst5]], [[cst6]], [[cst7]], [[cst8]] %7 = "tfl.abs"(%0) : (tensor) -> tensor %8 = "tfl.sin"(%1) : (tensor) -> tensor %9 = "tfl.cos"(%2) : (tensor) -> tensor %10 = "tfl.log"(%3) : (tensor) -> tensor %11 = "tfl.sqrt"(%4) : (tensor) -> tensor %12 = "tfl.rsqrt"(%5) : (tensor) -> tensor %13 = "tfl.square"(%6) : (tensor) -> tensor %14 = "tfl.tanh"(%one) : (tensor) -> tensor %15 = "tfl.logistic"(%one) : (tensor) -> tensor func.return %7, %8, %9, %10, %11, %12, %13, %14, %15 : tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor } // CHECK-LABEL: @cumsum_chained_pipeline func.func @cumsum_chained_pipeline() -> tensor<3xf32> { %input = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32> %axis = arith.constant dense<0> : tensor %cst_two = arith.constant dense<2.0> : tensor<3xf32> %one_tensor = arith.constant dense<1.0> : tensor<3xf32> // CHECK: %[[PIPELINE_RES:.*]] = arith.constant dense<[0.982013761, 0.999664664, 0.999999165]> : tensor<3xf32> // CHECK: return %[[PIPELINE_RES]] %0 = "tfl.cumsum"(%input, %axis) {exclusive = false, reverse = false} : (tensor<3xf32>, tensor) -> tensor<3xf32> %add = "tfl.add"(%0, %cst_two) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> %sub = "tfl.sub"(%add, %one_tensor) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> %mul = "tfl.mul"(%sub, %cst_two) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> %res_pipeline = "tfl.logistic"(%mul) : (tensor<3xf32>) -> tensor<3xf32> func.return %res_pipeline : tensor<3xf32> } // CHECK-LABEL: @rank func.func @rank() -> tensor<1xi32> { %cst = arith.constant dense<[[1], [2]]> : tensor<2x1xi32> // CHECK: %[[CST:.*]] = arith.constant dense<2> : tensor<1xi32> // CHECK: return %[[CST]] %0 = "tfl.rank"(%cst) : (tensor<2x1xi32>) -> tensor<1xi32> func.return %0 : tensor<1xi32> } // CHECK-LABEL: @rank_input_known_rank func.func @rank_input_known_rank(%arg0 : tensor<2x1xi32>) -> tensor<1xi32> { // CHECK: %[[CST:.*]] = arith.constant dense<2> : tensor<1xi32> // CHECK: return %[[CST]] %0 = "tfl.rank"(%arg0) : (tensor<2x1xi32>) -> tensor<1xi32> func.return %0 : tensor<1xi32> } // CHECK-LABEL: @reshape_dynamic_output func.func @reshape_dynamic_output() -> tensor { %input = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32> %shape = arith.constant dense<[4]> : tensor<1xi32> // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 2, 3, 4]> : tensor<4xi32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tfl.reshape"(%input, %shape) : (tensor<2x2xi32>, tensor<1xi32>) -> tensor func.return %0 : tensor } // CHECK-LABEL: @pseudo_const func.func @pseudo_const() -> tensor { // CHECK: %[[CST:.*]] = arith.constant dense<1> : tensor // CHECK: return %[[CST]] %0 = "tfl.pseudo_const"() {value = dense<1> : tensor} : () -> tensor func.return %0 : tensor } // CHECK-LABEL: @range_int func.func @range_int() -> tensor { %cst = arith.constant dense<0> : tensor %cst_1 = arith.constant dense<4> : tensor %cst_2 = arith.constant dense<1> : tensor // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[0, 1, 2, 3]> : tensor<4xi32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor } // CHECK-LABEL: @range_float func.func @range_float() -> tensor { %cst = arith.constant dense<0.0> : tensor %cst_1 = arith.constant dense<4.0> : tensor %cst_2 = arith.constant dense<1.0> : tensor // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<4xf32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor } // CHECK-LABEL: @range_float_neg_delta func.func @range_float_neg_delta() -> tensor { %cst = arith.constant dense<0.0> : tensor %cst_1 = arith.constant dense<-4.0> : tensor %cst_2 = arith.constant dense<-1.0> : tensor // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[0.000000e+00, -1.000000e+00, -2.000000e+00, -3.000000e+00]> : tensor<4xf32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor } // CHECK-LABEL: @range_float_nonzero_base func.func @range_float_nonzero_base() -> tensor { %cst = arith.constant dense<2.0> : tensor %cst_1 = arith.constant dense<7.0> : tensor %cst_2 = arith.constant dense<1.5> : tensor // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[2.000000e+00, 3.500000e+00, 5.000000e+00, 6.500000e+00]> : tensor<4xf32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor } // CHECK-LABEL: @transpose_dynamic func.func @transpose_dynamic() -> tensor { %cst = arith.constant dense<[1, 2, 3]> : tensor<3xi32> %cst_perm = arith.constant dense<0> : tensor<1xi32> // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<{{\[}}1, 2, 3]> : tensor<3xi32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tfl.transpose"(%cst, %cst_perm) : (tensor<3xi32>, tensor<1xi32>) -> tensor func.return %0 : tensor } // CHECK-LABEL: @add_dense_dense_int_same_shape_dynamic func.func @add_dense_dense_int_same_shape_dynamic() -> tensor { %0 = arith.constant dense<[15, 23, -44, -2]> : tensor<4xi32> %1 = arith.constant dense<[-10, -1, 42, 100]> : tensor<4xi32> %2 = "tfl.add"(%0, %1) {fused_activation_function = "NONE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor func.return %2 : tensor // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[5, 22, -2, 98]> : tensor<4xi32>}> : () -> tensor // CHECK: return %[[CST]] } // CHECK-LABEL: @concat_2_tensors_1_empty func.func @concat_2_tensors_1_empty() -> tensor<2xi32> { %1 = arith.constant dense<1> : tensor<2xi32> %2 = arith.constant dense<[]> : tensor<0xi32> %3 = "tfl.concatenation"(%1, %2) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<2xi32>, tensor<0xi32>) -> tensor<2xi32> func.return %3 : tensor<2xi32> // CHECK: %[[CST:.*]] = arith.constant dense<1> : tensor<2xi32> // CHECK: return %[[CST]] : tensor<2xi32> } // CHECK-LABEL: @concat_3_tensors_1_empty func.func @concat_3_tensors_1_empty() -> tensor { %0 = arith.constant dense<1> : tensor<2xi32> %1 = arith.constant dense<1> : tensor<2xi32> %2 = arith.constant dense<[]> : tensor<0xi32> %3 = "tfl.concatenation"(%0, %1, %2) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<2xi32>, tensor<2xi32>, tensor<0xi32>) -> tensor func.return %3 : tensor // CHECK: %0 = "tfl.concatenation"(%[[CST]], %[[CST]]) <{axis = 0 : i32, fused_activation_function = "NONE"}> // CHECK: return %0 : tensor } // CHECK-LABEL: @concatConstantTensorsFirstDim func.func @concatConstantTensorsFirstDim() -> tensor<2x2x3xi32> { %cst_0 = arith.constant dense<0> : tensor<1x2x3xi32> %cst_1 = arith.constant dense<1> : tensor<1x2x3xi32> %0 = "tfl.concatenation"(%cst_0, %cst_1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xi32>, tensor<1x2x3xi32>) -> tensor<2x2x3xi32> func.return %0 : tensor<2x2x3xi32> // CHECK: %[[CST:.*]] = arith.constant dense<[{{\[}}{{\[}}0, 0, 0], {{\[}}0, 0, 0]], {{\[}}{{\[}}1, 1, 1], {{\[}}1, 1, 1]]]> : tensor<2x2x3xi32> // CHECK-NOT: constant-dense // CHECK-NOT: "tfl.concatenation" // CHECK: return %[[CST]] } // CHECK-LABEL: @concatConstantTensorsMiddleDim func.func @concatConstantTensorsMiddleDim() -> tensor<1x4x3xi32> { %cst_0 = arith.constant dense<0> : tensor<1x2x3xi32> %cst_1 = arith.constant dense<1> : tensor<1x2x3xi32> %0 = "tfl.concatenation"(%cst_0, %cst_1) {axis = 1 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xi32>, tensor<1x2x3xi32>) -> tensor<1x4x3xi32> func.return %0 : tensor<1x4x3xi32> // CHECK: %[[CST:.*]] = arith.constant dense<[{{\[}}{{\[}}0, 0, 0], {{\[}}0, 0, 0], {{\[}}1, 1, 1], {{\[}}1, 1, 1]]]> : tensor<1x4x3xi32> // CHECK-NOT: constant-dense // CHECK-NOT: "tfl.concatenation" // CHECK: return %[[CST]] } // CHECK-LABEL: @concatConstantTensorsLastDim func.func @concatConstantTensorsLastDim() -> tensor<1x2x6xi32> { %cst_0 = arith.constant dense<0> : tensor<1x2x3xi32> %cst_1 = arith.constant dense<1> : tensor<1x2x3xi32> %0 = "tfl.concatenation"(%cst_0, %cst_1) {axis = 2 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xi32>, tensor<1x2x3xi32>) -> tensor<1x2x6xi32> func.return %0 : tensor<1x2x6xi32> // CHECK: %[[CST:.*]] = arith.constant dense<[{{\[}}{{\[}}0, 0, 0, 1, 1, 1], {{\[}}0, 0, 0, 1, 1, 1]]]> : tensor<1x2x6xi32> // CHECK-NOT: constant-dense // CHECK-NOT: "tfl.concatenation" // CHECK: return %[[CST]] } // CHECK-LABEL: @rsqrt_bf16 func.func @rsqrt_bf16() -> tensor { %cst = arith.constant dense<4.0> : tensor %0 = "tfl.rsqrt"(%cst) : (tensor) -> tensor func.return %0 : tensor // CHECK: %[[CST:.*]] = arith.constant dense<5.000000e-01> : tensor // CHECK: return %[[CST]] } // CHECK-LABEL: @cast_i64_to_i32 func.func @cast_i64_to_i32() -> tensor<5xi32> { %cst = arith.constant dense<[-1, 0, 1, 2147483647, 2147483648]> : tensor<5xi64> %0 = "tfl.cast"(%cst) : (tensor<5xi64>) -> tensor<5xi32> func.return %0 : tensor<5xi32> // CHECK: %[[CST:.*]] = arith.constant dense<[-1, 0, 1, 2147483647, -2147483648]> : tensor<5xi32> // CHECK: return %[[CST]] } // CHECK-LABEL: @cast_i32_to_i8 func.func @cast_i32_to_i8() -> tensor<6xi8> { %cst = arith.constant dense<[0, -1, 256, 127, -128, -129]> : tensor<6xi32> %0 = "tfl.cast"(%cst) : (tensor<6xi32>) -> tensor<6xi8> func.return %0 : tensor<6xi8> // CHECK: %[[CST:.*]] = arith.constant dense<[0, -1, 0, 127, -128, 127]> : tensor<6xi8> // CHECK: return %[[CST]] } // CHECK-LABEL: @cast_i8_to_i32 func.func @cast_i8_to_i32() -> tensor<4xi32> { %cst = arith.constant dense<[0, 128, -1, -128]> : tensor<4xi8> %0 = "tfl.cast"(%cst) : (tensor<4xi8>) -> tensor<4xi32> func.return %0 : tensor<4xi32> // CHECK: %[[CST:.*]] = arith.constant dense<[0, -128, -1, -128]> : tensor<4xi32> // CHECK: return %[[CST]] } // CHECK-LABEL: @cast_identity func.func @cast_identity(%arg0 : tensor<7xf32>) -> tensor<7xf32> { %0 = "tfl.cast"(%arg0) : (tensor<7xf32>) -> tensor<7xf32> func.return %0 : tensor<7xf32> // CHECK: return %arg0 : tensor<7xf32> } // CHECK-LABEL: @cast_i1_to_i8 func.func @cast_i1_to_i8() -> tensor<2xi8> { %cst = arith.constant dense<[false, true]> : tensor<2xi1> %0 = "tfl.cast"(%cst) : (tensor<2xi1>) -> tensor<2xi8> func.return %0 : tensor<2xi8> // CHECK: %[[CST:.*]] = arith.constant dense<[0, 1]> : tensor<2xi8> // CHECK: return %[[CST]] } // CHECK-LABEL: @cast_i8_to_i1 func.func @cast_i8_to_i1() -> tensor<4xi1> { %cst = arith.constant dense<[0, 1, 2, -1]> : tensor<4xi8> %0 = "tfl.cast"(%cst) : (tensor<4xi8>) -> tensor<4xi1> func.return %0 : tensor<4xi1> // CHECK: %[[CST:.*]] = arith.constant dense<[false, true, true, true]> : tensor<4xi1> // CHECK: return %[[CST]] } // CHECK-LABEL: @cast_f32_to_i32 func.func @cast_f32_to_i32() -> tensor<8xi32> { %cst = arith.constant dense<[-1.0, 0.0, 1.5, 0.99, 1.175494351e-38, 3.402823466e+38, -3.402823466e+38, -1.175494351e-38]> : tensor<8xf32> %0 = "tfl.cast"(%cst) : (tensor<8xf32>) -> tensor<8xi32> func.return %0 : tensor<8xi32> } // CHECK: %cst = arith.constant dense<[-1, 0, 1, 0, 0, 2147483647, -2147483648, 0]> : tensor<8xi32> // CHECK-LABEL: @cast_f32_to_i64 func.func @cast_f32_to_i64() -> tensor<4xi64> { %cst = arith.constant dense<[-1.0, 0.0, 1.5, 0.99]> : tensor<4xf32> %0 = "tfl.cast"(%cst) : (tensor<4xf32>) -> tensor<4xi64> func.return %0 : tensor<4xi64> } // CHECK: %cst = arith.constant dense<[-1, 0, 1, 0]> : tensor<4xi64> // CHECK-LABEL: @cast_i32_to_f32 func.func @cast_i32_to_f32() -> tensor<5xf32> { %cst = arith.constant dense<[-1, 0, 2, 2147483647, -2147483648]> : tensor<5xi32> %0 = "tfl.cast"(%cst) : (tensor<5xi32>) -> tensor<5xf32> func.return %0 : tensor<5xf32> } // CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 2.000000e+00, 2.14748365E+9, -2.14748365E+9]> : tensor<5xf32> // CHECK-LABEL: @cast_bool_to_f32 func.func @cast_bool_to_f32() -> tensor<2xf32> { %cst = arith.constant dense<[true, false]> : tensor<2xi1> %0 = "tfl.cast"(%cst) : (tensor<2xi1>) -> tensor<2xf32> func.return %0 : tensor<2xf32> } // CHECK: %cst = arith.constant dense<[1.000000e+00, 0.000000e+00]> : tensor<2xf32> // CHECK-LABEL: @cast_f64_to_f32 func.func @cast_f64_to_f32() -> tensor<4xf32> { %cst = arith.constant dense<[-1.0, 0.0, 1.5, 100.0]> : tensor<4xf64> %0 = "tfl.cast"(%cst) : (tensor<4xf64>) -> tensor<4xf32> func.return %0 : tensor<4xf32> } // CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.500000e+00, 1.000000e+02]> : tensor<4xf32> // CHECK-LABEL: @cast_f32_to_f64 func.func @cast_f32_to_f64() -> tensor<4xf64> { %cst = arith.constant dense<[-1.0, 0.0, 1.5, 100.0]> : tensor<4xf32> %0 = "tfl.cast"(%cst) : (tensor<4xf32>) -> tensor<4xf64> func.return %0 : tensor<4xf64> } // CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.500000e+00, 1.000000e+02]> : tensor<4xf64> // CHECK-LABEL: @cast_f32_to_f16 func.func @cast_f32_to_f16() -> tensor<4xf16> { %cst = arith.constant dense<[-1.0, 0.0, 1.5, 100.0]> : tensor<4xf32> %0 = "tfl.cast"(%cst) : (tensor<4xf32>) -> tensor<4xf16> func.return %0 : tensor<4xf16> } // CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.500000e+00, 1.000000e+02]> : tensor<4xf16> // CHECK-LABEL: @ConstantFoldFullyConnectedSmall func.func @ConstantFoldFullyConnectedSmall() -> tensor<3xf32> { %cst_input = arith.constant dense<[2.0, 3.0]> : tensor<2xf32> %cst_weights = arith.constant dense<[[5.0, 7.0], [11.0, 13.0], [17.0, 19.0]]> : tensor<3x2xf32> %cst_bias = arith.constant dense<[23.0, 29.0, 31.0]> : tensor<3xf32> %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2xf32>, tensor<3x2xf32>, tensor<3xf32>) -> tensor<3xf32> func.return %0 : tensor<3xf32> // [54, 90, 122] // CHECK: %[[CST:.*]] = arith.constant dense<[5.400000e+01, 9.000000e+01, 1.220000e+02]> : tensor<3xf32> // CHECK: return %[[CST]] } // CHECK-LABEL: @ConstantFoldFullyConnectedLarge func.func @ConstantFoldFullyConnectedLarge() -> tensor<1024xf32> { %cst_input = arith.constant dense<1.0> : tensor<512xf32> %cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32> %cst_bias = arith.constant dense<4.0> : tensor<1024xf32> %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<512xf32>, tensor<1024x512xf32>, tensor<1024xf32>) -> tensor<1024xf32> func.return %0 : tensor<1024xf32> // 1.0 * 2.0 * 512 + 4.0 = 1028.0 // CHECK: %[[CST:.*]] = arith.constant dense<1.028000e+03> : tensor<1024xf32> // CHECK: return %[[CST]] } // CHECK-LABEL: @ConstantFoldFullyConnectedNoBias func.func @ConstantFoldFullyConnectedNoBias() -> tensor<1024xf32> { %cst_input = arith.constant dense<1.0> : tensor<512xf32> %cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32> %cst_bias = "tfl.no_value"() {value = unit} : () -> none %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<512xf32>, tensor<1024x512xf32>, none) -> tensor<1024xf32> func.return %0 : tensor<1024xf32> // 1.0 * 2.0 * 512 = 1024.0 // CHECK: %[[CST:.*]] = arith.constant dense<1.024000e+03> : tensor<1024xf32> // CHECK: return %[[CST]] } // CHECK-LABEL: @NoFoldFullyConnectedNonFloat func.func @NoFoldFullyConnectedNonFloat() -> tensor<1024xf32> { %cst_input = arith.constant dense<1.0> : tensor<512xf32> %cst_weights = arith.constant dense<2> : tensor<1024x512xi8> %cst_bias = arith.constant dense<4.0> : tensor<1024xf32> %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<512xf32>, tensor<1024x512xi8>, tensor<1024xf32>) -> tensor<1024xf32> func.return %0 : tensor<1024xf32> // CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : tensor<512xf32> // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<2> : tensor<1024x512xi8> // CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<4.000000e+00> : tensor<1024xf32> // CHECK: %[[VAL:.*]] = "tfl.fully_connected"(%[[CST]], %[[CST_0]], %[[CST_1]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<512xf32>, tensor<1024x512xi8>, tensor<1024xf32>) -> tensor<1024xf32> // CHECK: return %[[VAL]] : tensor<1024xf32> } // CHECK-LABEL: @ConstantFoldFullyConnectedHighRank func.func @ConstantFoldFullyConnectedHighRank() -> tensor<2x1024xf32> { %cst_input = arith.constant dense<1.0> : tensor<2x512xf32> %cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32> %cst_bias = arith.constant dense<4.0> : tensor<1024xf32> %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x512xf32>, tensor<1024x512xf32>, tensor<1024xf32>) -> tensor<2x1024xf32> func.return %0 : tensor<2x1024xf32> // 1.0 * 2.0 * 512 + 4.0 = 1028.0 // CHECK: %[[CST:.*]] = arith.constant dense<1.028000e+03> : tensor<2x1024xf32> // CHECK: return %[[CST]] } // CHECK-LABEL: @ConstantFoldFullyConnectedCheckPrecision func.func @ConstantFoldFullyConnectedCheckPrecision() -> tensor<1xf32> { %cst_input = arith.constant dense<1.0> : tensor<4xf32> %cst_weights = arith.constant dense<[[1.0, 1.0e38, 1.0, -1.0e38]]> : tensor<1x4xf32> %cst_bias = arith.constant dense<0.0> : tensor<1xf32> %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4xf32>, tensor<1x4xf32>, tensor<1xf32>) -> tensor<1xf32> func.return %0 : tensor<1xf32> // CHECK: %[[CST:.*]] = arith.constant dense<2.000000e+00> : tensor<1xf32> // CHECK: return %[[CST]] } // CHECK-LABEL: fully_connected_with_unit_dim func.func @fully_connected_with_unit_dim() -> tensor<1x5xf32> { %0 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<1x5xf32>}> : () -> tensor<1x5xf32> %1 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<5x5xf32>}> : () -> tensor<5x5xf32> %2 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<1x5xf32>}> : () -> tensor<1x5xf32> %3 = "tfl.fully_connected"(%0, %1, %2) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"}> : (tensor<1x5xf32>, tensor<5x5xf32>, tensor<1x5xf32>) -> tensor<1x5xf32> return %3 : tensor<1x5xf32> } // CHECK: %cst = arith.constant dense<6.000000e+00> : tensor<1x5xf32> // CHECK-NOT: fully_connected // CHECK-LABEL: @ConstantFoldFullyConnectedBatched func.func @ConstantFoldFullyConnectedBatched() -> tensor<13x1536xf32> { %cst_input = arith.constant dense<1.0> : tensor<13x1536xf32> %cst_weights = arith.constant dense<1.0> : tensor<1536x1536xf32> %cst_bias = "tfl.no_value"() {value = unit} : () -> none %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"} : (tensor<13x1536xf32>, tensor<1536x1536xf32>, none) -> tensor<13x1536xf32> func.return %0 : tensor<13x1536xf32> // 1.0 * 1.0 * 1536 = 1536.0 // CHECK: %[[CST:.*]] = arith.constant dense<1.536000e+03> : tensor<13x1536xf32> // CHECK: return %[[CST]] } // CHECK-LABEL: @ShapeOpI32 func.func @ShapeOpI32(%arg0 : tensor<576x72xf32>) -> tensor<2xi32> { %0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi32> func.return %0 : tensor<2xi32> // CHECK: %[[CST:.*]] = arith.constant dense<[576, 72]> : tensor<2xi32> // CHECK: return %[[CST]] } // CHECK-LABEL: @ShapeOpI64 func.func @ShapeOpI64(%arg0 : tensor<576x72xf32>) -> tensor<2xi64> { %0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi64> func.return %0 : tensor<2xi64> // CHECK: %[[CST:.*]] = arith.constant dense<[576, 72]> : tensor<2xi64> // CHECK: return %[[CST]] } // CHECK-LABEL: @ConstFoldStridedSlice func.func @ConstFoldStridedSlice(%arg0 : tensor<15600xf32>) -> tensor<15600xf32> { %0 = "tfl.pseudo_const"() {value = dense<15600> : tensor<1xi32>} : () -> tensor<1xi32> %1 = "tfl.pseudo_const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> %2 = "tfl.pseudo_const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> %3 = "tfl.strided_slice"(%arg0, %1, %0, %2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false} : (tensor<15600xf32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<15600xf32> func.return %3 : tensor<15600xf32> // CHECK: return %arg0 } func.func @ConstFoldStridedSliceMultiDims(%arg0 : tensor<10x10x10xf32>) -> tensor<10x10x10xf32> { %0 = "tfl.pseudo_const"() {value = dense<[10, 10, 10]> : tensor<3xi32>} : () -> tensor<3xi32> %1 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32> %2 = "tfl.pseudo_const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32> %3 = "tfl.strided_slice"(%arg0, %1, %0, %2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false} : (tensor<10x10x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<10x10x10xf32> func.return %3 : tensor<10x10x10xf32> // CHECK: return %arg0 } func.func @NotFoldStridedSlice(%arg0 : tensor<10x10x10xf32>) -> tensor<9x9x9xf32> { %0 = "tfl.pseudo_const"() {value = dense<[9, 9, 9]> : tensor<3xi32>} : () -> tensor<3xi32> %1 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32> %2 = "tfl.pseudo_const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32> %3 = "tfl.strided_slice"(%arg0, %1, %0, %2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false} : (tensor<10x10x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<9x9x9xf32> func.return %3 : tensor<9x9x9xf32> // CHECK: %[[STRIDED_SLICE:.*]] = "tfl.strided_slice" // CHECK: return %[[STRIDED_SLICE]] } func.func @ConstFoldPad(%arg0: tensor<15600xf32>) -> tensor<15600xf32> { %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32> %1 = "tfl.pad"(%arg0, %0) : (tensor<15600xf32>, tensor<1x2xi32>) -> tensor<15600xf32> func.return %1 : tensor<15600xf32> // CHECK: return %arg0 } func.func @ConstFoldPadV2(%arg0: tensor<15600xf32>) -> tensor<15600xf32> { %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32> %1 = "tfl.pseudo_const"() {value = dense<0.0> : tensor} : () -> tensor %2 = "tfl.padv2"(%arg0, %0, %1) : (tensor<15600xf32>, tensor<1x2xi32>, tensor) -> tensor<15600xf32> func.return %2 : tensor<15600xf32> // CHECK: return %arg0 } // CHECK-LABEL: @ConstFoldEmbeddingLookup func.func @ConstFoldEmbeddingLookup() -> (tensor<5x2xf32>, tensor<3x2x2xf32>) { %index0 = "tfl.pseudo_const"() {value = dense<[2, 1, 0, 0, 2]> : tensor<5xi32>} : () -> tensor<5xi32> %index1 = "tfl.pseudo_const"() {value = dense<[0, 1, 0]> : tensor<3xi32>} : () -> tensor<3xi32> %value0 = "tfl.pseudo_const"() {value = dense<[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]> : tensor<3x2xf32>} : () -> tensor<3x2xf32> %value1 = "tfl.pseudo_const"() {value = dense<[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32> %lookup0 = "tfl.embedding_lookup"(%index0, %value0) : (tensor<5xi32>, tensor<3x2xf32>) -> tensor<5x2xf32> %lookup1 = "tfl.embedding_lookup"(%index1, %value1) : (tensor<3xi32>, tensor<2x2x2xf32>) -> tensor<3x2x2xf32> func.return %lookup0, %lookup1 : tensor<5x2xf32>, tensor<3x2x2xf32> // CHECK-DAG: %[[LOOKUP0:.*]] = arith.constant dense<{{\[\[}}5.000000e+00, 6.000000e+00], [3.000000e+00, 4.000000e+00], [1.000000e+00, 2.000000e+00], [1.000000e+00, 2.000000e+00], [5.000000e+00, 6.000000e+00]]> : tensor<5x2xf32> // CHECK-DAG: %[[LOOKUP1:.*]] = arith.constant 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]], {{\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00]]]> : tensor<3x2x2xf32> // CHECK: return %[[LOOKUP0]], %[[LOOKUP1]] : tensor<5x2xf32>, tensor<3x2x2xf32> } // CHECK-LABEL: @less_int_both_splat func.func @less_int_both_splat() -> tensor<4xi1> { %0 = arith.constant dense<3> : tensor<4xi32> %1 = arith.constant dense<10> : tensor<4xi32> %2 = "tfl.less"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense : tensor<4xi1> // CHECK-LABEL: @less_int_one_splat func.func @less_int_one_splat() -> tensor<4xi1> { %0 = arith.constant dense<3> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.less"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK:%cst = arith.constant dense<[true, false, false, false]> : tensor<4xi1> // CHECK-LABEL: @less_int func.func @less_int() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.less"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, false, false, true]> : tensor<4xi1> // CHECK-LABEL: @less_int64 func.func @less_int64() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64> %2 = "tfl.less"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, false, false, true]> : tensor<4xi1> // CHECK-LABEL: @less_float func.func @less_float() -> tensor<4xi1> { %0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32> %1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32> %2 = "tfl.less"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, false, false, true]> : tensor<4xi1> // CHECK-LABEL: @less_equal_int func.func @less_equal_int() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.less_equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, true, false, true]> : tensor<4xi1> // CHECK-LABEL: @less_equal_int64 func.func @less_equal_int64() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64> %2 = "tfl.less_equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, true, false, true]> : tensor<4xi1> // CHECK-LABEL: @less_equal_float func.func @less_equal_float() -> tensor<4xi1> { %0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32> %1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32> %2 = "tfl.less_equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, true, false, true]> : tensor<4xi1> // CHECK-LABEL: @greater_int func.func @greater_int() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.greater"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, false, true, false]> : tensor<4xi1> // CHECK-LABEL: @greater_int64 func.func @greater_int64() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64> %2 = "tfl.greater"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, false, true, false]> : tensor<4xi1> // CHECK-LABEL: @greater_float func.func @greater_float() -> tensor<4xi1> { %0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32> %1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32> %2 = "tfl.greater"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, false, true, false]> : tensor<4xi1> // CHECK-LABEL: @greater_equal_int func.func @greater_equal_int() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.greater_equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, true, true, false]> : tensor<4xi1> // CHECK-LABEL: @greater_equal_int64 func.func @greater_equal_int64() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64> %2 = "tfl.greater_equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, true, true, false]> : tensor<4xi1> // CHECK-LABEL: @greater_equal_float func.func @greater_equal_float() -> tensor<4xi1> { %0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32> %1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32> %2 = "tfl.greater_equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, true, true, false]> : tensor<4xi1> // CHECK-LABEL: @equal_int func.func @equal_int() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, true, false, false]> : tensor<4xi1> // CHECK-LABEL: @equal_int64 func.func @equal_int64() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64> %2 = "tfl.equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, true, false, false]> : tensor<4xi1> // CHECK-LABEL: @equal_float func.func @equal_float() -> tensor<4xi1> { %0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32> %1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32> %2 = "tfl.equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[false, true, false, false]> : tensor<4xi1> // CHECK-LABEL: @not_equal_int func.func @not_equal_int() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32> %2 = "tfl.not_equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, false, true, true]> : tensor<4xi1> // CHECK-LABEL: @not_equal_int64 func.func @not_equal_int64() -> tensor<4xi1> { %0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64> %1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64> %2 = "tfl.not_equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, false, true, true]> : tensor<4xi1> // CHECK-LABEL: @not_equal_float func.func @not_equal_float() -> tensor<4xi1> { %0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32> %1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32> %2 = "tfl.not_equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1> func.return %2 : tensor<4xi1> } // CHECK: %cst = arith.constant dense<[true, false, true, true]> : tensor<4xi1> // CHECK-LABEL: @logical_or func.func @logical_or() -> tensor<3xi1> { %0 = arith.constant dense<[true, false, true]> : tensor<3xi1> %1 = arith.constant dense<[false, false, true]> : tensor<3xi1> %2 = "tfl.logical_or"(%0, %1) : (tensor<3xi1>, tensor<3xi1>) -> tensor<3xi1> func.return %2 : tensor<3xi1> } // CHECK: %cst = arith.constant dense<[true, false, true]> : tensor<3xi1> // CHECK-LABEL: @logical_and func.func @logical_and() -> tensor<3xi1> { %0 = arith.constant dense<[true, false, true]> : tensor<3xi1> %1 = arith.constant dense<[false, false, true]> : tensor<3xi1> %2 = "tfl.logical_and"(%0, %1) : (tensor<3xi1>, tensor<3xi1>) -> tensor<3xi1> func.return %2 : tensor<3xi1> } // CHECK: %cst = arith.constant dense<[false, false, true]> : tensor<3xi1> // CHECK-LABEL: @select_splat_cond func.func @select_splat_cond() -> tensor<4xi32> { %cond = arith.constant dense : tensor<4xi1> %0 = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> %1 = arith.constant dense<[-1, -2, -3, -4]> : tensor<4xi32> %2 = "tfl.select"(%cond, %0, %1) : (tensor<4xi1>, tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> func.return %2 : tensor<4xi32> } // CHECK: %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> // CHECK-LABEL: select_splat_lhs func.func @select_splat_lhs() -> tensor<4xi32> { %cond = arith.constant dense<[true, true, false, false]> : tensor<4xi1> %0 = arith.constant dense<0> : tensor<4xi32> %1 = arith.constant dense<[-1, -2, -3, -4]> : tensor<4xi32> %2 = "tfl.select"(%cond, %0, %1) : (tensor<4xi1>, tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> func.return %2 : tensor<4xi32> } // CHECK: %cst = arith.constant dense<[0, 0, -3, -4]> : tensor<4xi32> // CHECK-LABEL: select_float func.func @select_float() -> tensor<4xf32> { %cond = arith.constant dense<[true, true, false, false]> : tensor<4xi1> %0 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32> %1 = arith.constant dense<[-1.0, -2.0, -3.0, -4.0]> : tensor<4xf32> %2 = "tfl.select"(%cond, %0, %1) : (tensor<4xi1>, tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> func.return %2 : tensor<4xf32> } // CHECK: %cst = arith.constant dense<[1.000000e+00, 2.000000e+00, -3.000000e+00, -4.000000e+00]> : tensor<4xf32 // CHECK-LABEL: ceil func.func @ceil() -> tensor<3xf32> { %cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf32> %0 = "tfl.ceil"(%cst) : (tensor<3xf32>) -> tensor<3xf32> func.return %0 : tensor<3xf32> } // CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.000000e+00]> : tensor<3xf32> // CHECK-LABEL: ceil_f64 func.func @ceil_f64() -> tensor<3xf64> { %cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf64> %0 = "tfl.ceil"(%cst) : (tensor<3xf64>) -> tensor<3xf64> func.return %0 : tensor<3xf64> } // CHECK: tfl.ceil // CHECK-LABEL: floor func.func @floor() -> tensor<3xf32> { %cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf32> %0 = "tfl.floor"(%cst) : (tensor<3xf32>) -> tensor<3xf32> func.return %0 : tensor<3xf32> } // CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 0.000000e+00]> : tensor<3xf32> // CHECK-LABEL: floor_f64 func.func @floor_f64() -> tensor<3xf64> { %cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf64> %0 = "tfl.floor"(%cst) : (tensor<3xf64>) -> tensor<3xf64> func.return %0 : tensor<3xf64> } // CHECK: tfl.floor // CHECK-LABEL: exp func.func @exp() -> tensor<4xf32> { %cst = arith.constant dense<[-1.0, 0.0, 0.99, 0.36787944117]> : tensor<4xf32> %0 = "tfl.exp"(%cst) : (tensor<4xf32>) -> tensor<4xf32> func.return %0 : tensor<4xf32> } // CHECK: %cst = arith.constant dense<[0.36787945, 1.000000e+00, 2.69123459, 1.44466782]> : tensor<4xf32> // CHECK-LABEL: exp_f64 func.func @exp_f64() -> tensor<4xf64> { %cst = arith.constant dense<[-1.0, 0.0, 0.99, 0.36787944117]> : tensor<4xf64> %0 = "tfl.exp"(%cst) : (tensor<4xf64>) -> tensor<4xf64> func.return %0 : tensor<4xf64> } // CHECK: tfl.exp // CHECK-LABEL: pow_float func.func @pow_float() -> tensor<3xf32> { %0 = arith.constant dense<[1.0, 0.0, 2.0]> : tensor<3xf32> %1 = arith.constant dense<[2.0, 3.0, -1.5]> : tensor<3xf32> %2 = "tfl.pow"(%0, %1) : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> func.return %2 : tensor<3xf32> } // CHECK: %cst = arith.constant dense<[1.000000e+00, 0.000000e+00, 0.353553385]> : tensor<3xf32> // CHECK-LABEL: pow_int func.func @pow_int() -> tensor<3xi32> { %0 = arith.constant dense<[1, 0, 2]> : tensor<3xi32> %1 = arith.constant dense<[2, 3, -1]> : tensor<3xi32> %2 = "tfl.pow"(%0, %1) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32> func.return %2 : tensor<3xi32> } // CHECK: %cst = arith.constant dense<[1, 0, 0]> : tensor<3xi32> // CHECK-LABEL: logical_not func.func @logical_not() -> tensor<3xi1> { %cst = arith.constant dense<[false, true, false]> : tensor<3xi1> %0 = "tfl.logical_not"(%cst) : (tensor<3xi1>) -> tensor<3xi1> func.return %0 : tensor<3xi1> } // CHECK: %cst = arith.constant dense<[true, false, true]> : tensor<3xi1> // CHECK-LABEL: logical_not_splat func.func @logical_not_splat() -> tensor<3xi1> { %cst = arith.constant dense : tensor<3xi1> %0 = "tfl.logical_not"(%cst) : (tensor<3xi1>) -> tensor<3xi1> func.return %0 : tensor<3xi1> } // CHECK: %cst = arith.constant dense : tensor<3xi1> // CHECK-LABEL: bitwise_xor_i32 func.func @bitwise_xor_i32() -> tensor<3xi32> { %0 = arith.constant dense<[0, 5, 3]> : tensor<3xi32> %1 = arith.constant dense<[5, 0, 7]> : tensor<3xi32> %2 = "tfl.bitwise_xor"(%0, %1) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32> func.return %2 : tensor<3xi32> } // CHECK: %cst = arith.constant dense<[5, 5, 4]> : tensor<3xi32> // CHECK-LABEL: bitwise_xor_i8 func.func @bitwise_xor_i8() -> tensor<3xi8> { %0 = arith.constant dense<[0, 5, 3]> : tensor<3xi8> %1 = arith.constant dense<[5, 0, 7]> : tensor<3xi8> %2 = "tfl.bitwise_xor"(%0, %1) : (tensor<3xi8>, tensor<3xi8>) -> tensor<3xi8> func.return %2 : tensor<3xi8> } // CHECK: %cst = arith.constant dense<[5, 5, 4]> : tensor<3xi8> // CHECK-LABEL: relu func.func @relu() -> tensor<3xf32> { %cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf32> %0 = "tfl.relu"(%cst) : (tensor<3xf32>) -> tensor<3xf32> func.return %0 : tensor<3xf32> } // CHECK: %cst = arith.constant dense<[0.000000e+00, 0.000000e+00, 9.900000e-01]> : tensor<3xf32> // CHECK-LABEL: slice func.func @slice_first_dim() -> tensor<1x1x5x6xf32> { %cst_0 = arith.constant dense<9.000000e+00> : tensor<2x1x5x6xf32> %cst_1 = arith.constant dense<0> : tensor<4xi32> %cst_2 = arith.constant dense<[1, 1, 5, 6]> : tensor<4xi32> %0 = "tfl.slice"(%cst_0, %cst_1, %cst_2) : (tensor<2x1x5x6xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x1x5x6xf32> func.return %0 : tensor<1x1x5x6xf32> } // CHECK %cst = arith.constant dense<9.000000e+00> : tensor<1x1x5x6xf32> // CHECK-LABEL: slice_trivial func.func @slice_trivial(%arg0: tensor<2x1x5x6xf32>) -> tensor<2x1x5x6xf32> { %cst_1 = arith.constant dense<0> : tensor<4xi32> %cst_2 = arith.constant dense<[2, 1, 5, 6]> : tensor<4xi32> %0 = "tfl.slice"(%arg0, %cst_1, %cst_2) : (tensor<2x1x5x6xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<2x1x5x6xf32> func.return %0 : tensor<2x1x5x6xf32> } // CHECK-NOT: tfl.slice // CHECK-LABEL: sum func.func @sum() -> tensor<2xf32> { %cst = arith.constant dense<[0, 1]> : tensor<2xi32> %cst_1 = arith.constant dense<[[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]]> : tensor<2x2x2xf32> %0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = false}> : (tensor<2x2x2xf32>, tensor<2xi32>) -> tensor<2xf32> func.return %0 : tensor<2xf32> } // CHECK: arith.constant dense<[1.200000e+01, 1.600000e+01]> : tensor<2xf32> // CHECK-LABEL: sum_keep_dims func.func @sum_keep_dims() -> tensor<1x1x2xf32> { %cst = arith.constant dense<[0, 1]> : tensor<2xi32> %cst_1 = arith.constant dense<[[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]]> : tensor<2x2x2xf32> %0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = true}> : (tensor<2x2x2xf32>, tensor<2xi32>) -> tensor<1x1x2xf32> func.return %0 : tensor<1x1x2xf32> } // CHECK-LITERAL: arith.constant dense<[[[1.200000e+01, 1.600000e+01]]]> : tensor<1x1x2xf32> // CHECK-LABEL: sum_general_reduction_dims func.func @sum_general_reduction_dims() -> tensor<3xf32> { %cst = arith.constant dense<[0, 2]> : tensor<2xi32> %cst_1 = arith.constant dense<2.000000e+00> : tensor<1x3x2xf32> %0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = false}> : (tensor<1x3x2xf32>, tensor<2xi32>) -> tensor<3xf32> func.return %0 : tensor<3xf32> } // CHECK: %cst = arith.constant dense<4.000000e+00> : tensor<3xf32> // CHECK-LABEL: sum_general_reduction_dims_keep_dims func.func @sum_general_reduction_dims_keep_dims() -> tensor<1x3x1xf32> { %cst = arith.constant dense<[0, 2]> : tensor<2xi32> %cst_1 = arith.constant dense<2.000000e+00> : tensor<1x3x2xf32> %0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = true}> : (tensor<1x3x2xf32>, tensor<2xi32>) -> tensor<1x3x1xf32> func.return %0 : tensor<1x3x1xf32> } // CHECK: %cst = arith.constant dense<4.000000e+00> : tensor<1x3x1xf32> // CHECK-LABEL: gather func.func @gather() -> (tensor<2x3x4x5xi16>, tensor<2x3x4x5xi16>) { %params = arith.constant dense<[ [[[1111, 1112, 1113, 1114, 1115], [1121, 1122, 1123, 1124, 1125], [1131, 1132, 1133, 1134, 1135], [1141, 1142, 1143, 1144, 1145], [1151, 1152, 1153, 1154, 1155], [1161, 1162, 1163, 1164, 1165]], [[1211, 1212, 1213, 1214, 1215], [1221, 1222, 1223, 1224, 1225], [1231, 1232, 1233, 1234, 1235], [1241, 1242, 1243, 1244, 1245], [1251, 1252, 1253, 1254, 1255], [1261, 1262, 1263, 1264, 1265]], [[1311, 1312, 1313, 1314, 1315], [1321, 1322, 1323, 1324, 1325], [1331, 1332, 1333, 1334, 1335], [1341, 1342, 1343, 1344, 1345], [1351, 1352, 1353, 1354, 1355], [1361, 1362, 1363, 1364, 1365]]], [[[2111, 2112, 2113, 2114, 2115], [2121, 2122, 2123, 2124, 2125], [2131, 2132, 2133, 2134, 2135], [2141, 2142, 2143, 2144, 2145], [2151, 2152, 2153, 2154, 2155], [2161, 2162, 2163, 2164, 2165]], [[2211, 2212, 2213, 2214, 2215], [2221, 2222, 2223, 2224, 2225], [2231, 2232, 2233, 2234, 2235], [2241, 2242, 2243, 2244, 2245], [2251, 2252, 2253, 2254, 2255], [2261, 2262, 2263, 2264, 2265]], [[2311, 2312, 2313, 2314, 2315], [2321, 2322, 2323, 2324, 2325], [2331, 2332, 2333, 2334, 2335], [2341, 2342, 2343, 2344, 2345], [2351, 2352, 2353, 2354, 2355], [2361, 2362, 2363, 2364, 2365]]]]> : tensor<2x3x6x5xi16> %indices = arith.constant dense<[[5, 4, 3, 2], [3, 2, 1, 0]]> : tensor<2x4xi64> %gathered = "tfl.gather"(%params, %indices) <{axis = 2 : i32, batch_dims = 1 : i32}> : (tensor<2x3x6x5xi16>, tensor<2x4xi64>) -> tensor<2x3x4x5xi16> // This is the same tensor as the one on the CHECK line %expected = arith.constant dense<[ [[[1161, 1162, 1163, 1164, 1165], [1151, 1152, 1153, 1154, 1155], [1141, 1142, 1143, 1144, 1145], [1131, 1132, 1133, 1134, 1135]], [[1261, 1262, 1263, 1264, 1265], [1251, 1252, 1253, 1254, 1255], [1241, 1242, 1243, 1244, 1245], [1231, 1232, 1233, 1234, 1235]], [[1361, 1362, 1363, 1364, 1365], [1351, 1352, 1353, 1354, 1355], [1341, 1342, 1343, 1344, 1345], [1331, 1332, 1333, 1334, 1335]]], [[[2141, 2142, 2143, 2144, 2145], [2131, 2132, 2133, 2134, 2135], [2121, 2122, 2123, 2124, 2125], [2111, 2112, 2113, 2114, 2115]], [[2241, 2242, 2243, 2244, 2245], [2231, 2232, 2233, 2234, 2235], [2221, 2222, 2223, 2224, 2225], [2211, 2212, 2213, 2214, 2215]], [[2341, 2342, 2343, 2344, 2345], [2331, 2332, 2333, 2334, 2335], [2321, 2322, 2323, 2324, 2325], [2311, 2312, 2313, 2314, 2315]]]]> : tensor<2x3x4x5xi16> func.return %gathered, %expected : tensor<2x3x4x5xi16>, tensor<2x3x4x5xi16> // CHECK-NOT: tfl.gather // CHECK: [[CST:%.*]] = arith.constant dense<"0x89048A048B048C048D047F048004810482048304750476047704780479046B046C046D046E046F04ED04EE04EF04F004F104E304E404E504E604E704D904DA04DB04DC04DD04CF04D004D104D204D304510552055305540555054705480549054A054B053D053E053F0540054105330534053505360537055D085E085F08600861085308540855085608570849084A084B084C084D083F084008410842084308C108C208C308C408C508B708B808B908BA08BB08AD08AE08AF08B008B108A308A408A508A608A708250926092709280929091B091C091D091E091F09110912091309140915090709080909090A090B09"> : tensor<2x3x4x5xi16> // If the return value is the same constant twice, the result is the same as expected // CHECK: return [[CST]], [[CST]] } // CHECK-LABEL: func @gather_nd_slices func.func @gather_nd_slices() -> tensor<2x2xi32> { %params = arith.constant dense<[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]> : tensor<2x2x2xi32> %indices = arith.constant dense<[[0, 1], [1, 0]]> : tensor<2x2xi64> %0 = "tfl.gather_nd"(%params, %indices) : (tensor<2x2x2xi32>, tensor<2x2xi64>) -> tensor<2x2xi32> return %0 : tensor<2x2xi32> // CHECK-NOT: tfl.gather_nd // CHECK: [[CST:%.*]] = arith.constant dense<{{\[\[}}3, 4], {{\[}}5, 6]]> : tensor<2x2xi32> // CHECK: return [[CST]] } // CHECK-LABEL: func @gather_nd_scalars func.func @gather_nd_scalars() -> tensor<4xf32> { %params = arith.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]> : tensor<3x3xf32> %indices = arith.constant dense<[[0, 0], [2, 2], [1, 0], [2, 0]]> : tensor<4x2xi32> %0 = "tfl.gather_nd"(%params, %indices) : (tensor<3x3xf32>, tensor<4x2xi32>) -> tensor<4xf32> return %0 : tensor<4xf32> // CHECK-NOT: tfl.gather_nd // CHECK: [[CST:%.+]] = arith.constant dense<[1.000000e+00, 9.000000e+00, 4.000000e+00, 7.000000e+00]> : tensor<4xf32> // CHECK: return [[CST]] } // CHECK-LABEL: reverse_2_dims func.func @reverse_2_dims() -> tensor<2x3x2xi32> { %input = "tfl.pseudo_const"() <{value = dense<[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]> : tensor<2x3x2xi32>}> : () -> tensor<2x3x2xi32> %axis = "tfl.pseudo_const"() <{value = dense<[0, 1]> : tensor<2xi32>}> : () -> tensor<2xi32> %reverse = "tfl.reverse_v2"(%input, %axis) : (tensor<2x3x2xi32>, tensor<2xi32>) -> tensor<2x3x2xi32> return %reverse : tensor<2x3x2xi32> } // CHECK-LITERAL: %cst = arith.constant dense<[[[11, 12], [9, 10], [7, 8]], [[5, 6], [3, 4], [1, 2]]]> : tensor<2x3x2xi32> // CHECK: return %cst : tensor<2x3x2xi32> // CHECK-LABEL: reverse_1_dim func.func @reverse_1_dim() -> tensor<2x3x2xi32> { %input = "tfl.pseudo_const"() <{value = dense<[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]> : tensor<2x3x2xi32>}> : () -> tensor<2x3x2xi32> %axis = "tfl.pseudo_const"() <{value = dense<[2]> : tensor<1xi32>}> : () -> tensor<1xi32> %reverse = "tfl.reverse_v2"(%input, %axis) : (tensor<2x3x2xi32>, tensor<1xi32>) -> tensor<2x3x2xi32> return %reverse : tensor<2x3x2xi32> } // CHECK-LITERAL: %cst = arith.constant dense<[[[2, 1], [4, 3], [6, 5]], [[8, 7], [10, 9], [12, 11]]]> : tensor<2x3x2xi32> // CHECK: return %cst : tensor<2x3x2xi32> // CHECK-LABEL: @slice_no_op func.func @slice_no_op() -> tensor<2x3x2xi32> { %cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32> %cst_2 = arith.constant dense<0> : tensor<3xi32> %cst_3 = arith.constant dense<[2, 3, 2]> : tensor<3xi32> %0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<2x3x2xi32> return %0 : tensor<2x3x2xi32> } // CHECK-LITERAL: %cst = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32> // CHECK-NOT: slice // CHECK-LABEL: @slice_some_dims func.func @slice_some_dims() -> tensor<2x2x1xi32> { %cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32> %cst_2 = arith.constant dense<[0, 1, 1]> : tensor<3xi32> %cst_3 = arith.constant dense<[2, 2, 1]> : tensor<3xi32> %0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<2x2x1xi32> return %0 : tensor<2x2x1xi32> } // CHECK-LITERAL: %cst = arith.constant dense<[[[3], [5]], [[9], [11]]]> : tensor<2x2x1xi32> // CHECK-NOT: slice // CHECK-LABEL: @slice_all_dims func.func @slice_all_dims() -> tensor<1x2x1xi32> { %cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32> %cst_2 = arith.constant dense<[1, 1, 1]> : tensor<3xi32> %cst_3 = arith.constant dense<[1, 2, 1]> : tensor<3xi32> %0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x2x1xi32> return %0 : tensor<1x2x1xi32> } // CHECK-LITERAL: %cst = arith.constant dense<[[[9], [11]]]> : tensor<1x2x1xi32> // CHECK-NOT: slice // CHECK-LABEL: @slice_some_dims_i64 func.func @slice_some_dims_i64() -> tensor<2x2x1xi32> { %cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32> %cst_2 = arith.constant dense<[0, 1, 1]> : tensor<3xi64> %cst_3 = arith.constant dense<[2, 2, 1]> : tensor<3xi64> %0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<2x2x1xi32> return %0 : tensor<2x2x1xi32> } // CHECK-LITERAL: %cst = arith.constant dense<[[[3], [5]], [[9], [11]]]> : tensor<2x2x1xi32> // CHECK-NOT: slice // CHECK-LABEL: @slice_big_float func.func @slice_big_float() -> tensor<1x1x1792x256xf32> { %cst_1 = arith.constant dense<9.000000e+00> : tensor<2x1x1792x256xf32> %cst_2 = arith.constant dense<[1, 0, 0, 0]> : tensor<4xi32> %cst_3 = arith.constant dense<[1, 1, 1792, 256]> : tensor<4xi32> %0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x1x1792x256xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x1x1792x256xf32> return %0 : tensor<1x1x1792x256xf32> } // CHECK-LITERAL: %cst = arith.constant dense<9.000000e+00> : tensor<1x1x1792x256xf32> // CHECK-NOT: slice