[mlir] Add a utility method on CallOpInterface for resolving the callable.

[rriddle]

Summary

This is the most common operation performed on a CallOpInterface. This just moves the existing functionality from the CallGraph so that other users can access it.

mlir/include/mlir/Analysis/CallGraph.h

 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
+class CallOpInterface;
 struct CallInterfaceCallable;
 class Operation;
 class Region;
 
 //===----------------------------------------------------------------------===//
 // CallGraphNode
 //===----------------------------------------------------------------------===//
 
public:
   CallGraphNode *lookupNode(Region *region) const;
 
   /// Return the callgraph node representing the indirect-external callee.
   CallGraphNode *getExternalNode() const {
     return const_cast<CallGraphNode *>(&externalNode);
   }
 
   /// Resolve the callable for given callee to a node in the callgraph, or the
-  /// external node if a valid node was not resolved. 'from' provides an anchor
-  /// for symbol table lookups, and is only required if the callable is a symbol
-  /// reference.
-  CallGraphNode *resolveCallable(CallInterfaceCallable callable,
-                                 Operation *from = nullptr) const;
+  /// external node if a valid node was not resolved.
+  CallGraphNode *resolveCallable(CallOpInterface call) const;
 
   /// An iterator over the nodes of the graph.
   using iterator = NodeIterator;
   iterator begin() const { return nodes.begin(); }
   iterator end() const { return nodes.end(); }
 
   /// Dump the graph in a human readable format.
   void dump() const;

mlir/include/mlir/Analysis/CallInterfaces.h

 //===- CallInterfaces.h - Call Interfaces for MLIR --------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the definitions of the call interfaces defined in
 // `CallInterfaces.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_CALLINTERFACES_H
 #define MLIR_ANALYSIS_CALLINTERFACES_H
 
-#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/PointerUnion.h"
 
 namespace mlir {
-
 /// A callable is either a symbol, or an SSA value, that is referenced by a
 /// call-like operation. This represents the destination of the call.
 struct CallInterfaceCallable : public PointerUnion<SymbolRefAttr, Value> {
   using PointerUnion<SymbolRefAttr, Value>::PointerUnion;
 };
 
 #include "mlir/Analysis/CallInterfaces.h.inc"
 } // end namespace mlir
 
 #endif // MLIR_ANALYSIS_CALLINTERFACES_H

mlir/include/mlir/Analysis/CallInterfaces.td

InterfaceMethod<[{
       "CallInterfaceCallable", "getCallableForCallee"
     >,
     InterfaceMethod<[{
         Returns the operands within this call that are used as arguments to the
         callee.
       }],
       "Operation::operand_range", "getArgOperands"
     >,
+    InterfaceMethod<[{
+        Resolve the callable operation for given callee to a
+        CallableOpInterface, or nullptr if a valid callable was not resolved.
+      }],
+      "Operation *", "resolveCallable", (ins), [{
+        // If the callable isn't a value, lookup the symbol reference.
+        CallInterfaceCallable callable = op.getCallableForCallee();
+        if (auto symbolRef = callable.dyn_cast<SymbolRefAttr>())
+          return SymbolTable::lookupNearestSymbolFrom(op, symbolRef);
+        return callable.get<Value>().getDefiningOp();
+      }]
+    >,
   ];
 }
 
 /// 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

mlir/lib/Analysis/CallGraph.cpp

 /// 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, 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));
+    if (resolveCalls && parentNode)
+      parentNode->addCallEdge(cg.resolveCallable(call));
     return;
   }
 
   // Compute the callgraph nodes and edges for each of the nested operations.
   if (CallableOpInterface callable = dyn_cast<CallableOpInterface>(op)) {
     if (auto *callableRegion = callable.getCallableRegion())
       parentNode = cg.getOrAddNode(callableRegion, parentNode);
     else

 /// registered.
 CallGraphNode *CallGraph::lookupNode(Region *region) const {
   auto it = nodes.find(region);
   return it == nodes.end() ? nullptr : it->second.get();
 }
 
 /// 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>())
-    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()))
-        return node;
-    }
-  }
+CallGraphNode *CallGraph::resolveCallable(CallOpInterface call) const {
+  Operation *callable = call.resolveCallable();
+  if (auto callableOp = dyn_cast_or_null<CallableOpInterface>(callable))
 
   // If we don't have a valid direct region, this is an external call.
   return getExternalNode();
 }
 
 //===----------------------------------------------------------------------===//
 // Printing
 

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)) {
-        CallInterfaceCallable callable = call.getCallableForCallee();
-
         // TODO(riverriddle) Support inlining nested call references.
+        CallInterfaceCallable callable = call.getCallableForCallee();
         if (SymbolRefAttr symRef = callable.dyn_cast<SymbolRefAttr>()) {
           if (!symRef.isa<FlatSymbolRefAttr>())
             continue;
         }
 
-        CallGraphNode *node = cg.resolveCallable(callable, &op);
+        CallGraphNode *node = cg.resolveCallable(call);
         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.