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/tfr/tests/canonicalize.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: tfr-opt %s -canonicalize="test-convergence" -verify-diagnostics -split-input-file | FileCheck %s
// Tests for ops with canonicalization patterns.
// CHECK-LABEL: get_real_shape
func.func @get_real_shape(%arg0: tensor<1x2xf32>) -> tensor<2xindex> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2xf32>) -> !tfr.tensor
%1 = tfr.get_shape %0 -> !shape.shape
%2 = shape.to_extent_tensor %1 : !shape.shape -> tensor<2xindex>
func.return %2 : tensor<2xindex>
// CHECK-NEXT: %[[e:.*]] = shape.const_shape [1, 2] : tensor<2xindex>
// CHECK-NEXT: return %[[e]] : tensor<2xindex>
}
// CHECK-LABEL: equal
func.func @equal() -> (i1, i1, i1, i1) {
%0 = tfr.constant f32 -> !tfr.attr
%1 = tfr.constant f32 -> !tfr.attr
%2 = tfr.constant i32 -> !tfr.attr
%same_type = tfr.equal %0,%1 -> i1
%diff_type = tfr.equal %0,%2 -> i1
%3 = tfr.constant "hello" -> !tfr.attr
%4 = tfr.constant "hello" -> !tfr.attr
%5 = tfr.constant "how are you" -> !tfr.attr
%same_str = tfr.equal %3,%4 -> i1
%diff_str = tfr.equal %3,%5 -> i1
func.return %same_type, %diff_type, %same_str, %diff_str : i1, i1, i1, i1
// CHECK-DAG: %true = arith.constant true
// CHECK-DAG: %false = arith.constant false
// CHECK-NEXT: return %true, %false, %true, %false : i1, i1, i1, i1
}
// -----
// CHECK-LABEL: constant_tensor_array
func.func @constant_tensor_array() -> !tfr.tensor {
%0 = tfr.constant [1, -1, 3] -> !tfr.attr
%1 = "tfr.constant_tensor"(%0) : (!tfr.attr) -> !tfr.tensor
func.return %1 : !tfr.tensor
// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() <{value = dense<[1, -1, 3]> : tensor<3xi64>}> : () -> tensor<3xi64>
// CHECK-NEXT: "tfr.cast"(%[[RES]]) : (tensor<3xi64>) -> !tfr.tensor
// CHECK-NEXT: return
}
// -----
// CHECK-LABEL: constant_tensor_scalar
func.func @constant_tensor_scalar() -> !tfr.tensor {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "tfr.constant_tensor"(%0) : (i32) -> !tfr.tensor
func.return %1 : !tfr.tensor
// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() <{value = dense<42> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: "tfr.cast"(%[[RES]]) : (tensor<i32>) -> !tfr.tensor
// CHECK-NEXT: return
}
// -----
// CHECK-LABEL: quant_raw_data
func.func @quant_raw_data(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:1>>) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:2>> {
%0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:1>>) -> !tfr.tensor
%1 = tfr.quant_raw_data(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = tfr.call @tf__risc(%1) : (!tfr.tensor) -> !tfr.tensor
%3 = "tfr.cast"(%2) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 0.3:3>>
%4 = "tfr.cast"(%3) : (tensor<1x10x!quant.uniform<i8:f32, 0.3:3>>) -> !tfr.tensor
%5 = tfr.quant_raw_data(%4) : (!tfr.tensor) -> !tfr.tensor
%6 = tfr.call @tf__risc(%5) : (!tfr.tensor) -> !tfr.tensor
%7 = "tfr.cast"(%6) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:2>>
func.return %7 : tensor<1x10x!quant.uniform<i8:f32, 0.2:2>>
// CHECK: %[[c0:.*]] = "tfr.cast"(%[[arg0:.*]]) : (tensor<1x10x!quant.uniform<i8:f32, 1.000000e-01:1>>) -> !tfr.tensor
// CHECK: %[[c1:.*]] = tfr.call @tf__risc(%[[c0]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK: %[[c2:.*]] = tfr.call @tf__risc(%[[c1]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK: %[[c3:.*]] = "tfr.cast"(%[[c2]]) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 2.000000e-01:2>>
// CHECK: return %[[c3]]
}
// -----
// CHECK-LABEL: quant_raw_data_with_list
func.func @quant_raw_data_with_list(%arg0: !tfr.tensor, %arg1: !tfr.tensor) -> !tfr.tensor {
%cst_1 = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
%1 = "tfr.cast"(%arg0) : (!tfr.tensor) -> tensor<1x4x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>
%2 = "tfr.cast"(%arg1) : (!tfr.tensor) -> tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>
%3 = "tfr.cast"(%2) : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>) -> !tfr.tensor
%4 = "tfr.cast"(%1) : (tensor<1x4x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>) -> !tfr.tensor
%5 = "tfr.build_list"(%3, %4) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%6 = tfr.quant_raw_data(%5) : (!tfr.tensor_list) -> !tfr.tensor_list
%7 = "tfr.cast"(%cst_1) : (tensor<i64>) -> !tfr.tensor
%8 = tfr.call @tf__concat(%7, %6) : (!tfr.tensor, !tfr.tensor_list) -> !tfr.tensor
func.return %8 : !tfr.tensor
// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
// CHECK: %[[BUILD_LIST_0:.*]] = "tfr.build_list"(%arg1, %arg0) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
// CHECK: %[[CAST_0:.*]] = "tfr.cast"(%[[CONST_0]]) : (tensor<i64>) -> !tfr.tensor
// CHECK: %[[CONCAT_O:.*]] = tfr.call @tf__concat(%[[CAST_0]], %[[BUILD_LIST_0]]) : (!tfr.tensor, !tfr.tensor_list) -> !tfr.tensor
// CHECK: return %[[CONCAT_O]] : !tfr.tensor
}
// -----
// CHECK-LABEL: cast_with_unranked_quant
func.func @cast_with_unranked_quant(%arg0: tensor<*xi8>, %arg1: tensor<*xi8>) -> tensor<*xf32> {
%0 = "tf.MaximumFloat"(%arg0, %arg1) : (tensor<*xi8>, tensor<*xi8>) -> tensor<*xi8>
%1 = "tfr.cast"(%0) : (tensor<*xi8>) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>
%3 = "tf.DequantizeFloat"(%2) : (tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>) -> tensor<*xf32>
func.return %3 : tensor<*xf32>
// The cast ops should not be removed in this case or it will result in an
// invalid DequantizeFloat op as following:
// %0 = "tf.MaximumFloat"(%arg0, %arg1) : (tensor<*xi8>, tensor<*xi8>) -> tensor<*xi8>
// %1 = "tf.DequantizeFloat"(%0) : (tensor<*xi8>) -> tensor<*xf32>
// CHECK: %[[MAXIMUMFLOAT_0:.*]] = "tf.MaximumFloat"(%arg0, %arg1) : (tensor<*xi8>, tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[CAST_0:.*]] = "tfr.cast"(%[[MAXIMUMFLOAT_0]]) : (tensor<*xi8>) -> !tfr.tensor
// CHECK: %[[CAST_1:.*]] = "tfr.cast"(%[[CAST_0]]) : (!tfr.tensor) -> tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>
// CHECK: %[[DEQUANTIZEFLOAT_0:.*]] = "tf.DequantizeFloat"(%[[CAST_1]]) : (tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>) -> tensor<*xf32>
// CHECK: return %[[DEQUANTIZEFLOAT_0]] : tensor<*xf32>
}
// -----
// CHECK-LABEL: quant_qparam
func.func @quant_qparam(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:42>>) -> (tensor<f32>, tensor<i32>) {
%0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:42>>) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
%1 = "tfr.cast"(%scale) : (!tfr.tensor) -> tensor<f32>
%2 = "tfr.cast"(%zp) : (!tfr.tensor) -> tensor<i32>
func.return %1, %2 : tensor<f32>, tensor<i32>
// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<1.000000e-01> : tensor<f32>}>
// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<42> : tensor<i32>}> : () -> tensor<i32>
// CHECK: return %[[scale]], %[[zp]]
}
// CHECK-LABEL: quant_qparam_per_channel
func.func @quant_qparam_per_channel(%arg0: tensor<1x3x!quant.uniform<i8:f32:1, {0.1:1, 0.2:2, 0.3:3}>>) -> (tensor<3xf32>, tensor<3xi32>) {
%0 = "tfr.cast"(%arg0) : (tensor<1x3x!quant.uniform<i8:f32:1, {0.1:1, 0.2:2, 0.3:3}>>) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
%1 = "tfr.cast"(%scale) : (!tfr.tensor) -> tensor<3xf32>
%2 = "tfr.cast"(%zp) : (!tfr.tensor) -> tensor<3xi32>
func.return %1, %2 : tensor<3xf32>, tensor<3xi32>
// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<[1.000000e-01, 2.000000e-01, 3.000000e-01]> : tensor<3xf32>}>
// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK: return %[[scale]], %[[zp]]
}
// CHECK-LABEL: quant_qparam_invalid
func.func @quant_qparam_invalid(%arg0: tensor<1x3x!quant.calibrated<f32<-1.0:1.0>>>) -> (!tfr.tensor, !tfr.tensor) {
%0 = "tfr.cast"(%arg0) : (tensor<1x3x!quant.calibrated<f32<-1.0:1.0>>>) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
func.return %scale, %zp: !tfr.tensor, !tfr.tensor
// CHECK: %[[scale:.*]], %[[zp:.*]] = tfr.quant_qparam(%[[input:.*]]) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
// CHECK: return %[[scale]], %[[zp]]
}
// -----
// CHECK-LABEL: redundant_cast_with_different_element_type
func.func @redundant_cast_with_different_element_type(%arg0: tensor<*xf32>) -> (tensor<*xi32>, tensor<2xi32>) {
%0 = "tfr.cast"(%arg0) : (tensor<*xf32>) -> !tfr.tensor
%1 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<*xi32>
%2 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<2xi32>
func.return %1, %2 : tensor<*xi32>, tensor<2xi32>
// CHECK: %[[tf_cast_unranked:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi32>
// CHECK: %[[ensure_shape:.*]] = "tf.EnsureShape"(%arg0) <{shape = #tf_type.shape<2>}> : (tensor<*xf32>) -> tensor<2xf32>
// CHECK: %[[tf_cast_ranked:.*]] = "tf.Cast"(%[[ensure_shape]]) <{Truncate = false}> : (tensor<2xf32>) -> tensor<2xi32>
// CHECK: return %[[tf_cast_unranked]], %[[tf_cast_ranked]] : tensor<*xi32>, tensor<2xi32>
}
// -----
// CHECK-LABEL: redundant_cast_with_quant_type
func.func @redundant_cast_with_quant_type(%arg0: tensor<10x!quant.uniform<i8:f32, 0.0039133410900831223:-128>>) -> (tensor<10xi32>) {
%0 = "tfr.cast"(%arg0) : (tensor<10x!quant.uniform<i8:f32, 0.0039133410900831223:-128>>) -> !tfr.tensor
%1 = tfr.quant_raw_data(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<10xi8>
%3 = "tf.Cast"(%2) {Truncate = false} : (tensor<10xi8>) -> tensor<10xi32>
func.return %3 : tensor<10xi32>
// CHECK: %[[CAST_0:.*]] = "tfr.cast"(%arg0) : (tensor<10x!quant.uniform<i8:f32, 0.0039133410900831223:-128>>) -> !tfr.tensor
// CHECK: %[[CAST_1:.*]] = "tfr.cast"(%[[CAST_0]]) : (!tfr.tensor) -> tensor<10xi8>
// CHECK: %[[CAST_2:.*]] = "tf.Cast"(%[[CAST_1]]) <{Truncate = false}> : (tensor<10xi8>) -> tensor<10xi32>
// CHECK: return %[[CAST_2]] : tensor<10xi32>
}
// -----
// CHECK-LABEL: build_const_list
func.func @build_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
func.return %2 : !tfr.attr
// CHECK-NEXT: %[[c:.*]] = tfr.constant [42 : i32, 41 : i32] -> !tfr.attr
// CHECK-NEXT: return %[[c]] : !tfr.attr
}
// -----
// CHECK-LABEL: build_high_dim_const_list
func.func @build_high_dim_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%3 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%4 = "tfr.build_list"(%2, %3) : (!tfr.attr, !tfr.attr) -> !tfr.attr
func.return %4 : !tfr.attr
// CHECK-NEXT: %[[c:.*]] = tfr.constant {{\[}}[42 : i32, 41 : i32], [42 : i32, 41 : i32]] -> !tfr.attr
// CHECK-NEXT: return %[[c]] : !tfr.attr
}
// -----
// CHECK-LABEL: get_length
func.func @get_length(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> index {
%0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
%1 = "tfr.get_length"(%0) : (!tfr.tensor_list) -> index
func.return %1 : index
// CHECK-NEXT: %[[c:.*]] = arith.constant 2 : index
// CHECK-NEXT: return %[[c]] : index
}
tensorflow/compiler/mlir/tfr/tests/control_flow.mlir
+71
View File
@@ -0,0 +1,71 @@
// 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: tfr-opt %s -tfr-decompose -verify-diagnostics -split-input-file | FileCheck %s
tfr.func @tf__my_pack(%values: !tfr.tensor_list,
%n: i32 {tfr.name="N"},
%axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i32
%eq = arith.cmpi eq, %n, %cst : i32
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp = tfr.call @tf__expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %temp : !tfr.tensor
} else {
%step = arith.index_cast %cst : i32 to index
%end = arith.index_cast %n : i32 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%temp) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp1 = tfr.call @tf__expand_dims(%v, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_concat(%reduce_iter, %temp1, %axis) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
// CHECK-LABEL: pack_one
func.func @pack_one(%arg0: tensor<2x3xf32>) -> tensor<1x2x3xf32> {
%0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<1x2x3xf32>
func.return %0 : tensor<1x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = arith.constant 0 : i32
// CHECK-NEXT: %[[CAST:.*]] = "tfr.cast"(%arg0) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[ED:.*]] = tfr.call @tf__expand_dims(%[[CAST]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[BACK:.*]] = "tfr.cast"(%[[ED]]) : (!tfr.tensor) -> tensor<1x2x3xf32>
// CHECK-NEXT: return %[[BACK]] : tensor<1x2x3xf32>
}
// CHECK-LABEL: pack_multiple
func.func @pack_multiple(%arg0: tensor<2x3xf32>,
%arg1: tensor<2x3xf32>,
%arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = arith.constant 0 : i32
// CHECK-NEXT: %[[CAST0:.*]] = "tfr.cast"(%arg0) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[CAST1:.*]] = "tfr.cast"(%arg1) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[CAST2:.*]] = "tfr.cast"(%arg2) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[EX0:.*]] = tfr.call @tf__expand_dims(%[[CAST0]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[EX1:.*]] = tfr.call @tf__expand_dims(%[[CAST1]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[CONCAT1:.*]] = tfr.call @tf__risc_concat(%[[EX0]], %[[EX1]], %c0_i32) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[EX2:.*]] = tfr.call @tf__expand_dims(%[[CAST2]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[CONCAT2:.*]] = tfr.call @tf__risc_concat(%[[CONCAT1]], %[[EX2]], %[[AXIS]]) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[BACK:.*]] = "tfr.cast"(%[[CONCAT2]]) : (!tfr.tensor) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[BACK]] : tensor<3x2x3xf32>
}
tensorflow/compiler/mlir/tfr/tests/decompose.mlir
+303
View File
@@ -0,0 +1,303 @@
// 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: tfr-opt %s -tfr-decompose -verify-diagnostics | FileCheck %s
// Definitions for ops that are being used in the tests.
// ex) tf.MyOp refers to tfr.func @tf__my_op
// CHECK-LABEL: @tf__fake_no_op
tfr.func @tf__fake_no_op(%arg0: !tfr.tensor) -> !tfr.tensor {
tfr.return %arg0 : !tfr.tensor
// CHECK-NEXT: tfr.return %arg0 : !tfr.tensor
}
// CHECK-LABEL: @tf__intermediate
tfr.func @tf__intermediate(%arg0: !tfr.tensor) -> !tfr.tensor {
%0 = tfr.call @tf__risc(%arg0) : (!tfr.tensor) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%arg0) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: tfr.return %[[id]] : !tfr.tensor
}
// CHECK-LABEL: @tf__fused_n
tfr.func @tf__fused_n(
%arg0: !tfr.tensor,
%arg1: !tfr.tensor_list,
%arg2: index {tfr.name="A",tfr.default=1:index})
-> !tfr.tensor_list {
%0 = tfr.call @tf__intermediate(%arg0) : (!tfr.tensor) -> !tfr.tensor
%1 = tfr.get_element %arg1[%arg2] : (!tfr.tensor_list, index) -> !tfr.tensor
%2 = tfr.call @tf__intermediate(%1) : (!tfr.tensor) -> !tfr.tensor
%3 = "tfr.build_list"(%0, %2) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
tfr.return %3 : !tfr.tensor_list
// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__intermediate(%arg0) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[ge:.*]] = tfr.get_element %arg1[%arg2] : (!tfr.tensor_list, index) -> !tfr.tensor
// CHECK-NEXT: %[[id2:.*]] = tfr.call @tf__intermediate(%[[ge]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[bl:.*]] = "tfr.build_list"(%[[id1]], %[[id2]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
// CHECK-NEXT: tfr.return %[[bl]] : !tfr.tensor_list
}
// CHECK-LABEL: @tf__my_max_pool
tfr.func @tf__my_max_pool(%input_: !tfr.tensor, %stride_w: i64{tfr.name="stride_w"}, %stride_h: i64{tfr.name="stride_h"}) -> (!tfr.tensor) {
%cst_1 = arith.constant 1 : i64
%stride = "tfr.build_list"(%cst_1, %stride_w, %stride_h, %cst_1) : (i64, i64, i64, i64) -> !tfr.attr
%filter = tfr.constant [1, 2, 2, 1] -> !tfr.attr
%padding = tfr.constant "VALID" -> !tfr.attr
%explicit_paddings = tfr.constant [] -> !tfr.attr
%data_format = tfr.constant "NHWC" -> !tfr.attr
%MaxPool = tfr.call @tf__max_pool(%input_, %stride, %filter, %padding, %explicit_paddings, %data_format) : (!tfr.tensor, !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr) -> (!tfr.tensor)
tfr.return %MaxPool : !tfr.tensor
// CHECK: tf__max_pool
}
// CHECK-LABEL: @tf__cast_float
tfr.func @tf__cast_float(%input_: !tfr.tensor, %out_type: !tfr.attr{tfr.name="out_type"}) -> (!tfr.tensor) {
%false = arith.constant false
%cast = tfr.call @tf__cast(%input_, %out_type, %false) : (!tfr.tensor, !tfr.attr, i1) -> (!tfr.tensor)
tfr.return %cast : !tfr.tensor
}
// end op definitions
// CHECK-LABEL: decompose_tf_no_op
func.func @decompose_tf_no_op(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tf.FakeNoOp"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
func.return %0 : tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: return %arg0
}
// CHECK-LABEL: decompose_tf_intermediate
func.func @decompose_tf_intermediate(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tf.Intermediate"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
func.return %0 : tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: %[[casted:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%[[casted]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id]]) : (!tfr.tensor) -> tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: return %[[back]]
}
// CHECK-LABEL: decompose_fused_n_default
func.func @decompose_fused_n_default(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[in2:.*]] = "tfr.cast"(%arg2) : (tensor<f32>) -> !tfr.tensor
// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[id2:.*]] = tfr.call @tf__risc(%[[in2]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id2]]) : (!tfr.tensor) -> tensor<f32>
// CHECK-NEXT: return %[[back]] : tensor<f32>
}
// CHECK-LABEL: decompose_fused_n
func.func @decompose_fused_n(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index} : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[in1:.*]] = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__risc(%[[in1]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id1]]) : (!tfr.tensor) -> tensor<f32>
// CHECK-NEXT: return %[[back]] : tensor<f32>
}
// CHECK-LABEL: attribute_propagate_direct
func.func @attribute_propagate_direct(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tf.Intermediate"(%arg0) {_tpu_replicate, device="hello"} : (tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
func.return %0 : tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: %[[casted:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%[[casted]]) {_tpu_replicate, device = "hello"}
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id]]) : (!tfr.tensor) -> tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: return %[[back]]
}
// CHECK-LABEL: attribute_propagate
func.func @attribute_propagate(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index, _tpu_replicate, device="hello"} : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[in1:.*]] = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) {_tpu_replicate, device = "hello"}
// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__risc(%[[in1]]) {_tpu_replicate, device = "hello"}
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id1]]) : (!tfr.tensor) -> tensor<f32>
// CHECK-NEXT: return %[[back]] : tensor<f32>
}
// CHECK: attribute_cast
func.func @attribute_cast(%arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
%0 = "tfr.cast"(%arg0) : (tensor<1x4x4x1xf32>) -> !tfr.tensor
%stride_i32 = arith.constant 2 : i32
%1 = tfr.call @tf__my_max_pool(%0, %stride_i32, %stride_i32) : (!tfr.tensor, i32, i32) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<1x2x2x1xf32>
func.return %2 : tensor<1x2x2x1xf32>
// CHECK: tf__max_pool
}
// CHECK-LABEL: no_tf_canonicalization
func.func @no_tf_canonicalization(%arg0: tensor<8xi1>, %arg1: tensor<8x3xf32>, %arg2: tensor<8x3xf32>) -> tensor<8x3xf32> {
%0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<8xi1>, tensor<8x3xf32>, tensor<8x3xf32>) -> tensor<8x3xf32>
func.return %0: tensor<8x3xf32>
// CHECK: "tf.Select"
}
// CHECK-LABEL: denied_attribute
func.func @denied_attribute(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// expected-error@+1 {{Denied unregistered attribute was found: denied_attr}}
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index, denied_attr} : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: "tf.FusedN"(%arg0, %arg1, %arg2) {A = 0 : index, denied_attr}
}
// CHECK-LABEL: quantized_tensor
func.func @quantized_tensor(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.0038396485615521669:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>> {
%0 = "tf.Intermediate"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.0038396485615521669:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>>
func.return %0 : tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>>
// CHECK: "tfr.cast"(%[[arg0:.*]]) : (tensor<1x10x!quant.uniform<i8:f32, 0.0038396485615521669:-128>>) -> !tfr.tensor
// CHECK: "tfr.cast"(%[[result:.*]]) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>>
}
// CHECK-LABEL: decompose_quant_act_range
func.func @decompose_quant_act_range() -> !tfr.tensor_list {
%scale = arith.constant 0.1 : f32
%zp = arith.constant 42 : i64
%none_attr = tfr.constant "NONE" -> !tfr.attr
%relu_attr = tfr.constant "RELU" -> !tfr.attr
%relu6_attr = tfr.constant "RELU6" -> !tfr.attr
%reluN1_1_attr = tfr.constant "RELU_N1_TO_1" -> !tfr.attr
%none:2 = "tfr.quant_act_range"(%none_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%relu:2 = "tfr.quant_act_range"(%relu_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%relu6:2 = "tfr.quant_act_range"(%relu6_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%reluN1_1:2 = "tfr.quant_act_range"(%reluN1_1_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%result = "tfr.build_list"(
%none#0, %none#1, %relu#0, %relu#1,
%relu6#0, %relu6#1, %reluN1_1#0, %reluN1_1#1) : (
!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor,
!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
func.return %result : !tfr.tensor_list
// CHECK-DAG: %[[N_128:.*]] = arith.constant -128 : i32
// CHECK-DAG: %[[N32:.*]] = arith.constant 32 : i32
// CHECK-DAG: %[[N42:.*]] = arith.constant 42 : i32
// CHECK-DAG: %[[N52:.*]] = arith.constant 52 : i32
// CHECK-DAG: %[[N102:.*]] = arith.constant 102 : i32
// CHECK-DAG: %[[N127:.*]] = arith.constant 127 : i32
// CHECK-NEXT: %[[none_min:.*]] = "tfr.constant_tensor"(%[[N_128]])
// CHECK-NEXT: %[[none_max:.*]] = "tfr.constant_tensor"(%[[N127]])
// CHECK-NEXT: %[[relu_min:.*]] = "tfr.constant_tensor"(%[[N42]])
// CHECK-NEXT: %[[relu_max:.*]] = "tfr.constant_tensor"(%[[N127]])
// CHECK-NEXT: %[[relu6_min:.*]] = "tfr.constant_tensor"(%[[N42]])
// CHECK-NEXT: %[[relu6_max:.*]] = "tfr.constant_tensor"(%[[N102]])
// CHECK-NEXT: %[[reluN1_1_min:.*]] = "tfr.constant_tensor"(%[[N32]])
// CHECK-NEXT: %[[reluN1_1_max:.*]] = "tfr.constant_tensor"(%[[N52]])
// CHECK-NEXT: %[[result:.*]] = "tfr.build_list"(%[[none_min]], %[[none_max]], %[[relu_min]], %[[relu_max]],
// CHECK-SAME: %[[relu6_min]], %[[relu6_max]], %[[reluN1_1_min]], %[[reluN1_1_max]]
// CHECK-NEXT: return %[[result]]
}
// CHECK-LABEL: decompose_quant_act_range_invalid
func.func @decompose_quant_act_range_invalid() -> (!tfr.tensor, !tfr.tensor) {
%scale = arith.constant 0.1 : f32
%zp = arith.constant 42 : i64
%elu_attr = tfr.constant "ELU" -> !tfr.attr
%min, %max = "tfr.quant_act_range"(%elu_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
func.return %min, %max : !tfr.tensor, !tfr.tensor
// CHECK: %[[elu_attr:.*]] = tfr.constant "ELU" -> !tfr.attr
// CHECK: %[[min:.*]], %[[max:.*]] = tfr.quant_act_range(%[[elu_attr]]
// CHECK: return %[[min]], %[[max]]
}
// CHECK-LABEL: decompose_quant_scale_factor
func.func @decompose_quant_scale_factor() -> (!tfr.tensor, !tfr.tensor) {
%output_scale = arith.constant 0.1 : f32
%input_scale = arith.constant 0.25 : f32
%filter_scale = arith.constant 0.4 : f32
%input_scale_tensor = "tfr.constant_tensor"(%input_scale) : (f32) -> !tfr.tensor
%filter_scale_tensor = "tfr.constant_tensor"(%filter_scale) : (f32) -> !tfr.tensor
%list = "tfr.build_list"(%input_scale_tensor, %filter_scale_tensor) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%out = "tfr.quant_scale_factor"(%output_scale, %list) : (f32, !tfr.tensor_list) -> !tfr.tensor
%perchannel_scale = arith.constant dense<[0.4, 4.0]> : tensor<2xf32>
%perchannel_scale_tensor = "tfr.cast"(%perchannel_scale) : (tensor<2xf32>) -> !tfr.tensor
%list2 = "tfr.build_list"(%input_scale_tensor, %perchannel_scale_tensor) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%perchannel = "tfr.quant_scale_factor"(%output_scale, %list2) : (f32, !tfr.tensor_list) -> !tfr.tensor
func.return %out, %perchannel : !tfr.tensor, !tfr.tensor
// CHECK-DAG: %[[scale_factors:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 1.000000e+01]> : tensor<2xf32>}> : () -> tensor<2xf32>
// CHECK-DAG: %[[scale_factor:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK: %[[cast:.*]] = "tfr.cast"(%[[scale_factor]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[cast_perchannel:.*]] = "tfr.cast"(%[[scale_factors]]) : (tensor<2xf32>) -> !tfr.tensor
// CHECK: return %[[cast]], %[[cast_perchannel]] : !tfr.tensor, !tfr.tensor
}
// CHECK-LABEL: decompose_quant_scale_factor_invalid
func.func @decompose_quant_scale_factor_invalid() -> !tfr.tensor {
%output_scale = arith.constant 0.1 : f32
%input_scale = arith.constant 0.25 : f32
%filter_scale = arith.constant 0.4 : f32
%input_scale_tensor = "tfr.constant_tensor"(%input_scale) : (f32) -> !tfr.tensor
%filter_scale_tensor = "tfr.constant_tensor"(%filter_scale) : (f32) -> !tfr.tensor
%list = "tfr.build_list"(%input_scale_tensor, %filter_scale_tensor, %input_scale_tensor) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%out = "tfr.quant_scale_factor"(%output_scale, %list) : (f32, !tfr.tensor_list) -> !tfr.tensor
func.return %out : !tfr.tensor
// CHECK-DAG: %[[cst_0:.*]] = arith.constant 1.000000e-01 : f32
// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<2.500000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK: %[[tfrcast0:.*]] = "tfr.cast"(%[[cst_1]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[tfrcast1:.*]] = "tfr.cast"(%[[cst_2]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[list:.*]] = "tfr.build_list"(%[[tfrcast0]], %[[tfrcast1]], %[[tfrcast0]]) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
// CHECK: %[[qsf:.*]] = tfr.quant_scale_factor(%[[cst_0]], %[[list]]) : (f32, !tfr.tensor_list) -> !tfr.tensor
// CHECK: return %[[qsf]] : !tfr.tensor
}
// CHECK-LABEL: decompose_quant_rescale
func.func @decompose_quant_rescale(%arg0: tensor<2xi32>) -> !tfr.tensor {
%zp = arith.constant 67 : i64
%cst = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
%scale_factor = "tfr.cast"(%cst) : (tensor<f32>) -> !tfr.tensor
%input = "tfr.cast"(%arg0) : (tensor<2xi32>) -> !tfr.tensor
%rescaled = "tfr.quant_rescale"(%input, %scale_factor, %zp) : (!tfr.tensor, !tfr.tensor, i64) -> !tfr.tensor
func.return %rescaled : !tfr.tensor
// CHECK-DAG: %[[f32:.*]] = tfr.constant f32 -> !tfr.attr
// CHECK-DAG: %[[i32:.*]] = tfr.constant i32 -> !tfr.attr
// CHECK-DAG: %[[scale_cst:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK-DAG: %false = arith.constant false
// CHECK-DAG: %[[zp_cst:.*]] = "tf.Const"() <{value = dense<67> : tensor<i64>}> : () -> tensor<i64>
// CHECK: %[[zp:.*]] = "tfr.cast"(%[[zp_cst]]) : (tensor<i64>) -> !tfr.tensor
// CHECK: %[[scale:.*]] = "tfr.cast"(%[[scale_cst]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[input:.*]] = "tfr.cast"(%arg0) : (tensor<2xi32>) -> !tfr.tensor
// CHECK: %[[cast:.*]] = tfr.call @tf__cast(%[[input]], %[[f32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
// CHECK: %[[rescaled:.*]] = tfr.call @tf__mul(%[[cast]], %[[scale]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
// CHECK: %[[rounded:.*]] = tfr.call @tf__round(%[[rescaled]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK: %[[zp_cast:.*]] = tfr.call @tf__cast(%[[zp]], %[[f32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
// CHECK: %[[recentered:.*]] = tfr.call @tf__add(%[[rounded]], %[[zp_cast]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
// CHECK: %[[cast_i32:.*]] = tfr.call @tf__cast(%[[recentered]], %[[i32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
// CHECK: return %[[cast_i32]] : !tfr.tensor
}
// CHECK-LABEL: decompose_output_type
func.func @decompose_output_type(%arg0: tensor<2xf32>) -> tensor<2xi32> {
%0 = "tf.CastFloat"(%arg0) : (tensor<2xf32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
// CHECK: %[[i32:.*]] = tfr.constant i32 -> !tfr.attr
// CHECK: tfr.call @tf__cast(%[[casted_arg:.*]], %[[i32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
}
tensorflow/compiler/mlir/tfr/tests/end2end.mlir
+249
View File
@@ -0,0 +1,249 @@
// 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: tfr-opt %s -tfr-decompose -tfr-raise-to-tf -canonicalize -verify-diagnostics -split-input-file | FileCheck %s
//=================> User models, from GraphDef <====================
// CHECK-LABEL: my_identity
func.func @my_identity(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = "tf.MyIdentity"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
func.return %0 : tensor<2x3xf32>
// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}
// CHECK-LABEL: my_rsqrt
func.func @my_rsqrt(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyRsqrt"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[RE:.*]] = "tf.RiscReciprocal"(%arg0) : (tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[SQRT:.*]] = "tf.RiscSqrt"(%[[RE]]) : (tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[SQRT]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_leaky_relu
func.func @my_leaky_relu(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyLeakyRelu"(%arg0) {alpha=3.0 : f32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_leaky_relu_with_default
func.func @my_leaky_relu_with_default(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyLeakyRelu"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_cast
func.func @my_cast(%arg0: tensor<2x3xf32>) -> tensor<2x3xi32> {
%0 = "tf.MyCast"(%arg0) {Tout=i32} : (tensor<2x3xf32>) -> tensor<2x3xi32>
func.return %0 : tensor<2x3xi32>
// CHECK-NEXT: %[[CAST:.*]] = "tf.RiscCast"(%arg0) {Tout = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CAST]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xi32>) -> tensor<2x3xi32>
// CHECK-NEXT: return %[[ES]] : tensor<2x3xi32>
}
// CHECK-LABEL: my_pack_single_input
func.func @my_pack_single_input(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ED]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_pack_multiple_inputs
func.func @my_pack_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED0:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED1:.*]] = "tf.ExpandDims"(%arg1, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC0:.*]] = "tf.RiscConcat"(%[[ED0]], %[[ED1]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED2:.*]] = "tf.ExpandDims"(%arg2, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC1:.*]] = "tf.RiscConcat"(%[[CC0]], %[[ED2]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CC1]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_add_n_single_input
func.func @my_add_n_single_input(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = "tf.MyAddN"(%arg0) {N=1:i32} : (tensor<2x3xf32>) -> tensor<2x3xf32>
func.return %0 : tensor<2x3xf32>
// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}
// CHECK-LABEL: my_add_n_multiple_inputs
func.func @my_add_n_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = "tf.MyAddN"(%arg0, %arg1, %arg2) {N=3:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
func.return %0 : tensor<2x3xf32>
// CHECK-NEXT: %[[ADD0:.*]] = "tf.RiscAdd"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ADD1:.*]] = "tf.RiscAdd"(%[[ADD0]], %arg2) : (tensor<*xf32>, tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ADD1]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xf32>) -> tensor<2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<2x3xf32>
}
// CHECK-LABEL: my_map_and_batch_dataset
func.func @my_map_and_batch_dataset(%input: tensor<*x!tf_type.variant>,
%other1: tensor<*xf32>,
%other2: tensor<*xi32>) -> tensor<*x!tf_type.variant> {
%0 = "tf.MyMapAndBatchDataset"(%input, %other1, %other2)
{batch_size=1000 : i64, num_parallel_calls = 8 : i64, drop_remainder = 0 : i1,
func = @"__some_func", output_types = [f32], output_shapes = [#tf_type.shape<>], preserve_cardinality = true}
: (tensor<*x!tf_type.variant>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.variant>
func.return %0 : tensor<*x!tf_type.variant>
// CHECK-DAG: %[[BATCH:.*]] = "tf.Const"() <{value = dense<1000> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[PARAL:.*]] = "tf.Const"() <{value = dense<8> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[KEEP:.*]] = "tf.Const"() <{value = dense<false> : tensor<i1>}> : () -> tensor<i1>
// CHECK: %[[CAST:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<*xi32>) -> tensor<*xf32>
// CHECK: %[[RET:.*]] = "tf.MapAndBatchDatasetV0"(%arg0, %[[BATCH]], %[[PARAL]], %[[KEEP]], %arg1, %[[CAST]])
// CHECK-SAME: {f = @__some_func, output_shapes = [#tf_type.shape<>], output_types = [f32], preserve_cardinality = true} : (tensor<*x!tf_type.variant>, tensor<i64>, tensor<i64>, tensor<i1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
// CHECK: return %[[RET]] : tensor<*x!tf_type.variant>
}
//=================> decomposition functions, translated from tf.compose api <====================
tfr.func @tf__my_identity(%value: !tfr.tensor) -> !tfr.tensor {
tfr.return %value : !tfr.tensor
}
tfr.func @tf__my_cast(%value: !tfr.tensor, %tout: !tfr.attr{tfr.name="Tout"}) -> !tfr.tensor {
%0 = tfr.call @tf__risc_cast(%value, %tout) : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
}
tfr.func @tf__my_rsqrt(%value: !tfr.tensor) -> !tfr.tensor {
%1 = tfr.call @tf__risc_reciprocal(%value) : (!tfr.tensor) -> !tfr.tensor
%2 = tfr.call @tf__risc_sqrt(%1) : (!tfr.tensor) -> !tfr.tensor
tfr.return %2 : !tfr.tensor
}
tfr.func @tf__my_leaky_relu(%value: !tfr.tensor, %alpha: f32 {tfr.name="alpha", tfr.default=0.2:f32}) -> !tfr.tensor {
%1 = tfr.call @tf__risc_shape(%value) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.constant_tensor"(%alpha) : (f32) -> tensor<f32>
%t = "tfr.cast"(%2) : (tensor<f32>) -> !tfr.tensor
%3 = tfr.call @tf__risc_broadcast(%t, %1) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
%4 = tfr.call @tf__risc_maximum(%value, %3) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
tfr.return %4 : !tfr.tensor
}
// TODO(fengliuai): use shape dialect to manipulate the shape then this can be decomposed further.
tfr.func @tf__my_expand_dims(%value: !tfr.tensor, %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
%axis_cst = "tfr.constant_tensor"(%axis) : (i32) -> tensor<i32>
%dim = "tfr.cast"(%axis_cst) : (tensor<i32>) -> !tfr.tensor
%0 = tfr.call @tf__expand_dims(%value, %dim) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
}
tfr.func @tf__my_pack(%values: !tfr.tensor_list,
%n: i32 {tfr.name="N"},
%axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i32
%eq = arith.cmpi eq, %n, %cst : i32
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp = tfr.call @tf__my_expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %temp : !tfr.tensor
} else {
%step = arith.index_cast %cst : i32 to index
%end = arith.index_cast %n : i32 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%temp) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp1 = tfr.call @tf__my_expand_dims(%v, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_concat(%reduce_iter, %temp1, %axis) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
%n: i32 {tfr.name="N"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i32
%eq = arith.cmpi eq, %n, %cst : i32
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %v1 : !tfr.tensor
} else {
%step = arith.index_cast %cst : i32 to index
%end = arith.index_cast %n : i32 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_add(%reduce_iter, %v) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_map_and_batch_dataset(
%input_dataset: !tfr.tensor,
%other_arguments: !tfr.tensor_list,
%batch_size: i64 {tfr.name="batch_size"},
%num_parallel_calls: i64 {tfr.name="num_parallel_calls"},
%drop_remainder: i1 {tfr.name="drop_remainder"},
%f: !tfr.attr {tfr.name="func"},
%output_types: !tfr.attr {tfr.name="output_types"},
%output_shapes: !tfr.attr {tfr.name="output_shapes"},
%preserve_cardinality: i1 {tfr.name="preserve_cardinality", tfr.default=false}) -> !tfr.tensor {
%batch = "tfr.constant_tensor"(%batch_size) : (i64) -> tensor<i64>
%batch1 = "tfr.cast"(%batch) : (tensor<i64>) -> !tfr.tensor
%calls = "tfr.constant_tensor"(%num_parallel_calls) : (i64) -> tensor<i64>
%calls1 = "tfr.cast"(%calls) : (tensor<i64>) -> !tfr.tensor
%drop = "tfr.constant_tensor"(%drop_remainder) : (i1) -> tensor<i1>
%drop1 = "tfr.cast"(%drop) : (tensor<i1>) -> !tfr.tensor
%ret = tfr.call @tf__map_and_batch_dataset_v0(%input_dataset, %batch1, %calls1, %drop1, %other_arguments, %f, %output_types, %output_shapes, %preserve_cardinality)
: (!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor_list, !tfr.attr, !tfr.attr, !tfr.attr, i1) -> !tfr.tensor
tfr.return %ret : !tfr.tensor
}
//=================> signatures of the primitive ops with kernels, modeled as external TFR function <==
tfr.func @tf__risc_cast_(!tfr.tensor, !tfr.attr{tfr.name="Tout"}) -> !tfr.tensor<Tout> attributes{Tout}
tfr.func @tf__risc_add_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_concat_(!tfr.tensor<T>, !tfr.tensor<T>, i32{tfr.name="axis"}) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_broadcast_(!tfr.tensor<T>, !tfr.tensor<Tidx>) -> !tfr.tensor<T> attributes{T, Tidx}
tfr.func @tf__risc_reciprocal_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_sqrt_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_shape_(!tfr.tensor, !tfr.attr{tfr.name="T", tfr.default=i32}) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_maximum_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__expand_dims_(!tfr.tensor<T>, !tfr.tensor<Tdim>) -> !tfr.tensor<T> attributes{T, Tdim}
tfr.func @tf__map_and_batch_dataset_v0_(!tfr.tensor<T>, !tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor_list<Targuments>,
!tfr.attr{tfr.name="f"}, !tfr.attr{tfr.name="output_types"}, !tfr.attr{tfr.name="output_shapes"}, i1{tfr.name="preserve_cardinality"})
-> !tfr.tensor<T> attributes{T, Targuments}
tensorflow/compiler/mlir/tfr/tests/ops.mlir
+442
View File
@@ -0,0 +1,442 @@
// 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: tfr-opt %s -verify-diagnostics -split-input-file | tfr-opt | FileCheck %s
// Tests for types, ops with custom constraints, verifiers, printer or parser
// methods.
// CHECK-LABEL: tensor_type_noconstraint
func.func private @tensor_type_noconstraint() -> !tfr.tensor
// -----
// CHECK-LABEL: tensor_type
func.func private @tensor_type() -> !tfr.tensor<T>
// -----
// CHECK-LABEL: tensor_list_type_noconstraint
func.func private @tensor_list_type_noconstraint() -> !tfr.tensor_list
// -----
// CHECK-LABEL: tensor_list_type_array_like
func.func private @tensor_list_type_array_like() -> !tfr.tensor_list<[N, T]>
// -----
// CHECK-LABEL: tensor_list_type_tuple_like
func.func private @tensor_list_type_tuple_like() -> !tfr.tensor_list<input_T>
// -----
// expected-error@+1 {{unbalanced '[' character in pretty dialect name}}
func.func private @tensor_invalid_1() -> !tfr.tensor<[N, T>
// -----
// expected-error@+1 {{unbalanced}}
func.func @tensor_invalid_2() -> !tfr.tensor<[N, T]
// -----
// CHECK-LABEL: call_op
func.func @call_op(%arg0: !tfr.tensor<T>, %arg1: !tfr.tensor_list<TL>, %arg2: i32) -> !tfr.tensor<K> {
%0 = tfr.call @Foo(%arg0, %arg1, %arg2) : (!tfr.tensor<T>, !tfr.tensor_list<TL>, i32) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
// CHECK-LABEL: call_op_arg_attr(%arg0: i32) -> !tfr.tensor<K>
func.func @call_op_arg_attr(%arg0: i32) -> !tfr.tensor<K> {
%0 = tfr.call @Bar(%arg0) : (i32) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
func.func @call_op_invalid_1(%arg0: tensor<?xf32>) -> !tfr.tensor<K> {
// expected-error@+1 {{got 'tensor<?xf32>'}}
%0 = tfr.call @Huu(%arg0) : (tensor<?xf32>) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
// CHECK-LABEL: get_shape
func.func @get_shape(%arg0: !tfr.tensor) -> (!shape.shape, !shape.shape) {
%0 = tfr.get_shape %arg0 -> !shape.shape
%1 = "tfr.get_shape"(%arg0) : (!tfr.tensor) -> !shape.shape
func.return %0, %1 : !shape.shape, !shape.shape
}
// -----
// CHECK-LABEL: get_real_shape
func.func @get_real_shape(%arg0: tensor<1x2xf32>) -> tensor<2xindex> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2xf32>) -> !tfr.tensor
%1 = tfr.get_shape %0 -> !shape.shape
%2 = shape.to_extent_tensor %1 : !shape.shape -> tensor<2xindex>
func.return %2 : tensor<2xindex>
}
// -----
func.func @get_element_type(%arg0: !tfr.tensor) -> (!tfr.attr, !tfr.attr) {
%0 = tfr.get_element_type %arg0 -> !tfr.attr
%1 = "tfr.get_element_type"(%arg0) : (!tfr.tensor) -> !tfr.attr
func.return %0, %1 : !tfr.attr, !tfr.attr
}
// -----
// CHECK-LABEL: from_tf_tensor
func.func @from_tf_tensor(%arg0: tensor<?xf32>) -> !tfr.tensor<K> {
%0 = "tfr.cast"(%arg0) : (tensor<?xf32>) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
// CHECK-LABEL: to_tf_tensor
func.func @to_tf_tensor(%arg0: !tfr.tensor<T>) -> tensor<?xi32> {
%0 = "tfr.cast"(%arg0) : (!tfr.tensor<T>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// -----
// CHECK-LABEL: constant
func.func @constant() -> (!tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr) {
%0 = tfr.constant f32 -> !tfr.attr
%1 = tfr.constant [f32, i32] -> !tfr.attr
%2 = "tfr.constant"() {value = f32} : () -> !tfr.attr
%3 = "tfr.constant"() {value = [f32, i32]} : () -> !tfr.attr
func.return %0, %1, %2, %3 : !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr
}
// -----
// CHECK-LABEL: equal
func.func @equal() -> (i1, i1, i1, i1) {
%0 = tfr.constant f32 -> !tfr.attr
%1 = tfr.constant f32 -> !tfr.attr
%2 = tfr.constant i32 -> !tfr.attr
%same_type = tfr.equal %0,%1 -> i1
%diff_type = tfr.equal %0,%2 -> i1
%3 = tfr.constant "hello" -> !tfr.attr
%4 = tfr.constant "hello" -> !tfr.attr
%5 = tfr.constant "how are you" -> !tfr.attr
%same_str = tfr.equal %3,%4 -> i1
%diff_str = tfr.equal %3,%5 -> i1
func.return %same_type, %diff_type, %same_str, %diff_str : i1, i1, i1, i1
}
// -----
// CHECK-LABEL: constant_tensor_scalar
func.func @constant_tensor_scalar(%arg0: i32) -> tensor<i32> {
%0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: constant_tensor_vector
func.func @constant_tensor_vector(%arg0: vector<1x2xi32>) -> tensor<1x2xi32> {
%0 = "tfr.constant_tensor"(%arg0) : (vector<1x2xi32>) -> tensor<1x2xi32>
func.return %0 : tensor<1x2xi32>
}
// -----
// CHECK-LABEL: constant_tensor_array
func.func @constant_tensor_array() -> !tfr.tensor {
%0 = tfr.constant [1, -1, 3] -> !tfr.attr
%1 = "tfr.constant_tensor"(%0) : (!tfr.attr) -> !tfr.tensor
func.return %1 : !tfr.tensor
}
// -----
// CHECK-LABEL: constant_tensor_scalar
func.func @constant_tensor_scalar() -> !tfr.tensor {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "tfr.constant_tensor"(%0) : (i32) -> !tfr.tensor
func.return %1 : !tfr.tensor
}
// -----
func.func @constant_tensor_invalid_0(%arg0: i32) -> tensor<f32> {
// expected-error@+1 {{input and output should have the same scalar types.}}
%0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
func.func @constant_tensor_invalid_1(%arg0: vector<1xi32>) -> tensor<?xi32> {
// expected-error@+1 {{output type should be static and ranked}}
%0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// -----
func.func @constant_tensor_invalid_2(%arg0: vector<1xi32>) -> tensor<1xf32> {
// expected-error@+1 {{input and output should have same shape and element type}}
%0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
func.func @constant_tensor_invalid_3(%arg0: vector<1xi32>) -> tensor<1x1xi32> {
// expected-error@+1 {{input and output should have same shape and element type}}
%0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1x1xi32>
func.return %0 : tensor<1x1xi32>
}
// -----
func.func @constant_tensor_invalid_4(%arg0: i32) -> tensor<1x1xi32> {
// expected-error@+1 {{input can not be converted to an output tensor}}
%0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<1x1xi32>
func.return %0 : tensor<1x1xi32>
}
// -----
// CHECK-LABEL: get_element
func.func @get_element(%arg0: !tfr.tensor_list<T>) -> !tfr.tensor {
%cst = "arith.constant"() {value = 1 : index} : () -> index
%0 = tfr.get_element %arg0[%cst] : (!tfr.tensor_list<T>, index) -> !tfr.tensor
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: build_list
func.func @build_list(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> !tfr.tensor_list {
%0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
func.return %0 : !tfr.tensor_list
}
// -----
// CHECK-LABEL: quant_act_range
func.func @quant_act_range(%arg0: !tfr.attr, %arg1: f32, %arg2: i64) -> !tfr.tensor {
%0:2 = "tfr.quant_act_range"(%arg0, %arg1, %arg2) : (!tfr.attr,f32,i64) -> (!tfr.tensor,!tfr.tensor)
func.return %0#0 : !tfr.tensor
}
// -----
// CHECK-LABEL: quant_rescale
func.func @quant_rescale(%arg0: !tfr.tensor, %arg1: !tfr.tensor, %arg2: i64) -> !tfr.tensor {
%0 = "tfr.quant_rescale"(%arg0, %arg1, %arg2) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor)
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: quant_raw_data
func.func @quant_raw_data(%arg0: !tfr.tensor) -> !tfr.tensor {
%0 = "tfr.quant_raw_data"(%arg0) : (!tfr.tensor) -> (!tfr.tensor)
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: quant_qparam
func.func @quant_qparam(%arg0: !tfr.tensor) -> (!tfr.tensor, !tfr.tensor) {
%scale, %zp = tfr.quant_qparam(%arg0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
func.return %scale, %zp : !tfr.tensor, !tfr.tensor
}
// -----
// CHECK-LABEL: quant_scale_factor
func.func @quant_scale_factor(%arg0: f32, %arg1: !tfr.tensor_list) -> (!tfr.tensor) {
%0 = "tfr.quant_scale_factor"(%arg0, %arg1) : (f32, !tfr.tensor_list) -> (!tfr.tensor)
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: build_const_list
func.func @build_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
func.return %2 : !tfr.attr
}
// -----
// CHECK-LABEL: build_high_dim_const_list
func.func @build_high_dim_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%3 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%4 = "tfr.build_list"(%2, %3) : (!tfr.attr, !tfr.attr) -> !tfr.attr
func.return %4 : !tfr.attr
}
// -----
// CHECK-LABEL: get_length
func.func @get_length(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> index {
%0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
%1 = "tfr.get_length"(%0) : (!tfr.tensor_list) -> index
func.return %1 : index
}
// -----
// CHECK-LABEL: tfr.func
tfr.func @External(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<C>,
%arg2: i32 {tfr.name = "A"},
%arg3: !tfr.attr {tfr.name = "T"})
-> (!tfr.tensor<A>, !tfr.tensor_list<C>)
attributes {A, C}
// -----
// CHECK-LABEL: tfr.func
tfr.func @Foo(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<C>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"})
-> (!tfr.tensor<A>, !tfr.tensor_list<C>)
attributes {A, C} {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<C>
}
// -----
// CHECK-LABEL: tfr.func
tfr.func @Bar(%arg0: !tfr.tensor<A>,
%arg2: i32 {tfr.name = "B"},
%arg3: vector<1xi32> {tfr.name = "C"})
-> (!tfr.tensor<A>, !tfr.tensor<A>)
attributes {A} {
tfr.return %arg0, %arg0 : !tfr.tensor<A>, !tfr.tensor<A>
}
// -----
// expected-error@+1 {{Undefined attributes are used: A}}
tfr.func @Foo_undefined_attr(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{3 attribute argument doesn't have a tfr.name attribute}}
tfr.func @Foo_unnamed_attr(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32>) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{tfr.tensor/tfr.tensor_list argument should be before non tensor arguments}}
tfr.func @Foo_invalid_arg_order(%arg0: !tfr.tensor<A>,
%arg2: i32 {tfr.name = "A"},
%arg1: !tfr.tensor_list<A>,
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
tfr.func @Foo_valid_arg_order0(
%arg1: !tfr.tensor_list,
%arg0: !tfr.tensor<T>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor, !tfr.tensor_list) attributes {T}{
tfr.return %arg0, %arg1 : !tfr.tensor<T>, !tfr.tensor_list
}
// -----
// expected-error@+1 {{tfr.tensor argument should be before tfr.tensor_list argument.}}
tfr.func @Foo_invalid_arg_order0(
%arg1: !tfr.tensor_list,
%arg0: !tfr.tensor<T>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor, !tfr.tensor_list) {
tfr.return %arg0, %arg1 : !tfr.tensor<T>, !tfr.tensor_list
}
// -----
// expected-error@+1 {{tfr.tensor result should be before tfr.tensor_list result}}
tfr.func @Foo_invalid_result_order(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor_list<A>, !tfr.tensor<A>) {
tfr.return %arg1, %arg0 : !tfr.tensor_list<A>, !tfr.tensor<A>
}
// -----
// expected-error@+1 {{More than one tfr.tensor_list argument isn't allowed}}
tfr.func @Foo_multiple_tensor_list_args(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: !tfr.tensor_list<A>,
%arg3: i32 {tfr.name = "A"},
%arg4: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{More than one tfr.tensor_list result isn't allowed}}
tfr.func @Foo_multiple_tensor_list_results(%arg0: !tfr.tensor<C>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor_list<A>, !tfr.tensor_list<A>) {
tfr.return %arg1, %arg1 : !tfr.tensor_list<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{None tfr.tensor/tfr.tensor_list results aren't allowed as a result}}
tfr.func @Foo_return_attr(%arg0: !tfr.tensor<C>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) -> i32 {
tfr.return %arg2 : i32
}
tensorflow/compiler/mlir/tfr/tests/raise_to_tf.mlir
+140
View File
@@ -0,0 +1,140 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tfr-opt %s -tfr-raise-to-tf -verify-diagnostics -split-input-file | FileCheck %s
tfr.func @tf__risc_same_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes {T}
tfr.func @tf__risc_concat_(!tfr.tensor_list<N, T>) -> !tfr.tensor<T> attributes {T, N}
tfr.func @tf__risc_split_(!tfr.tensor<T>, i32 {tfr.name="N"}) -> !tfr.tensor_list<N, T> attributes {T, N}
tfr.func @tf__risc_cast_(!tfr.tensor, !tfr.attr {tfr.name="K"}) -> !tfr.tensor<K> attributes {T, K}
tfr.func @tf__const_(!tfr.attr {tfr.name="value", tfr.type="tensor"},
!tfr.attr {tfr.name="K",tfr.type="dtype"}) -> !tfr.tensor<K> attributes {T, K}
tfr.func @tf__positive_(!tfr.tensor<T>) -> !tfr.tensor<i1_> attributes {T, i1_}
tfr.func @tf__invalid_type_op_(!tfr.tensor<T>) -> !tfr.tensor<i8_> attributes {T, i8_}
// CHECK-LABEL: decompose_tf_same
func.func @decompose_tf_same(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
%1 = tfr.call @tf__risc_same(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<1x2x3x4x!tf_type.string>
func.return %2 : tensor<1x2x3x4x!tf_type.string>
// CHECK: %[[id:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id]]) <{shape = #tf_type.shape<1x2x3x4>}> : (tensor<*x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
// CHECK: return %[[es]] : tensor<1x2x3x4x!tf_type.string>
}
// CHECK-LABEL: decompose_tf_consecutive
func.func @decompose_tf_consecutive(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
%1 = "tfr.cast"(%arg2) : (tensor<f32>) -> !tfr.tensor
%2 = tfr.call @tf__risc_same(%0) : (!tfr.tensor) -> !tfr.tensor
%3 = tfr.call @tf__risc_same(%1) : (!tfr.tensor) -> !tfr.tensor
%4 = "tfr.cast"(%3) : (!tfr.tensor) -> tensor<f32>
func.return %4 : tensor<f32>
// CHECK: %[[id0:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string>
// CHECK: %[[id2:.*]] = "tf.RiscSame"(%arg2) : (tensor<f32>) -> tensor<*xf32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id2]]) <{shape = #tf_type.shape<>}> : (tensor<*xf32>) -> tensor<f32>
// CHECK: return %[[es]] : tensor<f32>
}
// CHECK-LABEL: decompose_tf_concat_n
func.func @decompose_tf_concat_n(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<3xf32> {
%0 = "tfr.cast"(%arg0) : (tensor<f32>) -> !tfr.tensor
%1 = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
%2 = "tfr.cast"(%arg2) : (tensor<f32>) -> !tfr.tensor
%3 = "tfr.build_list"(%0, %1, %2) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%concat = tfr.call @tf__risc_concat(%3) : (!tfr.tensor_list) -> !tfr.tensor
%4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<3xf32>
func.return %4 : tensor<3xf32>
// CHECK: %[[concat:.*]] = "tf.RiscConcat"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[concat]]) <{shape = #tf_type.shape<3>}> : (tensor<*xf32>) -> tensor<3xf32>
// CHECK: return %[[es]] : tensor<3xf32>
}
// CHECK-LABEL: decompose_tf_split
func.func @decompose_tf_split(%arg0: tensor<3xf32>) -> (tensor<f32>) {
%0 = "tfr.cast"(%arg0) : (tensor<3xf32>) -> !tfr.tensor
%n = arith.constant 3: i32
%split = tfr.call @tf__risc_split(%0, %n) : (!tfr.tensor, i32) -> !tfr.tensor_list
%i0 = arith.constant 0: index
%s0 = tfr.get_element %split[%i0] : (!tfr.tensor_list, index) -> !tfr.tensor
%4 = "tfr.cast"(%s0) : (!tfr.tensor) -> tensor<f32>
func.return %4 : tensor<f32>
// CHECK: %[[split:.*]]:3 = "tf.RiscSplit"(%arg0) {N = 3 : i32} : (tensor<3xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>)
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[split]]#0) <{shape = #tf_type.shape<>}> : (tensor<*xf32>) -> tensor<f32>
// CHECK: return %[[es]] : tensor<f32>
}
// CHECK-LABEL: decompose_tf_cast
func.func @decompose_tf_cast(%arg0: tensor<f32>) -> tensor<i32> {
%0 = "tfr.cast"(%arg0) : (tensor<f32>) -> !tfr.tensor
%t = tfr.constant i32 -> !tfr.attr
%concat = tfr.call @tf__risc_cast(%0, %t) : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
%4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<i32>
func.return %4 : tensor<i32>
// CHECK: %[[tfcast:.*]] = "tf.RiscCast"(%arg0) {K = i32} : (tensor<f32>) -> tensor<*xi32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) <{shape = #tf_type.shape<>}> : (tensor<*xi32>) -> tensor<i32>
// CHECK: return %[[es]] : tensor<i32>
}
// CHECK-LABEL: convert_to_scalar_tensor
func.func @convert_to_scalar_tensor() -> tensor<f32> {
%0 = arith.constant 3.0: f32
%t = tfr.constant f32 -> !tfr.attr
%cst = tfr.call @tf__const(%0, %t) : (f32, !tfr.attr) -> !tfr.tensor
%4 = "tfr.cast"(%cst) : (!tfr.tensor) -> tensor<f32>
func.return %4 : tensor<f32>
// CHECK: %[[cst:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK: return %[[cst]] : tensor<f32>
}
// CHECK-LABEL: attribute_propagate
func.func @attribute_propagate(%arg0: tensor<f32>) -> tensor<i32> {
%0 = "tfr.cast"(%arg0) : (tensor<f32>) -> !tfr.tensor
%t = tfr.constant i32 -> !tfr.attr
%concat = tfr.call @tf__risc_cast(%0, %t) {device = "hello", _tpu_replicate} : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
%4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<i32>
func.return %4 : tensor<i32>
// CHECK: %[[tfcast:.*]] = "tf.RiscCast"(%arg0) {K = i32, _tpu_replicate, device = "hello"} : (tensor<f32>) -> tensor<*xi32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) <{shape = #tf_type.shape<>}> : (tensor<*xi32>) -> tensor<i32>
// CHECK: return %[[es]] : tensor<i32>
}
// CHECK-LABEL: fixed_element_attribute
func.func @fixed_element_attribute(%arg0: tensor<2xf32>) -> tensor<2xi1> {
%0 = "tfr.cast"(%arg0) : (tensor<2xf32>) -> !tfr.tensor
%1 = tfr.call @tf__positive(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<2xi1>
func.return %2 : tensor<2xi1>
// CHECK: %[[positive:.*]] = "tf.Positive"(%arg0) : (tensor<2xf32>) -> tensor<*xi1>
// CHECK: %[[res:.*]] = "tf.EnsureShape"(%[[positive]]) <{shape = #tf_type.shape<2>}> : (tensor<*xi1>) -> tensor<2xi1>
// CHECK: return %[[res]] : tensor<2xi1>
}
// CHECK-LABEL: fixed_element_attribute_invalid
func.func @fixed_element_attribute_invalid(%arg0: tensor<2xf32>) -> tensor<2xi8> {
%0 = "tfr.cast"(%arg0) : (tensor<2xf32>) -> !tfr.tensor
// expected-error@+1 {{type i8_ can't be resolved for the signature of the op}}
%1 = tfr.call @tf__invalid_type_op(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<2xi8>
func.return %2 : tensor<2xi8>
// CHECK: tfr.call @tf__invalid_type_op
}
tensorflow/compiler/mlir/tfr/tests/rewrite_quantized_io.mlir
+58
View File
@@ -0,0 +1,58 @@
// 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: tfr-opt %s -tfr-decompose -tfr-rewrite-quantized-io -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @tf__my_requantize
tfr.func @tf__my_requantize(%input: !tfr.tensor) -> !tfr.tensor {
%raw_data = tfr.quant_raw_data(%input) : (!tfr.tensor) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%input) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
%result = tfr.call @tf__requantize(%raw_data, %scale, %zp) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor
tfr.return %result : !tfr.tensor
}
// CHECK-LABEL: @tf__intermediate
tfr.func @tf__intermediate(%arg0: !tfr.tensor) -> !tfr.tensor {
%0 = tfr.call @tf__risc(%arg0) : (!tfr.tensor) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
}
// CHECK-LABEL: remove_quantized_io
func.func @remove_quantized_io(
%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:-128>>,
%arg1: tensor<1x5xf32>) -> (tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>, tensor<1x5xf32>) {
%0 = "tf.MyRequantize"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>
%1 = "tf.Intermediate"(%arg1) : (tensor<1x5xf32>) -> tensor<1x5xf32>
func.return %0, %1 : tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>, tensor<1x5xf32>
// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<1.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor<i32>}> : () -> tensor<i32>
// CHECK: %[[quant:.*]] = "tfr.cast"(%arg0) : (tensor<1x10xi8>) -> !tfr.tensor
// CHECK: %[[scale_cast:.*]] = "tfr.cast"(%[[scale]])
// CHECK: %[[zp_cast:.*]] = "tfr.cast"(%[[zp]])
// CHECK: %[[requant:.*]] = tfr.call @tf__requantize(%[[quant]], %[[scale_cast]], %[[zp_cast]])
// CHECK: %[[result:.*]] = "tfr.cast"(%[[requant]])
// CHECK-NOT: quant.uniform
// CHECK: return %[[result]], %[[float_resunt:.*]] : tensor<1x10xi8>, tensor<1x5xf32>
}
// CHECK-LABEL: quant_input_multiple_users
// expected-error@+1 {{The argument with type tensor<1x10x!quant.uniform<i8:f32, 1.000000e-01>> should have one user}}
func.func @quant_input_multiple_users(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1>>) -> (!tfr.tensor, !tfr.tensor) {
%0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1>>) -> !tfr.tensor
%1 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1>>) -> !tfr.tensor
func.return %0, %1 : !tfr.tensor, !tfr.tensor
}