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

[rriddle]

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

[merge_guards_bot]

Unit tests: pass. 61162 tests passed, 0 failed and 728 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[jpienaar]

LG, thanks

mlir/lib/Analysis/CallGraph.cpp

79

I'm trying to think when that would even happen where we aren't in some weird state, but can't, but also not sure it would be an error state so just skipping seems fine.

mlir/include/mlir/Analysis/CallInterfaces.td

 
 /// Interface for callable operations.
 def CallableOpInterface : OpInterface<"CallableOpInterface"> {
   let description = [{
     A callable operation is one who represents a potential sub-routine, and may
     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"
     >,
   ];
 }
 
 #endif // MLIR_CALLINTERFACES

mlir/include/mlir/IR/Function.h

public:
   /// Add a normal block to the end of the function's block list. The function
   /// should at least already have an entry block.
   Block *addBlock();
 
   //===--------------------------------------------------------------------===//
   // 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 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 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 callable region produces when executed.
+  ArrayRef<Type> getCallableResults() { return getType().getResults(); }
 
 private:
   // This trait needs access to the hooks defined below.
   friend class OpTrait::FunctionLike<FuncOp>;
 
   /// Returns the number of arguments. This is a hook for OpTrait::FunctionLike.
   unsigned getNumFuncArguments() { return getType().getInputs().size(); }
 

mlir/lib/Analysis/CallGraph.cpp

 
 //===----------------------------------------------------------------------===//
 // CallGraph
 //===----------------------------------------------------------------------===//
 
 /// 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 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)
+  if (CallableOpInterface callable = dyn_cast<CallableOpInterface>(op)) {
+    if (auto *callableRegion = callable.getCallableRegion())
+      parentNode = cg.getOrAddNode(callableRegion, parentNode);
+    else
       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.
 CallGraphNode *CallGraph::getOrAddNode(Region *region,
                                        CallGraphNode *parentNode) {
   assert(region && isa<CallableOpInterface>(region->getParentOp()) &&
          "expected parent operation to be callable");
   std::unique_ptr<CallGraphNode> &node = nodes[region];

 
 /// Resolve the callable for given callee to a node in the callgraph, or the
 /// external node if a valid node was not resolved.
 CallGraphNode *CallGraph::resolveCallable(CallInterfaceCallable callable,
                                           Operation *from) const {
   // 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;
     }
   }
 
   // If we don't have a valid direct region, this is an external call.
   return getExternalNode();
 }
 

mlir/lib/Transforms/Inliner.cpp

auto addToWorklist = [&](iterator_range<Region::iterator> blocks) {
     for (Block &block : blocks)
       worklist.push_back(&block);
   };
 
   addToWorklist(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;
       }
 
       // If this is not a call, traverse the nested regions. If
       // `traverseNestedCGNodes` is false, then don't traverse nested call graph
       // regions.

 // TODO(riverriddle) This pass should currently only be used for basic testing
 // of inlining functionality.
 namespace {
 struct InlinerPass : public OperationPass<InlinerPass> {
   void runOnOperation() override {
     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;
     for (auto *op : context->getRegisteredOperations())
       op->getCanonicalizationPatterns(canonPatterns, context);
 
     // Run the inline transform in post-order over the SCCs in the callgraph.
     Inliner inliner(context, cg);

mlir/lib/Transforms/Utils/InliningUtils.cpp

 LogicalResult mlir::inlineCall(InlinerInterface &interface,
                                CallOpInterface call,
                                CallableOpInterface callable, Region *src,
                                bool shouldCloneInlinedRegion) {
   // We expect the region to have at least one block.
   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.
   SmallVector<Value, 8> callOperands(call.getArgOperands());
   SmallVector<Value, 8> callResults(call.getOperation()->getResults());
   if (callOperands.size() != entryBlock->getNumArguments() ||
       callResults.size() != callableResultTypes.size())
     return failure();

mlir/test/Analysis/test-callgraph.mlir

-// 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"} {
 
   // CHECK: Node{{.*}}func_a
   func @func_a() {
     return
   }
%fn = "test.functional_region_op"() ({
       call @func_f() : () -> ()
       "test.return"() : () -> ()
     }) : () -> (() -> ())
 
     call_indirect %fn() : () -> ()
     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

 }
 
 //===----------------------------------------------------------------------===//
 // Test Call Interfaces
 //===----------------------------------------------------------------------===//
 
 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 = [{
     /// Get the argument operands to the called function.
     operand_range getArgOperands() { return inputs(); }
 
     /// Return the callee of this operation.
     CallInterfaceCallable getCallableForCallee() {
-      return getAttrOfType<FlatSymbolRefAttr>("callee");
+      return getAttrOfType<SymbolRefAttr>("callee");
     }
   }];
 }
 
 def FunctionalRegionOp : TEST_Op<"functional_region_op",
     [CallableOpInterface]> {
   let regions = (region AnyRegion:$body);
   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();
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Test Traits
 //===----------------------------------------------------------------------===//