op_deps_pass.cpp

// This LLVM pass takes LLVM bitcode / assembly as input and generates
// dependency graph among aten ops. From a set of root ops used by a model, we
// can calculate transitive closure of all dependent ops, then we can produce a
// custom LibTorch library with optimal build size which only registers and
// contains ops needed by the specific model - unregistered / unused ops can be
// stripped at link time.
//
// [Approach]
// To generate the dependency graph it searches for 3 types of connections in
// LLVM bitcode / assembly:
//  1) op registration: op name (schema string literal) -> registered function;
//  2) regular function call: function -> function;
//  3) op invocation: function -> op name (schema string literal)
//
// For #2 it uses similar algorithm as llvm::LazyCallGraph - not only looks into
// call/invoke instructions but also recursively searches for function pointers
// in each instruction's operands.
//
// For #1 and #3 it searches for connections between operator name string
// literals / function pointers and c10 op registration/invocation API calls in
// LLVM IR graph via "use" edges (bi-directional):
// 1. llvm::Value has "users()" method to get other llvm::Value nodes that use
//    the value;
// 2. most of types derive from llvm::User which has "operands()" method to get
//    other llvm::Value nodes being used by the value;
//
// [Limitation]
// For now the search doesn't go beyond the function boundary because the
// reference to op name string literals and c10 op registration/invocation
// APIs are almost always in the same function. If we create helper function
// around c10 API, we could simply add them to the regular expression used to identify c10 API.
//
// [Example]
// In the following example, it finds out:
//  1) the registered function for "quantized:add" operator;
//  2) one possible call path to at::empty() function;
//  3) the called operator name "aten::empty":
//
// - quantized::add
// - c10::detail::wrap_kernel_functor_unboxed_<at::native::(anonymous
//   namespace)::QAdd<false>, at::Tensor (at::Tensor, at::Tensor, double,
//   long)>::call(c10::OperatorKernel*, at::Tensor, at::Tensor, double, long)
// - at::native::(anonymous namespace)::QAdd<false>::operator()(at::Tensor,
//   at::Tensor, double, long)
// - void at::native::DispatchStub<void (*)(at::Tensor&, at::Tensor const&,
//   at::Tensor const&), at::native::qadd_stub>::operator()<at::Tensor&,
//   at::Tensor const&, at::Tensor const&>(c10::DeviceType, at::Tensor&,
//   at::Tensor const&, at::Tensor const&)
// - at::native::DispatchStub<void (*)(at::Tensor&, at::Tensor const&,
//   at::Tensor const&), at::native::qadd_stub>::choose_cpu_impl()
// - void at::native::(anonymous namespace)::qadd_kernel<false>(at::Tensor&,
//   at::Tensor const&, at::Tensor const&)
// - at::TensorIterator::binary_op(at::Tensor&, at::Tensor const&, at::Tensor
//   const&, bool)
// - at::TensorIterator::build()
// - at::TensorIterator::fast_set_up()
// - at::empty(c10::ArrayRef<long>, c10::TensorOptions const&,
//   c10::optional<c10::MemoryFormat>)
// - aten::empty

#include <deque>
#include <iostream>
#include <set>
#include <unordered_map>
#include <unordered_set>

#include "llvm/Demangle/Demangle.h"
#include "llvm/Analysis/LazyCallGraph.h"
#if LLVM_VERSION_MAJOR < 8
#include "llvm/IR/CallSite.h"
#endif
#include "llvm/IR/Constant.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Regex.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

namespace {

struct RegexOpt {
  std::shared_ptr<Regex> pattern;

  void operator=(const std::string& val) {
    if (val.empty()) {
      return;
    }
    pattern = std::make_shared<Regex>(val);
    std::string regexError;
    if (!pattern->isValid(regexError)) {
      report_fatal_error(
          "Invalid regular expression param: '" + val + "' err: " + regexError,
          false);
    }
  };
};

class RegexOptParser : public cl::basic_parser<RegexOpt> {
public:
  RegexOptParser(cl::Option& O) : basic_parser(O) {}
  virtual ~RegexOptParser() = default;

