experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[MLIR] Adopt Symbol setPublic/setPrivate/setNested functions.

Browse Source 
PiperOrigin-RevId: 342070484
Change-Id: I74b0592f28a7e4ce85b6ae67b986cb94a06961fb
This commit is contained in:
Rahul Joshi 2020-11-12 09:57:27 -08:00 committed by TensorFlower Gardener
parent 68f4b84637
commit 3ea78d0fdd
15 changed files with 16 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/compiler/mlir/lite/flatbuffer_import.cc
View File

@ -988,7 +988,7 @@ StatusOr<FuncOp> ConvertSubgraph(
} }
func.setAttr("tf.entry_function", builder.getDictionaryAttr(attributes)); func.setAttr("tf.entry_function", builder.getDictionaryAttr(attributes));
} else { } else {
func.setVisibility(FuncOp::Visibility::Private); func.setPrivate();
} }
absl::flat_hash_set<const tflite::OperatorT*> pruned_subgraph_ops; absl::flat_hash_set<const tflite::OperatorT*> pruned_subgraph_ops;

2
tensorflow/compiler/mlir/lite/transforms/legalize_tf_while.cc
View File

@ -62,7 +62,7 @@ void RunOnWhile(TF::WhileOp while_op) {
auto call = builder.create<CallOp>(while_op.getLoc(), func, new_operands); auto call = builder.create<CallOp>(while_op.getLoc(), func, new_operands);
builder.create<YieldOp>(while_op.getLoc(), call.getResults()); builder.create<YieldOp>(while_op.getLoc(), call.getResults());
// Mark old function as private so that it can be DCE'd if not called. // Mark old function as private so that it can be DCE'd if not called.
func.setVisibility(SymbolTable::Visibility::Private); func.setPrivate();
}; };
create_region_with_call(while_op.cond_function(), new_op.cond()); create_region_with_call(while_op.cond_function(), new_op.cond());
create_region_with_call(while_op.body_function(), new_op.body()); create_region_with_call(while_op.body_function(), new_op.body());

2
tensorflow/compiler/mlir/lite/transforms/while_loop_outline.cc
View File

@ -182,7 +182,7 @@ void WhileOutlinePass::OutlineWhile(WhileOp while_op) {
b.create<ReturnOp>(yield_op->getLoc(), args); b.create<ReturnOp>(yield_op->getLoc(), args);
yield_op->erase(); yield_op->erase();
symbol_table.insert(outlined_func); symbol_table.insert(outlined_func);
outlined_func.setVisibility(FuncOp::Visibility::Private); outlined_func.setPrivate();
return outlined_func; return outlined_func;
}; };

2
tensorflow/compiler/mlir/tensorflow/transforms/cluster_outlining.cc
View File

@ -68,7 +68,7 @@ FuncOp BuildFunction(llvm::ArrayRef<Value> live_ins,
// This function is not externally visible and marking it private would allow // This function is not externally visible and marking it private would allow
// symbol-dce pass to remove it when it is not referenced anymore. // symbol-dce pass to remove it when it is not referenced anymore.
outlined_func.setVisibility(FuncOp::Visibility::Private); outlined_func.setPrivate();
// Create function body. // Create function body.
Block* outlined_func_block = outlined_func.addEntryBlock(); Block* outlined_func_block = outlined_func.addEntryBlock();

2
tensorflow/compiler/mlir/tensorflow/transforms/cluster_tf_ops_pass.cc
View File

@ -207,7 +207,7 @@ void CreateFunctions(ModuleOp module_op,
StringAttr::get(metadata.result_devices[i], context)); StringAttr::get(metadata.result_devices[i], context));
} }
func_op.setVisibility(FuncOp::Visibility::Public); func_op.setPublic();
Block *block = func_op.addEntryBlock(); Block *block = func_op.addEntryBlock();
// Clones and moves the operations into the function's body. And the cloned // Clones and moves the operations into the function's body. And the cloned

2
tensorflow/compiler/mlir/tensorflow/transforms/executor_tpuv1_outline_tpu_island.cc
View File

