Home
last modified time | relevance | path

Searched refs:lowered_callee (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dstack_ops_decomposition.cc335 FuncOp lowered_callee = callee; in HandlePartitionedCallOp() local
338 lowered_callee = callee.clone(); in HandlePartitionedCallOp()
339 lowered_callee.setPrivate(); in HandlePartitionedCallOp()
346 ModifyFunctionSignature(lowered_callee, &callee_map, find_arg_stack_type); in HandlePartitionedCallOp()
350 if (lowered_callee != callee) { in HandlePartitionedCallOp()
351 lowered_callee.erase(); in HandlePartitionedCallOp()
354 info.decomposed_callee = lowered_callee; in HandlePartitionedCallOp()
360 if (lowered_callee != callee) { in HandlePartitionedCallOp()
362 lowered_callee.setName( in HandlePartitionedCallOp()
364 SymbolTable(module).insert(lowered_callee); in HandlePartitionedCallOp()
[all …]
Dtensor_list_ops_decomposition.cc479 FuncOp lowered_callee = callee; in HandlePartitionedCallOp() local
482 lowered_callee = callee.clone(); in HandlePartitionedCallOp()
483 lowered_callee.setPrivate(); in HandlePartitionedCallOp()
493 ModifyFunctionSignature(lowered_callee, cutil::GetSizeType(OpBuilder(call)), in HandlePartitionedCallOp()
498 &lowered_callee.front(), module, &callee_map, in HandlePartitionedCallOp()
503 ModifyFunctionReturn(lowered_callee, callee_map); in HandlePartitionedCallOp()
504 info.decomposed_callee = lowered_callee; in HandlePartitionedCallOp()
508 if (lowered_callee != callee) { in HandlePartitionedCallOp()
509 lowered_callee.setName(callee.getName()); in HandlePartitionedCallOp()
511 SymbolTable(module).insert(lowered_callee); in HandlePartitionedCallOp()
[all …]
Dtensor_array_ops_decomposition.cc784 FuncOp lowered_callee = callee; in HandlePartitionedCallOp() local
787 lowered_callee = callee.clone(); in HandlePartitionedCallOp()
788 lowered_callee.setPrivate(); in HandlePartitionedCallOp()
790 auto grads = AccessedGradients({lowered_callee}, module); in HandlePartitionedCallOp()
791 for (int64_t i = 0; i < lowered_callee.getNumArguments(); ++i) { in HandlePartitionedCallOp()
797 ChangeFunctionInputSignature(lowered_callee, grads, ta_arg_buffer_type, in HandlePartitionedCallOp()
799 if (failed(DecomposeTensorArrayOps(&lowered_callee.front(), module, in HandlePartitionedCallOp()
805 auto ret = lowered_callee.front().getTerminator()->getOperand(i); in HandlePartitionedCallOp()
812 info.decomposed_callee = lowered_callee; in HandlePartitionedCallOp()
813 if (lowered_callee != callee) { in HandlePartitionedCallOp()
[all …]