  // parse - Return true on error.
  bool parse(cl::Option&, StringRef, StringRef Arg, RegexOpt& Value) {
    Value = Arg.str();
    return false;
  }
  StringRef getValueName() const override {
    return "RegexOpt";
  }
};

RegexOpt FunctionSchemaPatternLoc;
cl::opt<RegexOpt, true, cl::parser<std::string>> FunctionSchemaPattern(
    "op_schema_pattern",
    cl::desc("Regular expression used to identify aten op schema strings. "
             "Example: -op_schema_pattern '^(aten|quantized)::[^ ]+'"),
    cl::location(FunctionSchemaPatternLoc),
    cl::Required,
    cl::ValueRequired);

RegexOpt OpRegistrationPatternLoc;
cl::opt<RegexOpt, true, cl::parser<std::string>> OpRegistrationPattern(
    "op_register_pattern",
    cl::desc("Regular expression used to identify c10 op registration API. "
             "Example: -op_register_pattern 'c10::RegisterOperators::op'"),
    cl::location(OpRegistrationPatternLoc),
    cl::Required,
    cl::ValueRequired);

RegexOpt OpInvocationPatternLoc;
cl::opt<RegexOpt, true, cl::parser<std::string>> OpInvocationPattern(
    "op_invoke_pattern",
    cl::desc("Regular expression used to identify c10 op invocation API. "
             "Example: -op_invoke_pattern 'c10::Dispatcher::findSchema'"),
    cl::location(OpInvocationPatternLoc),
    cl::Required,
    cl::ValueRequired);

// The `root_symbol_pattern` is used to specify the seeds of C++ symbols
// from which it searches for transitively reachable ops which need to be
// kept for these C++ APIs to be able to run.
//
// Why not dump ops that are reachable from any visible C++ symbols? Why
// limit it to a subset of root symbols?
// Because op registration callsites in static initializer are visible root
// symbols, too. It will dump ALL the registered ops without any filtering.
//
// Can we use some fixed entry point like `main()`?
// The target to be analyzed can be DSO that doesn't have a `main()`. And
// sometimes we want to get ops that could (but not yet) be called.
//
// This temporary flag will be deprecated by better alternatives in the future.
RegexOpt RootSymbolPatternLoc;
cl::opt<RegexOpt, true, cl::parser<std::string>> RootSymbolPattern(
    "root_symbol_pattern",
    cl::desc("Regular expression used to identify root symbols. It will insert "
             "an entry to the output graph with key = `__ROOT__` and value = "
             "set of ops reachable from root symbols, if the pattern is set. "
             "Example: -root_symbol_pattern 'torch::jit'"),
    cl::location(RootSymbolPatternLoc));

cl::list<RegexOpt, bool, RegexOptParser> TorchLibraryInitPattern(
    "torch_library_init_pattern",
    cl::desc("Regular expression used to identify TorchLibraryInit symbols "
             "that are generated by `TORCH_LIBRARY` macro. The first capturing "
             "group is used to extract namespace string. "
             "Example: -torch_library_init_pattern "
             "'^.*TORCH_LIBRARY_init_([^(]+)(\\(.*)?$'"),
    cl::ZeroOrMore);

cl::opt<int> Verbose(
    "v",
    cl::desc("Verbose level"),
    cl::Hidden,
    cl::init(0));

cl::opt<bool> DebugPath(
    "debug_path",
    cl::desc("Output path between two nodes."),
    cl::init(false));

using SET = std::set<std::string>;
using GRAPH = std::unordered_map<std::string, std::set<std::string>>;
using VALUE_MAP = std::unordered_map<Value*, Value*>;
using VALUE_SET = std::unordered_set<Value*>;

// SRC -> Inverse "tree" from all reachable destinations back to SRC, e.g.:
// (DEST-1 -> PREV_11, PREV_11 -> PREV_12, ..., PREV_1n -> SRC)
// (DEST-2 -> PREV_21, PREV_21 -> PREV_22, ..., PREV_2n -> SRC)
using PATH = std::unordered_map<std::string,
                                std::unordered_map<std::string, std::string>>;

inline std::string _name(const Value* V) {
  return V->getName().str();
}

// Referenced the logic in llvm-cxxfilt.cpp.
// Starting from LLVM 9 it provides a `demangle()` API. Here we keep our ad-hoc
// version for backward compatibility.
std::string _demangle(const std::string& mangled) {
  int status;
  const char* decorated = mangled.c_str();
  size_t decoratedLength = mangled.length();

  char *undecorated = itaniumDemangle(decorated, nullptr, nullptr, &status);

  if (!undecorated &&
      (decoratedLength > 6 && strncmp(decorated, "__imp_", 6) == 0)) {
    undecorated = itaniumDemangle(decorated + 6, nullptr, nullptr, &status);
  }
  std::string result(undecorated ? undecorated : mangled);
  free(undecorated);
  return result;
}

inline bool _isCallSite(Value* V) {
#if LLVM_VERSION_MAJOR >= 8
  return isa<CallBase>(V);
#else
  return !!CallSite(V);
#endif
}

inline Function* _getCalledFunction(Value* V) {
#if LLVM_VERSION_MAJOR >= 8
  return dyn_cast<CallBase>(V)->getCalledFunction();
#else
  return CallSite(V).getCalledFunction();
#endif
}

// LLVM_DEBUG needs opt to be built with debug support.
template<
    typename T,
    typename std::enable_if<std::is_base_of<Value, T>::value, int>::type = 0>
std::ostream& operator<<(std::ostream& out, T& I) {
  std::string str;
  raw_string_ostream O(str);
  O << I;
  return out << str;
}

class OpDependency : public ModulePass {
public:
  static char ID; // Pass identification, replacement for typeid

