wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 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/dtensor/mlir/tests/BUILD
+60
View File
@@ -0,0 +1,60 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
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",
# TODO(b/177569789): Fix below tests to use V2.
exclude = [
"layout_propagation.mlir",
"spmd_fill.mlir",
"spmd_metadata.mlir",
"spmd_reduction.mlir",
"spmd_tile.mlir",
],
tags_override = {
"move_compilation_to_host.mlir": ["no_oss"], # FIXME(b/264922760): The test fails on OSS.
"spmd_dtensor_ops.mlir": ["no_oss"], # FIXME(b/264922760): The test fails on OSS.
},
test_file_exts = ["mlir"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
":dtensor-opt",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
],
)
tf_cc_binary(
name = "dtensor-opt",
srcs = ["dtensor_mlir_opt_main.cc"],
deps = [
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:ops",
"//tensorflow/dtensor/cc:dtensor_ops",
"//tensorflow/dtensor/mlir:create_dtensor_mlir_passes",
"//tensorflow/dtensor/mlir:dtensor_mlir_passes",
"//tensorflow/dtensor/mlir:tf_dtensor_dialect",
"//tensorflow/dtensor/mlir/dtensor_dialect:Dialect",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:AllExtensions",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:MlirOptLib",
"@llvm-project//mlir:RegisterAllDialects",
"@llvm-project//mlir:RegisterAllExtensions",
"@llvm-project//mlir:RegisterAllPasses",
"@stablehlo//:register",
],
)
tensorflow/dtensor/mlir/tests/annotate_global_shape.mlir
+84
View File
@@ -0,0 +1,84 @@
// 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: dtensor-opt %s -dtensor-annotate-global-shape -split-input-file | FileCheck %s
// CHECK-LABEL: func @check_op_global_shape_annotated
func.func @check_op_global_shape_annotated() {
// CHECK: "tf.A"() {_global_shape = [#tf_type.shape<>]}
// CHECK-NEXT: "tf.B"() {_global_shape = [#tf_type.shape<64x64>, #tf_type.shape<2x8x8>]}
%1 = "tf.A"() : () -> tensor<i32>
%2, %3 = "tf.B"() : () -> (tensor<64x64xi64>, tensor<2x8x8xf32>)
func.return
}
// -----
// CHECK-LABEL: func @check_op_with_unranked_type_annotated
func.func @check_op_with_unranked_type_annotated() {
// CHECK: "tf.B"() {_global_shape = [#tf_type.shape<*>]}
%1 = "tf.B"() : () -> tensor<*xi32>
func.return
}
// -----
// CHECK-LABEL: func @check_op_with_non_static_shape
func.func @check_op_with_non_static_shape() {
// CHECK: "tf.B"() {_global_shape = [#tf_type.shape<4>, #tf_type.shape<?>]}
%1, %2 = "tf.B"() : () -> (tensor<4xi32>, tensor<?xi32>)
func.return
}
// -----
// CHECK-LABEL: func @check_function_arg_retval_annotated
// CHECK-SAME: %arg0: tensor<4x2xi32> {tf._global_shape = #tf_type.shape<4x2>}
// CHECK-SAME: (tensor<4x2xi32> {tf._global_shape = #tf_type.shape<4x2>})
func.func @check_function_arg_retval_annotated(%arg0: tensor<4x2xi32>) -> tensor<4x2xi32> {
%0 = "tf.Identity"(%arg0) : (tensor<4x2xi32>) -> tensor<4x2xi32>
func.return %0 : tensor<4x2xi32>
}
// -----
// CHECK-LABEL: func @check_function_callsites_annotated_properly
// CHECK-SAME: %arg0: tensor<4x2xi32> {tf._global_shape = #tf_type.shape<4x2>}
// CHECK-SAME: (tensor<4x2xi32> {tf._global_shape = #tf_type.shape<4x2>})
func.func @check_function_callsites_annotated_properly(%arg0: tensor<4x2xi32>) -> tensor<4x2xi32> {
// CHECK: "tf.StatefulPartitionedCall"
// CHECK-SAME: _global_shape = [#tf_type.shape<4x2>]
// CHECK-SAME: (tensor<4x2xi32>) -> tensor<4x2xi32>
%0 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @call_func} : (tensor<4x2xi32>) -> (tensor<4x2xi32>)
func.return %0 : tensor<4x2xi32>
}
// CHECK-LABEL: func @call_func
// CHECK-SAME: %arg0: tensor<4x2xi32> {tf._global_shape = #tf_type.shape<4x2>}
func.func @call_func(%arg0: tensor<4x2xi32>) -> tensor<4x2xi32> {
func.return %arg0 : tensor<4x2xi32>
}
// -----
// CHECK-LABEL: func @check_resource_type_shape
// CHECK-SAME: %arg1: tensor<!tf_type.resource<tensor<4x2xf32>>> {tf._global_shape = #tf_type.shape<4x2>
func.func @check_resource_type_shape(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<4x2xf32>>> {tf._layout = "sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = ""}} {
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
"tf.AssignVariableOp"(%arg1, %0) {device = ""} : (tensor<!tf_type.resource<tensor<4x2xf32>>>, tensor<f32>) -> ()
tf_device.return
}) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> ()
func.return
}
tensorflow/dtensor/mlir/tests/cluster_function_conversion.mlir
+127
View File
@@ -0,0 +1,127 @@
// 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: dtensor-opt %s -split-input-file -dtensor-cluster-function-conversion -verify-diagnostics | FileCheck %s
func.func @check_cluster_func_without_layout_disallowed() {
%1 = "tf.A"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> tensor<i32>
%2 = "tf.B"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> tensor<i32>
// expected-error @+1 {{requires _mesh attribute}}
%3 = "tf_device.cluster_func"(%1, %2) {func = @main_func1} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return
}
func.func @main_func1(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
func.return %arg0 : tensor<i32>
}
// -----
// CHECK-LABEL: func @check_layouts_retvals_attached_in_layout_op
func.func @check_layouts_retvals_attached_in_layout_op() -> tensor<i32> {
// CHECK-NOT: "tf_device.cluster_func"()
// CHECK: %[[SPC_OUT:.*]] = "tf.StatefulPartitionedCall"()
// CHECK-SAME: config = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: _layout = ["sharding_specs: mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%0 = "tf_device.cluster_func"() {func = @single_in_out, _mesh="|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @single_in_out() -> (tensor<i32>) {
%0 = "tf.Const"() {_layout = ["sharding_specs:scalar, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func @check_layouts_retval_attached_with_multi_in_op
func.func @check_layouts_retval_attached_with_multi_in_op(%arg0: tensor<i64>, %arg1: tensor<1xf32> {tf._layout = "sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3" }, %arg2: tensor<1xf32> {tf._layout = "mesh:CPU,x=2,y=2 layout:scalar" }) -> tensor<1xf32> {
// CHECK-NOT: "tf_device.cluster_func"()
// CHECK-NEXT: %[[SPC_OUT:.*]] = "tf.StatefulPartitionedCall"(%arg1, %arg2)
// CHECK-SAME: config = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%0 = "tf_device.cluster_func"(%arg1, %arg2) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", func = @multi_in} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
func.func @multi_in(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
%0 = "tf.Add"(%arg0, %arg1) {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// CHECK-LABEL: func @check_input_resource_layouts_attached_in_call_op
func.func @check_input_resource_layouts_attached_in_call_op() -> tensor<i32> {
// CHECK-NOT: "tf_device.cluster_func"()
// CHECK: %[[SPC_OUT:.*]] = "tf.StatefulPartitionedCall"()
// CHECK-SAME: config = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: _inferred_resource_indices = dense<1> : vector<1xi32>
// CHECK-SAME: _inferred_resource_layouts
// CHECK-SAME: "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%0 = "tf_device.cluster_func"() {func = @single_in_out, _mesh="|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", _inferred_resource_indices = dense<1> : vector<1xi32>,
_inferred_resource_layouts = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @single_in_out() -> (tensor<i32>) {
%0 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: func @check_nested_stateful_partitioned_call
func.func @check_nested_stateful_partitioned_call() -> (tensor<i32>, tensor<i32>) {
// CHECK-NOT: "tf_device.cluster_func"()
// CHECK: %[[SPC_OUT:.*]] = "tf.StatefulPartitionedCall"()
// CHECK-SAME: config = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%0:2 = "tf_device.cluster_func"() {func = @nested_stateful_partitioned_call, _mesh="|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>)
func.return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
func.func @nested_stateful_partitioned_call() -> (tensor<i32>, tensor<i32>) {
%0:2 = "tf.StatefulPartitionedCall()"() {config = "", config_proto = "", executor_type = "", f = @nested_cluster_func, _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> (tensor<i32>, tensor<i32>)
func.return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
func.func @nested_cluster_func() -> (tensor<i32>, tensor<i32>) {
%0:2 = "tf_device.cluster_func"() {func = @nested_func, _mesh="|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>)
func.return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
func.func @nested_func() -> (tensor<i32>, tensor<i32>) {
%0 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
func.return %0, %1 : tensor<i32>, tensor<i32>
}
// -----
// CHECK-LABEL: func @check_var_handle_op_skip_compilation
func.func @check_var_handle_op_skip_compilation() -> tensor<!tf_type.resource<tensor<i32>>> {
// CHECK-NOT: "tf_device.cluster_func"()
// CHECK: %[[SPC_OUT:.*]] = "tf.StatefulPartitionedCall"()
// CHECK-SAME: config = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
// CHECK-SAME: _layout = ["sharding_specs: mesh:TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"]
// CHECK-SAME: _skip_xla_compilation = true
%0 = "tf_device.cluster_func"() {func = @var_handle_op, _mesh="TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> tensor<!tf_type.resource<tensor<i32>>>
func.return %0 : tensor<!tf_type.resource<tensor<i32>>>
}
func.func @var_handle_op() -> (tensor<!tf_type.resource<tensor<i32>>>) {
%0 = "tf.VarHandleOp"() {_layout = ["sharding_specs:scalar, mesh:TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"], container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
func.return %0 : tensor<!tf_type.resource<tensor<i32>>>
}
tensorflow/dtensor/mlir/tests/constant_folding.mlir
+33
View File
@@ -0,0 +1,33 @@
// 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: dtensor-opt %s -split-input-file -dtensor-constant-folding | FileCheck %s
// Check that constants with same size/value are de-duplicated.
// CHECK-LABEL: func @check_constants_folded
func.func @check_constants_folded() {
// CHECK: %[[CONST_OUT_0:.*]] = "tf.Const"()
// CHECK-SAME: value = dense<[8, 128, 128]> : tensor<3xi32>
// CHECK-NEXT: %[[CONST_OUT_1:.*]] = "tf.Const"()
// CHECK-SAME: value = dense<[8, 128]> : tensor<2xi32>
// CHECK-NEXT: "tf.A"(%[[CONST_OUT_0]], %[[CONST_OUT_0]], %[[CONST_OUT_0]], %[[CONST_OUT_1]], %[[CONST_OUT_1]])
// CHECK-NEXT: return
%1 = "tf.Const"() {value = dense<[8, 128, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = "tf.Const"() {value = dense<[8, 128, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%3 = "tf.Const"() {value = dense<[8, 128, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%4 = "tf.Const"() {value = dense<[8, 128]> : tensor<2xi32>} : () -> tensor<2xi32>
%5 = "tf.Const"() {value = dense<[8, 128]> : tensor<2xi32>} : () -> tensor<2xi32>
"tf.A"(%1, %2, %3, %4, %5) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>, tensor<2xi32>, tensor<2xi32>) -> ()
func.return
}
tensorflow/dtensor/mlir/tests/cpu_layout.pbtxt
+18
View File
@@ -0,0 +1,18 @@
sharding_specs {
sharding_spec: "batch"
}
sharding_specs {
sharding_spec: "x"
}
mesh_config {
mesh_dimensions {
name: "batch"
size: 2
}
mesh_dimensions {
name: "x"
size: 2
}
devices: "/job:ps/device:CPU:0"
devices: "/job:ps/device:CPU:1"
}
tensorflow/dtensor/mlir/tests/decompose_controlflow.mlir
+284
View File
@@ -0,0 +1,284 @@
// 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: dtensor-opt %s -split-input-file -dtensor-decompose-controlflow -dtensor-merge-clusters -verify-diagnostics | FileCheck %s
// -----
// Check tf.If control flow ops are decomposed correctly.
// CHECK-LABEL: module @test_if_decomposed
module @test_if_decomposed {
func.func @main(%arg0: tensor<?xi32>) -> tensor<?xi32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[PREDICATE_RECV_OUT:.*]] = "tf.DTensorRecv"
// CHECK-SAME: key = "SendRecvKeyForControlflow_0"
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE_RECV_OUT]])
// CHECK-NEXT: "tf.D"
// CHECK-NEXT: "tf.Yield"
// CHECK: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"
// CHECK-SAME: () -> ()
// CHECK-NEXT: %[[CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: "tf.A"
// CHECK-NEXT: "tf.B"
// CHECK-NEXT: %[[PREDICATE_OUT:.*]] = "tf.G"
// CHECK-NEXT: "tf.F"
// CHECK-NEXT: "tf.DTensorSend"(%[[PREDICATE_OUT]])
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE_OUT]])
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: %[[E_OUT:.*]] = "tf.E"
// CHECK-NEXT: tf_device.return %[[E_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
// CHECK-NEXT: return %[[CLUSTER_OUT]]
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
%7 = "tf.F"() : () -> tensor<?xi32>
"tf.IfRegion"(%6) ({
"tf_device.cluster"() ({
"tf.D"() {} : () -> ()
tf_device.return
}) {_mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"} : () -> ()
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
}
// -----
// Check decomposing 2 tf_device.cluster ops inside then/else branch of tf.If.
// CHECK-LABEL: module @test_if_then_else_branches
module @test_if_then_else_branches {
// CHECK: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) -> tensor<?xi32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[PREDICATE_RECV_OUT:.*]] = "tf.DTensorRecv"
// CHECK-SAME: key = "SendRecvKeyForControlflow_0"
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE_RECV_OUT]])
// CHECK-NEXT: "tf.D"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"
// CHECK-SAME: () -> ()
// CHECK-NEXT: "tf_device.cluster"
// CHECK: %[[PREDICATE_RECV_OUT_2:.*]] = "tf.DTensorRecv"
// CHECK-SAME: key = "SendRecvKeyForControlflow_1"
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE_RECV_OUT_2]])
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.I"
// CHECK-NEXT: "tf.Yield"
// CHECK: tf_device.return
// CHECK-NEXT: _mesh = "TPU|a=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-SAME: () -> ()
// CHECK-NEXT: %[[CPU_CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: "tf.A"()
// CHECK-NEXT: "tf.B"()
// CHECK-NEXT: %[[PREDICATE_OUT:.*]] = "tf.G"()
// CHECK-NEXT: "tf.F"()
// CHECK-NEXT: "tf.DTensorSend"(%[[PREDICATE_OUT]])
// CHECK-SAME: key = "SendRecvKeyForControlflow_0"
// CHECK-NEXT: "tf.DTensorSend"(%[[PREDICATE_OUT]])
// CHECK-SAME: key = "SendRecvKeyForControlflow_1"
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE_OUT]])
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: %[[E_OUT:.*]] = "tf.E"()
// CHECK-NEXT: tf_device.return %[[E_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
// CHECK-NEXT: return %[[CPU_CLUSTER_OUT]]
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
%7 = "tf.F"() : () -> tensor<?xi32>
"tf.IfRegion"(%6) ({
"tf_device.cluster"() ({
"tf.D"() {} : () -> ()
tf_device.return
}) {_mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"} : () -> ()
"tf.Yield"() : () -> ()
}, {
"tf_device.cluster"() ({
"tf.I"() {} : () -> ()
tf_device.return
}) {_mesh = "TPU|a=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
}
// -----
// Check decomposing tf_device cluster inside tested tf.If op.
// CHECK-LABEL: module @test_nested_cluster_inside_if
// CHECK: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
module @test_nested_cluster_inside_if {
func.func @main(%arg0: tensor<i32>) -> tensor<?xi32> {
// CHECK: "tf_device.cluster"()
// CHECK: %[[OUTER_PREDICATE_RECV:.*]] = "tf.DTensorRecv"()
// CHECK-SAME: key = "SendRecvKeyForControlflow_1"
// CHECK-NEXT: "tf.IfRegion"(%[[OUTER_PREDICATE_RECV]])
// CHECK-NEXT: %[[INNER_PREDICATE_RECV:.*]] = "tf.DTensorRecv"()
// CHECK-SAME: key = "SendRecvKeyForControlflow_0"
// CHECK-NEXT: "tf.IfRegion"(%[[INNER_PREDICATE_RECV]])
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.I"
// CHECK-NEXT: "tf.D"
// CHECK: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"
// CHECK-SAME: () -> ()
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.A"
// CHECK-NEXT: "tf.B"
// CHECK-NEXT: %[[OUTER_PREDICATE:.*]] = "tf.G"
// CHECK-NEXT: "tf.DTensorSend"(%[[OUTER_PREDICATE]])
// CHECK-NEXT: "tf.IfRegion"(%[[OUTER_PREDICATE]])
// CHECK-NEXT: %[[INNER_PREDICATE:.*]] = "tf.H"
// CHECK-NEXT: "tf.DTensorSend"(%[[INNER_PREDICATE]])
// CHECK-NEXT: "tf.IfRegion"(%[[INNER_PREDICATE]])
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: %[[E_OUT:.*]] = "tf.E"
// CHECK-NEXT: tf_device.return %[[E_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
"tf.IfRegion"(%6) ({
%7 = "tf.H"(%4) : (tensor<?xi32>) -> (tensor<i1>)
"tf.IfRegion"(%7)({
"tf.Yield"() : () -> ()
},
{
"tf_device.cluster"() ({
%8 = "tf.I"() : () -> (tensor<?xi32>)
"tf.D"(%8) : (tensor<?xi32>) -> ()
tf_device.return
}) {_mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"} : () -> ()
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) { is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
}
// -----
// Check nested clusters with input edges are disallowed.
func.func @main(%arg0: tensor<?xi32>) -> tensor<?xi32> {
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
%7 = "tf.F"() : () -> tensor<?xi32>
"tf.IfRegion"(%6) ({
// expected-error @+1 {{found nested tf_device.Cluster op with inputs}}
"tf_device.cluster"() ({
"tf.D"(%4, %3, %7) {} : (tensor<?xi32>, tensor<?xi32>, tensor<?xi32>) -> ()
tf_device.return
}) {_mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"} : () -> ()
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
// -----
// Check nested clusters with outputs edges are disallowed.
func.func @main(%arg0: tensor<?xi32>) -> tensor<?xi32> {
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
%7 = "tf.F"() : () -> tensor<?xi32>
"tf.IfRegion"(%6) ({
// expected-error @+1 {{found nested tf_device.Cluster op with outputs}}
%9 = "tf_device.cluster"() ({
%8 = "tf.D"() : () -> tensor<?xi32>
tf_device.return %8 : tensor<?xi32>
}) {_mesh = "TPU|x=1|0|0|/job:localhost/task:0/device:TPU:0"} : () -> (tensor<?xi32>)
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
tensorflow/dtensor/mlir/tests/designate_resource_handle_mesh.mlir
+62
View File
@@ -0,0 +1,62 @@
// 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: dtensor-opt %s -split-input-file -dtensor-designate-resource-handle-mesh -verify-diagnostics | FileCheck %s
// Check that pass is no-op for tf_device.cluster ops that does not contain
// tf.VarHandle / tf.DestroyResource op.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<i32>) {
// CHECK: tf_device.cluster
// CHECK-NEXT: "tf.A"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) {_mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> (tensor<i32>)
func.return %0 : tensor<i32>
}
// -----
// Check that empty mesh is assigned to cluster containing tf.VarHandle op.
// CHECK-LABEL: func @check_empty_mesh_assigned_varhandle_op
func.func @check_empty_mesh_assigned_varhandle_op() -> (tensor<!tf_type.resource<tensor<i32>>>) {
// CHECK: tf_device.cluster
// CHECK-NEXT: %[[RESOURCE_OUT:.*]] = "tf.VarHandleOp"()
// CHECK-NEXT: tf_device.return %[[RESOURCE_OUT]]
// CHECK-NEXT: _mesh = "empty_mesh"
%1 = "tf_device.cluster"() ({
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
tf_device.return %0 : tensor<!tf_type.resource<tensor<i32>>>
}) : () -> (tensor<!tf_type.resource<tensor<i32>>>)
func.return %1 : tensor<!tf_type.resource<tensor<i32>>>
}
// -----
// Check that non-empty mesh is assigned to cluster containing tf.DestroyResource op.
// CHECK-LABEL: func @check_mesh_assigned_destroy_resource_op
func.func @check_mesh_assigned_destroy_resource_op(%arg0: tensor<!tf_type.resource>) -> () {
// CHECK: tf_device.cluster
// CHECK-NEXT: "tf.DestroyResourceOp"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
"tf_device.cluster"() ({
"tf.DestroyResourceOp"(%arg0) : (tensor<!tf_type.resource>) -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> ()
func.return
}
tensorflow/dtensor/mlir/tests/device_mesh_cluster_coarsening.mlir
+181
View File
@@ -0,0 +1,181 @@
// 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: dtensor-opt %s -dtensor-device-mesh-cluster-coarsening -split-input-file -verify-diagnostics | FileCheck %s
// The layout string is the text format escaped string from a simple 2 device, 1d mesh.
//
// The corresponding proto to CPU used in below tests:
// mesh_config {
// mesh_dimensions {
// dimension {
// name: "batch"
// }
// size: 2
// }
// devices: "/job:localhost/task:0/device:CPU:0"
// devices: "/job:localhost/task:0/device:CPU:1"
// }
//
// For TPU, just replace CPU string in devices with TPU.
// CHECK-LABEL: func @coarsen_cluster_with_same_device_config
func.func @coarsen_cluster_with_same_device_config() {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.B"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) {_mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<i32>)
%2 = "tf_device.cluster"() ({
%3 = "tf.B"() : () -> tensor<f32>
tf_device.return %3 : tensor<f32>
}) {_mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<f32>)
func.return
}
// -----
// CHECK-LABEL: func @coarsening_clusters_with_different_configs
func.func @coarsening_clusters_with_different_configs() {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.B"
// CHECK-NEXT: tf_device.return %[[A_OUT]], %[[B_OUT]]
// CHECK-NEXT: _mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
//
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[C_OUT:.*]] = "tf.C"
// CHECK-NEXT: %[[D_OUT:.*]] = "tf.D"
// CHECK-NEXT: tf_device.return %[[D_OUT]]
// CHECK: _mesh = "TPU|batch=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
//
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[E_OUT:.*]] = "tf.E"
// CHECK-NEXT: tf_device.return %[[E_OUT]]
// CHECK: _mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) {_mesh= "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<i32>)
%2 = "tf_device.cluster"() ({
%3 = "tf.B"() : () -> tensor<f32>
tf_device.return %3 : tensor<f32>
}) {_mesh= "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<f32>)
%5 = "tf_device.cluster"() ({
%4 = "tf.C"(%2, %0) : (tensor<f32>, tensor<i32>) -> tensor<f32>
tf_device.return %4 : tensor<f32>
}) {_mesh= "TPU|batch=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> (tensor<f32>)
%7 = "tf_device.cluster"() ({
%6 = "tf.D"(%0, %5) : (tensor<i32>, tensor<f32>) -> tensor<f32>
tf_device.return %6 : tensor<f32>
}) {_mesh= "TPU|batch=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> (tensor<f32>)
%9 = "tf_device.cluster"() ({
%8 = "tf.E"(%7) : (tensor<f32>) -> tensor<f32>
tf_device.return %8 : tensor<f32>
}) {_mesh= "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<i32>)
func.return
}
// -----
func.func @check_cluster_without_mesh_attribute_disallowed() {
// expected-error @+1 {{failed to merge mesh cluster as cluster does not have mesh attribute. This is likely due to problem in mesh propagation}}
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> (tensor<i32>)
%2 = "tf_device.cluster"() ({
%3 = "tf.B"() : () -> tensor<f32>
tf_device.return %3 : tensor<f32>
}) {_mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<f32>)
func.return
}
// -----
// Check ops in tf.WhileRegions are grouped into cluster correctly.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4xf32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3", tf._mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"})
-> (tensor<4xf32> {tf._default_layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.WhileRegion"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: constant
// CHECK-NEXT: "tf.NotEqual"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: constant
// CHECK-NEXT: "tf.Sub"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
// CHECK-NEXT: "tf.Yield"
// CHECK-NEXT: (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
%15:2 = "tf_device.cluster"() ({
%2:2 = "tf.WhileRegion"(%arg1, %arg0) ({
^bb0(%carg0: tensor<4xf32>, %carg1: tensor<i32>):
%11 = "tf_device.cluster"() ({
%limit = arith.constant dense<5> : tensor<i32>
tf_device.return %limit : tensor<i32>
}){_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"} : () -> tensor<i32>
%12 = "tf_device.cluster"() ({
%cond = "tf.NotEqual"(%carg1, %11) : (tensor<i32>, tensor<i32>) -> tensor<i1>
tf_device.return %cond : tensor<i1>
}) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"}: () -> tensor<i1>
"tf.Yield"(%12) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<4xf32>, %barg1: tensor<i32>):
%13 = "tf_device.cluster"() ({
%one = arith.constant dense<1.0> : tensor<4xf32>
tf_device.return %one: tensor<4xf32>
}) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"}: () -> tensor<4xf32>
%14 = "tf_device.cluster"() ({
%sub = "tf.Sub"(%barg0, %13) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
tf_device.return %sub: tensor<4xf32>
}) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"}: () -> tensor<4xf32>
"tf.Yield"(%14, %barg1) : (tensor<4xf32>, tensor<i32>) -> ()
}) {is_stateless = true} : (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
tf_device.return %2#0, %2#1 : tensor<4xf32>, tensor<i32>
}) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"} : () -> (tensor<4xf32>, tensor<i32>)
%16 = "tf_device.cluster"() ({
%5 = "tf.Identity"(%15#0) : (tensor<4xf32>) -> (tensor<4xf32>)
tf_device.return %5 : tensor<4xf32>
}){_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"} : () -> tensor<4xf32>
func.return %16 : tensor<4xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_all_gather.mlir
+45
View File
@@ -0,0 +1,45 @@
// 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: env LOWER_DTENSOR_GATHER_TO_COLLECTIVE_GATHER_V2=1 dtensor-opt %s -split-input-file -dtensor-all-gather-lowering -verify-diagnostics | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @lower_allgather_tpu_mesh
func.func @lower_allgather_tpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<2x2xf32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<2x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[UPDATED:.*]] = "tf.XlaDynamicUpdateSlice"(%{{[0-9]*}}, %arg1, %{{[0-9]*}})
// CHECK: %[[REDUCED:.*]] = "tf.DTensorAllReduce"(%[[UPDATED]]
// CHECK: tf_device.return %[[REDUCED]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllGather"(%arg1) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y
=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xf32>) -> tensor<2x4xf32>
tf_device.return %1 : tensor<2x4xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<2x4xf32>
func.return %0 : tensor<2x4xf32>
}
// CHECK-LABEL: func @lower_allgather_gpu_mesh
func.func @lower_allgather_gpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<2x2xf32> {tf._layout = "sharding_specs:x,y, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"}) -> tensor<2x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Transpose"(%arg1
// CHECK: %[[ALLGATHER_OUT:.*]] = "tf.CollectiveGatherV2"
// CHECK: "tf.Transpose"(%[[ALLGATHER_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllGather"(%arg1) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:GPU|x=2,y
=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>} : (tensor<2x2xf32>) -> tensor<2x4xf32>
tf_device.return %1 : tensor<2x4xf32>
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> tensor<2x4xf32>
func.return %0 : tensor<2x4xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_all_scatter.mlir
+37
View File
@@ -0,0 +1,37 @@
// 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: dtensor-opt %s -split-input-file -dtensor-all-scatter-lowering -verify-diagnostics | FileCheck %s --dump-input=fail
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<4x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[MOD_CONST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DIV_CONST:.*]] = "tf.Const"
// CHECK-NEXT: %[[PRE_MESH_COORDS:[0-9]*]] = "tf.Div"(%arg0, %[[DIV_CONST]])
// CHECK-NEXT: %[[MESH_COORDS:.*]] = "tf.FloorMod"(%[[PRE_MESH_COORDS]], %[[MOD_CONST]])
// CHECK-NEXT: %[[SLICE_SHAPE:.*]] = "tf.Const"
// CHECK-NEXT: %[[PRE_SLICE_OFFSET:.*]] = "tf.Const"
// CHECK-NEXT: %[[SLICE_OFFSET:[0-9]*]] = "tf.MatMul"(%[[MESH_COORDS]], %[[PRE_SLICE_OFFSET]])
// CHECK-NEXT: %[[SQUEEZED_OFFSET:[0-9]*]] = "tf.Squeeze"(%[[SLICE_OFFSET]])
// CHECK-NEXT: %[[SLICE:[0-9]*]] = "tf.Slice"(%arg1, %[[SQUEEZED_OFFSET]], %[[SLICE_SHAPE]])
// CHECK: tf_device.return %[[SLICE]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllScatter"(%arg1) {_layout = ["sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"],
input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>,
output_layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x2xf32>
tf_device.return %1 : tensor<4x2xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<4x2xf32>
func.return %0 : tensor<4x2xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_allreduce_combine_optimization.mlir
+289
View File
@@ -0,0 +1,289 @@
// 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: dtensor-opt %s -split-input-file -dtensor-allreduce-combine-optimization -verify-diagnostics | FileCheck %s
// Check that independent DTensorAllReduce ops of the same element type and group assignment are combined.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[VAL_1:.*]] = "tf.Const"
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<4x4xf32>}> : () -> tensor<4x4xf32>
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"()
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<2x2xi32>}> : () -> tensor<2x2xi32>
// CHECK: %[[VAL_2:.*]] = "tf.Const"
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<4x4xf32>}> : () -> tensor<4x4xf32>
//
// CHECK: %[[FILL:.*]] = "tf.Fill"
// CHECK: %[[FLATTEN_1:.*]] = "tf.Reshape"(%[[VAL_1]], %cst_{{[0-9]*}})
// CHECK: %[[UPDATE_1:.*]] = "tf.TensorStridedSliceUpdate"(%[[FILL]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[FLATTEN_1]])
// CHECK: %[[FLATTEN_2:.*]] = "tf.Reshape"(%[[VAL_2]], %cst_{{[0-9]*}})
// CHECK: %[[UPDATE_2:.*]] = "tf.TensorStridedSliceUpdate"(%[[UPDATE_1]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[FLATTEN_2]])
// CHECK: %[[ALL_REDUCE:.*]] = "tf.DTensorAllReduce"(%[[UPDATE_2]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[SLICE_1:.*]] = "tf.Slice"(%[[ALL_REDUCE]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[UNFLATTEN_1:.*]] = "tf.Reshape"(%[[SLICE_1]], %cst_{{[0-9]*}})
// CHECK: %[[SLICE_2:.*]] = "tf.Slice"(%[[ALL_REDUCE]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[UNFLATTEN_2:.*]] = "tf.Reshape"(%[[SLICE_2]], %cst_{{[0-9]*}})
//
// CHECK: %[[ADD:.*]] = "tf.Add"(%[[UNFLATTEN_1]], %[[UNFLATTEN_2]])
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%5 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%6 = "tf.DTensorAllReduce"(%4, %5) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%7 = "tf.Add"(%3, %6) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
"tf_device.return"(%7) : (tensor<4x4xf32>) -> ()
}) : () -> tensor<4x4xf32>
"func.return"() : () -> ()
}
// -----
// Check that two groups of interdependent DTensorAllReduce ops are combined layer by layer.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[VAL_1:.*]] = "tf.Const"
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<4x4xf32>}> : () -> tensor<4x4xf32>
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"()
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<2x2xi32>}> : () -> tensor<2x2xi32>
// CHECK: %[[VAL_2:.*]] = "tf.Const"
// CHECK-SAME: <{value = dense<{{1.0.*}}> : tensor<4x4xf32>}> : () -> tensor<4x4xf32>
//
//
// CHECK: %[[ALL_REDUCE_0:.*]] = "tf.DTensorAllReduce"(%[[VAL_2]], %[[GROUP_ASSIGNMENT]])
//
// CHECK: %[[FILL_1:.*]] = "tf.Fill"
// CHECK: %[[FLATTEN_1:.*]] = "tf.Reshape"(%[[VAL_1]], %cst_{{[0-9]*}})
// CHECK: %[[UPDATE_1:.*]] = "tf.TensorStridedSliceUpdate"(%[[FILL_1]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[FLATTEN_1]])
// CHECK: %[[FLATTEN_2:.*]] = "tf.Reshape"(%[[ALL_REDUCE_0]], %cst_{{[0-9]*}})
// CHECK: %[[UPDATE_2:.*]] = "tf.TensorStridedSliceUpdate"(%[[UPDATE_1]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[FLATTEN_2]])
// CHECK: %[[ALL_REDUCE_1:.*]] = "tf.DTensorAllReduce"(%[[UPDATE_2]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[SLICE_1:.*]] = "tf.Slice"(%[[ALL_REDUCE_1]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[UNFLATTEN_1:.*]] = "tf.Reshape"(%[[SLICE_1]], %cst_{{[0-9]*}})
// CHECK: %[[SLICE_2:.*]] = "tf.Slice"(%[[ALL_REDUCE_1]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[UNFLATTEN_2:.*]] = "tf.Reshape"(%[[SLICE_2]], %cst_{{[0-9]*}})
//
// CHECK: %[[FILL_2:.*]] = "tf.Fill"
// CHECK: %[[FLATTEN_3:.*]] = "tf.Reshape"(%[[UNFLATTEN_1]], %cst_{{[0-9]*}})
// CHECK: %[[UPDATE_3:.*]] = "tf.TensorStridedSliceUpdate"(%[[FILL_2]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[FLATTEN_3]])
// CHECK: %[[FLATTEN_4:.*]] = "tf.Reshape"(%[[UNFLATTEN_2]], %cst_{{[0-9]*}})
// CHECK: %[[UPDATE_4:.*]] = "tf.TensorStridedSliceUpdate"(%[[UPDATE_3]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[FLATTEN_4]])
// CHECK: %[[ALL_REDUCE_2:.*]] = "tf.DTensorAllReduce"(%[[UPDATE_4]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[SLICE_3:.*]] = "tf.Slice"(%[[ALL_REDUCE_2]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[UNFLATTEN_3:.*]] = "tf.Reshape"(%[[SLICE_3]], %cst_{{[0-9]*}})
// CHECK: %[[SLICE_4:.*]] = "tf.Slice"(%[[ALL_REDUCE_2]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[UNFLATTEN_4:.*]] = "tf.Reshape"(%[[SLICE_4]], %cst_{{[0-9]*}})
//
// CHECK: %[[ADD:.*]] = "tf.Add"(%[[UNFLATTEN_3]], %[[UNFLATTEN_4]])
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.DTensorAllReduce"(%3, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%5 = "tf.Const"() {value = dense<1.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%6 = "tf.DTensorAllReduce"(%5, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%7 = "tf.DTensorAllReduce"(%6, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%8 = "tf.DTensorAllReduce"(%7, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%9 = "tf.Add"(%4, %8) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
"tf_device.return"(%9) : (tensor<4x4xf32>) -> ()
}) : () -> tensor<4x4xf32>
"func.return"() : () -> ()
}
// -----
// Check that DTensorAllReduce ops across region boundaries are not combined.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[VAL:.*]] = "tf.Const"
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<4x4xf32>}> : () -> tensor<4x4xf32>
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"()
// CHECK-SAME: <{value = dense<{{.*}}> : tensor<2x2xi32>}> : () -> tensor<2x2xi32>
// CHECK: %[[ALL_REDUCE_1:.*]] = "tf.DTensorAllReduce"(%[[VAL]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
//
// CHECK: "tf.WhileRegion"
//
// CHECK: "tf.ToBool"
// CHECK-NEXT: "tf.Yield"
//
// CHECK: %[[WHILE_FILL:.*]] = "tf.Fill"
// CHECK: %[[WHILE_FLATTEN_1:.*]] = "tf.Reshape"(%[[VAL]], %cst_{{[0-9]*}})
// CHECK: %[[WHILE_UPDATE_1:.*]] = "tf.TensorStridedSliceUpdate"(%[[WHILE_FILL]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[WHILE_FLATTEN_1]])
// CHECK: %[[WHILE_FLATTEN_2:.*]] = "tf.Reshape"(%[[VAL]], %cst_{{[0-9]*}})
// CHECK: %[[WHILE_UPDATE_2:.*]] = "tf.TensorStridedSliceUpdate"(%[[WHILE_UPDATE_1]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %cst_{{[0-9]*}}, %[[WHILE_FLATTEN_2]])
// CHECK: %[[WHILE_ALL_REDUCE:.*]] = "tf.DTensorAllReduce"(%[[WHILE_UPDATE_2]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[WHILE_SLICE_1:.*]] = "tf.Slice"(%[[WHILE_ALL_REDUCE]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[WHILE_UNFLATTEN_1:.*]] = "tf.Reshape"(%[[WHILE_SLICE_1]], %cst_{{[0-9]*}})
// CHECK: %[[WHILE_SLICE_2:.*]] = "tf.Slice"(%[[WHILE_ALL_REDUCE]], %cst_{{[0-9]*}}, %cst_{{[0-9]*}})
// CHECK: %[[WHILE_UNFLATTEN_2:.*]] = "tf.Reshape"(%[[WHILE_SLICE_2]], %cst_{{[0-9]*}})
// CHECK: %[[WHILE_ADD:.*]] = "tf.Add"(%[[WHILE_UNFLATTEN_1]], %[[WHILE_UNFLATTEN_2]])
// CHECK: "tf.Yield"(%[[WHILE_ADD]])
//
// CHECK: %[[ALL_REDUCE_2:.*]] = "tf.DTensorAllReduce"(%[[VAL]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// CHECK: %[[ADD:.*]] = "tf.Add"(%[[ALL_REDUCE_1]], %[[ALL_REDUCE_2]])
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.WhileRegion"(%1) ({
^bb0(%arg: tensor<4x4xf32>):
%5 = "tf.ToBool"(%arg) : (tensor<4x4xf32>) -> tensor<i1>
"tf.Yield"(%5) : (tensor<i1>) -> ()
}, {
^bb0(%arg: tensor<4x4xf32>):
%5 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%6 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%7 = "tf.Add"(%5, %6) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
"tf.Yield"(%7) : (tensor<4x4xf32>) -> ()
}) {is_stateless = true} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%5 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%6 = "tf.Add"(%3, %5) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
"tf_device.return"(%6) : (tensor<4x4xf32>) -> ()
}) : () -> tensor<4x4xf32>
"func.return"() : () -> ()
}
// -----
module attributes {dtensor.all_reduce_combiner.num_ops_in_group = 2} {
// Check that when DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_GROUP_SIZE is set,
// independent DTensorAllReduce ops of the same element type and group
// assignment are combined no more than the specified size. Use of dummy All-
// Reduces (of the same input) gaurantees ops to be grouped together if envvar
// is not specified.
// The following scenario should have 3 groups *without* envvar set:
// group 1: 2 all reduces
// group 2: 3 all reduces
// group 3: 4 all reduces
// With DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_GROUP_SIZE=2, we expect to have
// the following 5 groups:
// group 1: 2 all reduces (original group, test for exact match of size)
// group 2: 2 all reduces (2/3 of original group 2, test of uneven split)
// group 3: 1 all reduces (1/3 of original group 2, test of uneven split)
// group 4: 2 all reduces (2/4 of original group 3, test of even split)
// group 5: 2 all reduces (2/4 of original group 3, test of even split)
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[ALL_REDUCE_1:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[ALL_REDUCE_2:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[ALL_REDUCE_3:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// CHECK: %[[ALL_REDUCE_4:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[ALL_REDUCE_5:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.Const"() {value = dense<1.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%4 = "tf.Const"() {value = dense<[[3, 2], [1, 0]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%5 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%6 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%7 = "tf.DTensorAllReduce"(%3, %4) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%8 = "tf.DTensorAllReduce"(%3, %4) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%9 = "tf.DTensorAllReduce"(%3, %4) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%10 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%11 = "tf.Const"() {value = dense<[[0, 1], [3, 2]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%12 = "tf.DTensorAllReduce"(%10, %11) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%13 = "tf.DTensorAllReduce"(%10, %11) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%14 = "tf.DTensorAllReduce"(%10, %11) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%15 = "tf.DTensorAllReduce"(%10, %11) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%16 = "tf.Add"(%9, %15) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
"tf_device.return"(%16) : (tensor<4x4xf32>) -> ()
}) : () -> tensor<4x4xf32>
"func.return"() : () -> ()
}
}
// -----
module attributes {dtensor.all_reduce_combiner.topological_distance = 2} {
// Check that when topologicial grouping is enabled in AllReduce combiner, the
// independent DTensorAllReduce ops of the same element type and group assign-
// ment are combined according to the topological distance between two ops.
//
// The following scenario would have 1 group of 7 AllReduces when topological
// distance is *not* set.
// - level 1: %4, %5 (case: <= topo_dist, simple case with same level)
// - level 2: %7 (case: <= topo_dist, simple case for eligible to group)
// - level 4: %16 (case: <= topo_dist, out of order, test for topo sort)
// - level 5: %15 (case: < topo_dist, out of order, test for topo sort)
// - level 8: %14 (case: > topo_dist, ineligible to group and out of order),
// %17 (case: > topo_dist, ineligible to group with 1st group,
// but should get grouped with %14)
//
// Detailed level computations are listed in the test below.
//
// With topological_distance set to 2, we expect the following grouping result
// - group 1: %4, %5, %7, %15, %16
// - group 2: %14, %17
//
// Note use of dummy AllReduces (with the same input) gaurantees ops to be
// grouped together if topologicial grouping is not enabled.
//
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[ALL_REDUCE_1:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[ALL_REDUCE_2:.*]] = "tf.DTensorAllReduce"
// CHECK-SAME: (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// CHECK: %[[ADD:.*]] = "tf.Add"
// CHECK-SAME: (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%0 = "tf_device.cluster"() ({
// topological level 0 for all tf.Const
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.Const"() {value = dense<1.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
// %4 topological_level: 1 = max(0, 0) + 1
%4 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %5 topological_level: 1 = max(0, 0) + 1
%5 = "tf.DTensorAllReduce"(%3, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %6 topological_level: 1 = max(0, 0) + 1
%6 = "tf.Add"(%1, %3) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
// %7 topological_level: 2 = max(1, 0) + 1
%7 = "tf.DTensorAllReduce"(%6, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// Dummy Adds to construct depth in compute graph
// %8 topological_level: 2 = max(1, 0) + 1
// %9 topological_level: 3 = max(2, 0) + 1
// %10 topological_level: 4 = max(3, 0) + 1
// %11 topological_level: 5 = max(4, 0) + 1
// %12 topological_level: 6 = max(5, 0) + 1
// %13 topological_level: 7 = max(6, 0) + 1
%8 = "tf.Add"(%6, %3) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%9 = "tf.Add"(%8, %1) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%10 = "tf.Add"(%9, %3) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%11 = "tf.Add"(%10, %1) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%12 = "tf.Add"(%11, %3) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%13 = "tf.Add"(%12, %1) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
// %14 topological_level: 8 = max(7, 0) + 1
%14 = "tf.DTensorAllReduce"(%13, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %15 topological_level: 5 = max(4, 0) + 1
%15 = "tf.DTensorAllReduce"(%10, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %16 topological_level: 4 = max(3, 0) + 1
%16 = "tf.DTensorAllReduce"(%9, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %17 topological_level: 8 = max(7, 0) + 1
%17 = "tf.DTensorAllReduce"(%13, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*GPU"], device_type = "GPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%18 = "tf.Add"(%15, %7) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
"tf_device.return"(%18) : (tensor<4x4xf32>) -> ()
}) : () -> tensor<4x4xf32>
"func.return"() : () -> ()
}
}
tensorflow/dtensor/mlir/tests/dtensor_allreduce_lowering.mlir
+55
View File
@@ -0,0 +1,55 @@
// 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: dtensor-opt -split-input-file -dtensor-all-reduce-lowering -verify-diagnostics %s| FileCheck %s --dump-input=fail
// Check the lowering of AllReduce on TPU with sum reduction.
// CHECK-LABEL: func @lower_allreduce_sum
func.func @lower_allreduce_sum() -> (tensor<4096x8192xf32>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[ALLREDUCE_OUT:.*]] = "tf.XlaAllReduce"(%[[CONST_OUT_1]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-NEXT return %[[ALLREDUCE_OUT]]
%0 = "tf.Const"() {value = dense<0.0> : tensor<4096x8192xf32>} : () -> tensor<4096x8192xf32>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>) -> tensor<4096x8192xf32>
func.return %2: tensor<4096x8192xf32>
}
// -----
// CHECK-LABEL: func @lower_all_reduce_gpu_mesh
func.func @lower_all_reduce_gpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<4096x8192xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"}) -> tensor<4096x8192xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[DEVICE_ID_RESHAPE:.*]] = "tf.Reshape"(%arg0
// CHECK: %[[RELATIVE_DEVICE_ID:.*]] = "tf.Sub"(%[[DEVICE_ID_RESHAPE]]
// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
// CHECK-DAG: %[[DEVICE_ID_TO_GROUP_KEY:.*]] = "tf.Const"() <{value = dense<[[[GROUP_KEYS:.*]]]> : tensor<8xi32>}>
// CHECK: %[[GROUP_KEY_SLICE:.*]] = "tf.Slice"(%[[DEVICE_ID_TO_GROUP_KEY]], %[[RELATIVE_DEVICE_ID]], %[[CONST_1]]
// CHECK: %[[GROUP_KEY_RESHAPE:.*]] = "tf.Reshape"(%[[GROUP_KEY_SLICE]]
// CHECK-DAG: %[[INSTANCE_KEY:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}>
// CHECK-DAG: %[[GROUP_SIZE:.*]] = "tf.Const"() <{value = dense<2> : tensor<i32>}>
// CHECK: %[[REDUCE_OUT:.*]] = "tf.CollectiveReduceV2"(%arg1, %[[GROUP_SIZE]], %[[GROUP_KEY_RESHAPE]], %[[INSTANCE_KEY]])
// CHECK-SAME: final_op = "Id"
// CHECK-SAME: merge_op = "Add"
// CHECK-NEXT tf_device.return %[[REDUCE_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%3 = "tf.DTensorAllReduce"(%arg1, %1) {_layout = ["sharding_specs:x,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"], device_type = "/job:localhost/replica:0/task:0/device:GPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>) -> tensor<4096x8192xf32>
tf_device.return %3 : tensor<4096x8192xf32>
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> tensor<4096x8192xf32>
func.return %0 : tensor<4096x8192xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_allreduce_scatter_optimization.mlir
+126
View File
@@ -0,0 +1,126 @@
// 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: dtensor-opt %s -split-input-file -dtensor-allreduce-scatter-optimization -verify-diagnostics | FileCheck %s
//
//
// CHECK-LABEL: func @all_reduce_only
func.func @all_reduce_only() {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// CHECK: "tf.DTensorAllReduce"
// CHECK-NOT: "tf.DTensorReduceScatter"
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*TPU"], device_type = "TPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
"tf_device.return"(%3) : (tensor<4x4xf32>) -> ()
}) : () -> tensor<4x4xf32>
"func.return"() : () -> ()
}
// -----
// CHECK-LABEL: func @all_reduce_scatter_2d_major_dim
func.func @all_reduce_scatter_2d_major_dim() {
// CHECK: %[[INPUT:.*]] = "tf.Const"() <{value = dense<0.0
// CHECK: %[[GROUP:.*]] = "tf.Const"() <{value =
// CHECK-SAME{LITERAL}: dense<[[0, 2], [1, 3]]>
// CHECK: %[[SCATTER_DIM:.*]] = "tf.Const"() <{value = dense<0>
// CHECK: "tf.DTensorReduceScatter"(%[[INPUT]], %[[GROUP]], %[[SCATTER_DIM]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-NOT: "tf.DTensorAllReduce"
// CHECK-NOT: "tf.DTensorAllScatter"
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 2], [1, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "TPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.DTensorAllScatter"(%3) {_layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU"], input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<2x4xf32>
"tf_device.return"(%4) : (tensor<2x4xf32>) -> ()
}) : () -> tensor<2x4xf32>
"func.return"() : () -> ()
}
// -----
// CHECK-LABEL: func @all_reduce_scatter_2d_minor_dim
func.func @all_reduce_scatter_2d_minor_dim() {
// CHECK: %[[INPUT:.*]] = "tf.Const"() <{value = dense<0.0
// CHECK: %[[GROUP:.*]] = "tf.Const"() <{value =
// CHECK-SAME{LITERAL}: dense<[[0, 2], [1, 3]]>
// CHECK: %[[SCATTER_DIM:.*]] = "tf.Const"() <{value = dense<1>
// CHECK: "tf.DTensorReduceScatter"(%[[INPUT]], %[[GROUP]], %[[SCATTER_DIM]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-NOT: "tf.DTensorAllReduce"
// CHECK-NOT: "tf.DTensorAllScatter"
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 2], [1, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "TPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.DTensorAllScatter"(%3) {_layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=2|*TPU"], input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, output_layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x2xf32>
"tf_device.return"(%4) : (tensor<4x2xf32>) -> ()
}) : () -> tensor<4x2xf32>
"func.return"() : () -> ()
}
// -----
// CHECK-LABEL: func @all_reduce_multiple_users
func.func @all_reduce_multiple_users() {
// CHECK: "tf.DTensorAllReduce"
// CHECK: "tf.DTensorAllScatter"
// CHECK-NOT: "tf.DTensorReduceScatter"
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 2], [1, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "TPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.DTensorAllScatter"(%3) {_layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU"], input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<2x4xf32>
%5 = "tf.Identity"(%3) : (tensor<4x4xf32>) -> tensor<4x4xf32>
"tf_device.return"(%4) : (tensor<2x4xf32>) -> ()
}) : () -> tensor<2x4xf32>
"func.return"() : () -> ()
}
// -----
// CHECK-LABEL: func @all_reduce_scatter_2d_mismatched_dim
func.func @all_reduce_scatter_2d_mismatched_dim() {
%0 = "tf_device.cluster"() ({
// CHECK: "tf.DTensorAllReduce"
// CHECK: "tf.DTensorAllScatter"
// CHECK-NOT: "tf.DTensorReduceScatter"
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "TPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.DTensorAllScatter"(%3) {_layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU"], input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<2x4xf32>
"tf_device.return"(%4) : (tensor<2x4xf32>) -> ()
}) : () -> tensor<2x4xf32>
"func.return"() : () -> ()
}
// -----
// CHECK-LABEL: func @all_reduce_scatter_2d_both_dims
func.func @all_reduce_scatter_2d_both_dims() {
// CHECK: "tf.DTensorAllReduce"
// CHECK: "tf.DTensorAllScatter"
// CHECK-NOT: "tf.DTensorReduceScatter"
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
%2 = "tf.Const"() {value = dense<[[0, 2], [1, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.DTensorAllReduce"(%1, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "TPU", reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
%4 = "tf.DTensorAllScatter"(%3) {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*TPU"], input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, output_layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<2x2xf32>
"tf_device.return"(%4) : (tensor<2x2xf32>) -> ()
}) : () -> tensor<2x2xf32>
"func.return"() : () -> ()
}
tensorflow/dtensor/mlir/tests/dtensor_allreduce_sum_optimization.mlir
+307
View File
@@ -0,0 +1,307 @@
// 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: dtensor-opt %s -split-input-file -dtensor-allreduce-sum-optimization -verify-diagnostics | FileCheck %s
// Check that DTensorAllReduce op with Add/AddN/AddV2 operations are optimized.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<916x8192xbf16>) {
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL1:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL2:.*]] = "tf.Const"
// CHECK-NEXT: %[[ADD_OUT:.*]] = "tf.AddN"(%[[VAL1]], %[[VAL2]])
// CHECK-NEXT: %[[REDUCTION_OUT:.*]] = "tf.DTensorAllReduce"(%[[ADD_OUT]], %[[GROUP_ASSIGNMENT]])
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.B"(%[[REDUCTION_OUT]])
%0 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%1 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%2 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () ->tensor<916x8192xbf16>
%3= "tf.DTensorAllReduce"(%1, %0) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%4= "tf.DTensorAllReduce"(%2, %0) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%5= "tf.AddN"(%3, %4) { device = ""} : (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
%6 = "tf.B"(%5) : (tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
// CHECK: %[[GROUP_ASSIGNMENT_2:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL3:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL4:.*]] = "tf.Const"
// CHECK-NEXT: %[[ADD_OUT_2:.*]] = "tf.Add"(%[[VAL3]], %[[VAL4]])
// CHECK-NEXT: %[[REDUCTION_OUT_2:.*]] = "tf.DTensorAllReduce"(%[[ADD_OUT_2]], %[[GROUP_ASSIGNMENT_2]])
// CHECK-NEXT: %[[C_OUT:.*]] = "tf.C"(%[[REDUCTION_OUT_2]])
%7 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%8 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%9 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () ->tensor<916x8192xbf16>
%10= "tf.DTensorAllReduce"(%8, %7) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%11= "tf.DTensorAllReduce"(%9, %7) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%12= "tf.Add"(%10, %11) { device = ""} : (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
%13 = "tf.C"(%12) : (tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
// CHECK: %[[GROUP_ASSIGNMENT_3:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL5:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL6:.*]] = "tf.Const"
// CHECK-NEXT: %[[ADD_OUT_3:.*]] = "tf.AddV2"(%[[VAL5]], %[[VAL6]])
// CHECK-NEXT: %[[REDUCTION_OUT_3:.*]] = "tf.DTensorAllReduce"(%[[ADD_OUT_3]], %[[GROUP_ASSIGNMENT_3]])
// CHECK-NEXT: %[[D_OUT:.*]] = "tf.D"(%[[REDUCTION_OUT_3]])
%14 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%15 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%16 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () ->tensor<916x8192xbf16>
%17= "tf.DTensorAllReduce"(%15, %14) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%18= "tf.DTensorAllReduce"(%16, %14) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%19= "tf.AddV2"(%17, %18) { device = ""} : (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
%20 = "tf.D"(%19) : (tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
func.return %6 : tensor<916x8192xbf16>
}
// -----
// Check that DTensorAllReduce op with operation with group assignment from different constant with same values are optimized correctly.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<916x8192xbf16>) {
// CHECK: %[[VAL_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL_2:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT_2:.*]] = "tf.Const"
// CHECK-NEXT: %[[ADD_OUT_1:.*]] = "tf.AddN"(%[[VAL_1]], %[[VAL_2]])
// CHECK-NEXT: %[[REDUCTION_OUT:.*]] = "tf.DTensorAllReduce"(%[[ADD_OUT_1]], %[[GROUP_ASSIGNMENT_1]])
// CHECK-NEXT: return %[[REDUCTION_OUT]]
%0 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%1 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%4= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%2 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%3= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%2, %3) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%6= "tf.AddN"(%4, %5): (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
func.return %6: tensor<916x8192xbf16>
}
// -----
// Check that DTensorAllReduce op with operation type that is not `sum` are not optimized.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<916x8192xbf16>) {
// CHECK: %[[REDUCE_OUT_1:.*]] = "tf.DTensorAllReduce"
// CHECK: %[[REDUCE_OUT_2:.*]] = "tf.DTensorAllReduce"
// CHECK: %[[ADD_OUT:.*]] = "tf.AddN"(%[[REDUCE_OUT_1]], %[[REDUCE_OUT_2]])
%0 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%1 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%4= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Mean"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%2 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%3= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%2, %3) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Max"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%6= "tf.AddN"(%4, %5): (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
func.return %6: tensor<916x8192xbf16>
}
// -----
// Check that DTensorAllReduce op with different group assignment are not optimized away.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<916x8192xbf16>) {
// CHECK: %[[REDUCE_OUT_1:.*]] = "tf.DTensorAllReduce"
// CHECK: %[[REDUCE_OUT_2:.*]] = "tf.DTensorAllReduce"
// CHECK: %[[ADD_OUT:.*]] = "tf.AddN"(%[[REDUCE_OUT_1]], %[[REDUCE_OUT_2]])
%0 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%1 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%4= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%2 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%3= "tf.Const"() {value = dense<1> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%2, %3) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%6= "tf.AddN"(%4, %5): (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
func.return %6: tensor<916x8192xbf16>
}
// -----
// Check that DTensorAllReduce op with malformed layout specification is disallowed.
func.func @main() -> (tensor<916x8192xbf16>) {
%0 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%1 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
// expected-error @+1 {{Malformed layout specification for DTensorAllReduce op found}}
%4= "tf.DTensorAllReduce"(%0, %1) {_layout = ["1234"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%2 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%3= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%2, %3) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%6= "tf.AddN"(%4, %5): (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
func.return %6: tensor<916x8192xbf16>
}
// -----
// Check that DTensorAllReduce op with missing layout specification is disallowed.
func.func @main() -> (tensor<916x8192xbf16>) {
%0 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%1 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
// expected-error @+1 {{DTensorAllReduce op must have layout specification}}
%4= "tf.DTensorAllReduce"(%0, %1) {device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%2 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xbf16>} : () -> tensor<916x8192xbf16>
%3= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%2, %3) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xbf16>, tensor<2x64xi32>) -> tensor<916x8192xbf16>
%6= "tf.AddN"(%4, %5): (tensor<916x8192xbf16>, tensor<916x8192xbf16>) -> tensor<916x8192xbf16>
func.return %6: tensor<916x8192xbf16>
}
// -----
// Check that DTensorAllReduce op inside while loop is optimized away correctly.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: %[[VAL_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[VAL_1_ID:.*]] = "tf.Identity"(%[[VAL_1]])
// CHECK-NEXT: %[[WHILE_OUT:.*]]:2 = "tf.WhileRegion"(%[[VAL_1_ID]], %[[ARG0]])
// CHECK: "tf.A"
// CHECK-NEXT: "tf.Yield"
// CHECK: ^bb0(%[[BARG0:.*]]: tensor<4xf32>, %[[BARG1:.*]]: tensor<i32>)
// CHECK: %[[INPUT0:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[GROUP:.*]] = "tf.Const"() <{value = dense<0> : tensor<2x64xi32>}> : () -> tensor<2x64xi32>
// CHECK-NEXT: %[[CAST_OUT:.*]] = "tf.Cast"(%[[INPUT0]])
// CHECK-NEXT: %[[ADD_OUT:.*]] = "tf.AddV2"(%[[CAST_OUT]], %[[BARG0]])
// CHECK-NEXT: %[[OUT:.*]] = "tf.Identity"(%[[ADD_OUT]])
// CHECK-NEXT: "tf.Yield"
// CHECK: %[[NEW_GROUP_ASSIGN:.*]] = "tf.Const"() <{value = dense<0> : tensor<2x64xi32>}> : () -> tensor<2x64xi32>
// CHECK: %[[ALL_REDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[WHILE_OUT]]#0, %[[NEW_GROUP_ASSIGN]])
%0 = "tf.Const"() {value = dense<0.0> : tensor<4xf32>} : () -> tensor<4xf32>
%2 = "tf.Identity"(%0) : (tensor<4xf32>) -> tensor<4xf32>
%9:2 = "tf.WhileRegion"(%2, %arg0) ({
^bb0(%carg0: tensor<4xf32>, %carg1: tensor<i32>):
%1 = "tf.A"(%carg1) : (tensor<i32>) -> (tensor<i1>)
"tf.Yield"(%1) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<4xf32>, %barg1: tensor<i32>):
%3 = "tf.Const"() {value = dense<0> : tensor<4xi32>} : () -> tensor<4xi32>
%4= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%3, %4) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4xi32>, tensor<2x64xi32>) -> tensor<4xi32>
%6 = "tf.Cast"(%5) : (tensor<4xi32>) -> tensor<4xf32>
%7 = "tf.AddV2"(%6, %barg0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
%8 = "tf.Identity"(%7) : (tensor<4xf32>) -> tensor<4xf32>
"tf.Yield"(%8, %barg1) : (tensor<4xf32>, tensor<i32>) -> ()
}) {is_stateless = true} : (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
"tf.C"(%9#0) : (tensor<4xf32>) -> ()
func.return
}
// -----
// Check that while op optimization is ignored if loop variant input is used for
// while loop condition logic.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: %[[WHILE_OUT:.*]]:2 = "tf.WhileRegion"
// CHECK: "tf.A"
// CHECK-NEXT: "tf.Yield"
// CHECK: ^bb0(%[[BARG0:.*]]: tensor<4xf32>, %[[BARG1:.*]]: tensor<i32>)
// CHECK: "tf.Const"()
// CHECK-NEXT: "tf.Const"()
// CHECK-NEXT: "tf.DTensorAllReduce"
// CHECK-NEXT: "tf.AddV2"
// CHECK-NEXT: "tf.Yield"
%0 = "tf.Const"() {value = dense<0.0> : tensor<4xf32>} : () -> tensor<4xf32>
%9:2 = "tf.WhileRegion"(%0, %arg0) ({
^bb0(%carg0: tensor<4xf32>, %carg1: tensor<i32>):
%1 = "tf.A"(%carg0) : (tensor<4xf32>) -> (tensor<i1>)
"tf.Yield"(%1) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<4xf32>, %barg1: tensor<i32>):
%3 = "tf.Const"() {value = dense<0.0> : tensor<4xf32>} : () -> tensor<4xf32>
%4= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%3, %4) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4xf32>, tensor<2x64xi32>) -> tensor<4xf32>
%6 = "tf.AddV2"(%5, %barg0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
"tf.Yield"(%6, %barg1) : (tensor<4xf32>, tensor<i32>) -> ()
}) {is_stateless = true} : (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
func.return
}
// -----
// Check that while op with input that is not constant zero will not trigger
// lazy all reduce optimization.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: %[[WHILE_OUT:.*]]:2 = "tf.WhileRegion"
// CHECK: "tf.A"
// CHECK-NEXT: "tf.Yield"
// CHECK: ^bb0(%[[BARG0:.*]]: tensor<4xf32>, %[[BARG1:.*]]: tensor<i32>)
// CHECK: "tf.Const"()
// CHECK-NEXT: "tf.Const"()
// CHECK-NEXT: "tf.DTensorAllReduce"
// CHECK-NEXT: "tf.AddV2"
// CHECK-NEXT: "tf.Yield"
%0 = "tf.Const"() {value = dense<[0.0, 1.0, 0.0, 0.0]> : tensor<4xf32>} : () -> tensor<4xf32>
%9:2 = "tf.WhileRegion"(%0, %arg0) ({
^bb0(%carg0: tensor<4xf32>, %carg1: tensor<i32>):
%1 = "tf.A"(%carg0) : (tensor<4xf32>) -> (tensor<i1>)
"tf.Yield"(%1) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<4xf32>, %barg1: tensor<i32>):
%3 = "tf.Const"() {value = dense<0.0> : tensor<4xf32>} : () -> tensor<4xf32>
%4= "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%5= "tf.DTensorAllReduce"(%3, %4) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4xf32>, tensor<2x64xi32>) -> tensor<4xf32>
%6 = "tf.AddV2"(%5, %barg0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
"tf.Yield"(%6, %barg1) : (tensor<4xf32>, tensor<i32>) -> ()
}) {is_stateless = true} : (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
func.return
}
// -----
// Check that optimization that does not reduce number of DTensorAllReduce is
// not applied.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<4096x8192xf32>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[REDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[CONST_OUT_1]], %[[GROUP_ASSIGNMENT]])
// CHECK-NEXT: "tf.Cast"(%[[REDUCE_OUT]])
%0 = "tf.Const"() {value = dense<0.0> : tensor<4096x8192xbf16>} : () -> tensor<4096x8192xbf16>
%1 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%2= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=64|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4096x8192xbf16>, tensor<2x64xi32>) -> tensor<4096x8192xbf16>
%3= "tf.Cast"(%2) {Truncate = false, device = ""} : (tensor<4096x8192xbf16>) -> tensor<4096x8192xf32>
func.return %3: tensor<4096x8192xf32>
}
// -----
// Check that DTensorAllReduce op moved after Identity-like operations.
// CHECK-LABEL: func @main
func.func @main() -> (tensor<916x8192xf32>) {
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[CONST_OUT_2:.*]] = "tf.Const"
// CHECK-NEXT: %[[CAST_OUT:.*]] = "tf.Cast"(%[[CONST_OUT_2]])
// CHECK-NEXT: %[[IDENTITY_OUT:.*]] = "tf.Identity"(%[[CAST_OUT]])
// CHECK-NEXT: %[[RESHAPE_CONST:.*]] = "tf.Const"
// CHECK-NEXT: %[[RESHAPE_OUT:.*]] = "tf.Reshape"(%[[IDENTITY_OUT]], %[[RESHAPE_CONST]])
// CHECK-NEXT: %[[ADD_OUT:.*]] = "tf.AddN"(%[[RESHAPE_OUT]], %[[CONST_OUT_1]])
// CHECK-NEXT: "tf.DTensorAllReduce"(%[[ADD_OUT]], %[[GROUP_ASSIGNMENT]])
%0 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>} : () -> tensor<2x64xi32>
%1 = "tf.Const"() {value = dense<0.0> : tensor<916x8192xf32>} : () -> tensor<916x8192xf32>
%2 = "tf.Const"() {value = dense<0.0> : tensor<8192x916xbf16>} : () ->tensor<8192x916xbf16>
%3= "tf.DTensorAllReduce"(%1, %0) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<916x8192xf32>, tensor<2x64xi32>) -> tensor<916x8192xf32>
%4= "tf.DTensorAllReduce"(%2, %0) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<8192x916xbf16>, tensor<2x64xi32>) -> tensor<8192x916xbf16>
%5 = "tf.Cast"(%4){Truncate = false, device = ""} : ( tensor<8192x916xbf16>) -> tensor<8192x916xf32>
%6 = "tf.Identity"(%5){Truncate = false, device = ""} : (tensor<8192x916xf32>) -> tensor<8192x916xf32>
%7 = "tf.Const"() {value = dense<[916,8192]> : tensor<2xi32>} : () -> tensor<2xi32>
%8 = "tf.Reshape"(%6, %7) : (tensor<8192x916xf32>, tensor<2xi32>) -> tensor<916x8192xf32>
%9= "tf.AddN"(%8, %3) { device = ""} : (tensor<916x8192xf32>, tensor<916x8192xf32>) -> tensor<916x8192xf32>
%10 = "tf.B"(%9) : (tensor<916x8192xf32>) -> tensor<916x8192xf32>
func.return %10 : tensor<916x8192xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_alltoall_lowering.mlir
+50
View File
@@ -0,0 +1,50 @@
// 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: dtensor-opt %s -split-input-file -dtensor-all-to-all-lowering -verify-diagnostics | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @lower_alltoall_tpu_mesh
func.func @lower_alltoall_tpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<4x2xf32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<2x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[ALLTOALL_OUT:.*]] = "tf.AllToAll"(%arg1
// CHECK: tf_device.return %[[ALLTOALL_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllToAll"(%arg1) {_layout = ["sharding_specs:x,unsharded, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"] ,
input_layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>,
output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y
=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x2xf32>) -> tensor<2x4xf32>
tf_device.return %1 : tensor<2x4xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<2x4xf32>
func.return %0 : tensor<2x4xf32>
}
// CHECK-LABEL: func @lower_alltoall_gpu_mesh
func.func @lower_alltoall_gpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<4x2xf32> {tf._layout = "sharding_specs:x,y, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"}) -> tensor<2x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[FLATTEN_OUT:.*]] = "tf.Reshape"(%arg1
// CHECK: %[[ALLTOALL_OUT:.*]] = "tf.CollectiveAllToAllV2"(%[[FLATTEN_OUT]]
// CHECK: %[[RESHAPE_1_OUT:.*]] = "tf.Reshape"(%[[ALLTOALL_OUT]]
// CHECK: %[[TRANSPOSE_OUT:.*]] = "tf.Transpose"(%[[RESHAPE_1_OUT]]
// CHECK: %[[RESHAPE_2_OUT:.*]] = "tf.Reshape"(%[[TRANSPOSE_OUT]]
// CHECK: tf_device.return %[[RESHAPE_2_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllToAll"(%arg1) {_layout = ["sharding_specs:x,unsharded, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"],
input_layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>,
output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>} : (tensor<4x2xf32>) -> tensor<2x4xf32>
tf_device.return %1 : tensor<2x4xf32>
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> tensor<2x4xf32>
func.return %0 : tensor<2x4xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_collective_type_lowering.mlir
+165
View File
@@ -0,0 +1,165 @@
// 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: dtensor-opt -split-input-file -dtensor-collective-type-lowering -verify-diagnostics %s| FileCheck %s --dump-input=fail
// Check the lowering of AllScatter on CPU with any complex reduction.
// CHECK-LABEL: func @lower_allgather_complex64
func.func @lower_allgather_complex64(%arg0: tensor<i32>,
%arg1: tensor<1x2xcomplex<f32>> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[REAL:.*]] = "tf.Real"(%arg1)
// CHECK-NEXT: %[[IMAG:.*]] = "tf.Imag"(%arg1)
// CHECK-NEXT: %[[ALLGATHER_OUT_REAL:.*]] = "tf.DTensorAllGather"(%[[REAL]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<1x2xf32>) -> tensor<2x2xf32>
// CHECK-NEXT: %[[ALLGATHER_OUT_IMAG:.*]] = "tf.DTensorAllGather"(%[[IMAG]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<1x2xf32>) -> tensor<2x2xf32>
// CHECK-NEXT: %[[OUTPUT:.*]] = "tf.Complex"(%[[ALLGATHER_OUT_REAL]], %[[ALLGATHER_OUT_IMAG]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xcomplex<f32>>
// CHECK-NEXT return %[[OUTPUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllGather"(%arg1) {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"], input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>, output_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<1x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
tf_device.return %1 : tensor<2x2xcomplex<f32>>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x2xcomplex<f32>>
func.return
}
// -----
// Check the lowering of DTensorAllToAll on TPU with any complex reduction.
// CHECK-LABEL: func @lower_all_to_all_complex128
func.func @lower_all_to_all_complex128(%arg0: tensor<i32>,
%arg1: tensor<1x2xcomplex<f64>> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[REAL:.*]] = "tf.Real"(%arg1)
// CHECK-NEXT: %[[IMAG:.*]] = "tf.Imag"(%arg1)
// CHECK-NEXT: %[[ALLTOALL_OUT_REAL:.*]] = "tf.DTensorAllToAll"(%[[REAL]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<1x2xf64>) -> tensor<2x1xf64>
// CHECK-NEXT: %[[ALLTOALL_OUT_IMAG:.*]] = "tf.DTensorAllToAll"(%[[IMAG]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<1x2xf64>) -> tensor<2x1xf64>
// CHECK-NEXT: %[[OUTPUT:.*]] = "tf.Complex"(%[[ALLTOALL_OUT_REAL]], %[[ALLTOALL_OUT_IMAG]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x1xf64>, tensor<2x1xf64>) -> tensor<2x1xcomplex<f64>>
// CHECK-NEXT return %[[OUTPUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorAllToAll"(%arg1) {_layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"], input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>, output_layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<1x2xcomplex<f64>>) -> tensor<2x1xcomplex<f64>>
tf_device.return %1 : tensor<2x1xcomplex<f64>>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x1xcomplex<f64>>
func.return
}
// -----
// Check the lowering of AllReduce on TPU with any boolean reduction.
// CHECK-LABEL: func @lower_allreduce_any_boolean
func.func @lower_allreduce_any_boolean() -> (tensor<4096x8192xi1>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[INPUT_CAST:.*]] = "tf.Cast"(%[[CONST_OUT_1]])
// CHECK-NEXT: %[[ALLREDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[INPUT_CAST]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: reduce_op = "Max"
// CHECK-NEXT: %[[OUTPUT_CAST:.*]] = "tf.Cast"(%[[ALLREDUCE_OUT]])
// CHECK-NEXT return %[[OUTPUT_CAST]]
%0 = "tf.Const"() {value = dense<1> : tensor<4096x8192xi1>} : () -> tensor<4096x8192xi1>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Any"} : (tensor<4096x8192xi1>, tensor<4x2xi32>) -> tensor<4096x8192xi1>
func.return %2: tensor<4096x8192xi1>
}
// Check the lowering of DTensorReduceScatter on TPU with any boolean reduction.
// CHECK-LABEL: func @lower_reduce_any_boolean_tpu
func.func @lower_reduce_any_boolean_tpu() -> (tensor<2048x8192xi1>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIMENSION:.*]] = "tf.Const"
// CHECK-NEXT: %[[INPUT_CAST:.*]] = "tf.Cast"(%[[CONST_OUT_1]])
// CHECK-NEXT: %[[REDUCE_SCATTER_OUT:.*]] = "tf.DTensorReduceScatter"(%[[INPUT_CAST]], %[[GROUP_ASSIGNMENT]], %[[SCATTER_DIMENSION]])
// CHECK-SAME: reduce_op = "Max"
// CHECK-NEXT: %[[OUTPUT_CAST:.*]] = "tf.Cast"(%[[REDUCE_SCATTER_OUT]])
// CHECK-NEXT return %[[OUTPUT_CAST]]
%0 = "tf.Const"() {value = dense<1> : tensor<4096x8192xi1>, _layout = ["sharding_specs:unsharded,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"]} : () -> tensor<4096x8192xi1>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%0, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Any"} : (tensor<4096x8192xi1>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xi1>
func.return %3: tensor<2048x8192xi1>
}
// -----
// Check for error of AllReduce on TPU with all boolean reduction.
func.func @lower_allreduce_sum_boolean() -> (tensor<4096x8192xi1>) {
%0 = "tf.Const"() {value = dense<1> : tensor<4096x8192xi1>} : () -> tensor<4096x8192xi1>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
// expected-error @+1 {{reduce for boolean only supports 'All'/'Min' or 'Any'/'Max' reduction}}
%2= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4096x8192xi1>, tensor<4x2xi32>) -> tensor<4096x8192xi1>
func.return %2: tensor<4096x8192xi1>
}
// -----
// Check for error of DTensorReduceScatter on TPU with sum boolean reduction.
func.func @lower_reduce_sum_boolean_tpu() -> (tensor<2048x8192xi1>) {
%0 = "tf.Const"() {value = dense<1> : tensor<4096x8192xi1>, _layout = ["sharding_specs:unsharded,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"]} : () -> tensor<4096x8192xi1>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// expected-error @+1 {{reduce for boolean only supports 'All'/'Min' or 'Any'/'Max' reduction}}
%3 = "tf.DTensorReduceScatter"(%0, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4096x8192xi1>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xi1>
func.return %3: tensor<2048x8192xi1>
}
// -----
// Check the lowering of AllReduce on TPU with all boolean reduction.
// CHECK-LABEL: func @lower_allreduce_all_boolean
func.func @lower_allreduce_all_boolean() -> (tensor<4096x8192xi1>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[INPUT_CAST:.*]] = "tf.Cast"(%[[CONST_OUT_1]])
// CHECK-NEXT: %[[ALLREDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[INPUT_CAST]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: reduce_op = "Min"
// CHECK-NEXT: %[[OUTPUT_CAST:.*]] = "tf.Cast"(%[[ALLREDUCE_OUT]])
// CHECK-NEXT return %[[OUTPUT_CAST]]
%0 = "tf.Const"() {value = dense<1> : tensor<4096x8192xi1>} : () -> tensor<4096x8192xi1>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2= "tf.DTensorAllReduce"(%0, %1) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "All"} : (tensor<4096x8192xi1>, tensor<4x2xi32>) -> tensor<4096x8192xi1>
func.return %2: tensor<4096x8192xi1>
}
// -----
// Tests unsupported integer types are promoted to i64.
// CHECK-LABEL: func @lower_all_reduce_i8_gpu_mesh
func.func @lower_all_reduce_i8_gpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<4096x8192xi8> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"}) -> tensor<4096x8192xi8> {
// CHECK: "tf_device.cluster"
// CHECK: %[[PRECAST_OUT:.*]] = "tf.Cast"(%arg1)
// CHECK-SAME: -> tensor<4096x8192xi64>
// CHECK: %[[REDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[PRECAST_OUT]],
// CHECK: %[[POSTCAST_OUT:.*]] = "tf.Cast"(%[[REDUCE_OUT]])
// CHECK-SAME: -> tensor<4096x8192xi8>
// CHECK-NEXT tf_device.return %[[POSTCAST_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%3 = "tf.DTensorAllReduce"(%arg1, %1) {_layout = ["sharding_specs:x,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"], device_type = "/job:localhost/replica:0/task:0/device:GPU", reduce_op = "Add"} : (tensor<4096x8192xi8>, tensor<4x2xi32>) -> tensor<4096x8192xi8>
tf_device.return %3 : tensor<4096x8192xi8>
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> tensor<4096x8192xi8>
func.return %0 : tensor<4096x8192xi8>
}
tensorflow/dtensor/mlir/tests/dtensor_layout_must_execute.mlir
+27
View File
@@ -0,0 +1,27 @@
// 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: dtensor-opt %s -split-input-file -sccp -canonicalize | FileCheck %s
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> tensor<2x4xi32> attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// COM: DTensorLayout Op for not used argument arg1 must not be removed
// CHECK: = "tf.DTensorLayout"(%arg1)
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%cst = "tf.Const"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"], _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", device = "", value = dense<[[1, 2, 3, 4], [5, 6, 7, 8]]> : tensor<2x4xi32>} : () -> tensor<2x4xi32>
%1 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%2 = "tf.Relayout"(%1) {device = "", layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
return %3 : tensor<2x4xi32>
}
tensorflow/dtensor/mlir/tests/dtensor_layout_to_xla_sharding_op.mlir
+69
View File
@@ -0,0 +1,69 @@
// 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: dtensor-opt %s -dtensor-layout-to-xla-sharding-op | FileCheck %s
// CHECK-LABEL: @check_layouts_are_converted_to_xla_sharding_op
func.func @check_layouts_are_converted_to_xla_sharding_op(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd"}) -> tensor<8x8xi32> {
// CHECK: [[tensor1:%[0-9]+]] = "tf.Identity"(%arg0)
// CHECK: [[tensor2:%[0-9]+]] = "tf.XlaSharding"([[tensor1]])
// CHECK-SAME: _XlaSharding = ""
// CHECK: return [[tensor2]]
%1 = "tf.Identity"(%arg0) {_global_shape = [#tf_type.shape<8x8>], device = ""} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %2 : tensor<8x8xi32>
}
// -----
// CHECK-LABEL: @no_xla_sharding_op_for_block_arg
func.func @no_xla_sharding_op_for_block_arg(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd"}) -> tensor<8x8xi32> {
// CHECK-NOT: "tf.DTensorLayout"
// CHECK-NOT: "tf.XlaSharding"
// CHECK: return %arg0
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %1 : tensor<8x8xi32>
}
// -----
// CHECK-LABEL: @no_xla_sharding_op_for_const_input
func.func @no_xla_sharding_op_for_const_input() -> tensor<8x8xi32> {
// CHECK: [[tensor:%[a-z0-9]+]] = "tf.Const"
// CHECK-NOT: "tf.DTensorLayout"
// CHECK-NOT: "tf.XlaSharding"
// CHECK: return [[tensor]]
%cst = "tf.Const"() {value = dense<-3> : tensor<i32>} : () -> tensor<8x8xi32>
%1 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %1 : tensor<8x8xi32>
}
// -----
// CHECK-LABEL: @no_xla_sharding_op_for_const_foldable_input
func.func @no_xla_sharding_op_for_const_foldable_input() -> tensor<8x8xi32> {
// CHECK: [[tensor:%[a-z0-9]+]] = "tf.Const"
// CHECK-NOT: "tf.DTensorLayout"
// CHECK-NOT: "tf.XlaSharding"
// CHECK-NOT: "tf.Reshape"
// CHECK: return [[tensor]]
%cst = "tf.Const"() {value = dense<-3> : tensor<i32>} : () -> tensor<8x8xi32>
%1 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%shape = arith.constant dense<8> : tensor<2xi32>
%2 = "tf.Reshape"(%1, %shape) {_global_shape = [#tf_type.shape<8x8>], device = ""} : (tensor<8x8xi32>, tensor<2xi32>) -> tensor<8x8xi32>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %3 : tensor<8x8xi32>
}
tensorflow/dtensor/mlir/tests/dtensor_mixed_precision_reduce.mlir
+125
View File
@@ -0,0 +1,125 @@
// 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: env DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE=4 dtensor-opt -- -split-input-file -dtensor-mixed-precision-reduce -verify-diagnostics %s | FileCheck %s
// Check bfloat16 AllReduce is upcasted for a sufficient group size.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<1x4xbf16>
func.func @main(
%arg0: tensor<1x4xbf16> {tf._global_shape = #tf_type.shape<8x4>, tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=8|*TPU"})
-> (tensor<4xbf16> {tf._global_shape = #tf_type.shape<4>}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"
// CHECK-NEXT: %[[SUM_OUT:.*]] = "tf.Sum"(%[[ARG0]], %[[AXIS]])
// CHECK-SAME: -> tensor<4xbf16>
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[UPCAST:.*]] = "tf.Cast"(%[[SUM_OUT]])
// CHECK-SAME: (tensor<4xbf16>) -> tensor<4xf32>
// CHECK-NEXT: %[[REDUCTION_OUT:.*]] = "tf.DTensorAllReduce"(%[[UPCAST]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: -> tensor<4xf32>
// CHECK-NEXT: %[[DOWNCAST:.*]] = "tf.Cast"(%[[REDUCTION_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:TPU|x=8|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<4xf32>) -> tensor<4xbf16>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {_global_shape = [#tf_type.shape<>], _layout = ["sharding_specs: mesh:TPU|x=8|*TPU"], value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Sum"(%arg0, %cst) {_global_shape = [#tf_type.shape<4>], device = "", keep_dims = false} : (tensor<1x4xbf16>, tensor<i32>) -> tensor<4xbf16>
%cst_0 = "tf.Const"() {value = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>} : () -> tensor<1x8xi32>
%2 = "tf.DTensorAllReduce"(%1, %cst_0) {_layout = ["sharding_specs:unsharded, mesh:TPU|x=8|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4xbf16>, tensor<1x8xi32>) -> tensor<4xbf16>
%3 = "tf.Identity"(%2) {_global_shape = [#tf_type.shape<4>], _layout = ["sharding_specs:unsharded, mesh:TPU|x=8|*TPU"], device = ""} : (tensor<4xbf16>) -> tensor<4xbf16>
tf_device.return %3 : tensor<4xbf16>
}) {_mesh = "TPU|x=8|*TPU"} : () -> tensor<4xbf16>
func.return %0 : tensor<4xbf16>
}
// -----
// Check that bfloat16 AllReduce is not upcasted for a small group size.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<1x4xbf16>
func.func @main(
%arg0: tensor<1x4xbf16> {tf._global_shape = #tf_type.shape<2x4>, tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2|*TPU"})
-> (tensor<4xbf16> {tf._global_shape = #tf_type.shape<4>}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"
// CHECK-NEXT: %[[SUM_OUT:.*]] = "tf.Sum"(%[[ARG0]], %[[AXIS]])
// CHECK-SAME: -> tensor<4xbf16>
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NOT: "tf.Cast"
// CHECK-NEXT: %[[REDUCTION_OUT:.*]] = "tf.DTensorAllReduce"(%[[SUM_OUT]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: -> tensor<4xbf16>
// CHECK-NOT: "tf.Cast"
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {_global_shape = [#tf_type.shape<>], _layout = ["sharding_specs: mesh:TPU|x=2|*TPU"], value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Sum"(%arg0, %cst) {_global_shape = [#tf_type.shape<4>], device = "", keep_dims = false} : (tensor<1x4xbf16>, tensor<i32>) -> tensor<4xbf16>
%cst_0 = "tf.Const"() {value = dense<[[0, 1]]> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
%2 = "tf.DTensorAllReduce"(%1, %cst_0) {_layout = ["sharding_specs:unsharded, mesh:TPU|x=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4xbf16>, tensor<1x2xi32>) -> tensor<4xbf16>
%3 = "tf.Identity"(%2) {_global_shape = [#tf_type.shape<4>], _layout = ["sharding_specs:unsharded, mesh:TPU|x=2|*TPU"], device = ""} : (tensor<4xbf16>) -> tensor<4xbf16>
tf_device.return %3 : tensor<4xbf16>
}) {_mesh = "TPU|x=2|*TPU"} : () -> tensor<4xbf16>
func.return %0 : tensor<4xbf16>
}
// -----
// Check bfloat16 ReduceScatter is upcasted for a sufficient group size.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<512x1024xbf16>
func.func @main(
%arg0: tensor<512x1024xbf16> {tf._global_shape = #tf_type.shape<4096x1024>, tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=8|*TPU"})
-> (tensor<512x1024xbf16> {tf._global_shape = #tf_type.shape<4096x1024>}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIM:.*]] = "tf.Const"
// CHECK-NEXT: %[[UPCAST:.*]] = "tf.Cast"(%[[ARG0]])
// CHECK-SAME: (tensor<512x1024xbf16>) -> tensor<512x1024xf32>
// CHECK-NEXT: %[[REDUCTION_OUT:.*]] = "tf.DTensorReduceScatter"(%[[UPCAST]], %[[GROUP_ASSIGNMENT]], %[[SCATTER_DIM]])
// CHECK-SAME: -> tensor<512x1024xf32>
// CHECK-NEXT: %[[DOWNCAST:.*]] = "tf.Cast"(%[[REDUCTION_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:TPU|x=8|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<512x1024xf32>) -> tensor<512x1024xbf16>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>} : () -> tensor<1x8xi32>
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.DTensorReduceScatter"(%arg0, %cst, %cst_0) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=8|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<512x1024xbf16>, tensor<1x8xi32>, tensor<i32>) -> tensor<512x1024xbf16>
%3 = "tf.Identity"(%2) {_global_shape = [#tf_type.shape<4096x1024>], _layout = ["sharding_specs:x,unsharded, mesh:TPU|x=8|*TPU"], device = ""} : (tensor<512x1024xbf16>) -> tensor<512x1024xbf16>
tf_device.return %3 : tensor<512x1024xbf16>
}) {_mesh = "TPU|x=8|*TPU"} : () -> tensor<512x1024xbf16>
func.return %0 : tensor<512x1024xbf16>
}
// -----
// Check that bfloat16 ReduceScatter is not upcasted for a small group size.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<512x1024xbf16>
func.func @main(
%arg0: tensor<512x1024xbf16> {tf._global_shape = #tf_type.shape<1024x1024>, tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2|*TPU"})
-> (tensor<512x1024xbf16> {tf._global_shape = #tf_type.shape<1024x1024>}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIM:.*]] = "tf.Const"
// CHECK-NOT: "tf.Cast"
// CHECK-NEXT: %[[REDUCTION_OUT:.*]] = "tf.DTensorReduceScatter"(%[[ARG0]], %[[GROUP_ASSIGNMENT]], %[[SCATTER_DIM]])
// CHECK-SAME: -> tensor<512x1024xbf16>
// CHECK-NOT: "tf.Cast"
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<[[0, 1]]> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
%cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.DTensorReduceScatter"(%arg0, %cst, %cst_0) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<512x1024xbf16>, tensor<1x2xi32>, tensor<i32>) -> tensor<512x1024xbf16>
%3 = "tf.Identity"(%2) {_global_shape = [#tf_type.shape<1024x1024>], _layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2|*TPU"], device = ""} : (tensor<512x1024xbf16>) -> tensor<512x1024xbf16>
tf_device.return %3 : tensor<512x1024xbf16>
}) {_mesh = "TPU|x=2|*TPU"} : () -> tensor<512x1024xbf16>
func.return %0 : tensor<512x1024xbf16>
}
tensorflow/dtensor/mlir/tests/dtensor_mlir_opt_main.cc
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2022 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 "llvm/Support/LogicalResult.h"
#include "mlir/InitAllDialects.h" // from @llvm-project
#include "mlir/InitAllExtensions.h" // from @llvm-project
#include "mlir/InitAllPasses.h" // from @llvm-project
#include "mlir/Tools/mlir-opt/MlirOptMain.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/init_mlir.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/dtensor/mlir/create_dtensor_mlir_passes.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
int main(int argc, char **argv) {
tensorflow::InitMlir y(&argc, &argv);
mlir::registerAllPasses();
tensorflow::dtensor::registerDTensorPasses();
mlir::DialectRegistry registry;
mlir::registerAllDialects(registry);
mlir::registerAllExtensions(registry);
mlir::RegisterAllTensorFlowDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
registry.insert<mlir::dtensor::DTensorDialect>();
mlir::TF::RegisterDTensorTFOps();
return failed(
mlir::MlirOptMain(argc, argv, "DTensor pass driver\n", registry));
}
tensorflow/dtensor/mlir/tests/dtensor_reduce_scatter_lowering.mlir
+125
View File
@@ -0,0 +1,125 @@
// 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: env DTENSOR_GPU_USE_NCCL_COMMUNICATION=1 dtensor-opt %s -split-input-file -dtensor-reduce-scatter-lowering -verify-diagnostics | FileCheck %s --dump-input=fail
// Check the lowering of DTensorReduceScatter on TPU with sum reduction.
// CHECK-LABEL: func @lower_reduce_scatter_sum_tpu
func.func @lower_reduce_scatter_sum_tpu() -> (tensor<2048x8192xf32>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIMENSION:.*]] = "tf.Const"
// CHECK-NEXT: %[[REDUCE_SCATTER_OUT:.*]] = "tf.XlaReduceScatter"(%[[CONST_OUT_1]], %[[GROUP_ASSIGNMENT]], %[[SCATTER_DIMENSION]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-NEXT return %[[REDUCE_SCATTER_OUT]]
%0 = "tf.Const"() {value = dense<0.0> : tensor<4096x8192xf32>, _layout = ["sharding_specs:unsharded,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"]} : () -> tensor<4096x8192xf32>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%0, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:tpu_mesh|x=2,y=4|*TPU"], device_type = "/job:localhost/replica:0/task:0/device:TPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xf32>
func.return %3: tensor<2048x8192xf32>
}
// Check the lowering of DTensorReduceScatter on CPU with sum reduction.
// CHECK-LABEL: func @lower_reduce_scatter_sum_cpu
func.func @lower_reduce_scatter_sum_cpu() -> (tensor<2048x8192xf32>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIMENSION:.*]] = "tf.Const"
// CHECK-NEXT: %[[ALL_REDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[CONST_OUT_1]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-NEXT: %[[ALL_SCATTER_OUT:.*]] = "tf.DTensorAllScatter"(%[[ALL_REDUCE_OUT]])
// CHECK-NEXT return %[[ALL_SCATTER_OUT]]
%0 = "tf.Const"() {value = dense<0.0> : tensor<4096x8192xf32>, _layout = ["sharding_specs:unsharded,unsharded, mesh:cpu_mesh|x=2,y=4|*CPU"]} : () -> tensor<4096x8192xf32>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%0, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:cpu_mesh|x=2,y=4|*CPU"], device_type = "/job:localhost/replica:0/task:0/device:CPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xf32>
func.return %3: tensor<2048x8192xf32>
}
// CHECK-LABEL: func @lower_reduce_scatter_gpu_mesh
func.func @lower_reduce_scatter_gpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<4096x8192xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"}) -> tensor<2048x8192xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[DEVICE_ID_RESHAPE:.*]] = "tf.Reshape"(%arg0
// CHECK: %[[RELATIVE_DEVICE_ID:.*]] = "tf.Sub"(%[[DEVICE_ID_RESHAPE]]
// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
// CHECK-DAG: %[[DEVICE_ID_TO_GROUP_KEY:.*]] = "tf.Const"() <{value = dense<[0, 0, 1, 1, 2, 2, 3, 3]> : tensor<8xi32>}>
// CHECK: %[[GROUP_KEY_SLICE:.*]] = "tf.Slice"(%[[DEVICE_ID_TO_GROUP_KEY]], %[[RELATIVE_DEVICE_ID]], %[[CONST_1]]
// CHECK: %[[GROUP_KEY_RESHAPE:.*]] = "tf.Reshape"(%[[GROUP_KEY_SLICE]]
// CHECK-DAG: %[[INSTANCE_KEY:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}>
// CHECK-DAG: %[[GROUP_SIZE:.*]] = "tf.Const"() <{value = dense<2> : tensor<i32>}>
// CHECK: %[[REDUCE_SCATTER_OUT:.*]] = "tf.CollectiveReduceScatterV2"(%arg1, %[[GROUP_SIZE]], %[[GROUP_KEY_RESHAPE]], %[[INSTANCE_KEY]])
// CHECK-SAME: final_op = "Id"
// CHECK-SAME: merge_op = "Add"
// CHECK-NEXT tf_device.return %[[REDUCE_SCATTER_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%arg1, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"], device_type = "/job:localhost/replica:0/task:0/device:GPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xf32>
tf_device.return %3 : tensor<2048x8192xf32>
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> tensor<2048x8192xf32>
func.return %0 : tensor<2048x8192xf32>
}
// CHECK-LABEL: func @lower_reduce_scatter_transpose_gpu_mesh
func.func @lower_reduce_scatter_transpose_gpu_mesh(%arg0: tensor<i32>,
%arg1: tensor<4096x8192xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"}) -> tensor<2048x8192xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "tf.Transpose"(%arg1
// CHECK: %[[REDUCE_SCATTER_OUT:.*]] = "tf.CollectiveReduceScatterV2"(%[[TRANSPOSED_INPUT]]
// CHECK: %[[TRANSPOSED_OUTPUT:.*]] = "tf.Transpose"(%[[REDUCE_SCATTER_OUT]]
// CHECK-NEXT tf_device.return %[[TRANSPOSED_OUTPUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%arg1, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:gpu_mesh|x=2,y=4|*GPU"], device_type = "/job:localhost/replica:0/task:0/device:GPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xf32>
tf_device.return %3 : tensor<2048x8192xf32>
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> tensor<2048x8192xf32>
func.return %0 : tensor<2048x8192xf32>
}
// Check the lowering of DTensorReduceScatter on CPU with any_boolean reduction.
// CHECK-LABEL: func @lower_reduce_scatter_any_boolean_cpu
func.func @lower_reduce_scatter_any_boolean_cpu() -> (tensor<2048x8192xi1>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIMENSION:.*]] = "tf.Const"
// CHECK-NEXT: %[[ALL_REDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[CONST_OUT_1]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: reduce_op = "Any"
// CHECK-NEXT: %[[ALL_SCATTER_OUT:.*]] = "tf.DTensorAllScatter"(%[[ALL_REDUCE_OUT]])
// CHECK-NEXT return %[[ALL_SCATTER_OUT]]
%0 = "tf.Const"() {value = dense<1> : tensor<4096x8192xi1>, _layout = ["sharding_specs:unsharded,unsharded, mesh:cpu_mesh|x=2,y=4|*CPU"]} : () -> tensor<4096x8192xi1>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%0, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:cpu_mesh|x=2,y=4|*CPU"], device_type = "/job:localhost/replica:0/task:0/device:CPU", reduce_op = "Any"} : (tensor<4096x8192xi1>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xi1>
func.return %3: tensor<2048x8192xi1>
}
// Check the lowering of DTensorReduceScatter without input layout.
// CHECK-LABEL: func @lower_reduce_scatter_no_input_layout
func.func @lower_reduce_scatter_no_input_layout() -> (tensor<2048x8192xf32>) {
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"
// CHECK-NEXT: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[SCATTER_DIMENSION:.*]] = "tf.Const"
// CHECK-NEXT: %[[ALL_REDUCE_OUT:.*]] = "tf.DTensorAllReduce"(%[[CONST_OUT_1]], %[[GROUP_ASSIGNMENT]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-NEXT: %[[ALL_SCATTER_OUT:.*]] = "tf.DTensorAllScatter"(%[[ALL_REDUCE_OUT]])
// CHECK-NEXT return %[[ALL_SCATTER_OUT]]
%0 = "tf.Const"() {value = dense<0.0> : tensor<4096x8192xf32>} : () -> tensor<4096x8192xf32>
%1 = "tf.Const"() {value = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorReduceScatter"(%0, %1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:cpu_mesh|x=2,y=4|*CPU"], device_type = "/job:localhost/replica:0/task:0/device:CPU", reduce_op = "Add"} : (tensor<4096x8192xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<2048x8192xf32>
func.return %3: tensor<2048x8192xf32>
}
tensorflow/dtensor/mlir/tests/dtensor_remove_dtensorlayout.mlir
+38
View File
@@ -0,0 +1,38 @@
// 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: dtensor-opt %s -split-input-file -dtensor-remove-dtensorlayout | FileCheck %s
// This test checks DTensorLayout ops are all removed, regardless of whether it
// has the `use_xla_spmd` attribute.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
%0 = "tf_device.cluster"() ({
// CHECK: "tf.Const"()
// CHECK-NOT: "tf.DTensorLayout"
// CHECK: "tf.Const"()
// CHECK-NOT: "tf.DTensorLayout"
// CHECK: "tf.Add"
// CHECK-NOT: "tf.DTensorLayout"
// CHECK-NEXT: tf_device.return
%1 = "tf.Const"() {value = dense<[[4, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3|use_xla_spmd>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
%3 = "tf.Const"() {value = dense<[[1, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
%5 = "tf.Add"(%2, %4): (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
%6 = "tf.DTensorLayout"(%5) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
tf_device.return %6 : tensor<2x2xi32>
}) {_mesh = "mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<2x2xi32>)
func.return
}
tensorflow/dtensor/mlir/tests/dtensor_replace_auxiliary_layout_op.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: dtensor-opt %s -split-input-file -dtensor-replace-auxiliary-layout-op | FileCheck %s
// Check auxiliary `tf.DTensorLayout` is replaced by `tf.Identity`.
// CHECK-LABEL: func @check_replace_auxiliary_layout_op
func.func @check_replace_auxiliary_layout_op(%arg0: tensor<8x8xi32>) -> tensor<8x8xi32> {
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: return
%0 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %1 : tensor<8x8xi32>
}
tensorflow/dtensor/mlir/tests/dtensor_replace_relayout_with_identity.mlir
+32
View File
@@ -0,0 +1,32 @@
// 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: dtensor-opt %s -split-input-file -dtensor-replace-relayout-with-identity | FileCheck %s
module {
// CHECK-LABEL: main
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<2x4xi32> {tf._global_shape = #tf_type.shape<2x4>, tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<2x4xi32> {tf._global_shape = #tf_type.shape<2x4>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<2x4>], global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
// CHECK: "tf.Identity"
// CHECK-NOT: "tf.Relayout"
// CHECK-SAME: _global_shape = [#tf_type.shape<2x4>], device = "", layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"
%2 = "tf.Relayout"(%1) {_global_shape = [#tf_type.shape<2x4>], device = "", layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%3 = "tf.DTensorLayout"(%2) {_global_shape = [#tf_type.shape<2x4>], global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
tf_device.return %3 : tensor<2x4xi32>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"} : () -> tensor<2x4xi32>
return %0 : tensor<2x4xi32>
}
}
tensorflow/dtensor/mlir/tests/dtensor_set_hlo_sharding.mlir
+40
View File
@@ -0,0 +1,40 @@
// 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: dtensor-opt %s -split-input-file -dtensor-set-hlo-sharding='check_layout_use_xla_spmd=true' --verify-diagnostics | FileCheck %s
// Check all inputs, outputs, and operations have sharding attributes, with `check_layout_use_xla_spmd` set to true.
// CHECK-LABEL: func @check_layouts_are_converted_to_xla_sharding_attributes
// CHECK-SAME: (%arg0: tensor<8x8xi32> {mhlo.sharding = "", tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd"}) -> (tensor<8x8xi32> {mhlo.sharding = ""}) {
func.func @check_layouts_are_converted_to_xla_sharding_attributes(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd"}) -> tensor<8x8xi32> {
// CHECK: "tf.DTensorLayout"
// CHECK: "tf.Identity"
// CHECK: "tf.DTensorLayout"
// CHECK-NEXT: return
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%2 = "tf.Identity"(%1) {_global_shape = [#tf_type.shape<8x8>], device = ""} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %3 : tensor<8x8xi32>
}
// -----
// Check that Layout ops not on XLA SPMD mesh is not allowed at this point.
func.func @check_layouts_not_xla_spmd_is_not_allowed(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<8x8xi32>) {
// expected-error @+1 {{'tf.DTensorLayout' op Found a layout operation that is not on XLA SPMD mesh during XLA SPMD integration.}}
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %1 : tensor<8x8xi32>
}
tensorflow/dtensor/mlir/tests/dtensor_set_hlo_sharding_default.mlir
+30
View File
@@ -0,0 +1,30 @@
// 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: dtensor-opt %s -split-input-file -dtensor-set-hlo-sharding | FileCheck %s
// Check all inputs and operations have sharding attributes, with `check_layout_use_xla_spmd` set to default value (false).
// CHECK-LABEL: func @check_layouts_are_converted_to_xla_sharding_attributes
// CHECK-SAME: (%arg0: tensor<8x8xi32> {mhlo.sharding = "", tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<8x8xi32> {mhlo.sharding = ""}) {
func.func @check_layouts_are_converted_to_xla_sharding_attributes(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"}) -> tensor<8x8xi32> {
// CHECK: "tf.DTensorLayout"
// CHECK: "tf.Identity"
// CHECK: "tf.DTensorLayout"
// CHECK-NEXT: return
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%2 = "tf.Identity"(%1) {_global_shape = [#tf_type.shape<8x8>], device = ""} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %3 : tensor<8x8xi32>
}
tensorflow/dtensor/mlir/tests/dtensor_xla_spmd_integration.mlir
+41
View File
@@ -0,0 +1,41 @@
// 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: dtensor-opt %s -split-input-file -dtensor-set-hlo-sharding='check_layout_use_xla_spmd=true' -dtensor-replace-auxiliary-layout-op -dtensor-layout-to-xla-sharding-op -verify-diagnostics | FileCheck %s
// Check after XLA SPMD integration, all inputs, outputs, and operations have sharding attributes.
// And all "tf.DTensorLayout" are removed.
// CHECK-LABEL: func @check_layouts_are_converted_to_xla_sharding_attributes
// CHECK-SAME: (%arg0: tensor<8x8xi32> {mhlo.sharding = "", tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd"}) -> (tensor<8x8xi32> {mhlo.sharding = ""})
func.func @check_layouts_are_converted_to_xla_sharding_attributes(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd"}) -> (tensor<8x8xi32>) {
// CHECK: [[tensor2:%[0-9]+]] = "tf.Identity"(%arg0)
// CHECK: [[tensor3:%[0-9]+]] = "tf.XlaSharding"([[tensor2]])
// CHECK-SAME: _XlaSharding = ""
// CHECK: return [[tensor3]]
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%2 = "tf.Identity"(%1) {_global_shape = [#tf_type.shape<8x8>], device = ""} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %3 : tensor<8x8xi32>
}
// -----
// Check that Layout ops not on XLA SPMD mesh is not allowed at this point.
func.func @check_layouts_not_xla_spmd_is_not_allowed(
%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<8x8xi32>) {
// expected-error @+1 {{'tf.DTensorLayout' op Found a layout operation that is not on XLA SPMD mesh during XLA SPMD integration.}}
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
return %1 : tensor<8x8xi32>
}
tensorflow/dtensor/mlir/tests/elide_identity_before_copy_to_mesh.mlir
+31
View File
@@ -0,0 +1,31 @@
// 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: dtensor-opt %s -split-input-file -dtensor-elide-identity-before-copy-to-mesh | FileCheck %s
// Check that identity before CopyToMeshGrad is elided.
// CHECK-LABEL: func @check_elide_identity
func.func @check_elide_identity() -> (tensor<4xi32>) {
// CHECK: %[[CONST:.*]] = "tf.Const"()
// CHECK-NEXT: %[[CONST_1:.*]] = "tf.Const"()
// CHECK-NEXT: "tf.CopyToMeshGrad"(%[[CONST]], %[[CONST_1]])
%cst = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32>} : () -> tensor<4xi32>
%cst_1 = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32>} : () -> tensor<4xi32>
%1 = "tf.Identity"(%cst) : (tensor<4xi32>) -> tensor<4xi32>
%2 = "tf.CopyToMeshGrad"(%1, %cst_1) {reference_layout=""}: (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %2 : tensor<4xi32>
}
tensorflow/dtensor/mlir/tests/function_renaming.mlir
+31
View File
@@ -0,0 +1,31 @@
// 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: dtensor-opt %s -split-input-file -dtensor-function-renaming -verify-diagnostics | FileCheck %s
module attributes {dtensor.cache_key = "_abc_def"} {
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<f32>, %arg1: tensor<4xf32>) -> (tensor<4xf32>) {
// CHECK: "tf.StatefulPartitionedCall"
// CHECK-SAME: f = @_func_0_abc_def
%0 = "tf.StatefulPartitionedCall"(%arg0, %arg1) {f = @_func_0, config = "", config_proto = "", executor_type = ""} : (tensor<f32>, tensor<4xf32>) -> (tensor<4xf32>)
func.return %0 : tensor<4xf32>
}
// CHECK-LABEL: func private @_func_0_abc_def
func.func private @_func_0(%arg0: tensor<f32>, %arg1: tensor<4xf32>) -> (tensor<4xf32>) {
%0 = "tf.AddV2"(%arg0, %arg1) : (tensor<f32>, tensor<4xf32>) -> (tensor<4xf32>)
func.return %0 : tensor<4xf32>
}
}
tensorflow/dtensor/mlir/tests/handle_cross_cluster_dependencies.mlir
+265
View File
@@ -0,0 +1,265 @@
// 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: dtensor-opt %s -dtensor-handle_cross_cluster_dependences -split-input-file -verify-diagnostics | FileCheck %s
// Check that CopyToMesh op must be used to send tensors across mesh clusters.
func.func @main() {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Neg"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
%2 = "tf_device.cluster"() ({
// expected-error @+1 {{CopyToMeshOp must be used to send data across mesh}}
%3 = "tf.Neg"(%0) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check that CopyToMesh inside incorrect mesh cluster is disallowed.
func.func @main() -> tensor<i32> {
// expected-error @+1 {{ Failed to extract mesh }}
%0:2 = "tf_device.cluster"() ({
%2 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Neg"(%2) : (tensor<i32>) -> tensor<i32>
tf_device.return %2, %3 : tensor<i32>, tensor<i32>
}) {_mesh="CPU|x=2,y=2|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>)
%2 = "tf_device.cluster"() ({
%3 = "tf.Relayout"(%0#0) { layout="sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : (tensor<i32>) -> (tensor<i32>)
%4 = "tf.Neg"(%3) : (tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return %0#1 : tensor<i32>
}
// -----
// Check that Const op is cloned across mesh to reduce data transfer.
// CHECK-LABEL: func @main
func.func @main() -> tensor<i32> {
// CHECK: %[[CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"()
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[CONST_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> tensor<i32>
%0:2 = "tf_device.cluster"() ({
%2 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Neg"(%2) : (tensor<i32>) -> tensor<i32>
tf_device.return %2, %3 : tensor<i32>, tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"()
// CHECK-NEXT: %[[LAYOUT_OUT:.*]] = "tf.Relayout"(%[[CONST_OUT]])
// CHECK-SAME: layout = "sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[LAYOUT_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%2 = "tf_device.cluster"() ({
%3 = "tf.Relayout"(%0#0) { layout ="sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : (tensor<i32>) -> (tensor<i32>)
%4 = "tf.Neg"(%3) : (tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return %0#1 : tensor<i32>
}
// -----
// Check that CopyToMesh op is lowered to DTensorSend/DTensorRecv op.
// CHECK-LABEL: func @main
func.func @main() -> tensor<i32> {
// CHECK: %[[CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"()
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[A_OUT]]
// CHECK-NEXT: "tf.DTensorSend"(%[[A_OUT]]
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-SAME: target_mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> tensor<i32>
%0:2 = "tf_device.cluster"() ({
%2 = "tf.A"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Neg"(%2) : (tensor<i32>) -> tensor<i32>
tf_device.return %2, %3 : tensor<i32>, tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf.DTensorRecv"()
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-SAME: mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: %[[RECVRELAYOUT_OUT:.*]] = "tf.Relayout"(%[[RECV_OUT]]
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[RECVRELAYOUT_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%2 = "tf_device.cluster"() ({
%3 = "tf.Relayout"(%0#0) { layout ="sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : (tensor<i32>) -> (tensor<i32>)
%4 = "tf.Neg"(%3) : (tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return %0#1 : tensor<i32>
}
// -----
// Check that tensor transfer from 2 clusters from same mesh without CopyToMesh is allowed.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"()
// CHECK-NEXT: %[[NEG_OUT_0:.*]] = "tf.Neg"(%[[CONST_OUT]]
// CHECK-NEXT: tf_device.return %[[NEG_OUT_0]]
// CHECK-NEXT: () -> tensor<i32>
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Neg"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"()
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
"tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[NEG_OUT_2:.*]] = "tf.Neg"(%[[CLUSTER_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%5 = "tf_device.cluster"() ({
%4 = "tf.Neg"(%0) : (tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check that CopyToMesh op with multiple usages is lowered to
// DTensorSend/DTensorRecv ops for each usages.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: %[[CPU_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"()
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[A_OUT]]
// CHECK-NEXT: "tf.DTensorSend"(%[[NEG_OUT]]
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-SAME: target_mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: "tf.DTensorSend"(%[[NEG_OUT]]
// CHECK-SAME: key = "communication_key_GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3_1"
// CHECK-SAME: target_mesh = #dtensor.mesh<GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> tensor<i32>
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Neg"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[RECV_OUT_1:.*]] = "tf.DTensorRecv"()
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-SAME: mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: %[[RELAYOUT_OUT_1:.*]] = "tf.Relayout"(%[[RECV_OUT_1]]
// CHECK-NEXT: %[[NEG_OUT_1:.*]] = "tf.Neg"(%[[RELAYOUT_OUT_1]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%2 = "tf_device.cluster"() ({
%3 = "tf.Relayout"(%0) { layout ="sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : (tensor<i32>) -> (tensor<i32>)
%4 = "tf.Neg"(%3) : (tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[RECV_OUT_2:.*]] = "tf.DTensorRecv"()
// CHECK-SAME: key = "communication_key_GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3_1"
// CHECK-SAME: mesh = #dtensor.mesh<GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>
// CHECK-NEXT: %[[RELAYOUT_OUT_2:.*]] = "tf.Relayout"(%[[RECV_OUT_2]]
// CHECK-NEXT: %[[NEG_OUT_2:.*]] = "tf.Neg"(%[[RELAYOUT_OUT_2]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%3 = "tf_device.cluster"() ({
%4 = "tf.Relayout"(%0) { layout ="sharding_specs:unsharded, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : (tensor<i32>) -> (tensor<i32>)
%5 = "tf.Neg"(%4) : (tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh="GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[NEG_OUT_3:.*]] = "tf.Neg"(%[[CPU_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%4 = "tf_device.cluster"() ({
%7 = "tf.Neg"(%0) : (tensor<i32>) -> tensor<i32>
tf_device.return %7 : tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check that CopyToMesh inside cluster is replaced with Relayout.
// CHECK-LABEL: func @main
func.func @main() -> tensor<i32> {
// CHECK: %[[CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"()
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[CONST_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> tensor<i32>
%0:2 = "tf_device.cluster"() ({
%2 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Neg"(%2) : (tensor<i32>) -> tensor<i32>
tf_device.return %2, %3 : tensor<i32>, tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"()
// CHECK-NEXT: %[[LAYOUT_OUT:.*]] = "tf.Relayout"(%[[CONST_OUT]])
// CHECK-SAME: layout = "sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-NEXT: %[[RELAYOUT_OUT:.*]] = "tf.Relayout"(%[[LAYOUT_OUT]])
// CHECK-SAME: layout = "sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"(%[[RELAYOUT_OUT]]
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: () -> ()
%2 = "tf_device.cluster"() ({
%3 = "tf.Relayout"(%0#0) { layout ="sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3" } : (tensor<i32>) -> (tensor<i32>)
%4 = "tf.Relayout"(%3) { layout ="sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3" } : (tensor<i32>) -> (tensor<i32>)
%5 = "tf.Neg"(%4) : (tensor<i32>) -> tensor<i32>
tf_device.return %5 : tensor<i32>
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return %0#1 : tensor<i32>
}
// -----
// Check that unused cluster results are removed.
// CHECK-LABEL: func @main
func.func @main() -> tensor<i32> {
%0:3 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
tf_device.return %1, %2, %3 : tensor<i32>, tensor<i32>, tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>, tensor<i32>, tensor<i32>)
func.return %0#2 : tensor<i32>
}
tensorflow/dtensor/mlir/tests/handle_sparsetensors.mlir
+95
View File
@@ -0,0 +1,95 @@
// 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: dtensor-opt %s -dtensor-sparse-tensor-to-dense-tensor -split-input-file -verify-diagnostics | FileCheck %s
// Check int32 SparseTensors expand to SparseToDenseOp.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<6x4xi32> {tf._layout = "sharding_specs:batch,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", tf._sparse = true}) -> tensor<6x4xi32> attributes {tf.entry_function = {inputs = "device_id,op_input_0"}}{
// CHECK: func @main(%arg0: tensor<i32>, %arg1: tensor<?x2xi64>, %arg2: tensor<2xi64>, %arg3: tensor<?xi32>) -> tensor<6x4xi32> attributes {tf.entry_function = {inputs = "device_id,op_input_sparse_indices_1,op_input_sparse_dense_shapes_1,op_input_sparse_values_1"}} {
// CHECK-NEXT: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DENSE:.*]] = "tf.SparseToDense"(%arg1, %arg2, %arg3, %[[CST]])
// CHECK-NEXT: %[[DENSE_OUT:.*]] = "tf.DTensorLayout"(%[[DENSE]])
// CHECK-NEXT: "tf.AddV2"(%[[DENSE_OUT:.*]], %[[DENSE_OUT:.*]])
// CHECK-NEXT: "tf.Identity"
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:batch,unsharded, mesh:|batch=2,x=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<6x4xi32>) -> tensor<6x4xi32>
%1 = "tf.AddV2"(%0, %0) {} : (tensor<6x4xi32>, tensor<6x4xi32>) -> tensor<6x4xi32>
%2 = "tf.Identity"(%1) {} : (tensor<6x4xi32>) -> tensor<6x4xi32>
func.return %2 : tensor<6x4xi32>
}
// -----
// Check float32 SparseTensors expand to SparseToDenseOp.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<6x4xf32> {tf._layout = "sharding_specs:batch,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", tf._sparse = true}) -> tensor<6x4xf32> attributes {tf.entry_function = {inputs = "device_id,op_input_0"}}{
// CHECK: func @main(%arg0: tensor<i32>, %arg1: tensor<?x2xi64>, %arg2: tensor<2xi64>, %arg3: tensor<?xf32>) -> tensor<6x4xf32> attributes {tf.entry_function = {inputs = "device_id,op_input_sparse_indices_1,op_input_sparse_dense_shapes_1,op_input_sparse_values_1"}} {
// CHECK-NEXT: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DENSE:.*]] = "tf.SparseToDense"(%arg1, %arg2, %arg3, %[[CST]])
// CHECK-NEXT: %[[DENSE_OUT:.*]] = "tf.DTensorLayout"(%[[DENSE]])
// CHECK-NEXT: "tf.AddV2"(%[[DENSE_OUT:.*]], %[[DENSE_OUT:.*]])
// CHECK-NEXT: "tf.Identity"
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:batch,unsharded, mesh:|batch=2,x=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%1 = "tf.AddV2"(%0, %0) {} : (tensor<6x4xf32>, tensor<6x4xf32>) -> tensor<6x4xf32>
%2 = "tf.Identity"(%1) {} : (tensor<6x4xf32>) -> tensor<6x4xf32>
func.return %2 : tensor<6x4xf32>
}
// -----
// Check int64 SparseTensors expand to SparseToDenseOp.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<6x4xi64> {tf._layout = "sharding_specs:batch,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", tf._sparse = true}) -> tensor<6x4xi64> attributes {tf.entry_function = {inputs = "device_id,op_input_0"}}{
// CHECK: func @main(%arg0: tensor<i32>, %arg1: tensor<?x2xi64>, %arg2: tensor<2xi64>, %arg3: tensor<?xi64>) -> tensor<6x4xi64> attributes {tf.entry_function = {inputs = "device_id,op_input_sparse_indices_1,op_input_sparse_dense_shapes_1,op_input_sparse_values_1"}} {
// CHECK-NEXT: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DENSE:.*]] = "tf.SparseToDense"(%arg1, %arg2, %arg3, %[[CST]])
// CHECK-NEXT: %[[DENSE_OUT:.*]] = "tf.DTensorLayout"(%[[DENSE]])
// CHECK-NEXT: "tf.AddV2"(%[[DENSE_OUT:.*]], %[[DENSE_OUT:.*]])
// CHECK-NEXT: "tf.Identity"
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:batch,unsharded, mesh:|batch=2,x=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<6x4xi64>) -> tensor<6x4xi64>
%1 = "tf.AddV2"(%0, %0) {device = ""} : (tensor<6x4xi64>, tensor<6x4xi64>) -> tensor<6x4xi64>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<6x4xi64>) -> tensor<6x4xi64>
func.return %2 : tensor<6x4xi64>
}
// -----
// Check the SparseTensor components are appended to the end of the block argument list.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<6x4xi64> {tf._layout = "sharding_specs:batch,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", tf._sparse = true}, %arg3: tensor<i32>) -> tensor<6x4xi64> attributes {tf.entry_function = {inputs = "device_id,op_input_0"}}{
// CHECK: func @main(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<?x2xi64>, %arg3: tensor<2xi64>, %arg4: tensor<?xi64>) -> tensor<6x4xi64> attributes {tf.entry_function = {inputs = "device_id,op_input_sparse_indices_1,op_input_sparse_dense_shapes_1,op_input_sparse_values_1"}} {
// CHECK-NEXT: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DENSE:.*]] = "tf.SparseToDense"(%arg2, %arg3, %arg4, %[[CST]])
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: %[[DENSE_OUT:.*]] = "tf.DTensorLayout"(%[[DENSE]])
// CHECK-NEXT: "tf.AddV2"(%[[DENSE_OUT:.*]], %[[DENSE_OUT:.*]])
%3 = "tf.Identity"(%arg3) {} : (tensor<i32>) -> tensor<i32>
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:batch,unsharded, mesh:|batch=2,x=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<6x4xi64>) -> tensor<6x4xi64>
%1 = "tf.AddV2"(%0, %0) {device = ""} : (tensor<6x4xi64>, tensor<6x4xi64>) -> tensor<6x4xi64>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<6x4xi64>) -> tensor<6x4xi64>
func.return %2 : tensor<6x4xi64>
}
// -----
// Check that a single SparseToDenseOp is created for all usages of a single SparseTensor
func.func @main(%arg0: tensor<i32>, %arg1: tensor<6x4xi64> {tf._layout = "sharding_specs:batch,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", tf._sparse = true}) -> tensor<6x4xi64> attributes {tf.entry_function = {inputs = "device_id,op_input_0"}}{
// CHECK: func @main(%arg0: tensor<i32>, %arg1: tensor<?x2xi64>, %arg2: tensor<2xi64>, %arg3: tensor<?xi64>) -> tensor<6x4xi64> attributes {tf.entry_function = {inputs = "device_id,op_input_sparse_indices_1,op_input_sparse_dense_shapes_1,op_input_sparse_values_1"}} {
// CHECK-NEXT: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DENSE:.*]] = "tf.SparseToDense"(%arg1, %arg2, %arg3, %[[CST]])
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: "tf.AddV2"
%1 = "tf.Identity"(%arg1) {device = ""} : (tensor<6x4xi64>) -> tensor<6x4xi64>
%2 = "tf.Identity"(%arg1) {device = ""} : (tensor<6x4xi64>) -> tensor<6x4xi64>
%3 = "tf.AddV2"(%1, %2) {device = ""} : (tensor<6x4xi64>, tensor<6x4xi64>) -> tensor<6x4xi64>
func.return %3 : tensor<6x4xi64>
}
tensorflow/dtensor/mlir/tests/layout_propagation.mlir
+549
View File
@@ -0,0 +1,549 @@
// 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: dtensor-opt %s -dtensor-annotate-global-shape -dtensor-layout-propagation -split-input-file -verify-diagnostics | FileCheck %s
// Check Unary op layout propagation.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK: %1 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<i32>
// CHECK: %2 = "tf.Neg"(%1)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<i32>) -> tensor<i32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Neg"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check elementwise op with operands having incompatible layouts is not
// allowed.
func.func @main() {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Const"() {_layout = ["sharding_specs:x,z, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
// expected-error @+1 {{mesh dimension not contained in mesh}}
%3 = "tf.Add"(%1, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) : () -> (tensor<i32>)
func.return
}
// -----
// Check elementwise op layout propagation with first operand missing layout.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK: %1 = "tf.Const"()
// CHECK-SAME: () -> tensor<2x2xi32>
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<2x2xi32>
// CHECK: %3 = "tf.Add"(%1, %2)
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<10> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%2 = "tf.Const"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%3 = "tf.Add"(%1, %2) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %3 : tensor<2x2xi32>
}) {_mesh = "CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<2x2xi32>)
func.return
}
// -----
// Check elementwise op layout propagation
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK: %1 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<i32>
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<i32>
// CHECK: %3 = "tf.Add"(%1, %2)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Add"(%1, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {_mesh = "mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check layout propagation of elementwise op with broadcast propagation.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK: %1 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<10x10xi32>
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs: mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<i32>
// CHECK: %3 = "tf.Add"(%1, %2)
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<10x10xi32>, tensor<i32>) -> tensor<10x10xi32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<10x10xi32>} : () -> tensor<10x10xi32>
%2 = "tf.Const"() {_layout = ["sharding_specs: mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Add"(%1, %2) : (tensor<10x10xi32>, tensor<i32>) -> tensor<10x10xi32>
tf_device.return %3 : tensor<10x10xi32>
}) {_mesh = "CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<10x10xi32>)
func.return
}
// -----
// Check layout propagation of elementwise op with multiple device cluster.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: () -> tensor<i32>
// CHECK: %3 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: () -> tensor<i32>
// CHECK: %4 = "tf.Add"(%2, %3)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%2 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%4 = "tf.Add"(%2, %3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
//
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: () -> tensor<i32>
//
// CHECK: %3 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: () -> tensor<i32>
//
// CHECK: %4 = "tf.Add"(%2, %3)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i32>
//
// CHECK: tf_device.return
%1 = "tf_device.cluster"() ({
%2 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value = dense<10> : tensor<i32>} : () -> tensor<i32>
%4 = "tf.Add"(%2, %3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %4 : tensor<i32>
}) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check layout propagation of elementwise op with multiple inputs.
// CHECK-LABEL: func @main
// CHECK: "tf_device.cluster"() ({
// CHECK-NEXT: %1 = "tf.Add"(%arg1, %arg2)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"]
// CHECK-SAME: (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: tf_device.return
func.func @main(%arg0: tensor<i64>,
%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"},
%arg2: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"}) -> tensor<1xf32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.Add"(%arg1, %arg2) {} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
tf_device.return %1 : tensor<1xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check layout propagation of pack op.
// CHECK-LABEL: func @main
func.func @main(%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3" }) -> tensor<1xf32> {
// CHECK: "tf.Pack"
// CHECK-SAME: _layout = ["sharding_specs:unsharded,x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -> tensor<4x10x10xf32>
// CHECK-NEXT: tf_device.return
// CHECK-SAME: tensor<4x10x10xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<10x10xf32>
%2 = "tf.A"() {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<10x10xf32>
%3 = "tf.A"() {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<10x10xf32>
%4 = "tf.A"() {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<10x10xf32>
%5 = "tf.Pack"(%1, %2, %3, %4) {} : (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -> tensor<4x10x10xf32>
tf_device.return %5 : tensor<4x10x10xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check layout propagation logic of tf.Pack op with a single operand.
// CHECK-LABEL: func @main
func.func @main(%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3" }) -> tensor<1xf32> {
// CHECK: "tf.Pack"(%1)
// CHECK-SAME: _layout = ["sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: axis = 0
// CHECK-SAME: (tensor<10x10xf32>) -> tensor<10x10xf32>
// CHECK-NEXT: tf_device.return
// CHECK-SAME: tensor<10x10xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<10x10xf32>
%2 = "tf.Pack"(%1) {axis = 0 : i64} : (tensor<10x10xf32>) -> tensor<10x10xf32>
tf_device.return %2 : tensor<10x10xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check layout propagation of pack op with non-matching layouts.
func.func @main(%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*TPU" }) -> tensor<1xf32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() {_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10xf32>
%2 = "tf.A"() {_layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10xf32>
%3 = "tf.A"() {_layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10xf32>
%4 = "tf.A"() {_layout = ["sharding_specs:y,x, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10xf32>
// expected-error @+1 {{'tf.Pack' op All arguments to pack must have the same layout.}}
%5 = "tf.Pack"(%1, %2, %3, %4) {} : (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -> tensor<4x10x10xf32>
tf_device.return %5 : tensor<4x10x10xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check layout propagation of reshape op with replicated inputs.
// CHECK-LABEL: func @main
func.func @main(%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*TPU" }) -> tensor<1xf32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const" () {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*TPU"], value = dense<10> : tensor<1xi32>}: () -> tensor<1xi32>
%2 = "tf.Const" () {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*TPU"], value = dense<-1> : tensor<1xi32>}: () -> tensor<1xi32>
%3 = "tf.Pack" (%1, %2) : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
%4 = "tf.A"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10xf32>
// CHECK: "tf.Reshape"(%4, %3)
// CHECK-SAME: (tensor<10x10xf32>, tensor<2xi32>) -> tensor<10x10xf32>
%5 = "tf.Reshape"(%4, %3) : (tensor<10x10xf32>, tensor<2xi32>) -> tensor<10x10xf32>
tf_device.return %5 : tensor<10x10xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check layout propagation of reshape op with replicated inputs with different rank.
// CHECK-LABEL: func @main
func.func @main(%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*TPU" }) -> tensor<1xf32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const" () {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*TPU"], value = dense<-1> : tensor<1xi32>}: () -> tensor<1xi32>
%2 = "tf.A"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10xf32>
// CHECK: "tf.Reshape"(%2, %1)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<10x10xf32>, tensor<1xi32>) -> tensor<100xf32>
%3 = "tf.Reshape"(%2, %1) : (tensor<10x10xf32>, tensor<1xi32>) -> tensor<100xf32>
tf_device.return %3 : tensor<100xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check reshape op with batch sharded inputs.
// CHECK-LABEL: func @main
func.func @main(%arg1: tensor<1xf32> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*TPU" }) -> tensor<1xf32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const" () {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*TPU"], value = dense<-1> : tensor<1xi32>}: () -> tensor<1xi32>
%2 = "tf.Const" () {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*TPU"], value = dense<100> : tensor<1xi32>}: () -> tensor<1xi32>
%3 = "tf.Pack" (%1, %2) : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
%4 = "tf.A"() {_layout = ["sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU"]} : () -> tensor<10x10x10xf32>
// CHECK: "tf.Reshape"
// CHECK-SAME: (tensor<10x10x10xf32>, tensor<2xi32>) -> tensor<10x100xf32>
%5 = "tf.Reshape"(%4, %3) : (tensor<10x10x10xf32>, tensor<2xi32>) -> tensor<10x100xf32>
tf_device.return %5 : tensor<10x100xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check that layout propagation of inputs that are sharded in non-batch dimension is disallowed.
func.func @main(%arg0: tensor<32x32xf32> { tf._layout = "sharding_specs:x,y, mesh:|x=2,y=2|*TPU"}, %arg1: tensor<32x32xf32> { tf._layout = "sharding_specs:x,y, mesh:|x=2,y=2|*TPU"}) {
"tf_device.cluster"() ({
// expected-error @+1 {{Features input to Softmax loss ops must be sharded only across batch dimension}}
"tf.SoftmaxCrossEntropyWithLogits"(%arg0, %arg1) : (tensor<32x32xf32>, tensor<32x32xf32>) -> (tensor<32x1xf32>, tensor<32x32xf32>)
tf_device.return
}) {layout = "sharding_specs:x,y, mesh:|x=2,y=2|*TPU"} : () -> ()
func.return
}
// -----
// Check layout propagation of read variable op.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<!tf_type.resource<tensor<1xf32>>> { tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU" }) -> tensor<1xf32> {
// CHECK: "tf_device.cluster"
//
// CHECK: %1 = "tf.ReadVariableOp"(%arg0)
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<!tf_type.resource<tensor<1xf32>>>) -> tensor<1xf32>
//
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.ReadVariableOp"(%arg0) : (tensor<!tf_type.resource<tensor<1xf32>>>) -> tensor<1xf32>
tf_device.return %1 : tensor<1xf32>
}) {} : () -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// Check layout propagation of const ops from it's consumers.
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
//
// CHECK: %1 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<22xi32>
//
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<22xi32>
//
// CHECK: %3 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<22xi32>
//
// CHECK: %4 = "tf.Add"(%1, %2)
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<22xi32>, tensor<22xi32>) -> tensor<22xi32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22]> : tensor<22xi32>} : () -> tensor<22xi32>
%2 = "tf.Const"() {value = dense<[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22]> : tensor<22xi32>} : () -> tensor<22xi32>
%3 = "tf.Const"() {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22]> : tensor<22xi32>} : () -> tensor<22xi32>
%4 = "tf.Add"(%1,%2) {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : (tensor<22xi32>,tensor<22xi32>) -> tensor<22xi32>
%5 = "tf.Add"(%2,%3) {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : (tensor<22xi32>,tensor<22xi32>) -> tensor<22xi32>
tf_device.return %3 : tensor<22xi32>
}) {_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<22xi32>)
func.return
}
// -----
// CHECK-LABEL: func @main
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK: %1 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<2xi32>
// CHECK: %2 = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: () -> tensor<2xi32>
// CHECK: %3 = "tf.Reshape"(%1, %2)
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%2 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tf.Reshape"(%1, %2) {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : (tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32>
tf_device.return %3 : tensor<1x2xi32>
}) {_layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> tensor<1x2xi32>
func.return
}
// -----
// Check layout propagation of fill op.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2xi32>, %arg2: tensor<f32>) -> (tensor<?x?xf32>{
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Fill"
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], device = ""}
// CHECK-SAME: tensor<2xi32>, tensor<f32>) -> tensor<?x?xf32>
tf._default_layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.Fill"(%arg1, %arg2) {device = ""} : (tensor<2xi32>, tensor<f32>) -> tensor<?x?xf32>
tf_device.return %1 : tensor<?x?xf32>
}) {} : () -> tensor<?x?xf32>
func.return %0 : tensor<?x?xf32>
}
// -----
// Check that layouts of ops in function definitions are propagated by inferring
// layouts from function default layout values.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> (tensor<?x?xi32>{
// CHECK: "tf_device.cluster"
tf._default_layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.PartitionedCall"(%arg1, %arg2) {f = @callee1, config = "", config_proto = "", executor_type = ""} : (tensor<i32>, tensor<i32>) -> tensor<?x?xi32>
tf_device.return %1 : tensor<?x?xi32>
}) { _mesh = "|x=2,y=2|*CPU" } : () -> tensor<?x?xi32>
func.return %0 : tensor<?x?xi32>
}
// CHECK-LABEL: func private @callee1
// CHECK-SAME: %arg0: tensor<i32>
// CHECK-SAME: %arg1: tensor<i32>
// CHECK: tf._default_layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
func.func private @callee1(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<?x?xi32> attributes {tf.signature.is_stateful} {
%1 = "tf_device.cluster"() ({
%0 = "tf.PartitionedCall"(%arg0, %arg1) {f = @callee2, config = "", config_proto = "", executor_type = ""} : (tensor<i32>, tensor<i32>) -> tensor<?x?xi32>
tf_device.return %0 : tensor<?x?xi32>
}) { _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|*CPU" } : () -> tensor<?x?xi32>
func.return %1 : tensor<?x?xi32>
}
// CHECK-LABEL: func private @callee2
// CHECK-SAME: %arg0: tensor<i32>
// CHECK-SAME: %arg1: tensor<i32>
// CHECK: tf._default_layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
func.func private @callee2(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<?x?xi32> attributes {tf.signature.is_stateful} {
%1 = "tf_device.cluster"() ({
%0 = "tf.Const"() { value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tf.Fill"(%0, %arg1) {device = ""} : (tensor<2xi32>, tensor<i32>) -> tensor<?x?xi32>
tf_device.return %1 : tensor<?x?xi32>
}) { _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3" } : () -> tensor<?x?xi32>
func.return %1 : tensor<?x?xi32>
}
// -----
// Check that layouts of ops in function definitions are propagated by inferring
// layouts from function argument layouts.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<?x?xi32>{ tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<?x?xi32>{ tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"}) -> (tensor<?x?xi32>) {
%0 = "tf_device.cluster"() ({
%1 = "tf.PartitionedCall"(%arg1, %arg2) {f = @callee1, config = "", config_proto = "", executor_type = ""} : (tensor<?x?xi32>, tensor<?x?xi32>) -> tensor<?x?xi32>
tf_device.return %1 : tensor<?x?xi32>
}) { _mesh = "|x=2,y=2|*CPU" } : () -> tensor<?x?xi32>
func.return %0 : tensor<?x?xi32>
}
// CHECK-LABEL: func private @callee1
// CHECK-SAME: %arg0: tensor<?x?xi32>
// CHECK-SAME: tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: %arg1: tensor<?x?xi32>
// CHECK-SAME: tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
func.func private @callee1(%arg0: tensor<?x?xi32>, %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tf.signature.is_stateful} {
%1 = "tf_device.cluster"() ({
%0 = "tf.PartitionedCall"(%arg0, %arg1) {f = @callee2, config = "", config_proto = "", executor_type = ""} : (tensor<?x?xi32>, tensor<?x?xi32>) -> tensor<?x?xi32>
tf_device.return %0 : tensor<?x?xi32>
}) { _mesh = "mesh:CPU,x=2,y=2" } : () -> tensor<?x?xi32>
func.return %1 : tensor<?x?xi32>
}
// CHECK-LABEL: func private @callee2
// CHECK-SAME: %arg0: tensor<?x?xi32>
// CHECK-SAME: tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: %arg1: tensor<?x?xi32>
// CHECK-SAME: tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
func.func private @callee2(%arg0: tensor<?x?xi32>, %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tf.signature.is_stateful} {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Add"
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%1 = "tf_device.cluster"() ({
%1 = "tf.Add"(%arg0, %arg1) {device = ""} : (tensor<?x?xi32>, tensor<?x?xi32>) -> tensor<?x?xi32>
tf_device.return %1 : tensor<?x?xi32>
}) { _mesh = "|x=2,y=2|*CPU" } : () -> tensor<?x?xi32>
func.return %1 : tensor<?x?xi32>
}
// -----
// Check that layouts of ops in functions with multiple outputs with different
// layouts are supported.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x2xi32>{
tf._layout = "sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|*CPU"},
%arg2: tensor<2x2xi32>{
tf._layout = "sharding_specs:batch,unsharded, mesh:|batch=2,x=2|*CPU"})
-> (tensor<2x2xi32>) {
%0 = "tf_device.cluster"() ({
// CHECK: "tf.PartitionedCall"
// CHECK-SAME: _layout
// CHECK-SAME: "sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
%1, %2 = "tf.PartitionedCall"(%arg1, %arg2) {f = @callee1, config = "", config_proto = "", executor_type = ""} : (tensor<2x2xi32>, tensor<2x2xi32>) -> (tensor<2x2xi32>, tensor<2x2xi32>)
tf_device.return %1 : tensor<2x2xi32>
}) {_mesh = "|batch=2,x=2|*CPU"} : () -> tensor<2x2xi32>
func.return %0 : tensor<2x2xi32>
}
// CHECK-LABEL: func private @callee1
// CHECK-SAME: %arg0: tensor<2x2xi32>
// CHECK-SAME: tf._layout = "sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK-SAME: %arg1: tensor<2x2xi32>
// CHECK-SAME: tf._layout = "sharding_specs:batch,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK: tf._default_layout = "sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
// CHECK: tf._default_layout = "sharding_specs:batch,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
func.func private @callee1(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>) -> (tensor<2x2xi32>, tensor<2x2xi32>) attributes {tf.signature.is_stateful} {
%5, %6 = "tf_device.cluster"() ({
// CHECK: "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%1 = "tf.Const"() {value = dense<10> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// CHECK: "tf.Add"
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%2 = "tf.Add"(%1, %arg0) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// CHECK: "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%3 = "tf.Const"() {value = dense<10> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// CHECK: "tf.Add"
// CHECK-SAME: _layout = ["sharding_specs:batch,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
%4 = "tf.Add"(%3, %arg1) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %2, %4 : tensor<2x2xi32>, tensor<2x2xi32>
}) {_mesh = "|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<2x2xi32>, tensor<2x2xi32>)
func.return %5, %6 : tensor<2x2xi32>, tensor<2x2xi32>
}
// -----
// Unimplemented op throws an error.
func.func @main() {
%0 = "tf_device.cluster"() ({
// expected-error @+1 {{does not implement layout propagation}}
%0 = "tf.A"() : () -> tensor<2xi32>
tf_device.return %0 : tensor<2xi32>
}) {_mesh = "|batch=2,x=2|*CPU"} : () -> tensor<2xi32>
func.return
}
tensorflow/dtensor/mlir/tests/layout_propagation_v2.mlir
+1087
View File
File diff suppressed because it is too large Load Diff
tensorflow/dtensor/mlir/tests/lower_send_recv.mlir
+92
View File
@@ -0,0 +1,92 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: dtensor-opt %s -split-input-file -dtensor-lower-send-recv -verify-diagnostics | FileCheck %s
// Check that Data transfer from CPU to TPU is lowered correctly.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[DEVICE_ID:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// COMMENT: Check Recv before Send
// CHECK-DAG: %[[RECV_DEVICE_ID:.*]] = "tf.Reshape"(%[[DEVICE_ID]], %[[RECV_SIZE_TYPE:[^)]*]])
// CHECK-DAG: %[[RECV_DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[RECV_ID_TO_ORDINAL:.*]], %[[RECV_DEVICE_ID]], %[[RECV_SLICE_SIZE:[^)]*]])
// CHECK-DAG: %[[RECV_DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[RECV_DEVICE_ORDINAL]], %[[RECV_SCALAR_TYPE:[^)]*]])
// CHECK-DAG: %[[RECV_DEVICE_ORDINAL_SCALAR_64:.*]] = "tf.Cast"(%[[RECV_DEVICE_ORDINAL_SCALAR]])
// CHECK-DAG: %[[RECV_ID_TO_ORDINAL]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}>
// CHECK-DAG: %[[RECV_SIZE_TYPE]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
// CHECK-DAG: %[[RECV_SLICE_SIZE]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
// CHECK-DAG: %[[RECV_SCALAR_TYPE]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}>
// COMMENT: Recv and Send seperated by the output tensor.
// CHECK: %[[PROGRAM_KEY:.*]] = "tf._XlaCompileMlirPlaceholderProgramKey"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}>
// CHECK-NEXT: %[[LAYOUT_OUT:.*]] = "tf.DTensorLayout"(%[[CONST_OUT]])
// CHECK-DAG: %[[SEND_DEVICE_ID:.*]] = "tf.Reshape"(%[[DEVICE_ID]], %[[SEND_SIZE_TYPE:[^)]*]])
// CHECK-DAG: %[[SEND_DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[SEND_ID_TO_ORDINAL:.*]], %[[SEND_DEVICE_ID]], %[[SEND_SLICE_SIZE:[^)]*]])
// CHECK-DAG: %[[SEND_DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[SEND_DEVICE_ORDINAL]], %[[SEND_SCALAR_TYPE:[^)]*]])
// CHECK-DAG: %[[SEND_DEVICE_ORDINAL_SCALAR_64:.*]] = "tf.Cast"(%[[SEND_DEVICE_ORDINAL_SCALAR]])
// CHECK-DAG: %[[SEND_ID_TO_ORDINAL]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}>
// CHECK-DAG: %[[SEND_SIZE_TYPE]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
// CHECK-DAG: %[[SEND_SLICE_SIZE]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
// CHECK-DAG: %[[SEND_SCALAR_TYPE]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}>
// CHECK: "tf._XlaSendFromHostV2"(%[[LAYOUT_OUT]], %[[PROGRAM_KEY]], %[[SEND_DEVICE_ORDINAL_SCALAR_64]])
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_KEY]], %[[RECV_DEVICE_ORDINAL_SCALAR_64]])
// CHECK-SAME: key = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_2"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: %[[TPU_RECV_OUT:.*]] = "tf.XlaRecvFromHost"()
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-NEXT: %[[TPU_LAYOUT_OUT:.*]] = "tf.Relayout"(%[[TPU_RECV_OUT]])
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: "tf.XlaSendToHost"(%[[A_OUT]])
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<1xi32>) -> tensor<1xi32>
"tf.DTensorSend"(%1) {key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0", target_mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<1xi32>) -> ()
%2 = "tf.DTensorRecv"() {key = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_2", mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>, shape = #tf_type.shape<>, _layout = ["sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"]} : () -> (tensor<1xi32>)
"tf.B"(%2) : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0", mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, shape = #tf_type.shape<>, _layout = ["sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"]} : () -> tensor<1xi32>
%1 = "tf.Relayout"(%0) {global_shape = #tf_type.shape<1>, layout = "sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : (tensor<1xi32>) -> tensor<1xi32>
%2 = "tf.A"(%1) : (tensor<1xi32>) -> tensor<1xi32>
"tf.DTensorSend"(%2) {key = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_2", target_mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
// -----
// Check that device id usages are added correctly.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[DEVICE_ID:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) -> tensor<1xi32> {
// CHECK: "tf_device.cluster"()
// CHECK-NEXT: "tf.Identity"(%[[DEVICE_ID]])
// CHECK-NEXT: "tf.XlaRecvFromHost"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
%2 = "tf_device.cluster"() ({
%0 = "tf.XlaRecvFromHost"() {_layout = ["sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"], key = "communication_key_0", shape = #tf_type.shape<1>} : () -> tensor<1xi32>
tf_device.return %0 : tensor<1xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<1xi32>
func.return %2 : tensor<1xi32>
}
tensorflow/dtensor/mlir/tests/merge_clusters.mlir
+326
View File
@@ -0,0 +1,326 @@
// 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: dtensor-opt %s -split-input-file -dtensor-merge-clusters -verify-diagnostics | FileCheck %s
// Check that multiple tf_device.Cluster ops with same mesh specification are
// merged correctly to a single global cluster.
// CHECK-LABEL: module @test_merge_same_mesh
module @test_merge_same_mesh {
func.func @main(%arg0: tensor<i32>) -> (tensor<1xi32>, tensor<i64>, tensor<1xi32>, tensor<i64>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Cast"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.FloorMod"
// CHECK-NEXT: "tf.XlaRecvFromHost"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Equal"
// CHECK-NEXT: "tf.IfRegion"
// CHECK: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK: %[[CLUSTER_OUT:.*]]:4 = "tf_device.cluster"
// CHECK: "tf._XlaCompileMlirPlaceholderProgramKey"
// CHECK-NEXT: %[[CONST_OUT1:.*]] = "tf.Const"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf._XlaSendFromHostV2"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf._XlaSendFromHostV2"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf._XlaSendFromHostV2"
// CHECK-NEXT: %[[CONST_OUT2:.*]] = "tf.Const"
// CHECK-NEXT: "tf._XlaSendFromHostV2"
// CHECK-NEXT: "tf._XlaCompileMlirPlaceholderProgramKey"
// CHECK-NEXT: %[[CAST_OUT:.*]] = "tf.Cast"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.FloorMod"
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf._XlaRecvAtHostV2"
// CHECK-NEXT: tf_device.return %[[CONST_OUT1]], %[[CONST_OUT2]], %[[RECV_OUT]], %[[CAST_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0
// CHECK-NEXT: return %[[CLUSTER_OUT]]#0, %[[CLUSTER_OUT]]#1, %[[CLUSTER_OUT]]#2, %[[CLUSTER_OUT]]#3
%7, %8 = "tf_device.cluster"() ({
%0 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"], value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%2 = "tf.Const"() {value = dense<0> : tensor<i64>} : () -> tensor<i64>
"tf._XlaSendFromHostV2"(%1, %0, %2) {key = "communication_key_sharding_specs:, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"} : (tensor<1xi32>, tensor<2x!tf_type.string>, tensor<i64>) -> ()
%3 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
"tf._XlaSendFromHostV2"(%1, %0, %3) {key = "communication_key_sharding_specs:, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"} : (tensor<1xi32>, tensor<2x!tf_type.string>, tensor<i64>) -> ()
%4 = "tf.Const"() {value = dense<2> : tensor<i64>} : () -> tensor<i64>
"tf._XlaSendFromHostV2"(%1, %0, %4) {key = "communication_key_sharding_specs:, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"} : (tensor<1xi32>, tensor<2x!tf_type.string>, tensor<i64>) -> ()
%5 = "tf.Const"() {value = dense<3> : tensor<i64>} : () -> tensor<i64>
"tf._XlaSendFromHostV2"(%1, %0, %5) {_layout = [], key = "communication_key_sharding_specs:, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"} : (tensor<1xi32>, tensor<2x!tf_type.string>, tensor<i64>) -> ()
tf_device.return %1, %5 : tensor<1xi32>, tensor<i64>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> (tensor<1xi32>, tensor<i64>)
"tf_device.cluster"() ({
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%1 = "tf.Const"() {value = dense<4> : tensor<i64>} : () -> tensor<i64>
%2 = "tf.FloorMod"(%0, %1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%3 = "tf.XlaRecvFromHost"() {_layout = ["sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"], key = "communication_key_sharding_specs:, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0", shape = #tf_type.shape<1>} : () -> tensor<1xi32>
%4 = "tf.Const"() {value = dense<0> : tensor<i64>} : () -> tensor<i64>
%5 = "tf.Equal"(%2, %4) {incompatible_shape_error = true} : (tensor<i64>, tensor<i64>) -> tensor<i1>
"tf.IfRegion"(%5) ({
"tf.XlaSendToHost"(%3) {key = "communication_key_sharding_specs:, CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_2"} : (tensor<1xi32>) -> ()
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) {_layout = [], is_stateless = false} : (tensor<i1>) -> ()
tf_device.return {_layout = []}
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
%9, %10 = "tf_device.cluster"() ({
%0 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%1 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%2 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
%3 = "tf.FloorMod"(%1, %2) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%4 = "tf._XlaRecvAtHostV2"(%0, %3) {_layout = ["sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"], key = "communication_key_sharding_specs:, CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_2"} : (tensor<2x!tf_type.string>, tensor<i64>) -> tensor<1xi32>
tf_device.return %4, %1 : tensor<1xi32>, tensor<i64>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> (tensor<1xi32>, tensor<i64>)
func.return %7, %8, %9, %10 : tensor<1xi32>, tensor<i64>, tensor<1xi32>, tensor<i64>
}
}
// -----
// Check that duplicate/nested tf_device.cluster ops are removed.
// CHECK-LABEL: module @test_nested_cluster_are_removed
module @test_nested_cluster_are_removed {
// CHECK: func @main
func.func @main(%arg0: tensor<?xi32>) -> tensor<?xi32> {
// CHECK: "tf_device.cluster"
// CHECK: "tf.A"
// CHECK-NEXT: "tf.B"
// CHECK-NEXT: "tf.G"
// CHECK-NEXT: "tf.F"
// CHECK-NEXT: "tf.IfRegion"
// CHECK-NEXT: %[[D_OUT:.*]] = "tf.D"
// CHECK-NEXT: %[[I_OUT:.*]] = "tf.I"(%[[D_OUT]])
// CHECK-NEXT: "tf.J"(%[[I_OUT]])
// CHECK-NEXT: "tf.Yield"
// CHECK: %[[E_OUT:.*]] = "tf.E"
// CHECK-NEXT: tf_device.return %[[E_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
%7 = "tf.F"() : () -> tensor<?xi32>
"tf.IfRegion"(%6) ({
%10 = "tf_device.cluster"() ({
%8 = "tf.D"(%4, %3, %7) {} : (tensor<?xi32>, tensor<?xi32>, tensor<?xi32>) -> (tensor<?xi32>)
%9 = "tf.I"(%8) : (tensor<?xi32>) -> (tensor<?xi32>)
tf_device.return %9 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> (tensor<?xi32>)
"tf.J"(%10) : (tensor<?xi32>) -> ()
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
}
// -----
// Check whether metadata attributes are cloned correctly during cluster
// merging.
// CHECK-LABEL: module @test_clone_metadata
module @test_clone_metadata {
func.func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<2x4xf32>>>) -> () {
// CHECK: "tf_device.cluster"
// CHECK: "tf.B"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-NEXT: "tf_device.cluster"
// CHECK: "tf.A"
// CHECK-NEXT: "tf.C"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _inferred_resource_indices = dense<1>
// CHECK-SAME: _inferred_resource_layouts = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]
// CHECK-SAME: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
"tf_device.cluster"() ({
"tf.A"() : () -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
"tf_device.cluster"() ({
"tf.B"() : () -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.C"() : () -> (tensor<2x4xf32>)
"tf.AssignVariableOp"(%arg1, %0) : (tensor<!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _inferred_resource_indices = dense<1> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
func.return
}
}
// -----
// Check whether metadata attributes are merged correctly.
// CHECK-LABEL: module @test_merge_metadata
module @test_merge_metadata {
func.func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<2x4xf32>>>, %arg2: tensor<!tf_type.resource<tensor<2x4xf32>>>) -> () {
// CHECK: "tf_device.cluster"
// CHECK: "tf.B"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-NEXT: "tf_device.cluster"
// CHECK: "tf.A"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: "tf.C"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _inferred_resource_indices = dense<[1, 2]>
// CHECK-SAME: _inferred_resource_layouts = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0", "sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]
// CHECK-SAME: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
"tf_device.cluster"() ({
%0 = "tf.A"() : () -> (tensor<2x4xf32>)
"tf.AssignVariableOp"(%arg2, %0) : (tensor<!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _inferred_resource_indices = dense<2> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
"tf_device.cluster"() ({
"tf.B"() : () -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.C"() : () -> (tensor<2x4xf32>)
"tf.AssignVariableOp"(%arg1, %0) : (tensor<!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _inferred_resource_indices = dense<1> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
func.return
}
}
// -----
// Check whether shape op metadata attributes are merged correctly.
// CHECK-LABEL: module @test_shape_merge
module @test_shape_merge {
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xf32>, %arg2: tensor<2x4xf32>) -> () {
// CHECK: "tf_device.cluster"
// CHECK: "tf.B"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-NEXT: "tf_device.cluster"
// CHECK: "tf.ShapeOp"
// CHECK-NEXT: "tf.ShapeOp"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
// CHECK-SAME: _shape_input_indices = dense<[1, 2]
// CHECK-SAME: _shape_input_layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0", "sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]
"tf_device.cluster"() ({
%0 = "tf.ShapeOp"(%arg1) : (tensor<2x4xf32>) -> (tensor<1xf32>)
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _shape_input_indices = dense<[1, 2]> : vector<2xi32>, _shape_input_layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0", "sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
"tf_device.cluster"() ({
"tf.B"() : () -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.ShapeOp"(%arg2) : (tensor<2x4xf32>) -> (tensor<1xf32>)
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _shape_input_indices = dense<[1]> : vector<1xi32>, _shape_input_layout = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
func.return
}
}
// -----
// Check whether conflicting metadata attributes disallowed.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<2x4xf32>>>, %arg2: tensor<!tf_type.resource<tensor<2x4xf32>>>) -> () {
"tf_device.cluster"() ({
%0 = "tf.A"() : () -> (tensor<2x4xf32>)
"tf.AssignVariableOp"(%arg2, %0) : (tensor<!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _inferred_resource_indices = dense<1> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:x,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
"tf_device.cluster"() ({
"tf.B"() : () -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
// expected-error @+1 {{Found conflicting metadata attributes while merging clusters}}
"tf_device.cluster"() ({
%0 = "tf.C"() : () -> (tensor<2x4xf32>)
"tf.AssignVariableOp"(%arg1, %0) : (tensor<!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0", _inferred_resource_indices = dense<1> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:unsharded,unsharded, mesh:CPU|x=1|0|0|CPU:0"]} : () -> ()
func.return
}
// -----
// Check that unused tf_device.cluster results are pruned away.
// CHECK-LABEL module @test_prune_unused
// CHECK: func @main
// CHECK-SAME: %[[DEVICE_ID:.*]]: tensor<i32>
module @test_prune_unused {
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"()
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"
// CHECK-SAME: () -> ()
%2 = "tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
tf_device.return %0 : tensor<i64>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<i64>
func.return
}
}
// -----
// Check clusters with no mesh specification are disallowed.
func.func @main(%arg0: tensor<?xi32>) -> tensor<?xi32> {
%2 = "tf_device.cluster"() ({
%3 = "tf.A"() : () -> (tensor<?xi32>)
%4 = "tf.B"() : () -> (tensor<?xi32>)
%6 = "tf.G"() : () -> (tensor<i1>)
%7 = "tf.F"() : () -> tensor<?xi32>
"tf.IfRegion"(%6) ({
// expected-error @+1 {{All clusters must have specified mesh}}
"tf_device.cluster"() ({
"tf.D"() : () -> ()
tf_device.return
}) : () -> ()
"tf.Yield"() : () -> ()
}, {
"tf.Yield"() : () -> ()
}) {is_stateless = false} : (tensor<i1>) -> ()
%5 = "tf.E"() : () -> tensor<?xi32>
tf_device.return %5 : tensor<?xi32>
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> tensor<?xi32>
func.return %2 : tensor<?xi32>
}
tensorflow/dtensor/mlir/tests/mesh_propagation.mlir
+600
View File
@@ -0,0 +1,600 @@
// 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: dtensor-opt %s -split-input-file -dtensor-mesh-propagation -verify-diagnostics | FileCheck %s
// Checks that default mesh is propagated.
// CHECK-LABEL: module @test_default_mesh
// CHECK-SAME: tf._default_mesh = "[[DEFAULT_MESH:.*]]"
module @test_default_mesh attributes {tf._default_mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3" } {
// CHECK: func.func @main
func.func @main(%arg0: tensor<i32>) -> tensor<2xi32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "[[DEFAULT_MESH]]"
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {device = "", value = dense<[3, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
tf_device.return %cst : tensor<2xi32>
}) : () -> tensor<2xi32>
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "[[DEFAULT_MESH]]"
%1 = "tf_device.cluster"() ({
%2 = "tf.Identity"(%0) {device = ""} : (tensor<2xi32>) -> tensor<2xi32>
tf_device.return %2 : tensor<2xi32>
}) : () -> tensor<2xi32>
return %1 : tensor<2xi32>
}
}
// -----
// Checks that input mesh is correctly propagated to it's consumers.
// CHECK-LABEL: module @test_input_mesh
module @test_input_mesh {
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) {_mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.B"
// CHECK-NEXT: tf_device.return %[[B_OUT]]
// CHECK-NEXT: _mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%2 = "tf_device.cluster"() ({
%3 = "tf.B"(%0) : (tensor<i32>) -> tensor<f32>
tf_device.return %3 : tensor<f32>
}) : () -> (tensor<f32>)
func.return
}
}
// -----
// Checks that mesh is propagated from inputs of `tf_device.Cluster` op if the
// inputs are arguments of enclosing function.
// CHECK-LABEL: module @test_args_of_enclosing_func
module @test_args_of_enclosing_func {
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32> {tf._mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"}, %arg2: tensor<1xf32> {tf._mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"}) -> () {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.A"(%arg1, %arg2) : (tensor<1xf32>, tensor<1xf32>) -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> (tensor<i32>)
func.return
}
}
// -----
// Check that error is raised if mesh from inputs is not unique.
module @test_error_on_different_meshes {
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32> {tf._mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"}, %arg2: tensor<1xf32> {tf._mesh = "TPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"}) -> () {
// COMMENT: Regex is used to do a multiline match.
// expected-error-re @+1 {{All inputs to `tf_device.Cluster` must have same mesh configuration{{.*}}List of found inputs:{{.*}}Input Cluster 0:{{.*}}Input Cluster 1:}}
%0 = "tf_device.cluster"() ({
%1 = "tf.A"(%arg1, %arg2) : (tensor<1xf32>, tensor<1xf32>) -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> (tensor<i32>)
func.return
}
}
// -----
// Checks that mesh is correctly propagated from `tf_device.Cluster` op's consumers.
// CHECK-LABEL: module @test_cluster_to_consumers
module @test_cluster_to_consumers {
func.func @main() {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> (tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.B"
// CHECK-NEXT: tf_device.return %[[B_OUT]]
// CHECK-NEXT: _mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%2 = "tf_device.cluster"() ({
%3 = "tf.B"(%0) : (tensor<i32>) -> tensor<f32>
tf_device.return %3 : tensor<f32>
}) {_mesh = "CPU|batch=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<f32>)
func.return
}
}
// -----
// Checks that mesh is correctly propagated from default layout of the enclosing function.
// CHECK-LABEL: module @test_default_layout
module @test_default_layout {
func.func @main() ->(tensor<i32>{tf._default_layout = "sharding_specs:unsharded, mesh:CPU|batch=2,x=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.A"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "CPU|batch=2,x=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> (tensor<i32>)
func.return %0 : tensor<i32>
}
}
// -----
// Check mesh is propagate from function body if no mesh can be find from inputs.
// CHECK-LABEL: module @test_no_mesh_from_inputs
module @test_no_mesh_from_inputs {
func.func @main(%arg0: tensor<i32>) -> tensor<f32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @mesh_from_func_body} : (tensor<i32>) -> tensor<f32>
tf_device.return %1 : tensor<f32>
// CHECK: _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
}) : () -> tensor<f32>
func.return %0 : tensor<f32>
}
func.func private @mesh_from_func_body(%arg0: tensor<i32>) -> tensor<f32> attributes {tf.signature.is_stateful} {
%0 = "tf_device.cluster"() ({
%3 = "tf.Const"() {_layout = ["sharding_specs: mesh:TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"], value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
tf_device.return %3 : tensor<0xi32>
}) {_mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> tensor<0xi32>
%1 = "tf_device.cluster"() ({
%3 = "tf.RandomUniform"(%0) {_layout = ["sharding_specs: mesh:TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"], seed = 0 : i64, seed2 = 0 : i64} : (tensor<0xi32>) -> tensor<f32>
tf_device.return %3 : tensor<f32>
}) {_mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> tensor<f32>
%2 = "tf_device.cluster"() ({
%3 = "tf.Identity"(%1) {} : (tensor<f32>) -> tensor<f32>
tf_device.return %3 : tensor<f32>
// CHECK: _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
}) : () -> tensor<f32>
func.return %2 : tensor<f32>
}
}
// -----
// CHECK-LABEL: module @test_return_const
module @test_return_const {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<8x128x128xf32> {tf._global_shape = #tf_type.shape<8x128x128>, tf._layout = "sharding_specs:x,unsharded,unsharded, mesh:CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1", tf._mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"}) -> (tensor<3xi32> {tf._global_shape = #tf_type.shape<3>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_global_shape = [#tf_type.shape<3>], value = dense<[8, 128, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
tf_device.return {_global_shape = []} %1 : tensor<3xi32>
}) {_global_shape = [#tf_type.shape<3>]} : () -> tensor<3xi32>
func.return %0 : tensor<3xi32>
}
}
// -----
// CHECK-LABEL: module @test_multi_mesh
module @test_multi_mesh {
func.func @main(%arg0: tensor<4xi32> {tf._layout = "sharding_specs:not_sharded mesh:CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"},
%arg1: tensor<4xi32> {tf._layout = "sharding_specs:not_sharded mesh:CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"}) -> (tensor<4xi32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[8, 8, 128, 128]> : tensor<4xi32>} : () -> tensor<4xi32>
tf_device.return %1 : tensor<4xi32>
}) : () -> tensor<4xi32>
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Add"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%2 = "tf_device.cluster"() ({
%3 = "tf.Add"(%arg0, %0) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
tf_device.return %3 : tensor<4xi32>
}) : () -> tensor<4xi32>
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "CPU|x=2,y=1|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"
%4 = "tf_device.cluster"() ({
%5 = "tf.Identity"(%arg1) : (tensor<4xi32>) -> tensor<4xi32>
tf_device.return %5 : tensor<4xi32>
}) : () -> tensor<4xi32>
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Relayout"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
%6 = "tf_device.cluster"() ({
%7 = "tf.Relayout"(%2) { layout = "sharding_specs:not_sharded mesh:TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : (tensor<4xi32>) -> tensor<4xi32>
tf_device.return %7 : tensor<4xi32>
}) { _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1" } : () -> tensor<4xi32>
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Relayout"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
%8 = "tf_device.cluster"() ({
%9 = "tf.Relayout"(%4) { layout = "sharding_specs:not_sharded mesh:TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : (tensor<4xi32>) -> tensor<4xi32>
tf_device.return %9 : tensor<4xi32>
}) { _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1" } : () -> tensor<4xi32>
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Add"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"
%10 = "tf_device.cluster"() ({
%11 = "tf.Add"(%6, %8) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
tf_device.return %11 : tensor<4xi32>
}) : () -> tensor<4xi32>
func.return %10 :tensor<4xi32>
}
}
// -----
// Check mesh propagation of ops inside tf.WhileRegion op.
// CHECK-LABEL: module @test_while
module @test_while {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4xf32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3", tf._mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"})
-> (tensor<4xf32> {tf._default_layout = "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.WhileRegion"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: constant
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.NotEqual"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: constant
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Sub"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
// CHECK-NEXT: "tf.Yield"
// CHECK: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"
%15:2 = "tf_device.cluster"() ({
%2:2 = "tf.WhileRegion"(%arg1, %arg0) ({
^bb0(%carg0: tensor<4xf32>, %carg1: tensor<i32>):
%11 = "tf_device.cluster"() ({
%limit = arith.constant dense<5> : tensor<i32>
tf_device.return %limit : tensor<i32>
}) : () -> tensor<i32>
%12 = "tf_device.cluster"() ({
%cond = "tf.NotEqual"(%carg1, %11) : (tensor<i32>, tensor<i32>) -> tensor<i1>
tf_device.return %cond : tensor<i1>
}) : () -> tensor<i1>
"tf.Yield"(%12) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<4xf32>, %barg1: tensor<i32>):
%13 = "tf_device.cluster"() ({
%one = arith.constant dense<1.0> : tensor<4xf32>
tf_device.return %one: tensor<4xf32>
}) : () -> tensor<4xf32>
%14 = "tf_device.cluster"() ({
%sub = "tf.Sub"(%barg0, %13) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
tf_device.return %sub: tensor<4xf32>
}) : () -> tensor<4xf32>
"tf.Yield"(%14, %barg1) : (tensor<4xf32>, tensor<i32>) -> ()
}) {is_stateless = true} : (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
tf_device.return %2#0, %2#1 : tensor<4xf32>, tensor<i32>
}) : () -> (tensor<4xf32>, tensor<i32>)
%16 = "tf_device.cluster"() ({
%5 = "tf.Identity"(%15#0) : (tensor<4xf32>) -> (tensor<4xf32>)
tf_device.return %5 : tensor<4xf32>
}) : () -> tensor<4xf32>
func.return %16 : tensor<4xf32>
}
}
// -----
// Check mesh propagation of while op with gpu and cpu
// CHECK-LABEL: module @test_while_gpu_cpu
module @test_while_gpu_cpu {
func.func @main(%arg0: tensor<i32> {tf._layout = "sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh="|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg1: tensor<4xf32> {tf._layout = "sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh="|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<4xf32> {tf._layout = "sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1", tf._mesh="|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"})
-> (tensor<4xf32> {tf._default_layout = "sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0,op_input_1", outputs = "op_output_0"}} {
// CHECK: tf.Sub
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK: tf.NotEqual
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK: tf.Add
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK: tf.Relayout
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"
// CHECK: tf.DTensorLayout
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"
// CHECK: tf.Mul
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"
%11 = "tf_device.cluster"() ({
%12 = "tf.Sub"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %12 : tensor<i32>
}) : () -> tensor<i32>
%1:3 = "tf_device.cluster"() ({
%2:3 = "tf.WhileRegion"(%11, %arg1, %arg2) ({
^bb0(%barg0: tensor<i32>, %barg1: tensor<4xf32>, %barg2: tensor<4xf32>):
%cond = "tf_device.cluster"() ({
%c = "tf.NotEqual"(%barg0, %barg0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
tf_device.return %c : tensor<i1>
}) : () -> tensor<i1>
"tf.Yield"(%cond) : (tensor<i1>) -> ()
}, {
^bb0(%carg0: tensor<i32>, %carg1: tensor<4xf32>, %carg2: tensor<4xf32>):
%3 = "tf_device.cluster"() ({
%4 = "tf.Add"(%carg1, %carg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
tf_device.return %4 : tensor<4xf32>
}) : () -> tensor<4xf32>
%5 = "tf_device.cluster"() ({
%6 = "tf.Relayout"(%3) {device="/job:localhost/replica:0/task:0/device:CUSTOM:0", layout="sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %6 : tensor<4xf32>
}) {_mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"} : () -> tensor<4xf32>
%7 = "tf_device.cluster"() ({
%8 = "tf.DTensorLayout"(%5) {device="", global_shape = #tf_type.shape<4>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1>} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %8 : tensor<4xf32>
}) {_mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1"} : () -> tensor<4xf32>
%9 = "tf_device.cluster"() ({
%10 = "tf.Mul"(%7, %carg2) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
tf_device.return %10 : tensor<4xf32>
}) : () -> tensor<4xf32>
"tf.Yield"(%carg0, %carg1, %9) : (tensor<i32>, tensor<4xf32>, tensor<4xf32>) -> ()
}) {is_stateless = true} : (tensor<i32>, tensor<4xf32>, tensor<4xf32>) -> (tensor<i32>, tensor<4xf32>, tensor<4xf32>)
tf_device.return %2#0, %2#1, %2#2 : tensor<i32>, tensor<4xf32>, tensor<4xf32>
}) : () -> (tensor<i32>, tensor<4xf32>, tensor<4xf32>)
func.return %1#2 : tensor<4xf32>
}
}
// -----
// Check mesh propagation of ops inside tf.IfRegion op.
// CHECK-LABEL: module @test_if
module @test_if {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4xf32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> tensor<4xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Sum"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Equal"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.IfRegion"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Sqrt"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Relayout"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf_device.cluster"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
// CHECK-NEXT: "tf.Yield"
// CHECK-NEXT: (tensor<i1>) -> (tensor<i1>, tensor<4xf32>)
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"
%0 = "tf_device.cluster"() ({
%10 = "tf.Const"() {value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
tf_device.return %10 : tensor<f32>
}) : () -> tensor<f32>
%1 = "tf_device.cluster"() ({
%10 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
tf_device.return %10 : tensor<1xi32>
}) : () -> tensor<1xi32>
%2 = "tf_device.cluster"() ({
%10 = "tf.Const"() {value = dense<0.000000e+00> : tensor<4xf32>} : () -> tensor<4xf32>
tf_device.return %10 : tensor<4xf32>
}) : () -> tensor<4xf32>
%3 = "tf_device.cluster"() ({
%10 = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
tf_device.return %10 : tensor<i1>
}) : () -> tensor<i1>
%5 = "tf_device.cluster"() ({
%10 = "tf.Sum"(%arg1, %1) {device = "", keep_dims = false} : (tensor<4xf32>, tensor<1xi32>) -> tensor<f32>
tf_device.return %10 : tensor<f32>
}) : () -> tensor<f32>
%6 = "tf_device.cluster"() ({
%10 = "tf.Equal"(%5, %0) {device = "", incompatible_shape_error = true} : (tensor<f32>, tensor<f32>) -> tensor<i1>
tf_device.return %10 : tensor<i1>
}) : () -> tensor<i1>
%7:2 = "tf_device.cluster"() ({
%10:2 = "tf.IfRegion"(%6) ({
%11 = "tf_device.cluster"() ({
%16 = "tf.Identity"(%3) {device = ""} : (tensor<i1>) -> tensor<i1>
tf_device.return %16 : tensor<i1>
}) : () -> tensor<i1>
%12 = "tf_device.cluster"() ({
%16 = "tf.Sqrt"(%arg1) {device = ""} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %16 : tensor<4xf32>
}) : () -> tensor<4xf32>
%13 = "tf_device.cluster"() ({
%16 = "tf.Relayout"(%12) {device = "", layout = "sharding_specs:x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %16 : tensor<4xf32>
}) : () -> tensor<4xf32>
%14 = "tf_device.cluster"() ({
%16 = "tf.DTensorLayout"(%13) {global_shape = #tf_type.shape<4>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %16 : tensor<4xf32>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<4xf32>
%15 = "tf_device.cluster"() ({
%16 = "tf.Identity"(%14) {device = ""} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %16 : tensor<4xf32>
}) : () -> tensor<4xf32>
"tf.Yield"(%11, %15) : (tensor<i1>, tensor<4xf32>) -> ()
}, {
%11 = "tf_device.cluster"() ({
%13 = "tf.Identity"(%3) {device = ""} : (tensor<i1>) -> tensor<i1>
tf_device.return %13 : tensor<i1>
}) : () -> tensor<i1>
%12 = "tf_device.cluster"() ({
%13 = "tf.Identity"(%2) {device = ""} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %13 : tensor<4xf32>
}) : () -> tensor<4xf32>
"tf.Yield"(%11, %12) : (tensor<i1>, tensor<4xf32>) -> ()
}) {_else_func_name = "cond_false_150", _lower_using_switch_merge = true, _read_only_resource_inputs = [], _then_func_name = "cond_true_140", device = "", is_stateless = true} : (tensor<i1>) -> (tensor<i1>, tensor<4xf32>)
tf_device.return %10#0, %10#1 : tensor<i1>, tensor<4xf32>
}) : () -> (tensor<i1>, tensor<4xf32>)
%8 = "tf_device.cluster"() ({
%10 = "tf.Identity"(%7#1) {device = ""} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %10 : tensor<4xf32>
}) : () -> tensor<4xf32>
%9 = "tf_device.cluster"() ({
%10 = "tf.Identity"(%8) {device = ""} : (tensor<4xf32>) -> tensor<4xf32>
tf_device.return %10 : tensor<4xf32>
}) : () -> tensor<4xf32>
func.return %9 : tensor<4xf32>
}
}
// -----
// Check mesh propagation of tf.WhileRegion inside tf.IfRegion op.
// This test only checks that the code doesn't crash under asan.
// Correctness check are covered by other tests.
// CHECK-LABEL: module @test_nested_while_inside_if
module @test_nested_while_inside_if {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4xf32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> tensor<4xf32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<0> : tensor<i1>} : () -> tensor<i1>
tf_device.return %1 : tensor<i1>
}) : () -> tensor<i1>
%7:1 = "tf_device.cluster"() ({
%10:1 = "tf.IfRegion"(%0) ({
%3:2 = "tf.WhileRegion"(%arg1, %arg0) ({
^bb0(%carg0: tensor<4xf32>, %carg1: tensor<i32>):
%11 = "tf_device.cluster"() ({
%limit = arith.constant dense<5> : tensor<i32>
tf_device.return %limit : tensor<i32>
}) : () -> tensor<i32>
%12 = "tf_device.cluster"() ({
%cond = "tf.NotEqual"(%carg1, %11) : (tensor<i32>, tensor<i32>) -> tensor<i1>
tf_device.return %cond : tensor<i1>
}) : () -> tensor<i1>
"tf.Yield"(%12) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<4xf32>, %barg1: tensor<i32>):
%13 = "tf_device.cluster"() ({
%one = arith.constant dense<1.0> : tensor<4xf32>
tf_device.return %one: tensor<4xf32>
}) : () -> tensor<4xf32>
%14 = "tf_device.cluster"() ({
%sub = "tf.Sub"(%barg0, %13) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
tf_device.return %sub: tensor<4xf32>
}) : () -> tensor<4xf32>
"tf.Yield"(%14, %barg1) : (tensor<4xf32>, tensor<i32>) -> ()
}) {is_stateless = true} : (tensor<4xf32>, tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
"tf.Yield"(%3#0) : (tensor<4xf32>) -> ()
}, {
"tf.Yield"(%arg1) : (tensor<4xf32>) -> ()
}) {_else_func_name = "cond_false_150", _lower_using_switch_merge = true, _read_only_resource_inputs = [], _then_func_name = "cond_true_140", device = "", is_stateless = true} : (tensor<i1>) -> (tensor<4xf32>)
tf_device.return %10#0 : tensor<4xf32>
}) : () -> (tensor<4xf32>)
func.return %7#0 : tensor<4xf32>
}
}
tensorflow/dtensor/mlir/tests/move_compilation_to_host.mlir
+401
View File
@@ -0,0 +1,401 @@
// 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: dtensor-opt %s -split-input-file -dtensor-move-compilation-to-host -verify-diagnostics | FileCheck %s
// Check that TPU Compilation ops are moved to host computation functions and
// Send/Recv ops are inserted to transfer program key.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 656 : i32}} {
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,%arg1: tensor<4xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<f32> {tf._global_shape = #tf_type.shape<>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// CHECK: "tf.StatefulPartitionedCall"
// CHECK-SAME: f = @_func_0
// CHECK-SAME: _mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"
// CHECK-NEXT: "tf.StatefulPartitionedCall"
// CHECK-SAME: f = @_func_1
// CHECK-SAME: _mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_layout = [], _mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_0} : (tensor<i32>, tensor<4xi32>) -> ()
%0 = "tf.StatefulPartitionedCall"(%arg0) {_layout = ["sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"], _mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config_proto = "", executor_type = "", f = @_func_1} : (tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func private @_func_0
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<4xi32>
func.func private @_func_0(%arg0: tensor<i32>, %arg1: tensor<4xi32>) {
// CHECK-NEXT: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<[0, 1]>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK: %[[PROGRAM_KEY:.*]] = "tf.Case"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: "tf_device.launch"()
// CHECK-NEXT: "tf.TPUExecute"(%[[ARG0]], %[[ARG1]], %[[PROGRAM_KEY]])
// CHECK-NEXT: tf_device.return
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%arg0, %arg1, %0#1) : (tensor<i32>, tensor<4xi32>, tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
// CHECK-LABEL: func private @_func_1
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func private @_func_1(%arg0: tensor<i32>) -> tensor<f32> {
// CHECK: %[[COMPILE_OUT:.*]]:2 = "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:CPU:0"}>
// CHECK-NEXT: %[[COMPILATION_STATUS:.*]], %[[PROGRAM_KEY:.*]] = "tf._TPUCompileMlir"()
// CHECK-NEXT: "tf._HostSend"(%[[PROGRAM_KEY]])
// CHECK-SAME: recv_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT: "tf._HostSend"(%[[PROGRAM_KEY]])
// CHECK-SAME: recv_device = "/job:localhost/replica:0/task:0/device:TPU:0"
// CHECK-SAME: send_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device_incarnation = 0
// CHECK-SAME: tensor_name = "compilation_send_recv_key_0
// CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT: "tf._HostSend"(%[[PROGRAM_KEY]])
// CHECK-SAME: recv_device = "/job:localhost/replica:0/task:0/device:TPU:1"
// CHECK-SAME: send_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device_incarnation = 0
// CHECK-SAME: tensor_name = "compilation_send_recv_key_1
// CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT: tf_device.return %[[COMPILATION_STATUS]], %[[PROGRAM_KEY]]
// CHECK: "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:CPU:0"}>
// CHECK-NEXT: "tf.TPUCompileSucceededAssert"(%[[COMPILE_OUT]]#0)
// CHECK-NEXT: tf_device.return
// CHECK: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[BROADCASTED_KEY:.*]] = "tf.Case"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Cast"
// CHECK-NEXT: %[[MOD_OUT:.*]] = "tf.FloorMod"
// CHECK-NEXT: "tf._XlaRecvAtHostV2"(%[[BROADCASTED_KEY]], %[[MOD_OUT]]
// CHECK-NEXT: "tf.Sqrt"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: return
%0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%2 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%3 = "tf.FloorMod"(%2, %0) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%4 = "tf._XlaRecvAtHostV2"(%1, %3) {key = "communication_key_sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0_0"} : (tensor<2x!tf_type.string>, tensor<i64>) -> tensor<f32>
%5 = "tf.Sqrt"(%4) : (tensor<f32>) -> tensor<f32>
%6 = "tf.Identity"(%5) : (tensor<f32>) -> tensor<f32>
%7 = "tf.Identity"(%6) : (tensor<f32>) -> tensor<f32>
func.return %7 : tensor<f32>
}
}
// -----
// Check that TPU Compilation ops are moved to host computation functions and
// Send/Recv ops are inserted to transfer program key for
// TPUExecuteAndUpdateVariables op
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 656 : i32}} {
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,%arg1: tensor<*x!tf_type.resource<tensor<4xf32>>> {tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<f32> {tf._global_shape = #tf_type.shape<>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// CHECK: "tf.StatefulPartitionedCall"
// CHECK-SAME: f = @_func_0
// CHECK-SAME: _mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"
// CHECK-NEXT: "tf.StatefulPartitionedCall"
// CHECK-SAME: f = @_func_1
// CHECK-SAME: _mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_layout = [], _mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_0} : (tensor<i32>, tensor<*x!tf_type.resource<tensor<4xf32>>>) -> ()
%0 = "tf.StatefulPartitionedCall"(%arg0) {_layout = ["sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0"], _mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config_proto = "", executor_type = "", f = @_func_1} : (tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
// CHECK-LABEL: func private @_func_0
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>
func.func private @_func_0(%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource<tensor<4xf32>>>) {
// CHECK-NEXT: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<[0, 1]>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK: %[[PROGRAM_KEY:.*]] = "tf.Case"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: "tf_device.launch"()
// CHECK-NEXT: "tf.TPUExecuteAndUpdateVariables"(%[[ARG1]], %[[PROGRAM_KEY]])
// CHECK-NEXT: tf_device.return
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecuteAndUpdateVariables"(%arg1, %0#1) {device_var_reads_indices = [0], device_var_updates_indices = [-1]} : (tensor<*x!tf_type.resource<tensor<4xf32>>>, tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
// CHECK-LABEL: func private @_func_1
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func private @_func_1(%arg0: tensor<i32>) -> tensor<f32> {
// CHECK: %[[COMPILE_OUT:.*]]:2 = "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:CPU:0"}>
// CHECK-NEXT: %[[COMPILATION_STATUS:.*]], %[[PROGRAM_KEY:.*]] = "tf._TPUCompileMlir"()
// CHECK-NEXT: "tf._HostSend"(%[[PROGRAM_KEY]])
// CHECK-SAME: recv_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device_incarnation = 0
// CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT: "tf._HostSend"(%[[PROGRAM_KEY]])
// CHECK-SAME: recv_device = "/job:localhost/replica:0/task:0/device:TPU:0"
// CHECK-SAME: send_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device_incarnation = 0
// CHECK-SAME: tensor_name = "compilation_send_recv_key_0
// CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT: "tf._HostSend"(%[[PROGRAM_KEY]])
// CHECK-SAME: recv_device = "/job:localhost/replica:0/task:0/device:TPU:1"
// CHECK-SAME: send_device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-SAME: send_device_incarnation = 0
// CHECK-SAME: tensor_name = "compilation_send_recv_key_1
// CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT: tf_device.return %[[COMPILATION_STATUS]], %[[PROGRAM_KEY]]
// CHECK: "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:CPU:0"}>
// CHECK-NEXT: "tf.TPUCompileSucceededAssert"(%[[COMPILE_OUT]]#0)
// CHECK-NEXT: tf_device.return
// CHECK: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[BROADCASTED_KEY:.*]] = "tf.Case"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Cast"
// CHECK-NEXT: %[[MOD_OUT:.*]] = "tf.FloorMod"
// CHECK-NEXT: "tf._XlaRecvAtHostV2"(%[[BROADCASTED_KEY]], %[[MOD_OUT]]
// CHECK-NEXT: "tf.Sqrt"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: "tf.Identity"
// CHECK-NEXT: return
%0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%2 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%3 = "tf.FloorMod"(%2, %0) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%4 = "tf._XlaRecvAtHostV2"(%1, %3) {key = "communication_key_sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0_0"} : (tensor<2x!tf_type.string>, tensor<i64>) -> tensor<f32>
%5 = "tf.Sqrt"(%4) : (tensor<f32>) -> tensor<f32>
%6 = "tf.Identity"(%5) : (tensor<f32>) -> tensor<f32>
%7 = "tf.Identity"(%6) : (tensor<f32>) -> tensor<f32>
func.return %7 : tensor<f32>
}
}
// -----
// Check that StatefulPartitionedCall op without mesh specification is
// disallowed.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 656 : i32}} {
func.func @main(%arg0: tensor<i32>,%arg1: tensor<4xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<f32> {tf._global_shape = #tf_type.shape<>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_0} : (tensor<i32>, tensor<4xi32>) -> ()
// expected-error @+1 {{StatefulPartitionCall op must have `_mesh` attribute specified}}
"tf.StatefulPartitionedCall"(%arg0) {config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_2} : (tensor<i32>) -> ()
%0 = "tf.StatefulPartitionedCall"(%arg0) {_mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config_proto = "", executor_type = "", f = @_func_1} : (tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
func.func private @_func_0(%arg0: tensor<i32>, %arg1: tensor<4xi32>) {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%arg0, %arg1, %0#1) : (tensor<i32>, tensor<4xi32>, tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
func.func private @_func_2(%arg0: tensor<i32>) {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%0#1) : (tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
func.func private @_func_1(%arg0: tensor<i32>) -> tensor<f32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%2 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%3 = "tf.FloorMod"(%2, %0) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%4 = "tf._XlaRecvAtHostV2"(%1, %3) {key = "communication_key_sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0_0"} : (tensor<2x!tf_type.string>, tensor<i64>) -> tensor<f32>
%5 = "tf.Identity"(%4) : (tensor<f32>) -> tensor<f32>
func.return %5 : tensor<f32>
}
}
// -----
// Check that multiple TPU cluster computations with same mesh is disallowed.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 656 : i32}} {
func.func @main(%arg0: tensor<i32>,%arg1: tensor<4xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<f32> {tf._global_shape = #tf_type.shape<>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_0} : (tensor<i32>, tensor<4xi32>) -> ()
// expected-error @+1 {{There should be exactly 1 function for each mesh in computation cluster}}
"tf.StatefulPartitionedCall"(%arg0) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_2} : (tensor<i32>) -> ()
%0 = "tf.StatefulPartitionedCall"(%arg0) {_mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config_proto = "", executor_type = "", f = @_func_1} : (tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
func.func private @_func_0(%arg0: tensor<i32>, %arg1: tensor<4xi32>) {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%arg0, %arg1, %0#1) : (tensor<i32>, tensor<4xi32>, tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
func.func private @_func_2(%arg0: tensor<i32>) {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%0#1) : (tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
func.func private @_func_1(%arg0: tensor<i32>) -> tensor<f32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%2 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%3 = "tf.FloorMod"(%2, %0) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%4 = "tf._XlaRecvAtHostV2"(%1, %3) {key = "communication_key_sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0_0"} : (tensor<2x!tf_type.string>, tensor<i64>) -> tensor<f32>
%5 = "tf.Identity"(%4) : (tensor<f32>) -> tensor<f32>
func.return %5 : tensor<f32>
}
}
// -----
// Check that at multiple TPU computations are disallowed.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 656 : i32}} {
// expected-error @+1 {{Only 1 XLA cluster for DTensor computation is supported for now}}
func.func @main(%arg0: tensor<i32>,%arg1: tensor<4xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<f32> {tf._global_shape = #tf_type.shape<>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func_0} : (tensor<i32>, tensor<4xi32>) -> ()
"tf.StatefulPartitionedCall"(%arg0) {_mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:TPU:0", config = "|x=1|0|0|/job:localhost/replica:0/task:0/device:TPU:0", config_proto = "", executor_type = "", f = @_func_2} : (tensor<i32>) -> ()
%0 = "tf.StatefulPartitionedCall"(%arg0) {_mesh = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config = "|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0", config_proto = "", executor_type = "", f = @_func_1} : (tensor<i32>) -> tensor<f32>
func.return %0 : tensor<f32>
}
func.func private @_func_0(%arg0: tensor<i32>, %arg1: tensor<4xi32>) {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%arg0, %arg1, %0#1) : (tensor<i32>, tensor<4xi32>, tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
func.func private @_func_2(%arg0: tensor<i32>) {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "...", mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
"tf_device.launch"() ({
"tf.TPUExecute"(%0#1) : (tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
func.return
}
func.func private @_func_1(%arg0: tensor<i32>) -> tensor<f32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<2x!tf_type.string>
%2 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<i32>) -> tensor<i64>
%3 = "tf.FloorMod"(%2, %0) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%4 = "tf._XlaRecvAtHostV2"(%1, %3) {key = "communication_key_sharding_specs: mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0_0"} : (tensor<2x!tf_type.string>, tensor<i64>) -> tensor<f32>
%5 = "tf.Identity"(%4) : (tensor<f32>) -> tensor<f32>
func.return %5 : tensor<f32>
}
}
tensorflow/dtensor/mlir/tests/multi_device_expansion.mlir
+253
View File
@@ -0,0 +1,253 @@
// 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: dtensor-opt %s -split-input-file -dtensor-multi-device-expansion -verify-diagnostics | FileCheck %s
module attributes {
tf._default_mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:CPU:1", "/job:localhost/replica:0/task:0/device:CPU:2", "/job:localhost/replica:0/task:0/device:CPU:3", "/job:localhost/replica:0/task:0/device:CPU:4", "/job:localhost/replica:0/task:0/device:CPU:5", "/job:localhost/replica:0/task:0/device:CPU:6", "/job:localhost/replica:0/task:0/device:CPU:7"},
dtensor.enable_multi_device_mode = true
} {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<8xi32> {tf._global_shape = #tf_type.shape<8>, tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", tf._mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"}) -> (tensor<8xi32> {tf._global_shape = #tf_type.shape<8>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
%1 = "tf.StatefulPartitionedCall"(%arg0, %arg1) {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"], _mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", config = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", config_proto = "", executor_type = "", f = @_test_func} : (tensor<i32>, tensor<8xi32>) -> tensor<8xi32>
return %1 : tensor<8xi32>
}
func.func private @_test_func(%arg0: tensor<i32>, %arg1: tensor<8xi32>) -> tensor<8xi32> {
return %arg1 : tensor<8xi32>
}
// CHECK-LABEL: func.func @main
// CHECK: %arg0: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK: %arg1: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:1"}
// CHECK: %arg2: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:2"}
// CHECK: %arg3: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:3"}
// CHECK: %arg4: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:4"}
// CHECK: %arg5: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:5"}
// CHECK: %arg6: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:6"}
// CHECK: %arg7: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:7"}
// CHECK: tf.entry_function = {inputs = "input_0,input_1,input_2,input_3,input_4,input_5,input_6,input_7", outputs = "output_0,output_1,output_2,output_3,output_4,output_5,output_6,output_7"
// CHECK: %[[RES:.*]]:8 = "tf.StatefulPartitionedCall"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7)
// CHECK-SAME: f = @_multi_device_func_16526499781244736596_7754837423929193503
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"]
// CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2, %[[RES]]#3, %[[RES]]#4, %[[RES]]#5, %[[RES]]#6, %[[RES]]#7
// CHECK-LABEL: func.func private @_multi_device_func_16526499781244736596_7754837423929193503(
// CHECK: %arg0: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}
// CHECK: %arg1: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:1"}
// CHECK: %arg2: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:2"}
// CHECK: %arg3: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:3"}
// CHECK: %arg4: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:4"}
// CHECK: %arg5: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:5"}
// CHECK: %arg6: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:6"}
// CHECK: %arg7: tensor<8xi32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:7"}
// CHECK: tf.entry_function = {inputs = "input_0,input_1,input_2,input_3,input_4,input_5,input_6,input_7", outputs = "output_0,output_1,output_2,output_3,output_4,output_5,output_6,output_7"
// CHECK: %[[CST0:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK: %[[CST1:.*]] = "tf.Const"
// CHECK: %[[CST2:.*]] = "tf.Const"
// CHECK: %[[CST3:.*]] = "tf.Const"
// CHECK: %[[CST4:.*]] = "tf.Const"
// CHECK: %[[CST5:.*]] = "tf.Const"
// CHECK: %[[CST6:.*]] = "tf.Const"
// CHECK: %[[CST7:.*]] = "tf.Const"
// CHECK: %[[RES0:.*]] = "tf.StatefulPartitionedCall"(%[[CST0]], %arg0)
// CHECK: %[[RES1:.*]] = "tf.StatefulPartitionedCall"(%[[CST1]], %arg1)
// CHECK: %[[RES2:.*]] = "tf.StatefulPartitionedCall"(%[[CST2]], %arg2)
// CHECK: %[[RES3:.*]] = "tf.StatefulPartitionedCall"(%[[CST3]], %arg3)
// CHECK: %[[RES4:.*]] = "tf.StatefulPartitionedCall"(%[[CST4]], %arg4)
// CHECK: %[[RES5:.*]] = "tf.StatefulPartitionedCall"(%[[CST5]], %arg5)
// CHECK: %[[RES6:.*]] = "tf.StatefulPartitionedCall"(%[[CST6]], %arg6)
// CHECK: %[[RES7:.*]] = "tf.StatefulPartitionedCall"(%[[CST7]], %arg7)
// CHECK: return %[[RES0]], %[[RES1]], %[[RES2]], %[[RES3]], %[[RES4]], %[[RES5]], %[[RES6]], %[[RES7]]
}
// -----
// Ensure that expansion doesn't occur when it's disabled.
module attributes { dtensor.enable_multi_device_mode = false } {
// CHECK-LABEL: func.func @main
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<8xi32> {tf._global_shape = #tf_type.shape<8>, tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", tf._mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"}) -> (tensor<8xi32> {tf._global_shape = #tf_type.shape<8>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// CHECK: %0 = "tf.StatefulPartitionedCall"(%arg0, %arg1)
%0 = "tf.StatefulPartitionedCall"(%arg0, %arg1) {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"], _mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", config = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", config_proto = "", executor_type = "", f = @_test_func} : (tensor<i32>, tensor<8xi32>) -> tensor<8xi32>
return %0 : tensor<8xi32>
}
func.func private @_test_func(%arg0: tensor<i32>, %arg1: tensor<8xi32>) -> tensor<8xi32> {
return %arg1 : tensor<8xi32>
}
}
// -----
// Foo and bar are not valid layouts or meshes, respectively.
module attributes {dtensor.enable_multi_device_mode = true} {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<8xi32> {tf._global_shape = #tf_type.shape<8>, tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", tf._mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"}) -> (tensor<8xi32> {tf._global_shape = #tf_type.shape<8>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// expected-error @+1 {{Failed to retrieve op mesh or layout.}}
%1 = "tf.StatefulPartitionedCall"(%arg0, %arg1) {_layout = ["foo"], _mesh = "bar", config = "", config_proto = "", executor_type = "", f = @_test_func} : (tensor<i32>, tensor<8xi32>) -> tensor<8xi32>
return %1 : tensor<8xi32>
}
func.func private @_test_func(%arg0: tensor<i32>, %arg1: tensor<8xi32>) -> tensor<8xi32> {
return %arg1 : tensor<8xi32>
}
}
// -----
module attributes {dtensor.enable_multi_device_mode = true} {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<8xi32> {tf._global_shape = #tf_type.shape<8>, tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", tf._mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"}) -> (tensor<8xi32> {tf._global_shape = #tf_type.shape<8>}) attributes {tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = "op_output_0"}} {
// expected-error @+1 {{Calls must be used by exactly one return op.}}
%1 = "tf.StatefulPartitionedCall"(%arg0, %arg1) {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"], _mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", config = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7", config_proto = "", executor_type = "", f = @_test_func} : (tensor<i32>, tensor<8xi32>) -> tensor<8xi32>
%2 = "tf.Identity"(%1) : (tensor<8xi32>) -> tensor<8xi32>
return %2 : tensor<8xi32>
}
func.func private @_test_func(%arg0: tensor<i32>, %arg1: tensor<8xi32>) -> tensor<8xi32> {
return %arg1 : tensor<8xi32>
}
}
// -----
// CHECK-LABEL: module @test_inferred_resource_attributes
// CHECK-LABEL: func.func @main
// CHECK: "tf.StatefulPartitionedCall"
// CHECK-SAME: _inferred_resource_indices = dense<[1, 2]>
// CHECK-SAME: _inferred_resource_layouts = ["sharding_specs:x,unsharded
// CHECK-SAME , "sharding_specs:unsharded,y
module @test_inferred_resource_attributes attributes {dtensor.all_reduce_combiner.num_ops_in_group = 0 : i64, dtensor.all_reduce_combiner.topological_distance = 0 : i64, dtensor.eager_operation_name = "AssignVariableOp", dtensor.enable_multi_device_mode = true, tf._default_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:CPU:1"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1555 : i32}} {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<!tf_type.resource<tensor<i32>>> {tf._assigned_resource_local_shape = #tf_type.shape<>, tf._global_shape = #tf_type.shape<>, tf._layout = "empty_layout", tf._mesh = "empty_mesh"}, %arg2: tensor<!tf_type.resource<tensor<i32>>> {tf._assigned_resource_local_shape = #tf_type.shape<>, tf._global_shape = #tf_type.shape<>, tf._layout = "empty_layout", tf._mesh = "empty_mesh"}) attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0,op_input_1", outputs = ""}} {
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_inferred_resource_indices = dense<1> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"], _layout = [], _mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", config = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", config_proto = "", executor_type = "", f = @_func} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> ()
"tf.StatefulPartitionedCall"(%arg0, %arg2) {_inferred_resource_indices = dense<2> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:unsharded,y, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"], _layout = [], _mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", config = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", config_proto = "", executor_type = "", f = @_func} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> ()
return
}
func.func private @_func(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<i32>>>) {
"tf.AssignVariableOp"(%arg1, %arg0) {_global_shape = [], _layout = [], device = "", validate_shape = false} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
return
}
}
// -----
// Tests TPU expansion when the computation returns values.
// CHECK-LABEL: func.func @main
// CHECK-LABEL: func.func private @_func
// CHECK-SAME: %arg0: tensor<1x2xi32> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"}
// CHECK-SAME: %arg1: tensor<1x2xi32> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:1"}
// CHECK-SAME: -> (tensor<2xi32>, tensor<2xi32>) {
// CHECK-NEXT: %0:2 = "tf_device.launch"() <{device = ""}> ({
// CHECK-NEXT: %compilation_status, %program = "tf._TPUCompileMlir"() <{metadata = ""}> : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
// CHECK-NEXT: tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<3x!tf_type.string>
// CHECK-NEXT: }) : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
// CHECK-NEXT: "tf_device.launch"() <{device = ""}> ({
// CHECK-NEXT: "tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: }) : () -> ()
// CHECK-NEXT: %1:2 = "tf_device.parallel_execute"() ({
// CHECK-NEXT: %2 = "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:TPU:0"}> ({
// CHECK-NEXT: %3 = "tf.TPUExecute"(%arg0, %0#1) : (tensor<1x2xi32>, tensor<3x!tf_type.string>) -> tensor<2xi32>
// CHECK-NEXT: tf_device.return %3 : tensor<2xi32>
// CHECK-NEXT: }) : () -> tensor<2xi32>
// CHECK-NEXT: tf_device.return %2 : tensor<2xi32>
// CHECK-NEXT: }, {
// CHECK-NEXT: %2 = "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:TPU:1"}> ({
// CHECK-NEXT: %3 = "tf.TPUExecute"(%arg1, %0#1) : (tensor<1x2xi32>, tensor<3x!tf_type.string>) -> tensor<2xi32>
// CHECK-NEXT: tf_device.return %3 : tensor<2xi32>
// CHECK-NEXT: }) : () -> tensor<2xi32>
// CHECK-NEXT: tf_device.return %2 : tensor<2xi32>
// CHECK-NEXT: }) : () -> (tensor<2xi32>, tensor<2xi32>)
// CHECK-NEXT: return %1#0, %1#1 : tensor<2xi32>, tensor<2xi32>
// CHECK-NEXT: }
module attributes {dtensor.all_reduce_combiner.num_ops_in_group = 0 : i64, dtensor.all_reduce_combiner.topological_distance = 0 : i64, dtensor.eager_operation_name = "Sum", dtensor.enable_multi_device_mode = true, tf._default_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:CPU:1", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1625 : i32}} {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<1x2xi32> {tf._global_shape = #tf_type.shape<2x2>, tf._layout = "sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}, %arg2: tensor<1xi32> {tf._global_shape = #tf_type.shape<1>, tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"}) -> (tensor<2xi32> {tf._global_shape = #tf_type.shape<2>}) attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0,op_input_1", outputs = "op_output_0"}} {
%0 = "tf.StatefulPartitionedCall"(%arg1) {_layout = ["sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"], _mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func} : (tensor<1x2xi32>) -> tensor<2xi32>
return %0 : tensor<2xi32>
}
func.func private @_func(%arg0: tensor<1x2xi32>) -> tensor<2xi32> {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = ""} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<3x!tf_type.string>
}) {device = ""} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
%1 = "tf_device.launch"() ({
%2 = "tf.TPUExecute"(%arg0, %0#1) : (tensor<1x2xi32>, tensor<3x!tf_type.string>) -> tensor<2xi32>
tf_device.return %2 : tensor<2xi32>
}) {device = ""} : () -> tensor<2xi32>
return %1 : tensor<2xi32>
}
}
// -----
// Tests TPU expansion when the computation has variable-related operations and does not return values.
// CHECK-LABEL: func.func @main
// CHECK-LABEL: func.func private @_func
// CHECK-SAME: %arg0: tensor<i32> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"}
// CHECK-SAME: %arg1: tensor<i32> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:1"}
// CHECK-SAME: %arg2: tensor<!tf_type.resource<tensor<i32>>> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"}
// CHECK-SAME: %arg3: tensor<!tf_type.resource<tensor<i32>>> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:1"}
// CHECK-NEXT: %0:2 = "tf_device.launch"() <{device = ""}> ({
// CHECK-NEXT: %compilation_status, %program = "tf._TPUCompileMlir"() <{metadata = ""}> : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
// CHECK-NEXT: tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<3x!tf_type.string>
// CHECK-NEXT: }) : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
// CHECK-NEXT: "tf_device.launch"() <{device = ""}> ({
// CHECK-NEXT: "tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: }) : () -> ()
// CHECK-NEXT: %1:2 = "tf_device.parallel_execute"() ({
// CHECK-NEXT: %2 = "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:TPU:0"}> ({
// CHECK-NEXT: %3 = "tf.TPUExecute"(%arg0, %0#1) : (tensor<i32>, tensor<3x!tf_type.string>) -> tensor<i32>
// CHECK-NEXT: tf_device.return %3 : tensor<i32>
// CHECK-NEXT: }) : () -> tensor<i32>
// CHECK-NEXT: tf_device.return %2 : tensor<i32>
// CHECK-NEXT: }, {
// CHECK-NEXT: %2 = "tf_device.launch"() <{device = "/job:localhost/replica:0/task:0/device:TPU:1"}> ({
// CHECK-NEXT: %3 = "tf.TPUExecute"(%arg1, %0#1) : (tensor<i32>, tensor<3x!tf_type.string>) -> tensor<i32>
// CHECK-NEXT: tf_device.return %3 : tensor<i32>
// CHECK-NEXT: }) : () -> tensor<i32>
// CHECK-NEXT: tf_device.return %2 : tensor<i32>
// CHECK-NEXT: }) : () -> (tensor<i32>, tensor<i32>)
// CHECK-NEXT: "tf.AssignVariableOp"(%arg2, %1#0) <{validate_shape = false}> {_global_shape = [], _layout = [], device = "/job:localhost/replica:0/task:0/device:TPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
// CHECK-NEXT: "tf.AssignVariableOp"(%arg3, %1#1) <{validate_shape = false}> {_global_shape = [], _layout = [], device = "/job:localhost/replica:0/task:0/device:TPU:1"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
// CHECK-NEXT: return
// CHECK-NEXT: }
module attributes {dtensor.all_reduce_combiner.num_ops_in_group = 0 : i64, dtensor.all_reduce_combiner.topological_distance = 0 : i64, dtensor.eager_operation_name = "Sum", dtensor.enable_multi_device_mode = true, tf._default_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:CPU:1", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1625 : i32}} {
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>}, %arg1: tensor<!tf_type.resource<tensor<i32>>> {tf._assigned_resource_local_shape = #tf_type.shape<>, tf._global_shape = #tf_type.shape<>, tf._layout = "empty_layout", tf._mesh = "empty_mesh"}) -> () attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "eager_operation", inputs = "device_id,op_input_0", outputs = ""}} {
"tf.StatefulPartitionedCall"(%arg0, %arg1) {_inferred_resource_indices = dense<1> : vector<1xi32>, _inferred_resource_layouts = ["sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"], _layout = [], _mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", config_proto = "", executor_type = "", f = @_func} : (tensor<i32>, tensor<!tf_type.resource<tensor<i32>>>) -> ()
return
}
func.func private @_func(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource<tensor<i32>>>) -> () {
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = ""} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<3x!tf_type.string>
}) {device = ""} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = ""} : () -> ()
%1 = "tf_device.launch"() ({
%2 = "tf.TPUExecute"(%arg0, %0#1) : (tensor<i32>, tensor<3x!tf_type.string>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {device = ""} : () -> tensor<i32>
"tf.AssignVariableOp"(%arg1, %1) {_global_shape = [], _layout = [], validate_shape = false} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
return
}
}
tensorflow/dtensor/mlir/tests/op_to_device_cluster.mlir
+86
View File
@@ -0,0 +1,86 @@
// 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: dtensor-opt %s -split-input-file -dtensor-op-to-device-cluster -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @check_device_cluster_with_mesh_attribute
func.func @check_device_cluster_with_mesh_attribute() -> tensor<i32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.Const"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
%0 = "tf.Const"() {value = dense<10> : tensor<i32>, _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: func @check_device_cluster_from_op_mesh
func.func @check_device_cluster_from_op_mesh() -> tensor<i32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.Const"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
%0 = "tf.Const"() {value = dense<10> : tensor<i32>, _layout = ["sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"], _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: func @check_device_cluster_from_dtensor_layout_op
func.func @check_device_cluster_from_dtensor_layout_op(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.DTensorLayout"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
%0 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs:scalar |x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>} : (tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: func @check_device_cluster_from_copy_to_mesh_op
func.func @check_device_cluster_from_copy_to_mesh_op(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.Relayout"
// CHECK-NEXT: tf_device.return %[[A_OUT]]
// CHECK-NEXT: _mesh = "|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"
%0 = "tf.Relayout"(%arg0) { layout = "sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : (tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: func @check_yield_op_ignored
func.func @check_yield_op_ignored(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.WhileRegion"
// CHECK-NEXT: bb0(%arg1: tensor<i32>):
// CHECK-NEXT: %[[CLUSTER_OUT:.*]] = "tf_device.cluster"
// CHECK-NEXT: %[[H_OUT:.*]] = "tf.H"
// CHECK-NEXT: tf_device.return %[[H_OUT]]
// CHECK-NEXT: () -> tensor<i1>
// CHECK-NEXT: "tf.Yield"(%[[CLUSTER_OUT]])
// CHECK: ^bb0(%arg1: tensor<i32>):
// CHECK-NEXT: "tf.Yield"
%0 = "tf.WhileRegion"(%arg0) ({
^bb0(%arg1: tensor<i32>):
%1 = "tf.H"(%arg1) : (tensor<i32>) -> tensor<i1>
"tf.Yield"(%1) : (tensor<i1>) -> ()
}, {
^bb0(%arg1: tensor<i32>):
"tf.Yield"(%arg1) : (tensor<i32>) -> ()
}) { is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
func.return %0 : tensor<i32>
}
tensorflow/dtensor/mlir/tests/propagate_default_layout.mlir
+124
View File
@@ -0,0 +1,124 @@
// 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: dtensor-opt %s -split-input-file -dtensor-propagate-default-layout | FileCheck %s
// Check that layouts attributes in function arguments are converted to layout
// ops.
// CHECK-LABEL: module @test_convert_arguments
module @test_convert_arguments {
// CHECK: func.func @main
// CHECK-SAME: %arg[[ARG_0:[a-z0-9]*]]: tensor<i32>
// CHECK-SAME: %arg[[ARG_1:[a-z0-9]*]]: tensor<i32>
func.func @main(
%arg1: tensor<i32>{ tf._layout = "sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"},
%arg2: tensor<i32>{ tf._layout = "sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}) -> (tensor<i32>) {
// CHECK: %[[ARG1_OUT:[a-z0-9]*]] = "tf.DTensorLayout"(%arg[[ARG_1]])
// CHECK-SAME: layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
// CHECK-SAME: dtensor.from_arg_index = [[ARG_1]]
// CHECK-SAME: (tensor<i32>) -> tensor<i32>
// CHECK-NEXT: %[[ARG0_OUT:[a-z0-9]*]] = "tf.DTensorLayout"(%arg[[ARG_0]])
// CHECK-SAME: layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
// CHECK-SAME: dtensor.from_arg_index = [[ARG_0]]
// CHECK-SAME: (tensor<i32>) -> tensor<i32>
// CHECK-NEXT: "tf.A"(%[[ARG0_OUT]], %[[ARG1_OUT]])
// CHECK-NEXT: "tf.B"(%[[ARG1_OUT]])
// CHECK-NEXT: "tf.C"(%[[ARG0_OUT]])
%1 = "tf.A"(%arg1, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"tf.B"(%arg2) : (tensor<i32>) -> ()
"tf.C"(%arg1) : (tensor<i32>) -> ()
func.return %1 : tensor<i32>
}
}
// -----
// Check that layouts attributes in function outputs are converted to layout
// ops.
// CHECK-LABEL: module @test_convert_outputs
module @test_convert_outputs {
// CHECK-NEXT: func @main
func.func @main() -> (tensor<i32>{tf._default_layout = "sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}) {
// CHECK: %[[A_OUT:.*]] = "tf.A"() : () -> tensor<i32>
// CHECK-NEXT: %[[LAYOUT_A_OUT:.*]] = "tf.DTensorLayout"(%[[A_OUT]])
// CHECK-SAME: layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
// CHECK-SAME: (tensor<i32>) -> tensor<i32>
%1 = "tf.A"() : () -> tensor<i32>
func.return %1 : tensor<i32>
}
}
// -----
// Check that layouts attributes of operations are correclty converted to layout
// op.
// CHECK-LABEL: module @test_convert_operations
module @test_convert_operations {
// CHECK-NEXT: func @main
func.func @main() -> (tensor<i32>) {
// CHECK: %[[A_OUT:.*]] = "tf.A"()
// CHECK-NEXT: "tf.DTensorLayout"(%[[A_OUT]])
// CHECK-SAME: layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
%1 = "tf.A"() {_layout = ["sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> tensor<i32>
func.return %1 : tensor<i32>
}
}
// -----
// Check that when creating DTensorLayout for operations, dtensor.from_arg_index is preserved.
// CHECK-LABEL: module @test_preserve_from_arg_index
module @test_preserve_from_arg_index {
// CHECK-NEXT: func @main
func.func @main() -> (tensor<i32>) {
// CHECK: %[[A_OUT:.*]] = "tf.A"()
// CHECK-NEXT: "tf.DTensorLayout"(%[[A_OUT]])
// CHECK-SAME: dtensor.from_arg_index = 3
%1 = "tf.A"() {dtensor.from_arg_index = 3, _layout = ["sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> tensor<i32>
func.return %1 : tensor<i32>
}
}
// -----
// CHECK-LABEL: module @test_convert_resource_args
module @test_convert_resource_args {
// Check that resource typed arg with layouts are correctly converted to DTesnorLayout with global shape.
// CHECK-NEXT: func @main
func.func @main(%arg0: tensor<!tf_type.resource<tensor<4x2xf32>>>{ tf._layout = "sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}) -> tensor<i32> {
// CHECK: "tf.DTensorLayout"(%arg0)
// CHECK-SAME: global_shape = #tf_type.shape<4x2>
// CHECK-SAME: layout = #dtensor.layout<sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
%1 = "tf.A"() {_layout = ["sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> tensor<i32>
func.return %1 : tensor<i32>
}
}
// -----
// CHECK-LABEL: module @test_convert_variant_args
module @test_convert_variant_args {
// Check that variant typed arg with layouts are correctly converted to DTesnorLayout.
// CHECK-NEXT: func @main
func.func @main(%arg0: tensor<!tf_type.variant<tensor<4x4xi32>>>{ tf._layout = "sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}) -> tensor<i32> {
// CHECK: "tf.DTensorLayout"(%arg0)
// CHECK-SAME: layout = #dtensor.layout<sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
%1 = "tf.A"() {_layout = ["sharding_specs:scalar mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : () -> tensor<i32>
func.return %1 : tensor<i32>
}
}
tensorflow/dtensor/mlir/tests/propagate_device_id_to_function.mlir
+49
View File
@@ -0,0 +1,49 @@
// 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: dtensor-opt %s -split-input-file -dtensor-propagate-device-id-to-function-args | FileCheck %s
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG_0:[a-z0-9]+]]: tensor<i32>
// CHECK-SAME: %[[ARG_1:[a-z0-9]+]]: tensor<i32>
// CHECK-SAME: %[[ARG_2:[a-z0-9]+]]: tensor<i32>
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<i32>{ tf._layout = ["\0A\00\12\C6\01\0A\05\0A\01x\10\02\0A\05\0A\01y\10\02\12,/job:localhost/replica:0/task:0/device:CPU:0\12,/job:localhost/replica:0/task:0/device:CPU:1\12,/job:localhost/replica:0/task:0/device:CPU:2\12,/job:localhost/replica:0/task:0/device:CPU:3"]},
%arg2: tensor<i32>{ tf._layout = ["\0A\00\12\C6\01\0A\05\0A\01x\10\02\0A\05\0A\01y\10\02\12,/job:localhost/replica:0/task:0/device:CPU:0\12,/job:localhost/replica:0/task:0/device:CPU:1\12,/job:localhost/replica:0/task:0/device:CPU:2\12,/job:localhost/replica:0/task:0/device:CPU:3"]}) -> (tensor<i32>) {
// CHECK: "tf.StatefulPartitionedCall"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]])
%1 = "tf.StatefulPartitionedCall"(%arg1, %arg2) {f = @callee1, config = "", config_proto = "", executor_type = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
// CHECK-LABEL: func private @callee1
// CHECK-SAME: %[[CALL1_ARG0:[a-z0-9]+]]: tensor<i32>
// CHECK-SAME: %[[CALL1_ARG1:[a-z0-9]+]]: tensor<i32>
// CHECK-SAME: %[[CALL1_ARG2:[a-z0-9]+]]: tensor<i32>
func.func private @callee1(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> attributes {tf.signature.is_stateful} {
// CHECK: "tf.PartitionedCall"(%[[CALL1_ARG0]], %[[CALL1_ARG1]], %[[CALL1_ARG2]])
// CHECK-SAME: {config = "", config_proto = "", executor_type = "", f = @callee2}
// CHECK-SAME: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%0 = "tf.PartitionedCall"(%arg0, %arg1) {f = @callee2, config = "", config_proto = "", executor_type = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: func private @callee2
// CHECK-SAME: %[[CALL2_ARG0:[a-z0-9]+]]: tensor<i32>
// CHECK-SAME: %[[CALL2_ARG1:[a-z0-9]+]]: tensor<i32>
// CHECK-SAME: %[[CALL2_ARG2:[a-z0-9]+]]: tensor<i32>
func.func private @callee2(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> attributes {tf.signature.is_stateful} {
// CHECK: "tf.Add"(%[[CALL2_ARG1]], %[[CALL2_ARG2]])
%1 = "tf.Add"(%arg0, %arg1) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %1 : tensor<i32>
}
tensorflow/dtensor/mlir/tests/restore_and_assign.mlir
+128
View File
@@ -0,0 +1,128 @@
// 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: dtensor-opt -- %s -split-input-file -dtensor-infer-shapes-for-restorev2-op -dtensor-layout-propagation-v2 -verify-diagnostics | FileCheck %s
// Check the combination of inferring shape for restorev2 op and layout
// propagation. After running both passes, all unknown shapes from RestoreV2
// should be made known, and output layouts of RestoreV2 should match the
// resource tensors being assigned to.
// Single mesh
func.func @main(
%arg0: tensor<i32>,
%arg1: tensor<!tf_type.string>,
%arg2: tensor<!tf_type.string>,
%arg3: tensor<!tf_type.string>,
%arg4: tensor<*x!tf_type.resource<tensor<4x8xf32>>>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: %[[RESTORE:.*]] = "tf.RestoreV2"(%0, %1, %2) : (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> tensor<4x8xf32>
// CHECK-NEXT: %[[DLAYOUT:.*]] = "tf.DTensorLayout"(%[[RESTORE]]) <{global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>}> : (tensor<4x8xf32>) -> tensor<4x8xf32>
// CHECK-NEXT: "tf.AssignVariableOp"(%3, %[[DLAYOUT]]) <{validate_shape = true}> : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<4x8xf32>) -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%1 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%2 = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%3 = "tf.DTensorLayout"(%arg4) {global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>) -> tensor<*x!tf_type.resource<tensor<4x8xf32>>>
%4 = "tf.RestoreV2"(%0, %1, %2): (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> (tensor<*xf32>)
"tf.AssignVariableOp"(%3, %4) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<*xf32>) -> ()
tf_device.return
}) {_mesh="CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
// -----
// Single mesh with ops between the Restore and Assign.
func.func @main(
%arg0: tensor<i32>,
%arg1: tensor<!tf_type.string>,
%arg2: tensor<!tf_type.string>,
%arg3: tensor<!tf_type.string>,
%arg4: tensor<*x!tf_type.resource<tensor<4x8xf64>>>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: %[[RESTORE:.*]] = "tf.RestoreV2"(%0, %1, %2) : (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> tensor<4x8xf32>
// CHECK-NEXT: %[[DLAYOUT:.*]] = "tf.DTensorLayout"(%[[RESTORE]]) <{global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>}> : (tensor<4x8xf32>) -> tensor<4x8xf32>
// CHECK-NEXT: %[[CAST:.*]] = "tf.Cast"(%[[DLAYOUT]]) <{Truncate = false}> : (tensor<4x8xf32>) -> tensor<4x8xf64>
// CHECK-NEXT: %[[DLAYOUT2:.*]] = "tf.DTensorLayout"(%[[CAST]]) <{global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>}> : (tensor<4x8xf64>) -> tensor<4x8xf64>
// CHECK-NEXT: "tf.AssignVariableOp"(%3, %[[DLAYOUT2]]) <{validate_shape = true}> : (tensor<*x!tf_type.resource<tensor<4x8xf64>>>, tensor<4x8xf64>) -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%1 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%2 = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%3 = "tf.DTensorLayout"(%arg4) {global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<*x!tf_type.resource<tensor<4x8xf64>>>) -> tensor<*x!tf_type.resource<tensor<4x8xf64>>>
%4 = "tf.RestoreV2"(%0, %1, %2): (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> (tensor<*xf32>)
%5 = "tf.Cast"(%4) {} : (tensor<*xf32>) -> tensor<*xf64>
"tf.AssignVariableOp"(%3, %5) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<4x8xf64>>>, tensor<*xf64>) -> ()
tf_device.return
}) {_mesh="CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
// -----
// Cross mesh with send/recv
func.func @main(
%arg0: tensor<i32>,
%arg1: tensor<!tf_type.string>,
%arg2: tensor<!tf_type.string>,
%arg3: tensor<!tf_type.string>,
%arg4: tensor<*x!tf_type.resource<tensor<4x8xf32>>>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[RESOURCE:.*]] = "tf.DTensorLayout"(%arg4)
// CHECK-NEXT: %[[RECV:.*]] = "tf.DTensorRecv"() <{
// CHECK-SAME: key = "communication_key_|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1"
// CHECK-SAME: mesh = #dtensor.mesh<|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>
// CHECK-SAME: shape = #tf_type.shape<4x8>
// CHECK-SAME: source_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>
// CHECK-SAME: () -> tensor<4x8xf32>
// CHECK-NEXT: %[[RECV_DL:.*]] = "tf.DTensorLayout"(%[[RECV]])
// CHECK-NEXT: %[[IDENTITY:.*]] = "tf.Identity"(%[[RECV_DL]]) : (tensor<4x8xf32>) -> tensor<4x8xf32>
// CHECK-NEXT: %[[IDENTITY_DL:.*]] = "tf.DTensorLayout"(%[[IDENTITY]])
// CHECK-NEXT: "tf.AssignVariableOp"(%[[RESOURCE]], %[[IDENTITY_DL]]) <{validate_shape = true}> : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<4x8xf32>) -> ()
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
%4 = "tf.DTensorLayout"(%arg4) {global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1>} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>) -> tensor<*x!tf_type.resource<tensor<4x8xf32>>>
%5 = "tf.DTensorRecv"() {key = "communication_key_|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", mesh = #dtensor.mesh<|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>, shape = #tf_type.shape<*>} : () -> tensor<*xf32>
%6 = "tf.Identity"(%5) : (tensor<*xf32>) -> tensor<*xf32>
"tf.AssignVariableOp"(%4, %6) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<*xf32>) -> ()
tf_device.return
}) {_mesh="TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> (tensor<i32>, tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[DL1:.*]] = "tf.DTensorLayout"(%arg1)
// CHECK-NEXT: %[[DL2:.*]] = "tf.DTensorLayout"(%arg2)
// CHECK-NEXT: %[[DL3:.*]] = "tf.DTensorLayout"(%arg3)
// CHECK-NEXT: %[[RESTORE:.*]] = "tf.RestoreV2"(%[[DL1]], %[[DL2]], %[[DL3]]) : (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> tensor<4x8xf32>
// CHECK-NEXT: "tf.DTensorLayout"(%[[RESTORE]]) <{global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>}> : (tensor<4x8xf32>) -> tensor<4x8xf32>
// CHECK-NEXT: "tf.DTensorSend"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%1 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%2 = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%3 = "tf.RestoreV2"(%0, %1, %2) {} : (tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> (tensor<*xf32>)
"tf.DTensorSend"(%3) {key = "communication_key_|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1", target_mesh = #dtensor.mesh<|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1>} : (tensor<*xf32>) -> ()
tf_device.return
}) {_mesh="CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> (tensor<i32>)
func.return
}
tensorflow/dtensor/mlir/tests/restore_shape_inference.mlir
+76
View File
@@ -0,0 +1,76 @@
// 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: dtensor-opt -- %s -split-input-file -dtensor-infer-shapes-for-restorev2-op -verify-diagnostics | FileCheck %s
// Check the tf.RestoreV2Op's and all connected ops' resulting types are inferred from the AssignVariableOps in a single mesh. All unknown shapes should be known after this pass.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.string>, %arg2: tensor<2x!tf_type.string>, %arg3: tensor<2x!tf_type.string>, %arg4: tensor<*x!tf_type.resource<tensor<4x8xf32>>>, %arg5: tensor<*x!tf_type.resource<tensor<i64>>>) {
// CHECK: %0:2 = "tf.RestoreV2"(%arg1, %arg2, %arg3) : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<4x8xf32>, tensor<i64>)
// CHECK-NEXT: "tf.AssignVariableOp"(%arg4, %0#0) <{validate_shape = true}> : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<4x8xf32>) -> ()
// CHECK: %1 = "tf.Identity"(%0#1) : (tensor<i64>) -> tensor<i64>
// CHECK-NEXT: "tf.AssignVariableOp"(%arg5, %1) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
%0:2 = "tf.RestoreV2"(%arg1, %arg2, %arg3): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xi64>)
"tf.AssignVariableOp"(%arg4, %0#0) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<*xf32>) -> ()
%1 = "tf.Identity"(%0#1) {} : (tensor<*xi64>) -> tensor<*xi64>
"tf.AssignVariableOp"(%arg5, %1) {validate_shape = false} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<*xi64>) -> ()
func.return
}
// -----
// Check the tf.RestoreV2Op's and all connected ops' resulting types are inferred from the AssignVariableOps in cross mesh cluster. All unknown shapes should be known after this pass.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.string>, %arg2: tensor<2x!tf_type.string>, %arg3: tensor<2x!tf_type.string>, %arg4: tensor<*x!tf_type.resource<tensor<4x8xf32>>>, %arg5: tensor<*x!tf_type.resource<tensor<i64>>>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %2 = "tf.DTensorRecv"() <{key = "communication_key_|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3", mesh = #dtensor.mesh<|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>, shape = #tf_type.shape<4x8>}> : () -> tensor<4x8xf32>
// CHECK-NEXT: %3 = "tf.Identity"(%2) : (tensor<4x8xf32>) -> tensor<4x8xf32>
// CHECK-NEXT: "tf.AssignVariableOp"(%arg4, %3) <{validate_shape = true}> : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<4x8xf32>) -> ()
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
%1 = "tf.DTensorRecv"() {key = "communication_key_|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3", mesh = #dtensor.mesh<|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>, shape = #tf_type.shape<*>} : () -> tensor<*xf32>
%2 = "tf.Identity"(%1) : (tensor<*xf32>) -> tensor<*xf32>
"tf.AssignVariableOp"(%arg4, %2) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<*xf32>) -> ()
tf_device.return
}) {_mesh="TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>, tensor<i32>)
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %2:2 = "tf.RestoreV2"(%arg1, %arg2, %arg3) : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<4x8xf32>, tensor<i64>)
// CHECK-NEXT: %3 = "tf.Identity"(%2#1) : (tensor<i64>) -> tensor<i64>
// CHECK-NEXT: "tf.AssignVariableOp"(%arg5, %3) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
// CHECK-NEXT: "tf.DTensorSend"(%2#0) <{key = "communication_key_|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3", target_mesh = #dtensor.mesh<|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>}> : (tensor<4x8xf32>) -> ()
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
%6:2 = "tf.RestoreV2"(%arg1, %arg2, %arg3) {} : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xi64>)
%7 = "tf.Identity"(%6#1) : (tensor<*xi64>) -> tensor<*xi64>
"tf.AssignVariableOp"(%arg5, %7) {validate_shape = false} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<*xi64>) -> ()
"tf.DTensorSend"(%6#0) {key = "communication_key_|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3", target_mesh = #dtensor.mesh<|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>} : (tensor<*xf32>) -> ()
tf_device.return
}) {_mesh="CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Check correctness of shape inference and element type propagation of a graph containing tf.Cast ops.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.string>, %arg2: tensor<2x!tf_type.string>, %arg3: tensor<2x!tf_type.string>, %arg4: tensor<*x!tf_type.resource<tensor<4x8xf32>>>, %arg5: tensor<*x!tf_type.resource<tensor<f32>>>) {
// CHECK: %0:2 = "tf.RestoreV2"(%arg1, %arg2, %arg3) : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<4x8xf32>, tensor<bf16>)
// CHECK-NEXT: "tf.AssignVariableOp"(%arg4, %0#0) <{validate_shape = true}> : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<4x8xf32>) -> ()
// CHECK: %1 = "tf.Cast"(%0#1) <{Truncate = false}> : (tensor<bf16>) -> tensor<f32>
// CHECK-NEXT: "tf.AssignVariableOp"(%arg5, %1) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<f32>>>, tensor<f32>) -> ()
%0:2 = "tf.RestoreV2"(%arg1, %arg2, %arg3): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xbf16>)
"tf.AssignVariableOp"(%arg4, %0#0) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<4x8xf32>>>, tensor<*xf32>) -> ()
%1 = "tf.Cast"(%0#1) {} : (tensor<*xbf16>) -> tensor<*xf32>
"tf.AssignVariableOp"(%arg5, %1) {validate_shape = false} : (tensor<*x!tf_type.resource<tensor<f32>>>, tensor<*xf32>) -> ()
func.return
}
tensorflow/dtensor/mlir/tests/set_default_sharding.mlir
+44
View File
@@ -0,0 +1,44 @@
// 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: dtensor-opt %s -dtensor-set-default-sharding | FileCheck %s
// CHECK-LABEL: func @empty_func
func.func @empty_func(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>) {
func.return %arg0, %arg1 : tensor<?xi32>, tensor<?xi32>
}
// CHECK-LABEL: func @check_default_sharding_set
func.func @check_default_sharding_set() {
%0 = "tf.A"() : () -> tensor<?xi32>
%1 = "tf.B"() : () -> tensor<?xi32>
// CHECK: tf_device.cluster_func
// CHECK-SAME: _tpu_replicate = "cluster0"
// CHECK-SAME: input_sharding_configuration = ["\08\01\1A\01\01\22\01\00", "\08\01\1A\01\01\22\01\00"],
// CHECK-SAME: output_sharding_configuration = ["\08\01\1A\01\01\22\01\00", "\08\01\1A\01\01\22\01\00"]
%2, %3 = "tf_device.cluster_func"(%0, %1) {_tpu_replicate = "cluster0", func = @empty_func} : (tensor<?xi32>, tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)
func.return
}
// CHECK-LABEL: func @check_non_tpu_cluster_func_ignored
func.func @check_non_tpu_cluster_func_ignored() {
%0 = "tf.A"() : () -> tensor<?xi32>
%1 = "tf.B"() : () -> tensor<?xi32>
// CHECK: tf_device.cluster_func
// CHECK-NOT: _tpu_replicate = "cluster0"
// CHECK-NOT: input_sharding_configuration
// CHECK-NOT: output_sharding_configuration
%2, %3 = "tf_device.cluster_func"(%0, %1) {func = @empty_func} : (tensor<?xi32>, tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)
func.return
}
tensorflow/dtensor/mlir/tests/sparse_expansion.mlir
+71
View File
@@ -0,0 +1,71 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: dtensor-opt -- %s -split-input-file -dtensor-annotate-global-shape -dtensor-sparse-expansion -verify-diagnostics | FileCheck %s
// Check SparseExpansion for tf.MatMul with 1 SparseTensor left operand expands to SparseTensorDenseMatMul op.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<4x16xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}, %arg2: tensor<?x2xi64>, %arg3: tensor<2xi64>, %arg4: tensor<?xf32>) -> tensor<8x16xf32> {
// CHECK: "tf.Const"
// CHECK-NEXT: "tf.SparseToDense"
// CHECK-NEXT: %[[DENSE:.*]] = "tf.DTensorLayout"(%arg1)
// CHECK-NEXT: "tf.DTensorLayout"
// CHECK-NEXT: "tf.SparseTensorDenseMatMul"(%arg2, %arg4, %arg3, %[[DENSE]])
%cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%0 = "tf.SparseToDense"(%arg2, %arg3, %arg4, %cst) : (tensor<?x2xi64>, tensor<2xi64>, tensor<?xf32>, tensor<f32>) -> tensor<8x4xf32>
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x16>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>} : (tensor<4x16xf32>) -> tensor<4x16xf32>
%2 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<8x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>} : (tensor<8x4xf32>) -> tensor<8x4xf32>
%3 = "tf.MatMul"(%2, %1) {device = "", transpose_a = false, transpose_b = false} : (tensor<8x4xf32>, tensor<4x16xf32>) -> tensor<8x16xf32>
%4 = "tf.Identity"(%3) {device = ""} : (tensor<8x16xf32>) -> tensor<8x16xf32>
%5 = "tf.Identity"(%2) {device = ""} : (tensor<8x4xf32>) -> tensor<8x4xf32>
func.return %4 : tensor<8x16xf32>
}
// -----
// Check that after SparseExpansion, unused SparseToDense ops are removed.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<4x16xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}, %arg2: tensor<?x2xi64>, %arg3: tensor<2xi64>, %arg4: tensor<?xf32>) -> tensor<8x16xf32> {
// CHECK: %[[DENSE:.*]] = "tf.DTensorLayout"(%arg1)
// CHECK-NEXT: "tf.SparseTensorDenseMatMul"(%arg2, %arg4, %arg3, %[[DENSE]])
%cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%0 = "tf.SparseToDense"(%arg2, %arg3, %arg4, %cst) : (tensor<?x2xi64>, tensor<2xi64>, tensor<?xf32>, tensor<f32>) -> tensor<8x4xf32>
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x16>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>} : (tensor<4x16xf32>) -> tensor<4x16xf32>
%2 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<8x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>} : (tensor<8x4xf32>) -> tensor<8x4xf32>
%3 = "tf.MatMul"(%2, %1) {device = "", transpose_a = false, transpose_b = false} : (tensor<8x4xf32>, tensor<4x16xf32>) -> tensor<8x16xf32>
%4 = "tf.Identity"(%3) {device = ""} : (tensor<8x16xf32>) -> tensor<8x16xf32>
func.return %4 : tensor<8x16xf32>
}
// -----
// Check SparseExpansion for tf.MatMul with 2 SparseTensor operands is a no-change. That is, the original op gets returned.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<?x2xi64>, %arg2: tensor<2xi64>, %arg3: tensor<?xf32>, %arg4: tensor<?x2xi64>, %arg5: tensor<2xi64>, %arg6: tensor<?xf32>) -> tensor<8x16xf32> {
// CHECK: "tf.Const"
// CHECK-NEXT: "tf.SparseToDense"
// CHECK: "tf.Const"
// CHECK: "tf.SparseToDense"
// CHECK: "tf.DTensorLayout"
// CHECK: "tf.DTensorLayout"
// CHECK: "tf.MatMul"
// CHECK: "tf.Identity"
%cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%0 = "tf.SparseToDense"(%arg4, %arg5, %arg6, %cst) : (tensor<?x2xi64>, tensor<2xi64>, tensor<?xf32>, tensor<f32>) -> tensor<4x16xf32>
%cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%1 = "tf.SparseToDense"(%arg1, %arg2, %arg3, %cst_0) : (tensor<?x2xi64>, tensor<2xi64>, tensor<?xf32>, tensor<f32>) -> tensor<8x4xf32>
%2 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<4x16>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|batch=2,x=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<4x16xf32>) -> tensor<4x16xf32>
%3 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<8x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|batch=2,x=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<8x4xf32>) -> tensor<8x4xf32>
%4 = "tf.MatMul"(%3, %2) {device = "", transpose_a = false, transpose_b = false} : (tensor<8x4xf32>, tensor<4x16xf32>) -> tensor<8x16xf32>
%5 = "tf.Identity"(%4) {device = ""} : (tensor<8x16xf32>) -> tensor<8x16xf32>
func.return %5 : tensor<8x16xf32>
}
tensorflow/dtensor/mlir/tests/spmd_batchparallel.mlir
+52
View File
@@ -0,0 +1,52 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s --dump-input=fail
// Test replicated layout.
func.func @main(%arg0: tensor<1xf32>,
%arg1: tensor<8x128x128x3xf32> {tf._layout = "sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<8x128x128x3xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.AdjustSaturation"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<1xf32>) -> tensor<1xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x128x128x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128x3xf32>) -> tensor<8x128x128x3xf32>
%3 = "tf.AdjustSaturation"(%2, %1) {} : (tensor<8x128x128x3xf32>, tensor<1xf32>) -> tensor<8x128x128x3xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<8x128x128x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128x3xf32>) -> tensor<8x128x128x3xf32>
tf_device.return %4 : tensor<8x128x128x3xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<8x128x128x3xf32>
func.return %0 : tensor<8x128x128x3xf32>
}
// -----
// Test batch sharded layout. Should emit Identity op.
func.func @main(%arg0: tensor<1xf32>,
%arg1: tensor<8x128x128x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<8x128x128x3xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.AdjustSaturation"
// CHECK-NEXT: "tf.IdentityN"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<1xf32>) -> tensor<1xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x128x128x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128x3xf32>) -> tensor<8x128x128x3xf32>
%3 = "tf.AdjustSaturation"(%2, %1) {} : (tensor<8x128x128x3xf32>, tensor<1xf32>) -> tensor<8x128x128x3xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<8x128x128x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128x3xf32>) -> tensor<8x128x128x3xf32>
tf_device.return %4 : tensor<8x128x128x3xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<8x128x128x3xf32>
func.return %0 : tensor<8x128x128x3xf32>
}
tensorflow/dtensor/mlir/tests/spmd_concat.mlir
+204
View File
@@ -0,0 +1,204 @@
// 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: dtensor-opt -- %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check that a dim sharded on all Concat inputs (which is not the concat dim)
// produces output layout with the same dim sharded.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<2x16x32xf32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<6x16x32xf32>
func.func @main(%arg0: tensor<2x32x32xf32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|*TPU"},
%arg1: tensor<6x32x32xf32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"()
// CHECK-NEXT: %[[CONCAT_OUT:.*]] = "tf.ConcatV2"(%[[ARG0]], %[[ARG1]], %[[AXIS]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<2x16x32xf32>, tensor<6x16x32xf32>, tensor<i32>) -> tensor<8x16x32xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<8x16x32xf32>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<2x32x32>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<2x32x32xf32>) -> tensor<2x32x32xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x32x32>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<6x32x32xf32>) -> tensor<6x32x32xf32>
%3 = "tf.ConcatV2"(%1, %2, %cst) : (tensor<2x32x32xf32>, tensor<6x32x32xf32>, tensor<i32>) -> tensor<8x32x32xf32>
%4 = "tf.Identity"(%3) : (tensor<8x32x32xf32>) -> tensor<8x32x32xf32>
tf_device.return %4 : tensor<8x32x32xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x32x32xf32>)
func.return
}
// -----
// Check that if the concat dim is sharded in any Concat inputs, then that dim
// is replicated in the output layout.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<8x4x32xf32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<8x2x32xf32>
// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: tensor<8x4x32xf32>
func.func @main(%arg0: tensor<8x4x32xf32> {tf._layout = "sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|*TPU"},
%arg1: tensor<8x8x32xf32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=4|*TPU"},
%arg2: tensor<8x16x32xf32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=4|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"()
// CHECK-NEXT: %[[ARG1_RELAYOUT:.*]] = "tf.DTensorAllGather"(%[[ARG1]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x2x32xf32>) -> tensor<8x8x32xf32>
// CHECK-NEXT: %[[ARG2_RELAYOUT:.*]] = "tf.DTensorAllGather"(%[[ARG2]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x4x32xf32>) -> tensor<8x16x32xf32>
// CHECK-NEXT: %[[CONCAT_OUT:.*]] = "tf.ConcatV2"(%[[ARG0]], %[[ARG1_RELAYOUT]], %[[ARG2_RELAYOUT]], %[[AXIS]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x4x32xf32>, tensor<8x8x32xf32>, tensor<8x16x32xf32>, tensor<i32>) -> tensor<8x28x32xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<8x28x32xf32>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x4x32>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:|x=4|*TPU>} : (tensor<8x4x32xf32>) -> tensor<8x4x32xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x32>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=4|*TPU>} : (tensor<8x8x32xf32>) -> tensor<8x8x32xf32>
%3 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x16x32>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=4|*TPU>} : (tensor<8x16x32xf32>) -> tensor<8x16x32xf32>
%4 = "tf.ConcatV2"(%1, %2, %3, %cst) : (tensor<8x4x32xf32>, tensor<8x8x32xf32>, tensor<8x16x32xf32>, tensor<i32>) -> tensor<8x28x32xf32>
%5 = "tf.Identity"(%4) : (tensor<8x28x32xf32>) -> tensor<8x28x32xf32>
tf_device.return %5 : tensor<8x28x32xf32>
}) {_mesh="|x=4|*TPU"} : () -> (tensor<8x28x32xf32>)
func.return
}
// -----
// Check that dims sharded on any Concat inputs (which is not the concat dim,
// and does not conflict with any other sharding) produces output layout with
// the same dims sharded.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<4x4x32xf32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<8x8x16xf32>
// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: tensor<4x16x16xf32>
func.func @main(%arg0: tensor<8x4x32xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%arg1: tensor<8x8x32xf32> {tf._layout = "sharding_specs:unsharded,unsharded,y, mesh:|x=2,y=2|*TPU"},
%arg2: tensor<8x16x32xf32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:|x=2,y=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"()
// CHECK-NEXT: %[[ARG0_RELAYOUT:.*]] = "tf.DTensorAllScatter"(%[[ARG0]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<4x4x32xf32>) -> tensor<4x4x16xf32>
// CHECK-NEXT: %[[ARG1_RELAYOUT:.*]] = "tf.DTensorAllScatter"(%[[ARG1]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x8x16xf32>) -> tensor<4x8x16xf32>
// CHECK-NEXT: %[[CONCAT_OUT:.*]] = "tf.ConcatV2"(%[[ARG0_RELAYOUT]], %[[ARG1_RELAYOUT]], %[[ARG2]], %[[AXIS]])
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<4x4x16xf32>, tensor<4x8x16xf32>, tensor<4x16x16xf32>, tensor<i32>) -> tensor<4x28x16xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<4x28x16xf32>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x4x32>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x4x32xf32>) -> tensor<8x4x32xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x32>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,y, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x32xf32>) -> tensor<8x8x32xf32>
%3 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x16x32>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:|x=2,y=2|*TPU>} : (tensor<8x16x32xf32>) -> tensor<8x16x32xf32>
%4 = "tf.ConcatV2"(%1, %2, %3, %cst) : (tensor<8x4x32xf32>, tensor<8x8x32xf32>, tensor<8x16x32xf32>, tensor<i32>) -> tensor<8x28x32xf32>
%5 = "tf.Identity"(%4) : (tensor<8x28x32xf32>) -> tensor<8x28x32xf32>
tf_device.return %5 : tensor<8x28x32xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x28x32xf32>)
func.return
}
// -----
// Check that any dims with conflicting sharding across the Concat inputs are
// deduplicated and the output layout is replicated in those dims.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<8x4x32xf32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<16x8x16xf32>
func.func @main(%arg0: tensor<8x8x32xf32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|*TPU"},
%arg1: tensor<16x8x32xf32> {tf._layout = "sharding_specs:unsharded,unsharded,x, mesh:|x=2,y=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"()
// CHECK-NEXT: %[[ARG0_RELAYOUT:.*]] = "tf.DTensorAllGather"(%[[ARG0]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x4x32xf32>) -> tensor<8x8x32xf32>
// CHECK-NEXT: %[[ARG1_RELAYOUT:.*]] = "tf.DTensorAllGather"(%[[ARG1]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded,x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<16x8x16xf32>) -> tensor<16x8x32xf32>
// CHECK-NEXT: %[[CONCAT_OUT:.*]] = "tf.ConcatV2"(%[[ARG0_RELAYOUT]], %[[ARG1_RELAYOUT]], %[[AXIS]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x8x32xf32>, tensor<16x8x32xf32>, tensor<i32>) -> tensor<24x8x32xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<24x8x32xf32>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8x32>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x32xf32>) -> tensor<8x8x32xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<16x8x32>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,x, mesh:|x=2,y=2|*TPU>} : (tensor<16x8x32xf32>) -> tensor<16x8x32xf32>
%3 = "tf.ConcatV2"(%1, %2, %cst) : (tensor<8x8x32xf32>, tensor<16x8x32xf32>, tensor<i32>) -> tensor<24x8x32xf32>
%4 = "tf.Identity"(%3) : (tensor<24x8x32xf32>) -> tensor<24x8x32xf32>
tf_device.return %4 : tensor<24x8x32xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<24x8x32xf32>)
func.return
}
// -----
// Check that if any Concat input is sharded on the concat dim, along with other
// inputs sharded on other dims, then relayout is correctly applied to those
// inputs and the order of DTensorAllScatter -> DTensorAllGather is correct.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<4x4x32xf32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<8x4x32xf32>
// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: tensor<4x8x32xf32>
func.func @main(%arg0: tensor<8x4x32xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%arg1: tensor<8x8x32xf32> {tf._layout = "sharding_specs:unsharded,y,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2: tensor<8x16x32xf32> {tf._layout = "sharding_specs:x,y,unsharded, mesh:|x=2,y=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"()
// CHECK-NEXT: %[[ARG1_SCATTER:.*]] = "tf.DTensorAllScatter"(%[[ARG1]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,y,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,y,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:x,y,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<8x4x32xf32>) -> tensor<4x4x32xf32>
// CHECK-NEXT: %[[ARG1_RELAYOUT:.*]] = "tf.DTensorAllGather"(%[[ARG1_SCATTER]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:x,y,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<4x4x32xf32>) -> tensor<4x8x32xf32>
// CHECK-NEXT: %[[ARG2_RELAYOUT:.*]] = "tf.DTensorAllGather"(%[[ARG2]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:x,y,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<4x8x32xf32>) -> tensor<4x16x32xf32>
// CHECK-NEXT: %[[CONCAT_OUT:.*]] = "tf.ConcatV2"(%[[ARG0]], %[[ARG1_RELAYOUT]], %[[ARG2_RELAYOUT]], %[[AXIS]])
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<4x4x32xf32>, tensor<4x8x32xf32>, tensor<4x16x32xf32>, tensor<i32>) -> tensor<4x28x32xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<4x28x32xf32>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x4x32>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x4x32xf32>) -> tensor<8x4x32xf32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x32>, layout = #dtensor.layout<sharding_specs:unsharded,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x32xf32>) -> tensor<8x8x32xf32>
%3 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x16x32>, layout = #dtensor.layout<sharding_specs:x,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x16x32xf32>) -> tensor<8x16x32xf32>
%4 = "tf.ConcatV2"(%1, %2, %3, %cst) : (tensor<8x4x32xf32>, tensor<8x8x32xf32>, tensor<8x16x32xf32>, tensor<i32>) -> tensor<8x28x32xf32>
%5 = "tf.Identity"(%4) : (tensor<8x28x32xf32>) -> tensor<8x28x32xf32>
tf_device.return %5 : tensor<8x28x32xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x28x32xf32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_conv.mlir
+715
View File
@@ -0,0 +1,715 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check that Conv2D uses input image layout as output layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"},
%arg2:tensor<8x3x3x3xf32>) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv2D"
// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]
// CHECK-SAME: (tensor<4x32x32x3xf32>, tensor<8x3x3x3xf32>) -> tensor<4x32x32x3xf32>
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU>} : (tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU>} : (tensor<8x3x3x3xf32>) -> tensor<8x3x3x3xf32>
%conv = "tf.Conv2D"(%img_layout, %filter_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} :
(tensor<8x32x32x3xf32>, tensor<8x3x3x3xf32>) -> tensor<8x32x32x3xf32>
tf_device.return %conv : tensor<8x32x32x3xf32>
}) {_mesh="|x=2,y=1|*TPU"} : () -> (tensor<8x32x32x3xf32>)
func.return
}
// -----
// Check that Conv3D uses input image layout as output layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"},
%arg2:tensor<8x3x3x3x3xf32>) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv3D"
// CHECK-SAME: data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]
// CHECK-SAME: (tensor<4x32x32x32x3xf32>, tensor<8x3x3x3x3xf32>) -> tensor<4x32x32x32x3xf32>
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU>} : (tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU>} : (tensor<8x3x3x3x3xf32>) -> tensor<8x3x3x3x3xf32>
%conv = "tf.Conv3D"(%img_layout, %filter_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} :
(tensor<8x32x32x32x3xf32>, tensor<8x3x3x3x3xf32>) -> tensor<8x32x32x32x3xf32>
tf_device.return %conv : tensor<8x32x32x32x3xf32>
}) {_mesh="|x=2,y=1|*TPU"} : () -> (tensor<8x32x32x32x3xf32>)
func.return
}
// -----
// Check that Conv2D backprop uses grads as output layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<1x3x3x3xf32>,
%arg2: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv2DBackpropInput"
// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]
// CHECK-SAME: (tensor<4xi32>, tensor<1x3x3x3xf32>, tensor<4x32x32x3xf32>) -> tensor<4x32x32x3xf32>
%0 = "tf_device.cluster"() ({
%img_shape = "tf.Const"() { value=dense<[8,32,32,3]> : tensor<4xi32>} : () -> tensor<4xi32>
%filter_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<1x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3xf32>) -> tensor<1x3x3x3xf32>
%grad_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%conv = "tf.Conv2DBackpropInput"(%img_shape, %filter_layout, %grad_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} :
(tensor<4xi32>, tensor<1x3x3x3xf32>, tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
tf_device.return %conv : tensor<8x32x32x3xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"}: () -> (tensor<8x32x32x3xf32>)
func.return
}
// -----
// Check that Conv3DBackPropInputV2 uses grads as output layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<1x3x3x3x3xf32>,
%arg2: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv3DBackpropInputV2"
// CHECK-SAME: data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]
// CHECK-SAME: (tensor<5xi32>, tensor<1x3x3x3x3xf32>, tensor<4x32x32x32x3xf32>) -> tensor<4x32x32x32x3xf32>
%0 = "tf_device.cluster"() ({
%img_shape = "tf.Const"() { value=dense<[8,32,32,32,3]> : tensor<5xi32>} : () -> tensor<5xi32>
%filter_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<1x3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3x3xf32>) -> tensor<1x3x3x3x3xf32>
%grad_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%conv = "tf.Conv3DBackpropInputV2"(%img_shape, %filter_layout, %grad_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} :
(tensor<5xi32>, tensor<1x3x3x3x3xf32>, tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
tf_device.return %conv : tensor<8x32x32x32x3xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"}: () -> (tensor<8x32x32x32x3xf32>)
func.return
}
// -----
// Check all reduce emitted in Conv2DBackpropFilter when image is batch sharded
// but input_shape is replicated const.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%input_img: tensor<2x9x9x1xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%grad: tensor<2x9x9x2xf32> {tf._layout = "sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU"}
) -> tensor<2x2x1x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[FILTER_SHAPE:.*]] = "tf.Const"()
// CHECK: %[[SLICE_OUT:.*]] = "tf.DTensorAllScatter"(%arg2)
// CHECK: %[[BACKPROP_OUT:.*]] = "tf.Conv2DBackpropFilter"(%arg1, %[[FILTER_SHAPE]], %[[SLICE_OUT]])
// CHECK: %[[TPU_GROUP:.*]] = "tf.Const"()
// CHECK: %[[XLA_ALL_REDUCE:.*]] = "tf.DTensorAllReduce"(%[[BACKPROP_OUT]], %[[TPU_GROUP]])
%0 = "tf_device.cluster"() ({
%filter_shape = "tf.Const"() { value = dense<[2, 2, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
%input_layout = "tf.DTensorLayout"(%input_img) {global_shape = #tf_type.shape<2x9x9x1>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<2x9x9x1xf32>) -> tensor<2x9x9x1xf32>
%grad_layout = "tf.DTensorLayout"(%grad) {global_shape = #tf_type.shape<2x9x9x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<2x9x9x2xf32>) -> tensor<2x9x9x2xf32>
%2 = "tf.Conv2DBackpropFilter"(%input_layout, %filter_shape, %grad_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<2x9x9x1xf32>, tensor<4xi32>, tensor<2x9x9x2xf32>) -> tensor<2x2x1x2xf32>
tf_device.return %2 : tensor<2x2x1x2xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"} : () -> tensor<2x2x1x2xf32>
func.return %0 : tensor<2x2x1x2xf32>
}
// -----
// Check all reduce emitted in Conv3DBackpropFilterV2 when image is batch
// sharded but input_shape is replicated const.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%input_img: tensor<2x9x9x9x1xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%grad: tensor<2x9x9x9x2xf32> {tf._layout = "sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU"}
) -> tensor<2x2x2x1x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[FILTER_SHAPE:.*]] = "tf.Const"()
// CHECK: %[[SLICE_OUT:.*]] = "tf.DTensorAllScatter"(%arg2)
// CHECK: %[[BACKPROP_OUT:.*]] = "tf.Conv3DBackpropFilterV2"(%arg1, %[[FILTER_SHAPE]], %[[SLICE_OUT]])
// CHECK: %[[TPU_GROUP:.*]] = "tf.Const"()
// CHECK: %[[XLA_ALL_REDUCE:.*]] = "tf.DTensorAllReduce"(%[[BACKPROP_OUT]], %[[TPU_GROUP]])
%0 = "tf_device.cluster"() ({
%filter_shape = "tf.Const"() { value = dense<[2, 2, 2, 1, 2]> : tensor<5xi32>} : () -> tensor<5xi32>
%input_layout = "tf.DTensorLayout"(%input_img) {global_shape = #tf_type.shape<2x9x9x9x1>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<2x9x9x9x1xf32>) -> tensor<2x9x9x9x1xf32>
%grad_layout = "tf.DTensorLayout"(%grad) {global_shape = #tf_type.shape<2x9x9x9x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<2x9x9x9x2xf32>) -> tensor<2x9x9x9x2xf32>
%2 = "tf.Conv3DBackpropFilterV2"(%input_layout, %filter_shape, %grad_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1, 1]} : (tensor<2x9x9x9x1xf32>, tensor<5xi32>, tensor<2x9x9x9x2xf32>) -> tensor<2x2x2x1x2xf32>
tf_device.return %2 : tensor<2x2x2x1x2xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"} : () -> tensor<2x2x2x1x2xf32>
func.return %0 : tensor<2x2x2x1x2xf32>
}
// -----
// Check that Conv2D with spatial partitioning requires kernel to be fully
// replicated.
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:unsharded,x,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2:tensor<8x3x3x3xf32>) {
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x3x3x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x3x3x3xf32>) -> tensor<8x3x3x3xf32>
// expected-error @+1 {{Filter for convolution must have fully replicated layout.}}
%conv = "tf.Conv2D"(%img_layout, %filter_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<8x32x32x3xf32>, tensor<8x3x3x3xf32>) -> tensor<8x32x32x3xf32>
tf_device.return %conv : tensor<8x32x32x3xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x32x32x3xf32>)
func.return
}
// -----
// Check that Conv3D with spatial partitioning requires kernel to be fully
// replicated.
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:unsharded,x,y,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2:tensor<8x3x3x3x3xf32>) {
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x3x3x3x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x3x3x3x3xf32>) -> tensor<8x3x3x3x3xf32>
// expected-error @+1 {{Filter for convolution must have fully replicated layout.}}
%conv = "tf.Conv3D"(%img_layout, %filter_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} : (tensor<8x32x32x32x3xf32>, tensor<8x3x3x3x3xf32>) -> tensor<8x32x32x32x3xf32>
tf_device.return %conv : tensor<8x32x32x32x3xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x32x32x32x3xf32>)
func.return
}
// -----
// Check that Conv2D with spatial partitioning requires input dimension size to
// be greater than halo size.
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x8x8x3xf32> {tf._layout = "sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2:tensor<15x15x3x3xf32>) {
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x8x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x8x3xf32>) -> tensor<8x8x8x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<15x15x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<15x15x3x3xf32>) -> tensor<15x15x3x3xf32>
// expected-error @+1 {{input shard tensor size of each processor must be greater than halo size}}
%conv = "tf.Conv2D"(%img_layout, %filter_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<8x8x8x3xf32>, tensor<15x15x3x3xf32>) -> tensor<8x8x8x3xf32>
tf_device.return %conv : tensor<8x8x8x3xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x8x8x3xf32>)
func.return
}
// -----
// Check that Conv3D with spatial partitioning requires input dimension size to
// be greater than halo size.
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x8x8x8x3xf32> {tf._layout = "sharding_specs:unsharded,unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2:tensor<15x3x15x3x3xf32>) {
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x8x8x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x8x8x3xf32>) -> tensor<8x8x8x8x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<15x3x15x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<15x3x15x3x3xf32>) -> tensor<15x3x15x3x3xf32>
// expected-error @+1 {{input shard tensor size of each processor must be greater than halo size}}
%conv = "tf.Conv3D"(%img_layout, %filter_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} : (tensor<8x8x8x8x3xf32>, tensor<15x3x15x3x3xf32>) -> tensor<8x8x8x8x3xf32>
tf_device.return %conv : tensor<8x8x8x8x3xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x8x8x8x3xf32>)
func.return
}
// -----
// Check that Conv2D with spatial partitioning using "SAME" padding produces
// begin and end halos on both spatial dimensions.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x8x8x3xf32> {tf._layout = "sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2:tensor<3x3x3x3xf32>) {
// CHECK: "tf_device.cluster"
// Build left halo on height dim.
// CHECK: %[[SLICE_H_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 3, 0, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT:.*]] = "tf.Slice"(%arg1, %[[SLICE_H_LEFT_BEGIN]], %[[SLICE_H_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_H_LEFT]], %[[PAIRS_H_LEFT]])
// CHECK-SAME: (tensor<8x1x4x3xf32>, tensor<4x2xi32>) -> tensor<8x1x4x3xf32>
// Build right halo on height dim.
// CHECK: %[[SLICE_H_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT:.*]] = "tf.Slice"(%arg1, %[[SLICE_H_RIGHT_BEGIN]], %[[SLICE_H_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_H_RIGHT]], %[[PAIRS_H_RIGHT]])
// CHECK-SAME: (tensor<8x1x4x3xf32>, tensor<4x2xi32>) -> tensor<8x1x4x3xf32>
// Concat the halos with the shard on the height dim.
// CHECK-NEXT: %[[CONCAT_H_AXIS:.*]] = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_H_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_H_LEFT]], %arg1, %[[EXCHANGED_HALO_H_RIGHT]], %[[CONCAT_H_AXIS]])
// CHECK-SAME: (tensor<8x1x4x3xf32>, tensor<8x4x4x3xf32>, tensor<8x1x4x3xf32>, tensor<i64>) -> tensor<8x6x4x3xf32>
// Build left halo on width dim.
// CHECK: %[[SLICE_W_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 0, 3, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 6, 1, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT:.*]] = "tf.Slice"(%[[CONCAT_H_TENSOR]], %[[SLICE_W_LEFT_BEGIN]], %[[SLICE_W_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x6x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x6x1x3xf32>
// CHECK-NEXT: %[[HALO_W_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_W_LEFT]], %[[PAIRS_W_LEFT]])
// CHECK-SAME: (tensor<8x6x1x3xf32>, tensor<4x2xi32>) -> tensor<8x6x1x3xf32>
// Build right halo on width dim.
// CHECK: %[[SLICE_W_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 6, 1, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT:.*]] = "tf.Slice"(%[[CONCAT_H_TENSOR]], %[[SLICE_W_RIGHT_BEGIN]], %[[SLICE_W_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x6x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x6x1x3xf32>
// CHECK-NEXT: %[[HALO_W_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_W_RIGHT]], %[[PAIRS_W_RIGHT]])
// CHECK-SAME: (tensor<8x6x1x3xf32>, tensor<4x2xi32>) -> tensor<8x6x1x3xf32>
// Concat the halos with the shard on the width dim.
// CHECK-NEXT: %[[CONCAT_W_AXIS:.*]] = "tf.Const"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_HW_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_W_LEFT]], %[[CONCAT_H_TENSOR]], %[[EXCHANGED_HALO_W_RIGHT]], %[[CONCAT_W_AXIS]])
// CHECK-SAME: (tensor<8x6x1x3xf32>, tensor<8x6x4x3xf32>, tensor<8x6x1x3xf32>, tensor<i64>) -> tensor<8x6x6x3xf32>
// CHECK-NEXT: "tf.Conv2D"(%[[CONCAT_HW_TENSOR]], %arg2)
// CHECK-SAME: padding = "VALID"
// CHECK-SAME: (tensor<8x6x6x3xf32>, tensor<3x3x3x3xf32>) -> tensor<8x4x4x3xf32>
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x8x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x8x3xf32>) -> tensor<8x8x8x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3xf32>
%conv = "tf.Conv2D"(%img_layout, %filter_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<8x8x8x3xf32>, tensor<3x3x3x3xf32>) -> tensor<8x8x8x3xf32>
tf_device.return %conv : tensor<8x8x8x3xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x8x8x3xf32>)
func.return
}
// -----
// Check that Conv3D with spatial partitioning using "SAME" padding produces
// begin and end halos on all spatial dimensions.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x8x8x8x3xf32> {tf._layout = "sharding_specs:unsharded,x,y,z,unsharded, mesh:|x=2,y=2,z=2|*TPU"},
%arg2:tensor<3x3x3x3x3xf32>) {
// CHECK: "tf_device.cluster"
// Build left halo on depth dim.
// CHECK: %[[SLICE_D_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 3, 0, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_LEFT:.*]] = "tf.Slice"(%arg1, %[[SLICE_D_LEFT_BEGIN]], %[[SLICE_D_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x1x4x4x3xf32>
// CHECK-NEXT: %[[HALO_D_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_D_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 4], [1, 5], [2, 6], [3, 7], [4, 0], [5, 1], [6, 2], [7, 3]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_D_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_D_LEFT]], %[[PAIRS_D_LEFT]])
// CHECK-SAME: (tensor<8x1x4x4x3xf32>, tensor<8x2xi32>) -> tensor<8x1x4x4x3xf32>
// Build right halo on depth dim.
// CHECK: %[[SLICE_D_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_RIGHT:.*]] = "tf.Slice"(%arg1, %[[SLICE_D_RIGHT_BEGIN]], %[[SLICE_D_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x1x4x4x3xf32>
// CHECK-NEXT: %[[HALO_D_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_D_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 4], [1, 5], [2, 6], [3, 7], [4, 0], [5, 1], [6, 2], [7, 3]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_D_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_D_RIGHT]], %[[PAIRS_D_RIGHT]])
// CHECK-SAME: (tensor<8x1x4x4x3xf32>, tensor<8x2xi32>) -> tensor<8x1x4x4x3xf32>
// Concat the halos with the shard on the depth dim.
// CHECK-NEXT: %[[CONCAT_D_AXIS:.*]] = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_D_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_D_LEFT]], %arg1, %[[EXCHANGED_HALO_D_RIGHT]], %[[CONCAT_D_AXIS]])
// CHECK-SAME: (tensor<8x1x4x4x3xf32>, tensor<8x4x4x4x3xf32>, tensor<8x1x4x4x3xf32>, tensor<i64>) -> tensor<8x6x4x4x3xf32>
// Build left halo on height dim.
// CHECK: %[[SLICE_H_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 0, 3, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 6, 1, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT:.*]] = "tf.Slice"(%[[CONCAT_D_TENSOR]], %[[SLICE_H_LEFT_BEGIN]], %[[SLICE_H_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x6x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x6x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1], [4, 6], [5, 7], [6, 4], [7, 5]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_H_LEFT]], %[[PAIRS_H_LEFT]])
// CHECK-SAME: (tensor<8x6x1x4x3xf32>, tensor<8x2xi32>) -> tensor<8x6x1x4x3xf32>
// Build right halo on height dim.
// CHECK: %[[SLICE_H_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 6, 1, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT:.*]] = "tf.Slice"(%[[CONCAT_D_TENSOR]], %[[SLICE_H_RIGHT_BEGIN]], %[[SLICE_H_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x6x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x6x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1], [4, 6], [5, 7], [6, 4], [7, 5]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_H_RIGHT]], %[[PAIRS_H_RIGHT]])
// CHECK-SAME: (tensor<8x6x1x4x3xf32>, tensor<8x2xi32>) -> tensor<8x6x1x4x3xf32>
// Concat the halos with the shard on the height dim.
// CHECK-NEXT: %[[CONCAT_H_AXIS:.*]] = "tf.Const"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_DH_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_H_LEFT]], %[[CONCAT_D_TENSOR]], %[[EXCHANGED_HALO_H_RIGHT]], %[[CONCAT_H_AXIS]])
// CHECK-SAME: (tensor<8x6x1x4x3xf32>, tensor<8x6x4x4x3xf32>, tensor<8x6x1x4x3xf32>, tensor<i64>) -> tensor<8x6x6x4x3xf32>
// Build left halo on width dim.
// CHECK: %[[SLICE_W_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 0, 0, 3, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 6, 6, 1, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT:.*]] = "tf.Slice"(%[[CONCAT_DH_TENSOR]], %[[SLICE_W_LEFT_BEGIN]], %[[SLICE_W_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x6x6x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x6x6x1x3xf32>
// CHECK-NEXT: %[[HALO_W_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2], [4, 5], [5, 4], [6, 7], [7, 6]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_W_LEFT]], %[[PAIRS_W_LEFT]])
// CHECK-SAME: (tensor<8x6x6x1x3xf32>, tensor<8x2xi32>) -> tensor<8x6x6x1x3xf32>
// Build right halo on width dim.
// CHECK: %[[SLICE_W_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 6, 6, 1, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT:.*]] = "tf.Slice"(%[[CONCAT_DH_TENSOR]], %[[SLICE_W_RIGHT_BEGIN]], %[[SLICE_W_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x6x6x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x6x6x1x3xf32>
// CHECK-NEXT: %[[HALO_W_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2], [4, 5], [5, 4], [6, 7], [7, 6]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_W_RIGHT]], %[[PAIRS_W_RIGHT]])
// CHECK-SAME: (tensor<8x6x6x1x3xf32>, tensor<8x2xi32>) -> tensor<8x6x6x1x3xf32>
// Concat the halos with the shard on the width dim.
// CHECK-NEXT: %[[CONCAT_W_AXIS:.*]] = "tf.Const"() <{value = dense<3> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_DHW_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_W_LEFT]], %[[CONCAT_DH_TENSOR]], %[[EXCHANGED_HALO_W_RIGHT]], %[[CONCAT_W_AXIS]])
// CHECK-SAME: (tensor<8x6x6x1x3xf32>, tensor<8x6x6x4x3xf32>, tensor<8x6x6x1x3xf32>, tensor<i64>) -> tensor<8x6x6x6x3xf32>
// CHECK-NEXT: "tf.Conv3D"(%[[CONCAT_DHW_TENSOR]], %arg2)
// CHECK-SAME: padding = "VALID"
// CHECK-SAME: (tensor<8x6x6x6x3xf32>, tensor<3x3x3x3x3xf32>) -> tensor<8x4x4x4x3xf32>
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x8x8x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,z,unsharded, mesh:|x=2,y=2,z=2|*TPU>} : (tensor<8x8x8x8x3xf32>) -> tensor<8x8x8x8x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2,z=8|*TPU>} : (tensor<3x3x3x3x3xf32>) -> tensor<3x3x3x3x3xf32>
%conv = "tf.Conv3D"(%img_layout, %filter_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} : (tensor<8x8x8x8x3xf32>, tensor<3x3x3x3x3xf32>) -> tensor<8x8x8x8x3xf32>
tf_device.return %conv : tensor<8x8x8x8x3xf32>
}) {_mesh="|x=2,y=2,z=2|*TPU"} : () -> (tensor<8x8x8x8x3xf32>)
func.return
}
// -----
// Check that Conv2D with spatial partitioning using "VALID" padding produces
// begin and end halos on both spatial dimensions and all necessary slice ops.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x8x8x3xf32> {tf._layout = "sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2:tensor<3x3x3x3xf32>) {
// CHECK: "tf_device.cluster"
// Build left halo on height dim.
// CHECK: %[[SLICE_H_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 3, 0, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT:.*]] = "tf.Slice"(%arg1, %[[SLICE_H_LEFT_BEGIN]], %[[SLICE_H_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_H_LEFT]], %[[PAIRS_H_LEFT]])
// CHECK-SAME: (tensor<8x1x4x3xf32>, tensor<4x2xi32>) -> tensor<8x1x4x3xf32>
// Build right halo on height dim.
// CHECK: %[[SLICE_H_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT:.*]] = "tf.Slice"(%arg1, %[[SLICE_H_RIGHT_BEGIN]], %[[SLICE_H_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_H_RIGHT]], %[[PAIRS_H_RIGHT]])
// CHECK-SAME: (tensor<8x1x4x3xf32>, tensor<4x2xi32>) -> tensor<8x1x4x3xf32>
// Concat the halos with the shard on the height dim.
// CHECK-NEXT: %[[CONCAT_H_AXIS:.*]] = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_H_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_H_LEFT]], %arg1, %[[EXCHANGED_HALO_H_RIGHT]], %[[CONCAT_H_AXIS]])
// CHECK-SAME: (tensor<8x1x4x3xf32>, tensor<8x4x4x3xf32>, tensor<8x1x4x3xf32>, tensor<i64>) -> tensor<8x6x4x3xf32>
// Dynamically slice the concatenated tensor to get correct size for VALID padding.
// CHECK-NEXT: %[[HALO_SIZES_H:.*]] = "tf.Const"() <{value = dense<[0, 1, 0, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[HALO_INCREMENTS_H:.*]] = "tf.Const"() <{value = dense<[0, 1, 0, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[VALID_OFFSET_H:.*]] = "tf.Mul"
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_H:.*]] = "tf.Sub"(%[[HALO_SIZES_H]], %[[VALID_OFFSET_H]])
// CHECK-NEXT: %[[VALID_SLICE_SIZE_H:.*]] = "tf.Const"() <{value = dense<[8, 5, 4, 3]> : tensor<4xi64>}> : () -> tensor<4xi64>
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_CAST_I64_H:.*]] = "tf.Cast"(%[[VALID_SLICE_BEGIN_H]]) <{Truncate = false}> : (tensor<4xi32>) -> tensor<4xi64>
// CHECK-NEXT: %[[VALID_SLICE_H_TENSOR:.*]] = "tf.Slice"(%[[CONCAT_H_TENSOR]], %[[VALID_SLICE_BEGIN_CAST_I64_H]], %[[VALID_SLICE_SIZE_H]])
// Build left halo on width dim.
// CHECK: %[[SLICE_W_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 0, 3, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 5, 1, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT:.*]] = "tf.Slice"(%[[VALID_SLICE_H_TENSOR]], %[[SLICE_W_LEFT_BEGIN]], %[[SLICE_W_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x5x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x5x1x3xf32>
// CHECK-NEXT: %[[HALO_W_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_W_LEFT]], %[[PAIRS_W_LEFT]])
// CHECK-SAME: (tensor<8x5x1x3xf32>, tensor<4x2xi32>) -> tensor<8x5x1x3xf32>
// Build right halo on width dim.
// CHECK: %[[SLICE_W_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 5, 1, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT:.*]] = "tf.Slice"(%[[VALID_SLICE_H_TENSOR]], %[[SLICE_W_RIGHT_BEGIN]], %[[SLICE_W_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x5x4x3xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<8x5x1x3xf32>
// CHECK-NEXT: %[[HALO_W_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_W_RIGHT]], %[[PAIRS_W_RIGHT]])
// CHECK-SAME: (tensor<8x5x1x3xf32>, tensor<4x2xi32>) -> tensor<8x5x1x3xf32>
// Concat the halos with the shard on the width dim.
// CHECK-NEXT: %[[CONCAT_W_AXIS:.*]] = "tf.Const"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_HW_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_W_LEFT]], %[[VALID_SLICE_H_TENSOR]], %[[EXCHANGED_HALO_W_RIGHT]], %[[CONCAT_W_AXIS]])
// CHECK-SAME: (tensor<8x5x1x3xf32>, tensor<8x5x4x3xf32>, tensor<8x5x1x3xf32>, tensor<i64>) -> tensor<8x5x6x3xf32>
// Dynamically slice the concatenated tensor to get correct size for VALID padding.
// CHECK-NEXT: %[[HALO_SIZES_W:.*]] = "tf.Const"() <{value = dense<[0, 0, 1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[HALO_INCREMENTS_W:.*]] = "tf.Const"() <{value = dense<[0, 0, 1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT: %[[VALID_OFFSET_W:.*]] = "tf.Mul"
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_W:.*]] = "tf.Sub"(%[[HALO_SIZES_W]], %[[VALID_OFFSET_W]])
// CHECK-NEXT: %[[VALID_SLICE_SIZE_W:.*]] = "tf.Const"() <{value = dense<[8, 5, 5, 3]> : tensor<4xi64>}> : () -> tensor<4xi64>
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_CAST_I64_W:.*]] = "tf.Cast"(%[[VALID_SLICE_BEGIN_W]]) <{Truncate = false}> : (tensor<4xi32>) -> tensor<4xi64>
// CHECK-NEXT: %[[VALID_SLICE_HW_TENSOR:.*]] = "tf.Slice"(%[[CONCAT_HW_TENSOR]], %[[VALID_SLICE_BEGIN_CAST_I64_W]], %[[VALID_SLICE_SIZE_W]])
// CHECK-NEXT: "tf.Conv2D"(%[[VALID_SLICE_HW_TENSOR]], %arg2)
// CHECK-SAME: padding = "VALID"
// CHECK-SAME: (tensor<8x5x5x3xf32>, tensor<3x3x3x3xf32>) -> tensor<8x3x3x3xf32>
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x8x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<8x8x8x3xf32>) -> tensor<8x8x8x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3xf32>
%conv = "tf.Conv2D"(%img_layout, %filter_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<8x8x8x3xf32>, tensor<3x3x3x3xf32>) -> tensor<8x6x6x3xf32>
tf_device.return %conv : tensor<8x6x6x3xf32>
}) {_mesh="|x=2,y=2|*TPU"} : () -> (tensor<8x6x6x3xf32>)
func.return
}
// -----
// Check that Conv3D with spatial partitioning using "VALID" padding produces
// begin and end halos on all spatial dimensions and all necessary slice ops.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x8x8x8x3xf32> {tf._layout = "sharding_specs:unsharded,x,y,z,unsharded, mesh:|x=2,y=2,z=2|*TPU"},
%arg2:tensor<3x3x3x3x3xf32>) {
// CHECK: "tf_device.cluster"
// Build left halo on depth dim.
// CHECK: %[[SLICE_D_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 3, 0, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_LEFT:.*]] = "tf.Slice"(%arg1, %[[SLICE_D_LEFT_BEGIN]], %[[SLICE_D_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x1x4x4x3xf32>
// CHECK-NEXT: %[[HALO_D_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_D_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 4], [1, 5], [2, 6], [3, 7], [4, 0], [5, 1], [6, 2], [7, 3]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_D_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_D_LEFT]], %[[PAIRS_D_LEFT]])
// CHECK-SAME: (tensor<8x1x4x4x3xf32>, tensor<8x2xi32>) -> tensor<8x1x4x4x3xf32>
// Build right halo on depth dim.
// CHECK: %[[SLICE_D_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 1, 4, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_D_RIGHT:.*]] = "tf.Slice"(%arg1, %[[SLICE_D_RIGHT_BEGIN]], %[[SLICE_D_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x4x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x1x4x4x3xf32>
// CHECK-NEXT: %[[HALO_D_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_D_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 4], [1, 5], [2, 6], [3, 7], [4, 0], [5, 1], [6, 2], [7, 3]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_D_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_D_RIGHT]], %[[PAIRS_D_RIGHT]])
// CHECK-SAME: (tensor<8x1x4x4x3xf32>, tensor<8x2xi32>) -> tensor<8x1x4x4x3xf32>
// Concat the halos with the shard on the depth dim.
// CHECK-NEXT: %[[CONCAT_D_AXIS:.*]] = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_D_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_D_LEFT]], %arg1, %[[EXCHANGED_HALO_D_RIGHT]], %[[CONCAT_D_AXIS]])
// CHECK-SAME: (tensor<8x1x4x4x3xf32>, tensor<8x4x4x4x3xf32>, tensor<8x1x4x4x3xf32>, tensor<i64>) -> tensor<8x6x4x4x3xf32>
// Dynamically slice the concatenated tensor to get correct size for VALID padding.
// CHECK-NEXT: %[[HALO_SIZES_D:.*]] = "tf.Const"() <{value = dense<[0, 1, 0, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[HALO_INCREMENTS_D:.*]] = "tf.Const"() <{value = dense<[0, 1, 0, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[VALID_OFFSET_D:.*]] = "tf.Mul"
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_D:.*]] = "tf.Sub"(%[[HALO_SIZES_D]], %[[VALID_OFFSET_D]])
// CHECK-NEXT: %[[VALID_SLICE_SIZE_D:.*]] = "tf.Const"() <{value = dense<[8, 5, 4, 4, 3]> : tensor<5xi64>}> : () -> tensor<5xi64>
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_CAST_I64_D:.*]] = "tf.Cast"(%[[VALID_SLICE_BEGIN_D]]) <{Truncate = false}> : (tensor<5xi32>) -> tensor<5xi64>
// CHECK-NEXT: %[[VALID_SLICE_D_TENSOR:.*]] = "tf.Slice"(%[[CONCAT_D_TENSOR]], %[[VALID_SLICE_BEGIN_CAST_I64_D]], %[[VALID_SLICE_SIZE_D]])
// Build left halo on height dim.
// CHECK: %[[SLICE_H_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 0, 3, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 5, 1, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_LEFT:.*]] = "tf.Slice"(%[[VALID_SLICE_D_TENSOR]], %[[SLICE_H_LEFT_BEGIN]], %[[SLICE_H_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x5x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x5x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1], [4, 6], [5, 7], [6, 4], [7, 5]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_H_LEFT]], %[[PAIRS_H_LEFT]])
// CHECK-SAME: (tensor<8x5x1x4x3xf32>, tensor<8x2xi32>) -> tensor<8x5x1x4x3xf32>
// Build right halo on height dim.
// CHECK: %[[SLICE_H_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 5, 1, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_H_RIGHT:.*]] = "tf.Slice"(%[[VALID_SLICE_D_TENSOR]], %[[SLICE_H_RIGHT_BEGIN]], %[[SLICE_H_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x5x4x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x5x1x4x3xf32>
// CHECK-NEXT: %[[HALO_H_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_H_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 2], [1, 3], [2, 0], [3, 1], [4, 6], [5, 7], [6, 4], [7, 5]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_H_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_H_RIGHT]], %[[PAIRS_H_RIGHT]])
// CHECK-SAME: (tensor<8x5x1x4x3xf32>, tensor<8x2xi32>) -> tensor<8x5x1x4x3xf32>
// Concat the halos with the shard on the height dim.
// CHECK-NEXT: %[[CONCAT_H_AXIS:.*]] = "tf.Const"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_DH_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_H_LEFT]], %[[VALID_SLICE_D_TENSOR]], %[[EXCHANGED_HALO_H_RIGHT]], %[[CONCAT_H_AXIS]])
// CHECK-SAME: (tensor<8x5x1x4x3xf32>, tensor<8x5x4x4x3xf32>, tensor<8x5x1x4x3xf32>, tensor<i64>) -> tensor<8x5x6x4x3xf32>
// Dynamically slice the concatenated tensor to get correct size for VALID padding.
// CHECK-NEXT: %[[HALO_SIZES_H:.*]] = "tf.Const"() <{value = dense<[0, 0, 1, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[HALO_INCREMENTS_H:.*]] = "tf.Const"() <{value = dense<[0, 0, 1, 0, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[VALID_OFFSET_H:.*]] = "tf.Mul"
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_H:.*]] = "tf.Sub"(%[[HALO_SIZES_H]], %[[VALID_OFFSET_H]])
// CHECK-NEXT: %[[VALID_SLICE_SIZE_H:.*]] = "tf.Const"() <{value = dense<[8, 5, 5, 4, 3]> : tensor<5xi64>}> : () -> tensor<5xi64>
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_CAST_I64_H:.*]] = "tf.Cast"(%[[VALID_SLICE_BEGIN_H]]) <{Truncate = false}> : (tensor<5xi32>) -> tensor<5xi64>
// CHECK-NEXT: %[[VALID_SLICE_DH_TENSOR:.*]] = "tf.Slice"(%[[CONCAT_DH_TENSOR]], %[[VALID_SLICE_BEGIN_CAST_I64_H]], %[[VALID_SLICE_SIZE_H]])
// Build left halo on width dim.
// CHECK: %[[SLICE_W_LEFT_BEGIN:.*]] = "tf.Const"() <{value = dense<[0, 0, 0, 3, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 5, 5, 1, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_LEFT:.*]] = "tf.Slice"(%[[VALID_SLICE_DH_TENSOR]], %[[SLICE_W_LEFT_BEGIN]], %[[SLICE_W_LEFT_SIZE]])
// CHECK-SAME: (tensor<8x5x5x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x5x5x1x3xf32>
// CHECK-NEXT: %[[HALO_W_LEFT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_LEFT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2], [4, 5], [5, 4], [6, 7], [7, 6]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_LEFT:.*]] = "tf.CollectivePermute"(%[[HALO_W_LEFT]], %[[PAIRS_W_LEFT]])
// CHECK-SAME: (tensor<8x5x5x1x3xf32>, tensor<8x2xi32>) -> tensor<8x5x5x1x3xf32>
// Build right halo on width dim.
// CHECK: %[[SLICE_W_RIGHT_BEGIN:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT_SIZE:.*]] = "tf.Const"() <{value = dense<[8, 5, 5, 1, 3]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[SLICE_W_RIGHT:.*]] = "tf.Slice"(%[[VALID_SLICE_DH_TENSOR]], %[[SLICE_W_RIGHT_BEGIN]], %[[SLICE_W_RIGHT_SIZE]])
// CHECK-SAME: (tensor<8x5x5x4x3xf32>, tensor<5xi32>, tensor<5xi32>) -> tensor<8x5x5x1x3xf32>
// CHECK-NEXT: %[[HALO_W_RIGHT:.*]] = "tf.SelectV2"
// CHECK-NEXT: %[[PAIRS_W_RIGHT:.*]] = "tf.Const"
// CHECK-SAME{LITERAL}: value = dense<[[0, 1], [1, 0], [2, 3], [3, 2], [4, 5], [5, 4], [6, 7], [7, 6]]>
// CHECK-NEXT: %[[EXCHANGED_HALO_W_RIGHT:.*]] = "tf.CollectivePermute"(%[[HALO_W_RIGHT]], %[[PAIRS_W_RIGHT]])
// CHECK-SAME: (tensor<8x5x5x1x3xf32>, tensor<8x2xi32>) -> tensor<8x5x5x1x3xf32>
// Concat the halos with the shard on the width dim.
// CHECK-NEXT: %[[CONCAT_W_AXIS:.*]] = "tf.Const"() <{value = dense<3> : tensor<i64>}> : () -> tensor<i64>
// CHECK-NEXT: %[[CONCAT_DHW_TENSOR:.*]] = "tf.ConcatV2"(%[[EXCHANGED_HALO_W_LEFT]], %[[VALID_SLICE_DH_TENSOR]], %[[EXCHANGED_HALO_W_RIGHT]], %[[CONCAT_W_AXIS]])
// CHECK-SAME: (tensor<8x5x5x1x3xf32>, tensor<8x5x5x4x3xf32>, tensor<8x5x5x1x3xf32>, tensor<i64>) -> tensor<8x5x5x6x3xf32>
// Dynamically slice the concatenated tensor to get correct size for VALID padding.
// CHECK-NEXT: %[[HALO_SIZES_W:.*]] = "tf.Const"() <{value = dense<[0, 0, 0, 1, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[HALO_INCREMENTS_W:.*]] = "tf.Const"() <{value = dense<[0, 0, 0, 1, 0]> : tensor<5xi32>}> : () -> tensor<5xi32>
// CHECK-NEXT: %[[VALID_OFFSET_W:.*]] = "tf.Mul"
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_W:.*]] = "tf.Sub"(%[[HALO_SIZES_W]], %[[VALID_OFFSET_W]])
// CHECK-NEXT: %[[VALID_SLICE_SIZE_W:.*]] = "tf.Const"() <{value = dense<[8, 5, 5, 5, 3]> : tensor<5xi64>}> : () -> tensor<5xi64>
// CHECK-NEXT: %[[VALID_SLICE_BEGIN_CAST_I64_W:.*]] = "tf.Cast"(%[[VALID_SLICE_BEGIN_W]]) <{Truncate = false}> : (tensor<5xi32>) -> tensor<5xi64>
// CHECK-NEXT: %[[VALID_SLICE_DHW_TENSOR:.*]] = "tf.Slice"(%[[CONCAT_DHW_TENSOR]], %[[VALID_SLICE_BEGIN_CAST_I64_W]], %[[VALID_SLICE_SIZE_W]])
// CHECK-NEXT: "tf.Conv3D"(%[[VALID_SLICE_DHW_TENSOR]], %arg2)
// CHECK-SAME: padding = "VALID"
// CHECK-SAME: (tensor<8x5x5x5x3xf32>, tensor<3x3x3x3x3xf32>) -> tensor<8x3x3x3x3xf32>
%0 = "tf_device.cluster"() ({
%img_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x8x8x8x3>, layout = #dtensor.layout<sharding_specs:unsharded,x,y,z,unsharded, mesh:|x=2,y=2,z=2|*TPU>} : (tensor<8x8x8x8x3xf32>) -> tensor<8x8x8x8x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2,z=8|*TPU>} : (tensor<3x3x3x3x3xf32>) -> tensor<3x3x3x3x3xf32>
%conv = "tf.Conv3D"(%img_layout, %filter_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "VALID", strides = [1, 1, 1, 1, 1]} : (tensor<8x8x8x8x3xf32>, tensor<3x3x3x3x3xf32>) -> tensor<8x6x6x6x3xf32>
tf_device.return %conv : tensor<8x6x6x6x3xf32>
}) {_mesh="|x=2,y=2,z=2|*TPU"} : () -> (tensor<8x6x6x6x3xf32>)
func.return
}
// -----
// Check that Conv2DBackpropInputV2 becomes Conv2DBackpropInput in SPMD expansion
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"},
%arg2: tensor<1x3x3x3xf32>,
%arg3: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv2DBackpropInput"
// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]
// CHECK-SAME: (tensor<4xi32>, tensor<1x3x3x3xf32>, tensor<4x32x32x3xf32>) -> tensor<4x32x32x3xf32>
%0 = "tf_device.cluster"() ({
%input_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<1x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3xf32>) -> tensor<1x3x3x3xf32>
%grad_layout = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%conv = "tf.Conv2DBackpropInputV2"(%input_layout, %filter_layout, %grad_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} :
(tensor<8x32x32x3xf32>, tensor<1x3x3x3xf32>, tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%conv_layout = "tf.DTensorLayout"(%conv) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
tf_device.return %conv_layout : tensor<8x32x32x3xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"}: () -> (tensor<8x32x32x3xf32>)
func.return
}
// -----
// Check that Conv3DBackpropInput becomes Conv3DBackpropInputV2 in SPMD expansion
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"},
%arg2: tensor<1x3x3x3x3xf32>,
%arg3: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv3DBackpropInputV2"
// CHECK-SAME: data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]
// CHECK-SAME: (tensor<5xi32>, tensor<1x3x3x3x3xf32>, tensor<4x32x32x32x3xf32>) -> tensor<4x32x32x32x3xf32>
%0 = "tf_device.cluster"() ({
%input_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<1x3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3x3xf32>) -> tensor<1x3x3x3x3xf32>
%grad_layout = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%conv = "tf.Conv3DBackpropInput"(%input_layout, %filter_layout, %grad_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} :
(tensor<8x32x32x32x3xf32>, tensor<1x3x3x3x3xf32>, tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%conv_layout = "tf.DTensorLayout"(%conv) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
tf_device.return %conv_layout : tensor<8x32x32x32x3xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"}: () -> (tensor<8x32x32x32x3xf32>)
func.return
}
// -----
// Check that Conv2DBackpropFilterV2 becomes Conv2DBackpropFilter in SPMD expansion
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"},
%arg2: tensor<1x3x3x3xf32>,
%arg3: tensor<8x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv2DBackpropFilter"
// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]
// CHECK-SAME: (tensor<4x32x32x3xf32>, tensor<4xi32>, tensor<4x32x32x3xf32>) -> tensor<1x3x3x3xf32>
%0 = "tf_device.cluster"() ({
%input_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<1x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3xf32>) -> tensor<1x3x3x3xf32>
%grad_layout = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<8x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x3xf32>) -> tensor<8x32x32x3xf32>
%conv = "tf.Conv2DBackpropFilterV2"(%input_layout, %filter_layout, %grad_layout) {data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} :
(tensor<8x32x32x3xf32>, tensor<1x3x3x3xf32>, tensor<8x32x32x3xf32>) -> tensor<1x3x3x3xf32>
%conv_layout = "tf.DTensorLayout"(%conv) {global_shape = #tf_type.shape<1x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3xf32>) -> tensor<1x3x3x3xf32>
tf_device.return %conv_layout : tensor<1x3x3x3xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"}: () -> (tensor<1x3x3x3xf32>)
func.return
}
// -----
// Check that Conv3DBackpropFilter becomes Conv3DBackpropFilterV2 in SPMD expansion
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"},
%arg2: tensor<1x3x3x3x3xf32>,
%arg3: tensor<8x32x32x32x3xf32> {tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=1|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONV_OUT:.*]] = "tf.Conv3DBackpropFilterV2"
// CHECK-SAME: data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]
// CHECK-SAME: (tensor<4x32x32x32x3xf32>, tensor<5xi32>, tensor<4x32x32x32x3xf32>) -> tensor<1x3x3x3x3xf32>
%0 = "tf_device.cluster"() ({
%input_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%filter_layout = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<1x3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3x3xf32>) -> tensor<1x3x3x3x3xf32>
%grad_layout = "tf.DTensorLayout"(%arg3) {global_shape = #tf_type.shape<8x32x32x32x3>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<8x32x32x32x3xf32>) -> tensor<8x32x32x32x3xf32>
%conv = "tf.Conv3DBackpropFilter"(%input_layout, %filter_layout, %grad_layout) {data_format = "NDHWC", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1, 1]} :
(tensor<8x32x32x32x3xf32>, tensor<1x3x3x3x3xf32>, tensor<8x32x32x32x3xf32>) -> tensor<1x3x3x3x3xf32>
%conv_layout = "tf.DTensorLayout"(%conv) {global_shape = #tf_type.shape<1x3x3x3x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} :
(tensor<1x3x3x3x3xf32>) -> tensor<1x3x3x3x3xf32>
tf_device.return %conv_layout : tensor<1x3x3x3x3xf32>
}) {_mesh="TPU|x=2,y=2|*TPU"}: () -> (tensor<1x3x3x3x3xf32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_dtensor_ops.mlir
+337
View File
@@ -0,0 +1,337 @@
// 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: dtensor-opt %s -split-input-file -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check that Data transfer from CPU to TPU is lowered correctly.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[PROGRAM_KEY:.*]] = "tf._XlaCompileMlirPlaceholderProgramKey"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"
// CHECK-NEXT: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[PREDICATE:.*]] = "tf.Equal"(%[[DEVICE_ORDINAL_SCALAR_64]], %[[ZERO]])
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE]])
// CHECK-NEXT: %[[ZERO_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: "tf._XlaSendFromHostV2"(%[[CONST_OUT]], %[[PROGRAM_KEY]], %[[ZERO_2]])
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[ID_TO_ORDINAL_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<[0, 1, 2, 3]>
// CHECK-NEXT: %[[SIZE_TYPE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID_2:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE_2]])
// CHECK-NEXT: %[[SLICE_SIZE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL_2:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL_2]], %[[DEVICE_ID_2]], %[[SLICE_SIZE_2]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_2:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL_2]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64_2:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR_2]])
// CHECK-NEXT: %[[ZERO_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[PREDICATE_2:.*]] = "tf.Equal"(%[[DEVICE_ORDINAL_SCALAR_64_2]], %[[ZERO_2]])
// CHECK-NEXT: %[[IF_OUT:.*]] = "tf.IfRegion"(%[[PREDICATE_2]])
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf.XlaRecvFromHost"()
// CHECK-SAME: key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0"
// CHECK-NEXT: "tf.Yield"(%[[RECV_OUT]])
// CHECK: %[[ZEROS_3:.*]] = "tf.Const"
// CHECK-NEXT: "tf.Yield"(%[[ZEROS_3]])
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[OUTPUT:.*]] = "tf.DTensorAllReduce"(%[[IF_OUT]], %[[GROUP_ASSIGNMENT]])
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<1xi32>) -> tensor<1xi32>
"tf.DTensorSend"(%1) {key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0", target_mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, target_layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3_0", mesh = #dtensor.mesh<TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>, shape = #tf_type.shape<>, source_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<1xi32>) -> tensor<1xi32>
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
// -----
// Check that Data transfer from TPU to CPU is lowered correctly.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"
// CHECK-NEXT: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<[0, 1, 2, 3]>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[PREDICATE:.*]] = "tf.Equal"(%[[DEVICE_ORDINAL_SCALAR_64]], %[[ZERO]])
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE]])
// CHECK-NEXT: "tf.XlaSendToHost"(%[[CONST_OUT]])
// CHECK-SAME: key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_0"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[PROGRAM_KEY:.*]] = "tf._XlaCompileMlirPlaceholderProgramKey"
// CHECK-NEXT: %[[ID_TO_ORDINAL_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[SIZE_TYPE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID_2:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE_2]])
// CHECK-NEXT: %[[SLICE_SIZE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL_2:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL_2]], %[[DEVICE_ID_2]], %[[SLICE_SIZE_2]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_2:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL_2]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64_2:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR_2]])
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_KEY]], %[[DEVICE_ORDINAL_SCALAR_64_2]])
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<1xi32>) -> tensor<1xi32>
"tf.DTensorSend"(%1) {key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_0", target_mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>, target_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_0", mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>, shape = #tf_type.shape<>, source_layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<1xi32>) -> tensor<1xi32>
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
func.return
}
// -----
// Check that tensor to send is converted to replicated layout before send.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"
// CHECK: %[[ALL_GATHER_OUT:.*]] = "tf.DTensorAllGather"
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<[0, 1, 2, 3]>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[PREDICATE:.*]] = "tf.Equal"(%[[DEVICE_ORDINAL_SCALAR_64]], %[[ZERO]])
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE]])
// CHECK-NEXT: "tf.XlaSendToHost"(%[[ALL_GATHER_OUT]])
// CHECK-SAME: key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_0"
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[PROGRAM_KEY:.*]] = "tf._XlaCompileMlirPlaceholderProgramKey"
// CHECK-NEXT: %[[ID_TO_ORDINAL_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[SIZE_TYPE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID_2:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE_2]])
// CHECK-NEXT: %[[SLICE_SIZE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL_2:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL_2]], %[[DEVICE_ID_2]], %[[SLICE_SIZE_2]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_2:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL_2]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64_2:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR_2]])
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_KEY]], %[[DEVICE_ORDINAL_SCALAR_64_2]])
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2xi32>) -> tensor<2xi32>
"tf.DTensorSend"(%1) {key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_0", target_mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>, target_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<2xi32>) -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0_0", mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>, shape = #tf_type.shape<>, source_layout = #dtensor.layout<sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : () -> tensor<2xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<2xi32>) -> tensor<2xi32>
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
func.return
}
// -----
// Check that send/recv to clusters with same mesh is disallowed.
func.func @main(%arg0: tensor<i32>) {
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=2|*CPU>} : (tensor<1xi32>) -> tensor<1xi32>
// expected-error @+1 {{Only use CopyToMesh to transfer data across different mesh cluster}}
"tf.DTensorSend"(%1) {key = "communication_key_CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1_0", target_mesh = #dtensor.mesh<CPU|x=2|*CPU>, target_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=2|*CPU>} : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=2|*CPU"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1_0", mesh = #dtensor.mesh<CPU|x=2|*CPU>, shape = #tf_type.shape<>, source_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=2|*CPU>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=2|*CPU>} : (tensor<1xi32>) -> tensor<1xi32>
tf_device.return
}) {_mesh = "CPU|x=2|*CPU"} : () -> ()
func.return
}
// -----
// Check that multi-mesh transfer between two non host clusters is disallowed.
func.func @main(%arg0: tensor<i32>) {
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<1xi32>) -> tensor<1xi32>
// expected-error @+1 {{f.CopyToMesh op must be used to send data from/to host mesh}}
"tf.DTensorSend"(%1) {key = "communication_key_GPU|x=2|0,1|0,1|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1_0", target_mesh = #dtensor.mesh<GPU|x=2|*GPU>, target_layout = #dtensor.layout<sharding_specs:unsharded, mesh:GPU|x=2|*GPU>} : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_GPU|x=2|0,1|0,1|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1_0", mesh = #dtensor.mesh<GPU|x=2|*GPU>, shape = #tf_type.shape<>, source_layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:GPU|x=2|*GPU>} : (tensor<1xi32>) -> tensor<1xi32>
tf_device.return
}) {_mesh = "GPU|x=2|*GPU"} : () -> ()
func.return
}
// -----
// Check that send/recv between two CPUs works.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"
// CHECK-NEXT: "tf._HostSend"(%[[CONST_OUT]])
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf._HostRecv"
// CHECK-SAME: tensor_name = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1_0"
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : (tensor<1xi32>) -> tensor<1xi32>
"tf.DTensorSend"(%1) {key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1_0", target_mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1>, target_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1>} : (tensor<1xi32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1_0", mesh = #dtensor.mesh<CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1>, shape = #tf_type.shape<1>, source_layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0>} : () -> tensor<1xi32>
%1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1>} : (tensor<1xi32>) -> tensor<1xi32>
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
// -----
// Check that Data transfer from CPU to GPU is lowered correctly.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[INPUT:.*]] = "tf.Const"
// CHECK-NEXT: "tf._HostSend"(%[[INPUT]])
// CHECK-NEXT: %[[ID_TO_ORDINAL:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[SIZE_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE]])
// CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL]], %[[DEVICE_ID]], %[[SLICE_SIZE]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR]])
// CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[PREDICATE:.*]] = "tf.Equal"(%[[DEVICE_ORDINAL_SCALAR_64]], %[[ZERO]])
// CHECK-NEXT: "tf.IfRegion"(%[[PREDICATE]])
// CHECK-NEXT: "tf.Yield"
// CHECK: "tf.Yield"
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[ID_TO_ORDINAL_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<[0, 1, 2, 3]>
// CHECK-NEXT: %[[SIZE_TYPE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ID_2:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_TYPE_2]])
// CHECK-NEXT: %[[SLICE_SIZE_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<1>
// CHECK-NEXT: %[[DEVICE_ORDINAL_2:.*]] = "tf.Slice"(%[[ID_TO_ORDINAL_2]], %[[DEVICE_ID_2]], %[[SLICE_SIZE_2]])
// CHECK-NEXT: %[[SCALAR_TYPE:.*]] = "tf.Const"
// CHECK-SAME: value = dense<>
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_2:.*]] = "tf.Reshape"(%[[DEVICE_ORDINAL_2]], %[[SCALAR_TYPE]])
// CHECK-NEXT: %[[DEVICE_ORDINAL_SCALAR_64_2:.*]] = "tf.Cast"(%[[DEVICE_ORDINAL_SCALAR_2]])
// CHECK-NEXT: %[[ZERO_2:.*]] = "tf.Const"
// CHECK-SAME: value = dense<0>
// CHECK-NEXT: %[[PREDICATE_2:.*]] = "tf.Equal"(%[[DEVICE_ORDINAL_SCALAR_64_2]], %[[ZERO_2]])
// CHECK-NEXT: %[[IF_OUT:.*]] = "tf.IfRegion"(%[[PREDICATE_2]])
// CHECK-NEXT: %[[RECV_OUT:.*]] = "tf._HostRecv"
// CHECK-NEXT: "tf.Yield"(%[[RECV_OUT]])
// CHECK: %[[ZEROS_3:.*]] = "tf.Const"
// CHECK-NEXT: "tf.Yield"(%[[ZEROS_3]])
// CHECK: %[[GROUP_ASSIGNMENT:.*]] = "tf.Const"
// CHECK-NEXT: %[[OUGPUT:.*]] = "tf.DTensorAllReduce"(%[[IF_OUT]], %[[GROUP_ASSIGNMENT]])
"tf_device.cluster"() ({
%0 = "tf.Const"() {value = dense<1.> : tensor<8x8xf32>} : () -> tensor<8x8xf32>
%1 = "tf.DTensorLayout"(%0) {_global_shape = [#tf_type.shape<8x8>], global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
"tf.DTensorSend"(%1) {key = "communication_key_GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3_0", target_mesh = #dtensor.mesh<GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>, target_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>} : (tensor<8x8xf32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=1|0|0|/job:localhost/task:0/device:CPU:0"} : () -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorRecv"() {key = "communication_key_GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3_0", mesh = #dtensor.mesh<GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>, shape = #tf_type.shape<8x8>, source_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : () -> tensor<8x8xf32>
%1 = "tf.DTensorLayout"(%0) {_global_shape = [#tf_type.shape<8x8>], global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
tf_device.return
}) {_mesh = "GPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:GPU:0,/job:localhost/task:0/device:GPU:1,/job:localhost/task:0/device:GPU:2,/job:localhost/task:0/device:GPU:3"} : () -> ()
func.return
}
tensorflow/dtensor/mlir/tests/spmd_einsum.mlir
+228
View File
@@ -0,0 +1,228 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Einsum (normal matrix multiplication)
// No AllToAll on input or output, only AllReduce on output.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4xi32> {tf._layout = "sharding_specs:unsharded,y, mesh:|x=2,y=2|*GPU"},
%arg2: tensor<4x2xi32> {tf._layout = "sharding_specs:y,unsharded, mesh:|x=2,y=2|*GPU"}) -> tensor<2x2xi32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[EINSUM_RESULT:.*]] = "tf.Einsum"(%arg1, %arg2)
// CHECK: %[[RETURN:.*]] = "tf.DTensorAllReduce"(%[[EINSUM_RESULT]]
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[RETURN]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:|x=2,y=2|*GPU>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x2>, layout = #dtensor.layout<sharding_specs:y,unsharded, mesh:|x=2,y=2|*GPU>} : (tensor<4x2xi32>) -> tensor<4x2xi32>
%3 = "tf.Einsum"(%1, %2) {equation="ab,bc->ac"} : (tensor<2x4xi32>, tensor<4x2xi32>) -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*GPU>} : (tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %4 : tensor<2x2xi32>
}) {_mesh = "|x=2,y=2|*GPU", _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*GPU"]} : () -> (tensor<2x2xi32>)
func.return %0 : tensor<2x2xi32>
}
// -----
// Replicated Einsum (normal matrix multiplication, no sharded dimensions reduced)
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4xi32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2: tensor<4x2xi32> {tf._layout = "sharding_specs:unsharded,y, mesh:|x=2,y=2|*TPU"}) -> tensor<2x2xi32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[EINSUM_RESULT:.*]] = "tf.Einsum"(%arg1, %arg2)
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[EINSUM_RESULT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x2>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:|x=2,y=2|*TPU>} : (tensor<4x2xi32>) -> tensor<4x2xi32>
%3 = "tf.Einsum"(%1, %2) {equation="ab,bc->ac"} : (tensor<2x4xi32>, tensor<4x2xi32>) -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:|x=2,y=2|*TPU>} : (tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %4 : tensor<2x2xi32>
}) {_mesh = "|x=2,y=2|*TPU", _layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*TPU"]} : () -> (tensor<2x2xi32>)
func.return %0 : tensor<2x2xi32>
}
// -----
// Einsum for transformer, no CRS
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<16x25x8xi32>{ tf._layout="sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU"},
%arg2: tensor<25x8x8x50xi32>{ tf._layout="sharding_specs:unsharded,unsharded,y,unsharded, mesh:|x=2,y=2|*TPU"}) -> tensor<16x25x8x50xi32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[EINSUM_RESULT:.*]] = "tf.Einsum"(%arg1, %arg2)
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[EINSUM_RESULT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<16x25x8>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<16x25x8xi32>) -> tensor<16x25x8xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<25x8x8x50>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<25x8x8x50xi32>) -> tensor<25x8x8x50xi32>
%3 = "tf.Einsum"(%1, %2) {equation="bse,sehq->bshq"} : (tensor<16x25x8xi32>, tensor<25x8x8x50xi32>) -> tensor<16x25x8x50xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<16x25x8x50>, layout = #dtensor.layout<sharding_specs:x,unsharded,y,unsharded, mesh:|x=2,y=2|*TPU>} : (tensor<16x25x8x50xi32>) -> tensor<16x25x8x50xi32>
tf_device.return %4 : tensor<16x25x8x50xi32>
}) {_mesh = "|x=2,y=2|*TPU"} : () -> (tensor<16x25x8x50xi32>)
func.return %0 : tensor<16x25x8x50xi32>
}
// -----
// Invalid equation
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<16x24x8xi32>{ tf._layout="sharding_specs:x,z,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"},
%arg2: tensor<24x8x8x50xi32>{ tf._layout="sharding_specs:y,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"}) -> tensor<16x24x8x50xi32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<16x24x8>, layout = #dtensor.layout<sharding_specs:x,z,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7>} : (tensor<16x24x8xi32>) -> tensor<16x24x8xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<24x8x8x50>, layout = #dtensor.layout<sharding_specs:y,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7>} : (tensor<24x8x8x50xi32>) -> tensor<24x8x8x50xi32>
// expected-error @+1 {{incompatible mesh dimensions in equation, label 's' is mapped to mesh dimension 'y' and 'z'}}
%3 = "tf.Einsum"(%1, %2) {equation="bse,sehq->bshq"} : (tensor<16x24x8xi32>, tensor<24x8x8x50xi32>) -> tensor<16x24x8x50xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<16x24x8x50>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7>} : (tensor<16x24x8x50xi32>) -> tensor<16x24x8x50xi32>
tf_device.return %4 : tensor<16x24x8x50xi32>
}) {_mesh = "TPU|x=2,y=2,z=2|*TPU", _layout = ["sharding_specs:x,unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2,z=2|*TPU"]} : () -> (tensor<16x24x8x50xi32>)
func.return %0 : tensor<16x24x8x50xi32>
}
// -----
// y,x . x,y -> *,y
// We unshard %arg1
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:y,x, mesh:TPU|x=2,y=2|*TPU"},
%arg2: tensor<4x4xf32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<4x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-NEXT: %[[EINSUM_RESULT:[0-9]*]] = "tf.Einsum"(%[[GATHERED]], %arg2)
// CHECK: %[[FINAL_REDUCE:[0-9]*]] = "tf.DTensorAllReduce"(%[[EINSUM_RESULT]], %cst)
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:y,x, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%3 = "tf.Einsum"(%1, %2) {equation="ab,bc->ac"} : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
tf_device.return %4 : tensor<4x4xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<4x4xf32>
func.return %0 : tensor<4x4xf32>
}
// -----
// *,x . x,* -> *,y
// We should slice arg2 before matmul rather than slicing the result.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|*TPU"},
%arg2: tensor<4x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<4x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[SLICE:[0-9]*]] = "tf.DTensorAllScatter"(%arg2)
// CHECK-NEXT: %[[EINSUM_RESULT:[0-9]*]] = "tf.Einsum"(%arg1, %[[SLICE]])
// CHECK: %[[FINAL_REDUCE:[0-9]*]] = "tf.DTensorAllReduce"(%[[EINSUM_RESULT]], %cst)
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%3 = "tf.Einsum"(%1, %2) {equation="ab,bc->ac"} : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
tf_device.return %4 : tensor<4x4xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<4x4xf32>
func.return %0 : tensor<4x4xf32>
}
// -----
// x,y . *,y -> x,y
// We unshard %arg1 on y.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"},
%arg2: tensor<4x4xf32> {tf._layout = "sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<4x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-NEXT: %[[EINSUM_RESULT:[0-9]*]] = "tf.Einsum"(%[[GATHERED]], %arg2)
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%3 = "tf.Einsum"(%1, %2) {equation="ab,bc->ac"} : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
tf_device.return %4 : tensor<4x4xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<4x4xf32>
func.return %0 : tensor<4x4xf32>
}
// -----
// Example from BERT 64 way sharding.
// bsd,dnh->bsnh x,*,y . *,y,* -> x,*,y,*
// Unshard arg1 along y.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<8x128x128xf32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU"},
%arg2: tensor<128x16x64xf32> {tf._layout = "sharding_specs:unsharded,y,unsharded, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<8x128x16x64xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-NEXT: %[[EINSUM_RESULT:[0-9]*]] = "tf.Einsum"(%[[GATHERED]], %arg2)
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[EINSUM_RESULT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x128x128>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128xf32>) -> tensor<8x128x128xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<128x16x64>, layout = #dtensor.layout<sharding_specs:unsharded,y,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<128x16x64xf32>) -> tensor<128x16x64xf32>
%3 = "tf.Einsum"(%1, %2) {equation="bsd,dnh->bsnh"} : (tensor<8x128x128xf32>, tensor<128x16x64xf32>) -> tensor<8x128x16x64xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<8x128x16x64>, layout = #dtensor.layout<sharding_specs:x,unsharded,y,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x16x64xf32>) -> tensor<8x128x16x64xf32>
tf_device.return %4 : tensor<8x128x16x64xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<8x128x16x64xf32>
func.return %0 : tensor<8x128x16x64xf32>
}
// -----
// Example from BERT 64 way sharding.
// bfd,bfi->id x,*,y . x,*,y -> y,*
// Unshard arg1 along y, reduce on output.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<8x128x128xf32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU"},
%arg2: tensor<8x128x256xf32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<256x128xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-NEXT: %[[EINSUM_RESULT:[0-9]*]] = "tf.Einsum"(%[[GATHERED]], %arg2)
// CHECK: %[[RETURN:.*]] = "tf.DTensorAllReduce"(%[[EINSUM_RESULT]]
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[RETURN]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x128x128>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128xf32>) -> tensor<8x128x128xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x128x256>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x256xf32>) -> tensor<8x128x256xf32>
%3 = "tf.Einsum"(%1, %2) {equation="bfd,bfi->id"} : (tensor<8x128x128xf32>, tensor<8x128x256xf32>) -> tensor<256x128xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<256x128>, layout = #dtensor.layout<sharding_specs:y,unsharded, mesh:TPU|x=2,y=2|*TPU>} : (tensor<256x128xf32>) -> tensor<256x128xf32>
tf_device.return %4 : tensor<256x128xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<256x128xf32>
func.return %0 : tensor<256x128xf32>
}
// -----
// Example from BERT 64 way sharding.
// bfi,bfd->di x,*,y . x,*,y -> *,y
// Unshard arg2 along y, reduce on output.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<8x128x256xf32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU"},
%arg2: tensor<8x128x128xf32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU"}) -> tensor<128x256xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg2)
// CHECK-NEXT: %[[EINSUM_RESULT:[0-9]*]] = "tf.Einsum"(%arg1, %[[GATHERED]])
// CHECK: %[[RETURN:.*]] = "tf.DTensorAllReduce"(%[[EINSUM_RESULT]]
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[RETURN]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x128x256>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x256xf32>) -> tensor<8x128x256xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x128x128>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<8x128x128xf32>) -> tensor<8x128x128xf32>
%3 = "tf.Einsum"(%1, %2) {equation="bfi,bfd->di"} : (tensor<8x128x256xf32>, tensor<8x128x128xf32>) -> tensor<128x256xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<128x256>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|*TPU>} : (tensor<128x256xf32>) -> tensor<128x256xf32>
tf_device.return %4 : tensor<128x256xf32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> tensor<128x256xf32>
func.return %0 : tensor<128x256xf32>
}
tensorflow/dtensor/mlir/tests/spmd_expansion.mlir
+852
View File
@@ -0,0 +1,852 @@
// 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: dtensor-opt -- %s -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
module @test_spmd_expect_error_no_layout {
func.func @main() {
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
%2 = "tf.B"() : () -> tensor<i32>
// expected-error @+1 {{No attached layout found for op : tf.Add}}
%3 = "tf.Add"(%1, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {_mesh = "TPU|x=2,y=1|*TPU"} : () -> (tensor<i32>)
func.return
}
}
// -----
module @test_spmd_malformed_layouts {
// Check that ops with malformed layouts are disallowed.
func.func @main() {
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<i32>
%2 = "tf.B"() : () -> tensor<i32>
// expected-error @+1 {{Expected 2 items but found}}
%3 = "tf.Add"(%1, %2) {_layout = [",,"]}: (tensor<i32>, tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {_mesh = "TPU|x=2,y=1|*TPU"} : () -> (tensor<i32>)
func.return
}
}
// -----
module @test_spmd_operands_without_layouts {
// Check operands without layouts are disallowed.
func.func @main(%arg0: tensor<i32>) {
%0 = "tf_device.cluster"() ({
%1 = "tf.A"() : () -> tensor<2x2xi32>
%2 = "tf.B"() : () -> tensor<2x2xi32>
// expected-error @+1 {{input layout of elementwise op must be known before SPMD expansion}}
%3 = "tf.Add"(%1, %2) {_layout = ["sharding_specs:x,y, TPU|x=2,y=1|*TPU"]}: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %3 : tensor<2x2xi32>
}) {_mesh = "TPU|x=2,y=1|*TPU"} : () -> (tensor<2x2xi32>)
func.return
}
}
// -----
// Check SPMD is skipped for layouts with XLA SPMD mesh.
//
// Arguments and ops and Retvals should remain in global shape.
module @test_spmd_skipped_for_xla {
func.func @main(%arg0: tensor<8x8xi32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7|use_xla_spmd"}) -> (tensor<8x8xi32>) {
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%2 = "tf.Identity"(%1) {_global_shape = [#tf_type.shape<8x8>], device = ""} : (tensor<8x8xi32>) -> tensor<8x8xi32>
%3= "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<8x8>, layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7|use_xla_spmd>} : (tensor<8x8xi32>) -> tensor<8x8xi32>
tf_device.return %3 : tensor<8x8xi32>
}) {_mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7|use_xla_spmd"} : () -> tensor<8x8xi32>
return %0 : tensor<8x8xi32>
}
}
// -----
// Check that elementwise batch parallel op SPMD expansion.
// CHECK-LABEL: module @test_spmd_batch_parallel
module @test_spmd_batch_parallel {
func.func @main(
%arg0: tensor<2x2xi32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|*TPU"},
%arg1: tensor<2x2xi32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|*TPU"}) {
%0 = "tf_device.cluster"() ({
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Add"
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
%3 = "tf.Add"(%arg0, %arg1) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|*TPU"]}: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %3 : tensor<2x2xi32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> (tensor<2x2xi32>)
func.return
}
}
// -----
// Check tf.Add SPMD with sharded inputs/outputs
// CHECK-LABEL: module @test_spmd_sharded_inputs
module @test_spmd_sharded_inputs {
func.func @main(
%arg0: tensor<2x2xi32> { tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"},
%arg1: tensor<2x2xi32> { tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[ADD_OUT:.*]] = "tf.Add"
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[ADD_OUT]]
%0 = "tf_device.cluster"() ({
%3 = "tf.Add"(%arg0, %arg1) {_layout = ["sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"]}: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %3 : tensor<2x2xi32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> (tensor<i32>)
func.return
}
}
// -----
// Check tf.Neg Op SPMD.
// CHECK-LABEL: module @test_spmd_neg_op_x_y
module @test_spmd_neg_op_x_y {
func.func @main(
%arg0: tensor<2x2xi32> { tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[NEG_OUT:.*]] = "tf.Neg"
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[NEG_OUT]]
%0 = "tf_device.cluster"() ({
%2 = "tf.Neg"(%arg0) {_layout = ["sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"]}: (tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %2 : tensor<2x2xi32>
}) {_mesh = "TPU|x=2,y=2|*TPU"} : () -> (tensor<i32>)
func.return
}
}
// -----
// Check replicated tf.Const op SPMD.
// CHECK-LABEL: module @test_spmd_const_op_unsharded
module @test_spmd_const_op_unsharded {
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.Const"
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.Const"
// CHECK-NEXT: %[[ADD_OUT:.*]] = "tf.Add"(%[[A_OUT]], %[[B_OUT]])
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[ADD_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value=dense<1> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
%2 = "tf.Const"() {_layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value=dense<1> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
%3 = "tf.Add"(%1, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]}: (tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<1x1xi32>
tf_device.return %3 : tensor<1x1xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<1x1xi32>
func.return
}
}
// -----
// Check sharded tf.Const op SPMD.
// CHECK-LABEL: module @test_spmd_const_op_sharded
module @test_spmd_const_op_sharded {
// CHECK: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<i32>
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<1>}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[A_OUT:.*]] = "tf.Const"
// CHECK-NEXT: %[[A_SLICE:[0-9]*]] = "tf.DTensorAllScatter"(%[[A_OUT]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: %[[IDENTITY_OUT:[0-9]*]] = "tf.IdentityN"(%[[A_SLICE]])
// CHECK-NEXT: %[[B_OUT:.*]] = "tf.Const"
// CHECK-NEXT: %[[ADD_OUT:[0-9]*]] = "tf.Add"(%[[IDENTITY_OUT]], %[[B_OUT]])
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[ADD_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2xi32>) -> (tensor<2xi32>)
%3 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<1>, layout = #dtensor.layout<sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<i32>) -> (tensor<i32>)
%5 = "tf.Add"(%2, %4) {_layout = ["sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]}: (tensor<2xi32>, tensor<i32>) -> tensor<2xi32>
tf_device.return %5 : tensor<2xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<2xi32>
func.return
}
}
// -----
// Check sharded tf.Const op SPMD with splat.
// CHECK-LABEL: module @test_spmd_const_op_sharded_with_splat
module @test_spmd_const_op_sharded_with_splat {
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CONST_OUT:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {{{.*}}} : () -> tensor<1xi32>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {_layout = ["sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], value=dense<1>: tensor<2xi32>} : () -> tensor<2xi32>
tf_device.return %1 : tensor<2xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<2xi32>
func.return
}
}
// -----
// Check replicated tf.BroadcastTo op SPMD.
// CHECK-LABEL: module @test_spmd_broadcast_replicated
module @test_spmd_broadcast_replicated {
func.func @main(%arg0: tensor<3xi32>) {
// CHECK: "tf.BroadcastTo"
// CHECK-SAME: tensor<3xi32>, tensor<2xi64>) -> tensor<3x3xi32>
%0 = "tf_device.cluster"() ({
%1 = arith.constant dense<[3, 3]> : tensor<2xi32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2xi32>) -> (tensor<2xi32>)
%3 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<3>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<3xi32>) -> (tensor<3xi32>)
%4 = "tf.BroadcastTo"(%3, %2) {_layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], device = ""} : (tensor<3xi32>, tensor<2xi32>) -> tensor<3x3xi32>
tf_device.return %4 : tensor<3x3xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<3x3xi32>
func.return
}
}
// -----
// Check replicated tf.range op SPMD.
// CHECK-LABEL: module @test_spmd_range
module @test_spmd_range {
func.func @main() {
// CHECK: "tf.Range"
// CHECK-SAME: tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<3xi32>
%0 = "tf_device.cluster"() ({
%0 = arith.constant dense<0> : tensor<i32>
%1 = arith.constant dense<3> : tensor<i32>
%2 = arith.constant dense<1> : tensor<i32>
%3 = "tf.Range"(%0, %1, %2) {_layout = ["sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<3xi32>
tf_device.return %3 : tensor<3xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<3xi32>
func.return
}
}
// -----
// Check tf.AssignVariable op SPMD
// CHECK-LABEL: module @test_spmd_assign_var
module @test_spmd_assign_var {
func.func @main(%arg0: tensor<32x32xi32> { tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}, %arg1: tensor<!tf_type.resource> { tf._layout = "empty_layout" }) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: tf_device.return
// CHECK-NEXT: _inferred_resource_indices = dense<1> : vector<1xi32>
// CHECK-SAME: _inferred_resource_layouts
// CHECK-SAME: "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
"tf_device.cluster"() ({
"tf.AssignVariableOp"(%arg1, %arg0) {dtype = i32} : (tensor<!tf_type.resource>, tensor<32x32xi32>) -> ()
tf_device.return
}) { _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check tf.Softmax op SPMD where input layout and output layout are different. DTensor needs a relayout.
// CHECK-LABEL: module @test_spmd_softmax_need_relayout
module @test_spmd_softmax_need_relayout {
func.func @main(%arg0: tensor<32x32xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Softmax"
// CHECK-NEXT: "tf.DTensorAllGather"
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<32x32>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<32x32xf32>) -> (tensor<32x32xf32>)
%2 = "tf.Softmax"(%1) {_global_shape = [#tf_type.shape<32x32>]} : (tensor<32x32xf32>) -> (tensor<32x32xf32>)
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<32x32>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<32x32xf32>) -> (tensor<32x32xf32>)
tf_device.return %3 : tensor<32x32xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<32x32xf32>)
func.return
}
}
// -----
// Check tf.Softmax op SPMD where last dimension is not sharded.
// CHECK-LABEL: module @test_spmd_softmax_last_dim_unsharded
module @test_spmd_softmax_last_dim_unsharded {
func.func @main(%arg0: tensor<32x32xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Softmax"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
"tf.Softmax"(%arg0) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : (tensor<32x32xf32>) -> (tensor<32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check tf.Softmax op with rank 3.
// CHECK-LABEL: module @test_spmd_softmax_rank_3
module @test_spmd_softmax_rank_3 {
func.func @main(%arg0: tensor<32x32x32xf32> { tf._layout = "sharding_specs:x,y,unsharded, mesh:TPU|x=2,y=2,z=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Softmax"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
"tf.Softmax"(%arg0) {_layout = ["sharding_specs:x,y,unsharded, mesh:TPU|x=2,y=2,z=2|*TPU"]} : (tensor<32x32x32xf32>) -> (tensor<32x32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2,z=2|*TPU"} : () -> ()
func.return
}
}
// -----
// Check SPMD expansion of softmax op with non-sharded last dimension.
// CHECK-LABEL: module @test_spmd_softmax_last_dim_unsharded1
module @test_spmd_softmax_last_dim_unsharded1 {
func.func @main(%arg0: tensor<32x32xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.LogSoftmax"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
"tf.LogSoftmax"(%arg0) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : (tensor<32x32xf32>) -> (tensor<32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check SPMD expansion of SoftMax op.
// CHECK-LABEL: module @test_spmd_log_softmax
module @test_spmd_log_softmax {
func.func @main(%arg0: tensor<32x32x32xf32> { tf._layout = "sharding_specs:x,y,unsharded, mesh:TPU|x=2,y=2,z=2|*TPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.LogSoftmax"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
"tf.LogSoftmax"(%arg0) {_layout = ["sharding_specs:x,y,unsharded, mesh:TPU|x=2,y=2,z=2|*TPU"]} : (tensor<32x32x32xf32>) -> (tensor<32x32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2,z=2|*TPU"} : () -> ()
func.return
}
}
// -----
// Check that Softmax op with last dimension sharded is supported on TPU's.
// CHECK-LABEL: module @test_spmd_softmax_last_dim_sharded
module @test_spmd_softmax_last_dim_sharded {
func.func @main(%arg0: tensor<32x32x32xf32> { tf._layout = "sharding_specs:x,y,z, mesh:TPU|x=2,y=2,z=2|*TPU"}) {
"tf_device.cluster"() ({
"tf.Softmax"(%arg0) {_layout = ["sharding_specs:x,y,z, mesh:TPU|x=2,y=2,z=2|*TPU"]} : (tensor<32x32x32xf32>) -> (tensor<32x32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2,z=2|*TPU"} : () -> ()
func.return
}
}
// -----
// Check that random uniform op with incompatible shape is disallowed.
module @test_spmd_random_op_with_incomplete_shape_disallowed {
func.func @main(%arg0: tensor<i32>) {
%0 = "tf_device.cluster"() ({
// %1 = "tf.Const"() <{value = dense<16> : tensor<2xi32>}> {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], device = ""} : () -> tensor<2xi32>
// %2 = "tf.Const"() <{value = dense<[123, 321]> : tensor<2xi32>}> {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], device = ""} : () -> tensor<2xi32>
%1 = arith.constant dense<[16]> : tensor<1xi32>
%2 = arith.constant dense<[2, 1]> : tensor<2xi32>
// expected-error @+1 {{Sharding dimension of random op does not match rank of the random op}}
%3 = "tf.StatelessRandomUniform"(%1, %2) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], device = ""} : (tensor<1xi32>, tensor<2xi32>) -> tensor<16x16xf32>
tf_device.return %3 : tensor<16x16xf32>
}) { _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<16x16xf32>
func.return
}
}
// -----
// Check Resource Apply op SPMD.
// CHECK-LABEL: module @test_spmd_resource
module @test_spmd_resource {
// CHECK: func @main
// CHECK-SAME: %arg0: tensor<f32>
// CHECK-SAME: %arg1: tensor<1x1xf32>
// CHECK-SAME: tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-SAME: %arg2: tensor<!tf_type.resource>
// CHECK-SAME: tf._layout = "empty_layout"
func.func @main(
%arg0: tensor<f32>,
%arg1: tensor<2x2xf32>{ tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<!tf_type.resource> { tf._layout = "empty_layout"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.ResourceApplyGradientDescent"(%arg2, %arg0, %arg1)
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
"tf.ResourceApplyGradientDescent"(%arg2, %arg0, %arg1) {_layout = ["sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"], use_locking = false} : (tensor<!tf_type.resource>, tensor<f32>, tensor<2x2xf32>) -> ()
tf_device.return
}) { _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check that function inputs are modified to reflect local input shapes.
// CHECK-LABEL: module @test_spmd_inputs_have_local_shapes
module @test_spmd_inputs_have_local_shapes {
// CHECK: func @main
// CHECK-SAME: %arg0: tensor<f32>
// CHECK-SAME: %arg1: tensor<1x1xf32>
// CHECK-SAME: tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-SAME: %arg2: tensor<!tf_type.resource>
// CHECK-SAME: tf._layout = "empty_layout"
func.func @main(
%arg0: tensor<f32>,
%arg1: tensor<2x2xf32>{ tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<!tf_type.resource> { tf._layout = "empty_layout"}) {
"tf_device.cluster"() ({
tf_device.return
}) {_mesh = "mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check that tf_device.Cluster op return values are updated to reflect local
// shape.
// CHECK-LABEL: module @test_spmd_returned_values_have_local_shapes
module @test_spmd_returned_values_have_local_shapes {
// CHECK: func @main
// CHECK-SAME: %arg0: tensor<f32>
// CHECK-SAME: %arg1: tensor<1x1xf32>
// CHECK-SAME: tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-SAME: %arg2: tensor<!tf_type.resource>
// CHECK-SAME: tf._layout = "empty_layout"
func.func @main(
%arg0: tensor<f32>,
%arg1: tensor<2x2xf32>{ tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<!tf_type.resource> { tf._layout = "empty_layout"}) {
"tf_device.cluster"() ({
// CHECK: tf_device.return
// CHECK-SAME: tensor<1x1xf32>
tf_device.return %arg1 : tensor<2x2xf32>
}) {_mesh = "mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<2x2xf32>)
func.return
}
}
// -----
// Check that function signature as well as return types of callsite operations
// are updated to reflect local shape.
// CHECK-LABEL: module @test_spmd_return_types_have_local_shapes_at_callsite
module @test_spmd_return_types_have_local_shapes_at_callsite {
// CHECK: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<f32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<1x1xf32>
// CHECK-SAME: tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"
// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: tensor<!tf_type.resource>
// CHECK-SAME: tf._layout = "empty_layout"
func.func @main(
%arg0: tensor<f32>,
%arg1: tensor<2x2xf32>{ tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<!tf_type.resource> { tf._layout = "empty_layout"}) {
"tf_device.cluster"() ({
// CHECK: "tf.StatefulPartitionedCall"(%[[ARG1]])
// CHECK-SAME: (tensor<1x1xf32>) -> tensor<1x1xf32>
%0 = "tf.StatefulPartitionedCall"(%arg1) {config = "", config_proto = "", executor_type = "", f = @pcall_func} : (tensor<2x2xf32>) -> (tensor<2x2xf32>)
tf_device.return %arg1 : tensor<2x2xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<2x2xf32>)
func.return
}
// CHECK: func @pcall_func
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<1x1xf32>
func.func @pcall_func(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
// CHECK: return %[[ARG0]] : tensor<1x1xf32>
func.return %arg0 : tensor<2x2xf32>
}
}
// -----
// Check DTensorLayout ops are removed after SPMD Expansion.
// CHECK-LABEL: module @test_spmd_dtensor_layout_ops_are_removed
module @test_spmd_dtensor_layout_ops_are_removed {
// CHECK: func @main
func.func @main(%arg0: tensor<i32>) {
%0 = "tf_device.cluster"() ({
// CHECK: "tf.Const"()
// CHECK-NOT: "tf.DTensorLayout"
// CHECK: "tf.Const"()
// CHECK-NOT: "tf.DTensorLayout"
// CHECK: "tf.Add"
// CHECK-NOT: "tf.DTensorLayout"
// CHECK-NEXT: tf_device.return
%1 = "tf.Const"() {value = dense<[[4, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
%3 = "tf.Const"() {value = dense<[[1, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
%5 = "tf.Add"(%2, %4): (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
%6 = "tf.DTensorLayout"(%5) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
tf_device.return %6 : tensor<2x2xi32>
}) {_mesh = "mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<2x2xi32>)
func.return
}
}
// -----
// CHECK-LABEL: module @test_spmd_neg_op_unsharded
module @test_spmd_neg_op_unsharded {
// CHECK: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<i32>
// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: tensor<2x2xf32>
func.func @main(
%arg0: tensor<i32>, %arg1: tensor<2x2xf32>) {
"tf_device.cluster"() ({
// CHECK: %[[ARG1_SLICE:[0-9]*]] = "tf.DTensorAllScatter"(%[[ARG1]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>
// CHECK-NEXT: %[[NEG_OUT:[0-9]*]] = "tf.Neg"(%[[ARG1_SLICE]])
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[NEG_OUT]]
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xf32>) -> (tensor<2x2xf32>)
%2 = "tf.Neg"(%1) : (tensor<2x2xf32>) -> tensor<2x2xf32>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xf32>) -> (tensor<2x2xf32>)
tf_device.return %3: tensor<2x2xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<2x2xf32>)
func.return
}
}
// -----
// A super tricky case where the DTensorLayout is out of the tf_device.cluster and somewhat gets casted to BlockArgument with a wild argument number 3.
// CHECK-LABEL: module @test_spmd_var_input
module @test_spmd_var_input {
// CHECK: func @main
// CHECK: %arg1: tensor<!tf_type.resource<tensor<1xf32>>>
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>},
%arg1: tensor<!tf_type.resource<tensor<2xf32>>> {tf._global_shape = #tf_type.shape<2>, tf._layout = "empty_layout", tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<2xf32> {tf._global_shape = #tf_type.shape<2>,
tf._layout = "sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg2) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xf32>) -> tensor<2xf32>
"tf.AssignVariableOp"(%0, %1) {_global_shape = [], device = ""} : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
tf_device.return
}) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> ()
func.return
}
}
// -----
// A super tricky case where the DTensorLayout is out of the tf_device.cluster and somewhat gets casted to BlockArgument with a wild argument number 3.
// CHECK-LABEL: module @test_spmd_assigned_value_is_input
module @test_spmd_assigned_value_is_input {
// CHECK: func @main
// CHECK: "tf.VarHandleOp"()
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
func.func @main(%arg0: tensor<2xf32> {tf._global_shape = #tf_type.shape<2>,
tf._layout = "sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// %1 = "tf.DTensorLayout"(%0) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf_device.cluster"() ({
%0 = "tf.VarHandleOp"() {_global_shape = [#tf_type.shape<2>], allowed_devices = [], container = "", device = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
%1 = "tf.DTensorLayout"(%0) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<!tf_type.resource<tensor<2xf32>>>
%2 = "tf.DTensorLayout"(%arg0) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xf32>) -> tensor<2xf32>
"tf.AssignVariableOp"(%1, %2) {_global_shape = [], device = ""} : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
tf_device.return
}) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> ()
func.return
}
}
// -----
// Check to ensure that the local shape of resource-type arguments are not double-calculated if they are assigned to a tensor value wihtin the function.
// CHECK-LABEL: module @test_spmd_var_arg_local_shapes
module @test_spmd_var_arg_local_shapes {
// CHECK: func @main
// CHECK-SAME: %arg0: tensor<i32>
// CHECK-SAME: %arg1: tensor<1x4xf32>
// CHECK-SAME: %arg2: tensor<!tf_type.resource<tensor<1x4xf32>>>
func.func @main(
%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>},
%arg1: tensor<2x4xf32> {
tf._global_shape = #tf_type.shape<2x4>,
tf._layout = "sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<*x!tf_type.resource<tensor<2x4xf32>>> {
tf._global_shape = #tf_type.shape<2x4>,
tf._layout = "sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"})
-> (tensor<2x4xf32> {tf._global_shape = #tf_type.shape<2x4>}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg2) {_global_shape = [#tf_type.shape<*>], global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<*x!tf_type.resource<tensor<2x4xf32>>>) -> tensor<*x!tf_type.resource<tensor<2x4xf32>>>
%2 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<2x4>], global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
"tf.AssignVariableOp"(%1, %2) {_global_shape = [], device = ""} : (tensor<*x!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
%3 = "tf.ReadVariableOp"(%1) {_global_shape = [#tf_type.shape<2x4>], device = ""} : (tensor<*x!tf_type.resource<tensor<2x4xf32>>>) -> tensor<2x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%5 = "tf.Identity"(%4) {_global_shape = [#tf_type.shape<2x4>], device = ""} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%6 = "tf.DTensorLayout"(%5) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
tf_device.return %6 : tensor<2x4xf32>
}) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x4xf32>
func.return %0 : tensor<2x4xf32>
}
}
// -----
// Check SPMD expansion of Cumsum op with sharding on axis dimension, should
// produce a replicated layout on that axis dimension, with allgather and
// allscatter for intermediate layout computation.
// CHECK-LABEL: module @test_spmd_cumsum_op_sharded
module @test_spmd_cumsum_op_sharded {
// CHECK: func @main
func.func @main(%arg0: tensor<32x32xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.DTensorAllGather"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Cumsum"
// CHECK-NEXT: "tf.DTensorAllScatter"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
%axis = "tf.Const"() { value = dense<0> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<i64>
"tf.Cumsum"(%arg0, %axis) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : (tensor<32x32xf32>, tensor<i64>) -> (tensor<32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check SPMD expansion of Cumsum op with no sharding on axis dim. This should
// not produce an allscatter or allgather for intermediate layout computation
// since no relayouts are happening.
// CHECK-LABEL: module @test_spmd_cumsum_op_no_sharding_on_axis_dim
module @test_spmd_cumsum_op_no_sharding_on_axis_dim {
// CHECK: func @main
func.func @main(%arg0: tensor<32x32xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Cumsum"
// CHECK-NEXT: tf_device.return
"tf_device.cluster"() ({
%axis = "tf.Const"() { value = dense<-1> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<i64>
"tf.Cumsum"(%arg0, %axis) {_layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : (tensor<32x32xf32>, tensor<i64>) -> (tensor<32x32xf32>)
tf_device.return
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
func.return
}
}
// -----
// Check Relayout for SparseTensors emits the appropriate ops required for relaying out a SparseTensor.
// We do this by doing a matmul (between a sparsetensor and a densetensor)
// tf.matmul (*,x) multiplied by (x,*) causes a relayout on the left operand.
// CHECK-LABEL: module @test_spmd_relayout_on_sparse_tensors
module @test_spmd_relayout_on_sparse_tensors {
// CHECK: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<4x16xf32> {tf._layout = "sharding_specs:unsharded,x, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"}, %arg2: tensor<?x2xi64>, %arg3: tensor<2xi64>, %arg4: tensor<?xf32>) -> tensor<8x16xf32> {
// CHECK: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[DENSE_0:.*]] = "tf.SparseToDense"(%arg2, %arg3, %arg4, %[[CST]])
// CHECK-NEXT: %[[RIGHT_OPERAND:.*]] = "tf.DTensorAllScatter"(%arg1)
// CHECK-NEXT: %[[ONE:.*]] = "tf.DTensorAllGather"(%[[DENSE_0]])
// CHECK-NEXT: %[[TWO:.*]] = "tf.ZerosLike"(%[[ONE]])
// CHECK-NEXT: %[[THREE:.*]] = "tf.NotEqual"(%[[ONE]], %[[TWO]])
// CHECK-NEXT: %[[WHERE:.*]] = "tf.Where"(%[[THREE]])
// CHECK-NEXT: %[[GATHER:.*]] = "tf.GatherNd"(%[[ONE]], %[[WHERE]])
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.Shape"(%[[ONE]])
// CHECK-NEXT: %[[CST_0:.*]] = "tf.Const"
// CHECK-NEXT: %[[LEFT_OPERAND:.*]] = "tf.SparseToDense"(%[[WHERE]], %[[SHAPE]], %[[GATHER]], %[[CST_0]])
// CHECK-NEXT: "tf.MatMul"(%[[LEFT_OPERAND]], %[[RIGHT_OPERAND]])
%0 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<?x2>, layout = #dtensor.layout<sharding_specs:x,batch, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3>} : (tensor<?x2xi64>) -> tensor<?x2xi64>
%1 = "tf_device.cluster"() ({
%cst = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%2 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3>} : (tensor<f32>) -> tensor<f32>
%3 = "tf.SparseToDense"(%0, %arg3, %arg4, %2) {_global_shape = [#tf_type.shape<8x4>]} : (tensor<?x2xi64>, tensor<2xi64>, tensor<?xf32>, tensor<f32>) -> tensor<8x4xf32>
%4 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<4x16>], global_shape = #tf_type.shape<4x16>, layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3>} : (tensor<4x16xf32>) -> tensor<4x16xf32>
%5 = "tf.DTensorLayout"(%3) {_global_shape = [#tf_type.shape<8x4>], global_shape = #tf_type.shape<8x4>, layout = #dtensor.layout<sharding_specs:x,batch, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3>} : (tensor<8x4xf32>) -> tensor<8x4xf32>
%6 = "tf.MatMul"(%5, %4) {_global_shape = [#tf_type.shape<8x16>], device = "", transpose_a = false, transpose_b = false} : (tensor<8x4xf32>, tensor<4x16xf32>) -> tensor<8x16xf32>
%7 = "tf.DTensorLayout"(%6) {global_shape = #tf_type.shape<8x16>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3>} : (tensor<8x16xf32>) -> tensor<8x16xf32>
%8 = "tf.Identity"(%7) {_global_shape = [#tf_type.shape<8x16>], device = ""} : (tensor<8x16xf32>) -> tensor<8x16xf32>
%9 = "tf.DTensorLayout"(%8) {global_shape = #tf_type.shape<8x16>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3>} : (tensor<8x16xf32>) -> tensor<8x16xf32>
tf_device.return %9 : tensor<8x16xf32>
}) {_mesh = "|batch=2,x=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3"} : () -> tensor<8x16xf32>
func.return %1 : tensor<8x16xf32>
}
}
// -----
// Check that relayout uses all-to-all for unsharded,x to x,unsharded.
// CHECK-LABEL: module @test_relayout_using_all_to_all
module @test_relayout_using_all_to_all {
// CHECK: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<32x32xf32> { tf._layout = "sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<32x32xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[CST:.*]] = "tf.Const"
// CHECK-NEXT: %[[BIAS_ADD_OUT:.*]] = "tf.BiasAdd"(%arg1, %cst)
// CHECK-NEXT: %[[ALL_TO_ALL_OUT:.*]] = "tf.DTensorAllToAll"(%[[BIAS_ADD_OUT]])
// CHECK-SAME: input_layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>
// CHECK-SAME: output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<32xf32>, _layout = ["sharding_specs:x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : () -> tensor<32xf32>
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<32x32>, layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>} : (tensor<32x32xf32>) -> tensor<32x32xf32>
%2 = "tf.BiasAdd"(%1, %cst) {global_shape = #tf_type.shape<32x32>} : (tensor<32x32xf32>, tensor<32xf32>) -> (tensor<32x32xf32>)
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<32x32>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3>} : (tensor<32x32xf32>) -> tensor<32x32xf32>
tf_device.return %3 : tensor<32x32xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<32x32xf32>)
func.return %0 : tensor<32x32xf32>
}
}
// -----
// Check SPMD expansion of TensorListReserve replicated and TensorListSet with a sharded tensor emits a gather to replicated.
// CHECK-LABEL: module @test_spmd_tensor_list_reserve_replicated
module @test_spmd_tensor_list_reserve_replicated {
// CHECK: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<4x4xi32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> (tensor<4x4xi32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.DTensorAllGather"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: "tf.TensorListReserve"
// CHECK-NEXT: "tf.TensorListSetItem"
// CHECK-NEXT: "tf.TensorListGetItem"
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<4> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%cst_0 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<4> : tensor<2xi32>} : () -> tensor<2xi32>
%2 = "tf.DTensorLayout"(%cst_0) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%cst_1 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<4> : tensor<i32>} : () -> tensor<i32>
%3 = "tf.DTensorLayout"(%cst_1) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<i32>) -> tensor<i32>
%cst_2 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<0> : tensor<i32>} : () -> tensor<i32>
%4 = "tf.DTensorLayout"(%cst_2) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<i32>) -> tensor<i32>
%cst_3 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<0> : tensor<i32>} : () -> tensor<i32>
%5 = "tf.DTensorLayout"(%cst_3) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<i32>) -> tensor<i32>
%6 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<4x4>], global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<4x4xi32>) -> tensor<4x4xi32>
%7 = "tf.TensorListReserve"(%1, %3) {_global_shape = [#tf_type.shape<>], device = ""} : (tensor<2xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<4x4xi32>>>
%8 = "tf.DTensorLayout"(%7) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.variant<tensor<4x4xi32>>>) -> tensor<!tf_type.variant<tensor<4x4xi32>>>
%9 = "tf.TensorListSetItem"(%8, %4, %6) {_global_shape = [#tf_type.shape<>], device = ""} : (tensor<!tf_type.variant<tensor<4x4xi32>>>, tensor<i32>, tensor<4x4xi32>) -> tensor<!tf_type.variant<tensor<4x4xi32>>>
%10 = "tf.DTensorLayout"(%9) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.variant<tensor<4x4xi32>>>) -> tensor<!tf_type.variant<tensor<4x4xi32>>>
%11 = "tf.TensorListGetItem"(%10, %5, %2) {_global_shape = [#tf_type.shape<4x4>], device = ""} : (tensor<!tf_type.variant<tensor<4x4xi32>>>, tensor<i32>, tensor<2xi32>) -> tensor<4x4xi32>
%12 = "tf.DTensorLayout"(%11) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<4x4xi32>) -> tensor<4x4xi32>
%13 = "tf.Identity"(%12) {_global_shape = [#tf_type.shape<4x4>], device = ""} : (tensor<4x4xi32>) -> tensor<4x4xi32>
%14 = "tf.DTensorLayout"(%13) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<4x4xi32>) -> tensor<4x4xi32>
tf_device.return %14 : tensor<4x4xi32>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<4x4xi32>
func.return %0 : tensor<4x4xi32>
}
}
// -----
// Check SPMD expansion of DisableCopyOnRead has correct shape.
// CHECK-LABEL: module @test_spmd_disable_copy_on_read
module @test_spmd_disable_copy_on_read {
// CHECK: func @main
func.func @main(
%arg0: tensor<i32>,
%arg1: tensor<!tf_type.resource<tensor<4x8xi32>>> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> () {
// CHECK: "tf_device.cluster"
// CHECK: "tf.DisableCopyOnRead"(%arg1) {_global_shape = [], _layout = [], device = ""} : (tensor<!tf_type.resource<tensor<2x8xi32>>>) -> ()
"tf_device.cluster"() ({
%0 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x8>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.resource<tensor<4x8xi32>>>) -> tensor<!tf_type.resource<tensor<4x8xi32>>>
"tf.DisableCopyOnRead"(%0) {device = ""} : (tensor<!tf_type.resource<tensor<4x8xi32>>>) -> ()
tf_device.return
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> ()
func.return
}
}
// -----
// Check SPMD expansion of ScatterNd op output is the sharding of updates
// tensor.
// CHECK-LABEL: module @test_spmd_scatter_nd_op
module @test_spmd_scatter_nd_op {
// CHECK: func @main
func.func @main(%arg0: tensor<2x4x4xi32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"}) -> (tensor<16x4x4xi32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"()
// CHECK-NEXT: %[[INDICES:.*]] = "tf.Const"()
// CHECK-NEXT: %[[NEW_SHAPE:.*]] = "tf.Const"() <{value = dense<[16, 2, 4]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK-NEXT: "tf.ScatterNd"(%[[INDICES]], %arg0, %[[NEW_SHAPE]])
%0 = "tf_device.cluster"() ({
%shape = "tf.Const"() {_global_shape = [#tf_type.shape<3>], value = dense<[16, 4, 4]> : tensor<3xi32>} : () -> tensor<3xi32>
%indices = "tf.Const"() {_global_shape = [#tf_type.shape<2x1>], value = dense<[[0], [15]]> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
%updates_with_layout = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<2x4x4>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<2x4x4xi32>) -> tensor<2x4x4xi32>
%indices_with_layout= "tf.DTensorLayout"(%indices) {global_shape = #tf_type.shape<2x1>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<2x1xi32>) -> tensor<2x1xi32>
%shape_with_layout = "tf.DTensorLayout"(%shape) {global_shape = #tf_type.shape<3>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<3xi32>) -> tensor<3xi32>
%4 = "tf.ScatterNd"(%indices_with_layout, %updates_with_layout, %shape_with_layout) {_global_shape = [#tf_type.shape<16x4x4>], device = ""} : (tensor<2x1xi32>, tensor<2x4x4xi32>, tensor<3xi32>) -> tensor<16x4x4xi32>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<16x4x4>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<16x4x4xi32>) -> tensor<16x4x4xi32>
tf_device.return %5 : tensor<16x4x4xi32>
}) {_mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"} : () -> tensor<16x4x4xi32>
return %0 : tensor<16x4x4xi32>
}
}
// -----
// Check SPMD expansion of ScatterNd op indices is relayout to replicated.
// CHECK-LABEL: module @test_spmd_scatter_nd_op_indices_layout
module @test_spmd_scatter_nd_op_indices_layout {
// CHECK: func @main
func.func @main(%arg0: tensor<2x4x4xi32> {tf._layout = "sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"},
%arg1: tensor<2x1xi32> {tf._global_shape = #tf_type.shape<2x1>, tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"}) -> (tensor<16x4x4xi32>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.DTensorAllGather"(%arg1)
// CHECK: "tf.ScatterNd"
%0 = "tf_device.cluster"() ({
%shape = "tf.Const"() {_global_shape = [#tf_type.shape<3>], value = dense<[16, 4, 4]> : tensor<3xi32>} : () -> tensor<3xi32>
%updates_with_layout = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<2x4x4>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<2x4x4xi32>) -> tensor<2x4x4xi32>
%indices_with_layout= "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x1>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<2x1xi32>) -> tensor<2x1xi32>
%shape_with_layout = "tf.DTensorLayout"(%shape) {global_shape = #tf_type.shape<3>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<3xi32>) -> tensor<3xi32>
%4 = "tf.ScatterNd"(%indices_with_layout, %updates_with_layout, %shape_with_layout) {_global_shape = [#tf_type.shape<16x4x4>], device = ""} : (tensor<2x1xi32>, tensor<2x4x4xi32>, tensor<3xi32>) -> tensor<16x4x4xi32>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<16x4x4>, layout = #dtensor.layout<sharding_specs:unsharded,x,unsharded, mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7>} : (tensor<16x4x4xi32>) -> tensor<16x4x4xi32>
tf_device.return %5 : tensor<16x4x4xi32>
}) {_mesh = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1,/job:localhost/replica:0/task:0/device:CPU:2,/job:localhost/replica:0/task:0/device:CPU:3,/job:localhost/replica:0/task:0/device:CPU:4,/job:localhost/replica:0/task:0/device:CPU:5,/job:localhost/replica:0/task:0/device:CPU:6,/job:localhost/replica:0/task:0/device:CPU:7"} : () -> tensor<16x4x4xi32>
return %0 : tensor<16x4x4xi32>
}
}
// -----
// Check stateful random operations raise error.
module @test_spmd_error_on_stateful_random_uniform_op {
func.func @main(
%arg0: tensor<2xi32> { tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"}) {
%0 = "tf_device.cluster"() ({
// expected-error @+1 {{Stateful random operations are not supported in DTensor.}}
%1 = "tf.RandomUniform"(%arg0) {_layout = ["sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU"]}: (tensor<2xi32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU>} : (tensor<4x4xf32>) -> (tensor<4x4xf32>)
tf_device.return %2 : tensor<4x4xf32>
}) {_mesh = "CPU|x=2,y=2|*CPU"} : () -> (tensor<4x4xf32>)
func.return
}
}
// -----
// Check stateful random operations raise error.
module @test_spmd_error_on_stateful_random_uniform_int_op {
func.func @main(
%arg0: tensor<2xi32> { tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"},
%arg1: tensor<1xi32> { tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"},
%arg2: tensor<1xi32> { tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"}
) {
%0 = "tf_device.cluster"() ({
// expected-error @+1 {{Stateful random operations are not supported in DTensor.}}
%1 = "tf.RandomUniformInt"(%arg0, %arg1, %arg2) {_layout = ["sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU"]}: (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4x4xi32>
%2 = "tf.DTensorLayout"(%1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=2,y=2|*TPU>} : (tensor<4x4xi32>) -> (tensor<4x4xi32>)
tf_device.return %2 : tensor<4x4xi32>
}) {_mesh = "CPU|x=2,y=2|*CPU"} : () -> (tensor<4x4xi32>)
func.return
}
}
tensorflow/dtensor/mlir/tests/spmd_fft.mlir
+164
View File
@@ -0,0 +1,164 @@
// 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: dtensor-opt %s -split-input-file -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s --dump-input=fail
// Check the SPMD expansion for FFT2D
// CHECK-LABEL: module @test_FFT2D
module @test_FFT2D {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4x8xcomplex<f32>> {tf._layout = "sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[FFT_OUT_1:.*]] = "tf.FFT"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x4x8xcomplex<f32>>) -> tensor<2x4x8xcomplex<f32>>
// CHECK-NEXT: %[[CONST_OUT_1:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_1:.*]] = "tf.Transpose"(%[[FFT_OUT_1]], %[[CONST_OUT_1]])
// CHECK-SAME: (tensor<2x4x8xcomplex<f32>>, tensor<3xi64>) -> tensor<2x8x4xcomplex<f32>>
// CHECK-NEXT: %[[IDENT_OUT:.*]] = "tf.Identity"(%[[TRANS_OUT_1]])
// CHECK-NEXT: %[[FFT_OUT_2:.*]] = "tf.FFT"(%[[IDENT_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x8x4xcomplex<f32>>) -> tensor<2x8x4xcomplex<f32>>
// CHECK-NEXT: %[[CONST_OUT_2:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_2:.*]] = "tf.Transpose"(%[[FFT_OUT_2]], %[[CONST_OUT_2]])
// CHECK-SAME: (tensor<2x8x4xcomplex<f32>>, tensor<3xi64>) -> tensor<2x4x8xcomplex<f32>>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4x8>, layout = #dtensor.layout<sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x8xcomplex<f32>>) -> tensor<2x4x8xcomplex<f32>>
%2 = "tf.FFT2D"(%1) : (tensor<2x4x8xcomplex<f32>>) -> tensor<2x4x8xcomplex<f32>>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<2x4x8>, layout = #dtensor.layout<sharding_specs:b,unsharded,x, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x8xcomplex<f32>>) -> tensor<2x4x8xcomplex<f32>>
tf_device.return %3 : tensor<2x4x8xcomplex<f32>>
}) {_mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x4x8xcomplex<f32>>
func.return
}
}
// -----
// Check the SPMD expansion for IFFT2D
// CHECK-LABEL: module @test_IFFT2D
module @test_IFFT2D {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4x8xcomplex<f64>> {tf._layout = "sharding_specs:b,x,y, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_1:.*]] = "tf.Transpose"(%arg1, %[[CONST_OUT_1]])
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>, tensor<3xi64>) -> tensor<2x4x4xcomplex<f64>>
// CHECK-NEXT: %[[IDENT_OUT:.*]] = "tf.Identity"(%[[TRANS_OUT_1]])
// CHECK-NEXT: %[[IFFT_OUT_1:.*]] = "tf.IFFT"(%[[IDENT_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>) -> tensor<2x4x4xcomplex<f64>>
// CHECK-NEXT: %[[CONST_OUT_2:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_2:.*]] = "tf.Transpose"(%[[IFFT_OUT_1]], %[[CONST_OUT_2]])
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>, tensor<3xi64>) -> tensor<2x4x4xcomplex<f64>>
// CHECK-NEXT: %[[ALLTOALL_OUT:.*]] = "tf.DTensorAllToAll"(%[[TRANS_OUT_2]])
// CHECK-SAME: _layout = ["sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>) -> tensor<2x2x8xcomplex<f64>>
// CHECK-NEXT: %[[IFFT_OUT_2:.*]] = "tf.IFFT"(%[[ALLTOALL_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x2x8xcomplex<f64>>) -> tensor<2x2x8xcomplex<f64>>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4x8>, layout = #dtensor.layout<sharding_specs:b,x,y, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x8xcomplex<f64>>) -> tensor<2x4x8xcomplex<f64>>
%2 = "tf.IFFT2D"(%1) : (tensor<2x4x8xcomplex<f64>>) -> tensor<2x4x8xcomplex<f64>>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<2x4x8>, layout = #dtensor.layout<sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x8xcomplex<f64>>) -> tensor<2x4x8xcomplex<f64>>
tf_device.return %3 : tensor<2x4x8xcomplex<f64>>
}) {_mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x4x8xcomplex<f64>>
func.return
}
}
// -----
// Check the SPMD expansion for RFFT2D
// CHECK-LABEL: module @test_RFFT2D
module @test_RFFT2D {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4x12xf64> {tf._layout = "sharding_specs:b,x,y, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-NEXT: %[[ALLGATHER_OUT:.*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK: %[[CONST_OUT_2:.*]] = "tf.Const"()
// CHECK: %[[RFFT_OUT:.*]] = "tf.RFFT"(%[[ALLGATHER_OUT]], %[[CONST_OUT_2]])
// CHECK-SAME: _layout = ["sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x4x12xf64>, tensor<1xi32>) -> tensor<2x4x6xcomplex<f64>>
// CHECK-NEXT: %[[CONST_OUT_3:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_1:.*]] = "tf.Transpose"(%[[RFFT_OUT]], %[[CONST_OUT_3]])
// CHECK-SAME: (tensor<2x4x6xcomplex<f64>>, tensor<3xi64>) -> tensor<2x6x4xcomplex<f64>>
// CHECK-NEXT: %[[IDENT_OUT:.*]] = "tf.Identity"(%[[TRANS_OUT_1]])
// CHECK-NEXT: %[[FFT_OUT:.*]] = "tf.FFT"(%[[IDENT_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:b,x,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x6x4xcomplex<f64>>) -> tensor<2x6x4xcomplex<f64>>
// CHECK-NEXT: %[[CONST_OUT_4:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_2:.*]] = "tf.Transpose"(%[[FFT_OUT]], %[[CONST_OUT_4]])
// CHECK-SAME: _layout = ["sharding_specs:b,unsharded,x, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x6x4xcomplex<f64>>, tensor<3xi64>) -> tensor<2x4x6xcomplex<f64>>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<[4, 10]> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4x12>, layout = #dtensor.layout<sharding_specs:b,x,y, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x12xf64>) -> tensor<2x4x12xf64>
%3 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%4 = "tf.RFFT2D"(%2, %3) : (tensor<2x4x12xf64>, tensor<2xi32>) -> tensor<2x4x6xcomplex<f64>>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<2x4x6>, layout = #dtensor.layout<sharding_specs:b,unsharded,x, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x6xcomplex<f64>>) -> tensor<2x4x6xcomplex<f64>>
tf_device.return %5 : tensor<2x4x6xcomplex<f64>>
}) {_mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x4x6xcomplex<f64>>
func.return
}
}
// -----
// Check the SPMD expansion for IRFFT2D
// CHECK-LABEL: module @test_IRFFT2D
module @test_IRFFT2D {
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4x8xcomplex<f64>> {tf._layout = "sharding_specs:b,unsharded,y, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1", tf._mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[CONST_OUT_1:.*]] = "tf.Const"()
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-NEXT: %[[CONST_OUT_2:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_1:.*]] = "tf.Transpose"(%arg1, %[[CONST_OUT_2]])
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>, tensor<3xi64>) -> tensor<2x4x4xcomplex<f64>>
// CHECK-NEXT: %[[IFFT_OUT:.*]] = "tf.IFFT"(%[[TRANS_OUT_1]])
// CHECK-SAME: _layout = ["sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>) -> tensor<2x4x4xcomplex<f64>>
// CHECK-NEXT: %[[CONST_OUT_3:.*]] = "tf.Const"()
// CHECK-NEXT: %[[TRANS_OUT_2:.*]] = "tf.Transpose"(%[[IFFT_OUT]], %[[CONST_OUT_3]])
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>, tensor<3xi64>) -> tensor<2x4x4xcomplex<f64>>
// CHECK-NEXT: %[[ALLTOALL_OUT:.*]] = "tf.DTensorAllToAll"(%[[TRANS_OUT_2]])
// CHECK-SAME: _layout = ["sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x4x4xcomplex<f64>>) -> tensor<2x2x8xcomplex<f64>>
// CHECK-NEXT: %[[CONST_OUT_4:.*]] = "tf.Const"()
// CHECK-NEXT: %[[IRFFT_OUT:.*]] = "tf.IRFFT"(%[[ALLTOALL_OUT]], %[[CONST_OUT_4]])
// CHECK-SAME: _layout = ["sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK-SAME: (tensor<2x2x8xcomplex<f64>>, tensor<1xi32>) -> tensor<2x2x8xf64>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {value = dense<[4, 8]> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%2 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4x8>, layout = #dtensor.layout<sharding_specs:b,unsharded,y, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x8xcomplex<f64>>) -> tensor<2x4x8xcomplex<f64>>
%3 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%4 = "tf.IRFFT2D"(%2, %3) : (tensor<2x4x8xcomplex<f64>>, tensor<2xi32>) -> tensor<2x4x8xf64>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<2x4x8>, layout = #dtensor.layout<sharding_specs:b,y,unsharded, mesh:|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4x8xf64>) -> tensor<2x4x8xf64>
tf_device.return %5 : tensor<2x4x8xf64>
}) {_mesh = "|b=1,x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2x4x8xf64>
func.return
}
}
tensorflow/dtensor/mlir/tests/spmd_fill.mlir
+67
View File
@@ -0,0 +1,67 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s --dump-input=fail
// Check Fill op on sharded default input as argument.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU" }, %arg2: tensor<f32> ) -> (tensor<?x?xf32> {tf._default_layout = "sharding_specs:x,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-SAME: dense<[2, 1]>
// CHECK-NEXT: "tf.Div"
// CHECK-NEXT: "tf.Fill"
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: tensor<2xi32>, tensor<f32>) -> tensor<?x?xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.Fill"(%arg1, %arg2) : (tensor<2xi32>, tensor<f32>) -> tensor<?x?xf32>
tf_device.return %1 : tensor<?x?xf32>
}) {} : () -> tensor<?x?xf32>
func.return %0 : tensor<?x?xf32>
}
// -----
// Check Fill op on sharded default input as ConstOp.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"}, %arg2: tensor<f32> ) -> (tensor<8x1xf32>{
tf._default_layout = "sharding_specs:x,unsharded, mesh:CPU|x=2,y=2|*CPU"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.Const"
// CHECK-NEXT: dense<[2, 1]>
// CHECK-NEXT: "tf.Div"
// CHECK-NEXT: "tf.Fill"
// CHECK-SAME: {_layout = ["sharding_specs:x,unsharded, mesh:CPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]}
// CHECK-SAME: tensor<2xi32>, tensor<f32>) -> tensor<?x?xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[8, 1]> : tensor<2xi32>, _layout = ["sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"]} : () -> tensor<2xi32>
%2 = "tf.Fill"(%1, %arg2) {device = ""} : (tensor<2xi32>, tensor<f32>) -> tensor<8x1xf32>
tf_device.return %2 : tensor<8x1xf32>
}) {} : () -> tensor<8x1xf32>
func.return %0 : tensor<8x1xf32>
}
// -----
// Check tf.Fill op with incompatible layout disallowed.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:CPU|x=2,y=2|*CPU"}, %arg2: tensor<f32> ) -> (tensor<8x1xf32>{
// expected-error @+4 {{The sharding spec for axis 0 splits among 3 values}}
tf._default_layout = "sharding_specs:x,unsharded, mesh:CPU|x=3,y=2|*CPU"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[8, 1]> : tensor<2xi32>, _layout = ["sharding_specs:unsharded, mesh:CPU|x=3,y=2|*CPU"]} : () -> tensor<2xi32>
%2 = "tf.Fill"(%1, %arg2) {device = ""} : (tensor<2xi32>, tensor<f32>) -> tensor<8x1xf32>
tf_device.return %2 : tensor<8x1xf32>
}) {} : () -> tensor<8x1xf32>
func.return %0 : tensor<8x1xf32>
}
tensorflow/dtensor/mlir/tests/spmd_io_ops.mlir
+35
View File
@@ -0,0 +1,35 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion | FileCheck %s
// Check ops registered to IO Op Expander only happen on Device 0.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource>) {
"tf_device.cluster"() ({
// CHECK: tf.NotEqual
// CHECK: tf.If
// CHECK-SAME: else_branch = @tf.[[ELSE:[a-zA-Z0-9_]*]]
// CHECK-SAME: then_branch = @tf.[[THEN:[a-zA-Z0-9_]*]]
// CHECK: func private @tf.[[THEN]]
// CHECK: tf.NoOp
// CHECK: func private @tf.[[ELSE]]
// CHECK: "tf.WriteSummary"
%3 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<""> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%2 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<2> : tensor<i32>} : () -> tensor<i32>
%1 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<1> : tensor<i64>} : () -> tensor<i64>
"tf.WriteSummary"(%arg1, %1, %2, %3, %3) {_global_shape = [], device = ""} : (tensor<*x!tf_type.resource>, tensor<i64>, tensor<i32>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> () tf_device.return
}) {_mesh = "CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
tensorflow/dtensor/mlir/tests/spmd_iterator.mlir
+121
View File
@@ -0,0 +1,121 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check that IteratorGetNextOp layout is set correctly based on iterator
// resource attribute `tf._element_layouts`.
// CHECK-LABEL: func @main
func.func @main(
%arg0: tensor<1xi32>,
%arg1: tensor<*x!tf_type.resource> {
tf._element_layouts = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"],
tf._layout = "sharding_specs: mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/task:0/device:CPU:7"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[ITER_OUT:.*]] = "tf.IteratorGetNext"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<*x!tf_type.resource>) -> tensor<8x16xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<8x16xf32>
%0 = "tf_device.cluster"() ({
%elem = "tf.IteratorGetNext"(%arg1) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<*x!tf_type.resource>) -> tensor<32x16xf32>
%identity = "tf.Identity"(%elem) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<32x16xf32>) -> tensor<32x16xf32>
tf_device.return %identity : tensor<32x16xf32>
}) {_mesh="|x=4,y=2|*TPU"} : () -> (tensor<32x16xf32>)
func.return
}
// -----
// Check that IteratorGetNextOp layout is set correctly based on iterator
// resource attribute `tf._element_layouts`, where a DTensorLayout op has been
// applied to the iterator resource tensor.
// CHECK-LABEL: func @main
func.func @main(
%arg0: tensor<1xi32>,
%arg1: tensor<*x!tf_type.resource> {
tf._element_layouts = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"],
tf._layout = "sharding_specs: mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/task:0/device:CPU:7"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[ITER_OUT:.*]] = "tf.IteratorGetNext"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<*x!tf_type.resource>) -> tensor<8x16xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<8x16xf32>
%0 = "tf_device.cluster"() ({
%elem_layout = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<*>, layout = #dtensor.layout<sharding_specs: mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/task:0/device:CPU:7>} : (tensor<*x!tf_type.resource>) -> tensor<*x!tf_type.resource>
%elem = "tf.IteratorGetNext"(%elem_layout) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<*x!tf_type.resource>) -> tensor<32x16xf32>
%identity = "tf.Identity"(%elem) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<32x16xf32>) -> tensor<32x16xf32>
tf_device.return %identity : tensor<32x16xf32>
}) {_mesh="|x=4,y=2|*TPU"} : () -> (tensor<32x16xf32>)
func.return
}
// -----
// Check that element layouts from iterator with optional output is set
// correctly based on iterator resource attribute `tf._element_layouts`.
// CHECK-LABEL: func @main
func.func @main(
%arg0: tensor<1xi32>,
%arg1: tensor<*x!tf_type.resource> {
tf._element_layouts = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"],
tf._layout = "sharding_specs: mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/task:0/device:CPU:7"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[ITER_OUT:.*]] = "tf.IteratorGetNext"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<*x!tf_type.resource>) -> tensor<8x16xf32>
// CHECK: "tf.WhileRegion"
// CHECK: %[[ITER_OPTIONAL_OUT:.*]] = "tf.IteratorGetNextAsOptional"(%arg1)
// CHECK-SAME: output_shapes = [#tf_type.shape<8x16>]
// CHECK-SAME: _layout = ["sharding_specs: mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<*x!tf_type.resource>) -> tensor<!tf_type.variant>
// CHECK-NEXT: %[[HAS_VALUE:.*]] = "tf.OptionalHasValue"(%[[ITER_OPTIONAL_OUT]])
// CHECK-SAME: _layout = ["sharding_specs: mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<!tf_type.variant>) -> tensor<i1>
// CHECK: %[[GET_VALUE:.*]] = "tf.OptionalGetValue"(%[[ITER_OPTIONAL_OUT]])
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:|x=4,y=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7"]
// CHECK-SAME: (tensor<!tf_type.variant>) -> tensor<8x16xf32>
// CHECK: tf_device.return
// CHECK-SAME: tensor<8x16xf32>
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<true> : tensor<i1>} : () -> tensor<i1>
%elem = "tf.IteratorGetNext"(%arg1) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<*x!tf_type.resource>) -> tensor<32x16xf32>
%while_region:2 = "tf.WhileRegion"(%cst, %elem) ({
^bb0(%arg2: tensor<i1>, %arg3: tensor<32x16xf32>):
%identity = "tf.Identity"(%arg2) {_global_shape = [#tf_type.shape<>]} : (tensor<i1>) -> tensor<i1>
"tf.Yield"(%identity) {_global_shape = []} : (tensor<i1>) -> ()
}, {
^bb0(%arg2: tensor<i1>, %arg3: tensor<32x16xf32>):
%iter_optional_out = "tf.IteratorGetNextAsOptional"(%arg1) {_global_shape = [#tf_type.shape<>], output_shapes = [#tf_type.shape<32x16>], output_types = [f32]} : (tensor<*x!tf_type.resource>) -> tensor<!tf_type.variant>
%has_value = "tf.OptionalHasValue"(%iter_optional_out) {_global_shape = [#tf_type.shape<>]} : (tensor<!tf_type.variant>) -> tensor<i1>
%if_region:2 = "tf.IfRegion"(%has_value) ({
%has_value_identity = "tf.Identity"(%has_value) {_global_shape = [#tf_type.shape<>]} : (tensor<i1>) -> tensor<i1>
%20 = "tf.OptionalGetValue"(%iter_optional_out) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<!tf_type.variant>) -> tensor<32x16xf32>
%22 = "tf.Identity"(%20) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<32x16xf32>) -> tensor<32x16xf32>
"tf.Yield"(%has_value_identity, %22) {_global_shape = []} : (tensor<i1>, tensor<32x16xf32>) -> ()
}, {
%arg_identity = "tf.Identity"(%arg3) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<32x16xf32>) -> tensor<32x16xf32>
%has_value_identity = "tf.Identity"(%has_value) {_global_shape = [#tf_type.shape<>]} : (tensor<i1>) -> tensor<i1>
"tf.Yield"(%has_value_identity, %arg_identity) {_global_shape = []} : (tensor<i1>, tensor<32x16xf32>) -> ()
}) {_global_shape = [#tf_type.shape<>, #tf_type.shape<32x16>], _lower_using_switch_merge = true, is_stateless = true} : (tensor<i1>) -> (tensor<i1>, tensor<32x16xf32>)
%1 = "tf.Identity"(%if_region#0) {_global_shape = [#tf_type.shape<>]} : (tensor<i1>) -> tensor<i1>
%2 = "tf.Identity"(%if_region#1) {_global_shape = [#tf_type.shape<32x16>]} : (tensor<32x16xf32>) -> tensor<32x16xf32>
"tf.Yield"(%1, %2) {_global_shape = []} : (tensor<i1>, tensor<32x16xf32>) -> ()
}
) {_global_shape = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<32x16>], _lower_using_switch_merge = true, is_stateless = false, shape_invariant} : (tensor<i1>, tensor<32x16xf32>) -> (tensor<i1>, tensor<32x16xf32>)
tf_device.return %while_region#1 : tensor<32x16xf32>
}) {_mesh="|x=4,y=2|*TPU"} : () -> (tensor<32x16xf32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_matmul.mlir
+194
View File
@@ -0,0 +1,194 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s --dump-input=fail
// Batch matmul, no batch dimensions
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<2x4xi32> {tf._layout = "sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<4x2xi32> {tf._layout = "sharding_specs:y,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[MATMUL_OUT:.*]] = "tf.BatchMatMulV2"(%arg1, %arg2)
// CHECK-SAME: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// CHECK: %[[GROUP_ID:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<2x2xi32>}> : () -> tensor<2x2xi32>
// CHECK: %[[SUM_OUT:.*]] = "tf.DTensorAllReduce"(%[[MATMUL_OUT]], %[[GROUP_ID]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[SUM_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x2>, layout = #dtensor.layout<sharding_specs:y,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x2xi32>) -> tensor<4x2xi32>
%3 = "tf.BatchMatMulV2"(%1, %2) {adj_x = false, adj_y = false}: (tensor<2x4xi32>, tensor<4x2xi32>) -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %4 : tensor<2x2xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Batch matmul with batch dims
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<8x2x4xi32> {tf._layout = "sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<8x4x2xi32> {tf._layout = "sharding_specs:x,y,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[MATMUL_OUT:.*]] = "tf.BatchMatMulV2"
// CHECK-SAME: (tensor<4x2x2xi32>, tensor<4x2x2xi32>) -> tensor<4x2x2xi32>
// CHECK: %[[GROUP_ID:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<2x2xi32>}> : () -> tensor<2x2xi32>
// CHECK: %[[SUM_OUT:.*]] = "tf.DTensorAllReduce"(%[[MATMUL_OUT]], %[[GROUP_ID]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<4x2x2xi32>, tensor<2x2xi32>) -> tensor<4x2x2xi32>
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[SUM_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<8x2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<8x2x4xi32>) -> tensor<8x2x4xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<8x4x2>, layout = #dtensor.layout<sharding_specs:x,y,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<8x4x2xi32>) -> tensor<8x4x2xi32>
%3 = "tf.BatchMatMulV2"(%1, %2) {adj_x = false, adj_y = false}: (tensor<8x2x4xi32>, tensor<8x4x2xi32>) -> tensor<8x2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<8x2x2>, layout = #dtensor.layout<sharding_specs:x,unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<8x2x2xi32>) -> tensor<8x2x2xi32>
tf_device.return %4 : tensor<8x2x2xi32>
}) { _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Batch matmul, with incompatible dimensions
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<2x4xi32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2,z=2|*TPU"},
%arg2: tensor<4x2xi32> {tf._layout = "sharding_specs:z,unsharded, mesh:TPU|x=2,y=2,z=2|*TPU"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x2>, layout = #dtensor.layout<sharding_specs:z,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7>} : (tensor<4x2xi32>) -> tensor<4x2xi32>
// expected-error @+1 {{Contracting dimension for matmul has sharding dimension y for the left input and z for the right input which are not equal.}}
%3 = "tf.BatchMatMulV2"(%1, %2) {adj_x = false, adj_y = false}: (tensor<2x4xi32>, tensor<4x2xi32>) -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7>} : (tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %4 : tensor<2x2xi32>
}) {_layout = ["sharding_specs:x,y, mesh:TPU|x=2,y=2,z=2|*TPU"], _mesh = "TPU|x=2,y=2|*TPU"} : () -> (tensor<i32>)
func.return
}
// -----
// Regular matmul
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<2x4xi32> {tf._layout = "sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<4x2xi32> {tf._layout = "sharding_specs:y,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[MATMUL_OUT:.*]] = "tf.MatMul"(%arg1, %arg2)
// CHECK-SAME: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// CHECK-NEXT: %[[GROUP_ID:.*]] = "tf.Const"()
// CHECK-NEXT: %[[SUM_OUT:.*]] = "tf.DTensorAllReduce"(%[[MATMUL_OUT]], %[[GROUP_ID]])
// CHECK-SAME: reduce_op = "Add"
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK-SAME: (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[SUM_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x4xi32>) -> tensor<2x4xi32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x2>, layout = #dtensor.layout<sharding_specs:y,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x2xi32>) -> tensor<4x2xi32>
%3 = "tf.MatMul"(%1, %2) {transpose_a = false, transpose_b = false}: (tensor<2x4xi32>, tensor<4x2xi32>) -> tensor<2x2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2x2xi32>) -> tensor<2x2xi32>
tf_device.return %4 : tensor<2x2xi32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Regular MatMul with one operand sharded
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<4x3xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<4x3xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[B_SLICE:[0-9]*]] = "tf.DTensorAllScatter"(%arg2)
// CHECK: %[[MATMUL_OUT:.*]] = "tf.MatMul"(%arg1, %[[B_SLICE]])
// CHECK: %[[COLL_OUT:.*]] = "tf.DTensorAllReduce"
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[COLL_OUT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x3xf32>) -> tensor<4x3xf32>
%3 = "tf.MatMul"(%1, %2) {transpose_a = false, transpose_b = false} : (tensor<4x4xf32>, tensor<4x3xf32>) -> tensor<4x3xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x3>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x3xf32>) -> tensor<4x3xf32>
tf_device.return %4 : tensor<4x3xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<4x3xf32>
func.return %0 : tensor<4x3xf32>
}
// -----
// y,x . x,y -> *,y
// We unshard %arg1
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:y,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<4x4xf32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<4x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-NEXT: %[[MATMUL_RESULT:[0-9]*]] = "tf.MatMul"(%[[GATHERED]], %arg2)
// CHECK: %[[FINAL_REDUCE:[0-9]*]] = "tf.DTensorAllReduce"(%[[MATMUL_RESULT]], %cst)
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:y,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%3 = "tf.MatMul"(%1, %2) {transpose_a = false, transpose_b = false} : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
tf_device.return %4 : tensor<4x4xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<4x4xf32>
func.return %0 : tensor<4x4xf32>
}
// -----
// *,x . x,* -> *,y
// We should slice arg2 before matmul rather than slicing the result.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<4x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<4x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[SLICE:[0-9]*]] = "tf.DTensorAllScatter"(%arg2)
// CHECK-NEXT: %[[MATMUL_RESULT:[0-9]*]] = "tf.MatMul"(%arg1, %[[SLICE]])
// CHECK: %[[FINAL_REDUCE:[0-9]*]] = "tf.DTensorAllReduce"(%[[MATMUL_RESULT]], %cst)
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%3 = "tf.MatMul"(%1, %2) {transpose_a = false, transpose_b = false} : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
tf_device.return %4 : tensor<4x4xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<4x4xf32>
func.return %0 : tensor<4x4xf32>
}
// -----
// x,y . *,y -> x,y
// We unshard %arg1 on y.
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<4x4xf32> {tf._layout = "sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<4x4xf32> {tf._layout = "sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) -> tensor<4x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK-NEXT: %[[MATMUL_RESULT:[0-9]*]] = "tf.MatMul"(%[[GATHERED]], %arg2)
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:unsharded,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
%3 = "tf.MatMul"(%1, %2) {transpose_a = false, transpose_b = false} : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<4x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
tf_device.return %4 : tensor<4x4xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> tensor<4x4xf32>
func.return %0 : tensor<4x4xf32>
}
tensorflow/dtensor/mlir/tests/spmd_metadata.mlir
+125
View File
@@ -0,0 +1,125 @@
// 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: dtensor-opt -- %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation -dtensor-spmd-expansion | FileCheck %s
// SPMD of shape of with replicated layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Shape"
// CHECK-NOT: "tf.Const"()
// CHECK-NOT: %[[C:.*]] = "tf.Mul"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%a = "tf.Const"() {value = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]> : tensor<2x10xi32>,
_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"]}: () -> tensor<2x10xi32>
%shape_a = "tf.Shape"(%a) : (tensor<2x10xi32>) -> tensor<2xi32>
tf_device.return %shape_a : tensor<2xi32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<i32>)
func.return
}
// -----
// Check SPMD of shape op with 2D input.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Shape"
// CHECK: "tf.Const"() {value = dense<[2, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
// CHECK: %[[C:.*]] = "tf.Mul"(%[[A:.*]], %[[B:.*]]) {_layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%a = "tf.Const"() {value = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]> : tensor<2x10xi32>,
_layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"]}: () -> tensor<2x10xi32>
%shape_a = "tf.Shape"(%a) : (tensor<2x10xi32>) -> tensor<2xi32>
tf_device.return %shape_a : tensor<2xi32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<i32>)
func.return
}
// -----
// Check SPMD of shape op with 3D input.
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: tensor<i32>
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[MOD_CONST:[0-9]*]] = "tf.Const"
// CHECK-NEXT: %[[DIV_CONST:[0-9]*]] = "tf.Const"
// CHECK-NEXT: %[[PRE_MESH_COORDS:[0-9]*]] = "tf.Div"(%[[ARG0]], %[[DIV_CONST]])
// CHECK-NEXT: %[[MESH_COORDS:[0-9]*]] = "tf.FloorMod"(%[[PRE_MESH_COORDS]], %[[MOD_CONST]])
// CHECK-NEXT: %[[TENSOR:[0-9]*]] = "tf.Const"
// CHECK-NEXT: %[[SLICE_SHAPE:[0-9]*]] = "tf.Const"
// CHECK-NEXT: %[[PRE_SLICE_OFFSET:[0-9]*]] = "tf.Const"
// CHECK-NEXT: %[[SLICE_OFFSET:[0-9]*]] = "tf.MatMul"(%[[MESH_COORDS]], %[[PRE_SLICE_OFFSET]])
// CHECK-NEXT: %[[SQUEEZED_OFFSET:[0-9]*]] = "tf.Squeeze"(%[[SLICE_OFFSET]])
// CHECK-NEXT: %[[TENSOR_SLICE:[0-9]*]] = "tf.Slice"(%[[TENSOR]], %[[SQUEEZED_OFFSET]], %[[SLICE_SHAPE]])
// CHECK-NEXT: %[[TENSOR_SLICE_IDENTITY:[0-9]*]] = "tf.IdentityN"(%[[TENSOR_SLICE]])
// CHECK-NEXT: %[[TENSOR_SLICE_SHAPE:[0-9]*]] = "tf.Shape"(%[[TENSOR_SLICE_IDENTITY]])
// CHECK-NEXT: %[[TENSOR_SPLIT_SIZES:[0-9]*]] = "tf.Const"
// CHECK-NEXT: %[[TENSOR_SHAPE:[0-9]*]] = "tf.Mul"(%[[TENSOR_SLICE_SHAPE]], %[[TENSOR_SPLIT_SIZES]])
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]> : tensor<2x2x2xi32>,
_layout = ["sharding_specs:x,y,unsharded, mesh:|x=2,y=2|*CPU"]}: () -> tensor<2x2x2xi32>
%2 = "tf.Shape"(%1) : (tensor<2x2x2xi32>) -> tensor<3xi32>
tf_device.return %2 : tensor<3xi32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<3xi32>)
func.return
}
// -----
// Check SPMD of rank op with 3D input.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Rank"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%b = "tf.Const"() {value = dense<[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]> : tensor<2x2x2xi32>,
_layout = ["sharding_specs:x,y,unsharded, mesh:|x=2,y=2|*CPU"]}: () -> tensor<2x2x2xi32>
%rank_b = "tf.Rank"(%b) : (tensor<2x2x2xi32>) -> tensor<i32>
tf_device.return %rank_b : tensor<i32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<i32>)
func.return
}
// -----
// Check SPMD BroadcastGradientArgs op.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<8x3x3x3xf32>{tf._layout = "sharding_specs:x,unsharded,unsharded,unsharded, mesh:|x=2,y=2|*CPU"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[OUT1:.*]], %[[OUT2:.*]] = "tf.BroadcastGradientArgs"
// CHECK-SAME: _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3", "sharding_specs:unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<4xi32>, tensor<4xi32>) -> (tensor<?xi32>, tensor<?xi32>)
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[OUT1]]
%0 = "tf_device.cluster"() ({
%dimension = "tf.Const"() { value = dense<-1> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:|x=2,y=2|*CPU"]} : () -> tensor<i64>
%v2 = "tf.Sum"(%arg1, %dimension) {keep_dims = true}: (tensor<8x3x3x3xf32>, tensor<i64>) -> tensor<8x3x3x1xf32>
%s1 = "tf.Shape"(%arg1): (tensor<8x3x3x3xf32>) -> tensor<4xi32>
%s2 = "tf.Shape"(%v2): (tensor<8x3x3x1xf32>) -> tensor<4xi32>
%b1, %b2 = "tf.BroadcastGradientArgs"(%s1, %s2): (tensor<4xi32>, tensor<4xi32>) -> (tensor<4xi32>, tensor<4xi32>)
tf_device.return %b1 : tensor<4xi32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<i32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_random.mlir
+76
View File
@@ -0,0 +1,76 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s --dump-input=fail
// Random with no sharding
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.Const"()
// CHECK-SAME: value = dense<[32, 32, 64]>
// CHECK-NEXT: %[[RANDOM:.*]] = "tf.StatelessRandomUniform"(%[[SHAPE]], %arg1)
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[RANDOM]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<2xi32>) -> tensor<2xi32>
%2 = "tf.Const"() {value = dense<[32, 32, 64]> : tensor<3xi64>} : () -> tensor<3xi64>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<3>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<3xi64>) -> tensor<3xi64>
%4 = "tf.StatelessRandomUniform"(%3, %1) : (tensor<3xi64>, tensor<2xi32>) -> tensor<32x32x64xf32>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<32x32x64>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<32x32x64xf32>) -> tensor<32x32x64xf32>
tf_device.return %5 : tensor<32x32x64xf32>
}) {_mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return
}
// -----
// Random with x,z sharding
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:TPU|x=4,y=2,z=2|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:10,/job:localhost/task:0/device:TPU:11,/job:localhost/task:0/device:TPU:12,/job:localhost/task:0/device:TPU:13,/job:localhost/task:0/device:TPU:14,/job:localhost/task:0/device:TPU:15"}) {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: %[[MESH_SIZES:.*]] = "tf.Const"()
// CHECK-SAME: 4, 2, 2
// CHECK-NEXT: %[[MESH_SIZES_RUNNING_PRODUCT:.*]] = "tf.Const"() <{value =
// CHECK-SAME: 4, 2, 1
// CHECK-NEXT: %[[MESH_COORDS_PRE_MOD:.*]] = "tf.Div"(%arg0, %[[MESH_SIZES_RUNNING_PRODUCT]])
// CHECK-NEXT: %[[MESH_COORDS:.*]] = "tf.FloorMod"(%[[MESH_COORDS_PRE_MOD]], %[[MESH_SIZES]])
// CHECK-SAME: _mesh_coordinates =
// CHECK-NEXT: %[[MESH_MULTIPLER:.*]] = "tf.Const"()
// CHECK-SAME [65536], [0], [262144]
// CHECK-NEXT: %[[DEVICE_SEED:.*]] = "tf.MatMul"(%[[MESH_COORDS]], %[[MESH_MULTIPLER]])
// CHECK-NEXT: %[[PRIME:.*]] = "tf.Const"() <{value = dense<65521>
// CHECK-NEXT: %[[DEVICE_SEED_PRIME:.*]] = "tf.AddV2"(%[[DEVICE_SEED]], %[[PRIME]])
// CHECK-NEXT: %[[DEVICE_SEED_SQUEEZE:.*]] = "tf.Squeeze"(%[[DEVICE_SEED_PRIME]]) <{
// CHECK-NOT: dtensor.device_seed_for_mesh_dims
// CHECK-SAME: }
// CHECK-NEXT: %[[OLD_SHAPE:.*]] = "tf.Const"(
// CHECK-NEXT: %[[DEVICE_SEED_CAST:.*]] = "tf.Cast"(%[[DEVICE_SEED_SQUEEZE]])
// CHECK-NEXT: %[[NEW_SEED:.*]] = "tf.BitwiseXor"(%arg1, %[[DEVICE_SEED_CAST]])
// CHECK-NEXT: %[[NEW_SHAPE:.*]] = "tf.Const"() <{value = dense<[8, 32, 32]>
// CHECK-NEXT: %[[RANDOM:.*]] = "tf.StatelessRandomUniform"(%[[NEW_SHAPE]], %[[NEW_SEED]])
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[RANDOM]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=4,y=2,z=2|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:10,/job:localhost/task:0/device:TPU:11,/job:localhost/task:0/device:TPU:12,/job:localhost/task:0/device:TPU:13,/job:localhost/task:0/device:TPU:14,/job:localhost/task:0/device:TPU:15>} : (tensor<2xi32>) -> tensor<2xi32>
%2 = "tf.Const"() {value = dense<[32, 32, 64]> : tensor<3xi64>} : () -> tensor<3xi64>
%3 = "tf.DTensorLayout"(%2) {global_shape = #tf_type.shape<3>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:TPU|x=4,y=2,z=2|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:10,/job:localhost/task:0/device:TPU:11,/job:localhost/task:0/device:TPU:12,/job:localhost/task:0/device:TPU:13,/job:localhost/task:0/device:TPU:14,/job:localhost/task:0/device:TPU:15>} : (tensor<3xi64>) -> tensor<3xi64>
%4 = "tf.StatelessRandomUniform"(%3, %1) : (tensor<3xi64>, tensor<2xi32>) -> tensor<32x32x64xf32>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<32x32x64>, layout = #dtensor.layout<sharding_specs:x,unsharded,z, mesh:TPU|x=4,y=2,z=2|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:10,/job:localhost/task:0/device:TPU:11,/job:localhost/task:0/device:TPU:12,/job:localhost/task:0/device:TPU:13,/job:localhost/task:0/device:TPU:14,/job:localhost/task:0/device:TPU:15>} : (tensor<32x32x64xf32>) -> tensor<32x32x64xf32>
tf_device.return %5 : tensor<32x32x64xf32>
}) {_mesh = "TPU|x=4,y=2,z=2|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3,/job:localhost/task:0/device:TPU:4,/job:localhost/task:0/device:TPU:5,/job:localhost/task:0/device:TPU:6,/job:localhost/task:0/device:TPU:7,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:8,/job:localhost/task:0/device:TPU:10,/job:localhost/task:0/device:TPU:11,/job:localhost/task:0/device:TPU:12,/job:localhost/task:0/device:TPU:13,/job:localhost/task:0/device:TPU:14,/job:localhost/task:0/device:TPU:15"} : () -> (tensor<i32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_reduction.mlir
+105
View File
@@ -0,0 +1,105 @@
// 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: dtensor-opt --nouse_layout_propagation_v2 -- %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation -dtensor-spmd-expansion | FileCheck %s
// Check SPMD expansion of reduction op with replicated input.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Sum"
// CHECK-NOT: "tf.DTensorAllReduce"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%value = "tf.Const"() {value = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]> : tensor<2x10xi32>,
_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"]}: () -> tensor<2x10xi32>
%dimension = "tf.Const"() { value = dense<1> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:|x=2,y=2|*CPU"]} : () -> tensor<i64>
%sum = "tf.Sum"(%value, %dimension) {keep_dims=true}: (tensor<2x10xi32>, tensor<i64>) -> tensor<2x1xi32>
tf_device.return %sum : tensor<2x1xi32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<i32>)
func.return
}
// -----
// Check SPMD expansion of reduction op on TPU mesh with replicated input.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Sum"
// CHECK-NOT: tf.DTensorAllReduce
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%value = "tf.Const"() {value = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]> : tensor<2x10xi32>,
_layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*TPU"]}: () -> tensor<2x10xi32>
%dimension = "tf.Const"() { value = dense<1> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:|x=2,y=2|*TPU"] } : () -> tensor<i64>
%sum = "tf.Sum"(%value, %dimension) {keep_dims=true}: (tensor<2x10xi32>, tensor<i64>) -> tensor<2x1xi32>
tf_device.return %sum : tensor<2x1xi32>
}) {_mesh = ["|x=2,y=2|*TPU"]} : () -> (tensor<i32>)
func.return
}
// -----
// Check SPMD expansion of reduce op with sharded inputs on TPU mesh.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK: tf.DTensorAllReduce
%0 = "tf_device.cluster"() ({
%value = "tf.Const"() {value = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]> : tensor<2x10xi32>,
_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*TPU"]}: () -> tensor<2x10xi32>
%dimension = "tf.Const"() { value = dense<1> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:|x=2,y=2|*TPU"] } : () -> tensor<i64>
%sum = "tf.Sum"(%value, %dimension) {keep_dims=true}: (tensor<2x10xi32>, tensor<i64>) -> tensor<2x1xi32>
tf_device.return %sum : tensor<2x1xi32>
}) {_mesh = "|x=2,y=2|*TPU"} : () -> (tensor<i32>)
func.return
}
// -----
// Check SPMD reduction of reduce op with sharded inputs.
// CHECK-LABEL: func @main
// CHECK: "tf_device.cluster"
// CHECK: "tf.DTensorAllReduce"
func.func @main(%arg0: tensor<i32>) {
%0 = "tf_device.cluster"() ({
%value = "tf.Const"() {value = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]> : tensor<2x10xi32>,
_layout = ["sharding_specs:x,y, mesh:|x=2,y=2|*CPU"]}: () -> tensor<2x10xi32>
%dimension = "tf.Const"() { value = dense<1> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:|x=2,y=2|*TPU"]} : () -> tensor<i64>
%sum = "tf.Sum"(%value, %dimension) {keep_dims=true}: (tensor<2x10xi32>, tensor<i64>) -> tensor<2x1xi32>
tf_device.return %sum : tensor<2x1xi32>
}) {_mesh = "|x=2,y=2|*CPU"} : () -> (tensor<2x1xi32>)
func.return
}
// -----
// Check that reduction over an unsharded dimension, should not emit an
// all-reduce.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
// CHECK: "tf_device.cluster"
// CHECK-NOT: tf.DTensorAllReduce
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%value = "tf.Const"() {value = dense<[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]> : tensor<2x2x2xi32>,
_layout = ["sharding_specs:x,y,unsharded, mesh:|x=2,y=2|*TPU"]}: () -> tensor<2x2x2xi32>
%dimension = "tf.Const"() { value = dense<2> : tensor<i64>, _layout = ["sharding_specs:scalar, mesh:|x=2,y=2|*TPU"] } : () -> tensor<i64>
%sum = "tf.Sum"(%value, %dimension) {keep_dims=true}: (tensor<2x2x2xi32>, tensor<i64>) -> tensor<2x2x1xi32>
tf_device.return %sum : tensor<2x2x1xi32>
}) {_mesh = "|x=2,y=2|*TPU"} : () -> (tensor<i32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_save_restore.mlir
+203
View File
@@ -0,0 +1,203 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion | FileCheck %s
// Check SPMD for save op for replicated tensor only happens on device 0.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
"tf_device.cluster"() ({
// CHECK: "tf.Case"
// CHECK-SAME: branches = [@tf.[[D0:.*]], @tf.[[D1:.*]]], is_stateless = false
// CHECK: func private @tf.[[D0]]
// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<"_dev-0-of-2">
// CHECK: "tf.Add"(%arg0, %[[CST]])
// CHECK: ""
// CHECK: func private @tf.[[D1]]
// CHECK: "tf.NoOp"
%0 = "tf.Const"() {value = dense<"/dev/null"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%1 = "tf.Const"() {value = dense<"t1"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%2 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%3 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
"tf.SaveV2"(%0, %1, %2, %3) : (tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<2xi32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
// -----
// Check SPMD for save op for sharded tensor.
// The following should generate a switch case on device id, and 2 save ops in each branch.
// One for device 0, another for device 1.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
"tf_device.cluster"() ({
// CHECK: tf.Case
// CHECK-SAME: branches = [@tf.[[D0:.*]], @tf.[[D1:.*]]], is_stateless = false
// CHECK: func private @tf.[[D0]]
// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<"_dev-0-of-2">
// CHECK: "tf.Add"(%arg0, %[[CST]])
// CHECK: "2 0,1"
// CHECK: func private @tf.[[D1]]
// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<"_dev-1-of-2">
// CHECK: "tf.Add"(%arg0, %[[CST]])
// CHECK: "2 1,1"
%0 = "tf.Const"() {value = dense<"/dev/null"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%1 = "tf.Const"() {value = dense<"t1"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%2 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%3 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:x, mesh:CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
"tf.SaveV2"(%0, %1, %2, %4) : (tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<2xi32>) -> ()
tf_device.return
}) {_mesh = "CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
// -----
// Check MergeV2 only happens on Device 0 for DTensor Checkpointing V1 and
// a dtensor all-reduce is introduced to the graph.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) {
"tf_device.cluster"() ({
// CHECK: tf.DTensorAllReduce
// CHECK: tf.NotEqual
// CHECK: tf.If
// CHECK-SAME: else_branch = @tf.[[ELSE:[a-zA-Z0-9_]*]]
// CHECK-SAME: then_branch = @tf.[[THEN:[a-zA-Z0-9_]*]]
// CHECK: func private @tf.[[THEN]]
// CHECK: tf.NoOp
// CHECK: func private @tf.[[ELSE]]
// CHECK: "tf.Const"() <{value = dense<"_dev-0-of-2">
// CHECK: "tf.Add"
// CHECK: "tf.Const"() <{value = dense<"_dev-1-of-2">
// CHECK: "tf.Add"
// CHECK: "tf.Concat"
// CHECK: "tf.MergeV2Checkpoints"
%0 = "tf.Const"() {value = dense<"/dev/null/device"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
%1 = "tf.Const"() {value = dense<"/dev/null/destination"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
"tf.MergeV2Checkpoints"(%0, %1) {allow_missing_files = true, delete_old_dirs = false} : (tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> ()
tf_device.return
}) {_mesh = "CPU|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1"} : () -> ()
func.return
}
// -----
// Check DTensorRestoreV2 does local restore with slice_spec.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>},
%arg1: tensor<!tf_type.string> {
tf._layout = "sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<2x!tf_type.string> {
tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg3: tensor<2x!tf_type.string> {
tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg4: tensor<8x2xf32> {
tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg5: tensor<2x4xf32> {
tf._layout = "sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> (tensor<*xf32> , tensor<*xf32> ) {
%0:2 = "tf_device.cluster"() ({
// CHECK: %[[CONDITION:.*]] = "tf.Equal"(%[[LOCAL_DEVICE_IDS:.*]], %arg0)
// CHECK: %[[IDX_TENSOR:.*]] = "tf.Where"(%[[CONDITION]])
// CHECK-SAME: tensor<1x1xi64>
// CHECK: %[[BRANCH_IDX:.*]] = "tf.Cast"(%[[IDX_TENSOR]]
// CHECK: "tf.Reshape"(%[[BRANCH_IDX]]
// CHECK-SAME: (tensor<1x1xi32>, tensor<0xi32>) -> tensor<i32>
// CHECK: tf.Case
// CHECK-SAME: branches = [@tf.[[D0:.*]], @tf.[[D1:.*]]], is_stateless = false
// CHECK: func private @tf.[[D0]]
// CHECK: "tf.Const"() <{value = dense<["", "2 4 0,1:-"]>
// CHECK: func private @tf.[[D1]]
// CHECK: "tf.Const"() <{value = dense<["", "2 4 1,1:-"]>
%1 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%2 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<["model/r/.ATTRIBUTES/VARIABLE_VALUE", "model/s/.ATTRIBUTES/VARIABLE_VALUE"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%3 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<"/dev/null/ckpt-0"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%4 = "tf.DTensorLayout"(%3) {_global_shape = [#tf_type.shape<>], global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%5 = "tf.DTensorLayout"(%2) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x!tf_type.string>) -> tensor<2x!tf_type.string>
%6 = "tf.DTensorLayout"(%1) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x!tf_type.string>) -> tensor<2x!tf_type.string>
%7:2 = "tf.DTensorRestoreV2"(%4, %5, %6) {_global_shape = [#tf_type.shape<*>, #tf_type.shape<*>], device = "", input_dtypes = [f32, f32],
input_shapes=[#tf_type.shape<8x2>, #tf_type.shape<2x4>],
input_layouts=["sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
"sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]} : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xf32>)
%8 = "tf.DTensorLayout"(%7#0) {global_shape = #tf_type.shape<*>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<*xf32>) -> tensor<*xf32>
%9 = "tf.DTensorLayout"(%7#1) {global_shape = #tf_type.shape<*>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<*xf32>) -> tensor<*xf32>
tf_device.return %8, %9 : tensor<*xf32>, tensor<*xf32>
}) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> (tensor<*xf32>, tensor<*xf32>)
func.return %0#0, %0#1 : tensor<*xf32>, tensor<*xf32>
}
// -----
// Check RestoreV2 does local restore with correct shape_and_slice spec.
// Restores a replicated 8x2 tensor and a x,unsharded 2x4 tensor.
// The expansion of a RestoreV2 should be the same expansion as a
// DTensorRestoreV2.
//
// To check correctness of the expansion, we just need to check that the
// correct `shape_and_slices` constant string is produced for each
// device_id function.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> {tf._global_shape = #tf_type.shape<>},
%arg1: tensor<!tf_type.string> {
tf._layout = "sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg2: tensor<2x!tf_type.string> {
tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg3: tensor<2x!tf_type.string> {
tf._layout = "sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg4: tensor<*x!tf_type.resource<tensor<8x2xf32>>> {
tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"},
%arg5: tensor<*x!tf_type.resource<tensor<2x4xf32>>> {
tf._layout = "sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1",
tf._mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> (tensor<8x2xf32> , tensor<2x4xf32> ) {
%0:2 = "tf_device.cluster"() ({
// CHECK: %[[CONDITION:.*]] = "tf.Equal"(%[[LOCAL_DEVICE_IDS:.*]], %arg0)
// CHECK: %[[IDX_TENSOR:.*]] = "tf.Where"(%[[CONDITION]])
// CHECK-SAME: tensor<1x1xi64>
// CHECK: %[[BRANCH_IDX:.*]] = "tf.Cast"(%[[IDX_TENSOR]]
// CHECK: "tf.Reshape"(%[[BRANCH_IDX]]
// CHECK-SAME: (tensor<1x1xi32>, tensor<0xi32>) -> tensor<i32>
// CHECK: tf.Case
// CHECK-SAME: branches = [@tf.[[D0:.*]], @tf.[[D1:.*]]], is_stateless = false
// CHECK: func private @tf.[[D0]]
// CHECK: "tf.Const"() <{value = dense<["", "2 4 0,1:-"]>
// CHECK: func private @tf.[[D1]]
// CHECK: "tf.Const"() <{value = dense<["", "2 4 1,1:-"]>
%1 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%2 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<["model/r/.ATTRIBUTES/VARIABLE_VALUE", "model/s/.ATTRIBUTES/VARIABLE_VALUE"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
%3 = "tf.Const"() {_global_shape = [#tf_type.shape<>], value = dense<"/dev/null/ckpt-0"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%4 = "tf.DTensorLayout"(%3) {_global_shape = [#tf_type.shape<>], global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
%5 = "tf.DTensorLayout"(%2) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x!tf_type.string>) -> tensor<2x!tf_type.string>
%6 = "tf.DTensorLayout"(%1) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x!tf_type.string>) -> tensor<2x!tf_type.string>
%7:2 = "tf.RestoreV2"(%4, %5, %6) {input_dtypes = [f32, f32]} : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<8x2xf32>, tensor<2x4xf32>)
"tf.AssignVariableOp"(%arg4, %7#0) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<8x2xf32>>>, tensor<8x2xf32>) -> ()
"tf.AssignVariableOp"(%arg5, %7#1) {validate_shape = true} : (tensor<*x!tf_type.resource<tensor<2x4xf32>>>, tensor<2x4xf32>) -> ()
%8 = "tf.DTensorLayout"(%7#0) {global_shape = #tf_type.shape<*>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<8x2xf32>) -> tensor<8x2xf32>
%9 = "tf.DTensorLayout"(%7#1) {global_shape = #tf_type.shape<*>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
tf_device.return %8, %9 : tensor<8x2xf32>, tensor<2x4xf32>
}) {_mesh = "|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> (tensor<8x2xf32>, tensor<2x4xf32>)
func.return %0#0, %0#1 : tensor<8x2xf32>, tensor<2x4xf32>
}
tensorflow/dtensor/mlir/tests/spmd_segment_sum.mlir
+43
View File
@@ -0,0 +1,43 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// UnsortedSegmentSum data and segment sum same layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32>,
%arg1: tensor<16x2xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"},
%arg2: tensor<16xi32> {tf._layout = "sharding_specs:x, mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
// CHECK: "tf_device.cluster"
// CHECK: %[[NUM_SEGMENTS:.*]] = "tf.Const"
// CHECK-SAME: () -> tensor<i32>
// CHECK: %[[LOCAL_RESULT:.*]] = "tf.UnsortedSegmentSum"(%arg1, %arg2, %[[NUM_SEGMENTS]])
// CHECK-SAME: (tensor<4x2xf32>, tensor<4xi32>, tensor<i32>) -> tensor<8x2xf32>
// CHECK: %[[RESULT:.*]] = "tf.DTensorAllReduce"(%[[LOCAL_RESULT]]
// CHECK-SAME: reduce_op = "Add"
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]
// CHECK: %[[FINAL_RESULT:.*]] = "tf.DTensorAllScatter"(%[[RESULT]]
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[FINAL_RESULT]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<16x2>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<16x2xf32>) -> tensor<16x2xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<16>, layout = #dtensor.layout<sharding_specs:x, mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<16xi32>) -> tensor<16xi32>
%3 = "tf.Const"() {value = dense<8> : tensor<i32>} : () -> tensor<i32>
%4 = "tf.DTensorLayout"(%3) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs: mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<i32>) -> tensor<i32>
%5 = "tf.UnsortedSegmentSum"(%1, %2, %4) : (tensor<16x2xf32>, tensor<16xi32>, tensor<i32>) -> tensor<8x2xf32>
%6 = "tf.DTensorLayout"(%5) {global_shape = #tf_type.shape<8x2>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3>} : (tensor<8x2xf32>) -> tensor<8x2xf32>
tf_device.return %6 : tensor<8x2xf32>
}) {_mesh = "TPU|x=4|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> (tensor<i32>)
func.return
}
tensorflow/dtensor/mlir/tests/spmd_slice.mlir
+169
View File
@@ -0,0 +1,169 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s --dump-input=always
// Check SPMD of splice op with replicated input.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<2x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Slice"(%arg1, %cst, %cst_1)
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xf32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%2 = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%4 = "tf.Slice"(%1, %2, %3) : (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xf32>
tf_device.return %4 : tensor<2x2xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<2x2xf32>)
func.return %0 : tensor<2x2xf32>
}
// -----
// Check that the slice on sharded x dimension is from 1, which requires a relayout to a fully replicated layout.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<1x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[GATHERED:[0-9]*]] = "tf.DTensorAllGather"(%arg1)
// CHECK: "tf.Slice"(%[[GATHERED]], %cst, %cst_1)
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xf32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%2 = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%4 = "tf.Slice"(%1, %2, %3) : (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xf32>
tf_device.return %4 : tensor<1x2xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<1x2xf32>)
func.return %0 : tensor<1x2xf32>
}
// -----
// Check that the slice on sharded x dimension is from 0, which can operate on and produce sharded tensors.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<2x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Slice"(%arg1, %cst, %cst_1)
// CHECK-SAME: _layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<1x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xf32>
// CHECK: tf_device.return
// CHECK-SAME: _layout = []
// CHECK-SAME: tensor<1x2xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%2 = "tf.Const"() {value = dense<[0, 0]> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%4 = "tf.Slice"(%1, %2, %3) : (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xf32>
tf_device.return %4 : tensor<2x2xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<2x2xf32>)
func.return %0 : tensor<2x2xf32>
}
// -----
// Slice a sharded input but produce a replicated output. This would crash without the fix for b/181933405.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<2x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: "tf.Slice"(%arg1, %cst, %cst_1)
// CHECK-SAME: (tensor<1x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xf32>
// CHECK: tf_device.return
// CHECK-NEXT: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"], _mesh = "|x=2,y=2|*CPU"
// CHECK-SAME: () -> tensor<2x2xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%2 = "tf.Const"() {value = dense<[0, 0]> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%4 = "tf.Slice"(%1, %2, %3) : (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xf32>
%5 = "tf.DTensorLayout"(%4) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2x2xf32>) -> tensor<2x2xf32>
tf_device.return %5 : tensor<2x2xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<2x2xf32>)
func.return %0 : tensor<2x2xf32>
}
// -----
// Check SPMD expansion slice with a dynamic begins and sharded input on non
// full slice dimensions.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x4xf32> {tf._layout = "sharding_specs:unsharded,x, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<2xi64> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<1x4xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi64>}> : () -> tensor<2xi64>
// CHECK-NEXT: %[[SLICE:.*]] = "tf.Slice"(%arg1, %arg2, %[[SLICE_SIZE]])
// CHECK-SAME: _layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<2x2xf32>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x2xf32>
// CHECK-NEXT: tf_device.return
// CHECK-SAME: %[[SLICE]]
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,x, mesh:|x=2,y=2|*CPU>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2xi64>) -> tensor<2xi64>
%3 = "tf.Const"() {value = dense<[1, 4]> : tensor<2xi64>} : () -> tensor<2xi64>
%4 = "tf.Slice"(%1, %2, %3) : (tensor<2x4xf32>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x4xf32>
tf_device.return %4 : tensor<1x4xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded,x, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<1x4xf32>)
func.return %0 : tensor<1x4xf32>
}
// -----
// Check SPMD expansion of strided slice op with replicated input.
func.func @main(%arg0: tensor<i32>, %arg1: tensor<2x4xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<2x2xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %cst_2 = "tf.Const"() <{value = dense<2> : tensor<2xi32>}>
// CHECK: "tf.StridedSlice"(%arg1, %cst, %cst_2, %cst_1)
// CHECK-SAME: _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3"]
// CHECK-SAME: (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xf32>
// CHECK: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<2x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<2x4xf32>) -> tensor<2x4xf32>
%2 = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tf.Const"() {value = dense<[2, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%4 = "tf.Const"() {value = dense<[1, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
%5 = "tf.StridedSlice"(%1, %2, %3, %4) {begin_mask = 0 : i64, device = "", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<2x4xf32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xf32>
tf_device.return %5 : tensor<2x2xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<2x2xf32>)
func.return %0 : tensor<2x2xf32>
}
// -----
// Check layout propagation and spmd expansion of strided slice grad op.
func.func @main(%arg0: tensor<15x12xf32>) -> tensor<15x197x12xf32> {
// CHECK: "tf_device.cluster"
%0 = "tf_device.cluster"() ({
%cst = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_0 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<[0, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_1 = "tf.Const"() {_global_shape = [#tf_type.shape<2>], value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_2 = "tf.Const"() {_global_shape = [#tf_type.shape<3>], value = dense<[15, 197, 12]> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tf.DTensorLayout"(%arg0) {_global_shape = [#tf_type.shape<15x12>], global_shape = #tf_type.shape<15x12>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<15x12xf32>) -> tensor<15x12xf32>
%2 = "tf.DTensorLayout"(%cst_2) {_global_shape = [#tf_type.shape<3>], global_shape = #tf_type.shape<3>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<3xi32>) -> tensor<3xi32>
%3 = "tf.DTensorLayout"(%cst_1) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%4 = "tf.DTensorLayout"(%cst_0) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%5 = "tf.DTensorLayout"(%cst) {_global_shape = [#tf_type.shape<2>], global_shape = #tf_type.shape<2>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2xi32>) -> tensor<2xi32>
%6 = "tf.StridedSliceGrad"(%2, %3, %4, %5, %1) {_global_shape = [#tf_type.shape<15x197x12>], begin_mask = 1 : i64, device = "", ellipsis_mask = 0 : i64, end_mask = 1 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64} : (tensor<3xi32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>, tensor<15x12xf32>) -> tensor<15x197x12xf32>
%7 = "tf.DTensorLayout"(%6) {_global_shape = [#tf_type.shape<15x197x12>], global_shape = #tf_type.shape<15x197x12>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded,unsharded, mesh:|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<15x197x12xf32>) -> tensor<15x197x12xf32>
tf_device.return %7 : tensor<15x197x12xf32>
}) {_mesh = "|x=1,y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<15x197x12xf32>
func.return %0 : tensor<15x197x12xf32>
}
tensorflow/dtensor/mlir/tests/spmd_softmax_loss.mlir
+148
View File
@@ -0,0 +1,148 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check SPMD of Softmax with no sharding.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<6x4xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<6x4xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<6xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.SoftmaxCrossEntropyWithLogits"
// CHECK-NEXT: "tf.IdentityN"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%loss, %backprop = "tf.SoftmaxCrossEntropyWithLogits"(%1, %2) : (tensor<6x4xf32>, tensor<6x4xf32>) -> (tensor<6xf32>, tensor<6x4xf32>)
%3 = "tf.DTensorLayout"(%loss) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6xf32>) -> tensor<6xf32>
%4 = "tf.DTensorLayout"(%backprop) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
tf_device.return %3 : tensor<6xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<6xf32>)
func.return %0 : tensor<6xf32>
}
// -----
// Check SPMD of Softmax with batch sharding but no class sharding.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<6x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<6x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<6xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.SoftmaxCrossEntropyWithLogits"
// CHECK-NEXT: "tf.IdentityN"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%loss, %backprop = "tf.SoftmaxCrossEntropyWithLogits"(%1, %2) : (tensor<6x4xf32>, tensor<6x4xf32>) -> (tensor<6xf32>, tensor<6x4xf32>)
%3 = "tf.DTensorLayout"(%loss) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2,y=2|*CPU>} : (tensor<6xf32>) -> tensor<6xf32>
%4 = "tf.DTensorLayout"(%backprop) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
tf_device.return %3 : tensor<6xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:x, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<6xf32>)
func.return %0 : tensor<6xf32>
}
// -----
// Check SPMD of Softmax with batch sharding and class sharding,
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<6x4xf32> {tf._layout = "sharding_specs:x,y, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<6x4xf32> {tf._layout = "sharding_specs:x,y, mesh:|x=2,y=2|*CPU"}) -> tensor<6xf32> {
// CHECK: "tf_device.cluster"
// CHECK: %[[LOCAL_MAX:.*]] = "tf.Max"(%arg1,
// CHECK: %[[MAX:.*]] = "tf.DTensorAllReduce"(%[[LOCAL_MAX]],
// CHECK-SAME: "Max"
// CHECK: %[[SHIFTED_LOGITS:.*]] = "tf.Sub"(%arg1, %[[MAX]])
// CHECK: %[[EXP_LOGITS:.*]] = "tf.Exp"(%[[SHIFTED_LOGITS]])
// CHECK: %[[LOCAL_SUM:.*]] = "tf.Sum"(%[[EXP_LOGITS]],
// CHECK: %[[SUM:.*]] = "tf.DTensorAllReduce"(%[[LOCAL_SUM]],
// CHECK-SAME: "Add"
// CHECK: %[[LOG_SUM:.*]] = "tf.Log"(%[[SUM]])
// CHECK: %[[LOG_SOFTMAX:.*]] = "tf.Sub"(%[[SHIFTED_LOGITS]], %[[LOG_SUM]])
// CHECK: %[[SOFTMAX:.*]] = "tf.Div"(%[[EXP_LOGITS]], %[[SUM]])
// CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%arg2,
// CHECK: %[[SAFE_LOG_SOFTMAX:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[ZERO:.*]], %[[LOG_SOFTMAX]])
// CHECK: %[[PROD:.*]] = "tf.Mul"(%arg2, %[[SAFE_LOG_SOFTMAX]])
// CHECK: %[[LOCAL_NEG_LOSS:.*]] = "tf.Sum"(%[[PROD]],
// CHECK: %[[NEG_LOSS:.*]] = "tf.DTensorAllReduce"(%[[LOCAL_NEG_LOSS]],
// CHECK-SAME: "Add"
// CHECK: %[[SQUEEZED_NEG_LOSS:.*]] = "tf.Squeeze"(%[[NEG_LOSS]])
// CHECK: %[[LOSS:.*]] = "tf.Neg"(%[[SQUEEZED_NEG_LOSS]])
// CHECK: %[[BACKPROP:.*]] = "tf.Sub"(%[[SOFTMAX]], %arg2)
// CHECK-NEXT: "tf.IdentityN"(%[[LOSS]], %[[BACKPROP]])
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%loss, %backprop = "tf.SoftmaxCrossEntropyWithLogits"(%1, %2) : (tensor<6x4xf32>, tensor<6x4xf32>) -> (tensor<6xf32>, tensor<6x4xf32>)
%3 = "tf.DTensorLayout"(%loss) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2,y=2|*CPU>} : (tensor<6xf32>) -> tensor<6xf32>
%4 = "tf.DTensorLayout"(%backprop) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,y, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
tf_device.return %3 : tensor<6xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<6xf32>)
func.return %0 : tensor<6xf32>
}
// -----
// Check SPMD of SparseSoftmax with no sharding.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<6x4xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<6xi32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"}) -> tensor<6xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.SparseSoftmaxCrossEntropyWithLogits"
// CHECK-NEXT: "tf.IdentityN"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6xi32>) -> tensor<6xi32>
%loss, %backprop = "tf.SparseSoftmaxCrossEntropyWithLogits"(%1, %2) : (tensor<6x4xf32>, tensor<6xi32>) -> (tensor<6xf32>, tensor<6x4xf32>)
%3 = "tf.DTensorLayout"(%loss) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6xf32>) -> tensor<6xf32>
%4 = "tf.DTensorLayout"(%backprop) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
tf_device.return %3 : tensor<6xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<6xf32>)
func.return %0 : tensor<6xf32>
}
// -----
// Check SPMD of Softmax with batch sharding but no class sharding.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<6x4xf32> {tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<6xi32> {tf._layout = "sharding_specs:x, mesh:|x=2,y=2|*CPU"}) -> tensor<6xf32> {
// CHECK: "tf_device.cluster"
// CHECK-NEXT: "tf.SparseSoftmaxCrossEntropyWithLogits"
// CHECK-NEXT: "tf.IdentityN"
// CHECK-NEXT: tf_device.return
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
%2 = "tf.DTensorLayout"(%arg2) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2,y=2|*CPU>} : (tensor<6xi32>) -> tensor<6xi32>
%loss, %backprop = "tf.SparseSoftmaxCrossEntropyWithLogits"(%1, %2) : (tensor<6x4xf32>, tensor<6xi32>) -> (tensor<6xf32>, tensor<6x4xf32>)
%3 = "tf.DTensorLayout"(%loss) {global_shape = #tf_type.shape<6>, layout = #dtensor.layout<sharding_specs:x, mesh:|x=2,y=2|*CPU>} : (tensor<6xf32>) -> tensor<6xf32>
%4 = "tf.DTensorLayout"(%backprop) {global_shape = #tf_type.shape<6x4>, layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU>} : (tensor<6x4xf32>) -> tensor<6x4xf32>
tf_device.return %3 : tensor<6xf32>
}) {_mesh = "|x=2,y=2|*CPU", _layout = ["sharding_specs:x, mesh:|x=2,y=2|*CPU"]} : () -> (tensor<6xf32>)
func.return %0 : tensor<6xf32>
}
tensorflow/dtensor/mlir/tests/spmd_squeeze.mlir
+69
View File
@@ -0,0 +1,69 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check Squeeze with positive index.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> , %arg1: tensor<2x1xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> tensor<2xf32> {
// CHECK: "tf.Squeeze"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=1|0,1|0,1|
// CHECK-SAME: (tensor<1x1xf32>) -> tensor<1xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<2x1>], global_shape = #tf_type.shape<2x1>,
layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x1xf32>) -> tensor<2x1xf32>
%2 = "tf.Squeeze"(%1) {_global_shape = [#tf_type.shape<2>], device = "", squeeze_dims = [1]} : (tensor<2x1xf32>) -> tensor<2xf32>
tf_device.return %2 : tensor<2xf32>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// -----
// Check Squeeze with negative index.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> , %arg1: tensor<2x1xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> tensor<2xf32> {
// CHECK: "tf.Squeeze"(%arg1)
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=1|0,1|0,1|
// CHECK-SAME: (tensor<1x1xf32>) -> tensor<1xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<2x1>], global_shape = #tf_type.shape<2x1>,
layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x1xf32>) -> tensor<2x1xf32>
%2 = "tf.Squeeze"(%1) {_global_shape = [#tf_type.shape<2>], device = "", squeeze_dims = [-1]} : (tensor<2x1xf32>) -> tensor<2xf32>
tf_device.return %2 : tensor<2xf32>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// -----
// Check Squeeze that does not locally squeeze the dim with local shape 1.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32> , %arg1: tensor<2x1xf32> { tf._layout = "sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}) -> tensor<2xf32> {
// CHECK: "tf.Squeeze"(%arg1)
// CHECK-SAME: squeeze_dims = [1]
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2,y=1|0,1|0,1|
// CHECK-SAME: (tensor<1x1xf32>) -> tensor<1xf32>
%0 = "tf_device.cluster"() ({
%1 = "tf.DTensorLayout"(%arg1) {_global_shape = [#tf_type.shape<2x1>], global_shape = #tf_type.shape<2x1>,
layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<2x1xf32>) -> tensor<2x1xf32>
%2 = "tf.Squeeze"(%1) {_global_shape = [#tf_type.shape<2>], device = "", squeeze_dims = []} : (tensor<2x1xf32>) -> tensor<2xf32>
tf_device.return %2 : tensor<2xf32>
}) {_mesh = "|x=2,y=1|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"} : () -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
tensorflow/dtensor/mlir/tests/spmd_tile.mlir
+76
View File
@@ -0,0 +1,76 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check TileOp on sharded const input.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<2x1xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"}) -> (tensor<4x3xf32>{
// CHECK: "tf_device.cluster"
// CHECK: %[[MULTIPLES:.*]] = "tf.Const"() {value = dense<[1, 3]> : tensor<2xi32>} : () -> tensor<2xi32>
// CHECK: "tf.Tile"(%arg1, %[[MULTIPLES]])
// CHECK-SAME: (tensor<2x1xf32>, tensor<2xi32>) -> tensor<2x3xf32>
tf._default_layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[2, 3]> : tensor<2xi32>, _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> tensor<2xi32>
%2 = "tf.Tile"(%arg1, %1) {device = ""} : (tensor<2x1xf32>, tensor<2xi32>) -> tensor<4x3xf32>
tf_device.return %2 : tensor<4x3xf32>
}) {} : () -> tensor<4x3xf32>
func.return %0 : tensor<4x3xf32>
}
// -----
// Check TileOp on sharded const input with partial shape.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<?x1xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"}) -> (tensor<?x3xf32>{
// CHECK: "tf_device.cluster"
// CHECK: %[[MULTIPLES:.*]] = "tf.Const"()
// CHECK-NEXT: dense<[1, 3]>
// CHECK: "tf.Tile"
// CHECK-SAME: (tensor<?x1xf32>, tensor<2xi32>) -> tensor<?x3xf32>
tf._default_layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[2, 3]> : tensor<2xi32>, _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> tensor<2xi32>
%2 = "tf.Tile"(%arg1, %1) {device = ""} : (tensor<?x1xf32>, tensor<2xi32>) -> tensor<?x3xf32>
tf_device.return %2 : tensor<?x3xf32>
}) {} : () -> tensor<?x3xf32>
func.return %0 : tensor<?x3xf32>
}
// -----
// Check TileOp on sharded const input with partial shape.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>,
%arg1: tensor<?x1xf32> {tf._layout = "sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"},
%arg2: tensor<2xi32> {tf._layout = "sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"}) -> (tensor<?x3xf32>{
// CHECK: "tf_device.cluster"
// CHECK: %[[MULTIPLES:.*]] = "tf.Const"()
// CHECK-NEXT: dense<[1, 3]>
// CHECK: "tf.Tile"
// CHECK-SAME: (tensor<?x1xf32>, tensor<2xi32>) -> tensor<?x3xf32>
tf._default_layout = "sharding_specs:x,unsharded, mesh:|x=2,y=2|*CPU"}) {
%0 = "tf_device.cluster"() ({
%1 = "tf.Const"() {value = dense<[2, 3]> : tensor<2xi32>, _layout = ["sharding_specs:unsharded, mesh:|x=2,y=2|*CPU"]} : () -> tensor<2xi32>
%2 = "tf.Tile"(%arg1, %1) {device = ""} : (tensor<?x1xf32>, tensor<2xi32>) -> tensor<?x3xf32>
tf_device.return %2 : tensor<?x3xf32>
}) {} : () -> tensor<?x3xf32>
func.return %0 : tensor<?x3xf32>
}
tensorflow/dtensor/mlir/tests/spmd_var_handle.mlir
+47
View File
@@ -0,0 +1,47 @@
// 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: dtensor-opt %s -split-input-file -dtensor-annotate-global-shape -dtensor-layout-propagation-v2 -dtensor-spmd-expansion -verify-diagnostics | FileCheck %s
// Check layout propagation for tf.VarHandleOp followed by Relayout.
func.func @main() -> (tensor<2xi32>) {
// CHECK: tf_device.cluster
// CHECK-NEXT: "tf.VarHandleOp"()
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
// CHECK: "tf.ReadVariableOp"
// CHECK-SAME: _layout = ["sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"]
%0 = "tf_device.cluster"() ({
%0 = "tf.VarHandleOp"() {_global_shape = [#tf_type.shape<>], allowed_devices = [], container = "", device = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<2xi32>>>
%1 = "tf.Relayout"(%0) {_global_shape = [#tf_type.shape<>], device = "", layout = "sharding_specs:x, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1"}
: (tensor<!tf_type.resource<tensor<2xi32>>>) -> tensor<!tf_type.resource<tensor<2xi32>>>
%2 = "tf.ReadVariableOp"(%0) {_global_shape = [#tf_type.shape<2>], device = ""} : (tensor<!tf_type.resource<tensor<2xi32>>>) -> tensor<2xi32>
tf_device.return %2 : tensor<2xi32>
}) {_mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> (tensor<2xi32>)
func.return %0 : tensor<2xi32>
}
// -----
// Check layout propagation for tf.VarHandleOp without a Relayout.
func.func @main() -> (tensor<!tf_type.resource<tensor<2xi32>>>) {
// CHECK: tf_device.cluster
// CHECK-NEXT: "tf.VarHandleOp"()
// CHECK-SAME: _layout = ["empty_layout"]
%0 = "tf_device.cluster"() ({
%0 = "tf.VarHandleOp"() {_global_shape = [#tf_type.shape<>], allowed_devices = [], container = "", device = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<2xi32>>>
tf_device.return %0 : tensor<!tf_type.resource<tensor<2xi32>>>
}) {_mesh = "TPU|x=2|0,1|0,1|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1"} : () -> (tensor<!tf_type.resource<tensor<2xi32>>>)
func.return %0 : tensor<!tf_type.resource<tensor<2xi32>>>
}
tensorflow/dtensor/mlir/tests/tf_dtensor_ops.mlir
+123
View File
@@ -0,0 +1,123 @@
// 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: dtensor-opt %s -split-input-file -verify-diagnostics
// Check that a DTensorLayout op with mismatched rank between layout and input
// value is disallowed.
func.func @invalid_rank_disallowed(%arg0: tensor<i32>) {
// expected-error@+1 {{requires matching rank for layout and input, but got 2}}
%0 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<>, layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU>} : (tensor<i32>) -> tensor<i32>
func.return
}
// -----
// Check that a DTensorLayout op with sharding configuration that cannot evenly
// divide the dimension of the input value is disallowed.
func.func @invalid_sharding_dim_disallowed(%arg0: tensor<2x2xi32>) {
// expected-error@+1 {{requires dimension 0 to be divisible by sharding specified in DTensorLayout, but got dimension size=2 is not divisible by number of shards in layout for this dimension=8.}}
%0 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<2x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>} : (tensor<2x2xi32>) -> (tensor<2x2xi32>)
func.return
}
// -----
func.func @dtensor_layout_with_sharding(%arg0: tensor<16x2xi32>) {
// CHECK: "tf.DTensorLayout"
// CHECK-SAME: layout = #dtensor.layout<sharding_specs:x,y mesh:CPU|x=8,y=2|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3,CPU:4,CPU:5,CPU:6,CPU:7,CPU:8,CPU:9,CPU:10,CPU:11,CPU:12,CPU:13,CPU:14,CPU:15>
// CHECK-SAME: (tensor<16x2xi32>) -> (tensor<16x2xi32>)
%0 = "tf.DTensorLayout"(%arg0) {global_shape = #tf_type.shape<16x2>, layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16x2xi32>)
func.return
}
// -----
func.func @dtensor_all_gather_unequal_tensor_ranks(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{received input and output layouts of unequal ranks 2 and 1}}
%0 = "tf.DTensorAllGather"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16xi32>)
func.return
}
// -----
func.func @dtensor_all_gather_bad_layouts(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{dimension 1 of output layout has sharding spec y which is more sharded then the input layout spec unsharded}}
%0 = "tf.DTensorAllGather"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16x2xi32>)
func.return
}
// -----
func.func @dtensor_all_gather_bad_input_rank(%arg0: tensor<16xi32>) {
// expected-error@+1 {{input layout rank 2 is not equal to input rank 1}}
%0 = "tf.DTensorAllGather"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16xi32>) -> (tensor<16x2xi32>)
func.return
}
// -----
func.func @dtensor_all_gather_bad_output_rank(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{output layout rank 2 is not equal to output rank 1}}
%0 = "tf.DTensorAllGather"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16xi32>)
func.return
}
// -----
func.func @dtensor_all_gather_bad_output_shape(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{computed output shape 4 at dimension 1 is not equal to actual output shape 2}}
%0 = "tf.DTensorAllGather"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16x2xi32>)
func.return
}
// -----
func.func @dtensor_all_scatter_unequal_tensor_ranks(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{received input and output layouts of unequal ranks 2 and 1}}
%0 = "tf.DTensorAllScatter"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16xi32>)
func.return
}
// -----
func.func @dtensor_all_scatter_bad_layouts(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{dimension 1 of input layout has sharding spec y which is more sharded then the output layout spec unsharded}}
%0 = "tf.DTensorAllScatter"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16x2xi32>)
func.return
}
// -----
func.func @dtensor_all_scatter_bad_input_rank(%arg0: tensor<16xi32>) {
// expected-error@+1 {{input layout rank 2 is not equal to input rank 1}}
%0 = "tf.DTensorAllScatter"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16xi32>) -> (tensor<16x2xi32>)
func.return
}
// -----
func.func @dtensor_all_scatter_bad_output_rank(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{output layout rank 2 is not equal to output rank 1}}
%0 = "tf.DTensorAllScatter"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16xi32>)
func.return
}
// -----
func.func @dtensor_all_scatter_bad_output_shape(%arg0: tensor<16x2xi32>) {
// expected-error@+1 {{computed output shape 1 at dimension 1 is not equal to actual output shape 2}}
%0 = "tf.DTensorAllScatter"(%arg0) {input_layout = #dtensor.layout<sharding_specs:x,unsharded, mesh:CPU|x=8,y=2|*CPU>, output_layout = #dtensor.layout<sharding_specs:x,y, mesh:CPU|x=8,y=2|*CPU>} : (tensor<16x2xi32>) -> (tensor<16x2xi32>)
func.return
}
tensorflow/dtensor/mlir/tests/tpu_add_resource_device_attribute.mlir
+88
View File
@@ -0,0 +1,88 @@
// 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: dtensor-opt %s -split-input-file -dtensor-tpu-add-resource-device-attribute | FileCheck %s
// Test that tf.ReadVariable op and tf.AssignVariable op has device attribute
// added that is consistent with device attribute of TPUExecute op.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<!tf_type.resource<tensor<300x128xf32>>>, %arg1: tensor<!tf_type.resource<tensor<300x128xf32>>>) -> (tensor<300x128xf32>) {
%0 = "tf.StatefulPartitionedCall"(%arg0, %arg1) {config = "mesh:TPU,x=2,y=2", config_proto = "", executor_type = "", f = @tpu_func} : (tensor<!tf_type.resource<tensor<300x128xf32>>>, tensor<!tf_type.resource<tensor<300x128xf32>>>) -> (tensor<300x128xf32>)
func.return %0 :tensor<300x128xf32>
}
// CHECK-LABEL: func @tpu_func
// CHECK-SAME: %arg0: tensor<!tf_type.resource<tensor<300x128xf32>>>
// CHECK-SAME: tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"
// CHECK-SAME: %arg1: tensor<!tf_type.resource<tensor<300x128xf32>>>
// CHECK-SAME: tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"
func.func @tpu_func(%arg0: tensor<!tf_type.resource<tensor<300x128xf32>>>, %arg1: tensor<!tf_type.resource<tensor<300x128xf32>>>) -> tensor<300x128xf32> {
%0 = "tf.ReadVariableOp"(%arg0) {_global_shape = [#tf_type.shape<300x128>], _layout = ["mesh:TPU,x=2,y=2 layout:x,unsharded,"]} : (tensor<!tf_type.resource<tensor<300x128xf32>>>) -> tensor<300x128xf32>
%1:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {
metadata = "...",
mlir_module = ".."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%1#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
%2 = "tf_device.launch"() ({
%3 = "tf.TPUExecute"(%0, %1#1) : (tensor<300x128xf32>, tensor<2x!tf_type.string>) -> tensor<300x128xf32>
tf_device.return %3 : tensor<300x128xf32>
}) {device = "/job:localhost/replica:0/task:0/device:TPU:0"} : () -> tensor<300x128xf32>
"tf.AssignVariableOp"(%arg1, %2) : (tensor<!tf_type.resource<tensor<300x128xf32>>>, tensor<300x128xf32>) -> ()
func.return %2 : tensor<300x128xf32>
}
// -----
// Test that device attribute to resource input to TPU computation is correctly
// added.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<!tf_type.resource<tensor<300x128xf32>>>) -> (tensor<300x128xf32>) {
%0 = "tf.StatefulPartitionedCall"(%arg0) {config = "mesh:TPU,x=2,y=2", config_proto = "", executor_type = "", f = @tpu_func} : (tensor<!tf_type.resource<tensor<300x128xf32>>>) -> (tensor<300x128xf32>)
func.return %0 :tensor<300x128xf32>
}
// CHECK-LABEL: func @tpu_func
// CHECK-SAME: %arg0: tensor<!tf_type.resource<tensor<300x128xf32>>>
// CHECK-SAME: tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"
func.func @tpu_func(%arg0: tensor<!tf_type.resource<tensor<300x128xf32>>>) -> tensor<300x128xf32> {
%1:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {
metadata = "...",
mlir_module = ".."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%1#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
%2 = "tf_device.launch"() ({
%3 = "tf.TPUExecute"(%arg0, %1#1) : (tensor<!tf_type.resource<tensor<300x128xf32>>>, tensor<2x!tf_type.string>) -> tensor<300x128xf32>
tf_device.return %3 : tensor<300x128xf32>
}) {device = "/job:localhost/replica:0/task:0/device:TPU:0"} : () -> tensor<300x128xf32>
func.return %2 : tensor<300x128xf32>
}
tensorflow/dtensor/mlir/tests/tpu_integration.mlir
+50
View File
@@ -0,0 +1,50 @@
// 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: dtensor-opt %s -dtensor-tpu-integration | FileCheck %s
// Test that tf_device.Cluster op is created for tf.StatefulPartitionedCall that
// runs on TPU's.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0, %1 = "tf.StatefulPartitionedCall"(%arg0) {config = "|x=2,y=2|*TPU", config_proto = "", executor_type = "", f = @tpu_func} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
%2, %3 = "tf.StatefulPartitionedCall"(%arg0) {config = "|x=2,y=2|*TPU", config_proto = "", executor_type = "", f = @cpu_func} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.return %0, %1 : tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @tpu_func
func.func @tpu_func(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
// CHECK: tf_device.cluster
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Add
// CHECK-NEXT: tf_device.return
// CHECK: _tpu_replicate
// CHECK-SAME: device_assignment = []
// CHECK-SAME: num_cores_per_replica = 1
// CHECK-SAME: padding_map = []
// CHECK-SAME: step_marker_location = ""
// CHECK-SAME: topology = ""
// CHECK-SAME: use_spmd_for_xla_partitioning = false
%1 = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
%2 = "tf.Add"(%1, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
func.return %2, %arg0: tensor<i32>, tensor<i32>
}
// CHECK-LABEL: func @cpu_func
func.func @cpu_func(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
// CHECK-NOT: tf_device.Cluster
%0, %1 = "tf.A"(%arg0) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.return %0, %1: tensor<i32>, tensor<i32>
}
tensorflow/dtensor/mlir/tests/undo_merge_const_across_mesh.mlir
+31
View File
@@ -0,0 +1,31 @@
// 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: dtensor-opt %s -split-input-file -dtensor-undo-merge-const-across-mesh | FileCheck %s
// Check that constants with different meshes are duplicated.
// CHECK-LABEL: func @check_undo_sccp
func.func @check_undo_sccp() -> (tensor<4xi32>, tensor<4xi32>) {
// CHECK-DAG: "tf.DTensorLayout"(%[[CONST_A:.*]]) <{global_shape = #tf_type.shape<4>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>}> : (tensor<4xi32>) -> tensor<4xi32>
// CHECK-DAG: %[[CONST_A]] = "tf.Const"()
// CHECK-DAG: "tf.DTensorLayout"(%[[CONST_B:.*]]) <{global_shape = #tf_type.shape<4>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>}> : (tensor<4xi32>) -> tensor<4xi32>
// CHECK-DAG: %[[CONST_B]] = "tf.Const"()
%cst = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32>} : () -> tensor<4xi32>
%2 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<4>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|x=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<4xi32>) -> tensor<4xi32>
%3 = "tf.DTensorLayout"(%cst) {global_shape = #tf_type.shape<4>, layout = #dtensor.layout<sharding_specs:unsharded, mesh:|y=2|0,1|0,1|/job:localhost/replica:0/task:0/device:CPU:0,/job:localhost/replica:0/task:0/device:CPU:1>} : (tensor<4xi32>) -> tensor<4xi32>
func.return %2, %3 : tensor<4xi32>, tensor<4xi32>
}
tensorflow/dtensor/mlir/tests/update_tpu_metadata.mlir
+138
View File
@@ -0,0 +1,138 @@
// 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: dtensor-opt %s -split-input-file -dtensor-update-tpu-metadata -verify-diagnostics | FileCheck %s
// Check that TPUCompileMetadata proto is updated with correct number of replicas.
// CHECK-LABEL: func @main
func.func @main() {
"tf.StatefulPartitionedCall"() {config = ":|x=2,y=2|*TPU", config_proto = "", executor_type = "", f = @f_callee} : () -> ()
func.return
}
func.func @f_callee() {
// CHECK: tf_device.launch
// CHECK: device = ""
// CHECK: "tf._TPUCompileMlir"
// CHECK-SAME: metadata = "\0A\09\08\01\12\05\12\03\08\80\01\18\04 \01"
%0:2 = "tf_device.launch"() ({
%1, %2 = "tf._TPUCompileMlir"() {
NumDynamicShapes = 0 : i64,
metadata = "\0A\09\08\01\12\05\12\03\08\80\01\18\01 \01",
mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %1, %2 : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "tpu_host:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
// CHECK: "tf.TPUExecute"
"tf.TPUExecute"(%0#1) : (tensor<2x!tf_type.string>) -> ()
func.return
}
// -----
// Check that device placement of _TPUCompileMlir/TPUExeute operation is removed/updated properly.
// CHECK-LABEL: func @main
func.func @main() {
"tf.StatefulPartitionedCall"() {config = "|x=2,y=2|*TPU", config_proto = "", executor_type = "", f = @f_callee} : () -> ()
func.return
}
func.func @f_callee() {
// CHECK: tf_device.launch
// CHECK: device = ""
// CHECK: "tf._TPUCompileMlir"
%0:2 = "tf_device.launch"() ({
%1, %2 = "tf._TPUCompileMlir"() {
NumDynamicShapes = 0 : i64,
metadata = "\0A\09\08\01\12\05\12\03\08\80\01\18\01 \01",
mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %1, %2 : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "tpu_host:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
// CHECK: tf_device.launch
// CHECK: device = ""
// CHECK: "tf.TPUExecute"
"tf_device.launch"() ({
"tf.TPUExecute"(%0#1) : (tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:TPU:0"} : () -> ()
func.return
}
// -----
// Check that unparable TPUCompilaMetadataProto is disallowed.
func.func @main() {
"tf.StatefulPartitionedCall"() {config = "|x=2,y=2|*TPU", config_proto = "", executor_type = "", f = @f_callee} : () -> ()
func.return
}
func.func @f_callee() {
%0:2 = "tf_device.launch"() ({
// expected-error @+1 {{unable to parse TPUCompileMetadata}}
%1, %2 = "tf._TPUCompileMlir"() {
NumDynamicShapes = 0 : i64,
metadata = "\0A\0B\0C",
mlir_module = "..."} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
tf_device.return %1, %2 : tensor<!tf_type.string>, tensor<2x!tf_type.string>
}) {device = "tpu_host:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUExecute"(%0#1) : (tensor<2x!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:TPU:0"} : () -> ()
func.return
}
// -----
// Check for Xla Spmd mesh that TPUCompileOp has correct metadata proto and
// number of program outputs is equal to number of devices on mesh.
// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i32>, %arg1: tensor<12x24xf32>) -> (tensor<12x24xf32>) {
%0 = "tf.StatefulPartitionedCall"(%arg1) {
config = "|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/replica:0/task:0/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1,/job:localhost/replica:0/task:0/device:TPU:2,/job:localhost/replica:0/task:0/device:TPU:3,/job:localhost/replica:0/task:0/device:TPU:4,/job:localhost/replica:0/task:0/device:TPU:5,/job:localhost/replica:0/task:0/device:TPU:6,/job:localhost/replica:0/task:0/device:TPU:7|use_xla_spmd",
config_proto = "",
executor_type = "",
f = @_xla_spmd_func} : (tensor<12x24xf32>) -> tensor<12x24xf32>
return %0 : tensor<12x24xf32>
}
func.func private @_xla_spmd_func(%arg0: tensor<12x24xf32>) -> tensor<12x24xf32> {
// CHECK: tf_device.launch
// CHECK: device = ""
// CHECK: %compilation_status, %program:8 = "tf._TPUCompileMlir"
// CHECK-SAME: metadata = "\0A\10\08\01\12\08\12\02\08\0C\12\02\08\18\18\01\22\00\12\02\0A\00\18\01 \08x\01\88\01\ED\91\DC\F5\C3\8C\95\B5\90\01"
%0:2 = "tf_device.launch"() ({
%compilation_status, %program = "tf._TPUCompileMlir"() {metadata = "\0A\18\08\01\12\08\12\02\08\0C\12\02\08\18\18\01\22\08\08\01\1A\01\01\22\01\00\12\0A\0A\08\08\01\1A\01\01\22\01\00\18\01 \01\88\01\ED\91\DC\F5\C3\8C\95\B5\90\01", mlir_module = "#loc = loc(unknown)\0Amodule attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1345 : i32}} {\0A func.func @main(%arg0: tensor<12x24xf32> {mhlo.sharding = \22\22} loc(unknown)) -> (tensor<12x24xf32> {mhlo.sharding = \22\22}) {\0A %0 = \22tf.Identity\22(%arg0) : (tensor<12x24xf32>) -> tensor<12x24xf32> loc(#loc)\0A return %0 : tensor<12x24xf32> loc(#loc)\0A } loc(#loc)\0A} loc(#loc)\0A"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<3x!tf_type.string>
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
"tf_device.launch"() ({
"tf.TPUCompileSucceededAssert"(%0#0) : (tensor<!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()
%1 = "tf_device.launch"() ({
// CHECK: "tf.TPUExecute"
%2 = "tf.TPUExecute"(%arg0, %0#1) : (tensor<12x24xf32>, tensor<3x!tf_type.string>) -> tensor<12x24xf32>
tf_device.return %2 : tensor<12x24xf32>
}) {device = "/job:localhost/replica:0/task:0/device:TPU:0"} : () -> tensor<12x24xf32>
return %1 : tensor<12x24xf32>
}
func.func private @_func(%arg0: tensor<12x24xf32> {mhlo.sharding = ""}) -> (tensor<12x24xf32> {mhlo.sharding = ""}) {
%0 = "tf.Identity"(%arg0) : (tensor<12x24xf32>) -> tensor<12x24xf32>
return %0 : tensor<12x24xf32>
}