Searched refs:compile_op (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/dtensor/mlir/ |
| D | move_compilation_to_host.cc | 89 mlir::TF::_TPUCompileMlirOp compile_op, in CreateSendRecvOpsToTransferProgramKey() argument
93 mlir::Value compilation_key = *compile_op.program().begin(); in CreateSendRecvOpsToTransferProgramKey()
108 builder.setInsertionPointAfter(compile_op); in CreateSendRecvOpsToTransferProgramKey()
112 compile_op->getLoc(), compilation_key, tensor_name, in CreateSendRecvOpsToTransferProgramKey()
134 compile_op.getLoc(), in CreateSendRecvOpsToTransferProgramKey()
143 compile_op->getLoc(), in CreateSendRecvOpsToTransferProgramKey()
167 compile_op.getLoc(), in CreateSendRecvOpsToTransferProgramKey()
189 mlir::TF::_TPUCompileMlirOp compile_op, in SendRecvCompilationKey() argument
199 info.recv_insertion_point, compile_op, compile_launch_op, in SendRecvCompilationKey()
263 mlir::TF::_TPUCompileMlirOp compile_op = *compile_ops.begin(); in HandleCompilationOps() local
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/ |
| D | fuse_tpu_compile_and_execute_ops.cc | 83 auto compile_op = ::llvm::dyn_cast<mlir::TF::_TPUCompileMlirOp>( in runOnOperation() local
85 if (!compile_op) { in runOnOperation()
91 builder.setInsertionPointAfter(compile_op); in runOnOperation()
129 compile_op.mlir_module(), compile_op.metadata(), producer_name); in runOnOperation()
132 for (auto program_result : compile_op.program()) { in runOnOperation()
139 assert(compile_op.use_empty()); in runOnOperation()
140 compile_op.erase(); in runOnOperation()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 389 auto compile_op = builder->create<TF::_TPUCompileMlirOp>( in BuildCompileOp() local
395 return WrapOpInLaunch(builder, compile_op.getLoc(), compile_op, in BuildCompileOp()
510 Operation* compile_op, tf_device::ClusterFuncOp cluster_func, in AddToParallelExecuteOp() argument
562 execute_inputs.emplace_back(compile_op->getResult(core + 1)); in AddToParallelExecuteOp()
601 void BuildTPUCompileSucceededAssertOp(Operation* compile_op, in BuildTPUCompileSucceededAssertOp() argument
606 compile_op->getLoc(), result_id->getResult(0)); in BuildTPUCompileSucceededAssertOp()
607 WrapOpInLaunch(builder, compile_op->getLoc(), assert_op, compilation_device); in BuildTPUCompileSucceededAssertOp()
769 Operation* compile_op = in Rewrite() local
774 if (!compile_op) return failure(); in Rewrite()
783 key_op.replaceAllUsesWith(compile_op->getResult(1)); in Rewrite()
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu.py | 1884 compile_op, outputs = split_compile_and_replicate(
1902 return compile_op, [outputs[0]]
1939 return compile_op, results
|