[MLIR][LivenessAnalysis] Add custom visitCallOperation

mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h

@@ -82,6 +82,11 @@ class LivenessAnalysis : public SparseBackwardDataFlowAnalysis<Liveness> {
   LogicalResult visitOperation(Operation *op, ArrayRef<Liveness *> operands,
                                ArrayRef<const Liveness *> results) override;
 
+  LogicalResult
+  visitCallOperation(CallOpInterface call,
+                     ArrayRef<Liveness *> operandLattices,
+                     ArrayRef<const Liveness *> resultLattices) override;
+
   void visitBranchOperand(OpOperand &operand) override;
 
   void visitCallOperand(OpOperand &operand) override;

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

@@ -431,6 +431,11 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
       CallOpInterface call, ArrayRef<AbstractSparseLattice *> operandLattices,
       ArrayRef<const AbstractSparseLattice *> resultLattices) = 0;
 
+  /// The transfer function for call operations.
+  virtual LogicalResult visitCallOperationImpl(
+      CallOpInterface call, ArrayRef<AbstractSparseLattice *> operandLattices,
+      ArrayRef<const AbstractSparseLattice *> resultLattices) = 0;
+
   // Visit operands on branch instructions that are not forwarded.
   virtual void visitBranchOperand(OpOperand &operand) = 0;
 
@@ -500,6 +505,7 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   SmallVector<const AbstractSparseLattice *>
   getLatticeElementsFor(ProgramPoint *point, ValueRange values);
 
+protected:
   SymbolTableCollection &symbolTable;
 };
 
@@ -529,6 +535,13 @@ class SparseBackwardDataFlowAnalysis
                                        ArrayRef<StateT *> operands,
                                        ArrayRef<const StateT *> results) = 0;
 
+  // implementation-specific hook for visitCallOperation.
+  virtual LogicalResult visitCallOperation(CallOpInterface call,
+                                       ArrayRef<StateT *> operands,
+                                       ArrayRef<const StateT *> results) {
+                                        return failure();
+                                       }
+
   /// Visit a call to an external function. This function is expected to set
   /// lattice values of the call operands. By default, calls `visitCallOperand`
   /// for all operands.
@@ -583,6 +596,17 @@ class SparseBackwardDataFlowAnalysis
         {reinterpret_cast<const StateT *const *>(resultLattices.begin()),
          resultLattices.size()});
   }
+  LogicalResult visitCallOperationImpl(
+      CallOpInterface call, ArrayRef<AbstractSparseLattice *> operandLattices,
+      ArrayRef<const AbstractSparseLattice *> resultLattices) override {
+    return visitCallOperation(
+        call,
+        {reinterpret_cast<StateT *const *>(operandLattices.begin()),
+         operandLattices.size()},
+        {reinterpret_cast<const StateT *const *>(resultLattices.begin()),
+         resultLattices.size()});
+  }
+
 };
 
 } // end namespace dataflow

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include <cassert>
 #include <mlir/Analysis/DataFlow/LivenessAnalysis.h>
 
@@ -115,6 +116,27 @@ LivenessAnalysis::visitOperation(Operation *op, ArrayRef<Liveness *> operands,
   return success();
 }
 
+LogicalResult
+LivenessAnalysis::visitCallOperation(CallOpInterface call, ArrayRef<Liveness *> operands,
+                                 ArrayRef<const Liveness *> results) {
+  LDBG() << "[visitCallOperation] inspecting " <<  call;
+  // For thread-safety, check config first before accessing symbolTable.
+  if (!getSolverConfig().isInterprocedural()) {
+      visitExternalCall(call, operands, results);
+      return success();
+  }
+  Operation *callableOp = call.resolveCallableInTable(&symbolTable);
+  if (auto funcOp = dyn_cast<FunctionOpInterface>(callableOp)) {
+    if (funcOp.isPublic()) {
+      LDBG() << "[visitCallOperation] encounter a public function " <<  funcOp
+             << " Treat it as external.";
+      visitExternalCall(call, operands, results);
+      return success();
+    }
+  }
+  return failure();
+}
+
 void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   LDBG() << "Visiting branch operand: " << operand.get()
          << " in op: " << *operand.getOwner();

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

@@ -496,25 +496,33 @@ AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
   // operands of the call op that are forwarded to these arguments.
   if (auto call = dyn_cast<CallOpInterface>(op)) {
     LDBG() << "Processing CallOpInterface operation";
+
+    if (visitCallOperationImpl(call, operandLattices, resultLattices).succeeded()) {
+      return success();
+    }
+
+    //  Treat any function as external due to config.
+    if (!getSolverConfig().isInterprocedural()) {
+        visitExternalCallImpl(call, operandLattices, resultLattices);
+        return success();
+    }
+
     Operation *callableOp = call.resolveCallableInTable(&symbolTable);
     if (auto callable = dyn_cast_or_null<CallableOpInterface>(callableOp)) {
       // Not all operands of a call op forward to arguments. Such operands are
       // 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.
+      // If the call invokes an external function, 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() ||
-          !getSolverConfig().isInterprocedural()) {
+      if (!region || region->empty()) {
         visitExternalCallImpl(call, operandLattices, resultLattices);
         return success();
       }
-
       // Otherwise, propagate information from the entry point of the function
       // back to operands whenever possible.
       Block &block = region->front();