[mlir][DCA] Fix visiting call ops when run at function scopes

[Mogball]

Summary

When dead-code analysis is run at the scope of a function, call ops to
other functions at the same level were being marked as unreachable,
since the analysis optimistically assumes the call op to have no known
predecessors and that all predecessors are known, but the callee would
never get visited.

This patch fixes the bug by checking if a referenced function is above
the top-level op of the analysis, and is thus considered an external
callable.

Fixes #56830

[zero9178]

LGTM!
Tested with my downstream compiler which was originally affected as well and it works perfectly now.

(Here is the corresponding GitHub issue you wanted as well, so that you can refer to it in the commit message: https://github.com/llvm/llvm-project/issues/56830)

mlir/include/mlir/Analysis/DataFlow/DeadCodeAnalysis.h

239

mlir/include/mlir/Analysis/DataFlow/DeadCodeAnalysis.h

private:
 
   /// Mark the entry blocks of the operation as executable.
   void markEntryBlocksLive(Operation *op);
 
   /// Get the constant values of the operands of the operation. Returns none if
   /// any of the operand lattices are uninitialized.
   Optional<SmallVector<Attribute>> getOperandValues(Operation *op);
 
+  /// The top-level operation the analysis is running on. This is used to detect
+  /// if a callable is outside the scope of the analysis and thus must be
+  /// considered an external callable.
+  Operation *analysisScope;
+
   /// A symbol table used for O(1) symbol lookups during simplification.
   SymbolTableCollection symbolTable;
 };
 
 } // end namespace dataflow
 } // end namespace mlir
 
 #endif // MLIR_ANALYSIS_DATAFLOW_DEADCODEANALYSIS_H

mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp

LogicalResult DeadCodeAnalysis::initialize(Operation *top) {
   // Mark as overdefined the predecessors of symbol callables with potentially
   // unknown predecessors.
   initializeSymbolCallables(top);
 
   return initializeRecursively(top);
 }
 
 void DeadCodeAnalysis::initializeSymbolCallables(Operation *top) {
+  analysisScope = top;
   auto walkFn = [&](Operation *symTable, bool allUsesVisible) {
     Region &symbolTableRegion = symTable->getRegion(0);
     Block *symbolTableBlock = &symbolTableRegion.front();
 
     bool foundSymbolCallable = false;
     for (auto callable : symbolTableBlock->getOps<CallableOpInterface>()) {
       Region *callableRegion = callable.getCallableRegion();
       if (!callableRegion)
if (auto branch = dyn_cast<BranchOpInterface>(op)) {
         markEdgeLive(op->getBlock(), successor);
     }
   }
 
   return success();
 }
 
 void DeadCodeAnalysis::visitCallOperation(CallOpInterface call) {
-  Operation *callableOp = nullptr;
-  if (Value callableValue = call.getCallableForCallee().dyn_cast<Value>())
-    callableOp = callableValue.getDefiningOp();
-  else
-    callableOp = call.resolveCallable(&symbolTable);
+  Operation *callableOp = call.resolveCallable(&symbolTable);
 
   // A call to a externally-defined callable has unknown predecessors.
-  const auto isExternalCallable = [](Operation *op) {
+  const auto isExternalCallable = [this](Operation *op) {
+    // A callable outside the analysis scope is an external callable.
+    if (!analysisScope->isAncestor(op))
+      return true;
+    // Otherwise, check if the callable region is defined.
     if (auto callable = dyn_cast<CallableOpInterface>(op))
       return !callable.getCallableRegion();
     return false;
   };
 
   // TODO: Add support for non-symbol callables when necessary. If the
   // callable has non-call uses we would mark as having reached pessimistic
   // fixpoint, otherwise allow for propagating the return values out.

mlir/test/Transforms/sccp-callgraph.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -sccp -split-input-file | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(sccp)" -split-input-file | FileCheck %s --check-prefix=NESTED
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="func.func(sccp)" -split-input-file | FileCheck %s --check-prefix=FUNC
 
 /// Check that a constant is properly propagated through the arguments and
 /// results of a private function.
 
 // CHECK-LABEL: func private @private(
 func.func private @private(%arg0 : i32) -> i32 {
   // CHECK: %[[CST:.*]] = arith.constant 1 : i32
   // CHECK: return %[[CST]] : i32
func.func private @unreferenced_private_function() -> i32 {
   // CHECK: %[[RES:.*]] = arith.select
   // CHECK: return %[[RES]] : i32
   %true = arith.constant true
   %cst0 = arith.constant 0 : i32
   %cst1 = arith.constant 1 : i32
   %result = arith.select %true, %cst0, %cst1 : i32
   return %result : i32
 }
+
+// -----
+
+/// Check that callables outside the analysis scope are marked as external.
+
+func.func private @foo() -> index {
+  %0 = arith.constant 10 : index
+  return %0 : index
+}
+
+// CHECK-LABEL: func @bar
+// FUNC-LABEL: func @bar
+func.func @bar(%arg0: index) -> index {
+  // CHECK: %[[C10:.*]] = arith.constant 10
+  %c0 = arith.constant 0 : index
+  %1 = arith.constant 420 : index
+  %7 = arith.cmpi eq, %arg0, %c0 : index
+  cf.cond_br %7, ^bb1(%1 : index), ^bb2
+
+// CHECK: ^bb1(%[[ARG:.*]]: index):
+// FUNC: ^bb1(%[[ARG:.*]]: index):
+^bb1(%8: index):  // 2 preds: ^bb0, ^bb4
+  // CHECK-NEXT: return %[[ARG]]
+  // FUNC-NEXT: return %[[ARG]]
+  return %8 : index
+
+// CHECK: ^bb2
+// FUNC: ^bb2
+^bb2:
+  // FUNC-NEXT: %[[FOO:.*]] = call @foo
+  %13 = call @foo() : () -> index
+  // CHECK: cf.br ^bb1(%[[C10]]
+  // FUNC: cf.br ^bb1(%[[FOO]]
+  cf.br ^bb1(%13 : index)
+}