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

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

View File

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

View File

@ -62,7 +62,7 @@ void RunOnWhile(TF::WhileOp while_op) {
auto call = builder.create<CallOp>(while_op.getLoc(), func, new_operands);
builder.create<YieldOp>(while_op.getLoc(), call.getResults());
// 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.body_function(), new_op.body());

View File

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

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
// symbol-dce pass to remove it when it is not referenced anymore.
outlined_func.setVisibility(FuncOp::Visibility::Private);
outlined_func.setPrivate();
// Create function body.
Block* outlined_func_block = outlined_func.addEntryBlock();

View File

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

View File

@ -108,7 +108,7 @@ void TPUBridgeExecutorIslandOutlining::runOnOperation() {
auto outlined_func =
OpBuilder(ctx).create<FuncOp>(island_op.getLoc(), name, func_type);
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
// new function later.

View File

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

View File

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

View File

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

View File

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

View File

@ -337,7 +337,7 @@ LogicalResult HandlePartitionedCallOp(
if (!callee.isPrivate()) {
// Clone non-private callee in case of signature change.
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 it = data_var_to_size_var.find(call.getOperand(index));

View File

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

View File

@ -481,7 +481,7 @@ LogicalResult HandlePartitionedCallOp(
if (!callee.isPrivate()) {
// Clone non-private callee in case of signature change.
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 it = buffer_to_size->find(call.getOperand(index));

View File

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

View File

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