  OpDependency() : ModulePass(ID) {}
  ~OpDependency() = default;

  bool runOnModule(Module& M) override {
    // Scan all functions and instructions to construct function -> function
    // dependency graph and to find out:
    // - visible functions matching `root_symbol_pattern` option;
    // - instructions that might register or invoke operators, respectively.
    GRAPH deps;
    VALUE_SET visibleFuncs, opRegistrationInsts, opInvocationInsts;
    scanAllFunctions(
        M, &deps, &visibleFuncs, &opRegistrationInsts, &opInvocationInsts);

    // "Key nodes" are nodes we want to keep in output graph. They are usually
    // op-schema strings.
    SET keyNodes;

    // Insert a dummy root node with links to function nodes matching the
    // "root symbol" regex pattern and with default visibility. The goal is to
    // find aten ops that are possibly called via torch C++ APIs.
    insertRoot(visibleFuncs, &deps, &keyNodes);

    // Scan op registration/invocation API calls to construct the link between
    // op name (a.k.a op schema string) and related functions.
    // Dump the op-schema -> function and function -> op-schema mappings into
    // the same `deps` graph with function -> function mappings as they will
    // be processed together next.
    scanOpRegistration(opRegistrationInsts, &keyNodes, &deps);
    scanOpInvocation(opInvocationInsts, &keyNodes, &deps);

    // Shrink the graph by removing intermediate nodes (functions) while
    // maintaining transitive dependency between operators (schema strings).
    GRAPH result;
    std::shared_ptr<PATH> path = DebugPath ? std::make_shared<PATH>() : nullptr;
    simplifyGraph(deps, keyNodes, &result, path.get());

    printAsYAML(std::cout, keyNodes, result, path.get());
    return false;
  }

private:
  static void insertRoot(
      const VALUE_SET& visibleFuncs, GRAPH* deps, SET* keyNodes) {
    if (!RootSymbolPatternLoc.pattern) {
      return;
    }
    SET roots;
    for (const auto& F : visibleFuncs) {
      std::string name = _name(F);
      auto demangled = _demangle(name);
      if (RootSymbolPatternLoc.pattern->match(demangled)) {
        roots.insert(name);
        if (Verbose) {
          std::cerr << "[DEBUG][ROOT_FUNC] " << demangled << std::endl;
        }
      }
    }

    static const std::string ROOT_NODE{"__ROOT__"};
    deps->emplace(ROOT_NODE, std::move(roots));
    keyNodes->insert(ROOT_NODE);
  }

  // Scan the entire IR graph to construct function -> function dependency graph
  // as well as instructions that might register or invoke operators.
  static void scanAllFunctions(
      Module& M, GRAPH* deps, VALUE_SET* visibleFuncs,
      VALUE_SET* opRegistrationInsts, VALUE_SET* opInvocationInsts) {
    for (Function& F : M) {
      if (F.hasDefaultVisibility()) {
        visibleFuncs->insert(&F);
      }
      std::string caller = _name(&F);
      std::string callerDemangled = _demangle(caller);
      for (BasicBlock& BB : F) {
        for (Instruction& I : BB) {
          scanReferredFunctions(I, [&](Function* func) -> void {
            std::string callee = _name(func);
            std::string calleeDemangled = _demangle(callee);
            (*deps)[caller].insert(callee);
            if (Verbose > 1) {
              std::cerr << "[DEBUG][FUNC_CALL] " << callerDemangled << " => "
                        << calleeDemangled << std::endl;
            }
            // One registration/invocation API might call another registration/
            // invocation API in which case we can skip processing the nested
            // call. This is a simple trick to avoid "cannot find registered/
            // invoked op" warning and doesn't affect correctness, because
            // later in scanOpRegistration we'll walk the transitively reachable
            // IR graph again from each registration instance.
            if (!OpRegistrationPatternLoc.pattern->match(callerDemangled) &&
                OpRegistrationPatternLoc.pattern->match(calleeDemangled)) {
              (*opRegistrationInsts).insert(&I);
            }
            if (!OpInvocationPatternLoc.pattern->match(callerDemangled) &&
                OpInvocationPatternLoc.pattern->match(calleeDemangled)) {
              (*opInvocationInsts).insert(&I);
            }
          });
        }
      }
    }
  }

