[MLIR] Adopt Symbol setPublic/setPrivate/setNested functions.
PiperOrigin-RevId: 342070484 Change-Id: I74b0592f28a7e4ce85b6ae67b986cb94a06961fb
This commit is contained in:
parent
68f4b84637
commit
3ea78d0fdd
@ -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;
|
||||
|
@ -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());
|
||||
|
@ -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;
|
||||
};
|
||||
|
||||
|
@ -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();
|
||||
|
@ -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
|
||||
|
@ -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.
|
||||
|
@ -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";
|
||||
|
@ -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.
|
||||
|
@ -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;
|
||||
|
@ -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;
|
||||
|
@ -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));
|
||||
|
@ -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) {
|
||||
|
@ -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));
|
||||
|
@ -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);
|
||||
}
|
||||
|
@ -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;
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user