[mlir][SCCP] Don't visit private callables unless they are used when tracking interprocedural arguments/results

[rriddle]

Summary

This fixes a subtle bug where SCCP could incorrectly optimize a private callable while waiting for its arguments to be resolved.

Fixes PR#48457

mlir/lib/Transforms/SCCP.cpp

private:
 
   /// 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 entry block of the given region as executable. Returns false if
+  /// the block was already marked executable. If `markArgsOverdefined` is true,
+  /// the arguments of the entry block are also set to overdefined.
+  bool markEntryBlockExecutable(Region *region, bool markArgsOverdefined);
+
   /// Mark the given block as executable. Returns false if the block was already
   /// marked executable.
   bool markBlockExecutable(Block *block);
 
   /// Returns true if the given block is executable.
   bool isBlockExecutable(Block *block) const;
 
   /// Mark the edge between 'from' and 'to' as executable.
private:
   /// A symbol table used for O(1) symbol lookups during simplification.
   SymbolTableCollection symbolTable;
 };
 } // end anonymous namespace
 
 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());
+    // Mark the entry block as executable. The values passed to these regions
+    // are also invisible, so mark any arguments as overdefined.
+    markEntryBlockExecutable(&region, /*markArgsOverdefined=*/true);
   }
   initializeSymbolCallables(op);
 }
 
 void SCCPSolver::solve() {
   while (!blockWorklist.empty() || !opWorklist.empty()) {
     // Process any operations in the op worklist.
     while (!opWorklist.empty())
for (auto callable : symbolTableBlock->getOps<CallableOpInterface>()) {
       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->getArguments());
+      if (symbol.isPublic() || (!allUsesVisible && symbol.isNested())) {
+        for (Region &region : callable->getRegions())
+          markEntryBlockExecutable(&region, /*markArgsOverdefined=*/true);
+      }
     }
     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.
for (const SymbolTable::SymbolUse &use : *uses) {
           if (callOp->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());
+      if (it != callableLatticeState.end()) {
+        for (Region &region : it->first->getRegions())
+          markEntryBlockExecutable(&region, /*markArgsOverdefined=*/true);
+      }
     }
   };
   SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
                                 walkFn);
 }
 
 LogicalResult SCCPSolver::replaceWithConstant(OpBuilder &builder,
                                               OperationFolder &folder,
void SCCPSolver::visitOperation(Operation *op) {
 
   // 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))
+  if (isa<CallableOpInterface>(op)) {
+    // If this callable has a tracked lattice state, it will be visited by calls
+    // that reference it instead. This way, we don't assume that it is
+    // executable unless there is a proper reference to it.
+    if (callableLatticeState.count(op))
+      return;
     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)
for (unsigned i = 0, e = foldResults.size(); i != e; ++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.
+  // Mark the regions as executable. If we aren't tracking lattice state for
+  // this callable, mark all of the region arguments as overdefined.
   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());
-  }
+  for (Region &region : op->getRegions())
+    markEntryBlockExecutable(&region, !isTrackingLatticeState);
 
   // 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) {
void SCCPSolver::visitCallOperation(CallOpInterface op) {
   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);
   }
 
+  // Visit the callable.
+  visitCallableOperation(callableOp);
+
   // 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()) {
-      if (region.empty())
-        continue;
-      Block *entryBlock = &region.front();
-      markBlockExecutable(entryBlock);
-      markAllOverdefined(entryBlock->getArguments());
-    }
+    for (Region &region : op->getRegions())
+      markEntryBlockExecutable(&region, /*markArgsOverdefined=*/true);
 
     // Don't try to simulate the results of a region operation as we can't
     // guarantee that folding will be out-of-place. We don't allow in-place
     // folds as the desire here is for simulated execution, and not general
     // folding.
     return markAllOverdefined(op, op->getResults());
   }
 
if (argLattice.isOverdefined())
       break;
   }
 
   // If the lattice was updated, visit any executable users of the argument.
   if (updatedLattice)
     visitUsers(arg);
 }
 
+bool SCCPSolver::markEntryBlockExecutable(Region *region,
+                                          bool markArgsOverdefined) {
+  if (!region->empty()) {
+    if (markArgsOverdefined)
+      markAllOverdefined(region->front().getArguments());
+    return markBlockExecutable(&region->front());
+  }
+  return false;
+}
+
 bool SCCPSolver::markBlockExecutable(Block *block) {
   bool marked = executableBlocks.insert(block).second;
   if (marked)
     blockWorklist.push_back(block);
   return marked;
 }
 
 bool SCCPSolver::isBlockExecutable(Block *block) const {

mlir/test/Transforms/sccp-callgraph.mlir

func @unknown_terminator() -> i32 {
   return %result : i32
 }
 
 // -----
 
 /// Check that return values are overdefined when the constant conflicts.
 
 func private @callable(%arg0 : i32) -> i32 {
-  "unknown.return"(%arg0) : (i32) -> ()
+  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
 
func @non_symbol_defining_callable() -> i32 {
 
   %fn = "test.functional_region_op"() ({
     %1 = constant 1 : i32
     "test.return"(%1) : (i32) -> ()
   }) : () -> (() -> i32)
   %res = call_indirect %fn() : () -> (i32)
   return %res : i32
 }
+
+// -----
+
+/// Check that private callables don't get processed if they have no uses.
+
+// CHECK-LABEL: func private @unreferenced_private_function
+func private @unreferenced_private_function() -> i32 {
+  // CHECK: %[[RES:.*]] = select
+  // CHECK: return %[[RES]] : i32
+  %true = constant true
+  %cst0 = constant 0 : i32
+  %cst1 = constant 1 : i32
+  %result = select %true, %cst0, %cst1 : i32
+  return %result : i32
+}