  // llvm::CallGraph only searches for functions referenced by "CallSites" (i.e.
  // by call/invoke instructions). However functions can be referenced by
  // non-call/invoke instructions as well (being passed as function pointer),
  // e.g.:
  // ```
  // store i64 ptrtoint (void (%"class.at::Tensor"*, %"class.at::Tensor"*)*
  //                     @at::foo_op(at::Tensor const&) to i64), i64* %14, ...
  // ```
  // "@at::foo_op" is a operand of "ptrtoint", which in turn is a constant
  // operand of "store" instruction. The stored function pointer can be called
  // indirectly later on.
  //
  // Sometimes directly called functions can be in ConstExpr as well, e.g.:
  // ```
  // invoke void bitcast (
  //    void (ty1*, ...)* @c10::Dispatcher::findSchema(...) to
  //    void (ty2*, ...)*)(...)
  // ```
  // In above case, "CallSite(I).getCalledFunction()" won't return "findSchema"
  // as it's nested in "bitcast" instruction.
  //
  // To cover these cases this method recursively traverses all operands of the
  // input instruction "I" to search for directly/indirectly referenced function
  // pointers by the instruction. The referenced functions might NOT actually be
  // called (which is fine for our use case). llvm::LazyCallGraph has similar
  // logic and we reuse its "visitReferences" method to traverse all operands.
  static void scanReferredFunctions(
      Instruction& I, const std::function<void(Function*)>& CB) {
    SmallVector<Constant*, 16> worklist;
    SmallPtrSet<Constant*, 16> visited;

    if (_isCallSite(&I)) {
      Function* callee = _getCalledFunction(&I);
      if (callee && !callee->isIntrinsic() && visited.insert(callee).second) {
        CB(callee);
      }
    }

    for (Value* op : I.operand_values()) {
      Constant* C = dyn_cast<Constant>(op);
      if (C && visited.insert(C).second) {
        worklist.push_back(C);
      }
    }

    LazyCallGraph::visitReferences(worklist, visited, [&](Function& F) {
      if (!F.isIntrinsic()) {
        CB(&F);
      }
    });
  }

  // Naive connectivity analysis to find out all nodes that are reachable from a
  // specific node in IR graph by following each node's "use" edges (link to its
  // operands and users).
  // This is the core algorithm we use to find the connection between op name
  // string literals and registered/invoked functions - there should be a path
  // to connect them to the c10 op registration/invocation APIs.
  // For now the search doesn't go beyond the function boundary because the
  // reference to op name string literals and c10 op registration/invocation
  // APIs are almost always in the same function.
  static void scanConnectedNodes(
      Value* src,
      VALUE_SET blocked,
      const std::function<void(Value*)>& CB, VALUE_MAP* debugPath) {
    std::deque<Value*> worklist;
    SmallPtrSet<Value*, 16> visited;

    auto insert = [&](Value* cur, Value* parent) -> void {
      if (!blocked.count(cur) && visited.insert(cur).second) {
        worklist.push_back(cur);
        if (debugPath) {
          (*debugPath).emplace(cur, parent);
        }
      }
    };

    auto expandOperands = [&](Value* V) -> void {
      // Stops if it doesn't have operands (!isa<User>) or it is a function.
      if (!isa<User>(V) || isa<Function>(V)) {
        return;
      }
      auto node = dyn_cast<User>(V);
      for (auto& O : node->operands()) {
        insert(O, node);
      }
    };

    auto blockSiblingOperands = [&](User* U, Value* V) -> void {
      // This is to handle a special case only appears in LLVM 9 (not in 5 - 8
      // and 10), where it can falsely associate unrelated PyTorch op
      // registrations.
      //
      // If the value `V` is used by a PHI-node `U`, then we should stop
      // crawling `U`'s operands, i.e. `V`'s siblings in `U`. E.g.:
      //
      //   114:                                            ; preds = %111, %109
      //     %115 = phi i32 [ %110, %109 ], [ %112, %111 ]
      //
      // `%115` might take the value of `%110` or `%112`, depending on from
      // which label it comes. Assuming `V` is `%110` and `U` is `%115`, we can
      // continue to scan `%115` but should not crawl `%112` as it does not
      // directly pass data from `%110` to `%112` (and vice versa).
      //
      // NB: we probably should do the same for other LLVM instructions with
      // this kind of selective semantics. But for the purpose of analyzing
      // PyTorch registrations it seems to be sufficent for now.
      if (isa<PHINode>(U)) {
        for (auto& S : U->operands()) {
          blocked.insert(S);
        }
      }
    };

    auto expandUsers = [&](Value* V) -> void {
      // If the value is not constant, then the user of the value might pass
      // other value into it, e.g.:
      //   store @.str.15, %10
      //   invoke @c10.reg_op, %10, @foo
      // The store instruction, which is the user of "%10", passes "@.str.15" to
      // "%10" which in turn is passed to "@c10.reg_op" API function.
      // Users of constants are not interesting as they cannot change the state
      // of the constant. We skip users of functions as well assuming
      // interesting values (op names and function pointers) are not set via
      // other invocations of the function.
      if (!isa<User>(V) || isa<Constant>(V) || isa<Function>(V)) {
        return;
      }
      for (auto U : V->users()) {
        insert(U, V);
        blockSiblingOperands(U, V);
      }
    };

    auto expand = [&](Value* V) -> void {
      expandOperands(V);
      expandUsers(V);
    };

    expand(src);
    while (!worklist.empty()) {
      auto cur = worklist.front();
      worklist.pop_front();
      expand(cur);

      if (isa<Function>(cur) || isa<Constant>(cur)) {
        CB(cur);
      }
    }
  }

