[mlir] support non-interprocedural dataflow analyses (#75583)

The core implementation of the dataflow anlysis framework is
interpocedural by design. While this offers better analysis precision,
it also comes with additional cost as it takes longer for the analysis
to reach the fixpoint state. Add a configuration mechanism to the
dataflow solver to control whether it operates inteprocedurally or not
to offer clients a choice.

As a positive side effect, this change also adds hooks for explicitly
processing external/opaque function calls in the dataflow analyses,
e.g., based off of attributes present in the the function declaration or
call operation such as alias scopes and modref available in the LLVM
dialect.

This change should not affect existing analyses and the default solver
configuration remains interprocedural.

Co-authored-by: Jacob Peng <jacobmpeng@gmail.com>

1 files changed, 38 insertions, 19 deletions

diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 9f544d656df9..b47bba16fd90 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -116,8 +116,27 @@ void AbstractSparseForwardDataFlowAnalysis::visitOperation(Operation *op) {
                                  resultLattices);
   }
 
-  // The results of a call operation are determined by the callgraph.
+  // Grab the lattice elements of the operands.
+  SmallVector<const AbstractSparseLattice *> operandLattices;
+  operandLattices.reserve(op->getNumOperands());
+  for (Value operand : op->getOperands()) {
+    AbstractSparseLattice *operandLattice = getLatticeElement(operand);
+    operandLattice->useDefSubscribe(this);
+    operandLattices.push_back(operandLattice);
+  }
+
   if (auto call = dyn_cast<CallOpInterface>(op)) {
+    // If the call operation is to an external function, attempt to infer the
+    // results from the call arguments.
+    auto callable =
+        dyn_cast_if_present<CallableOpInterface>(call.resolveCallable());
+    if (!getSolverConfig().isInterprocedural() ||
+        (callable && !callable.getCallableRegion())) {
+      return visitExternalCallImpl(call, operandLattices, resultLattices);
+    }
+
+    // Otherwise, the results of a call operation are determined by the
+    // callgraph.
     const auto *predecessors = getOrCreateFor<PredecessorState>(op, call);
     // If not all return sites are known, then conservatively assume we can't
     // reason about the data-flow.
@@ -129,15 +148,6 @@ void AbstractSparseForwardDataFlowAnalysis::visitOperation(Operation *op) {
     return;
   }
 
-  // Grab the lattice elements of the operands.
-  SmallVector<const AbstractSparseLattice *> operandLattices;
-  operandLattices.reserve(op->getNumOperands());
-  for (Value operand : op->getOperands()) {
-    AbstractSparseLattice *operandLattice = getLatticeElement(operand);
-    operandLattice->useDefSubscribe(this);
-    operandLattices.push_back(operandLattice);
-  }
-
   // Invoke the operation transfer function.
   visitOperationImpl(op, operandLattices, resultLattices);
 }
@@ -168,8 +178,10 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) {
       const auto *callsites = getOrCreateFor<PredecessorState>(block, callable);
       // If not all callsites are known, conservatively mark all lattices as
       // having reached their pessimistic fixpoints.
-      if (!callsites->allPredecessorsKnown())
+      if (!callsites->allPredecessorsKnown() ||
+          !getSolverConfig().isInterprocedural()) {
         return setAllToEntryStates(argLattices);
+      }
       for (Operation *callsite : callsites->getKnownPredecessors()) {
         auto call = cast<CallOpInterface>(callsite);
         for (auto it : llvm::zip(call.getArgOperands(), argLattices))
@@ -433,19 +445,26 @@ void AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
       // stored in `unaccounted`.
       BitVector unaccounted(op->getNumOperands(), true);
 
+      // If the call invokes an external function (or a function treated as
+      // external due to config), defer to the corresponding extension hook.
+      // By default, it just does `visitCallOperand` for all operands.
       OperandRange argOperands = call.getArgOperands();
       MutableArrayRef<OpOperand> argOpOperands =
           operandsToOpOperands(argOperands);
       Region *region = callable.getCallableRegion();
-      if (region && !region->empty()) {
-        Block &block = region->front();
-        for (auto [blockArg, argOpOperand] :
-             llvm::zip(block.getArguments(), argOpOperands)) {
-          meet(getLatticeElement(argOpOperand.get()),
-               *getLatticeElementFor(op, blockArg));
-          unaccounted.reset(argOpOperand.getOperandNumber());
-        }
+      if (!region || region->empty() || !getSolverConfig().isInterprocedural())
+        return visitExternalCallImpl(call, operandLattices, resultLattices);
+
+      // Otherwise, propagate information from the entry point of the function
+      // back to operands whenever possible.
+      Block &block = region->front();
+      for (auto [blockArg, argOpOperand] :
+           llvm::zip(block.getArguments(), argOpOperands)) {
+        meet(getLatticeElement(argOpOperand.get()),
+             *getLatticeElementFor(op, blockArg));
+        unaccounted.reset(argOpOperand.getOperandNumber());
       }
+
       // Handle the operands of the call op that aren't forwarded to any
       // arguments.
       for (int index : unaccounted.set_bits()) {