[mlir] Update the CallGraph for nested symbol references, and simplify CallableOpInterface

Summary:
This enables tracking calls that cross symbol table boundaries. It also simplifies some of the implementation details of CallableOpInterface, i.e. there can only be one region within the callable operation.

Depends On D72042

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D72043

Author: River707

mlir/include/mlir/Analysis/CallInterfaces.td

@@ -54,29 +54,23 @@ def CallableOpInterface : OpInterface<"CallableOpInterface"> {
     be a target for a call-like operation (those providing the CallOpInterface
     above). These operations may be traditional functional operation
     `func @foo(...)`, as well as function producing operations
-    `%foo = dialect.create_function(...)`. These operations may produce multiple
-    callable regions, or subroutines.
+    `%foo = dialect.create_function(...)`. These operations may only contain a
+    single region, or subroutine.
   }];
 
   let methods = [
     InterfaceMethod<[{
-        Returns a region on the current operation that the given callable refers
-        to. This may return null in the case of an external callable object,
-        e.g. an external function.
+        Returns the region on the current operation that is callable. This may
+        return null in the case of an external callable object, e.g. an external
+        function.
       }],
-      "Region *", "getCallableRegion", (ins "CallInterfaceCallable":$callable)
+      "Region *", "getCallableRegion"
     >,
     InterfaceMethod<[{
-        Returns all of the callable regions of this operation.
-      }],
-      "void", "getCallableRegions",
-              (ins "SmallVectorImpl<Region *> &":$callables)
-    >,
-    InterfaceMethod<[{
-        Returns the results types that the given callable region produces when
+        Returns the results types that the callable region produces when
         executed.
       }],
-      "ArrayRef<Type>", "getCallableResults", (ins "Region *":$callable)
+      "ArrayRef<Type>", "getCallableResults"
     >,
   ];
 }

mlir/include/mlir/IR/Function.h

@@ -122,26 +122,13 @@ class FuncOp : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
   // CallableOpInterface
   //===--------------------------------------------------------------------===//
 
-  /// Returns a region on the current operation that the given callable refers
-  /// to. This may return null in the case of an external callable object, e.g.
-  /// an external function.
-  Region *getCallableRegion(CallInterfaceCallable callable) {
-    assert(callable.get<SymbolRefAttr>().getLeafReference() == getName());
-    return isExternal() ? nullptr : &getBody();
-  }
-
-  /// Returns all of the callable regions of this operation.
-  void getCallableRegions(SmallVectorImpl<Region *> &callables) {
-    if (!isExternal())
-      callables.push_back(&getBody());
-  }
+  /// Returns the region on the current operation that is callable. This may
+  /// return null in the case of an external callable object, e.g. an external
+  /// function.
+  Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }
 
-  /// Returns the results types that the given callable region produces when
-  /// executed.
-  ArrayRef<Type> getCallableResults(Region *region) {
-    assert(!isExternal() && region == &getBody() && "invalid callable");
-    return getType().getResults();
-  }
+  /// Returns the results types that the callable region produces when executed.
+  ArrayRef<Type> getCallableResults() { return getType().getResults(); }
 
 private:
   // This trait needs access to the hooks defined below.

mlir/lib/Analysis/CallGraph.cpp