  // Calculate transitive closure and remove intermediate (non-key) nodes.
  // Note that there are two type of nodes in the dependency graph:
  // 1) String literals in source files, e.g.:
  //    "aten::cos_(Tensor(a!) self) -> Tensor(a!)", which represents operator
  //    "schema";
  // 2) Function symbols in object files, e.g.:
  //    "at::CPUType::(anonymous namespace)::cos_(at::Tensor&)";
  // Both of them are added to the dependency graph as std::string. Ultimately
  // we only care about #1 as that's what we use to prune registered ops via
  // codegen, then #2 will be stripped by linker automatically. So the goal is
  // to remove #2 from the graph while maintaining the transitive dependency
  // between #1. #1 is called "key nodes" in this method.
  static void simplifyGraph(
      const GRAPH& input, SET& keyNodes, GRAPH* output, PATH* path) {
    // Starting from every key node, use BFS to traverse all nodes that are
    // transitively reachable from the node in the sparse graph.
    for (auto& key : keyNodes) {
      std::deque<std::string> queue;
      SET visited;  // has some runtime issue with std::unordered_set
      auto expand = [&](const std::string& curNode) -> void {
        auto it = input.find(curNode);
        if (it == input.end()) {
          return;
        }
        for (const auto& next : it->second) {
          if (!visited.insert(next).second) {
            continue;
          }
          queue.push_back(next);
          if (path) {
            (*path)[key].emplace(next, curNode);
          }
        }
      };

      expand(key);
      while (!queue.empty()) {
        auto curNode = queue.front();
        queue.pop_front();
        if (keyNodes.count(curNode)) {
          // Output links between key nodes.
          (*output)[key].insert(curNode);
          // Stop expanding key nodes.
          continue;
        }
        expand(curNode);
      }
    }
  }

  // Find out operator names and function pointers that are transitively
  // connected to the same 'src' instruction.
  static void scanOpSchemaStrAndFunction(
      Instruction* src, const VALUE_SET& blocked,
      const std::string& contextualNamespace,
      SET* visitedOps, SET* visitedFunctions) {
    std::shared_ptr<VALUE_MAP> debugPath =
        (Verbose > 2 ? std::make_shared<VALUE_MAP>() : nullptr);
    auto callback = [&](Value* V) -> void {
      if (auto schemaStr = extractOpSchema(contextualNamespace, V)) {
        if (visitedOps) {
          // NB: Some debug string constants might be connected to the
          // registration instruction, e.g.: "Lambda". Since we have factored
          // out namespace from op schema string, there is no longer a simple
          // way to identify these fake ops. For now we simply take the first
          // instance as the real op name is closest to the seed instruction
          // in BFS order.
          if (!visitedOps->empty()) {
            if (Verbose) {
              std::cerr << "[INFO] ignore extra op schema str: " << *schemaStr
                        << " in: " << _demangle(_name(src->getFunction()))
                        << ", because already found valid op schema str: "
                        << *visitedOps->begin() << std::endl;
            }
          } else {
            (*visitedOps).insert(*schemaStr);
          }
        }
        if (Verbose > 1) {
          std::cerr << "[DEBUG][OP_SCHEMA] " << *schemaStr << std::endl;
          printDebugPath(debugPath.get(), src, V);
        }
      } else if (auto F = dyn_cast<Function>(V)) {
        if (F->isIntrinsic()) {
          return;
        }
        if (visitedFunctions) {
          (*visitedFunctions).insert(_name(F));
        }
        if (Verbose > 1) {
          std::cerr << "[DEBUG][FUNC] " << _demangle(_name(F)) << std::endl;
          printDebugPath(debugPath.get(), src, V);
        }
      }
    };
    scanConnectedNodes(src, blocked, callback, debugPath.get());
  }

