Searched refs:compile_launch (Results 1 – 3 of 3) sorted by relevance
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
D | tpu_dynamic_layout_pass.cc | 143 tf_device::LaunchOp compile_launch, in BuildGetLayout() argument 146 compile_launch.getLoc(), in BuildGetLayout() 158 tf_device::LaunchOp compile_launch, in BuildCopyWithLayout() argument 169 tf_device::LaunchOp compile_launch) { in HandleInput() argument 170 OpBuilder builder = CreateBuilderAfterOp(compile_launch); in HandleInput() 172 compile_launch, &builder); in HandleInput() 174 auto copy_with_layout = BuildCopyWithLayout(execute_launch, compile_launch, in HandleInput() 184 tf_device::LaunchOp execute_launch, tf_device::LaunchOp compile_launch, in HandleReplicatedInputs() argument 197 OpBuilder builder = CreateBuilderAfterOp(compile_launch); in HandleReplicatedInputs() 199 compile_launch, &builder); in HandleReplicatedInputs() [all …]
|
D | tpu_variable_runtime_reformatting.cc | 148 tf_device::LaunchOp compile_launch) { in AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping() argument 175 Operation& compile = compile_launch.GetBody().front(); in AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping() 408 auto compile_launch = llvm::dyn_cast<tf_device::LaunchOp>(compile); in HandleReplicateOp() local 409 if (!compile_launch || !compile_launch.WrapsSingleOp() || in HandleReplicateOp() 410 !llvm::isa<TF::_TPUCompileMlirOp>(compile_launch.GetBody().front())) in HandleReplicateOp() 416 while_op, replicate, execute, compile_launch); in HandleReplicateOp() 454 reformat_operands.push_back(compile_launch.getResult(1)); in HandleReplicateOp()
|
D | tpu_merge_variables_with_execute.cc | 428 auto compile_launch = llvm::dyn_cast_or_null<tf_device::LaunchOp>( in GetTPUCompileOp() local 430 if (!compile_launch) return {}; in GetTPUCompileOp() 432 compile_launch.GetBody().front()); in GetTPUCompileOp()
|