@@ -74,67 +74,38 @@ void CallGraphNode::addEdge(CallGraphNode *node, Edge::Kind kind) {
 /// Recursively compute the callgraph edges for the given operation. Computed
 /// edges are placed into the given callgraph object.
 static void computeCallGraph(Operation *op, CallGraph &cg,
-                             CallGraphNode *parentNode);
-
-/// Compute the set of callgraph nodes that are created by regions nested within
-/// 'op'.
-static void computeCallables(Operation *op, CallGraph &cg,
-                             CallGraphNode *parentNode) {
-  if (op->getNumRegions() == 0)
+                             CallGraphNode *parentNode, bool resolveCalls) {
+  if (CallOpInterface call = dyn_cast<CallOpInterface>(op)) {
+    // If there is no parent node, we ignore this operation. Even if this
+    // operation was a call, there would be no callgraph node to attribute it
+    // to.
+    if (!resolveCalls || !parentNode)
+      return;
+    parentNode->addCallEdge(
+        cg.resolveCallable(call.getCallableForCallee(), op));
     return;
-  if (auto callableOp = dyn_cast<CallableOpInterface>(op)) {
-    SmallVector<Region *, 1> callables;
-    callableOp.getCallableRegions(callables);
-    for (auto *callableRegion : callables)
-      cg.getOrAddNode(callableRegion, parentNode);
   }
-}
 
-/// Recursively compute the callgraph edges within the given region. Computed
-/// edges are placed into the given callgraph object.
-static void computeCallGraph(Region &region, CallGraph &cg,
-                             CallGraphNode *parentNode) {
-  // Iterate over the nested operations twice:
-  /// One to fully create nodes in the for each callable region of a nested
-  /// operation;
-  for (auto &block : region)
-    for (auto &nested : block)
-      computeCallables(&nested, cg, parentNode);
-
-  /// And another to recursively compute the callgraph.
-  for (auto &block : region)
-    for (auto &nested : block)
-      computeCallGraph(&nested, cg, parentNode);
-}
-
-/// Recursively compute the callgraph edges for the given operation. Computed
-/// edges are placed into the given callgraph object.
-static void computeCallGraph(Operation *op, CallGraph &cg,
-                             CallGraphNode *parentNode) {
   // Compute the callgraph nodes and edges for each of the nested operations.
-  auto isCallable = isa<CallableOpInterface>(op);
-  for (auto &region : op->getRegions()) {
-    // Check to see if this region is a callable node, if so this is the parent
-    // node of the nested region.
-    CallGraphNode *nestedParentNode;
-    if (!isCallable || !(nestedParentNode = cg.lookupNode(&region)))
-      nestedParentNode = parentNode;
-    computeCallGraph(region, cg, nestedParentNode);
+  if (CallableOpInterface callable = dyn_cast<CallableOpInterface>(op)) {
+    if (auto *callableRegion = callable.getCallableRegion())
+      parentNode = cg.getOrAddNode(callableRegion, parentNode);
+    else
+      return;
   }
 
-  // If there is no parent node, we ignore this operation. Even if this
-  // operation was a call, there would be no callgraph node to attribute it to.
-  if (!parentNode)
-    return;
-
-  // If this is a call operation, resolve the callee.
-  if (auto call = dyn_cast<CallOpInterface>(op))
-    parentNode->addCallEdge(
-        cg.resolveCallable(call.getCallableForCallee(), op));
+  for (Region &region : op->getRegions())
+    for (Block &block : region)
+      for (Operation &nested : block)
+        computeCallGraph(&nested, cg, parentNode, resolveCalls);
 }
 
 CallGraph::CallGraph(Operation *op) : externalNode(/*callableRegion=*/nullptr) {
-  computeCallGraph(op, *this, /*parentNode=*/nullptr);
+  // Make two passes over the graph, one to compute the callables and one to
+  // resolve the calls. We split these up as we may have nested callable objects
+  // that need to be reserved before the calls.
+  computeCallGraph(op, *this, /*parentNode=*/nullptr, /*resolveCalls=*/false);
+  computeCallGraph(op, *this, /*parentNode=*/nullptr, /*resolveCalls=*/true);
 }
 
 /// Get or add a call graph node for the given region.
@@ -175,17 +146,15 @@ CallGraphNode *CallGraph::resolveCallable(CallInterfaceCallable callable,
   // Get the callee operation from the callable.
   Operation *callee;
   if (auto symbolRef = callable.dyn_cast<SymbolRefAttr>())
-    // TODO(riverriddle) Support nested references.
-    callee = SymbolTable::lookupNearestSymbolFrom(from,
-                                                  symbolRef.getRootReference());
+    callee = SymbolTable::lookupNearestSymbolFrom(from, symbolRef);
   else
     callee = callable.get<Value>().getDefiningOp();
 
   // If the callee is non-null and is a valid callable object, try to get the
   // called region from it.
   if (callee && callee->getNumRegions()) {
     if (auto callableOp = dyn_cast_or_null<CallableOpInterface>(callee)) {
-      if (auto *node = lookupNode(callableOp.getCallableRegion(callable)))
+      if (auto *node = lookupNode(callableOp.getCallableRegion()))
         return node;
     }
   }

mlir/lib/Transforms/Inliner.cpp

@@ -86,8 +86,15 @@ static void collectCallOps(iterator_range<Region::iterator> blocks,
   while (!worklist.empty()) {
     for (Operation &op : *worklist.pop_back_val()) {
       if (auto call = dyn_cast<CallOpInterface>(op)) {
-        CallGraphNode *node =
-            cg.resolveCallable(call.getCallableForCallee(), &op);
+        CallInterfaceCallable callable = call.getCallableForCallee();
+
+        // TODO(riverriddle) Support inlining nested call references.
+        if (SymbolRefAttr symRef = callable.dyn_cast<SymbolRefAttr>()) {
+          if (!symRef.isa<FlatSymbolRefAttr>())
+            continue;
+        }
+
+        CallGraphNode *node = cg.resolveCallable(callable, &op);
         if (!node->isExternal())
           calls.emplace_back(call, node);
         continue;
@@ -274,6 +281,15 @@ struct InlinerPass : public OperationPass<InlinerPass> {
     CallGraph &cg = getAnalysis<CallGraph>();
     auto *context = &getContext();
 
+    // The inliner should only be run on operations that define a symbol table,
+    // as the callgraph will need to resolve references.
+    Operation *op = getOperation();
+    if (!op->hasTrait<OpTrait::SymbolTable>()) {
+      op->emitOpError() << " was scheduled to run under the inliner, but does "
+                           "not define a symbol table";
+      return signalPassFailure();
+    }
+
     // Collect a set of canonicalization patterns to use when simplifying
     // callable regions within an SCC.
     OwningRewritePatternList canonPatterns;

mlir/lib/Transforms/Utils/InliningUtils.cpp

@@ -284,7 +284,7 @@ LogicalResult mlir::inlineCall(InlinerInterface &interface,
   if (src->empty())
     return failure();
   auto *entryBlock = &src->front();
-  ArrayRef<Type> callableResultTypes = callable.getCallableResults(src);
+  ArrayRef<Type> callableResultTypes = callable.getCallableResults();
 
   // Make sure that the number of arguments and results matchup between the call
   // and the region.

mlir/test/Analysis/test-callgraph.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -test-print-callgraph 2>&1 | FileCheck %s --dump-input-on-failure
+// RUN: mlir-opt %s -test-print-callgraph -split-input-file 2>&1 | FileCheck %s --dump-input-on-failure
 
 // CHECK-LABEL: Testing : "simple"
 module attributes {test.name = "simple"} {
@@ -50,3 +50,22 @@ module attributes {test.name = "simple"} {
     return
   }
 }
+
+// -----
+
+// CHECK-LABEL: Testing : "nested"
+module attributes {test.name = "nested"} {
+  module @nested_module {
+    // CHECK: Node{{.*}}func_a
+    func @func_a() {
+      return
+    }
+  }
+
+  // CHECK: Node{{.*}}func_b
+  // CHECK: Call-Edge{{.*}}func_a
+  func @func_b() {
+    "test.conversion_call_op"() { callee = @nested_module::@func_a } : () -> ()
+    return
+  }
+}

mlir/test/lib/TestDialect/TestOps.td

@@ -230,7 +230,7 @@ def SizedRegionOp : TEST_Op<"sized_region_op", []> {
 
 def ConversionCallOp : TEST_Op<"conversion_call_op",
     [CallOpInterface]> {
-  let arguments = (ins Variadic<AnyType>:$inputs, FlatSymbolRefAttr:$callee);
+  let arguments = (ins Variadic<AnyType>:$inputs, SymbolRefAttr:$callee);
   let results = (outs Variadic<AnyType>);
 
   let extraClassDeclaration = [{
@@ -239,7 +239,7 @@ def ConversionCallOp : TEST_Op<"conversion_call_op",
 
     /// Return the callee of this operation.
     CallInterfaceCallable getCallableForCallee() {
-      return getAttrOfType<FlatSymbolRefAttr>("callee");
+      return getAttrOfType<SymbolRefAttr>("callee");
     }
   }];
 }
@@ -250,11 +250,8 @@ def FunctionalRegionOp : TEST_Op<"functional_region_op",
   let results = (outs FunctionType);
 
   let extraClassDeclaration = [{
-    Region *getCallableRegion(CallInterfaceCallable) { return &body(); }
-    void getCallableRegions(SmallVectorImpl<Region *> &callables) {
-      callables.push_back(&body());
-    }
-    ArrayRef<Type> getCallableResults(Region *) {
+    Region *getCallableRegion() { return &body(); }
+    ArrayRef<Type> getCallableResults() {
       return getType().cast<FunctionType>().getResults();
     }
   }];