| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
| D | tpu_extract_head_tail_outside_compilation.mlir | 8 // CHECK: "tf_device.launch"
11 // CHECK-NEXT: tf_device.return
14 // CHECK: "tf_device.cluster"
17 // CHECK-NEXT: tf_device.return
18 "tf_device.cluster"() ( {
22 tf_device.return
29 // CHECK: %[[LAUNCH_OUT:.*]] = "tf_device.launch"
32 // CHECK-NEXT: tf_device.return %[[A_OUT]]
35 // CHECK: "tf_device.cluster"
38 // CHECK-NEXT: tf_device.return
[all …]
|
| D | tf_device_ops.mlir | 5 "tf_device.launch"() ( {
6 // CHECK: tf_device.return
7 tf_device.return
15 %result = "tf_device.launch"() ( {
16 // CHECK: tf_device.return %[[ARG_0]] : tensor<*xf32>
17 tf_device.return %arg_0 : tensor<*xf32>
25 %result:2 = "tf_device.launch"() ( {
26 // CHECK: tf_device.return %[[ARG_0]], %[[ARG_1]] : tensor<*xf32>, tensor<*xi32>
27 tf_device.return %arg_0, %arg_1 : tensor<*xf32>, tensor<*xi32>
34 tf_device.replicate {n = 2 : i32} {
[all …]
|
| D | tf_device_ops_invalid.mlir | 5 tf_device.replicate([%arg0] as %input0: tensor<*xf32>) {n = 1 : i32} {
6 // expected-error@-1 {{'tf_device.replicate' expects 'n' to be at least 2, got 1}}
16 tf_device.replicate([%arg0, %arg0, %arg0] as %input0: tensor<*xf32>) {n = 2 : i32} {
17 // expected-error@-1 {{'tf_device.replicate' expects number of operands for replicated input 0 to b…
28 tf_device.replicate([%arg0, %arg0] as %input0: tensor<*xf32>) {n = 2 : i32} {
38 tf_device.replicate() {n = 2 : i32} {
39 // expected-error@-1 {{custom op 'tf_device.replicate' expects a single block region}}
42 tf_device.return
51 tf_device.replicate() {n = 2 : i32} {
52 // expected-error@-1 {{custom op 'tf_device.replicate' expects a tf_device.return terminator}}
[all …]
|
| D | outside_compiled_to_host_launch.mlir | 7 // CHECK-NOT: "tf_device.launch"
9 %0 = "tf_device.cluster"() ( {
12 tf_device.return %2 : tensor<?xi32>
23 // CHECK: "tf_device.launch"
26 // CHECK-NEXT: tf_device.return
29 "tf_device.cluster"() ( {
33 tf_device.return
43 // CHECK: tf_device.replicate
44 // CHECK-NEXT: "tf_device.cluster"
46 // CHECK-NEXT: "tf_device.launch"
[all …]
|
| D | replicate_invariant_op_hoisting.mlir | 6 %0:4 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*xf32>) {n = 2: i32} {
9 tf_device.return %1, %2 : tensor<?xi32>, tensor<*xi32>
15 // CHECK: tf_device.replicate([%[[ARG_0]], %[[ARG_1]]] as %[[RI:[a-z0-9]*]]: tensor<*xf32>)
17 // CHECK: tf_device.return %[[SHAPE]], %[[OP_A]]
23 %0:2 = tf_device.replicate([%arg0, %arg0] as %ri: tensor<*xf32>) {n = 2: i32} {
25 tf_device.return %1 : tensor<?xi32>
31 // CHECK: tf_device.replicate
32 // CHECK: tf_device.return %[[SHAPE]]
38 %0:6 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*x!tf.resource>) {n = 2: i32} {
42 tf_device.return %1, %2, %3 : tensor<*xf32>, tensor<?xi32>, tensor<*xi32>
[all …]
|
| D | tpu_extract_outside_compilation.mlir | 10 %0 = "tf_device.cluster"() ( {
13 tf_device.return %2 : tensor<?xi32>
18 // CHECK-NOT: "tf_device.parallel_execute"
24 // CHECK: "tf_device.parallel_execute"
25 // CHECK-NEXT: "tf_device.launch"
28 // CHECK-NEXT: tf_device.return
30 // CHECK: "tf_device.cluster"
33 "tf_device.cluster"() ( {
37 tf_device.return
46 // CHECK: "tf_device.parallel_execute"
[all …]
|
| D | tpu-dynamic-layout-pass.mlir | 7 // CHECK: %[[COMPILE:.*]]:2 = "tf_device.launch"
9 %compile:2 = "tf_device.launch"() ( {
15 tf_device.return %1#0, %1#1 : tensor<!tf.string>, tensor<2x!tf.string>
22 // CHECK: "tf_device.launch"
24 "tf_device.launch"() ( {
26 tf_device.return
30 // CHECK: "tf_device.launch"
32 %execute = "tf_device.launch"() ( {
35 tf_device.return %3 : tensor<i32>
47 %compile:2 = "tf_device.launch"() ( {
[all …]
|
| D | tpu_colocate_composite_resource_ops.mlir | 4 // tf_device.Cluster.
9 // CHECK: tf_device.replicate
11 tf_device.replicate(%arg0 as %arg1: tensor<*x!tf.resource<tensor<4xf32>>>) {
15 // CHECK: %[[RESOURCE_OUT:.*]] = "tf_device.launch"()
17 // CHECK-NEXT: tf_device.return %[[READ_OUT]]
21 "tf_device.launch"() ( {
23 tf_device.return
25 "tf_device.launch"() ( {
28 tf_device.return
30 tf_device.return
[all …]
|
| D | tpu_parallel_execute_sink_resource_write.mlir | 6 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute"
7 %0:2 = "tf_device.parallel_execute"() ( {
8 tf_device.return %arg0 : tensor<i1>
10 tf_device.return %arg0 : tensor<i1>
22 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute"
23 %0:2 = "tf_device.parallel_execute"() ( {
24 tf_device.return %arg0 : tensor<i1>
26 tf_device.return %arg0 : tensor<i1>
37 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute"
38 %0:2 = "tf_device.parallel_execute"() ( {
[all …]
|
| D | tpu_host_computation_expansion.mlir | 7 // CHECK: "tf_device.cluster"
10 "tf_device.cluster"() ( {
14 tf_device.return
21 // CHECK: "tf_device.cluster"
24 "tf_device.cluster"() ( {
28 tf_device.return
35 // CHECK: "tf_device.cluster"
41 "tf_device.cluster"() ( {
46 tf_device.return
53 // CHECK: "tf_device.cluster"
[all …]
|
| D | tpu_dynamic_padding_mapper.mlir | 12 …tf_device.replicate([%arg0, %arg0] as %ri_0: tensor<i1>, [%arg0, %arg0] as %ri_1: tensor<i1>) {_re…
13 …"tf_device.cluster_func"(%ri_0, %ri_1) {func = @func0, padding_map = ["\10\02\18\01"]} : (tensor<i…
14 tf_device.return
39 …tf_device.replicate([%arg0, %arg0] as %ri_0: tensor<i1>, [%arg0, %arg0] as %ri_1: tensor<i1>, [%ar…
40 …"tf_device.cluster_func"(%ri_0, %ri_1, %ri_2) {func = @func1, padding_map = ["\10\02\18\01", "\10\…
41 tf_device.return
71 …tf_device.replicate([%arg0, %arg0] as %ri_0: tensor<i1>, [%arg0, %arg0] as %ri_1: tensor<i1>, [%ar…
72 …"tf_device.cluster_func"(%ri_0, %ri_1, %ri_2, %ri_3, %ri_4) {func = @func2, padding_map = ["\10\02…
73 tf_device.return
92 …tf_device.replicate([%arg0, %arg0] as %ri_0: tensor<i1>, [%arg0, %arg0] as %ri_1: tensor<i1>) {_re…
[all …]
|
| D | launch_to_device_attribute.mlir | 5 // by parent `tf_device.launch`.
11 %launch:2 = "tf_device.launch"() ( {
13 tf_device.return %b#1, %b#0 : tensor<f32>, tensor<i32>
27 // CHECK-NOT: "tf_device.launch"
32 // assigned by parent `tf_device.launch`.
38 %launch:2 = "tf_device.launch"() ( {
41 tf_device.return %c, %b : tensor<f32>, tensor<i32>
57 // CHECK-NOT: "tf_device.launch"
66 %launch:2 = "tf_device.launch"() ( {
68 tf_device.return %a#1, %a#0 : tensor<f32>, tensor<i32>
[all …]
|
| D | cluster_outlining.mlir | 3 // Tests simple case of a single `tf_device.cluster`.
13 …// CHECK: %[[CLUSTER_OUTPUT:[0-9]*]] = "tf_device.cluster_func"(%[[A_OUTPUT]]) {func = @[[CLUSTER:…
14 %3 = "tf_device.cluster"() ( {
16 tf_device.return %4 : tensor<?xi32>
34 // Tests that multiple `tf_device.cluster` that depend on each other are
45 …// CHECK: %[[CLUSTER_0_OUTPUT:[0-9]*]] = "tf_device.cluster_func"(%[[A_OUTPUT]]) {func = @[[CLUSTE…
46 %3 = "tf_device.cluster"() ( {
48 tf_device.return %6 : tensor<?xi32>
54 …// CHECK: %[[CLUSTER_1_OUTPUT:[0-9]*]] = "tf_device.cluster_func"(%[[CLUSTER_0_OUTPUT]], %[[D_OUTP…
55 %5 = "tf_device.cluster"() ( {
[all …]
|
| D | tpu-merge-variables-with-execute.mlir | 20 // CHECK: %[[COMPILE:.*]]:2 = "tf_device.launch"
21 %compile:2 = "tf_device.launch"() ( {
30 tf_device.return %0#0, %0#1 : tensor<!tf.string>, tensor<2x!tf.string>
32 // CHECK: %[[EXE:.*]] = "tf_device.launch"
36 %execute:2 = "tf_device.launch"() ( {
41 tf_device.return %0#0, %0#1 : tensor<32xf32>, tensor<16xf32>
43 // CHECK-NEXT: tf_device.return
65 // CHECK: %[[COMPILE:.*]]:2 = "tf_device.launch"
66 %compile:2 = "tf_device.launch"() ( {
74 tf_device.return %0#0, %0#1 : tensor<!tf.string>, tensor<2x!tf.string>
[all …]
|
| D | replicate_to_island.mlir | 11 tf_device.replicate {n = 2 : i32} {
12 tf_device.return
32 tf_device.replicate {n = 2 : i32} {
33 "tf_device.launch"() ( {
35 tf_device.return
37 tf_device.return
57 tf_device.replicate {n = 2 : i32, devices = {CORE_0 = ["/CPU:0", "/GPU:1"]}} {
58 "tf_device.launch"() ( {
60 tf_device.return
62 tf_device.return
[all …]
|
| D | tpu_rewrite.mlir | 8 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
23 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
33 // Tests `tf_device.cluster_func` with missing `num_cores_per_replicas`
39 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, step_marker_location …
49 // Tests `tf_device.cluster_func` with bad `num_cores_per_replicas` attribute.
54 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
64 // Tests `tf_device.cluster_func` with missing `step_marker_location` attribute.
69 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
79 // Tests `tf_device.cluster_func` with bad `step_marker_location` attribute.
84 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
[all …]
|
| D | tpu-variable-runtime-reformatting.mlir | 39 // CHECK: %[[COMPILE:.*]]:2 = "tf_device.launch"
41 %compile:2 = "tf_device.launch"() ( {
47 tf_device.return %b2#0, %b2#1 : tensor<!tf.string>, tensor<2x!tf.string>
49 "tf_device.launch"() ( {
51 tf_device.return
53 // CHECK: tf_device.replicate
58 … %rep:2 = tf_device.replicate([%arg0, %arg1] as %arg30: tensor<*x!tf.resource<tensor<f32>>>,
63 // CHECK: "tf_device.launch"
65 // CHECK-NEXT: tf_device.return
68 "tf_device.launch"() ( {
[all …]
|
| D | mark_ops_for_outside_compilation.mlir | 5 %0 = "tf_device.cluster"() ( {
12 tf_device.return %2 : tensor<i32>
19 %0 = "tf_device.cluster"() ( {
26 tf_device.return %2 : tensor<i32>
33 %0 = "tf_device.cluster"() ( {
44 tf_device.return %2 : tensor<i32>
51 %0 = "tf_device.cluster"() ( {
58 tf_device.return %2 : tensor<i32>
65 %0 = "tf_device.cluster"() ( {
76 tf_device.return %4 : tensor<f32>
[all …]
|
| D | tpu_identity_pruning.mlir | 9 // CHECK: "tf_device.cluster"
10 // CHECK-NEXT: tf_device.return [[ARG0]]
11 %0 = "tf_device.cluster"() ( {
13 tf_device.return %1 : tensor<i32>
24 // CHECK: "tf_device.cluster"
25 // CHECK-NEXT: tf_device.return [[ARG0]], [[ARG1]]
26 %0:2 = "tf_device.cluster"() ( {
28 tf_device.return %1#0, %1#1 : tensor<i32>, tensor<f32>
38 // CHECK: "tf_device.cluster"
41 %0 = "tf_device.cluster"() ( {
[all …]
|
| D | parallel_execute_to_islands.mlir | 7 "tf_device.parallel_execute"() ({
8 tf_device.return
10 tf_device.return
35 %4:2 = "tf_device.parallel_execute"() ({
37 tf_device.return %5 : tensor<i1>
40 tf_device.return %5 : tensor<i32>
68 %3:2 = "tf_device.parallel_execute"() ({
70 tf_device.return %4 : tensor<i1>
73 tf_device.return %4 : tensor<i32>
96 %1:2 = "tf_device.parallel_execute"() ({
[all …]
|
| D | merge_control_flow.mlir | 7 // CHECK: tf_device.cluster
10 "tf_device.cluster"() ( {
25 tf_device.return
34 // CHECK: tf_device.cluster
37 "tf_device.cluster"() ( {
61 tf_device.return
70 // CHECK: tf_device.cluster
76 "tf_device.cluster"() ( {
90 tf_device.return
99 // CHECK: tf_device.cluster
[all …]
|
| D | cluster_formation.mlir | 12 // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
19 // CHECK: tf_device.return %[[C_OUTPUT]]
42 // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
49 // CHECK: tf_device.return %[[C_OUTPUT]]
73 // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
80 // CHECK: tf_device.return %[[B_OUTPUT]]
106 // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
113 // CHECK: tf_device.return %[[B_OUTPUT]]
137 // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
141 // CHECK: tf_device.return %[[A_OUTPUT]]
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tf_passes.td | 148 `tf_device.cluster`, and copies over attributes from the associated
149 `tf.TPUReplicateMetadata` op to the newly created `tf_device.cluster`. If the
151 not copied over but instead the `tf_device.cluster` is further wrapped with a
152 `tf_device.replicate`, and associated `tf.TPUReplicatedInput` and
153 `tf.TPUReplicatedOutput` ops are replaced as the `tf_device.replicate` operands
156 the `tf_device.cluster`.
176 %cluster = "tf_device.cluster"() ( {
178 tf_device.return %identity : tensor<i32>
200 %replicate:2 = tf_device.replicate([%arg0, %arg1] as %replicated_input) {n = 2 : i32} {
201 %cluster = "tf_device.cluster"() ( {
[all …]
|
| D | tpu_extract_head_tail_outside_compilation.cc | 86 tf_device::LaunchOp CreateLaunchForBlock(OpBuilder* builder, Operation* op, in CreateLaunchForBlock()
108 auto launch = builder->create<tf_device::LaunchOp>( in CreateLaunchForBlock()
113 builder->create<tf_device::ReturnOp>(op->getLoc(), launch_results); in CreateLaunchForBlock()
132 tf_device::ClusterOp cluster) { in FindOutsideCompiledOpsAtHead()
181 void CreateHeadComputation(OpBuilder* builder, tf_device::ClusterOp cluster, in CreateHeadComputation()
190 tf_device::LaunchOp launch = CreateLaunchForBlock( in CreateHeadComputation()
203 const mlir::TF::RuntimeDevices& devices, tf_device::ClusterOp cluster, in LiftHeadOutsideCompiledOps()
226 tf_device::ClusterOp cluster, in FindOutsideCompiledOpsAtTailAndClusterResults()
289 void CreateTailComputation(OpBuilder* builder, tf_device::ClusterOp cluster, in CreateTailComputation()
298 tf_device::LaunchOp launch = CreateLaunchForBlock( in CreateTailComputation()
[all …]
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | device_util.py | 23 from tensorflow.python.framework import device as tf_device unknown
49 d = tf_device.DeviceSpec.from_string(d.name)
51 d = tf_device.DeviceSpec.from_string(d)
56 result = tf_device.DeviceSpec(
62 host_cpu = tf_device.DeviceSpec.from_string(
72 tf_device.DeviceSpec.from_string(default))
124 spec = tf_device.DeviceSpec.from_string(device)
125 return tf_device.DeviceSpec(
|