@ -108,7 +108,7 @@ void TPUBridgeExecutorIslandOutlining::runOnOperation() {
auto outlined_func = auto outlined_func =
OpBuilder(ctx).create<FuncOp>(island_op.getLoc(), name, func_type); OpBuilder(ctx).create<FuncOp>(island_op.getLoc(), name, func_type);
outlined_symbol_table.insert(outlined_func); outlined_symbol_table.insert(outlined_func);
outlined_func.setVisibility(FuncOp::Visibility::Nested); outlined_func.setNested();
// We will "steal" the body of the island and replace it with a call to the // We will "steal" the body of the island and replace it with a call to the
// new function later. // new function later.

2
tensorflow/compiler/mlir/tensorflow/transforms/guarantee_all_funcs_one_use.cc
View File

@ -93,7 +93,7 @@ class GuaranteeAllFuncsOneUse
} }
auto new_func = func.clone(); auto new_func = func.clone();
symbol_table.insert(new_func); symbol_table.insert(new_func);
new_func.setVisibility(SymbolTable::Visibility::Private); new_func.setPrivate();
if (failed(symbol_table.replaceAllSymbolUses(func, new_func.getName(), if (failed(symbol_table.replaceAllSymbolUses(func, new_func.getName(),
use.getUser()))) { use.getUser()))) {
return func.emitError() << "could not replace symbol use"; return func.emitError() << "could not replace symbol use";

2
tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc
View File

@ -150,7 +150,7 @@ void ExtractSingleBlockRegion(Region& region, StringRef name,
builder.create<ReturnOp>(terminator->getLoc(), return_values); builder.create<ReturnOp>(terminator->getLoc(), return_values);
terminator->erase(); terminator->erase();
outlined_func.setVisibility(FuncOp::Visibility::Private); outlined_func.setPrivate();
// Add the outlined function to the worklist in case its body has // Add the outlined function to the worklist in case its body has
// IfRegion or WhileRegion ops that need to converted. // IfRegion or WhileRegion ops that need to converted.

2
tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc
View File

@ -1111,7 +1111,7 @@ LogicalResult HandlePartitionedCallOpCallee(
name_base += "_resource_lifted"; name_base += "_resource_lifted";
auto name = name_base; auto name = name_base;
callee = callee.clone(); callee = callee.clone();
callee.setVisibility(SymbolTable::Visibility::Private); callee.setPrivate();
callee.setName(name); callee.setName(name);
SymbolTable(module).insert(callee); SymbolTable(module).insert(callee);
result->lifted_callee = callee; result->lifted_callee = callee;

2
tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting_cleanup.cc
View File

@ -122,7 +122,7 @@ FuncOp CloneFunctionIfNeeded(FuncOp func) {
if (func_uses.hasValue() && llvm::hasSingleElement(func_uses.getValue())) if (func_uses.hasValue() && llvm::hasSingleElement(func_uses.getValue()))
return func; return func;
FuncOp cloned = func.clone(); FuncOp cloned = func.clone();
cloned.setVisibility(SymbolTable::Visibility::Private); cloned.setPrivate();
cloned.setName(func.getName().str() + "_lifted"); cloned.setName(func.getName().str() + "_lifted");
SymbolTable(module).insert(cloned); SymbolTable(module).insert(cloned);
return cloned; return cloned;

2
tensorflow/compiler/mlir/tensorflow/transforms/stack_ops_decomposition.cc
View File

@ -337,7 +337,7 @@ LogicalResult HandlePartitionedCallOp(
if (!callee.isPrivate()) { if (!callee.isPrivate()) {
// Clone non-private callee in case of signature change. // Clone non-private callee in case of signature change.
lowered_callee = callee.clone(); lowered_callee = callee.clone();
lowered_callee.setVisibility(SymbolTable::Visibility::Private); lowered_callee.setPrivate();
} }
auto find_arg_stack_type = [&](int64_t index) -> llvm::Optional<Type> { auto find_arg_stack_type = [&](int64_t index) -> llvm::Optional<Type> {
auto it = data_var_to_size_var.find(call.getOperand(index)); auto it = data_var_to_size_var.find(call.getOperand(index));

2
tensorflow/compiler/mlir/tensorflow/transforms/tensor_array_ops_decomposition.cc
View File

@ -772,7 +772,7 @@ LogicalResult HandlePartitionedCallOp(
if (!callee.isPrivate()) { if (!callee.isPrivate()) {
// Clone non-private callee in case of signature change. // Clone non-private callee in case of signature change.
lowered_callee = callee.clone(); lowered_callee = callee.clone();
lowered_callee.setVisibility(SymbolTable::Visibility::Private); lowered_callee.setPrivate();
} }
auto grads = AccessedGradients({lowered_callee}, module); auto grads = AccessedGradients({lowered_callee}, module);
for (int64_t i = 0; i < lowered_callee.getNumArguments(); ++i) { for (int64_t i = 0; i < lowered_callee.getNumArguments(); ++i) {

2
tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc
View File

@ -481,7 +481,7 @@ LogicalResult HandlePartitionedCallOp(
if (!callee.isPrivate()) { if (!callee.isPrivate()) {
// Clone non-private callee in case of signature change. // Clone non-private callee in case of signature change.
lowered_callee = callee.clone(); lowered_callee = callee.clone();
lowered_callee.setVisibility(SymbolTable::Visibility::Private); lowered_callee.setPrivate();
} }
auto find_arg_buffer_type = [&](int64_t index) -> llvm::Optional<Type> { auto find_arg_buffer_type = [&](int64_t index) -> llvm::Optional<Type> {
auto it = buffer_to_size->find(call.getOperand(index)); auto it = buffer_to_size->find(call.getOperand(index));

4
tensorflow/compiler/mlir/tensorflow/transforms/tpu_rewrite_pass.cc
View File

@ -147,9 +147,9 @@ LogicalResult EncapsulateFuncAndSerialize(FuncOp entry_func,
// We can simply change name of TPU program's main function because there // We can simply change name of TPU program's main function because there
// should be no other reference to it. // should be no other reference to it.
clone.setName("main"); clone.setName("main");
clone.setVisibility(FuncOp::Visibility::Public); clone.setPublic();
} else { } else {
clone.setVisibility(FuncOp::Visibility::Private); clone.setPrivate();
} }
symbol_table.insert(clone); symbol_table.insert(clone);
} }

2
tensorflow/compiler/mlir/xla/transforms/legalize_tf_communication.cc
View File

@ -201,7 +201,7 @@ LogicalResult GetFunctionsToRewrite(
if (func.getSecond().original.isPublic() && if (func.getSecond().original.isPublic() &&
!func.getSecond().original.symbolKnownUseEmpty(module)) { !func.getSecond().original.symbolKnownUseEmpty(module)) {
auto clone = func.getSecond().original.clone(); auto clone = func.getSecond().original.clone();
clone.setVisibility(SymbolTable::Visibility::Private); clone.setPrivate();
symbol_table.insert(clone); symbol_table.insert(clone);
func.getSecond().clone = clone; func.getSecond().clone = clone;
} }