  // This method looks for op schema strings and function pointers that connect
  // to the same c10 op registration API call via "use" edges (bi-directional)
  // in IR graph - exploring both nodes being used (operands) by the node and
  // nodes using (users) the node.
  //
  // It assumes that the function pointers are needed (registered) for the op.
  //
  // For example, from op name "aten::add" to registration API call:
  // [OP_SCHEMA] aten::add
  // [PATH][1][CONST] [70 x i8] c"aten::add.Scalar(Tensor self...\00"
  // [PATH][2][CONST] @.str.55.20575 = private unnamed_addr constant [70 x i8]
  //                  c"aten::add.Scalar(Tensor self, ...\00", align 1
  // [PATH][3][CONST] i8* getelementptr inbounds ([70 x i8], [70 x i8]*
  //                  @.str.55.20575, i64 0, i64 0)
  // [PATH][4][INST]  invoke void @std::basic_string<...>::basic_string(...)
  //                  (%"class.std::basic_string"* ... %1477,
  //                  i8* getelementptr ... @.str.55.20575 ...)
  // [PATH][5][INST]  %1477 = alloca %"class.std::basic_string" ...
  // [PATH][6][INST]  %4086 = invoke ...
  //                  @c10::RegisterOperators::Options::schema(... %1477)
  // [PATH][7][INST]  %4088 = invoke ... @...catchAllKernel...(... %4086, ...
  //                  @at::TypeDefault::add(at::Tensor const&...))
  // [PATH][8][INST]  %4090 = invoke ...
  //                  &&(%"class.c10::RegisterOperators::Options"*... %4088 ...)
  // [PATH][9][INST]  invoke void
  //                  @c10::RegisterOperators::checkSchemaAndRegisterOp_(...
  //                  %"class.c10::RegisterOperators::Options"* ... %4090)
  //
  // From function pointer to registration API call:
  // [FUNC] at::TypeDefault::add(at::Tensor const&, c10::Scalar, c10::Scalar)
  // [PATH][1][FUNC]  at::TypeDefault::add(at::Tensor const&...)
  // [PATH][2][INST]  %4088 = invoke ... @...catchAllKernel...(... %4086, ...
  //                  @at::TypeDefault::add(at::Tensor const&...))
  // [PATH][3][INST]  %4090 = invoke ...
  //                  &&(%"class.c10::RegisterOperators::Options"*... %4088 ...)
  // [PATH][4][INST]  invoke void
  //                  @c10::RegisterOperators::checkSchemaAndRegisterOp_(...
  //                  %"class.c10::RegisterOperators::Options"* ... %4090)
  static void scanOpRegistration(
      VALUE_SET& instructions, SET* opSchemaStrs, GRAPH* schemaStrToFunctions) {
    for (auto V : instructions) {
      auto I = dyn_cast<Instruction>(V);
      // We only need to process call/invoke instructions.
      if (!I || !_isCallSite(I)) {
        continue;
      }
      auto contextualNamespace = inferContextualNamespace(I);
      if (Verbose && !contextualNamespace.empty()) {
        std::cerr << "[DEBUG][REG][NAMESPACE] " << contextualNamespace
                  << std::endl;
      }
      if (Verbose > 2) {
        std::cerr << "[DEBUG][REG][INST] " << *I << std::endl;
      }
      SET visitedOps, visitedFunctions;
      // Pass in "instructions" set as "blocked" set - all operator registration
      // calls are connected to global op registry object so we should avoid
      // going from one op registration call to another op registration call via
      // the global registry object.
      scanOpSchemaStrAndFunction(
          I, instructions, contextualNamespace, &visitedOps, &visitedFunctions);
      if (visitedOps.size() != 1) {
        std::cerr << "[WARNING] found " << visitedOps.size() << " ops ( ";
        for (auto& op : visitedOps) {
          std::cerr << op << " ";
        }
        std::cerr << ") in a registration call in function: "
                  << _demangle(_name(I->getFunction()))
                  << " contextualNamespace: " << contextualNamespace
                  << std::endl;
      }
      for (const auto& op : visitedOps) {
        opSchemaStrs->insert(op);
        if (visitedFunctions.empty()) {
          std::cerr << "[WARNING] could not find registered function for op: "
                    << op << " in function: "
                    << _demangle(_name(I->getFunction()))
                    << " contextualNamespace: " << contextualNamespace
                    << std::endl;
        }
        for (const auto& func : visitedFunctions) {
          (*schemaStrToFunctions)[op].insert(func);
          if (Verbose) {
            std::cerr << "[DEBUG][OP_REG] " << op << " => "
                      << _demangle(func) << std::endl;
          }
        }
      }
    }
  }

