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/tf2xla/tests/BUILD
+33
View File
@@ -0,0 +1,33 @@
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
size_override = {
"legalize-tf-binary-elementwise.mlir": "medium",
"legalize-tf-include-tf2xla-fallback.mlir": "medium",
"legalize-tf-prefer-tf2xla.mlir": "medium",
"legalize-tf.mlir": "medium",
},
test_file_exts = [
"mlir",
"hlotxt",
],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir:tf-opt",
"@llvm-project//llvm:FileCheck",
],
)
tensorflow/compiler/mlir/tf2xla/tests/adjust-layout.mlir
+28
View File
@@ -0,0 +1,28 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt -pass-pipeline='builtin.module(func.func(infeed-ops-xla-adjust-layout))' %s | env FILECHECK_OPTS="" FileCheck %s
func.func @infeed_dequeue_tuple() -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>) {
// CHECK: [[TOKEN:%.*]] = mhlo.create_token : !mhlo.token
%0 = "mhlo.create_token"() : () -> !mhlo.token
// CHECK: [[INFEED:%.*]]:3 = "mhlo.infeed"([[TOKEN]]) <{
// CHECK-SAME{LITERAL}: infeed_config = "", layout = [[1, 3, 2, 0], [1, 2, 0]]
// CHECK-SAME: }> : (!mhlo.token) -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>, !mhlo.token)
%1:3 = "mhlo.infeed"(%0) {infeed_config = ""} : (!mhlo.token) -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>, !mhlo.token)
// CHECK: return [[INFEED]]#0, [[INFEED]]#1
func.return %1#0, %1#1 : tensor<1x8x4x4xi32>, tensor<1x100x1xf32>
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-BatchMatMulV2.mlir
+103
View File
@@ -0,0 +1,103 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt -xla-legalize-tf %s | FileCheck %s
//===----------------------------------------------------------------------===//
// tf.BatchMatMulV2 op legalizations.
//===----------------------------------------------------------------------===//
func.func @batchmatmulv2_basic(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL: func @batchmatmulv2_basic
// CHECK-SAME: ([[LHS:%.*]]: tensor<1x4x2xf32>, [[RHS:%.*]]: tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
// CHECK: [[LHSSHAPE:%.*]] = shape.shape_of [[LHS]] : tensor<1x4x2xf32>
// CHECK: [[RHSSHAPE:%.*]] = shape.shape_of [[RHS]] : tensor<3x2x4xf32>
// CHECK: [[CM2:%.*]] = arith.constant -2 : index
// CHECK: [[LHSHEAD:%.*]], [[LHSTAIL:%.*]] = "shape.split_at"([[LHSSHAPE]], [[CM2]])
// CHECK: [[RHSHEAD:%.*]], [[RHSTAIL:%.*]] = "shape.split_at"([[RHSSHAPE]], [[CM2]])
// CHECK: [[BCASTHEAD:%.*]] = shape.broadcast [[LHSHEAD]], [[RHSHEAD]]
// CHECK: [[LHSBCASTSHAPE:%.*]] = shape.concat [[BCASTHEAD]], [[LHSTAIL]]
// CHECK: [[LHSSHAPEEXTENTS:%.*]] = shape.to_extent_tensor [[LHSBCASTSHAPE]]
// CHECK: [[LHSBCAST:%.*]] = "mhlo.dynamic_broadcast_in_dim"([[LHS]], [[LHSSHAPEEXTENTS]]) <{broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}> : (tensor<1x4x2xf32>, tensor<3xindex>) -> tensor<3x4x2xf32>
// CHECK: [[RHSBCASTSHAPE:%.*]] = shape.concat [[BCASTHEAD]], [[RHSTAIL]]
// CHECK: [[RESULT:%.*]] = "mhlo.dot_general"([[LHSBCAST]], [[RHS]])
// CHECK: return [[RESULT]] : tensor<3x4x4xf32>
// CHECK: }
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
func.return %0 : tensor<3x4x4xf32>
}
func.func @batchmatmulv2_lhs_batch(%arg0: tensor<3x4x2xf32>, %arg1: tensor<2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL: func @batchmatmulv2_lhs_batch
// CHECK: "mhlo.dynamic_broadcast_in_dim"({{.*}}, {{.*}}) <{broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}>
// CHECK: "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-SAME: lhs_batching_dimensions = [0]
// CHECK-SAME: rhs_batching_dimensions = [0]
// CHECK-SAME: lhs_contracting_dimensions = [2]
// CHECK-SAME: rhs_contracting_dimensions = [1]
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, device = ""} : (tensor<3x4x2xf32>, tensor<2x4xf32>) -> tensor<3x4x4xf32>
func.return %0 : tensor<3x4x4xf32>
}
func.func @batchmatmulv2_rhs_batch(%arg0: tensor<4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL: func @batchmatmulv2_rhs_batch
// CHECK: "mhlo.dynamic_broadcast_in_dim"({{.*}}, {{.*}}) <{broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}>
// CHECK: "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-SAME: lhs_batching_dimensions = [0]
// CHECK-SAME: rhs_batching_dimensions = [0]
// CHECK-SAME: lhs_contracting_dimensions = [2]
// CHECK-SAME: rhs_contracting_dimensions = [1]
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, device = ""} : (tensor<4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
func.return %0 : tensor<3x4x4xf32>
}
func.func @batchmatmulv2_dynamic(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
// CHECK-LABEL: func @batchmatmulv2_dynamic
// CHECK: "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-SAME: lhs_batching_dimensions = [0]
// CHECK-SAME: rhs_batching_dimensions = [0]
// CHECK-SAME: lhs_contracting_dimensions = [2]
// CHECK-SAME: rhs_contracting_dimensions = [1]
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, device = ""} : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
func.return %0 : tensor<?x?x?xf32>
}
func.func @batchmatmulv2_adj_real(%arg0: tensor<2x5xf32>, %arg1: tensor<4x2xf32>) -> tensor<5x4xf32> {
// CHECK-LABEL: func @batchmatmulv2_adj_real
// CHECK: "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-NOT: lhs_batching_dimensions
// CHECK-NOT: rhs_batching_dimensions
// CHECK-SAME: lhs_contracting_dimensions = [0]
// CHECK-SAME: rhs_contracting_dimensions = [1]
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = true, adj_y = true, device = ""} : (tensor<2x5xf32>, tensor<4x2xf32>) -> tensor<5x4xf32>
func.return %0 : tensor<5x4xf32>
}
func.func @batchmatmulv2_adj_complex(%arg0: tensor<2x5xcomplex<f32>>, %arg1: tensor<4x2xcomplex<f32>>) -> tensor<5x4xcomplex<f32>> {
// CHECK-LABEL: func @batchmatmulv2_adj_complex(
// CHECK-SAME: [[LHS:%.*]]: tensor<2x5xcomplex<f32>>, [[RHS:%.*]]: tensor<4x2xcomplex<f32>>) -> tensor<5x4xcomplex<f32>> {
// CHECK: [[LHSRE:%.*]] = mhlo.real [[LHS]]
// CHECK: [[LHSIM:%.*]] = mhlo.imag [[LHS]]
// CHECK: [[LHSIMNEG:%.*]] = mhlo.negate [[LHSIM]]
// CHECK: [[LHSCONJ:%.*]] = mhlo.complex [[LHSRE]], [[LHSIMNEG]]
// CHECK: [[RHSRE:%.*]] = mhlo.real [[RHS]]
// CHECK: [[RHSIM:%.*]] = mhlo.imag [[RHS]]
// CHECK: [[RHSIMNEG:%.*]] = mhlo.negate [[RHSIM]]
// CHECK: [[RHSCONJ:%.*]] = mhlo.complex [[RHSRE]], [[RHSIMNEG]]
// CHECK: shape.shape_of [[LHSCONJ]]
// CHECK: shape.shape_of [[RHSCONJ]]
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = true, adj_y = true, device = ""} : (tensor<2x5xcomplex<f32>>, tensor<4x2xcomplex<f32>>) -> tensor<5x4xcomplex<f32>>
func.return %0 : tensor<5x4xcomplex<f32>>
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-binary-elementwise.mlir
+396
View File
@@ -0,0 +1,396 @@
// 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.
// ==============================================================================
// Note that binary elementwise tests are run with chlo legalization enabled
// (unlike the rest), since this is the primary use case for such ops and
// verification of shapes and broadcasts is desired.
// RUN: tf-opt "-xla-legalize-tf=legalize-chlo=true" -canonicalize %s | FileCheck %s
// RUN: tf-opt "-xla-legalize-tf=legalize-chlo=false" %s | FileCheck --check-prefix CHLO %s
//===----------------------------------------------------------------------===//
// Binary op legalizations.
// Most of these expand from the same pattern. Full semantics are
// verified for tf.Add and pattern application only for the rest.
//===----------------------------------------------------------------------===//
// CHECK-LABEL: func @add
func.func @add(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %[[SUM0:.*]] = mhlo.add %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %[[SUM0]] : tensor<2xi32>
%1 = "tf.AddV2"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %1: tensor<2xi32>
}
// CHECK-LABEL: func @broadcast_add
// TODO(laurenzo): Change this to a (5 + 2x1) shaped add to make the check
// patterns unambiguous and more interesting (once broadcastable trait is
// fixed upstream).
func.func @broadcast_add(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi32> {
// CHECK-NEXT: %[[LHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<1> : tensor<1xi64>}>
// CHECK-NEXT: mhlo.add %[[LHS_BCAST]], %arg1
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi32>
func.return %0: tensor<1x2xi32>
}
// CHECK-LABEL: func @broadcast_multi_dim_add
// TODO(laurenzo): Change this to a (4x1x1 + 1x4x4x4) shaped add once upstream
// broadcastable bug is fixed (helps make the CHECK matching unambiguous)
func.func @broadcast_multi_dim_add(%arg0: tensor<4x1x1xi32>, %arg1: tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32> {
// CHECK-NEXT: %[[LHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>}>
// CHECK-NEXT: mhlo.add %[[LHS_BCAST]], %arg1
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<4x1x1xi32>, tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32>
func.return %0: tensor<4x4x4x4xi32>
}
// CHECK-LABEL: func @add_dynamic
func.func @add_dynamic(%arg0: tensor<?xi32>, %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> {
// CHECK-DAG: %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0
// CHECK-DAG: %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1
// CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]]
// CHECK-NEXT: shape.assuming %[[WITNESS:.+]]
// CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
// CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1
// CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.broadcast %[[LHS_SHAPE]], %[[RHS_SHAPE]] : tensor<1xindex>, tensor<2xindex> -> tensor<2xindex>
// CHECK-NEXT: %[[LHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}>
// CHECK-NEXT: %[[RHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}>
// CHECK-NEXT: %[[RESULT:.+]] = mhlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] : tensor<?x?xi32>
// CHECK-NEXT: shape.assuming_yield %[[RESULT]]
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<?xi32>, tensor<?x?xi32>) -> tensor<?x?xi32>
func.return %0: tensor<?x?xi32>
}
// CHECK-LABEL: func @broadcast_add_unranked
// CHLO-LABEL: func @broadcast_add_unranked
func.func @broadcast_add_unranked(%arg0: tensor<1xi32>, %arg1: tensor<*xi32>) -> tensor<*xi32> {
// CHECK: tf.Add
// CHLO: chlo.broadcast_add %arg0, %arg1
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<1xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %0: tensor<*xi32>
}
// CHECK-LABEL: func @div
func.func @div(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %0 = mhlo.divide %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
%0 = "tf.Div"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
}
// CHECK-LABEL: func @shift_left
func.func @shift_left(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
// CHECK: mhlo.shift_left %arg0, %arg1 : tensor<4xi32>
%0 = "tf.LeftShift"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %0 : tensor<4xi32>
}
// CHECK-LABEL: func @div_unranked
func.func @div_unranked(%arg0: tensor<*xi32>, %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> {
// CHECK-NEXT: tf.Div
%0 = "tf.Div"(%arg0, %arg1) : (tensor<*xi32>, tensor<?x?xi32>) -> tensor<?x?xi32>
func.return %0: tensor<?x?xi32>
}
// CHECK-LABEL: func @maximum
func.func @maximum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
// CHECK-NEXT: mhlo.maximum %arg0, %arg1 : tensor<4xf32>
%0 = "tf.Maximum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: func @minimum
func.func @minimum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
// CHECK-NEXT: mhlo.minimum %arg0, %arg1 : tensor<4xf32>
%0 = "tf.Minimum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: func @mod
// CHLO-LABEL: func @mod
func.func @mod(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
// CHECK-NEXT: mhlo.remainder %arg0, %arg1 : tensor<4xf32>
// CHLO: chlo.broadcast_remainder
%0 = "tf.Mod"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: func @mul
func.func @mul(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %0 = mhlo.multiply %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
%0 = "tf.Mul"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
}
// CHECK-LABEL: func @real_div
func.func @real_div(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %0 = mhlo.divide %arg0, %arg0 : tensor<2xi32>
%0 = "tf.RealDiv"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
}
// CHECK-LABEL: func @sub
func.func @sub(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %0 = mhlo.subtract %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
%0 = "tf.Sub"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
}
// CHECK-LABEL: func @shift_right
func.func @shift_right(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
// CHECK: mhlo.shift_right_arithmetic %arg0, %arg1 : tensor<4xi32>
%0 = "tf.RightShift"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %0 : tensor<4xi32>
}
// CHECK-LABEL: func @shift_right_unsigned
func.func @shift_right_unsigned(%arg0: tensor<4xui8>, %arg1: tensor<4xui8>) -> tensor<4xui8> {
// CHECK: mhlo.shift_right_logical %arg0, %arg1 : tensor<4xui8>
%0 = "tf.RightShift"(%arg0, %arg1) : (tensor<4xui8>, tensor<4xui8>) -> tensor<4xui8>
func.return %0 : tensor<4xui8>
}
// CHECK-LABEL: func @broadcast_shift_right_unsigned
func.func @broadcast_shift_right_unsigned(%arg0: tensor<4xui8>, %arg1: tensor<2x4xui8>) -> tensor<2x4xui8> {
// CHECK: %[[BROADCAST:.*]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<4xui8>) -> tensor<2x4xui8>
// CHECK: mhlo.shift_right_logical %[[BROADCAST]], %arg1 : tensor<2x4xui8>
%0 = "tf.RightShift"(%arg0, %arg1) : (tensor<4xui8>, tensor<2x4xui8>) -> tensor<2x4xui8>
func.return %0 : tensor<2x4xui8>
}
// CHECK-LABEL: func @and
func.func @and(%arg0: tensor<2xi1>, %arg1: tensor<2xi1>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.and
%0 = "tf.LogicalAnd"(%arg0, %arg1) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: func @and_unranked
func.func @and_unranked(%arg0: tensor<*xi1>, %arg1: tensor<*xi1>) -> tensor<*xi1> {
// CHECK: tf.LogicalAnd
%0 = "tf.LogicalAnd"(%arg0, %arg1) : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
func.return %0: tensor<*xi1>
}
// CHECK-LABEL: func @or
func.func @or(%arg0: tensor<2xi1>, %arg1: tensor<2xi1>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.or
%0 = "tf.LogicalOr"(%arg0, %arg1) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: func @bitwise_or
func.func @bitwise_or(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
// CHECK-NEXT: mhlo.or
%0 = "tf.BitwiseOr"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %0: tensor<4xi32>
}
// CHECK-LABEL: func @bitwise_or_unsigned
func.func @bitwise_or_unsigned(%arg0: tensor<4xui32>, %arg1: tensor<4xui32>) -> tensor<4xui32> {
// CHECK-NEXT: mhlo.or
%0 = "tf.BitwiseOr"(%arg0, %arg1) : (tensor<4xui32>, tensor<4xui32>) -> tensor<4xui32>
func.return %0: tensor<4xui32>
}
// CHECK-LABEL: func @bitwise_xor
func.func @bitwise_xor(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
// CHECK-NEXT: mhlo.xor
%0 = "tf.BitwiseXor"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %0: tensor<4xi32>
}
// CHECK-LABEL: func @bitwise_xor_unsigned
func.func @bitwise_xor_unsigned(%arg0: tensor<4xui32>, %arg1: tensor<4xui32>) -> tensor<4xui32> {
// CHECK-NEXT: mhlo.xor
%0 = "tf.BitwiseXor"(%arg0, %arg1) : (tensor<4xui32>, tensor<4xui32>) -> tensor<4xui32>
func.return %0: tensor<4xui32>
}
// CHECK-LABEL: func @bitwise_and
func.func @bitwise_and(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
// CHECK-NEXT: mhlo.and
%0 = "tf.BitwiseAnd"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %0: tensor<4xi32>
}
// CHECK-LABEL: func @bitwise_and_unsigned
func.func @bitwise_and_unsigned(%arg0: tensor<4xui32>, %arg1: tensor<4xui32>) -> tensor<4xui32> {
// CHECK-NEXT: mhlo.and
%0 = "tf.BitwiseAnd"(%arg0, %arg1) : (tensor<4xui32>, tensor<4xui32>) -> tensor<4xui32>
func.return %0: tensor<4xui32>
}
// CHECK-LABEL: func @pow
func.func @pow(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NEXT: mhlo.power
%0 = "tf.Pow"(%arg0, %arg0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
func.return %0: tensor<2xf32>
}
//===----------------------------------------------------------------------===//
// Equality op legalizations.
// tf.Equal and tf.NotEqual expand from the same pattern. Full semantics are
// verified for tf.Equal and pattern application only for tf.NotEqual
//===----------------------------------------------------------------------===//
// CHECK-LABEL: func @equal
func.func @equal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.compare EQ, %arg0, %arg1
%0 = "tf.Equal"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: func @equal_dynamic
func.func @equal_dynamic(%arg0: tensor<?xi32>, %arg1: tensor<1xi32>) -> tensor<?xi1> {
// TODO(jpienaar): Uncomment below when fallout from https://reviews.llvm.org/D83194 fixed.
// NOT-CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
// NOT-CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1]
// NOT-CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]]
// NOT-CHECK-NEXT: shape.assuming %[[WITNESS]] -> (tensor<?xi1>) {
// NOT-CHECK-DAG: %[[LHS_SHAPE1:.+]] = shape.shape_of %arg0
// NOT-CHECK-NEXT: %[[RESULT_SHAPE:.+]] = shape.broadcast %[[LHS_SHAPE1]], %[[RHS_SHAPE]] : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
// NOT-CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = tensor.cast %[[RESULT_SHAPE]] : tensor<?xindex> to tensor<1xindex>
// NOT-CHECK-DAG: %[[LHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<0> : tensor<1xi64>}>
// NOT-CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<0> : tensor<1xi64>}>
// NOT-CHECK-NEXT: %[[RESULT:.+]] = mhlo.compare EQ, %[[LHS_BCAST]], %[[RHS_BCAST]]
// NOT-CHECK-NEXT: shape.assuming_yield %[[RESULT]]
%0 = "tf.Equal"(%arg0, %arg1) : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi1>
func.return %0: tensor<?xi1>
}
// CHECK-LABEL: func @equal_broadcast
func.func @equal_broadcast(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> {
// CHECK-DAG: %[[LHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<1> : tensor<1xi64>}>
// CHECK-NEXT: mhlo.compare EQ, %[[LHS_BCAST]], %arg1
%0 = "tf.Equal"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1>
func.return %0: tensor<1x2xi1>
}
// CHECK-LABEL: func @equal_broadcast_no_incompatible_shapes_error
func.func @equal_broadcast_no_incompatible_shapes_error(%arg0: tensor<2xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> {
// CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = true}>
%0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<2xi32>, tensor<1x2xi32>) -> tensor<1x2xi1>
func.return %0: tensor<1x2xi1>
}
// CHECK-LABEL: func @equal_incompatible_shape_broadcastable
func.func @equal_incompatible_shape_broadcastable(%arg0: tensor<?xi32>, %arg1: tensor<1xi32>) -> tensor<?xi1> {
// CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = true}>
%0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi1>
func.return %0: tensor<?xi1>
}
// CHECK-LABEL: func @equal_incompatible_shape_dynamic
func.func @equal_incompatible_shape_dynamic(%arg0: tensor<2xi32>, %arg1: tensor<?xi32>) -> tensor<*xi1> {
// CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = false}>
%0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<2xi32>, tensor<?xi32>) -> tensor<*xi1>
func.return %0: tensor<*xi1>
}
// CHECK-LABEL: func @equal_incompatible_shape_both_dynamic
func.func @equal_incompatible_shape_both_dynamic(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<*xi1> {
// CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = false}>
%0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<?xi32>, tensor<?xi32>) -> tensor<*xi1>
func.return %0: tensor<*xi1>
}
// CHECK-LABEL: func @equal_unranked
func.func @equal_unranked(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> tensor<*xi1> {
// CHECK: "tf.Equal"
// CHLO: chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction EQ>}
%0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi1>
func.return %0: tensor<*xi1>
}
// CHECK-LABEL: func @equal_unsupported_type
func.func @equal_unsupported_type(%arg0: tensor<*x!tf_type.string>, %arg1: tensor<*x!tf_type.string>) -> tensor<*xi1> {
// CHECK: "tf.Equal"
%0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>) -> tensor<*xi1>
func.return %0: tensor<*xi1>
}
// CHECK-LABEL: func @notequal
func.func @notequal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.compare NE, %arg0, %arg1
%0 = "tf.NotEqual"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
//===----------------------------------------------------------------------===//
// Compare op legalizations.
// These expand from the same pattern. Full semantics are checked for
// tf.Greater. Others just check that the pattern applied.
//===----------------------------------------------------------------------===//
// CHECK-LABEL: func @greater
func.func @greater(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK: mhlo.compare GT, %arg0, %arg1
%0 = "tf.Greater"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: func @broadcast_greater
func.func @broadcast_greater(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> {
// CHECK-NEXT: %[[LHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<1> : tensor<1xi64>}>
// CHECK-NEXT: mhlo.compare GT, %[[LHS_BCAST]], %arg1
%0 = "tf.Greater"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1>
func.return %0: tensor<1x2xi1>
}
// CHECK-LABEL: func @greater_dynamic
func.func @greater_dynamic(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<?xi1> {
// CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
// CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1
// CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]]
// CHECK-NEXT: shape.assuming %[[WITNESS]]
// CHECK-DAG: %[[LHS_SHAPE1:.+]] = shape.shape_of %arg0
// CHECK-DAG: %[[RHS_SHAPE1:.+]] = shape.shape_of %arg1
// CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.broadcast %[[LHS_SHAPE1]], %[[RHS_SHAPE1]] : tensor<1xindex>, tensor<1xindex> -> tensor<1xindex>
// CHECK-DAG: %[[LHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<0> : tensor<1xi64>}>
// CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<0> : tensor<1xi64>}>
// CHECK-NEXT: mhlo.compare GT, %[[LHS_BCAST]], %[[RHS_BCAST]]
%0 = "tf.Greater"(%arg0, %arg1) : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi1>
func.return %0: tensor<?xi1>
}
// CHECK-LABEL: func @greater_uranked
func.func @greater_uranked(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> tensor<*xi1> {
// CHECK: "tf.Greater"
// CHLO: chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction GT>}
%0 = "tf.Greater"(%arg0, %arg1) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi1>
func.return %0: tensor<*xi1>
}
// CHECK-LABEL: func @greater_equal
func.func @greater_equal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.compare GE, %arg0, %arg1
%0 = "tf.GreaterEqual"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: func @less
func.func @less(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.compare LT, %arg0, %arg1
%0 = "tf.Less"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: func @less_equal
func.func @less_equal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.compare LE, %arg0, %arg1
%0 = "tf.LessEqual"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-collective.mlir
+299
View File
@@ -0,0 +1,299 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt -split-input-file -verify-diagnostics -xla-legalize-tf-collective -xla-legalize-tf %s | FileCheck %s
// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 0
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_size = 2
// CHECK-LABEL: func @all_reduce_cross_replica
func.func @all_reduce_cross_replica(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: "mhlo.all_reduce"
// CHECK{LITERAL}: replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
// CHECK-NOT: channel_handle
// CHECK: mhlo.add
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 0
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_size = 2
// CHECK-LABEL: func @all_reduce_cross_replica_and_partition
func.func @all_reduce_cross_replica_and_partition(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: "mhlo.all_reduce"
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 10001, type = 1>
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
// CHECK: mhlo.add
// CHECK: mhlo.return
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplicaAndPartition"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
// CHECK: "mhlo.all_reduce"
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 10000, type = 1>
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
// CHECK: mhlo.add
// CHECK: mhlo.return
%1 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplicaAndPartition"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
%2 = "tf.Add"(%0, %1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %2 : tensor<f32>
}
// -----
// CHECK-LABEL: func @xla_all_reduce_add
func.func @xla_all_reduce_add(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: "mhlo.all_reduce"
// CHECK: mhlo.add
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @xla_all_reduce_max
func.func @xla_all_reduce_max(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: "mhlo.all_reduce"
// CHECK: mhlo.maximum
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Max", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @xla_all_reduce_mean
func.func @xla_all_reduce_mean(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<1.000000e+00>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK: mhlo.add
// CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
// CHECK-NEXT: return %[[RESULT]]
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Mean", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @xla_all_reduce_min
func.func @xla_all_reduce_min(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: "mhlo.all_reduce"
// CHECK: mhlo.minimum
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Min", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @xla_all_reduce_mul
func.func @xla_all_reduce_mul(%input: tensor<f32>) -> tensor<f32> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// CHECK: "mhlo.all_reduce"
// CHECK: mhlo.mul
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Mul", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
func.func @xla_all_reduce_tuple(%input: tuple<tensor<f32>, tensor<f32>>) -> tuple<tensor<f32>, tensor<f32>> {
%group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
// expected-error@+1 {{'tf.XlaAllReduce' op operand #0 must be tensor of bfloat16 or 16-bit float or 32-bit float or 32-bit integer or 32-bit unsigned integer values, but got 'tuple<tensor<f32>, tensor<f32>>'}}
%0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplica"} : (tuple<tensor<f32>, tensor<f32>>, tensor<2x1xi32>) -> tuple<tensor<f32>, tensor<f32>>
func.return %0 : tuple<tensor<f32>, tensor<f32>>
}
// -----
// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 1
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_size = 2
// CHECK-LABEL: func @collective_reduce_v2
func.func @collective_reduce_v2(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: "mhlo.all_reduce"
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 10001, type = 1>
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.add
// CHECK: mhlo.return
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
// CHECK: "mhlo.all_reduce"
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 10000, type = 1>
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.add
// CHECK: mhlo.return
%1 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
%2 = "tf.Add"(%0, %1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %2 : tensor<f32>
}
// -----
// CHECK-LABEL: func @collective_reduce_v2_add_id
func.func @collective_reduce_v2_add_id(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.add
// CHECK: mhlo.return
// CHECK: return %[[REDUCE]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_max_id
func.func @collective_reduce_v2_max_id(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.maximum
// CHECK: mhlo.return
// CHECK: return %[[REDUCE]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Max", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_min_id
func.func @collective_reduce_v2_min_id(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.minimum
// CHECK: mhlo.return
// CHECK: return %[[REDUCE]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Min", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_mul_id
func.func @collective_reduce_v2_mul_id(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.mul
// CHECK: mhlo.return
// CHECK: return %[[REDUCE]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Mul", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_add_div
func.func @collective_reduce_v2_add_div(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<2.000000e+00>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.add
// CHECK: mhlo.return
// CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
// CHECK-NEXT: return %[[RESULT]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Div"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_max_div
func.func @collective_reduce_v2_max_div(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<2.000000e+00>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.maximum
// CHECK: mhlo.return
// CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
// CHECK-NEXT: return %[[RESULT]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Max", final_op = "Div"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_min_div
func.func @collective_reduce_v2_min_div(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<2.000000e+00>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.minimum
// CHECK: mhlo.return
// CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
// CHECK-NEXT: return %[[RESULT]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Min", final_op = "Div"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @collective_reduce_v2_mul_div
func.func @collective_reduce_v2_mul_div(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<2.000000e+00>
// CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
// CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK: mhlo.mul
// CHECK: mhlo.return
// CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
// CHECK-NEXT: return %[[RESULT]]
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Mul", final_op = "Div"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 0
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_size = 2
// CHECK-LABEL: func @collective_assign_group_v2
func.func @collective_assign_group_v2(%input: tensor<f32>) -> tensor<f32> {
%rank = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
%key_base = "tf.Const"() { value = dense<10> : tensor<i32> } : () -> tensor<i32>
%group_assignment = "tf.Const"() { value = dense<[[0, 1]]> : tensor<1x2xi32> } : () -> tensor<1x2xi32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
%group_size, %group_key = "tf.CollectiveAssignGroupV2"(%group_assignment, %rank, %key_base) {} : (tensor<1x2xi32>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
// CHECK-NOT: "tf.CollectiveAssignGroupV2"
// CHECK: "mhlo.all_reduce"
// CHECK: mhlo.add
// CHECK{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// CHECK-NOT: "tf.CollectiveAssignGroupV2"
%0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
func.func @inconsistent_collective_info(%input: tensor<f32>) -> tensor<f32> {
%group_key = "tf.Const"() { value = dense<11> : tensor<i32> } : () -> tensor<i32>
%group_size1 = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
%group_size2 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// expected-error@below {{op module already contains an attribute tf2xla.collective_info.group_size=2, overwritting to a new value 1 is not allowed.}}
%0 = "tf.CollectiveReduceV2"(%input, %group_size1, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
%1 = "tf.CollectiveReduceV2"(%input, %group_size2, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
%2 = "tf.Add"(%0, %1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
func.return %2 : tensor<f32>
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-communication.mlir
+930
View File
@@ -0,0 +1,930 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt -split-input-file -verify-diagnostics -xla-legalize-tf-communication %s | FileCheck %s
// Test legalization of `tf._XlaHostComputeMlir` expands into individual
// `mhlo.send` per operand and `mhlo.recv` per result. Channel Id's are uniquely
// assigned per mhlo communcation op, and frontend attributes (modified keys)
// and op shardings are added. Sink tokens are created
// if there are more than one operand or more than one result.
//
// The following op sharding assigns full tensors to TPU core 0.
// Proto debug string:
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 0
// Serialized string:
// "\08\01\1A\01\01\22\01\00"
// CHECK-LABEL: func @host_compute
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>, [[ARG1:%.*]]: tensor<i64>)
func.func @host_compute(%arg0: tensor<i32>, %arg1: tensor<i64>) -> (tensor<f32>, tensor<f64>) {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[SEND_ARG0_TOKEN:%.*]] = "mhlo.send"([[ARG0]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 1, type = 2>
// CHECK-SAME: is_host_transfer = true
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "host_compute_channel_send_dtoh_0"}
// CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
// CHECK-SAME: (tensor<i32>, !mhlo.token) -> !mhlo.token
// CHECK: [[SEND_ARG1_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 2, type = 2>
// CHECK-SAME: is_host_transfer = true
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "host_compute_channel_send_dtoh_1"}
// CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
// CHECK-SAME: (tensor<i64>, !mhlo.token) -> !mhlo.token
// CHECK: [[SEND_SINK_TOKEN:%.*]] = mhlo.after_all [[SEND_ARG0_TOKEN]], [[SEND_ARG1_TOKEN]]
// CHECK: [[RECV_RETVAL0_TUPLE:%.*]]:2 = "mhlo.recv"([[SEND_SINK_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 3, type = 3>
// CHECK-SAME: is_host_transfer = true
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "host_compute_channel_recv_htod_0"}
// CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
// CHECK-SAME: (!mhlo.token) -> (tensor<f32>, !mhlo.token)
// CHECK: [[RECV_RETVAL1_TUPLE:%.*]]:2 = "mhlo.recv"([[SEND_SINK_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 4, type = 3>
// CHECK-SAME: is_host_transfer = true
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "host_compute_channel_recv_htod_1"}
// CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
// CHECK-SAME: (!mhlo.token) -> (tensor<f64>, !mhlo.token)
// CHECK: [[RECV_SINK_TOKEN:%.*]] = mhlo.after_all [[RECV_RETVAL0_TUPLE]]#1, [[RECV_RETVAL1_TUPLE]]#1
%0:2 = "tf._XlaHostComputeMlir"(%arg0, %arg1) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : (tensor<i32>, tensor<i64>) -> (tensor<f32>, tensor<f64>)
// CHECK: return [[RECV_RETVAL0_TUPLE]]#0, [[RECV_RETVAL1_TUPLE]]#0 : tensor<f32>, tensor<f64>
func.return %0#0, %0#1 : tensor<f32>, tensor<f64>
}
// -----
// Tests `tf._XlaHostComputeMlir` with no operands simply forwards the input
// token to its generated `mhlo.recv`.
// CHECK-LABEL: func @host_compute_no_operands_one_result
func.func @host_compute_no_operands_one_result() {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK-NOT: "mhlo.send"
// CHECK-NOT: "mhlo.after_all"
// CHECK: "mhlo.recv"([[INIT_TOKEN]])
%0 = "tf._XlaHostComputeMlir"() {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : () -> tensor<i32>
func.return
}
// -----
// Tests `tf._XlaHostComputeMlir` with no results simply forwards its token from
// the generated `mhlo.send`.
// CHECK-LABEL: func @host_compute_one_operand_no_results
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @host_compute_one_operand_no_results(%arg0: tensor<i32>) {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[SEND_TOKEN:%.*]] = "mhlo.send"([[ARG0]], [[INIT_TOKEN]])
// CHECK-NOT: "mhlo.after_all"
"tf._XlaHostComputeMlir"(%arg0) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : (tensor<i32>) -> ()
// CHECK: "mhlo.recv"([[SEND_TOKEN]])
%0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>
func.return
}
// -----
// Tests `tf._XlaHostComputeMlir` with one operand and one result does not
// create any `mhlo.after_all` ops.
// CHECK-LABEL: func @host_compute_single_operand_result
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @host_compute_single_operand_result(%arg0: tensor<i32>) {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[SEND_TOKEN:%.*]] = "mhlo.send"([[ARG0]], [[INIT_TOKEN]])
// CHECK-NOT: "mhlo.after_all"
// CHECK: "mhlo.recv"([[SEND_TOKEN]])
// CHECK-NOT: "mhlo.after_all"
%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : (tensor<i32>) -> tensor<i32>
func.return
}
// -----
// Test legalization of `tf.XlaSendToHost` expands into a `mhlo.send` op.
// CHECK-LABEL: func @send_to_host
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @send_to_host(%arg0: tensor<i32>) {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: "mhlo.send"([[ARG0]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: is_host_transfer = true
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_key_dtoh_0"}
// CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
// CHECK-SAME: (tensor<i32>, !mhlo.token) -> !mhlo.token
"tf.XlaSendToHost"(%arg0) {key = "send_key"} : (tensor<i32>) -> ()
func.return
}
// -----
// Test legalization of `tf.XlaRecvFromHost` expands into a `mhlo.recv` op.
// CHECK-LABEL: func @recv_from_host
func.func @recv_from_host() -> tensor<i32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: is_host_transfer = true
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_key_htod_0"}
// CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
// CHECK-SAME: (!mhlo.token) -> (tensor<i32>, !mhlo.token)
%0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>
// CHECK: return [[RECV_TUPLE]]#0 : tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// Test legalization of multiple TF/XLA communication ops are sequenced with
// their generated tokens. Channel Id's are also uniquely assigned.
// CHECK-LABEL: func @multiple_consecutive_ops
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @multiple_consecutive_ops(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[SEND0_ARG0_TOKEN:%.*]] = "mhlo.send"([[ARG0]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send0_dtoh_0"}
// CHECK: [[RECV0_RETVAL0_TUPLE:%.*]]:2 = "mhlo.recv"([[SEND0_ARG0_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv0_htod_0"}
%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "recv0", send_key = "send0", host_mlir_module = ""} : (tensor<i32>) -> tensor<i32>
// CHECK: [[SEND1_ARG0_TOKEN:%.*]] = "mhlo.send"([[RECV0_RETVAL0_TUPLE]]#0, [[RECV0_RETVAL0_TUPLE]]#1)
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send1_dtoh_0"}
// CHECK: [[RECV1_RETVAL0_TUPLE:%.*]]:2 = "mhlo.recv"([[SEND1_ARG0_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv1_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%0) {recv_key = "recv1", send_key = "send1", host_mlir_module = ""} : (tensor<i32>) -> tensor<i32>
// CHECK: return [[RECV1_RETVAL0_TUPLE]]#0 : tensor<i32>
func.return %1 : tensor<i32>
}
// -----
// Test private function with TF/XLA communication op being called by another
// function gets rewritten with an extra token argument and an extra token
// result, and the caller passes in a token. The top level function not called
// (or public) will be updated to create a token.
// CHECK: func @main([[MAIN_ARG0:%.*]]: tensor<i32>) -> tensor<i32>
func.func @main(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: [[MAIN_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[CALL:%.*]]:2 = call @callee([[MAIN_ARG0]], [[MAIN_TOKEN]])
// CHECK-SAME: (tensor<i32>, !mhlo.token) -> (tensor<i32>, !mhlo.token)
%0 = func.call @callee(%arg0) : (tensor<i32>) -> tensor<i32>
// CHECK: return [[CALL]]#0
func.return %0 : tensor<i32>
}
// CHECK: func private @callee([[CALLEE_ARG0:%.*]]: tensor<i32>, [[CALLEE_ARG1:%.*]]: !mhlo.token) -> (tensor<i32>, !mhlo.token)
func.func private @callee(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK-NOT: mhlo.create_token
// CHECK: [[SEND_ARG0_TOKEN:%.*]] = "mhlo.send"([[CALLEE_ARG0]], [[CALLEE_ARG1]])
// CHECK: [[RECV_RETVAL0_TUPLE:%.*]]:2 = "mhlo.recv"([[SEND_ARG0_TOKEN]])
%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "recv", send_key = "send", host_mlir_module = ""} : (tensor<i32>) -> tensor<i32>
// CHECK: return [[RECV_RETVAL0_TUPLE]]#0, [[RECV_RETVAL0_TUPLE]]#1
func.return %0 : tensor<i32>
}
// -----
// Test public function with TF/XLA communication op being called by another
// function. The original public function will be modified to create a token,
// while the function is cloned and rewritten with an extra token argument and
// an extra token result. All callers to the original function are updated to
// point to the cloned function and the function the caller is in is updated to
// pass a token or create a token.
// CHECK: func @main([[MAIN_ARG0:%.*]]: tensor<i32>) -> tensor<i32>
func.func @main(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: [[MAIN_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[CALL:%.*]]:2 = call [[CALLEE_CLONE:@.*]]([[MAIN_ARG0]], [[MAIN_TOKEN]])
// CHECK-SAME: (tensor<i32>, !mhlo.token) -> (tensor<i32>, !mhlo.token)
%0 = func.call @callee(%arg0) : (tensor<i32>) -> tensor<i32>
// CHECK: return [[CALL]]#0 : tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK: func @callee([[CALLEE_ARG0:%.*]]: tensor<i32>) -> tensor<i32>
func.func @callee(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: [[CALLEE_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[SEND_ARG0_TOKEN:%.*]] = "mhlo.send"([[CALLEE_ARG0]], [[CALLEE_TOKEN]])
// CHECK: [[RECV_RETVAL0_TUPLE:%.*]]:2 = "mhlo.recv"([[SEND_ARG0_TOKEN]])
%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "recv", send_key = "send", host_mlir_module = ""} : (tensor<i32>) -> tensor<i32>
// CHECK: return [[RECV_RETVAL0_TUPLE]]#0
func.return %0 : tensor<i32>
}
// CHECK: func private [[CALLEE_CLONE]]([[CALLEE_CLONE_ARG0:%.*]]: tensor<i32>, [[CALLEE_CLONE_ARG1:%.*]]: !mhlo.token) -> (tensor<i32>, !mhlo.token)
// CHECK-NOT: mhlo.create_token
// CHECK: [[CLONE_SEND_ARG0_TOKEN:%.*]] = "mhlo.send"([[CALLEE_CLONE_ARG0]], [[CALLEE_CLONE_ARG1]])
// CHECK: [[CLONE_RECV_RETVAL0_TUPLE:%.*]]:2 = "mhlo.recv"([[CLONE_SEND_ARG0_TOKEN]])
// CHECK: return [[CLONE_RECV_RETVAL0_TUPLE]]#0, [[CLONE_RECV_RETVAL0_TUPLE]]#1
// -----
// Tests generated tokens are passed into a function call that also has TF/XLA
// communication ops.
// CHECK: func @main([[MAIN_ARG0:%.*]]: tensor<i32>)
func.func @main(%arg0: tensor<i32>) {
// CHECK: [[MAIN_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[MAIN_SEND0_TOKEN:%.*]] = "mhlo.send"([[MAIN_ARG0]], [[MAIN_TOKEN]])
"tf.XlaSendToHost"(%arg0) {key = "send0"} : (tensor<i32>) -> ()
// CHECK: [[CALL_TOKEN:%.*]] = call @callee([[MAIN_SEND0_TOKEN]])
// CHECK-SAME: (!mhlo.token) -> !mhlo.token
func.call @callee() : () -> ()
// CHECK: [[MAIN_SEND2_TOKEN:%.*]] = "mhlo.send"([[MAIN_ARG0]], [[CALL_TOKEN]])
"tf.XlaSendToHost"(%arg0) {key = "send2"} : (tensor<i32>) -> ()
func.return
}
// CHECK: func private @callee([[CALLEE_ARG0:%.*]]: !mhlo.token) -> !mhlo.token
func.func private @callee() {
// CHECK-NOT: mhlo.create_token
// CHECK: [[ZERO:%.*]] = mhlo.constant dense<0>
%0 = mhlo.constant dense<0> : tensor<i32>
// CHECK: [[CALLEE_SEND_TOKEN:%.*]] = "mhlo.send"([[ZERO]], [[CALLEE_ARG0]])
"tf.XlaSendToHost"(%0) {key = "send1"} : (tensor<i32>) -> ()
// CHECK: return [[CALLEE_SEND_TOKEN]]
func.return
}
// -----
// Test only the top level function generates a token.
// CHECK: func private @callee0()
func.func private @callee0() {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: call @callee1([[INIT_TOKEN]])
func.call @callee1() : () -> ()
func.return
}
// CHECK: func private @callee1([[CALLEE1_ARG0:%.*]]: !mhlo.token) -> !mhlo.token
func.func private @callee1() {
// CHECK-NOT: mhlo.create_token
// CHECK: [[CALL_2:%.*]] = call @callee2([[CALLEE1_ARG0]])
func.call @callee2() : () -> ()
// CHECK: return [[CALL_2]]
func.return
}
// CHECK: func private @callee2([[CALLEE2_ARG0:%.*]]: !mhlo.token) -> !mhlo.token
func.func private @callee2() {
// CHECK-NOT: mhlo.create_token
// CHECK: [[RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[CALLEE2_ARG0]])
%0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>
// CHECK: return [[RECV_TUPLE]]#1
func.return
}
// -----
// Test cloned function rewrite also checks transitive function calls to
// TF/XLA communication ops.
// CHECK: func @callee3()
func.func @callee3() {
// CHECK: [[CALLEE3_INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: call @callee4{{.+}}([[CALLEE3_INIT_TOKEN]])
func.call @callee4() : () -> ()
func.return
}
// CHECK: func @callee4()
func.func @callee4() {
// CHECK: [[CALLEE4_INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[CALL_5:%.*]] = call @callee5([[CALLEE4_INIT_TOKEN]])
func.call @callee5() : () -> ()
// CHECK: return
func.return
}
// CHECK: func private @callee5([[CALLEE5_ARG0:%.*]]: !mhlo.token) -> !mhlo.token
func.func private @callee5() {
// CHECK-NOT: mhlo.create_token
// CHECK: [[RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[CALLEE5_ARG0]])
%0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>
// CHECK: return [[RECV_TUPLE]]#1
func.return
}
// CHECK: func private @callee4{{.+}}([[CALLEE4_ARG0:%.*]]: !mhlo.token) -> !mhlo.token
// CHECK-NOT: mhlo.create_token
// CHECK: [[CALL_5:%.*]] = call @callee5([[CALLEE4_ARG0]])
// CHECK: return [[CALL_5]]
// -----
// Tests `mhlo.if` with branches populated with TF/XLA communication ops.
// CHECK-LABEL: func @if_both_branches
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>, [[ARG2:%.*]]: tensor<f32>)
func.func @if_both_branches(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[IF:%.*]]:2 = "mhlo.if"([[ARG0]])
%0 = "mhlo.if"(%arg0) ({
// CHECK: [[TRUE_SEND_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_if_true_dtoh_0"}
// CHECK: [[TRUE_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[TRUE_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_if_true_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg1) {recv_key = "recv_if_true", send_key = "send_if_true", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: mhlo.return [[TRUE_RECV_TUPLE]]#0, [[TRUE_RECV_TUPLE]]#1
"mhlo.return"(%1) : (tensor<f32>) -> ()
}, {
// CHECK: [[FALSE_SEND_TOKEN:%.*]] = "mhlo.send"([[ARG2]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_if_false_dtoh_0"}
// CHECK: [[FALSE_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[FALSE_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_if_false_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg2) {recv_key = "recv_if_false", send_key = "send_if_false", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: mhlo.return [[FALSE_RECV_TUPLE]]#0, [[FALSE_RECV_TUPLE]]#1
"mhlo.return"(%1) : (tensor<f32>) -> ()
// CHECK: (tensor<i1>) -> (tensor<f32>, !mhlo.token)
}) : (tensor<i1>) -> tensor<f32>
// CHECK: return [[IF]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.if` with only the `true` branch populated with TF/XLA
// communication ops.
// CHECK-LABEL: func @if_true_branch
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>, [[ARG2:%.*]]: tensor<f32>)
func.func @if_true_branch(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[IF:%.*]]:2 = "mhlo.if"([[ARG0]])
%0 = "mhlo.if"(%arg0) ({
// CHECK: [[TRUE_SEND_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_if_true_dtoh_0"}
// CHECK: [[TRUE_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[TRUE_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_if_true_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg1) {recv_key = "recv_if_true", send_key = "send_if_true", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: mhlo.return [[TRUE_RECV_TUPLE]]#0, [[TRUE_RECV_TUPLE]]#1
"mhlo.return"(%1) : (tensor<f32>) -> ()
}, {
// CHECK: mhlo.return [[ARG2]], [[INIT_TOKEN]]
"mhlo.return"(%arg2) : (tensor<f32>) -> ()
// CHECK: (tensor<i1>) -> (tensor<f32>, !mhlo.token)
}) : (tensor<i1>) -> tensor<f32>
// CHECK: return [[IF]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.if` with only the `false` branch populated with TF/XLA
// communication ops.
// CHECK-LABEL: func @if_false_branch
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>, [[ARG2:%.*]]: tensor<f32>)
func.func @if_false_branch(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[IF:%.*]]:2 = "mhlo.if"([[ARG0]])
%0 = "mhlo.if"(%arg0) ({
// CHECK: mhlo.return [[ARG1]], [[INIT_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}, {
// CHECK: [[FALSE_SEND_TOKEN:%.*]] = "mhlo.send"([[ARG2]], [[INIT_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_if_false_dtoh_0"}
// CHECK: [[FALSE_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[FALSE_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_if_false_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg2) {recv_key = "recv_if_false", send_key = "send_if_false", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: mhlo.return [[FALSE_RECV_TUPLE]]#0, [[FALSE_RECV_TUPLE]]#1
"mhlo.return"(%1) : (tensor<f32>) -> ()
// CHECK: (tensor<i1>) -> (tensor<f32>, !mhlo.token)
}) : (tensor<i1>) -> tensor<f32>
// CHECK: return [[IF]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.if` with tuple arg from a `mhlo.tuple` only used by `mhlo.if` is
// replaced.
// CHECK-LABEL: func @if_replace_tuple_arg
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>, [[ARG2:%.*]]: tensor<f32>)
func.func @if_replace_tuple_arg(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// CHECK-NOT: mhlo.tuple [[ARG1]], [[ARG2]]
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[IF:%.*]] = "mhlo.if"([[ARG0]])
%1 = "mhlo.if"(%arg0) ({
"tf.XlaSendToHost"(%arg1) {key = "send_key"} : (tensor<f32>) -> ()
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}, {
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}) : (tensor<i1>) -> tensor<f32>
func.return %1 : tensor<f32>
}
// -----
// Tests `mhlo.if` with tuple arg not from a `mhlo.tuple` is unpacked.
// CHECK-LABEL: func @if_unpack_tuple_arg
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tuple<tensor<f32>, tensor<f32>>)
func.func @if_unpack_tuple_arg(%arg0: tensor<i1>, %arg1: tuple<tensor<f32>, tensor<f32>>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK-DAG: [[IF_ARG_ELEMENT0:%.*]] = mhlo.get_tuple_element [[ARG1]][0]
// CHECK-DAG: [[IF_ARG_ELEMENT1:%.*]] = mhlo.get_tuple_element [[ARG1]][1]
%0 = "mhlo.get_tuple_element"(%arg1) {index = 0 : i32} : (tuple<tensor<f32>, tensor<f32>>) -> tensor<f32>
%1 = "mhlo.get_tuple_element"(%arg1) {index = 1 : i32} : (tuple<tensor<f32>, tensor<f32>>) -> tensor<f32>
// CHECK: [[IF:%.*]] = "mhlo.if"([[ARG0]])
%2 = "mhlo.if"(%arg0) ({
"tf.XlaSendToHost"(%0) {key = "send_key"} : (tensor<f32>) -> ()
"mhlo.return"(%0) : (tensor<f32>) -> ()
}, {
"mhlo.return"(%0) : (tensor<f32>) -> ()
}) : (tensor<i1>) -> tensor<f32>
func.return %2 : tensor<f32>
}
// -----
// Tests `mhlo.if` tuple result is extended with a `mhlo.token`.
// CHECK-LABEL: func @if_extend_tuple_result
func.func @if_extend_tuple_result(%arg0: tensor<i1>, %arg1: tuple<tensor<f32>, tensor<f32>>) -> tuple<tensor<f32>, tensor<f32>> {
%0 = "mhlo.get_tuple_element"(%arg1) {index = 0 : i32} : (tuple<tensor<f32>, tensor<f32>>) -> tensor<f32>
%1 = "mhlo.get_tuple_element"(%arg1) {index = 1 : i32} : (tuple<tensor<f32>, tensor<f32>>) -> tensor<f32>
// CHECK: [[IF:%.*]]:3 = "mhlo.if"
%2:2 = "mhlo.if"(%arg0) ({
"tf.XlaSendToHost"(%0) {key = "send_key"} : (tensor<f32>) -> ()
"mhlo.return"(%0, %1) : (tensor<f32>, tensor<f32>) -> ()
}, {
"mhlo.return"(%0, %1) : (tensor<f32>, tensor<f32>) -> ()
// CHECK: (tensor<i1>) -> (tensor<f32>, tensor<f32>, !mhlo.token)
}) : (tensor<i1>) -> (tensor<f32>, tensor<f32>)
// CHECK: [[IF_SUBTUPLE_RESULT:%.*]] = mhlo.tuple [[IF]]#0, [[IF]]#1
%3 = "mhlo.tuple"(%2#0, %2#1) : (tensor<f32>, tensor<f32>) -> tuple<tensor<f32>, tensor<f32>>
// CHECK: return [[IF_SUBTUPLE_RESULT]]
func.return %3 : tuple<tensor<f32>, tensor<f32>>
}
// -----
// Tests nested `mhlo.if` containing TF/XLA communication ops.
// CHECK-LABEL: func @if_nested
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>)
func.func @if_nested(%arg0: tensor<i1>, %arg1: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[IF:%.*]]:2 = "mhlo.if"([[ARG0]])
%0 = "mhlo.if"(%arg0) ({
%1 = mhlo.constant dense<false> : tensor<i1>
// CHECK: [[INNER_IF:%.*]]:2 = "mhlo.if"({{%.*}})
%2 = "mhlo.if"(%1) ({
// CHECK: [[SEND_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
"tf.XlaSendToHost"(%arg1) {key = "send_key"} : (tensor<f32>) -> ()
// CHECK: mhlo.return [[ARG1]], [[SEND_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
// CHECK-NEXT: }, {
}, {
// CHECK: mhlo.return [[ARG1]], [[INIT_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
// CHECK-NEXT: (tensor<i1>) -> (tensor<f32>, !mhlo.token)
}) : (tensor<i1>) -> tensor<f32>
// CHECK: mhlo.return [[ARG1]], [[INNER_IF]]#1
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
// CHECK-NEXT: }, {
}, {
// CHECK: mhlo.return [[ARG1]], [[INIT_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
// CHECK-NEXT: (tensor<i1>) -> (tensor<f32>, !mhlo.token)
}) : (tensor<i1>) -> tensor<f32>
// CHECK-NEXT: return [[IF]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.if` containing a function call to TF/XLA communication ops.
// CHECK-LABEL: func @if_function_call
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>)
func.func @if_function_call(%arg0: tensor<i1>, %arg1: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: "mhlo.if"
%0 = "mhlo.if"(%arg0) ({
// CHECK: [[CALL_TOKEN:%.*]] = func.call @callee([[ARG1]], [[INIT_TOKEN]])
func.call @callee(%arg1) : (tensor<f32>) -> ()
// CHECK: mhlo.return [[ARG1]], [[CALL_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}, {
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}) : (tensor<i1>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func private @callee
// CHECK-SAME: ([[CALLEE_ARG0:%.*]]: tensor<f32>, [[CALLEE_ARG1:%.*]]: !mhlo.token) -> !mhlo.token
func.func private @callee(%arg0: tensor<f32>) {
// CHECK: [[SEND_TOKEN:%.*]] = "mhlo.send"
"tf.XlaSendToHost"(%arg0) {key = "send_key"} : (tensor<f32>) -> ()
// CHECK: return [[SEND_TOKEN]]
func.return
}
// -----
// Tests `mhlo.if` containing multiple TF/XLA communication ops.
// CHECK-LABEL: func @if_region_multiple_ops
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>)
func.func @if_region_multiple_ops(%arg0: tensor<i1>, %arg1: tensor<f32>) {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: "mhlo.if"
%0 = "mhlo.if"(%arg0) ({
// CHECK: [[SEND0_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
"tf.XlaSendToHost"(%arg1) {key = "send_key0"} : (tensor<f32>) -> ()
// CHECK: [[SEND1_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[SEND0_TOKEN]])
"tf.XlaSendToHost"(%arg1) {key = "send_key1"} : (tensor<f32>) -> ()
// CHECK: mhlo.return [[ARG1]], [[SEND1_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}, {
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}) : (tensor<i1>) -> tensor<f32>
func.return
}
// -----
// Tests `mhlo.if` containing TF/XLA communication ops followed by other TF/XLA
// communication ops.
// CHECK-LABEL: func @if_followed_by_communication_op
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i1>, [[ARG1:%.*]]: tensor<f32>)
func.func @if_followed_by_communication_op(%arg0: tensor<i1>, %arg1: tensor<f32>) {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK-NEXT: [[IF:%.*]]:2 = "mhlo.if"
%0 = "mhlo.if"(%arg0) ({
// CHECK-NEXT: [[SEND0_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
"tf.XlaSendToHost"(%arg1) {key = "send_key0"} : (tensor<f32>) -> ()
// CHECK-NEXT: mhlo.return [[ARG1]], [[SEND0_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
// CHECK-NEXT: }, {
}, {
// CHECK-NEXT: mhlo.return [[ARG1]], [[INIT_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
// CHECK-NEXT: })
}) : (tensor<i1>) -> tensor<f32>
// CHECK-NEXT: "mhlo.send"({{.*}}, [[IF]]#1)
"tf.XlaSendToHost"(%arg1) {key = "send_key1"} : (tensor<f32>) -> ()
func.return
}
// -----
// Tests `mhlo.while` with cond and body populated with TF/XLA communication
// ops.
// CHECK-LABEL: func @while_cond_body
// CHECK-SAME: ([[ARG0:%.*]]: tensor<f32>)
func.func @while_cond_body(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[WHILE:%.*]]:2 = mhlo.while([[ITER_ARG_VALUE:.*]] = [[ARG0]], [[ITER_ARG_TOKEN:.*]] = [[INIT_TOKEN]])
%0 = "mhlo.while"(%arg0) ({
^bb0(%arg1: tensor<f32>):
// CHECK: [[COND_SEND_TOKEN:%.*]] = "mhlo.send"([[ITER_ARG_VALUE]], [[ITER_ARG_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_while_cond_dtoh_0"}
// CHECK: [[COND_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[COND_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_while_cond_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg1) {recv_key = "recv_while_cond", send_key = "send_while_cond", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: [[COND_COMPARE:%.*]] = mhlo.compare LT, [[COND_RECV_TUPLE]]#0, [[COND_RECV_TUPLE]]#0
%2 = "mhlo.compare"(%1, %1) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
// CHECK: mhlo.return [[COND_COMPARE]]
"mhlo.return"(%2) : (tensor<i1>) -> ()
}, {
^bb0(%arg1: tensor<f32>):
// CHECK: [[BODY_SEND_TOKEN:%.*]] = "mhlo.send"([[ITER_ARG_VALUE]], [[ITER_ARG_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_while_body_dtoh_0"}
// CHECK: [[BODY_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[BODY_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_while_body_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg1) {recv_key = "recv_while_body", send_key = "send_while_body", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: mhlo.return [[BODY_RECV_TUPLE]]#0, [[BODY_RECV_TUPLE]]#1
"mhlo.return"(%1) : (tensor<f32>) -> ()
}) : (tensor<f32>) -> tensor<f32>
// CHECK: return [[WHILE]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.while` with only the `cond` region populated with TF/XLA
// communication ops.
// CHECK-LABEL: func @while_cond
// CHECK-SAME: ([[ARG0:%.*]]: tensor<f32>)
func.func @while_cond(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[WHILE:%.*]]:2 = mhlo.while([[ITER_ARG_VALUE:.*]] = [[ARG0]], [[ITER_ARG_TOKEN:.*]] = [[INIT_TOKEN]])
%0 = "mhlo.while"(%arg0) ({
^bb0(%arg1: tensor<f32>):
// CHECK: [[COND_SEND_TOKEN:%.*]] = "mhlo.send"([[ITER_ARG_VALUE]], [[ITER_ARG_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_while_cond_dtoh_0"}
// CHECK: [[COND_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[COND_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_while_cond_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg1) {recv_key = "recv_while_cond", send_key = "send_while_cond", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: [[COND_COMPARE:%.*]] = mhlo.compare LT, [[COND_RECV_TUPLE]]#0, [[COND_RECV_TUPLE]]#0
%2 = "mhlo.compare"(%1, %1) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
// CHECK: mhlo.return [[COND_COMPARE]]
"mhlo.return"(%2) : (tensor<i1>) -> ()
}, {
^bb0(%arg1: tensor<f32>):
// CHECK: mhlo.return [[ITER_ARG_VALUE]], [[ITER_ARG_TOKEN]]
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}) : (tensor<f32>) -> tensor<f32>
// CHECK: return [[WHILE]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.while` with only the `body` region populated with TF/XLA
// communication ops.
// CHECK-LABEL: func @while_body
// CHECK-SAME: ([[ARG0:%.*]]: tensor<f32>)
func.func @while_body(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token
// CHECK: [[WHILE:%.*]]:2 = mhlo.while([[ITER_ARG_VALUE:.*]] = [[ARG0]], [[ITER_ARG_TOKEN:.*]] = [[INIT_TOKEN]])
%0 = "mhlo.while"(%arg0) ({
^bb0(%arg1: tensor<f32>):
// CHECK: [[COND_COMPARE:%.*]] = mhlo.compare LT, [[ITER_ARG_VALUE]], [[ITER_ARG_VALUE]]
%2 = "mhlo.compare"(%arg1, %arg1) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
// CHECK: mhlo.return [[COND_COMPARE]]
"mhlo.return"(%2) : (tensor<i1>) -> ()
}, {
^bb0(%arg1: tensor<f32>):
// CHECK: [[BODY_SEND_TOKEN:%.*]] = "mhlo.send"([[ITER_ARG_VALUE]], [[ITER_ARG_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 2
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "send_while_body_dtoh_0"}
// CHECK: [[BODY_RECV_TUPLE:%.*]]:2 = "mhlo.recv"([[BODY_SEND_TOKEN]])
// CHECK-SAME: channel_handle = #mhlo.channel_handle
// CHECK-SAME: handle =
// CHECK-SAME: type = 3
// CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_while_body_htod_0"}
%1 = "tf._XlaHostComputeMlir"(%arg1) {recv_key = "recv_while_body", send_key = "send_while_body", host_mlir_module = ""} : (tensor<f32>) -> tensor<f32>
// CHECK: mhlo.return [[BODY_RECV_TUPLE]]#0, [[BODY_RECV_TUPLE]]#1
"mhlo.return"(%1) : (tensor<f32>) -> ()
}) : (tensor<f32>) -> tensor<f32>
// CHECK: return [[WHILE]]#0
func.return %0 : tensor<f32>
}
// -----
// Tests `mhlo.while` containing TF/XLA communication ops followed by other
// TF/XLA communication ops.
// CHECK-LABEL: func @while_followed_by_communication_op
func.func @while_followed_by_communication_op(%arg0: tensor<f32>) {
// CHECK: [[WHILE:%.*]]:2 = mhlo.while
%0 = "mhlo.while"(%arg0) ({
^bb0(%arg1: tensor<f32>):
"tf.XlaSendToHost"(%arg1) {key = "send_key0"} : (tensor<f32>) -> ()
%1 = "mhlo.compare"(%arg1, %arg1) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
"mhlo.return"(%1) : (tensor<i1>) -> ()
}, {
^bb0(%arg1: tensor<f32>):
"mhlo.return"(%arg1) : (tensor<f32>) -> ()
}) : (tensor<f32>) -> tensor<f32>
// CHECK: "mhlo.send"({{.*}}, [[WHILE]]#1)
"tf.XlaSendToHost"(%arg0) {key = "send_key1"} : (tensor<f32>) -> ()
func.return
}
// -----
// Tests unsupported parent of TF/XLA communication op.
func.func @unsupported_ancestor(%arg0: tensor<?x?xf32>, %arg1: tensor<f32>) {
%0 = "mhlo.reduce"(%arg0, %arg1) ({
^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
%1 = mhlo.add %arg2, %arg3 : tensor<f32>
// expected-error@+1 {{expects ancestor(s) to be of ['mhlo.if', 'func.func']}}
"tf._XlaHostComputeMlir"() {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : () -> ()
"mhlo.return"(%1) : (tensor<f32>) -> ()
}) {dimensions = dense<[1]> : tensor<1xi64>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?xf32>
func.return
}
// -----
// Tests transitive unsupported parent of TF/XLA communication op.
func.func @unsupported_ancestor(%arg0: tensor<?x?xf32>, %arg1: tensor<f32>) {
%0 = "mhlo.reduce"(%arg0, %arg1) ({
^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
%1 = mhlo.add %arg2, %arg3 : tensor<f32>
// expected-error@+1 {{expects ancestor(s) to be of ['mhlo.if', 'func.func']}}
func.call @callee() : () -> ()
"mhlo.return"(%1) : (tensor<f32>) -> ()
}) {dimensions = dense<[1]> : tensor<1xi64>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?xf32>
func.return
}
func.func private @callee() {
"tf._XlaHostComputeMlir"() {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : () -> ()
func.return
}
// -----
// Tests function with more than one block that is to be rewritten emits an
// error instead.
// expected-error@+1 {{'func.func' ops with more than one block are not supported}}
func.func @multi_block_func() {
cf.br ^bb1
^bb1:
%0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>
func.return
}
// -----
// CHECK-LABEL: func @host_compute_manual_sharding
func.func @host_compute_manual_sharding(%arg0: tensor<i32>) {
// CHECK: "mhlo.send"
// CHECK-SAME: mhlo.sharding = "\08\04"
// CHECK: "mhlo.recv"
// CHECK-SAME: mhlo.sharding = "\08\04"
%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = "", manual_sharding = true} : (tensor<i32>) -> tensor<i32>
func.return
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-include-tf2xla-fallback.mlir
+72
View File
@@ -0,0 +1,72 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true device-type=XLA_CPU_JIT" -verify-diagnostics %s | FileCheck --check-prefix SUPPORTED_FALLBACK_DEVICE %s
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true" %s | FileCheck --check-prefix UNSPECIFIED_FALLBACK_DEVICE %s
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true device-type=INVALID_DEVICE_TYPE" %s | FileCheck --check-prefix UNSUPPORTED_FALLBACK_DEVICE %s
// We run this test four times:
// 1) Legalize without using TF2XLA fallback (ops cannot be legalized).
// 2) Use fallback with a device that supports all ops (ops can be legalized).
// 3) Use fallback with unspecified device (ops cannot be legalized).
// 4) Use fallback with specified but unsupported device (ops cannot be legalized).
//
// Note: For 3) and 4) we do not use `-verify-diagnostics` because these cases
// produce remarks that don't occur for 1) and 2) and there is no way to check
// the remarks only for 3) and 4) (except using two files).
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
// CHECK-LABEL: non_max_suppression_v4
func.func @non_max_suppression_v4(%arg0: tensor<3x4xf32>, %arg1: tensor<3xf32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<2xi32> {
%max_size = mhlo.constant dense<2> : tensor<i32>
// NO_FALLBACK: tf.NonMaxSuppressionV4
// SUPPORTED_FALLBACK_DEVICE-NOT: tf.NonMaxSuppressionV4
// UNSPECIFIED_FALLBACK_DEVICE: tf.NonMaxSuppressionV4
// UNSUPPORTED_FALLBACK_DEVICE: tf.NonMaxSuppressionV4
%0:2 = "tf.NonMaxSuppressionV4"(%arg0, %arg1, %max_size, %arg2, %arg3) {pad_to_max_output_size = true}: (tensor<3x4xf32>, tensor<3xf32>, tensor<i32>, tensor<f32>, tensor<f32>) -> (tensor<2xi32>, tensor<i32>)
func.return %0#0 : tensor<2xi32>
}
// CHECK-LABEL: mirror_pad
func.func @mirror_pad(%arg0: tensor<2x3xcomplex<f64>>) -> tensor<4x7xcomplex<f64>> {
%0 = mhlo.constant dense<[[1, 1], [2, 2]]> : tensor<2x2xi32>
// NO_FALLBACK: tf.MirrorPad
// SUPPORTED_FALLBACK_DEVICE-NOT: tf.MirrorPad
// UNSPECIFIED_FALLBACK_DEVICE: tf.MirrorPad
// UNSUPPORTED_FALLBACK_DEVICE: tf.MirrorPad
%1 = "tf.MirrorPad"(%arg0, %0) {mode = "SYMMETRIC"} : (tensor<2x3xcomplex<f64>>, tensor<2x2xi32>) -> tensor<4x7xcomplex<f64>>
func.return %1 : tensor<4x7xcomplex<f64>>
}
// BatchMatMulV2 has native as well as fallback lowering patterns available.
// The fallback pattern uses dot_general without broadcast on operands and then
// transposes the output which is faster. However, the fallback pattern doesn't
// support dynamic shaped operands like the native lowering. Verify that
// fallback lowering is preferred for static shaped operands when available.
// CHECK-LABEL: batchmatmulv2
func.func @batchmatmulv2(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// NO_FALLBACK: mhlo.dynamic_broadcast_in_dim
// NO_FALLBACK: mhlo.dot_general
// SUPPORTED_FALLBACK_DEVICE: mhlo.reduce
// SUPPORTED_FALLBACK_DEVICE: mhlo.dot_general
// SUPPORTED_FALLBACK_DEVICE: mhlo.transpose
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
func.return %0 : tensor<3x4x4xf32>
}
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-prefer-tf2xla.mlir
+244
View File
@@ -0,0 +1,244 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT legalize-chlo=false use-tf2xla-fallback=true prefer-tf2xla=true" %s | FileCheck %s
// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT legalize-chlo=false prefer-tf2xla=true" %s | FileCheck --check-prefix NOFALLBACK %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
// CHECK-LABEL: @abs
func.func @abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NOT: tf.Abs
%0 = "tf.Abs"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// -----
// CHECK-LABEL: func @testBroadcastGradientArgs
func.func @testBroadcastGradientArgs(%s0: tensor<4xi32>, %s1: tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>) {
// CHECK: tf.BroadcastGradientArgs
%r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) : (tensor<4xi32>, tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>)
func.return %r0, %r1 : tensor<1xi32>, tensor<0xi32>
}
// -----
// CHECK-LABEL: @acos
func.func @acos(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NOT: tf.Acos
%0 = "tf.Acos"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// -----
// NOFALLBACK-LABEL: @xla_svd
func.func @xla_svd(%arg0: tensor<1x1xf32>) -> (tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) {
// NOFALLBACK: XlaSvd
%s, %u, %v = "tf.XlaSvd"(%arg0) {max_iter = 1, epsilon = 1.0E-09 : f32, precision_config = ""} : (tensor<1x1xf32>) -> (tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>)
func.return %s, %u, %v : tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>
}
//===----------------------------------------------------------------------===//
// Random op legalizations.
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @random_uniform_simple
func.func @random_uniform_simple(%arg0: tensor<3xi32>) -> tensor<12x?x64xf32> {
// expected-remark@+1 {{lowering requires operand #0 to be a constant}}
%0 = "tf.RandomUniform"(%arg0) {device = "", seed = 0 : i64, seed2 = 0 : i64} : (tensor<3xi32>) -> tensor<12x?x64xf32>
func.return %0 : tensor<12x?x64xf32>
}
// -----
// CHECK-LABEL: func @random_uniform_with_seeds
func.func @random_uniform_with_seeds(%arg0: tensor<4xi32>) -> tensor<32x12x12x64xf32> {
// CHECK: %0 = mhlo.constant dense<[32, 12, 12, 64]> : tensor<4xi32>
// CHECK-NEXT: %1 = mhlo.constant dense<[32, 12, 12, 64]> : tensor<4xi32>
// CHECK-NEXT: %2 = "mhlo.slice"(%1) <{limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1xi32>
// CHECK-NEXT: %3 = mhlo.reshape %2 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %4 = mhlo.convert %3 : tensor<i32>
// CHECK-NEXT: %5 = "mhlo.slice"(%1) <{limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1xi32>
// CHECK-NEXT: %6 = mhlo.reshape %5 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %7 = mhlo.convert %6 : tensor<i32>
// CHECK-NEXT: %8 = "mhlo.slice"(%1) <{limit_indices = dense<3> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1xi32>
// CHECK-NEXT: %9 = mhlo.reshape %8 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %10 = mhlo.convert %9 : tensor<i32>
// CHECK-NEXT: %11 = "mhlo.slice"(%1) <{limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<3> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1xi32>
// CHECK-NEXT: %12 = mhlo.reshape %11 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %13 = mhlo.convert %12 : tensor<i32>
// CHECK-NEXT: %14 = mhlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-NEXT: %15 = mhlo.constant dense<1.000000e+00> : tensor<f32>
// CHECK-NEXT: %16 = mhlo.constant dense<[32, 12, 12, 64]> : tensor<4xi64>
// CHECK-NEXT: %17 = "mhlo.rng"(%14, %15, %16) <{rng_distribution = #mhlo.rng_distribution<UNIFORM>}> : (tensor<f32>, tensor<f32>, tensor<4xi64>) -> tensor<32x12x12x64xf32>
%cst = "tf.Const"() {value = dense<[32, 12, 12, 64]> : tensor<4xi32>} : () -> tensor<4xi32>
%0 = "tf.RandomUniform"(%cst) {seed = 87654321 : i64, seed2 = 0 : i64} : (tensor<4xi32>) -> tensor<32x12x12x64xf32>
// CHECK: return %17 : tensor<32x12x12x64xf32>
func.return %0 : tensor<32x12x12x64xf32>
}
// CHECK-LABEL: slice_variable_start
func.func @slice_variable_start(%arg0: tensor<3x4xi32>, %arg1: tensor<2xi64>) -> tensor<1x4xi32> {
// CHECK: %[[SLICED_START1:.*]] = "mhlo.slice"(%arg1)
// CHECK-DAG-SAME: {limit_indices = dense<1> : tensor<1xi64>,
// CHECK-DAG-SAME: start_indices = dense<0> : tensor<1xi64>,
// CHECK-DAG-SAME: strides = dense<1> : tensor<1xi64>} : (tensor<2xi64>) -> tensor<1xi64>
// CHECK: %[[RESHAPED_START1:.*]] = mhlo.reshape %[[SLICED_START1]] : (tensor<1xi64>) -> tensor<i64>
// CHECK: %[[SLICED_START2:.*]] = "mhlo.slice"(%arg1)
// CHECK-DAG-SAME: {limit_indices = dense<2> : tensor<1xi64>,
// CHECK-DAG-SAME: start_indices = dense<1> : tensor<1xi64>,
// CHECK-DAG-SAME: strides = dense<1> : tensor<1xi64>} : (tensor<2xi64>) -> tensor<1xi64>
// CHECK: %[[RESHAPED_START2:.*]] = mhlo.reshape %[[SLICED_START2]] : (tensor<1xi64>) -> tensor<i64>
// CHECK: %[[RESULT:.*]] = "mhlo.dynamic_slice"(%arg0, %[[RESHAPED_START1]], %[[RESHAPED_START2]]) <{slice_sizes = dense<[1, 4]> : tensor<2xi64>}> : (tensor<3x4xi32>, tensor<i64>, tensor<i64>) -> tensor<1x4xi32>
// CHECK: return %[[RESULT]] : tensor<1x4xi32>
%sizes = "tf.Const"() {value = dense<[1, 4]> : tensor<2xi64>} : () -> (tensor<2xi64>)
%0 = "tf.Slice"(%arg0, %arg1, %sizes) : (tensor<3x4xi32>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x4xi32>
func.return %0 : tensor<1x4xi32>
}
// CHECK-LABEL: slice_variable_start_negsize
func.func @slice_variable_start_negsize(%arg0: tensor<3x4xi32>, %arg1: tensor<2xi32>) -> tensor<1x4xi32> {
// CHECK: %0 = mhlo.constant dense<[1, -1]> : tensor<2xi32>
// CHECK-NEXT: %1 = mhlo.constant dense<[1, -1]> : tensor<2xi32>
// CHECK-NEXT: %2 = "mhlo.slice"(%1) <{limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<2xi32>) -> tensor<1xi32>
// CHECK-NEXT: %3 = mhlo.reshape %2 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %4 = mhlo.constant dense<[1, -1]> : tensor<2xi32>
// CHECK-NEXT: %5 = "mhlo.slice"(%4) <{limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<2xi32>) -> tensor<1xi32>
// CHECK-NEXT: %6 = mhlo.reshape %5 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %7 = mhlo.constant dense<3> : tensor<i32>
// CHECK-NEXT: %8 = "mhlo.set_dimension_size"(%arg0, %7) <{dimension = 0 : i64}> : (tensor<3x4xi32>, tensor<i32>) -> tensor<3x4xi32>
// CHECK-NEXT: %9 = mhlo.constant dense<4> : tensor<i32>
// CHECK-NEXT: %10 = "mhlo.set_dimension_size"(%8, %9) <{dimension = 1 : i64}> : (tensor<3x4xi32>, tensor<i32>) -> tensor<3x4xi32>
// CHECK-NEXT: %11 = mhlo.constant dense<0> : tensor<i32>
// CHECK-NEXT: %12 = "mhlo.pad"(%10, %11) <{edge_padding_high = dense<[3, 4]> : tensor<2xi64>, edge_padding_low = dense<0> : tensor<2xi64>, interior_padding = dense<0> : tensor<2xi64>}> : (tensor<3x4xi32>, tensor<i32>) -> tensor<6x8xi32>
// CHECK-NEXT: %13 = "mhlo.slice"(%arg1) <{limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<2xi32>) -> tensor<1xi32>
// CHECK-NEXT: %14 = mhlo.reshape %13 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %15 = "mhlo.slice"(%arg1) <{limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<2xi32>) -> tensor<1xi32>
// CHECK-NEXT: %16 = mhlo.reshape %15 : (tensor<1xi32>) -> tensor<i32>
// CHECK-NEXT: %17 = "mhlo.dynamic_slice"(%12, %14, %16) <{slice_sizes = dense<[3, 4]> : tensor<2xi64>}> : (tensor<6x8xi32>, tensor<i32>, tensor<i32>) -> tensor<3x4xi32>
// CHECK-NEXT: %18 = "mhlo.slice"(%17) <{limit_indices = dense<[1, 4]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}> : (tensor<3x4xi32>) -> tensor<1x4xi32>
// CHECK-NEXT: %19 = mhlo.constant dense<4> : tensor<i32>
// CHECK-NEXT: %20 = mhlo.subtract %19, %16 : tensor<i32>
// CHECK-NEXT: %21 = "mhlo.set_dimension_size"(%18, %20) <{dimension = 1 : i64}> : (tensor<1x4xi32>, tensor<i32>) -> tensor<1x?xi32, #mhlo.type_extensions<bounds = [?, 4]>>
// CHECK-NEXT: %cast = tensor.cast %21 : tensor<1x?xi32, #mhlo.type_extensions<bounds = [?, 4]>> to tensor<1x4xi32>
// CHECK-NEXT: return %cast : tensor<1x4xi32>
%sizes = "tf.Const"() {value = dense<[1, -1]> : tensor<2xi32>} : () -> (tensor<2xi32>)
%0 = "tf.Slice"(%arg0, %arg1, %sizes) : (tensor<3x4xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x4xi32>
func.return %0 : tensor<1x4xi32>
}
//===----------------------------------------------------------------------===//
// StridedSlice op legalizations.
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: simple_strided_slice
func.func @simple_strided_slice(%input: tensor<4x8xf32>) -> tensor<3x2xf32> {
// CHECK: %0 = mhlo.constant dense<[0, 1]> : tensor<2xi32>
// CHECK-NEXT: %1 = mhlo.constant dense<[3, 7]> : tensor<2xi32>
// CHECK-NEXT: %2 = mhlo.constant dense<[1, 3]> : tensor<2xi32>
// CHECK-NEXT: %3 = "mhlo.slice"(%arg0)
%begin = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi32>} : () -> (tensor<2xi32>)
%end = "tf.Const"() {value = dense<[3, 7]> : tensor<2xi32>} : () -> (tensor<2xi32>)
%strides = "tf.Const"() {value = dense<[1, 3]> : tensor<2xi32>} : () -> (tensor<2xi32>)
%output = "tf.StridedSlice"(%input, %begin, %end, %strides)
: (tensor<4x8xf32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<3x2xf32>
func.return %output : tensor<3x2xf32>
// CHECK: return %4 : tensor<3x2xf32>
}
// -----
// CHECK-LABEL: strided_slice_with_uknown_begin
func.func @strided_slice_with_uknown_begin(%input: tensor<4x8xf32>, %begin: tensor<1xi32>) -> tensor<4x8xf32> {
%end = "tf.Const"() {value = dense<[3]> : tensor<1xi32>} : () -> (tensor<1xi32>)
%strides = "tf.Const"() {value = dense<[1]> : tensor<1xi32>} : () -> (tensor<1xi32>)
// CHECK: %[[DYNAMIC_SLICE:.*]] = "mhlo.dynamic_slice"
%output = "tf.StridedSlice"(%input, %begin, %end, %strides)
: (tensor<4x8xf32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4x8xf32>
func.return %output : tensor<4x8xf32>
}
//===----------------------------------------------------------------------===//
// Fused op legalizations.
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: fused_conv2d
func.func @fused_conv2d(%input: tensor<1x300x300x40xi8>,
%filter: tensor<3x3x40x40xi8>,
%bias: tensor<40xf32>,
%act: tensor<0xi8>) -> tensor<1x300x300x40xi8> {
// CHECK: %[[v0:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
// CHECK-NEXT: %[[v1:.*]] = mhlo.constant dense<2.000000e+00> : tensor<f32>
// CHECK-NEXT: %[[v2:.*]] = mhlo.constant dense<2.000000e+00> : tensor<f32>
// CHECK-NEXT: %[[v3:.*]] = mhlo.constant dense<-1.280000e+02> : tensor<f32>
// CHECK-NEXT: %[[v4:.*]] = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<1x300x300x40xf32>
// CHECK-NEXT: %[[v5:.*]] = mhlo.convert %arg0 : (tensor<1x300x300x40xi8>) -> tensor<1x300x300x40xf32>
// CHECK-NEXT: %[[v6:.*]] = mhlo.convert %arg1 : (tensor<3x3x40x40xi8>) -> tensor<3x3x40x40xf32>
// CHECK: %[[v7:.*]] = mhlo.convolution(%[[v5]], %[[v6]])
// CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK-SAME{LITERAL}: window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1], reverse = [0, 0]}
// CHECK-SAME: batch_group_count = 1
// CHECK-SAME: feature_group_count = 1
// CHECK-NEXT: %[[v8:.*]] = mhlo.convert %7 : tensor<1x300x300x40xf32>
// CHECK-NEXT: %[[v9:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
// CHECK-NEXT: %[[v10:.*]] = "mhlo.broadcast_in_dim"(%9) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<1x300x300x40xf32>
// CHECK-NEXT: %11 = mhlo.multiply %8, %10 : tensor<1x300x300x40xf32>
// CHECK-NEXT: %12 = mhlo.convert %arg2 : tensor<40xf32>
// CHECK-NEXT: %13 = "mhlo.broadcast_in_dim"(%12) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<40xf32>) -> tensor<1x300x300x40xf32>
// CHECK-NEXT: %14 = mhlo.add %11, %13 : tensor<1x300x300x40xf32>
// CHECK-NEXT: %15 = mhlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-NEXT: %16 = "mhlo.broadcast_in_dim"(%15) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<1x300x300x40xf32>
// CHECK-NEXT: %17 = mhlo.maximum %14, %16 : tensor<1x300x300x40xf32>
// CHECK-NEXT: %18 = mhlo.constant dense<1.270000e+02> : tensor<f32>
// CHECK-NEXT: %19 = "mhlo.broadcast_in_dim"(%18) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<1x300x300x40xf32>
// CHECK-NEXT: %20 = mhlo.clamp %4, %17, %19 : tensor<1x300x300x40xf32>
// CHECK-NEXT: %21 = mhlo.round_nearest_even %20 : tensor<1x300x300x40xf32>
// CHECK-NEXT: %22 = mhlo.convert %21 : (tensor<1x300x300x40xf32>) -> tensor<1x300x300x40xi8>
// CHECK-NEXT: return %22 : tensor<1x300x300x40xi8>
%input_scale = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
%side_input_scale = "tf.Const"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
%conv2d = "tf._FusedConv2D"(%input, %filter, %bias, %act, %input_scale, %side_input_scale) {
data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 9.99999974E-5 : f32, explicit_paddings = [], filter_format = "HWIO", fused_ops = ["BiasAdd", "Relu"], leakyrelu_alpha = 2.000000e-01 : f32, num_args = 2 : i64, operandSegmentSizes = array<i32: 1, 1, 2, 2>, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true
} : (tensor<1x300x300x40xi8>, tensor<3x3x40x40xi8>, tensor<40xf32>, tensor<0xi8>, tensor<f32>, tensor<f32>) -> tensor<1x300x300x40xi8>
func.return %conv2d : tensor<1x300x300x40xi8>
}
//===----------------------------------------------------------------------===//
// tf.UniqueV2 legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: @doesnt_legalize_uniquev2
func.func @doesnt_legalize_uniquev2(%arg0: tensor<3xf32>) -> (tensor<?xf32>, tensor<3xi32>) {
//CHECK: %0 = mhlo.constant dense<0> : tensor<1xi32>
//CHECK-NEXT: %y, %idx = "tf.UniqueV2"(%arg0, %0) {device = ""} : (tensor<3xf32>, tensor<1xi32>) -> (tensor<?xf32>, tensor<3xi32>)
//CHECK-NEXT: return %y, %idx : tensor<?xf32>, tensor<3xi32>
%cst = "tf.Const"() {device = "", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%y, %idx = "tf.UniqueV2"(%arg0, %cst) {device = ""} : (tensor<3xf32>, tensor<1xi32>) -> (tensor<?xf32>, tensor<3xi32>)
func.return %y, %idx : tensor<?xf32>, tensor<3xi32>
}
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-quant.mlir
+576
View File
@@ -0,0 +1,576 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt -convert-tf-quant-types -quant-convert-tf-quant-ops-to-mhlo -canonicalize "-xla-legalize-tf=legalize-chlo=false" -split-input-file %s | env FILECHECK_OPTS="" FileCheck %s
//===----------------------------------------------------------------------===//
// tf.UniformQuantizedDotHybrid legalization
//===----------------------------------------------------------------------===//
// CHECK-LABEL: func @quantized_matmul_fn
func.func @quantized_matmul_fn(%input: tensor<?xf32>) -> tensor<?xf32> {
%weight = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2032207D207D2074656E736F725F636F6E74656E743A20225C3030315C3030325C3030335C30303422"> : tensor<2x2x!tf_type.qint8> } : () -> tensor<2x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[CONST:.*]] = mhlo.constant()
// CHECK-SAME{LITERAL}: value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi8>
// CHECK-SAME: tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: "mhlo.dot"(%arg0, %[[CONST]]) : (tensor<?xf32>, tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?xf32>
%0 = "tf.UniformQuantizedDotHybrid"(%input, %weight, %weight_scales, %weight_zps) {rhs_quantization_axis = -1 : i64, rhs_quantization_min_val = -128 : i64, rhs_quantization_max_val = 127 : i64} : (tensor<?xf32>, tensor<2x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
//===----------------------------------------------------------------------===//
// tf.UniformQuantizedConvolutionHybrid legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantized_convolution_hybrid
func.func @uniform_quantized_convolution_hybrid(%input: tensor<1x6x6x3xf32>) -> tensor<1x4x1x2xf32> {
%weight = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2032207D207D20696E745F76616C3A20313237"> : tensor<2x3x3x2x!tf_type.qint8>} : () -> tensor<2x3x3x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[CONST:.*]] = mhlo.constant()
// CHECK-SAME{LITERAL} value = dense<127> : tensor<2x3x3x2xi8>
// CHECK-SAME: tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: mhlo.convolution(%arg0, %[[CONST]])
// CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK-SAME{LITERAL}: window = {stride = [1, 2], pad = [[0, 0], [0, 0]], lhs_dilate = [1, 1], rhs_dilate = [2, 2]}
// CHECK-SAME{LITERAL}: batch_group_count = 1 : i64, feature_group_count = 1 : i64
// CHECK-SAME: (tensor<1x6x6x3xf32>, tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<1x4x1x2xf32>
%0 = "tf.UniformQuantizedConvolutionHybrid"(%input, %weight, %weight_scales, %weight_zps) {
window_strides = [1, 2],
padding = "VALID",
explicit_padding = [],
lhs_dilation = [1, 1],
rhs_dilation = [2, 2],
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -128 : i64,
rhs_quantization_max_val = 127 : i64
} : (tensor<1x6x6x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x4x1x2xf32>
func.return %0 : tensor<1x4x1x2xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_convolution_hybrid_same
func.func @uniform_quantized_convolution_hybrid_same(%input: tensor<1x2x2x3xf32>) -> tensor<1x2x1x2xf32> {
%weight = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2032207D207D20696E745F76616C3A20313237"> : tensor<2x3x3x2x!tf_type.qint8>} : () -> tensor<2x3x3x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[CONST:.*]] = mhlo.constant()
// CHECK-SAME{LITERAL} value = dense<127> : tensor<2x3x3x2xi8>
// CHECK-SAME: tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: mhlo.convolution(%arg0, %[[CONST]])
// CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK-SAME{LITERAL}: window = {stride = [1, 2], pad = [[1, 1], [1, 2]], lhs_dilate = [1, 1], rhs_dilate = [2, 2]}
// CHECK-SAME{LITERAL}: batch_group_count = 1 : i64, feature_group_count = 1 : i64
// CHECK-SAME: (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<1x2x1x2xf32>
%0 = "tf.UniformQuantizedConvolutionHybrid"(%input, %weight, %weight_scales, %weight_zps) {
window_strides = [1, 2],
padding = "SAME",
explicit_padding = [],
lhs_dilation = [1, 1],
rhs_dilation = [2, 2],
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -128 : i64,
rhs_quantization_max_val = 127 : i64
} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x2x1x2xf32>
func.return %0 : tensor<1x2x1x2xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_convolution_hybrid_explicit
func.func @uniform_quantized_convolution_hybrid_explicit(%input: tensor<1x2x2x3xf32>) -> tensor<1x3x3x2xf32> {
%weight = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2032207D207D20696E745F76616C3A20313237"> : tensor<2x3x3x2x!tf_type.qint8>} : () -> tensor<2x3x3x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[CONST:.*]] = mhlo.constant()
// CHECK-SAME{LITERAL} value = dense<127> : tensor<2x3x3x2xi8>
// CHECK-SAME: tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: mhlo.convolution(%arg0, %[[CONST]])
// CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK-SAME{LITERAL}: window = {stride = [1, 2], pad = [[1, 2], [3, 4]], lhs_dilate = [1, 1], rhs_dilate = [2, 2]}
// CHECK-SAME{LITERAL}: batch_group_count = 1 : i64, feature_group_count = 1 : i64
// CHECK-SAME: (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<1x3x3x2xf32>
%0 = "tf.UniformQuantizedConvolutionHybrid"(%input, %weight, %weight_scales, %weight_zps) {
window_strides = [1, 2],
padding = "EXPLICIT",
explicit_padding = [1, 2, 3, 4],
lhs_dilation = [1, 1],
rhs_dilation = [2, 2],
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -128 : i64,
rhs_quantization_max_val = 127 : i64
} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x3x3x2xf32>
func.return %0 : tensor<1x3x3x2xf32>
}
//===----------------------------------------------------------------------===//
// tf.UniformQuantize and tf.UniformDequantize legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantize_and_dequantize
func.func @uniform_quantize_and_dequantize(%arg0 : tensor<2xf32>) -> tensor<2xf32> {
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[QUANTIZE:.*]] = mhlo.uniform_quantize %arg0 : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANTIZE]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xi8>
// CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<2xi8>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_2]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xf32>
// CHECK: return %[[DEQUANTIZE]] : tensor<2xf32>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x!tf_type.qint8>
%1 = "tf.UniformDequantize"(%0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<2xf32>
func.return %1 : tensor<2xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_and_dequantize_per_axis
func.func @uniform_quantize_and_dequantize_per_axis(%arg0 : tensor<2x2xf32>) -> tensor<2x2xf32> {
%scales = "tf.Const"() { value = dense<[1.0, 2.0]> : tensor<2xf32> } : () -> tensor<2xf32>
%zps = "tf.Const"() { value = dense<[3, 4]> : tensor<2xi32> } : () -> tensor<2xi32>
// CHECK: %[[QUANTIZE:.*]] = mhlo.uniform_quantize %arg0 : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>
// CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANTIZE]] : (tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>) -> tensor<2x2xi8>
// CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<2x2xi8>) -> tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>
// CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_2]] : (tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>) -> tensor<2x2xf32>
// CHECK: return %[[DEQUANTIZE]] : tensor<2x2xf32>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = 0 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2xf32>, tensor<2xf32>, tensor<2xi32>) -> tensor<2x2x!tf_type.qint8>
%1 = "tf.UniformDequantize"(%0, %scales, %zps) {
quantization_axis = 0 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2x!tf_type.qint8>, tensor<2xf32>, tensor<2xi32>) -> tensor<2x2xf32>
func.return %1 : tensor<2x2xf32>
}
//===----------------------------------------------------------------------===//
// tf.UniformRequantize legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantize_requantize_and_dequantize
func.func @uniform_quantize_requantize_and_dequantize(%arg0 : tensor<4xf32>) -> tensor<4xf32> {
%scales_0 = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps_0 = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
%scales_1 = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%zps_1 = "tf.Const"() { value = dense<5> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[QUANTIZE:.*]] = mhlo.uniform_quantize %arg0 : (tensor<4xf32>) -> tensor<4x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANTIZE]] : (tensor<4x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<4xi8>
// CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<4xi8>) -> tensor<4x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: %[[REQUANTIZE:.*]] = mhlo.uniform_quantize %[[CONVERT_2]] : (tensor<4x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<4x!quant.uniform<i8:f32, 2.000000e+00:5>>
// CHECK: %[[CONVERT_3:.*]] = mhlo.bitcast_convert %[[REQUANTIZE]] : (tensor<4x!quant.uniform<i8:f32, 2.000000e+00:5>>) -> tensor<4xi8>
// CHECK: %[[CONVERT_4:.*]] = mhlo.bitcast_convert %[[CONVERT_3]] : (tensor<4xi8>) -> tensor<4x!quant.uniform<i8:f32, 2.000000e+00:5>>
// CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_4]] : (tensor<4x!quant.uniform<i8:f32, 2.000000e+00:5>>) -> tensor<4xf32>
// CHECK: return %[[DEQUANTIZE]] : tensor<4xf32>
%0 = "tf.UniformQuantize"(%arg0, %scales_0, %zps_0) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<4xf32>, tensor<f32>, tensor<i32>) -> tensor<4x!tf_type.qint8>
%1 = "tf.UniformRequantize"(%0, %scales_0, %zps_0, %scales_1, %zps_1) {
input_quantization_axis = -1 : i64, input_quantization_min_val = -128 : i64, input_quantization_max_val = 127 : i64,
output_quantization_axis = -1 : i64, output_quantization_min_val = -128 : i64, output_quantization_max_val = 127 : i64
} : (tensor<4x!tf_type.qint8>, tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>) -> tensor<4x!tf_type.qint8>
%2 = "tf.UniformDequantize"(%1, %scales_1, %zps_1) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<4x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<4xf32>
func.return %2 : tensor<4xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_requantize_and_dequantize_per_axis
func.func @uniform_quantize_requantize_and_dequantize_per_axis(%arg0 : tensor<2x2xf32>) -> tensor<2x2xf32> {
%scales_0 = "tf.Const"() { value = dense<[1.0, 2.0]> : tensor<2xf32> } : () -> tensor<2xf32>
%zps_0 = "tf.Const"() { value = dense<[3, 4]> : tensor<2xi32> } : () -> tensor<2xi32>
%scales_1 = "tf.Const"() { value = dense<[3.0, 4.0]> : tensor<2xf32> } : () -> tensor<2xf32>
%zps_1 = "tf.Const"() { value = dense<[5, 6]> : tensor<2xi32> } : () -> tensor<2xi32>
// CHECK: %[[QUANTIZE:.*]] = mhlo.uniform_quantize %arg0 : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>
// CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANTIZE]] : (tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>) -> tensor<2x2xi8>
// CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<2x2xi8>) -> tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>
// CHECK: %[[REQUANTIZE:.*]] = mhlo.uniform_quantize %[[CONVERT_2]] : (tensor<2x2x!quant.uniform<i8:f32:0, {1.000000e+00:3,2.000000e+00:4}>>) -> tensor<2x2x!quant.uniform<i8:f32:0, {3.000000e+00:5,4.000000e+00:6}>>
// CHECK: %[[CONVERT_3:.*]] = mhlo.bitcast_convert %[[REQUANTIZE]] : (tensor<2x2x!quant.uniform<i8:f32:0, {3.000000e+00:5,4.000000e+00:6}>>) -> tensor<2x2xi8>
// CHECK: %[[CONVERT_4:.*]] = mhlo.bitcast_convert %[[CONVERT_3]] : (tensor<2x2xi8>) -> tensor<2x2x!quant.uniform<i8:f32:0, {3.000000e+00:5,4.000000e+00:6}>>
// CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_4]] : (tensor<2x2x!quant.uniform<i8:f32:0, {3.000000e+00:5,4.000000e+00:6}>>) -> tensor<2x2xf32>
// CHECK: return %[[DEQUANTIZE]] : tensor<2x2xf32>
%0 = "tf.UniformQuantize"(%arg0, %scales_0, %zps_0) {
quantization_axis = 0 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2xf32>, tensor<2xf32>, tensor<2xi32>) -> tensor<2x2x!tf_type.qint8>
%1 = "tf.UniformRequantize"(%0, %scales_0, %zps_0, %scales_1, %zps_1) {
input_quantization_axis = 0 : i64, input_quantization_min_val = -128 : i64, input_quantization_max_val = 127 : i64,
output_quantization_axis = 0 : i64, output_quantization_min_val = -128 : i64, output_quantization_max_val = 127 : i64
} : (tensor<2x2x!tf_type.qint8>, tensor<2xf32>, tensor<2xi32>, tensor<2xf32>, tensor<2xi32>) -> tensor<2x2x!tf_type.qint8>
%2 = "tf.UniformDequantize"(%1, %scales_1, %zps_1) {
quantization_axis = 0 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2x!tf_type.qint8>, tensor<2xf32>, tensor<2xi32>) -> tensor<2x2xf32>
func.return %2 : tensor<2x2xf32>
}
//===----------------------------------------------------------------------===//
// tf.UniformQuantizedDot legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantized_dot
func.func @uniform_quantized_dot(%input: tensor<?xf32>) -> tensor<?x!tf_type.qint32> {
%input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%weight = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2032207D207D2074656E736F725F636F6E74656E743A20225C3030315C3030325C3030335C30303422"> : tensor<2x2x!tf_type.qint8> } : () -> tensor<2x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
%output_scales = "tf.Const"() { value = dense<3.0> : tensor<f32> } : () -> tensor<f32>
%output_zps = "tf.Const"() { value = dense<5> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[RHS:.*]] = mhlo.constant()
// CHECK-SAME{LITERAL}: <{value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi8>}> : () -> tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK-DAG: %[[LHS:.*]] = mhlo.uniform_quantize %arg0 : (tensor<?xf32>) -> tensor<?x!quant.uniform<i8:f32, 2.000000e+00:4>>
// CHECK-DAG: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[LHS]] : (tensor<?x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<?xi8>
// CHECK-DAG: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<?xi8>) -> tensor<?x!quant.uniform<i8:f32, 2.000000e+00:4>>
// CHECK: "mhlo.dot"(%[[CONVERT_2]], %[[RHS]]) : (tensor<?x!quant.uniform<i8:f32, 2.000000e+00:4>>, tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>)
// CHECK-SAME: -> tensor<?x!quant.uniform<i32:f32, 3.000000e+00:5>>
%0 = "tf.UniformQuantize"(%input, %input_scales, %input_zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint8>
%1 = "tf.UniformQuantizedDot"(
%0, %weight,
%input_scales, %input_zps,
%weight_scales, %weight_zps,
%output_scales, %output_zps) {
lhs_quantization_axis = -1 : i64,
lhs_quantization_min_val = -128 : i64,
lhs_quantization_max_val = 127 : i64,
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -128 : i64,
rhs_quantization_max_val = 127 : i64,
output_quantization_axis = -1 : i64,
output_quantization_min_val = -2147483648 : i64,
output_quantization_max_val = 2147483647 : i64} : (
tensor<?x!tf_type.qint8>, tensor<2x2x!tf_type.qint8>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint32>
func.return %1 : tensor<?x!tf_type.qint32>
}
//===----------------------------------------------------------------------===//
// tf.UniformQuantizedConvolution legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantized_convolution
func.func @uniform_quantized_convolution(%input: tensor<1x6x6x3xf32>) -> tensor<1x4x1x2x!tf_type.qint32> {
%input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%weight = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2033207D2064696D207B2073697A653A2032207D207D20696E745F76616C3A20313237"> : tensor<2x3x3x2x!tf_type.qint8>} : () -> tensor<2x3x3x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
%output_scales = "tf.Const"() { value = dense<3.0> : tensor<f32> } : () -> tensor<f32>
%output_zps = "tf.Const"() { value = dense<5> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[RHS:.*]] = mhlo.constant()
// CHECK-SAME{LITERAL}: <{value = dense<127> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK-DAG: %[[LHS:.*]] = mhlo.uniform_quantize %arg0 : (tensor<1x6x6x3xf32>) -> tensor<1x6x6x3x!quant.uniform<i8:f32, 2.000000e+00:4>>
// CHECK-DAG: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[LHS]] : (tensor<1x6x6x3x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x6x6x3xi8>
// CHECK-DAG: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<1x6x6x3xi8>) -> tensor<1x6x6x3x!quant.uniform<i8:f32, 2.000000e+00:4>>
// CHECK: mhlo.convolution(%[[CONVERT_2]], %[[RHS]])
// CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK-SAME{LITERAL}: window = {stride = [1, 2], pad = [[0, 0], [0, 0]], lhs_dilate = [1, 1], rhs_dilate = [2, 2]}
// CHECK-SAME{LITERAL}: batch_group_count = 1 : i64, feature_group_count = 1 : i64
// CHECK-SAME: (tensor<1x6x6x3x!quant.uniform<i8:f32, 2.000000e+00:4>>, tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<1x4x1x2x!quant.uniform<i32:f32, 3.000000e+00:5>>
%0 = "tf.UniformQuantize"(%input, %input_scales, %input_zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x6x6x3xf32>, tensor<f32>, tensor<i32>) -> tensor<1x6x6x3x!tf_type.qint8>
%1 = "tf.UniformQuantizedConvolution"(
%0, %weight,
%input_scales, %input_zps,
%weight_scales, %weight_zps,
%output_scales, %output_zps) {
window_strides = [1, 2],
padding = "VALID",
explicit_padding = [],
lhs_dilation = [1, 1],
rhs_dilation = [2, 2],
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
lhs_quantization_axis = -1 : i64,
lhs_quantization_min_val = -128 : i64,
lhs_quantization_max_val = 127 : i64,
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -128 : i64,
rhs_quantization_max_val = 127 : i64,
output_quantization_axis = -1 : i64,
output_quantization_min_val = -2147483648 : i64,
output_quantization_max_val = 2147483647 : i64} : (
tensor<1x6x6x3x!tf_type.qint8>, tensor<2x3x3x2x!tf_type.qint8>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>) -> tensor<1x4x1x2x!tf_type.qint32>
func.return %1 : tensor<1x4x1x2x!tf_type.qint32>
}
//===----------------------------------------------------------------------===//
// tf.UniformQuantizedAdd legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantized_add
func.func @uniform_quantized_add(%arg0: tensor<3x2x!tf_type.qint32>) -> tensor<3x2x!tf_type.qint32> {
%input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// tensor_proto that points to dense<127> of type !tf_type.qint32.
%bias = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<2x!tf_type.qint32> } : () -> tensor<2x!tf_type.qint32>
%bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[LHS:.*]] = mhlo.bitcast_convert %arg0 : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
// CHECK-DAG: %[[RHS:.*]] = mhlo.constant() <{value = dense<127> : tensor<2xi32>}> : () -> tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>
// CHECK: %[[RES:.*]] = chlo.broadcast_add %[[LHS]], %[[RHS]] {broadcast_dimensions = array<i64: 1>} :
// CHECK-SAME: (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>, tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>)
// CHECK-SAME: -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
// CHECK: %[[RES_INT:.*]] = mhlo.bitcast_convert %[[RES]] : (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>) -> tensor<3x2xi32>
// CHECK: return %[[RES_INT]] : tensor<3x2xi32>
%0 = "tf.UniformQuantizedAdd"(
%arg0, %bias,
%input_scales, %input_zps,
%bias_scales, %bias_zps,
%output_scales, %output_zps) {
lhs_quantization_axis = -1 : i64,
lhs_quantization_min_val = -2147483648 : i64,
lhs_quantization_max_val = 2147483647 : i64,
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -2147483648 : i64,
rhs_quantization_max_val = 2147483647 : i64,
output_quantization_axis = -1 : i64,
output_quantization_min_val = -2147483648 : i64,
output_quantization_max_val = 2147483647 : i64} : (
tensor<3x2x!tf_type.qint32>, tensor<2x!tf_type.qint32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>) -> tensor<3x2x!tf_type.qint32>
func.return %0 : tensor<3x2x!tf_type.qint32>
}
//===----------------------------------------------------------------------===//
// tf.UniformQuantizedClipByValue legalization
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @uniform_quantized_clip_by_value
func.func @uniform_quantized_clip_by_value(%input: tensor<3x2xf32>) -> tensor<3x2x!tf_type.qint32> {
%scales = "tf.Const"() { value = dense<2.0> : tensor<2xf32> } : () -> tensor<2xf32>
%zps = "tf.Const"() { value = dense<4> : tensor<2xi32> } : () -> tensor<2xi32>
// tensor_proto that points to dense<127> of type !tf_type.qint32.
// CHECK-DAG: %[[MIN_MAX:.*]] = mhlo.constant() <{value = dense<127> : tensor<3x2xi32>}> : () -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
%min = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<3x2x!tf_type.qint32> } : () -> tensor<3x2x!tf_type.qint32>
%max = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<3x2x!tf_type.qint32> } : () -> tensor<3x2x!tf_type.qint32>
// CHECK-DAG: %[[OPERAND:.*]] = mhlo.uniform_quantize %arg0 : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
%0 = "tf.UniformQuantize"(%input, %scales, %zps) {
quantization_axis = 1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
} : (tensor<3x2xf32>, tensor<2xf32>, tensor<2xi32>) -> tensor<3x2x!tf_type.qint32>
// CHECK-DAG: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[OPERAND]] : (tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>) -> tensor<3x2xi32>
// CHECK-DAG: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK: %[[MIN_CLIPPED:.*]] = chlo.broadcast_maximum %[[CONVERT_2]], %[[MIN_MAX]] {broadcast_dimensions = array<i64: 1>} :
// CHECK-SAME: (tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>, tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>)
// CHECK-SAME: -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK: %[[MAX_CLIPPED:.*]] = chlo.broadcast_minimum %[[MIN_CLIPPED]], %[[MIN_MAX]] {broadcast_dimensions = array<i64: 1>} :
// CHECK-SAME: (tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>, tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>)
// CHECK-SAME: -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK: %[[RESULT:.*]] = mhlo.bitcast_convert %[[MAX_CLIPPED]] : (tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>) -> tensor<3x2xi32>
// CHECK: return %[[RESULT]] : tensor<3x2xi32>
%1 = "tf.UniformQuantizedClipByValue"(%0, %min, %max, %scales, %zps) {
quantization_axis = 1 : i64,
quantization_min_val = -2147483648 : i64,
quantization_max_val = 2147483647 : i64
} : (tensor<3x2x!tf_type.qint32>, tensor<3x2x!tf_type.qint32>, tensor<3x2x!tf_type.qint32>, tensor<2xf32>, tensor<2xi32>) -> tensor<3x2x!tf_type.qint32>
func.return %1 : tensor<3x2x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_clip_by_value_min_not_const
func.func @uniform_quantized_clip_by_value_min_not_const(%input: tensor<3x2x!tf_type.qint32>, %min: tensor<3x2x!tf_type.qint32>) -> tensor<3x2x!tf_type.qint32> {
%scales = "tf.Const"() { value = dense<2.0> : tensor<2xf32> } : () -> tensor<2xf32>
%zps = "tf.Const"() { value = dense<4> : tensor<2xi32> } : () -> tensor<2xi32>
// tensor_proto that points to dense<127> of type !tf_type.qint32.
%max = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<3x2x!tf_type.qint32> } : () -> tensor<3x2x!tf_type.qint32>
// CHECK-DAG: %[[INPUT:.*]] = mhlo.bitcast_convert %arg0 : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK-DAG: %[[MIN:.*]] = mhlo.bitcast_convert %arg1 : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK: chlo.broadcast_maximum %[[INPUT]], %[[MIN]]
%res = "tf.UniformQuantizedClipByValue"(%input, %min, %max, %scales, %zps) {
quantization_axis = 1 : i64,
quantization_min_val = -2147483648 : i64,
quantization_max_val = 2147483647 : i64
} : (tensor<3x2x!tf_type.qint32>, tensor<3x2x!tf_type.qint32>, tensor<3x2x!tf_type.qint32>, tensor<2xf32>, tensor<2xi32>) -> tensor<3x2x!tf_type.qint32>
func.return %res : tensor<3x2x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_clip_by_value_max_not_const
func.func @uniform_quantized_clip_by_value_max_not_const(%input: tensor<3x2x!tf_type.qint32>, %max: tensor<3x2x!tf_type.qint32>) -> tensor<3x2x!tf_type.qint32> {
%scales = "tf.Const"() { value = dense<2.0> : tensor<2xf32> } : () -> tensor<2xf32>
%zps = "tf.Const"() { value = dense<4> : tensor<2xi32> } : () -> tensor<2xi32>
// tensor_proto that points to dense<127> of type !tf_type.qint32.
%min = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<3x2x!tf_type.qint32> } : () -> tensor<3x2x!tf_type.qint32>
// CHECK-DAG: %[[INPUT:.*]] = mhlo.bitcast_convert %arg0 : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK-DAG: %[[MAX:.*]] = mhlo.bitcast_convert %arg1 : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32:1, {2.000000e+00:4,2.000000e+00:4}>>
// CHECK-DAG: %[[INPUT_1:.*]] = chlo.broadcast_maximum
// CHECK: chlo.broadcast_minimum %[[INPUT_1]], %[[MAX]]
%res = "tf.UniformQuantizedClipByValue"(%input, %min, %max, %scales, %zps) {
quantization_axis = 1 : i64,
quantization_min_val = -2147483648 : i64,
quantization_max_val = 2147483647 : i64
} : (tensor<3x2x!tf_type.qint32>, tensor<3x2x!tf_type.qint32>, tensor<3x2x!tf_type.qint32>, tensor<2xf32>, tensor<2xi32>) -> tensor<3x2x!tf_type.qint32>
func.return %res : tensor<3x2x!tf_type.qint32>
}
//===----------------------------------------------------------------------===//
// quant.uniform type handling with control flow ops
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: func @while_region_with_quant
func.func @while_region_with_quant(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<?xf32> {
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
%one = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[QUANT0:.*]] = mhlo.uniform_quantize %[[ARG:.*]] : (tensor<?xf32>) -> tensor<?x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANT0]] : (tensor<?x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?xi8>
// CHECK: mhlo.while()
// CHECK: cond
// CHECK: %[[CHECK_RES:.*]] = chlo.broadcast_compare
// CHECK: mhlo.return %[[CHECK_RES]] : tensor<i1>
// CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<?xi8>) -> tensor<?x!quant.uniform<i8:f32, 1.000000e+00:3>>
// CHECK: %[[RET:.*]] = mhlo.uniform_dequantize %[[CONVERT_2]] : (tensor<?x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?xf32>
// CHECK: return %[[RET]] : tensor<?xf32>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint8>
%1 = "tf.WhileRegion"(%0) ({
^bb0(%carg0: tensor<?x!tf_type.qint8>):
%check = "tf.Equal"(%arg1, %one) : (tensor<i32>, tensor<i32>) -> tensor<i1>
"tf.Yield"(%check) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<?x!tf_type.qint8>):
%id = "tf.Identity"(%barg0) : (tensor<?x!tf_type.qint8>) -> tensor<?x!tf_type.qint8>
"tf.Yield"(%id) : (tensor<?x!tf_type.qint8>) -> ()
}) {is_stateless = false} : (tensor<?x!tf_type.qint8>) -> tensor<?x!tf_type.qint8>
%2 = "tf.UniformDequantize"(%1, %scales, %zps) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<?x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?xf32>
func.return %2 : tensor<?xf32>
}
// -----
// CHECK-LABEL: func @while_region_with_quant_two_args
func.func @while_region_with_quant_two_args(%arg0: tensor<2x2xf32>, %arg1: tensor<i32>) -> (tensor<2x?xf32>, tensor<?x2xf32>) {
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps2 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%zps4 = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%one = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[QUANT0:.*]] = mhlo.uniform_quantize %[[ARG:.*]] : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:2>>
// CHECK: %[[INT0:.*]] = mhlo.bitcast_convert %[[QUANT0]] : (tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:2>>) -> tensor<2x2xi8>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps2) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint8>
// CHECK: %[[QUANT1:.*]] = mhlo.uniform_quantize %[[ARG:.*]] : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:4>>
// CHECK: %[[INT1:.*]] = mhlo.bitcast_convert %[[QUANT1]] : (tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:4>>) -> tensor<2x2xi8>
%1 = "tf.UniformQuantize"(%arg0, %scales, %zps4) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint8>
// CHECK: %[[WHILE_RESULT:.*]]:2 = mhlo.while(%[[ARG0:.*]] = %[[INT0]], %[[ARG1:.*]] = %[[INT1]])
// CHECK-SAME: tensor<2x2xi8>, tensor<2x2xi8>
// CHECK: cond
// CHECK: do
// CHECK: mhlo.return %[[ARG0]], %[[ARG1]] : tensor<2x?xi8>, tensor<?x2xi8>
%2:2 = "tf.WhileRegion"(%0, %1) ({
^bb0(%carg0: tensor<2x2x!tf_type.qint8>, %carg1: tensor<2x2x!tf_type.qint8>):
%check = "tf.Equal"(%arg1, %one) : (tensor<i32>, tensor<i32>) -> tensor<i1>
"tf.Yield"(%check) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<2x?x!tf_type.qint8>, %barg1: tensor<?x2x!tf_type.qint8>):
%id = "tf.Identity"(%barg0) : (tensor<2x?x!tf_type.qint8>) -> tensor<2x?x!tf_type.qint8>
"tf.Yield"(%id, %barg1) : (tensor<2x?x!tf_type.qint8>, tensor<?x2x!tf_type.qint8>) -> ()
}) {is_stateless = false} : (tensor<2x2x!tf_type.qint8>, tensor<2x2x!tf_type.qint8>) -> (tensor<2x?x!tf_type.qint8>, tensor<?x2x!tf_type.qint8>)
// %[[RESULT0:.*]] = mhlo.uniform_dequantize %[[WHILE_RESULT]]#0
%3 = "tf.UniformDequantize"(%2#0, %scales, %zps2) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<2x?x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<2x?xf32>
// %[[RESULT1:.*]] = mhlo.uniform_dequantize %[[WHILE_RESULT]]#0
%4 = "tf.UniformDequantize"(%2#1, %scales, %zps4) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<?x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?x2xf32>
// return %[[RESULT0]], %[[RESULT1]]
func.return %3, %4 : tensor<2x?xf32>, tensor<?x2xf32>
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-with-tf2xla-hlo-importer.mlir
+724
View File
@@ -0,0 +1,724 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT prefer-tf2xla=true use-tf2xla-fallback=true" %s -verify-diagnostics -mlir-disable-threading | FileCheck %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
// CHECK-LABEL: binary_op
func.func @binary_op(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: mhlo.atan2
%0 = "tf.Atan2"(%arg0, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: multiple_return_values
func.func @multiple_return_values(%arg0: tensor<3xi64>) -> tensor<i64> {
%0:3 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<3xi64>) -> (tensor<i64>, tensor<i64>, tensor<i64>)
// CHECK: return %1 : tensor<i64>
func.return %0#0 : tensor<i64>
}
// CHECK-LABEL: constant_parameter
func.func @constant_parameter(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = "tf.Const"() {value = dense<1.42> : tensor<2xf32>} : () -> tensor<2xf32>
// CHECK: mhlo.atan2 %arg0, %2
%1 = "tf.Atan2"(%arg0, %0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: uses_translated_return_type
func.func @uses_translated_return_type(%arg0: tensor<3xf32>) -> tensor<?xf32> {
// CHECK: tensor.cast %{{[0-9]+}} : tensor<?xf32, #mhlo.type_extensions<bounds = [3]>> to tensor<?xf32>
%y, %idx = "tf.Unique"(%arg0) {device = ""} : (tensor<3xf32>) -> (tensor<?xf32>, tensor<3xi32>)
return %y : tensor<?xf32>
}
// CHECK-LABEL: @abs
func.func @abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NOT: tf.Abs
%0 = "tf.Abs"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: func @testBroadcastGradientArgs
func.func @testBroadcastGradientArgs(%s0: tensor<4xi32>, %s1: tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>) {
// CHECK: tf.BroadcastGradientArgs
%r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) : (tensor<4xi32>, tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>)
func.return %r0, %r1 : tensor<1xi32>, tensor<0xi32>
}
// CHECK-LABEL: @acos
func.func @acos(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NOT: tf.Acos
%0 = "tf.Acos"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: unknown_op
func.func @unknown_op(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: tf.CustomTestOp
%0 = "tf.CustomTestOp"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: add_v2
func.func @add_v2(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK: mhlo.add %arg0, %arg0 : tensor<2xi32>
%0 = "tf.AddV2"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
}
// CHECK-LABEL: not_allowlisted_op
func.func @not_allowlisted_op(%arg0: tensor<3xi32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<?x?x?xf32> {
// CHECK: tf.TensorListReserve
%0 = "tf.TensorListReserve"(%arg0, %arg1) : (tensor<3xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?x?x?xf32>>>
// CHECK: tf.TensorListGetItem
%1 = "tf.TensorListGetItem"(%0, %arg2, %arg0) : (tensor<!tf_type.variant<tensor<?x?x?xf32>>>, tensor<i32>, tensor<3xi32>) -> tensor<?x?x?xf32>
func.return %1 : tensor<?x?x?xf32>
}
// CHECK-LABEL: unranked_operand
func.func @unranked_operand(%arg0: tensor<*xf32>) -> tensor<*xf32> {
// CHECK: tf.Atan2
// expected-remark@+1 {{lowering requires bounded tensor operands}}
%0 = "tf.Atan2"(%arg0, %arg0) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
// CHECK-LABEL: dynamic_operand
func.func @dynamic_operand(%arg0: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: tf.Atan2
// expected-remark@+1 {{lowering requires bounded tensor operands}}
%0 = "tf.Atan2"(%arg0, %arg0) : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
// CHECK-LABEL: tuple_type
func.func @tuple_type(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
// Verifies that the pass can handle operands of non-tensor type like tuple
// from non TensorFlow ops.
%0 = "mhlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<f32>, tensor<i32>>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: unsupported_dtype
func.func @unsupported_dtype(%arg0: tensor<2x!tf_type.variant>) -> tensor<2x!tf_type.variant> {
// CHECK: tf.AddN
// expected-remark@+1 {{skipping legalization due to unsupported type 'tensor<2x!tf_type.variant>'}}
%0 = "tf.AddN"(%arg0, %arg0) : (tensor<2x!tf_type.variant>, tensor<2x!tf_type.variant>) -> tensor<2x!tf_type.variant>
func.return %0 : tensor<2x!tf_type.variant>
}
// CHECK-LABEL: multiple_dialect_ops
func.func @multiple_dialect_ops(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: mhlo.negate
%0 = "mhlo.negate"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
// CHECK: mhlo.atan2
%1 = "tf.Atan2"(%arg0, %0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
func.return %1 : tensor<2xf32>
}
// CHECK-LABEL: binary_op_broadcast
func.func @binary_op_broadcast(%arg0: tensor<4x1xf32>, %arg1: tensor<4x1x4xf32>) -> tensor<4x4x4xf32> {
// CHECK: %[[BROADCAST0:.*]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}> : (tensor<4x1xf32>) -> tensor<4x4x1xf32>
// CHECK: %[[RESHAPE0:.*]] = mhlo.reshape %[[BROADCAST0]] : (tensor<4x4x1xf32>) -> tensor<4x4xf32>
// CHECK: %[[UPDATED_ARG0:.*]] = "mhlo.broadcast_in_dim"(%[[RESHAPE0]]) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}> : (tensor<4x4xf32>) -> tensor<4x4x4xf32>
// CHECK: %[[RESHAPE1:.*]] = mhlo.reshape %arg1 : (tensor<4x1x4xf32>) -> tensor<4x4xf32>
// CHECK: %[[UPDATED_ARG1:.*]] = "mhlo.broadcast_in_dim"(%[[RESHAPE1]]) <{broadcast_dimensions = dense<[0, 2]> : tensor<2xi64>}> : (tensor<4x4xf32>) -> tensor<4x4x4xf32>
// CHECK: %[[RESULT:.*]] = mhlo.atan2 %[[UPDATED_ARG0]], %[[UPDATED_ARG1]] : tensor<4x4x4xf32>
// CHECK: return %[[RESULT]] : tensor<4x4x4xf32>
%0 = "tf.Atan2"(%arg0, %arg1) : (tensor<4x1xf32>, tensor<4x1x4xf32>) -> tensor<4x4x4xf32>
func.return %0: tensor<4x4x4xf32>
}
// CHECK-LABEL: func @ternary_op
func.func @ternary_op(%arg0: tensor<2xi1>, %arg1: tensor<2xi32>, %arg2: tensor<2xi32>) -> tensor<2xi32> {
// CHECK: mhlo.select %arg0, %arg1, %arg2
%0 = "tf.SelectV2"(%arg0, %arg1, %arg2) : (tensor<2xi1>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
}
// CHECK-LABEL: func @convert
func.func @convert(%arg0: tensor<2xi32>) -> tensor<2xf32> {
// CHECK: mhlo.convert %arg0 : (tensor<2xi32>) -> tensor<2xf32>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<2xi32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: func @constant
func.func @constant(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: %[[ONE:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
// CHECK: %[[RESULT:.*]] = mhlo.divide %[[ONE]], %arg0 : tensor<f32>
// CHECK: return %[[RESULT]]
%0 = "tf.Inv"(%arg0) : (tensor<f32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func @const_inputs
// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x2xf64>, %[[ARG1:.*]]: tensor<f64>,
func.func @const_inputs(%arg0: tensor<2x2xf64>, %arg1: tensor<f64>, %arg2: tensor<2xi32>, %arg3: tensor<2xi32>, %arg4: tensor<2xi32>) -> tensor<6x5xf64> {
// CHECK: "mhlo.pad"(%[[ARG0]], %[[ARG1]])
// CHECK-SAME-DAG: edge_padding_high = dense<[1, 2]> : tensor<2xi64>
// CHECK-SAME-DAG: edge_padding_low = dense<[2, 1]> : tensor<2xi64>
// CHECK-SAME-DAG: interior_padding = dense<[1, 0]> : tensor<2xi64>
%0 = mhlo.constant dense<[2, 1]> : tensor<2xi32>
%1 = mhlo.constant dense<[1, 2]> : tensor<2xi32>
%2 = mhlo.constant dense<[1, 0]> : tensor<2xi32>
%3 = "tf.XlaPad"(%arg0, %arg1, %0, %1, %2) : (tensor<2x2xf64>, tensor<f64>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<6x5xf64>
func.return %3 : tensor<6x5xf64>
}
func.func @non_const_inputs(%arg0: tensor<2x2xf64>, %arg1: tensor<f64>, %arg2: tensor<2xi32>, %arg3: tensor<2xi32>, %arg4: tensor<2xi32>) -> tensor<6x5xf64> {
// expected-remark@+1 {{compilation to HLO failed: INVALID_ARGUMENT: Input 2 to node `tf.XlaPad` with op XlaPad must be a compile-time constant.}}
%0 = "tf.XlaPad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<2x2xf64>, tensor<f64>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<6x5xf64>
func.return %0 : tensor<6x5xf64>
}
// CHECK-LABEL: dynamic_result_type
func.func @dynamic_result_type(%arg0: tensor<2xf32>) -> tensor<*xf32> {
// CHECK: %[[RESULT:.*]] = mhlo.atan2 %arg0, %arg0 : tensor<2xf32>
// CHECK: tensor.cast %[[RESULT]] : tensor<2xf32> to tensor<*xf32>
%0 = "tf.Atan2"(%arg0, %arg0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<*xf32>
// return %[[RESULT]]
func.return %0 : tensor<*xf32>
}
func.func @truncated_normal() -> tensor<2x2xf32> {
// CHECK-NOT: tf.TruncatedNormal
%0 = mhlo.constant dense<[2, 2]> : tensor<2xi32>
%1 = "tf.TruncatedNormal"(%0) {T = i32, device = "", dtype = f32, seed = 0 : i64, seed2 = 1950157571 : i64} : (tensor<2xi32>) -> tensor<2x2xf32>
func.return %1 : tensor<2x2xf32>
}
// CHECK-LABEL: dynamic_update_slice
// CHECK-SAME: (%[[ARG0:.*]]: tensor<3x4xi32>, %[[ARG1:.*]]: tensor<2x2xi32>, %[[ARG2:.*]]: tensor<2xi32>
func.func @dynamic_update_slice(%arg0: tensor<3x4xi32>, %arg1: tensor<2x2xi32>, %arg2: tensor<2xi32>) -> tensor<3x4xi32> {
// CHECK: %[[SLICE0:.*]] = "mhlo.slice"(%[[ARG2]])
// CHECK-DAG-SAME: start_indices = dense<0> : tensor<1xi64>
// CHECK-DAG-SAME: limit_indices = dense<1> : tensor<1xi64>
// CHECK-DAG-SAME: strides = dense<1> : tensor<1xi64>
// CHECK-SAME: (tensor<2xi32>) -> tensor<1xi32>
// CHECK: %[[DIM0:.*]] = mhlo.reshape %[[SLICE0]] : (tensor<1xi32>) -> tensor<i32>
// CHECK: %[[SLICE1:.*]] = "mhlo.slice"(%[[ARG2]])
// CHECK-DAG-SAME: start_indices = dense<1> : tensor<1xi64>
// CHECK-DAG-SAME: limit_indices = dense<2> : tensor<1xi64>
// CHECK-DAG-SAME: strides = dense<1> : tensor<1xi64>
// CHECK-SAME: (tensor<2xi32>) -> tensor<1xi32>
// CHECK: %[[DIM1:.*]] = mhlo.reshape %[[SLICE1]] : (tensor<1xi32>) -> tensor<i32>
// CHECK: mhlo.dynamic_update_slice %[[ARG0]], %[[ARG1]], %[[DIM0]], %[[DIM1]]
%0 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %arg2) : (tensor<3x4xi32>, tensor<2x2xi32>, tensor<2xi32>) -> tensor<3x4xi32>
func.return %0: tensor<3x4xi32>
}
// CHECK-LABEL: @sparse_to_dense
// CHECK-SAME: (%[[ARG0:.*]]: tensor<3x2xi32>, %[[ARG1:.*]]: tensor<3xf32>, %[[ARG2:.*]]: tensor<f32>)
func.func @sparse_to_dense(%arg0: tensor<3x2xi32>, %arg1: tensor<3xf32>, %arg2: tensor<f32>) -> tensor<3x3xf32> {
// CHECK: %[[DEFAULT:.*]] = "mhlo.broadcast_in_dim"(%[[ARG2]]) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<3x3xf32>
// CHECK: %[[RESULT:.*]] = "mhlo.scatter"(%[[DEFAULT]], %[[ARG0]], %[[ARG1]])
// CHECK-SAME: indices_are_sorted = false
// CHECK-SAME: scatter_dimension_numbers
// CHECK-SAME: inserted_window_dims = [0, 1]
// CHECK-SAME: scatter_dims_to_operand_dims = [0, 1]
// CHECK-SAME: index_vector_dim = 1
// CHECK-SAME: unique_indices = false
// CHECK-NEXT: ^bb0(%[[ARG3:.*]]: tensor<f32>, %[[ARG4:.*]]: tensor<f32>):
// CHECK: mhlo.return %[[ARG4]] : tensor<f32>
// CHECK: })
// CHECK-SAME: (tensor<3x3xf32>, tensor<3x2xi32>, tensor<3xf32>) -> tensor<3x3xf32>
// return %[[RESULT]] : tensor<3x3xf32>
%cst = mhlo.constant dense<3> : tensor<2xi32>
%0 = "tf.SparseToDense"(%arg0, %cst, %arg1, %arg2) {validate_indices = true}: (tensor<3x2xi32>, tensor<2xi32>, tensor<3xf32>, tensor<f32>) -> tensor<3x3xf32>
func.return %0 : tensor<3x3xf32>
}
// CHECK-LABEL: reverse_sequence
func.func @reverse_sequence(%arg0: tensor<4x2x3x1x1xi32>, %arg1: tensor<3xi32>) -> tensor<4x2x3x1x1xi32> {
// CHECK-NOT: tf.ReverseSequence
%0 = "tf.ReverseSequence"(%arg0, %arg1) {batch_dim = 2 : i64, seq_dim = 0 : i64}: (tensor<4x2x3x1x1xi32>, tensor<3xi32>) -> tensor<4x2x3x1x1xi32>
func.return %0 : tensor<4x2x3x1x1xi32>
}
// CHECK-LABEL: mirror_pad
func.func @mirror_pad(%arg0: tensor<2x3xcomplex<f64>>) -> tensor<4x7xcomplex<f64>> {
%0 = mhlo.constant dense<[[1, 1], [2, 2]]> : tensor<2x2xi32>
// CHECK-NOT: tf.MirrorPad
%1 = "tf.MirrorPad"(%arg0, %0) {mode = "SYMMETRIC"} : (tensor<2x3xcomplex<f64>>, tensor<2x2xi32>) -> tensor<4x7xcomplex<f64>>
func.return %1 : tensor<4x7xcomplex<f64>>
}
// CHECK-LABEL: bucketize
func.func @bucketize(%arg0: tensor<2x5xf32>) -> tensor<2x5xi32> {
// CHECK-NOT: tf.Bucketize
%0 = "tf.Bucketize"(%arg0) {boundaries = [0.000000e+00 : f32, 3.000000e+00 : f32, 8.000000e+00 : f32, 1.100000e+01 : f32]} : (tensor<2x5xf32>) -> tensor<2x5xi32>
func.return %0 : tensor<2x5xi32>
}
// CHECK-LABEL: arg_min
func.func @arg_min(%arg0: tensor<6xf64>) -> tensor<i32> {
// CHECK-NOT: ArgMin
%0 = mhlo.constant dense<0> : tensor<i32>
%1 = "tf.ArgMin"(%arg0, %0) : (tensor<6xf64>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: non_max_suppression_v4
func.func @non_max_suppression_v4(%arg0: tensor<3x4xf32>, %arg1: tensor<3xf32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<2xi32> {
%max_size = mhlo.constant dense<2> : tensor<i32>
// CHECK-NOT: tf.NonMaxSuppressionV4
%0:2 = "tf.NonMaxSuppressionV4"(%arg0, %arg1, %max_size, %arg2, %arg3) {pad_to_max_output_size = true}: (tensor<3x4xf32>, tensor<3xf32>, tensor<i32>, tensor<f32>, tensor<f32>) -> (tensor<2xi32>, tensor<i32>)
func.return %0#0 : tensor<2xi32>
}
// CHECK-LABEL: bessel_i0e
func.func @bessel_i0e(%arg0: tensor<3xf16>, %arg1: tensor<3xf32>, %arg2: tensor<3xf64>) -> (tensor<3xf16>, tensor<3xf32>, tensor<3xf64>) {
// CHECK-NOT: tf.BesselI0e
%0 = "tf.BesselI0e"(%arg0) : (tensor<3xf16>) -> (tensor<3xf16>)
%1 = "tf.BesselI0e"(%arg1) : (tensor<3xf32>) -> (tensor<3xf32>)
%2 = "tf.BesselI0e"(%arg2) : (tensor<3xf64>) -> (tensor<3xf64>)
func.return %0, %1, %2 : tensor<3xf16>, tensor<3xf32>, tensor<3xf64>
}
// CHECK-LABEL: bessel_i1e
func.func @bessel_i1e(%arg0: tensor<3xf16>, %arg1: tensor<3xf32>, %arg2: tensor<3xf64>) -> (tensor<3xf16>, tensor<3xf32>, tensor<3xf64>) {
// CHECK-NOT: tf.BesselI1e
%0 = "tf.BesselI1e"(%arg0) : (tensor<3xf16>) -> (tensor<3xf16>)
%1 = "tf.BesselI1e"(%arg1) : (tensor<3xf32>) -> (tensor<3xf32>)
%2 = "tf.BesselI1e"(%arg2) : (tensor<3xf64>) -> (tensor<3xf64>)
func.return %0, %1, %2 : tensor<3xf16>, tensor<3xf32>, tensor<3xf64>
}
// CHECK-LABEL: diag
func.func @diag(%arg0: tensor<2xf32>) -> tensor<2x2xf32> {
// CHECK-NOT: tf.Diag
%0 = "tf.Diag"(%arg0) : (tensor<2xf32>) -> tensor<2x2xf32>
func.return %0 : tensor<2x2xf32>
}
// CHECK-LABEL: random_uniform_int
func.func @random_uniform_int(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<1000xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
// CHECK-NOT: tf.RandomUniformInt
%1 = "tf.RandomUniformInt"(%0, %arg0, %arg1) {seed = 0 : i64, seed2 = 0 : i64} : (tensor<1xi32>, tensor<i32>, tensor<i32>) -> tensor<1000xi32>
func.return %1 : tensor<1000xi32>
}
// CHECK-LABEL: multinomial
func.func @multinomial(%arg0: tensor<2x4xf32>, %seed: tensor<i32>, %seed2: tensor<i32>) -> tensor<2x10xi32> {
// CHECK-NOT: tf.Multinomial
%samples = "tf.Const"() { value = dense<10> : tensor<i32> } : () -> tensor<i32>
%1 = "tf.Multinomial"(%arg0, %samples) {seed = 0, seed2 = 0}: (tensor<2x4xf32>, tensor<i32>) -> tensor<2x10xi32>
func.return %1 : tensor<2x10xi32>
}
// CHECK-LABEL: @set_dynamic_dimension_size
func.func @set_dynamic_dimension_size(%input: tensor<4xf32>, %size: tensor<i32>) -> tensor<?xf32> {
%dimension = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
// CHECK: mhlo.set_dimension_size
// CHECK-SAME: <{dimension = 0 : i64}> : (tensor<4xf32>, tensor<i32>) -> tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>
%0 = "tf.XlaSetDynamicDimensionSize"(%input, %dimension, %size) : (tensor<4xf32>, tensor<i32>, tensor<i32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
// CHECK-LABEL: unique
func.func @unique(%arg0: tensor<5xf32>) -> (tensor<?xf32>, tensor<?xi32>) {
// CHECK-NOT: tf.Unique
%0, %1 = "tf.Unique"(%arg0) : (tensor<5xf32>) -> (tensor<?xf32>, tensor<?xi32>)
func.return %0, %1 : tensor<?xf32> , tensor<?xi32>
}
// CHECK-LABEL: @erfinv
func.func @erfinv(%input: tensor<4xf32>) -> tensor<4xf32> {
// CHECK-NOT: tf.Erfinv
%0 = "tf.Erfinv"(%input) : (tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: @ndtri
func.func @ndtri(%input: tensor<4xf32>) -> tensor<4xf32> {
// CHECK-NOT: tf.Ndtri
%0 = "tf.Ndtri"(%input) : (tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: @fake_param
func.func @fake_param() -> tensor<4xf32> {
// CHECK-NOT: tf.FakeParam
%0 = "tf.FakeParam"() {shape = #tf_type.shape<4>} : () -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: @parameterized_truncated_normal
func.func @parameterized_truncated_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<10000000xf32> {
%0 = "tf.Const"() {value = dense<10000000> : tensor<1xi32>} : () -> tensor<1xi32>
// CHECK-NOT: tf.ParameterizedTruncatedNormal
%1 = "tf.ParameterizedTruncatedNormal"(%0, %arg0, %arg1, %arg2, %arg3) {seed = 0 : i64, seed2 = 0 : i64} : (tensor<1xi32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<10000000xf32>
func.return %1 : tensor<10000000xf32>
}
// Check XlaSpmdFullToShardShape's conversion from split sharding to manual
// sharding.
// The split sharding is:
// type: OTHER
// tile_assignment_dimensions: 2
// tile_assignment_dimensions: 1
// tile_assignment_devices: 0
// tile_assignment_devices: 1
// Serialized string:
// "\08\03\1A\02\02\01\22\02\00\01"
// The manual sharding is:
// type: MANUAL
// Serialized string:
// "\08\04"
// CHECK-LABEL: @xla_spmd_full_to_shard_shape
func.func @xla_spmd_full_to_shard_shape(%arg0: tensor<2x2xi64>) -> (tensor<1x2xi64>) {
// CHECK: %[[SHARDING:.*]] = mhlo.custom_call @Sharding(%arg0) {backend_config = "", mhlo.sharding = "{devices=[2,1]0,1}"} : (tensor<2x2xi64>) -> tensor<2x2xi64>
// CHECK: %[[MANUAL:.*]] = mhlo.custom_call @SPMDFullToShardShape(%[[SHARDING]]) {backend_config = "", mhlo.sharding = "{manual}"} : (tensor<2x2xi64>) -> tensor<1x2xi64>
// CHECK: return %[[MANUAL]]
%0 = "tf.XlaSpmdFullToShardShape"(%arg0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xi64>) -> tensor<1x2xi64>
func.return %0 : tensor<1x2xi64>
}
// Check XlaSpmdShardToFullShape's conversion from manual sharding to split
// sharding.
// The manual sharding is:
// type: MANUAL
// Serialized string:
// "\08\04"
// The split sharding is:
// type: OTHER
// tile_assignment_dimensions: 2
// tile_assignment_dimensions: 1
// tile_assignment_devices: 0
// tile_assignment_devices: 1
// Serialized string:
// "\08\03\1A\02\02\01\22\02\00\01"
// CHECK-LABEL: @xla_spmd_shard_to_full_shape
func.func @xla_spmd_shard_to_full_shape(%arg0: tensor<1x2xi64>) -> (tensor<2x2xi64>) {
// CHECK: %[[SHARDING:.*]] = mhlo.custom_call @Sharding(%arg0) {backend_config = "", mhlo.sharding = "{manual}"} : (tensor<1x2xi64>) -> tensor<1x2xi64>
// CHECK: %[[FULL:.*]] = mhlo.custom_call @SPMDShardToFullShape(%[[SHARDING]]) {backend_config = "", mhlo.sharding = "{devices=[2,1]0,1}"} : (tensor<1x2xi64>) -> tensor<2x2xi64>
// CHECK: return %[[FULL]]
%0 = "tf.XlaSpmdShardToFullShape"(%arg0) {dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<1x2xi64>) -> tensor<2x2xi64>
func.return %0 : tensor<2x2xi64>
}
// CHECK-LABEL: @xla_svd
func.func @xla_svd(%arg0: tensor<1x1xf32>) -> (tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) {
// CHECK-NOT: XlaSvd
%s, %u, %v = "tf.XlaSvd"(%arg0) {max_iter = 1, epsilon = 1.0E-09 : f32, precision_config = ""} : (tensor<1x1xf32>) -> (tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>)
func.return %s, %u, %v : tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>
}
func.func @identity(%arg0: f32) -> f32 {
func.return %arg0 : f32
}
// This test verifies that legalization for ops with symbol reference attribute
// is not attempted even if they are in allow-list. XLA op kernels for these
// ops compile the function to HLO on-demand which won't work in our case as it
// may contain unsupported ops in the fallback nor we provide XlaCompiler to
// the kernel. Using a allowed op Atan2 to protect against future addition of a
// new op with a symbol ref.
// CHECK-LABEL: @atan2_with_symbol_ref
func.func @atan2_with_symbol_ref(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: tf.Atan2
// expected-remark@+1 {{ops with symbol references are not supported}}
%0 = "tf.Atan2"(%arg0, %arg0) {_body = @identity} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
func.func private @branch0(tensor<2xf32>) -> tensor<2xf32>
func.func private @branch1(tensor<2xf32>) -> tensor<2xf32>
func.func @case_with_symbol_ref(%arg0: tensor<i32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: tf.Case
// expected-remark@+1 {{ops with symbol references are not supported}}
%0 = "tf.Case"(%arg0, %arg1) {branches = [@branch0, @branch1], is_stateless = false} : (tensor<i32>, tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK-LABEL: const
func.func @const() -> tensor<2xf32> {
// CHECK: mhlo.const
%cst = "tf.Const"() {value = dense<2.0> : tensor<2xf32>} : () -> tensor<2xf32>
func.return %cst : tensor<2xf32>
}
// CHECK-LABEL: @bounds_propagation
func.func @bounds_propagation(%input: tensor<4xf32>, %size: tensor<i32>) -> tensor<?xf32> {
%dimension = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[BOUNDED:.*]] = "mhlo.set_dimension_size"
// CHECK-SAME: <{dimension = 0 : i64}> : (tensor<4xf32>, tensor<i32>) -> tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>
%0 = "tf.XlaSetDynamicDimensionSize"(%input, %dimension, %size) : (tensor<4xf32>, tensor<i32>, tensor<i32>) -> tensor<?xf32>
%axis = "tf.Const"() { value = dense<0> : tensor<1xi32> } : () -> tensor<1xi32>
// CHECK: %[[REVERSED:.*]] = "mhlo.reverse"(%[[BOUNDED]])
// CHECK-SAME: {dimensions = dense<0> : tensor<1xi64>}
// CHECK-SAME: (tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>) -> tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>
%1 = "tf.ReverseV2"(%0, %axis) : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
// CHECK: %[[RESULT:.*]] = tensor.cast %[[REVERSED]] : tensor<?xf32, #mhlo.type_extensions<bounds = [4]>> to tensor<?xf32>
// CHECK: return %[[RESULT]] : tensor<?xf32>
func.return %1 : tensor<?xf32>
}
// CHECK-LABEL: @bounds_propagation_skip_symbol_ref_ops
func.func @bounds_propagation_skip_symbol_ref_ops(%input: tensor<4xf32>, %size: tensor<i32>) -> tensor<?xf32> {
%dimension = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[BOUNDED:.*]] = "mhlo.set_dimension_size"
// CHECK-SAME: <{dimension = 0 : i64}> : (tensor<4xf32>, tensor<i32>) -> tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>
%0 = "tf.XlaSetDynamicDimensionSize"(%input, %dimension, %size) : (tensor<4xf32>, tensor<i32>, tensor<i32>) -> tensor<?xf32>
// CHECK: %[[ORIGINAL:.*]] = tensor.cast %[[BOUNDED]] : tensor<?xf32, #mhlo.type_extensions<bounds = [4]>> to tensor<?xf32>
%axis = "tf.Const"() { value = dense<0> : tensor<1xi32> } : () -> tensor<1xi32>
// CHECK: tf.ReverseV2
// CHECK-SAME: (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
// expected-remark@+1 {{lowering requires bounded tensor operands}}
%1 = "tf.ReverseV2"(%0, %axis) {_body = @identity} : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
func.return %1 : tensor<?xf32>
}
// CHECK-LABEL: func @set_bound
func.func @set_bound(%arg0: tensor<i32>) -> tensor<i32> {
%bound = "tf.Const"() {value = dense<16> : tensor<i32>} : () -> tensor<i32>
// CHECK: %[[RESULT:.*]] = mhlo.custom_call @SetBound(%arg0) {backend_config = "", mhlo.literal = dense<16> : tensor<i32>}
%bounded = "tf.XlaSetBound"(%arg0, %bound) : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: return %[[RESULT]]
func.return %bounded : tensor<i32>
}
// CHECK-LABEL: func @greater
func.func @greater(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> {
// CHECK-NEXT: mhlo.compare GT, %arg0, %arg1
%0 = "tf.Greater"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
func.return %0: tensor<2xi1>
}
// CHECK-LABEL: batchmatmulv2
func.func @batchmatmulv2(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK: mhlo.reduce
// CHECK: mhlo.dot_general
// CHECK: mhlo.transpose
%0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
func.return %0 : tensor<3x4x4xf32>
}
// CHECK-LABEL: approx_topk
func.func @approx_topk(%arg0: tensor<!tf_type.resource<tensor<10x500xbf16>>> {tf._user_specified_name = "db"}) -> (tensor<10x10xbf16>) {
%0 = "tf.ReadVariableOp"(%arg0) {device = ""} : (tensor<!tf_type.resource<tensor<10x500xbf16>>>) -> tensor<10x500xbf16>
// CHECK: mhlo.compare GT
%values, %indices = "tf.ApproxTopK"(%0) {aggregate_to_topk = true, device = "", is_max_k = true, k = 10 : i64, recall_target = 0.949999988 : f32, reduction_dimension = -1 : i64, reduction_input_size_override = -1 : i64} : (tensor<10x500xbf16>) -> (tensor<10x10xbf16>, tensor<10x10xi32>)
return %values : tensor<10x10xbf16>
}
// CHECK-LABEL: fusedBatchNormV3_noTraining
func.func @fusedBatchNormV3_noTraining(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>) -> (tensor<8x8x8x8xf32>) {
// CHECK: "mhlo.batch_norm_inference"({{.*}}, %arg1, %arg2, %arg3, %arg4) <{epsilon = 1.000000e-03 : f32, feature_index = 3 : i64}> : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> tensor<8x8x8x8xf32>
%0:6 = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg2, %arg3, %arg4) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)
func.return %0#0 : tensor<8x8x8x8xf32>
}
// CHECK-LABEL: fusedBatchNormV3_training
func.func @fusedBatchNormV3_training(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>) -> (tensor<8x8x8x8xf32>) {
// CHECK: %[[OUT:.*]], %[[MEAN:.*]], %[[VAR:.*]] = "mhlo.batch_norm_training"({{.*}}, %arg1, %arg2) <{epsilon = 1.000000e-03 : f32, feature_index = 3 : i64}> : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>)
%0:6 = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg2, %arg3, %arg4) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, exponential_avg_factor = 1.0 : f32, is_training = true} : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)
func.return %0#0 : tensor<8x8x8x8xf32>
}
// CHECK-LABEL: fusedBatchNormGradV3_noTraining
func.func @fusedBatchNormGradV3_noTraining(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8x8x8x8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>, %arg5: tensor<8xf32>) -> (tensor<8x8x8x8xf32>) {
// CHECK-NEXT: %[[grad:.*]] = mhlo.convert %arg0 : tensor<8x8x8x8xf32>
// CHECK-NEXT: %[[act:.*]] = mhlo.convert %arg1 : tensor<8x8x8x8xf32>
// CHECK: %[[scr1:.*]] = mhlo.rsqrt
// CHECK: %[[bcast_arg3:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg3, {{.*}}) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<8xf32>, tensor<4xindex>) -> tensor<8x8x8x8xf32>
// CHECK: %[[sub:.*]] = mhlo.subtract %[[act]], %[[bcast_arg3]] : tensor<8x8x8x8xf32>
// CHECK: %[[mul:.*]] = mhlo.multiply %[[grad]], %[[sub]] : tensor<8x8x8x8xf32>
// CHECK: mhlo.constant dense<[0, 1, 2]> : tensor<3xi64>
// CHECK-NEXT: %[[cmul:.*]] = mhlo.convert %[[mul]] : tensor<8x8x8x8xf32>
// CHECK-NEXT: %[[init:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-NEXT: %[[convert_init:.*]] = mhlo.convert %[[init]] : tensor<f32>
// CHECK: %[[red1:.*]] = mhlo.reduce(%[[cmul]] init: %[[convert_init]]) applies mhlo.add across dimensions = [0, 1, 2] : (tensor<8x8x8x8xf32>, tensor<f32>) -> tensor<8xf32>
// CHECK: %[[scr2:.*]] = mhlo.convert %[[red1]] : tensor<8xf32>
// CHECK: %[[mul2:.*]] = mhlo.multiply %arg2, %[[scr1]] : tensor<8xf32>
// CHECK: %[[bcast_mul2:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[mul2]], {{.*}}) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<8xf32>, tensor<4xindex>) -> tensor<8x8x8x8xf32>
// CHECK: %[[mul3:.*]] = mhlo.multiply %[[grad]], %[[bcast_mul2]] : tensor<8x8x8x8xf32>
// CHECK: %[[scale_backprop:.*]] = mhlo.multiply %[[scr1]], %[[scr2]] : tensor<8xf32>
// CHECK: mhlo.constant dense<[0, 1, 2]> : tensor<3xi64>
// CHECK: %[[cgrad:.*]] = mhlo.convert %[[grad]] : tensor<8x8x8x8xf32>
// CHECK: %[[init2:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-NEXT: %[[convert_init2:.*]] = mhlo.convert %[[init2]] : tensor<f32>
// CHECK: %[[red2:.*]] = mhlo.reduce(%[[cgrad]] init: %[[convert_init2]]) applies mhlo.add across dimensions = [0, 1, 2] : (tensor<8x8x8x8xf32>, tensor<f32>) -> tensor<8xf32>
// CHECK: %[[offset_backprop:.*]] = mhlo.convert %[[red2]] : tensor<8xf32>
// CHECK: %[[x_backprop:.*]] = mhlo.convert %[[mul3]] : tensor<8x8x8x8xf32>
// CHECK: return %[[x_backprop]] : tensor<8x8x8x8xf32>
%0:5 = "tf.FusedBatchNormGradV3"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)
func.return %0#0 : tensor<8x8x8x8xf32>
}
// CHECK-LABEL: fusedBatchNormGradV3_Training
func.func @fusedBatchNormGradV3_Training(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8x8x8x8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>, %arg5: tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<0xf32>, tensor<*xf32>) {
// CHECK-NEXT: %[[grad:.*]] = mhlo.convert %arg0 : tensor<8x8x8x8xf32>
// CHECK-NEXT: %[[act:.*]] = mhlo.convert %arg1 : tensor<8x8x8x8xf32>
// CHECK-NEXT: %[[grad_operand:.*]], %[[grad_scale:.*]], %[[grad_offset:.*]] = "mhlo.batch_norm_grad"(%[[act]], %arg2, %arg3, %arg4, %[[grad]]) <{epsilon = 1.000000e-03 : f32, feature_index = 3 : i64}> : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8x8x8x8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>)
// CHECK-NEXT: %[[x_backprop:.*]] = mhlo.convert %[[grad_operand]] : tensor<8x8x8x8xf32>
// CHECK: return %[[x_backprop]]
// CHECK-SAME: tensor<8x8x8x8xf32>
%0:5 = "tf.FusedBatchNormGradV3"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = true} : (tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<0xf32>, tensor<*xf32>)
func.return %0#0, %0#3, %0#4 : tensor<8x8x8x8xf32>, tensor<0xf32>, tensor<*xf32>
}
// CHECK-LABEL: @max_pool_grad_valid
// CHECK-SAME: %[[INPUT:.*]]: tensor<10x24x24x64xf32>, %arg1: tensor<10x12x12x64xf32>, %[[GRAD:.*]]: tensor<10x12x12x64xf32>
func.func @max_pool_grad_valid(%orig_input: tensor<10x24x24x64xf32>, %orig_output: tensor<10x12x12x64xf32>, %grad: tensor<10x12x12x64xf32>) -> tensor<10x24x24x64xf32> {
// CHECK: %[[ZERO:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK: %[[RESULT:.*]] = "mhlo.select_and_scatter"(%[[INPUT]], %[[GRAD]], %[[ZERO]]) <{
// CHECK-SAME: padding = dense<0> : tensor<4x2xi64>, window_dimensions = dense<[1, 2, 2, 1]> : tensor<4xi64>, window_strides = dense<[1, 2, 2, 1]> : tensor<4xi64>
// CHECK-SAME }> ({
// CHECK: ^bb0(%[[VALUE_A:.*]]: tensor<f32>, %[[VALUE_B:.*]]: tensor<f32>):
// CHECK: %[[SELECT_RESULT:.*]] = mhlo.compare GE, %[[VALUE_A]], %[[VALUE_B]] : (tensor<f32>, tensor<f32>) -> tensor<i1>
// CHECK: mhlo.return %[[SELECT_RESULT]] : tensor<i1>
// CHECK: }, {
// CHECK: ^bb0(%[[VALUE_A:.*]]: tensor<f32>, %[[VALUE_B:.*]]: tensor<f32>):
// CHECK: %[[SELECT_RESULT:.*]] = mhlo.add %[[VALUE_A]], %[[VALUE_B]] : tensor<f32>
// CHECK: mhlo.return %[[SELECT_RESULT]] : tensor<f32>
// CHECK: }) : (tensor<10x24x24x64xf32>, tensor<10x12x12x64xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
// CHECK: return %[[RESULT]] : tensor<10x24x24x64xf32>
%result = "tf.MaxPoolGrad"(%orig_input, %orig_output, %grad) {
data_format = "NHWC",
explicit_paddings = [],
ksize = [1, 2, 2, 1],
padding = "VALID",
strides = [1, 2, 2, 1]
} : (tensor<10x24x24x64xf32>, tensor<10x12x12x64xf32>, tensor<10x12x12x64xf32>) -> tensor<10x24x24x64xf32>
func.return %result : tensor<10x24x24x64xf32>
}
// CHECK-LABEL: @max_pool_grad_same
func.func @max_pool_grad_same(%orig_input: tensor<2x13x25x7xf32>, %orig_output: tensor<2x4x7x7xf32>, %grad: tensor<2x4x7x7xf32>) -> tensor<2x13x25x7xf32> {
// CHECK: padding = dense<{{\[\[}}0, 0], [0, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>
%result = "tf.MaxPoolGrad"(%orig_input, %orig_output, %grad) {
data_format = "NHWC",
explicit_paddings = [],
ksize = [1, 2, 3, 1],
padding = "SAME",
strides = [1, 4, 4, 1]
} : (tensor<2x13x25x7xf32>, tensor<2x4x7x7xf32>, tensor<2x4x7x7xf32>) -> tensor<2x13x25x7xf32>
func.return %result : tensor<2x13x25x7xf32>
}
//===--------------------------------------------------------------------===//
// tf.XlaReduceScatter legalization
//===--------------------------------------------------------------------===//
// CHECK-LABEL: func @xla_reduce_scatter
func.func @xla_reduce_scatter(%arg0: tensor<128x128xf32>) -> tensor<64x128xf32> {
%cst = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%cst_0 = "tf.Const"() {value = dense<[[0, 4], [1, 5], [2, 6], [3, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
// CHECK: "mhlo.reduce_scatter"(%arg0)
// CHECK{LITERAL}: replica_groups = dense<[[0, 4], [1, 5], [2, 6], [3, 7]]>
// CHECK-SAME: scatter_dimension = 0
//
%1 = "tf.XlaReduceScatter"(%arg0, %cst_0, %cst) {reduce_op = "Add"} : (tensor<128x128xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<64x128xf32>
func.return %1 : tensor<64x128xf32>
}
// CHECK-LABEL: func @tf_mod
func.func @tf_mod(%arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
%cst = "tf.Const"() {value = dense<7.000000e+00> : tensor<f32>} : () -> tensor<f32>
// CHECK: mhlo.broadcast_in_dim
// CHECK: mhlo.remainder
%6 = "tf.Mod"(%arg1, %cst) {_global_shape = [#tf_type.shape<4x8>], device = ""} : (tensor<2x2xf32>, tensor<f32>) -> tensor<2x2xf32>
return %6 : tensor<2x2xf32>
}
// CHECK-LABEL: func @concat_v2
func.func @concat_v2(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
// CHECK: "mhlo.concatenate"({{.*}}) <{dimension = 0 : i64}> : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32>
%axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
%1 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor<i64>) -> tensor<6x3xf32>
func.return %1 : tensor<6x3xf32>
}
// CHECK-LABEL: func @matrix_diag_part_v3
// CHECK-SAME: %[[ARG:.*]]: tensor<7x140x128xi32>
func.func @matrix_diag_part_v3(%arg0: tensor<7x140x128xi32>) -> tensor<7x22x128xi32> {
%0 = mhlo.constant dense<42> : tensor<i32> // padding value
%1 = mhlo.constant dense<[-10, 11]> : tensor<2xi32> // k
// CHECK: mhlo.iota
// CHECK: mhlo.reshape
// CHECK: mhlo.concatenate
// CHECK: mhlo.gather
// CHECK: mhlo.broadcast
// CHECK: mhlo.select
%2 = "tf.MatrixDiagPartV3"(%arg0, %1, %0) {
T = i32, align = "RIGHT_LEFT"
} : (tensor<7x140x128xi32>, tensor<2xi32>, tensor<i32>) -> tensor<7x22x128xi32>
func.return %2: tensor<7x22x128xi32>
}
// CHECK-LABEL: func @xla_call_module
func.func @xla_call_module(%arg0: tensor<f32>) -> tensor<*xf32> {
// Equivalent to the following:
//
// module @jit_sin {
// func.func public @main(%arg0: tensor<f32>) -> tensor<f32> {
// %0 = mhlo.sine %arg0 : tensor<f32>
// return %0 : tensor<f32>
// }
// }
// CHECK: call @main.1
%0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device = "", dim_args_spec = [], function_list = [], disabled_checks = [], has_token_input_output = false, module = "ML\EFR\03MLIRxxx-trunk\00\01\17\05\01\05\01\03\05\03\07\07\t\0B\03K5\07\01\1B\07\0B\13\0B3\0B\0B\0B\0B\0F\0B\13\0B\03\1B\0F\1B\0B\0B\0B\0B\0B\0F\13\0B\0B\0B\0B\03\07\0F\17\07\02\A7\1F\05\0D\03\03\03\07\05\0F\03\0B\0B\1B\0D'\0F)\031\113\05\11\05\13\05\15\05\17\1D\15\17\05\19\17\19\EF\01\05\1B\03\03\1D\0D\05\1F!#%\1D\1D\1D\1F\1D!\1D##\03\03\03+\0D\03-/\1D%\1D'\1D)\1D+)\01\05\11\03\01\03\01\t\04A\05\01\11\01\05\07\03\01\05\03\11\01\t\05\03\05\0B\03\01\01\05\06\13\03\01\03\01\07\04\01\03\03\06\03\01\05\01\00\9A\04-\0F\0B\03!\1B\1D\05\1B\83/\1F\15\1D\15\11\13\15\11\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00sine_v1\00return_v1\00sym_name\00jit_sin\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00jit(sin)/jit(main)/sin\00third_party/py/jax/experimental/jax2tf/tests/back_compat_test.py\00jax.arg_info\00x\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00", platforms = ["CPU"], version = 6 : i64} : (tensor<f32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
// Verifies that the following functions are added from xla_call_module. Note this must be at the end of the file.
// CHECK: func.func private @main.1(%arg0: tensor<f32> {mhlo.sharding = "{replicated}"}) -> tensor<f32> {
// CHECK: %0 = mhlo.sine %arg0 : tensor<f32>
// CHECK: return %0 : tensor<f32>
// CHECK: }
}
tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir
+2783
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD
+60
View File
@@ -0,0 +1,60 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
)
cc_library(
name = "graph_to_tf_executor_registration",
srcs = [
"graph_to_tf_executor_registration.cc",
],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow/translate/tools:file_tf_mlir_translate",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/compiler/mlir/tools:translate_cl_options",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
"@tsl//tsl/platform:protobuf",
"@xla//xla/client:client_library",
"@xla//xla/client:compile_only_client",
"@xla//xla/service/cpu:cpu_compiler",
"@xla//xla/service/cpu:cpu_transfer_manager",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/stream_executor/host:host_platform",
"@xla//xla/stream_executor/host:host_platform_id",
"@xla//xla/tsl/platform:status",
],
alwayslink = 1,
)
tf_cc_test(
name = "graph_to_tf_executor_registration_test",
size = "small",
srcs = ["graph_to_tf_executor_registration_test.cc"],
deps = [
":graph_to_tf_executor_registration",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
],
)
tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc
+193
View File
@@ -0,0 +1,193 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <utility>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/tools/file_tf_mlir_translate.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include "tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/stream_executor/host/host_platform_id.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
using tsl::Status;
static constexpr char kMlirToGraphCompilationCheckName[] =
"mlir-to-graph-compilation-check";
// Use CPU arbitrarily in order to check that a graph compiles at all
static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT";
static Status CompileGraph(tensorflow::Graph* graph,
xla::CompileOnlyClient* client) {
if (!graph || !client) {
return Status(absl::StatusCode::kInvalidArgument,
"Invalid graph or client");
}
tensorflow::FunctionDefLibrary flib;
auto flib_def = std::make_unique<tensorflow::FunctionLibraryDefinition>(
tensorflow::OpRegistry::Global(), flib);
tensorflow::XlaCompiler::Options options;
options.device_type = tensorflow::DeviceType(kArbitraryDeviceName);
options.client = client;
options.flib_def = flib_def.get();
tensorflow::XlaCompiler compiler(options);
std::unique_ptr<tensorflow::Graph> graph_copy =
std::make_unique<tensorflow::Graph>(tensorflow::OpRegistry::Global());
tensorflow::CopyGraph(*graph, graph_copy.get());
tensorflow::XlaCompiler::CompileOptions compile_options;
tensorflow::XlaCompiler::CompilationResult result;
return compiler.CompileGraph(compile_options,
kMlirToGraphCompilationCheckName,
std::move(graph_copy), {}, &result);
}
static mlir::OwningOpRef<mlir::ModuleOp> GraphdefToMlirTranslateFunction(
llvm::StringRef input, mlir::MLIRContext* context) {
tensorflow::GraphdefToMlirOptions options{
debug_info_file, xla_compile_device_type,
prune_unused_nodes, convert_legacy_fed_inputs,
graph_as_function, upgrade_legacy,
enable_shape_inference, unconditionally_use_set_output_shapes,
enable_soft_placement, set_original_tf_func_name};
auto module_or = tensorflow::GraphdefToMlirTranslateFunction(
input, input_arrays, input_dtypes, input_shapes, output_arrays,
control_output_arrays, options, context);
if (!module_or.status().ok()) return nullptr;
return std::move(module_or).value();
}
static mlir::TranslateToMLIRRegistration GraphdefToMlirTranslate(
"graphdef-to-mlir", "graphdef-to-mlir", GraphdefToMlirTranslateFunction);
static mlir::LogicalResult MlirToGraphTranslateFunction(
mlir::ModuleOp module, llvm::raw_ostream& output) {
if (!module) return mlir::failure();
tensorflow::GraphExportConfig confs;
confs.export_entry_func_to_flib = export_entry_func_to_flib;
confs.export_original_tf_func_name = export_original_tf_func_name;
std::unique_ptr<tensorflow::FunctionLibraryDefinition> flib_def;
auto graph =
std::make_unique<tensorflow::Graph>(tensorflow::OpRegistry::Global());
absl::flat_hash_set<tensorflow::Node*> control_ret_nodes;
auto status = tensorflow::tf2xla::v2::ConvertTfExecutorToGraph(
module, confs, &graph, flib_def.get(), &control_ret_nodes);
if (!status.ok()) {
LOG(ERROR) << "Export to Graph failed: " << status;
return mlir::failure();
}
// Use Host platform, which should always exist, to make sure graphs compile.
auto platform = stream_executor::PlatformManager::PlatformWithId(
stream_executor::host::kHostPlatformId);
if (!platform.ok()) {
return mlir::failure();
}
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform.value());
tensorflow::XlaOpRegistry::RegisterCompilationKernels();
// Verify that the resulting graph can compile.
if (client.ok() && !CompileGraph(graph.get(), client.value()).ok()) {
return mlir::failure();
}
auto graphdef = std::make_unique<tensorflow::GraphDef>();
// Print the graph to the output after going through GraphDef conversion.
// The DumpGraphToFile would do this anyway so just skip straight to it.
graph->ToGraphDef(graphdef.get());
output << tsl::LegacyUnredactedDebugString(*graphdef);
return mlir::success();
}
static mlir::TranslateFromMLIRRegistration mlir_to_graph_translate(
/*name=*/"mlir-to-graph", /*description=*/"convert mlir to graph",
MlirToGraphTranslateFunction, [](mlir::DialectRegistry& registry) {
mlir::RegisterAllTensorFlowDialects(registry);
});
static llvm::LogicalResult MlirToGraphdefTranslateFunction(
mlir::ModuleOp module, llvm::raw_ostream& output) {
if (!module) return mlir::failure();
tensorflow::GraphExportConfig confs;
confs.export_entry_func_to_flib = export_entry_func_to_flib;
confs.export_original_tf_func_name = export_original_tf_func_name;
tensorflow::FunctionLibraryDefinition flib_def(
tensorflow::OpRegistry::Global(), tensorflow::FunctionDefLibrary());
auto graph =
std::make_unique<tensorflow::Graph>(tensorflow::OpRegistry::Global());
absl::flat_hash_set<tensorflow::Node*> control_ret_nodes;
auto status = tensorflow::tf2xla::v2::ConvertTfExecutorToGraph(
module, confs, &graph, &flib_def, &control_ret_nodes);
if (!status.ok()) {
LOG(ERROR) << "Export to Graph failed: " << status;
return mlir::failure();
}
tensorflow::GraphDef graphdef;
graph->ToGraphDef(&graphdef);
output << tsl::LegacyUnredactedDebugString(graphdef);
return mlir::success();
}
static mlir::TranslateFromMLIRRegistration mlir_to_graphdef_translate(
"mlir-to-graphdef", "mlir-to-graphdef", MlirToGraphdefTranslateFunction,
[](mlir::DialectRegistry& registry) {
mlir::RegisterAllTensorFlowDialects(registry);
});
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration_test.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/match.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
using mlir::LogicalResult;
using mlir::StringRef;
using mlir::Translation;
using mlir::TranslationParser;
class MlirTranslationTest : public ::testing::Test {
private:
static constexpr char kMlirToGraphFlag[] = "-mlir-to-graph";
public:
MlirTranslationTest() : translation_(RegisterTranslation()) {
// Create fake command line args so that the parser gets chosen.
std::vector<const char*> argv = {""};
argv.push_back(kMlirToGraphFlag);
llvm::cl::ParseCommandLineOptions(argv.size(), &argv[0],
"TF MLIR translation test\n");
}
LogicalResult Translate(StringRef source, std::string& sink) {
auto source_manager = std::make_shared<llvm::SourceMgr>();
auto source_buffer = llvm::MemoryBuffer::getMemBuffer(source);
source_manager->AddNewSourceBuffer(std::move(source_buffer), llvm::SMLoc());
mlir::MLIRContext context;
llvm::raw_string_ostream os(sink);
return (**translation_)(source_manager, os, &context);
}
private:
llvm::cl::opt<const Translation*, false, TranslationParser>*
RegisterTranslation() {
// Can only register once per process.
static const auto requested_translation =
new llvm::cl::opt<const Translation*, false, TranslationParser>(
llvm::cl::desc("Translation to perform"));
return requested_translation;
}
llvm::cl::opt<const Translation*, false, TranslationParser>* translation_;
};
TEST_F(MlirTranslationTest, TranslatesMlirToGraph) {
static constexpr char kMlirSource[] = R"(
func.func @main() -> (tensor<1x2xf16>, tensor<2xf16>) {
%graph:2 = tf_executor.graph {
%0:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<1.0> : tensor<1x2xf16>} : () -> tensor<1x2xf16> loc("const1")
%1:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<[1.0, 2.0]> : tensor<2xf16>} : () -> tensor<2xf16> loc("const2")
tf_executor.fetch %0#0, %1#0 : tensor<1x2xf16>, tensor<2xf16>
}
func.return %graph#0, %graph#1 : tensor<1x2xf16>, tensor<2xf16>
})";
std::string result;
auto status = Translate(kMlirSource, result);
ASSERT_TRUE(status.succeeded());
EXPECT_TRUE(absl::StrContains(result, "node {"));
}
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_cpu.mlir
+26
View File
@@ -0,0 +1,26 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "--tfxla-device-specific-transforms=device-type=XLA_CPU_JIT" -verify-diagnostics -split-input-file %s | FileCheck -dump-input=fail %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1399 : i32}} {
// CHECK-LABEL: stateless_op
func.func @stateless_op() -> tensor<i32> {
// CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
%0 = "tf.StatelessRandomGetAlg"() {device = ""} : () -> tensor<i32>
return %0 : tensor<i32>
}
}
tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_gpu.mlir
+26
View File
@@ -0,0 +1,26 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "--tfxla-device-specific-transforms=device-type=XLA_GPU_JIT" -verify-diagnostics -split-input-file %s | FileCheck -dump-input=fail %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1399 : i32}} {
// CHECK-LABEL: stateless_op
func.func @stateless_op() -> tensor<i32> {
// CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
%0 = "tf.StatelessRandomGetAlg"() {device = ""} : () -> tensor<i32>
return %0 : tensor<i32>
}
}
tensorflow/compiler/mlir/tf2xla/tests/verify-tfxla-legalization-no-chlo.mlir
+25
View File
@@ -0,0 +1,25 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "-tfxla-verify-legalization=legalize-chlo=false" -verify-diagnostics -split-input-file %s | FileCheck %s --dump-input=fail
// Tests the VerifyTFXLALegalization Pass, that just ensures we don't have
// any illegal ops at the end of the pipeline. This runs with
// legalize-chlo=false since errors can't be mixed with the legalize-chlo=True
// version.
// CHECK-LABEL: allows_chlo
func.func @allows_chlo(%arg0: tensor<1x32x10x32xi32>, %arg1: tensor<32xi32>) -> tensor<1x32x10x32xi32> {
%0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = array<i64: 3>} : (tensor<1x32x10x32xi32>, tensor<32xi32>) -> tensor<1x32x10x32xi32>
func.return %0 : tensor<1x32x10x32xi32>
}
tensorflow/compiler/mlir/tf2xla/tests/verify-tfxla-legalization.mlir
+73
View File
@@ -0,0 +1,73 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tf-opt "-tfxla-verify-legalization=legalize-chlo=true" -verify-diagnostics -split-input-file %s | FileCheck -dump-input=fail %s
// Tests the VerifyTFXLALegalization Pass, that just ensures we don't have
// any illegal ops at the end of the pipeline.
// CHECK-LABEL: allowsMHLO
func.func @allowsMHLO() -> (tensor<8x64x32x4xcomplex<f32>> {mhlo.sharding = ""}) {
%0 = mhlo.constant dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>
%1 = mhlo.constant dense<(1.000000e+00,1.000000e+00)> : tensor<8x64x128xcomplex<f32>>
%2 = "mhlo.einsum"(%1, %0) <{einsum_config = "abc,cde->abde"}> : (tensor<8x64x128xcomplex<f32>>, tensor<128x32x4xcomplex<f32>>) -> tensor<8x64x32x4xcomplex<f32>>
return %2 : tensor<8x64x32x4xcomplex<f32>>
}
// -----
func.func @invalid_non_mhlo() -> ( tensor<128x32x4xcomplex<f32>> {mhlo.sharding = ""}) {
// expected-error @+1 {{Could not legalize op: tf.Const}}
%cst = "tf.Const"() {value = dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>} : () -> tensor<128x32x4xcomplex<f32>>
return %cst : tensor<128x32x4xcomplex<f32>>
}
// -----
func.func @invalid_mixed_mhlo() -> (tensor<8x64x128xcomplex<f32>> {mhlo.sharding = ""}) {
%0 = mhlo.constant dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>
// expected-error @+1 {{Could not legalize op: tf.Const}}
%cst_0 = "tf.Const"() {value = dense<(1.000000e+00,1.000000e+00)> : tensor<8x64x128xcomplex<f32>>} : () -> tensor<8x64x128xcomplex<f32>>
return %cst_0 : tensor<8x64x128xcomplex<f32>>
}
// -----
func.func @fails_chlo(%arg0: tensor<1x32x10x32xi32>, %arg1: tensor<32xi32>) -> tensor<1x32x10x32xi32> {
// expected-error @+1 {{Could not legalize op: chlo.broadcast_add}}
%0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = array<i64: 3>} : (tensor<1x32x10x32xi32>, tensor<32xi32>) -> tensor<1x32x10x32xi32>
func.return %0 : tensor<1x32x10x32xi32>
}
// -----
func.func @multiple_failures() -> (tensor<8x64x128xcomplex<f32>> {mhlo.sharding = ""}) {
%0 = mhlo.constant dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>
// expected-error @+1 {{Could not legalize op: tf.Const}}
%cst_0 = "tf.Const"() {value = dense<(1.000000e+00,1.000000e+00)> : tensor<8x64x128xcomplex<f32>>} : () -> tensor<8x64x128xcomplex<f32>>
// expected-error @+1 {{Could not legalize op: tf.XlaEinsum}}
%1 = "tf.XlaEinsum"(%cst_0, %0) {equation = "abc,cde->abde"} : (tensor<8x64x128xcomplex<f32>>, tensor<128x32x4xcomplex<f32>>) -> tensor<8x64x32x4xcomplex<f32>>
return %cst_0 : tensor<8x64x128xcomplex<f32>>
}
// -----
func.func @nonstatic_shape_mhlo() -> tensor<?xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "i,j", outputs = "identity_RetVal"}} {
%0 = mhlo.constant dense<1.000000e+00> : tensor<f64>
%1 = mhlo.convert %0 : (tensor<f64>) -> tensor<i64>
%2 = mhlo.reshape %1 : (tensor<i64>) -> tensor<1xi64>
// expected-error @+1 {{Node `mhlo.dynamic_iota` must have compile-time constant}}
%3 = "mhlo.dynamic_iota"(%2) <{iota_dimension = 0 : i64}> : (tensor<1xi64>) -> tensor<?xi32>
%4 = mhlo.multiply %3, %3 : tensor<?xi32>
return %4 : tensor<?xi32>
}