Home
last modified time | relevance | path

Searched refs:tpu_core (Results 1 – 10 of 10) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf_communication.cc214 void SetOpSharding(Operation* op, int64_t tpu_core) { in SetOpSharding() argument
216 ::xla::sharding_builder::AssignDevice(tpu_core).SerializeAsString(); in SetOpSharding()
256 const Optional<int64_t>& tpu_core, Value token) { in CreateSendOp() argument
268 if (tpu_core) SetOpSharding(send, *tpu_core); in CreateSendOp()
277 const Optional<int64_t>& tpu_core, Value token) { in CreateRecvOp() argument
292 if (tpu_core) SetOpSharding(recv, *tpu_core); in CreateRecvOp()
296 if (tpu_core) SetOpSharding(get_tuple_element, *tpu_core); in CreateRecvOp()
302 if (tpu_core) SetOpSharding(new_token, *tpu_core); in CreateRecvOp()
331 int64_t tpu_core = host_compute.tpu_coreAttr().getInt(); in RewriteHostComputeOp() local
337 host_compute.send_key(), operand.index(), tpu_core, token); in RewriteHostComputeOp()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/
Dlegalize-tf-communication.mlir6 // and op shardings (based on `tpu_core`) are added. Sink tokens are created
73 … = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 0 : i64} : (ten…
81 // Tests `tf._XlaHostComputeMlir` with `tpu_core` assigns the correct op
103 … = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 1 : i64} : (ten…
119 … = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 0 : i64} : () -…
135 … = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 0 : i64} : (ten…
156 … = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 0 : i64} : (ten…
229 …%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "recv0", send_key = "send0", tpu_core = 0 : i64} …
244 …%1 = "tf._XlaHostComputeMlir"(%0) {recv_key = "recv1", send_key = "send1", tpu_core = 0 : i64} : (…
279 …%0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "recv", send_key = "send", tpu_core = 0 : i64} : …
[all …]
/external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/
DXlaHostCompute.pbtxt47 name: "tpu_core"
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/
Dpod_tpu_driver.cc305 for (auto& tpu_core : *tpu_chip.mutable_core()) { in PodTpuDriver()
306 int current_core = tpu_core.id(); in PodTpuDriver()
314 core_to_driver_core_.insert({current_core, tpu_core.id()}); in PodTpuDriver()
316 tpu_core.set_id(current_core); in PodTpuDriver()
317 tpu_core.set_core_on_host_index(current_core); in PodTpuDriver()
318 *(pod_info_.add_local_core()) = tpu_core; in PodTpuDriver()
/external/tensorflow/tensorflow/core/api_def/base_api/
Dapi_def_XlaHostCompute.pbtxt60 name: "tpu_core"
/external/tensorflow/tensorflow/compiler/jit/
Dextract_outside_compilation_pass_test.cc315 int tpu_core; in TEST_F() local
316 TF_CHECK_OK(GetNodeAttr(host_compute_0->attrs(), "tpu_core", &tpu_core)); in TEST_F()
317 EXPECT_EQ(tpu_core, 1); in TEST_F()
318 TF_CHECK_OK(GetNodeAttr(host_compute_1->attrs(), "tpu_core", &tpu_core)); in TEST_F()
319 EXPECT_EQ(tpu_core, 0); in TEST_F()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/
Dtpu_extract_outside_compilation.mlir587 // CHECK-SAME: tpu_core = 0
690 // CHECK-SAME: tpu_core = 0
754 // CHECK-SAME: tpu_core = 0
823 // CHECK-SAME: tpu_core = 0
1233 // CHECK-SAME: tpu_core = 0
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtf_passes.td380 …mpute_channel_0_0_retvals", send_key = "host_compute_channel_0_0_args", tpu_core = 0 : i64} : (ten…
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_generated_ops.td17830 DefaultValuedAttr<I64Attr, "0">:$tpu_core
18597 DefaultValuedAttr<I64Attr, "0">:$tpu_core
/external/tensorflow/tensorflow/core/ops/
Dops.pbtxt58666 name: "tpu_core"