Searched refs:host_compute (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_communication.cc | 327 TF::_XlaHostComputeMlirOp host_compute, in RewriteHostComputeOp() argument
329 builder.setInsertionPoint(host_compute); in RewriteHostComputeOp()
330 Location loc = host_compute.getLoc(); in RewriteHostComputeOp()
331 int64_t tpu_core = host_compute.tpu_coreAttr().getInt(); in RewriteHostComputeOp()
334 for (auto operand : llvm::enumerate(host_compute.inputs())) { in RewriteHostComputeOp()
337 host_compute.send_key(), operand.index(), tpu_core, token); in RewriteHostComputeOp()
343 for (auto result : llvm::enumerate(host_compute.outputs())) { in RewriteHostComputeOp()
346 host_compute.recv_key(), result.index(), tpu_core, token); in RewriteHostComputeOp()
351 host_compute.erase(); in RewriteHostComputeOp()
771 if (auto host_compute = dyn_cast<TF::_XlaHostComputeMlirOp>(curr_op)) { in RewriteFunction() local
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_extract_outside_compilation.cc | 245 auto host_compute = builder.create<TF::_XlaHostComputeMlirOp>( in CreateHostCompute() local
250 return host_compute; in CreateHostCompute()
292 auto host_compute = CreateHostCompute( in MoveOpsToHost() local
337 for (auto result : llvm::zip(external_outputs, host_compute.getResults())) { in MoveOpsToHost()
343 host_compute.erase(); in MoveOpsToHost()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/tests/ |
| D | legalize-tf-communication.mlir | 17 // CHECK-LABEL: func @host_compute
19 func @host_compute(%arg0: tensor<i32>, %arg1: tensor<i64>) -> (tensor<f32>, tensor<f64>) {
|