  static std::string inferContextualNamespace(Instruction* I) {
    auto functionName = _demangle(_name(I->getFunction()));
    for (auto& pattern : TorchLibraryInitPattern) {
      if (!pattern.pattern->match(functionName)) {
        continue;
      }
      if (Verbose) {
        std::cerr << "[DEBUG][REG][INIT_FUNC] " << functionName << std::endl;
      }
      return pattern.pattern->sub("\\1", functionName) + "::";
    }
    return {};
  }

  // Similar as scanOpRegistration - it searches for op schema strings that
  // connect to c10 op invocation API call and assume the parent function of the
  // API call invokes the operator.
  //
  // For example, from op name "aten::empty" to invocation API call:
  // [OP_SCHEMA] aten::empty
  // [PATH][1][CONST] [12 x i8] c"aten::empty\00"
  // [PATH][2][CONST] @.str.69.1990 = private unnamed_addr constant [12 x i8]
  //                  c"aten::empty\00", align 1
  // [PATH][3][CONST] i8* getelementptr inbounds ([12 x i8], [12 x i8]*
  //                  @.str.69.1990, i64 0, i64 0)
  // [PATH][4][INST]  invoke void @std::basic_string<...>::basic_string(...
  //                  (%"class.std::basic_string"* nonnull %19,
  //                  i8* getelementptr inbounds ([12 x i8], [12 x i8]*
  //                  @.str.69.1990, i64 0, i64 0) ...
  // [PATH][5][INST]  %19 = alloca %"class.std::basic_string", align 8
  // [PATH][6][INST]  %53 = bitcast %"class.std::basic_string"* %19 to i64*
  // [PATH][7][INST]  %54 = load i64, i64* %53, align 8, !tbaa !4
  // [PATH][8][INST]  store i64 %54, i64* %55, align 8, !tbaa !4
  // [PATH][9][INST]  %55 = bitcast %"struct.c10::OperatorName"* %18 to i64*
  // [PATH][10][INST] %18 = alloca %"struct.c10::OperatorName", align 8
  // [PATH][11][INST] invoke void @c10::Dispatcher::findSchema(c10::OperatorName
  //                  const&)(%"class.c10::optional.105"* nonnull sret %17,
  //                  %"class.c10::Dispatcher.6320"* nonnull %45,
  //                  %"struct.c10::OperatorName"* nonnull dereferenceable(16)
  //                  %18)
  static void scanOpInvocation(
      VALUE_SET& instructions, SET* opSchemaStrs, GRAPH* functionToSchemaStrs) {
    for (auto V : instructions) {
      auto I = dyn_cast<Instruction>(V);
      // We only need to process call/invoke instructions.
      if (!I || !_isCallSite(I)) {
        continue;
      }
      if (Verbose > 2) {
        std::cerr << "[DEBUG][CALL][INST] " << *I << std::endl;
      }
      std::string caller = _name(I->getFunction());
      SET visitedOps;
      scanOpSchemaStrAndFunction(I, {}, {}, &visitedOps, nullptr);
      if (visitedOps.size() != 1) {
        std::cerr << "[WARNING] found " << visitedOps.size() << " ops ( ";
        for (auto& op : visitedOps) {
          std::cerr << op << " ";
        }
        std::cerr << ") in a invocation call in function: "
                  << _demangle(caller) << std::endl;
      }
      for (const auto& op : visitedOps) {
        opSchemaStrs->insert(op);
        (*functionToSchemaStrs)[caller].insert(op);
        if (Verbose) {
          std::cerr << "[DEBUG][OP_CALL] " << _demangle(caller) << " => "
                    << op << std::endl;
        }
      }
    }
  }

