[mlir][SCCP] Add support for propagating across symbol based calls

[rriddle]

Summary

This revision adds support for propagating constants across symbol-based callgraph edges. It uses the existing Call/CallableOpInterfaces to detect the dataflow edges, and propagates constants through arguments and out of returns.

Depends On D78522

[bondhugula]

Looks great!

It assumes that the arguments of callable can be tracked so long as we know that all of the uses are calls.

When can the arguments not be tracked? And shouldn't the CallableOpInterface guarantee that?

mlir/lib/Transforms/SCCP.cpp

446

Nit: `/*...=*/!op->getBlock()

[rriddle]

In D78592#1996224, @bondhugula wrote:

Looks great!

It assumes that the arguments of callable can be tracked so long as we know that all of the uses are calls.

When can the arguments not be tracked? And shouldn't the CallableOpInterface guarantee that?

I thought I had a reason originally, but after thinking about it the constraints of symbol visibility already guarantee this today.

[rriddle]

Resolve comments

[bondhugula]

mlir/lib/IR/SymbolTable.cpp

230

Nit: nested -> nestedOp

233

was a symbol table -> had a symbol table trait

mlir/lib/Transforms/SCCP.cpp

774

Callable things can never have results, right? Can't you assert for the trait instead? Also, this should be moved to before the check right above.

[bondhugula]

mlir/lib/Transforms/SCCP.cpp

601

/*owner=*/ /*to=*/ /*from=*/

776

region-like -> return-like

777

return values -> call site results ?

797

Likewise.

[bondhugula]

adds some initial support for propagating constants across callgraph edges

Could you say a line or two on what 'initial' is or what other things you want to add in the near term?

[rriddle]

Resolve comments

[rriddle]

mlir/lib/Transforms/SCCP.cpp

774

non-symbol defining callables return the callable as a result value. This check is intended to intended to guard against the missing support for those callables.

[rriddle]

Anything else here Uday?

[bondhugula]

In D78592#1996277, @rriddle wrote:

In D78592#1996224, @bondhugula wrote:

Looks great!

It assumes that the arguments of callable can be tracked so long as we know that all of the uses are calls.

When can the arguments not be tracked? And shouldn't the CallableOpInterface guarantee that?

I thought I had a reason originally, but after thinking about it the constraints of symbol visibility already guarantee this today.

It'd be good to incorporate this understanding somewhere as a comment.

In D78592#2002333, @rriddle wrote:

Anything else here Uday?

Sorry about the delay - I should be able to get back on this by later today.

[bondhugula]

Looks great overall. A bunch of mostly minor comments on tweaking doc/comments. Feel free to ignore the minor ones if not necessary. Many of them are also due to the order in which I looked at the changes and "local" reading.

mlir/include/mlir/IR/SymbolTable.h

92

Rephrase to clarify which uses are being referred to.

93

"visible" <- clarify as to visible to whom? (the walk?)

all uses -> all their uses?

mlir/lib/Transforms/SCCP.cpp

120

Nit: "a callable operation".

124

of results. -> of results initialized to the default lattice value (Unknown).

139

Are there examples of callables that don't define symbols in the codebase?

143

calls referenced by callables? Did you mean "calls referencing this callable"? Also, shouldn't this just return empty if the callable doesn't define a symbol?

151

Nit: results -> return values?

152

don't have proper SSA values -> aren't SSA values? (since the parent op is declarative)?

179

Nit: symbol callables -> symbol defining callables

386

Did you just mean "propagate within" or "propagate within and out of"?

395

Nit: symbol callables -> symbol defining callables (I know that 'symbol' has now almost been established as an adjective for 'callable' in most of the code, but this may still be useful for new reviewers/readers.)

411

Do you have a test case for non-call references?

412

which functions? You mean 'callables'?

413

Nit: all of the calls -> all of its calls

566

The second statement "This is simply detecting ..." is not really clear. What's the connection b/w non-symbol callables and non-call uses?

575

I think the CallOpInterface / getCallableForCallee comment should have a line explaining why the callee (CallInterfaceCallable) can be an SSA value.

581

Will the default constructed value from 'lookup' above be nullptr?

777

But this terminator might not be the only terminator in the region of the callable op. Why mark call site results overdefined? You'd just propagate to successors the same way as for non-callable op region terminators that aren't return-like. ? You do have a test case ("complex_inner_if") where you have a cond_br in the callable. I was expecting that to not work as a result of this.

mlir/test/Transforms/sccp-callgraph.mlir

69

Prefix repetition.

107

But this function a 'call user'. The non-call user is further below.

184

Nit: Move this below the next function - closer to its use.

233

This is a nice test case.

[bondhugula]

mlir/lib/Transforms/SCCP.cpp

411

I found this one among the test cases.

[rriddle]

Resolve comments

mlir/lib/Transforms/SCCP.cpp

581

Yes.

777

This branch is only taken when there are no CFG successors. This is the same behavior as the region case, i.e. if we see an "exiting" terminator that isn't return-like we assume all exit values are overdefined.

mlir/test/Transforms/sccp-callgraph.mlir

107

The function name is more of a test case name.

[bondhugula]

mlir/lib/Transforms/SCCP.cpp

777

I now see that far above at line 714 - the conditional over the call itself. Locally, this would be misleading including the method name visitCallableTerminatorOperation since these are for term ops that have 0 successors. Could you either update the doc comment or have a comment here to remind the reader? Adding ZeroSucc to the method name will be too long.

[rriddle]

Resolve comments

mlir/include/mlir/IR/SymbolTable.h

public:
   static Visibility getSymbolVisibility(Operation *symbol);
   /// Sets the visibility of the given symbol operation.
   static void setSymbolVisibility(Operation *symbol, Visibility vis);
 
   /// Returns the nearest symbol table from a given operation `from`. Returns
   /// nullptr if no valid parent symbol table could be found.
   static Operation *getNearestSymbolTable(Operation *from);
 
+  /// Walks all symbol table operations nested within, and including, `op`. For
+  /// each symbol table operation, the provided callback is invoked with the op
+  /// and a boolean signifying if the symbols within that symbol table can be
+  /// treated as if all uses within the IR are visible to the caller.
+  /// `allSymUsesVisible` identifies whether all of the symbol uses of symbols
+  /// within `op` are visible.
+  static void walkSymbolTables(Operation *op, bool allSymUsesVisible,
+                               function_ref<void(Operation *, bool)> callback);
+
   /// Returns the operation registered with the given symbol name with the
   /// regions of 'symbolTableOp'. 'symbolTableOp' is required to be an operation
   /// with the 'OpTrait::SymbolTable' trait.
   static Operation *lookupSymbolIn(Operation *op, StringRef symbol);
   static Operation *lookupSymbolIn(Operation *op, SymbolRefAttr symbol);
   /// A variant of 'lookupSymbolIn' that returns all of the symbols referenced
   /// by a given SymbolRefAttr. Returns failure if any of the nested references
   /// could not be resolved.

mlir/include/mlir/Interfaces/CallInterfaces.td

let description = [{
     indirect calls to other operations `call_indirect %foo`. An operation that
     uses this interface, must *not* also provide the `CallableOpInterface`.
   }];
 
   let methods = [
     InterfaceMethod<[{
         Returns the callee of this call-like operation. A `callee` is either a
         reference to a symbol, via SymbolRefAttr, or a reference to a defined
-        SSA value.
+        SSA value. If the reference is an SSA value, the SSA value corresponds
+        to a region of a lambda-like operation.
       }],
       "CallInterfaceCallable", "getCallableForCallee"
     >,
     InterfaceMethod<[{
         Returns the operands within this call that are used as arguments to the
         callee.
       }],
       "Operation::operand_range", "getArgOperands"

mlir/lib/IR/SymbolTable.cpp

while (!from->hasTrait<OpTrait::SymbolTable>()) {
 
     // Check that this is a valid op and isn't an unknown symbol table.
     if (!from || isPotentiallyUnknownSymbolTable(from))
       return nullptr;
   }
   return from;
 }
 
+/// Walks all symbol table operations nested within, and including, `op`. For
+/// each symbol table operation, the provided callback is invoked with the op
+/// and a boolean signifying if the symbols within that symbol table can be
+/// treated as if all uses are visible. `allSymUsesVisible` identifies whether
+/// all of the symbol uses of symbols within `op` are visible.
+void SymbolTable::walkSymbolTables(
+    Operation *op, bool allSymUsesVisible,
+    function_ref<void(Operation *, bool)> callback) {
+  bool isSymbolTable = op->hasTrait<OpTrait::SymbolTable>();
+  if (isSymbolTable) {
+    SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
+    allSymUsesVisible |= !symbol || symbol.isPrivate();
+  } else {
+    // Otherwise if 'op' is not a symbol table, any nested symbols are
+    // guaranteed to be hidden.
+    allSymUsesVisible = true;
+  }
+
+  for (Region &region : op->getRegions())
+    for (Block &block : region)
+      for (Operation &nestedOp : block)
+        walkSymbolTables(&nestedOp, allSymUsesVisible, callback);
+
+  // If 'op' had the symbol table trait, visit it after any nested symbol
+  // tables.
+  if (isSymbolTable)
+    callback(op, allSymUsesVisible);
+}
+
 /// Returns the operation registered with the given symbol name with the
 /// regions of 'symbolTableOp'. 'symbolTableOp' is required to be an operation
 /// with the 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol
 /// was found.
 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
                                        StringRef symbol) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
 

mlir/lib/Transforms/Inliner.cpp

 #define DEBUG_TYPE "inlining"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Tracking
 //===----------------------------------------------------------------------===//
 
-/// Walk all of the symbol table operations nested with 'op' along with a
-/// boolean signifying if the symbols within can be treated as if all uses are
-/// visible. The provided callback is invoked with the symbol table operation,
-/// and a boolean signaling if all of the uses within the symbol table are
-/// visible.
-static void walkSymbolTables(Operation *op, bool allSymUsesVisible,
-                             function_ref<void(Operation *, bool)> callback) {
-  if (op->hasTrait<OpTrait::SymbolTable>()) {
-    SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
-    allSymUsesVisible = allSymUsesVisible || !symbol || symbol.isPrivate();
-    callback(op, allSymUsesVisible);
-  } else {
-    // Otherwise if 'op' is not a symbol table, any nested symbols are
-    // guaranteed to be hidden.
-    allSymUsesVisible = true;
-  }
-
-  for (Region &region : op->getRegions())
-    for (Block &block : region)
-      for (Operation &nested : block)
-        walkSymbolTables(&nested, allSymUsesVisible, callback);
-}
-
 /// Walk all of the used symbol callgraph nodes referenced with the given op.
 static void walkReferencedSymbolNodes(
     Operation *op, CallGraph &cg,
     DenseMap<Attribute, CallGraphNode *> &resolvedRefs,
     function_ref<void(CallGraphNode *, Operation *)> callback) {
   auto symbolUses = SymbolTable::getSymbolUses(op);
   assert(symbolUses && "expected uses to be valid");
 
for (Block &block : symbolTableOp->getRegion(0)) {
           }
         }
         // Otherwise, check for any referenced nodes. These will be always-live.
         walkReferencedSymbolNodes(&op, cg, alwaysLiveNodes,
                                   [](CallGraphNode *, Operation *) {});
       }
     }
   };
-  walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(), walkFn);
+  SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
+                                walkFn);
 
   // Drop the use information for any discardable nodes that are always live.
   for (auto &it : alwaysLiveNodes)
     discardableSymNodeUses.erase(it.second);
 
   // Compute the uses for each of the callable nodes in the graph.
   for (CallGraphNode *node : cg)
     recomputeUses(node, cg);

mlir/lib/Transforms/SCCP.cpp

private:
   llvm::PointerIntPair<Attribute, 2, Kind> constantAndTag;
 
   /// The dialect the constant originated from. This is only valid if the
   /// lattice is a constant. This is not used as part of the key, and is only
   /// needed to materialize the held constant if necessary.
   Dialect *constantDialect;
 };
 
+/// This class contains various state used when computing the lattice of a
+/// callable operation.
+class CallableLatticeState {
+public:
+  /// Build a lattice state with a given callable region, and a specified number
+  /// of results to be initialized to the default lattice value (Unknown).
+  CallableLatticeState(Region *callableRegion, unsigned numResults)
+      : callableArguments(callableRegion->front().getArguments()),
+        resultLatticeValues(numResults) {}
+
+  /// Returns the arguments to the callable region.
+  Block::BlockArgListType getCallableArguments() const {
+    return callableArguments;
+  }
+
+  /// Returns the lattice value for the results of the callable region.
+  MutableArrayRef<LatticeValue> getResultLatticeValues() {
+    return resultLatticeValues;
+  }
+
+  /// Add a call to this callable. This is only used if the callable defines a
+  /// symbol.
+  void addSymbolCall(Operation *op) { symbolCalls.push_back(op); }
+
+  /// Return the calls that reference this callable. This is only used
+  /// if the callable defines a symbol.
+  ArrayRef<Operation *> getSymbolCalls() const { return symbolCalls; }
+
+private:
+  /// The arguments of the callable region.
+  Block::BlockArgListType callableArguments;
+
+  /// The lattice state for each of the results of this region. The return
+  /// values of the callable aren't SSA values, so we need to track them
+  /// separately.
+  SmallVector<LatticeValue, 4> resultLatticeValues;
+
+  /// The calls referencing this callable if this callable defines a symbol.
+  /// This removes the need to recompute symbol references during propagation.
+  /// Value based references are trivial to resolve, so they can be done
+  /// in-place.
+  SmallVector<Operation *, 4> symbolCalls;
+};
+
 /// This class represents the solver for the SCCP analysis. This class acts as
 /// the propagation engine for computing which values form constants.
 class SCCPSolver {
 public:
-  /// Initialize the solver with a given set of regions.
-  SCCPSolver(MutableArrayRef<Region> regions);
+  /// Initialize the solver with the given top-level operation.
+  SCCPSolver(Operation *op);
 
   /// Run the solver until it converges.
   void solve();
 
   /// Rewrite the given regions using the computing analysis. This replaces the
   /// uses of all values that have been computed to be constant, and erases as
   /// many newly dead operations.
   void rewrite(MLIRContext *context, MutableArrayRef<Region> regions);
 
 private:
+  /// Initialize the set of symbol defining callables that can have their
+  /// arguments and results tracked. 'op' is the top-level operation that SCCP
+  /// is operating on.
+  void initializeSymbolCallables(Operation *op);
+
   /// Replace the given value with a constant if the corresponding lattice
   /// represents a constant. Returns success if the value was replaced, failure
   /// otherwise.
   LogicalResult replaceWithConstant(OpBuilder &builder, OperationFolder &folder,
                                     Value value);
 
   /// Visit the users of the given IR that reside within executable blocks.
   template <typename T>
   void visitUsers(T &value) {
     for (Operation *user : value.getUsers())
       if (isBlockExecutable(user->getBlock()))
         visitOperation(user);
   }
 
   /// Visit the given operation and compute any necessary lattice state.
   void visitOperation(Operation *op);
 
+  /// Visit the given call operation and compute any necessary lattice state.
+  void visitCallOperation(CallOpInterface op);
+
+  /// Visit the given callable operation and compute any necessary lattice
+  /// state.
+  void visitCallableOperation(Operation *op);
+
   /// Visit the given operation, which defines regions, and compute any
   /// necessary lattice state. This also resolves the lattice state of both the
   /// operation results and any nested regions.
   void visitRegionOperation(Operation *op,
                             ArrayRef<Attribute> constantOperands);
 
   /// Visit the given set of region successors, computing any necessary lattice
   /// state. The provided function returns the input operands to the region at
   /// the given index. If the index is 'None', the input operands correspond to
   /// the parent operation results.
   void visitRegionSuccessors(
       Operation *parentOp, ArrayRef<RegionSuccessor> regionSuccessors,
       function_ref<OperandRange(Optional<unsigned>)> getInputsForRegion);
 
   /// Visit the given terminator operation and compute any necessary lattice
   /// state.
   void visitTerminatorOperation(Operation *op,
                                 ArrayRef<Attribute> constantOperands);
 
+  /// Visit the given terminator operation that exits a callable region. These
+  /// are terminators with no CFG successors.
+  void visitCallableTerminatorOperation(Operation *callable,
+                                        Operation *terminator);
+
   /// Visit the given block and compute any necessary lattice state.
   void visitBlock(Block *block);
 
   /// Visit argument #'i' of the given block and compute any necessary lattice
   /// state.
   void visitBlockArgument(Block *block, int i);
 
   /// Mark the given block as executable. Returns false if the block was already
private:
   /// The set of control flow edges that are known to execute.
   DenseSet<std::pair<Block *, Block *>> executableEdges;
 
   /// A worklist containing blocks that need to be processed.
   SmallVector<Block *, 64> blockWorklist;
 
   /// A worklist of operations that need to be processed.
   SmallVector<Operation *, 64> opWorklist;
+
+  /// The callable operations that have their argument/result state tracked.
+  DenseMap<Operation *, CallableLatticeState> callableLatticeState;
+
+  /// A map between a call operation and the resolved symbol callable. This
+  /// avoids re-resolving symbol references during propagation. Value based
+  /// callables are trivial to resolve, so they can be done in-place.
+  DenseMap<Operation *, Operation *> callToSymbolCallable;
 };
 } // end anonymous namespace
 
-SCCPSolver::SCCPSolver(MutableArrayRef<Region> regions) {
-  for (Region &region : regions) {
+SCCPSolver::SCCPSolver(Operation *op) {
+  /// Initialize the solver with the regions within this operation.
+  for (Region &region : op->getRegions()) {
     if (region.empty())
       continue;
     Block *entryBlock = &region.front();
 
     // Mark the entry block as executable.
     markBlockExecutable(entryBlock);
 
     // The values passed to these regions are invisible, so mark any arguments
     // as overdefined.
     markAllOverdefined(entryBlock->getArguments());
   }
+  initializeSymbolCallables(op);
 }
 
 void SCCPSolver::solve() {
   while (!blockWorklist.empty() || !opWorklist.empty()) {
     // Process any operations in the op worklist.
     while (!opWorklist.empty())
       visitUsers(*opWorklist.pop_back_val());
 
for (Operation &op : llvm::make_early_inc_range(*block)) {
       }
 
       // Add any the regions of this operation to the worklist.
       addToWorklist(op.getRegions());
     }
   }
 }
 
+void SCCPSolver::initializeSymbolCallables(Operation *op) {
+  // Initialize the set of symbol callables that can have their state tracked.
+  // This tracks which symbol callable operations we can propagate within and
+  // out of.
+  auto walkFn = [&](Operation *symTable, bool allUsesVisible) {
+    Region &symbolTableRegion = symTable->getRegion(0);
+    Block *symbolTableBlock = &symbolTableRegion.front();
+    for (auto callable : symbolTableBlock->getOps<CallableOpInterface>()) {
+      // We won't be able to track external callables.
+      Region *callableRegion = callable.getCallableRegion();
+      if (!callableRegion)
+        continue;
+      // We only care about symbol defining callables here.
+      auto symbol = dyn_cast<SymbolOpInterface>(callable.getOperation());
+      if (!symbol)
+        continue;
+      callableLatticeState.try_emplace(callable, callableRegion,
+                                       callable.getCallableResults().size());
+
+      // If not all of the uses of this symbol are visible, we can't track the
+      // state of the arguments.
+      if (symbol.isPublic() || (!allUsesVisible && symbol.isNested()))
+        markAllOverdefined(callableRegion->front().getArguments());
+    }
+    if (callableLatticeState.empty())
+      return;
+
+    // After computing the valid callables, walk any symbol uses to check
+    // for non-call references. We won't be able to track the lattice state
+    // for arguments to these callables, as we can't guarantee that we can see
+    // all of its calls.
+    Optional<SymbolTable::UseRange> uses =
+        SymbolTable::getSymbolUses(&symbolTableRegion);
+    if (!uses) {
+      // If we couldn't gather the symbol uses, conservatively assume that
+      // we can't track information for any nested symbols.
+      op->walk([&](CallableOpInterface op) { callableLatticeState.erase(op); });
+      return;
+    }
+
+    for (const SymbolTable::SymbolUse &use : *uses) {
+      // If the use is a call, track it to avoid the need to recompute the
+      // reference later.
+      if (auto callOp = dyn_cast<CallOpInterface>(use.getUser())) {
+        Operation *symCallable = callOp.resolveCallable();
+        auto callableLatticeIt = callableLatticeState.find(symCallable);
+        if (callableLatticeIt != callableLatticeState.end()) {
+          callToSymbolCallable.try_emplace(callOp, symCallable);
+
+          // We only need to record the call in the lattice if it produces any
+          // values.
+          if (callOp.getOperation()->getNumResults())
+            callableLatticeIt->second.addSymbolCall(callOp);
+        }
+        continue;
+      }
+      // This use isn't a call, so don't we know all of the callers.
+      auto *symbol = SymbolTable::lookupSymbolIn(op, use.getSymbolRef());
+      auto it = callableLatticeState.find(symbol);
+      if (it != callableLatticeState.end())
+        markAllOverdefined(it->second.getCallableArguments());
+    }
+  };
+  SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
+                                walkFn);
+}
+
 LogicalResult SCCPSolver::replaceWithConstant(OpBuilder &builder,
                                               OperationFolder &folder,
                                               Value value) {
   auto it = latticeValues.find(value);
   auto attr = it == latticeValues.end() ? nullptr : it->second.getConstant();
   if (!attr)
     return failure();
 
if (operandLattice.isUnknown())
       return;
     operandConstants.push_back(operandLattice.getConstant());
   }
 
   // If this is a terminator operation, process any control flow lattice state.
   if (op->isKnownTerminator())
     visitTerminatorOperation(op, operandConstants);
 
+  // Process call operations. The call visitor processes result values, so we
+  // can exit afterwards.
+  if (CallOpInterface call = dyn_cast<CallOpInterface>(op))
+    return visitCallOperation(call);
+
+  // Process callable operations. These are specially handled region operations
+  // that track dataflow via calls.
+  if (isa<CallableOpInterface>(op))
+    return visitCallableOperation(op);
+
   // Process region holding operations. The region visitor processes result
   // values, so we can exit afterwards.
   if (op->getNumRegions())
     return visitRegionOperation(op, operandConstants);
 
   // If this op produces no results, it can't produce any constants.
   if (op->getNumResults() == 0)
     return;
for (unsigned i = 0, e = foldResults.size(); i != e; ++i) {
     OpFoldResult foldResult = foldResults[i];
     if (Attribute foldAttr = foldResult.dyn_cast<Attribute>())
       meet(op, resultLattice, LatticeValue(foldAttr, opDialect));
     else
       meet(op, resultLattice, latticeValues[foldResult.get<Value>()]);
   }
 }
 
+void SCCPSolver::visitCallableOperation(Operation *op) {
+  // Mark the regions as executable.
+  bool isTrackingLatticeState = callableLatticeState.count(op);
+  for (Region &region : op->getRegions()) {
+    if (region.empty())
+      continue;
+    Block *entryBlock = &region.front();
+    markBlockExecutable(entryBlock);
+
+    // If we aren't tracking lattice state for this callable, mark all of the
+    // region arguments as overdefined.
+    if (!isTrackingLatticeState)
+      markAllOverdefined(entryBlock->getArguments());
+  }
+
+  // TODO: Add support for non-symbol callables when necessary. If the callable
+  // has non-call uses we would mark overdefined, otherwise allow for
+  // propagating the return values out.
+  markAllOverdefined(op, op->getResults());
+}
+
+void SCCPSolver::visitCallOperation(CallOpInterface op) {
+  ResultRange callResults = op.getOperation()->getResults();
+
+  // Resolve the callable operation for this call.
+  Operation *callableOp = nullptr;
+  if (Value callableValue = op.getCallableForCallee().dyn_cast<Value>())
+    callableOp = callableValue.getDefiningOp();
+  else
+    callableOp = callToSymbolCallable.lookup(op);
+
+  // The callable of this call can't be resolved, mark any results overdefined.
+  if (!callableOp)
+    return markAllOverdefined(op, callResults);
+
+  // If this callable is tracking state, merge the argument operands with the
+  // arguments of the callable.
+  auto callableLatticeIt = callableLatticeState.find(callableOp);
+  if (callableLatticeIt == callableLatticeState.end())
+    return markAllOverdefined(op, callResults);
+
+  OperandRange callOperands = op.getArgOperands();
+  auto callableArgs = callableLatticeIt->second.getCallableArguments();
+  for (auto it : llvm::zip(callOperands, callableArgs)) {
+    BlockArgument callableArg = std::get<1>(it);
+    if (latticeValues[callableArg].meet(latticeValues[std::get<0>(it)]))
+      visitUsers(callableArg);
+  }
+
+  // Merge in the lattice state for the callable results as well.
+  auto callableResults = callableLatticeIt->second.getResultLatticeValues();
+  for (auto it : llvm::zip(callResults, callableResults))
+    meet(/*owner=*/op, /*to=*/latticeValues[std::get<0>(it)],
+         /*from=*/std::get<1>(it));
+}
+
 void SCCPSolver::visitRegionOperation(Operation *op,
                                       ArrayRef<Attribute> constantOperands) {
   // Check to see if we can reason about the internal control flow of this
   // region operation.
   auto regionInterface = dyn_cast<RegionBranchOpInterface>(op);
   if (!regionInterface) {
     // If we can't, conservatively mark all regions as executable.
     for (Region &region : op->getRegions()) {
for (const RegionSuccessor &it : regionSuccessors) {
     }
   }
 }
 
 void SCCPSolver::visitTerminatorOperation(
     Operation *op, ArrayRef<Attribute> constantOperands) {
   // If this operation has no successors, we treat it as an exiting terminator.
   if (op->getNumSuccessors() == 0) {
-    // Check to see if the parent tracks region control flow.
     Region *parentRegion = op->getParentRegion();
     Operation *parentOp = parentRegion->getParentOp();
+
+    // Check to see if this is a terminator for a callable region.
+    if (isa<CallableOpInterface>(parentOp))
+      return visitCallableTerminatorOperation(parentOp, op);
+
+    // Otherwise, check to see if the parent tracks region control flow.
     auto regionInterface = dyn_cast<RegionBranchOpInterface>(parentOp);
     if (!regionInterface || !isBlockExecutable(parentOp->getBlock()))
       return;
 
     // Query the set of successors from the current region.
     SmallVector<RegionSuccessor, 1> regionSuccessors;
     regionInterface.getSuccessorRegions(parentRegion->getRegionNumber(),
                                         constantOperands, regionSuccessors);
void SCCPSolver::visitTerminatorOperation(
   }
 
   // Otherwise, conservatively treat all edges as executable.
   Block *block = op->getBlock();
   for (Block *succ : op->getSuccessors())
     markEdgeExecutable(block, succ);
 }
 
+void SCCPSolver::visitCallableTerminatorOperation(Operation *callable,
+                                                  Operation *terminator) {
+  // If there are no exiting values, we have nothing to track.
+  if (terminator->getNumOperands() == 0)
+    return;
+
+  // If this callable isn't tracking any lattice state there is nothing to do.
+  auto latticeIt = callableLatticeState.find(callable);
+  if (latticeIt == callableLatticeState.end())
+    return;
+  assert(callable->getNumResults() == 0 && "expected symbol callable");
+
+  // If this terminator is not "return-like", conservatively mark all of the
+  // call-site results as overdefined.
+  auto callableResultLattices = latticeIt->second.getResultLatticeValues();
+  if (!terminator->hasTrait<OpTrait::ReturnLike>()) {
+    for (auto &it : callableResultLattices)
+      it.markOverdefined();
+    for (Operation *call : latticeIt->second.getSymbolCalls())
+      markAllOverdefined(call, call->getResults());
+    return;
+  }
+
+  // Merge the terminator operands into the results.
+  bool anyChanged = false;
+  for (auto it : llvm::zip(terminator->getOperands(), callableResultLattices))
+    anyChanged |= std::get<1>(it).meet(latticeValues[std::get<0>(it)]);
+  if (!anyChanged)
+    return;
+
+  // If any of the result lattices changed, update the callers.
+  for (Operation *call : latticeIt->second.getSymbolCalls())
+    for (auto it : llvm::zip(call->getResults(), callableResultLattices))
+      meet(call, latticeValues[std::get<0>(it)], std::get<1>(it));
+}
+
 void SCCPSolver::visitBlock(Block *block) {
   // If the block is not the entry block we need to compute the lattice state
   // for the block arguments. Entry block argument lattices are computed
   // elsewhere, such as when visiting the parent operation.
   if (!block->isEntryBlock()) {
     for (int i : llvm::seq<int>(0, block->getNumArguments()))
       visitBlockArgument(block, i);
   }
struct SCCP : public SCCPBase<SCCP> {
   void runOnOperation() override;
 };
 } // end anonymous namespace
 
 void SCCP::runOnOperation() {
   Operation *op = getOperation();
 
   // Solve for SCCP constraints within nested regions.
-  SCCPSolver solver(op->getRegions());
+  SCCPSolver solver(op);
   solver.solve();
 
   // Cleanup any operations using the solver analysis.
   solver.rewrite(&getContext(), op->getRegions());
 }
 
 std::unique_ptr<Pass> mlir::createSCCPPass() {
   return std::make_unique<SCCP>();
 }

mlir/test/Transforms/sccp-callgraph.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -sccp -split-input-file | FileCheck %s -dump-input-on-failure
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="module(sccp)" -split-input-file | FileCheck %s --check-prefix=NESTED -dump-input-on-failure
+
+/// Check that a constant is properly propagated through the arguments and
+/// results of a private function.
+
+// CHECK-LABEL: func @private(
+func @private(%arg0 : i32) -> i32 attributes { sym_visibility = "private" } {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: return %[[CST]] : i32
+
+  return %arg0 : i32
+}
+
+// CHECK-LABEL: func @simple_private(
+func @simple_private() -> i32 {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: return %[[CST]] : i32
+
+  %1 = constant 1 : i32
+  %result = call @private(%1) : (i32) -> i32
+  return %result : i32
+}
+
+// -----
+
+/// Check that a constant is properly propagated through the arguments and
+/// results of a visible nested function.
+
+// CHECK: func @nested(
+func @nested(%arg0 : i32) -> i32 attributes { sym_visibility = "nested" } {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: return %[[CST]] : i32
+
+  return %arg0 : i32
+}
+
+// CHECK-LABEL: func @simple_nested(
+func @simple_nested() -> i32 {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: return %[[CST]] : i32
+
+  %1 = constant 1 : i32
+  %result = call @nested(%1) : (i32) -> i32
+  return %result : i32
+}
+
+// -----
+
+/// Check that non-visible nested functions do not track arguments.
+module {
+  // NESTED-LABEL: module @nested_module
+  module @nested_module attributes { sym_visibility = "public" } {
+
+    // NESTED: func @nested(
+    func @nested(%arg0 : i32) -> (i32, i32) attributes { sym_visibility = "nested" } {
+      // NESTED: %[[CST:.*]] = constant 1 : i32
+      // NESTED: return %[[CST]], %arg0 : i32, i32
+
+      %1 = constant 1 : i32
+      return %1, %arg0 : i32, i32
+    }
+
+    // NESTED: func @nested_not_all_uses_visible(
+    func @nested_not_all_uses_visible() -> (i32, i32) {
+      // NESTED: %[[CST:.*]] = constant 1 : i32
+      // NESTED: %[[CALL:.*]]:2 = call @nested
+      // NESTED: return %[[CST]], %[[CALL]]#1 : i32, i32
+
+      %1 = constant 1 : i32
+      %result:2 = call @nested(%1) : (i32) -> (i32, i32)
+      return %result#0, %result#1 : i32, i32
+    }
+  }
+}
+
+// -----
+
+/// Check that public functions do not track arguments.
+
+// CHECK-LABEL: func @public(
+func @public(%arg0 : i32) -> (i32, i32) attributes { sym_visibility = "public" } {
+  %1 = constant 1 : i32
+  return %1, %arg0 : i32, i32
+}
+
+// CHECK-LABEL: func @simple_public(
+func @simple_public() -> (i32, i32) {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: %[[CALL:.*]]:2 = call @public
+  // CHECK: return %[[CST]], %[[CALL]]#1 : i32, i32
+
+  %1 = constant 1 : i32
+  %result:2 = call @public(%1) : (i32) -> (i32, i32)
+  return %result#0, %result#1 : i32, i32
+}
+
+// -----
+
+/// Check that functions with non-call users don't have arguments tracked.
+
+func @callable(%arg0 : i32) -> (i32, i32) attributes { sym_visibility = "private" } {
+  %1 = constant 1 : i32
+  return %1, %arg0 : i32, i32
+}
+
+// CHECK-LABEL: func @non_call_users(
+func @non_call_users() -> (i32, i32) {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: %[[CALL:.*]]:2 = call @callable
+  // CHECK: return %[[CST]], %[[CALL]]#1 : i32, i32
+
+  %1 = constant 1 : i32
+  %result:2 = call @callable(%1) : (i32) -> (i32, i32)
+  return %result#0, %result#1 : i32, i32
+}
+
+"live.user"() {uses = [@callable]} : () -> ()
+
+// -----
+
+/// Check that return values are overdefined in the presence of an unknown terminator.
+
+func @callable(%arg0 : i32) -> i32 attributes { sym_visibility = "private" } {
+  "unknown.return"(%arg0) : (i32) -> ()
+}
+
+// CHECK-LABEL: func @unknown_terminator(
+func @unknown_terminator() -> i32 {
+  // CHECK: %[[CALL:.*]] = call @callable
+  // CHECK: return %[[CALL]] : i32
+
+  %1 = constant 1 : i32
+  %result = call @callable(%1) : (i32) -> i32
+  return %result : i32
+}
+
+// -----
+
+/// Check that return values are overdefined when the constant conflicts.
+
+func @callable(%arg0 : i32) -> i32 attributes { sym_visibility = "private" } {
+  "unknown.return"(%arg0) : (i32) -> ()
+}
+
+// CHECK-LABEL: func @conflicting_constant(
+func @conflicting_constant() -> (i32, i32) {
+  // CHECK: %[[CALL1:.*]] = call @callable
+  // CHECK: %[[CALL2:.*]] = call @callable
+  // CHECK: return %[[CALL1]], %[[CALL2]] : i32, i32
+
+  %1 = constant 1 : i32
+  %2 = constant 2 : i32
+  %result = call @callable(%1) : (i32) -> i32
+  %result2 = call @callable(%2) : (i32) -> i32
+  return %result, %result2 : i32, i32
+}
+
+// -----
+
+/// Check that return values are overdefined when the constant conflicts with a
+/// non-constant.
+
+func @callable(%arg0 : i32) -> i32 attributes { sym_visibility = "private" } {
+  "unknown.return"(%arg0) : (i32) -> ()
+}
+
+// CHECK-LABEL: func @conflicting_constant(
+func @conflicting_constant(%arg0 : i32) -> (i32, i32) {
+  // CHECK: %[[CALL1:.*]] = call @callable
+  // CHECK: %[[CALL2:.*]] = call @callable
+  // CHECK: return %[[CALL1]], %[[CALL2]] : i32, i32
+
+  %1 = constant 1 : i32
+  %result = call @callable(%1) : (i32) -> i32
+  %result2 = call @callable(%arg0) : (i32) -> i32
+  return %result, %result2 : i32, i32
+}
+
+// -----
+
+/// Check a more complex interaction with calls and control flow.
+
+// CHECK-LABEL: func @complex_inner_if(
+func @complex_inner_if(%arg0 : i32) -> i32 attributes { sym_visibility = "private" } {
+  // CHECK-DAG: %[[TRUE:.*]] = constant 1 : i1
+  // CHECK-DAG: %[[CST:.*]] = constant 1 : i32
+  // CHECK: cond_br %[[TRUE]], ^bb1
+
+  %cst_20 = constant 20 : i32
+  %cond = cmpi "ult", %arg0, %cst_20 : i32
+  cond_br %cond, ^bb1, ^bb2
+
+^bb1:
+  // CHECK: ^bb1:
+  // CHECK: return %[[CST]] : i32
+
+  %cst_1 = constant 1 : i32
+  return %cst_1 : i32
+
+^bb2:
+  %cst_1_2 = constant 1 : i32
+  %arg_inc = addi %arg0, %cst_1_2 : i32
+  return %arg_inc : i32
+}
+
+func @complex_cond() -> i1
+
+// CHECK-LABEL: func @complex_callee(
+func @complex_callee(%arg0 : i32) -> i32 attributes { sym_visibility = "private" } {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+
+  %loop_cond = call @complex_cond() : () -> i1
+  cond_br %loop_cond, ^bb1, ^bb2
+
+^bb1:
+  // CHECK: ^bb1:
+  // CHECK-NEXT: return %[[CST]] : i32
+  return %arg0 : i32
+
+^bb2:
+  // CHECK: ^bb2:
+  // CHECK: call @complex_inner_if(%[[CST]]) : (i32) -> i32
+  // CHECK: call @complex_callee(%[[CST]]) : (i32) -> i32
+  // CHECK: return %[[CST]] : i32
+
+  %updated_arg = call @complex_inner_if(%arg0) : (i32) -> i32
+  %res = call @complex_callee(%updated_arg) : (i32) -> i32
+  return %res : i32
+}
+
+// CHECK-LABEL: func @complex_caller(
+func @complex_caller(%arg0 : i32) -> i32 {
+  // CHECK: %[[CST:.*]] = constant 1 : i32
+  // CHECK: return %[[CST]] : i32
+
+  %1 = constant 1 : i32
+  %result = call @complex_callee(%1) : (i32) -> i32
+  return %result : i32
+}
+
+// -----
+
+/// Check that non-symbol defining callables currently go to overdefined.
+
+// CHECK-LABEL: func @non_symbol_defining_callable
+func @non_symbol_defining_callable() -> i32 {
+  // CHECK: %[[RES:.*]] = call_indirect
+  // CHECK: return %[[RES]] : i32
+
+  %fn = "test.functional_region_op"() ({
+    %1 = constant 1 : i32
+    "test.return"(%1) : (i32) -> ()
+  }) : () -> (() -> i32)
+  %res = call_indirect %fn() : () -> (i32)
+  return %res : i32
+}

mlir/test/lib/Dialect/Test/TestOps.td

def TestRecursiveRewriteOp : TEST_Op<"recursive_rewrite"> {
   let assemblyFormat = "$depth attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
 // Test Type Legalization
 //===----------------------------------------------------------------------===//
 
 def TestRegionBuilderOp : TEST_Op<"region_builder">;
-def TestReturnOp : TEST_Op<"return", [Terminator]>,
+def TestReturnOp : TEST_Op<"return", [ReturnLike, Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 def TestCastOp : TEST_Op<"cast">,
   Arguments<(ins Variadic<AnyType>)>, Results<(outs AnyType)>;
 def TestInvalidOp : TEST_Op<"invalid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 def TestTypeProducerOp : TEST_Op<"type_producer">,
   Results<(outs AnyType)>;
 def TestTypeConsumerOp : TEST_Op<"type_consumer">,