  static void extractStringValue(
      Value* V, const std::function<void(const std::string&)>& CB) {
    if (isa<UndefValue>(V)) {
      // UndefValue inherits from ConstantValue, but don't contain any data
      // See: https://llvm.org/docs/LangRef.html#undefined-values
      return;
    }

    if (auto array = dyn_cast<ConstantDataArray>(V)) {
      // Normal case for c-style string literal and "std::basic_string".
      if (array->isCString()) {
        CB(array->getAsCString().str());
      } else if (array->isString()) {
        std::cerr << "[WARNING] ignore non-C string: "
                  << array->getAsString().str() << std::endl;
      }
    } else if (auto CI = dyn_cast<ConstantInt>(V)) {
      // Short string literal might be encoded into constant integer, e.g.:
      // "aten::AA" => 4702103508586165345 (0x41413A3A6E657461)
      // This can be tricky as it depends on consistent endianness/size.
      // Seen this case for "std::__1::basic_string" ABI.
      uint64_t intValue = CI->getZExtValue();
      auto data = reinterpret_cast<const char*>(&intValue);
      CB({data, data + sizeof(uint64_t)/sizeof(char)});
    } else if (auto C = dyn_cast<Constant>(V)) {
      // Short string literal might be in a constant vector, e.g.:
      // store <2 x i64> <i64 8, i64 4702103508586165345>, <2 x i64>* %25
      // Recursively extract each element to cover this case.
      // Seen this case for "std::__cxx11::basic_string" ABI.
      for (unsigned i = 0; auto elem = C->getAggregateElement(i); ++i) {
        extractStringValue(elem, CB);
      }
    }
  }

  static std::shared_ptr<std::string> extractOpSchema(
      const std::string& contextualNamespace, Value* V) {
    std::vector<std::string> schemaStrs;
    extractStringValue(V, [&](const std::string& str) {
      // NB: some operator names might contain namespace. If this occurs, we
      // MUST NOT use the contextual namespace. Fortunately, it's easy to tell
      // if a namespace is included: a double colon will be present.
      // In particular, this will occur with TORCH_SELECTIVE_NAME.
      const std::string& schemaStr =
          (contextualNamespace.empty() || str.find("::") != std::string::npos)
          ? str : contextualNamespace + str;
      if (FunctionSchemaPatternLoc.pattern->match(schemaStr)) {
        schemaStrs.push_back(schemaStr);
      }
    });
    if (schemaStrs.empty()) {
      return {};
    }
    if (schemaStrs.size() > 1) {
      std::cerr << "[WARNING] found " << schemaStrs.size()
                << " op schema strings in one value!" << std::endl;
    }
    const std::string schemaStr = schemaStrs[0];
    auto pos = schemaStr.find_first_of(".(");
    return std::make_shared<std::string>(
        pos == std::string::npos ? schemaStr : schemaStr.substr(0, pos));
  }

  static void printDebugPath(
      const VALUE_MAP* debugPath, Value* src, Value* dest) {
    if (!debugPath) {
      return;
    }
    int depth = 0;
    for (auto N = dest; ; N = debugPath->at(N)) {
      std::cerr << "[DEBUG][PATH][" << ++depth << "]";
      printDebugValue(N);
      std::cerr << std::endl;
      if (N == src) {
        break;
      }
    }
  }

  static void printDebugValue(Value* V) {
    if (auto F = dyn_cast<Function>(V)) {
      std::cerr << "[FUNC] " << _demangle(_name(F));
    } else if (isa<Constant>(V)) {
      std::cerr << "[CONST] " << *V;
    } else if (isa<Instruction>(V)) {
      std::cerr << "[INST] " << *V;
    } else if (V) {
      std::cerr << "[VALUE] " << *V;
    } else {
      std::cerr << "NULL";
    }
  }

  static void printAsYAML(
      std::ostream& out, const SET& keys, const GRAPH& graph,
      const PATH* path) {
    for (const auto& K : keys) {
      out << "- name: " << _demangle(K) << std::endl;
      auto it = graph.find(K);
      if (it == graph.end() || it->second.empty()) {
        continue;
      }
      out << "  depends:" << std::endl;
      for (const auto& value : it->second) {
        out << "  - name: " << _demangle(value) << std::endl;
        if (path) {
          std::vector<std::string> rpath;
          for (std::string prev = value;
               rpath.push_back(prev), prev != K;
               prev = path->at(K).at(prev));
          out << "    path:" << std::endl;
          for (auto pit = rpath.rbegin(); pit != rpath.rend(); ++pit) {
            out << "    - " << _demangle(*pit) << std::endl;
          }
        }
      }
    }
  }
};

} // namespace

char OpDependency::ID = 0;
static RegisterPass<OpDependency> X("op_dependency", "Op Dependency Pass");