Implement nested slicing matcher & enhance MatchFinder class

	- nested slicing matcher
	- enhance MatchFinder class
	- rename getSlice static method to avoid collision with SliceAnalysis::getSlice

Author: chios202

mlir/include/mlir/IR/Matchers.h

@@ -21,8 +21,6 @@
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Query/Matcher/MatchersInternal.h"
 #include "mlir/Query/Query.h"
-#include "llvm/ADT/SetVector.h"
-
 namespace mlir {
 
 namespace detail {
@@ -366,21 +364,14 @@ struct RecursivePatternMatcher {
   std::tuple<OperandMatchers...> operandMatchers;
 };
 
-/// Fills `backwardSlice` with the computed backward slice (i.e.
-/// all the transitive defs of op)
-///
-/// The implementation traverses the def chains in postorder traversal for
-/// efficiency reasons: if an operation is already in `backwardSlice`, no
-/// need to traverse its definitions again. Since use-def chains form a DAG,
-/// this terminates.
-///
-/// Upon return to the root call, `backwardSlice` is filled with a
-/// postorder list of defs. This happens to be a topological order, from the
-/// point of view of the use-def chains.
+/// A matcher encapsulating the initial `getBackwardSlice` method from
+/// SliceAnalysis.h
+/// Additionally, it limits the slice computation to a certain depth level using
+/// a custom filter
 ///
-/// Example starting from node 8
+/// Example starting from node 9, assuming the matcher
+/// computes the slice for the first two depth levels
 /// ============================
-///
 ///    1       2      3      4
 ///    |_______|      |______|
 ///    |   |             |
@@ -393,240 +384,52 @@ struct RecursivePatternMatcher {
 ///              9
 ///
 /// Assuming all local orders match the numbering order:
-///    {1, 2, 5, 3, 4, 6}
-///
-
+///     {5, 7, 6, 8, 9}
 class BackwardSliceMatcher {
 public:
-  BackwardSliceMatcher(mlir::query::matcher::DynMatcher &&innerMatcher,
+  BackwardSliceMatcher(query::matcher::DynMatcher &&innerMatcher,
                        int64_t maxDepth)
       : innerMatcher(std::move(innerMatcher)), maxDepth(maxDepth) {}
-
   bool match(Operation *op, SetVector<Operation *> &backwardSlice,
-             mlir::query::QueryOptions &options) {
+             query::QueryOptions &options) {
 
     if (innerMatcher.match(op) &&
         matches(op, backwardSlice, options, maxDepth)) {
-      if (!options.inclusive) {
-        // Don't insert the top level operation, we just queried on it and don't
-        // want it in the results.
-        backwardSlice.remove(op);
-      }
       return true;
     }
     return false;
   }
 
 private:
-  bool matches(Operation *op, SetVector<Operation *> &backwardSlice,
-               mlir::query::QueryOptions &options, int64_t remainingDepth) {
-
-    if (op->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
-      return false;
-    }
-
-    auto processValue = [&](Value value) {
-      // We need to check the current depth level;
-      // if we have reached level 0, we stop further traversing
-      if (remainingDepth == 0) {
-        return;
-      }
-      if (auto *definingOp = value.getDefiningOp()) {
-        // We omit traversing the same operations
-        if (backwardSlice.count(definingOp) == 0)
-          matches(definingOp, backwardSlice, options, remainingDepth - 1);
-      } else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
-        if (options.omitBlockArguments)
-          return;
-        Block *block = blockArg.getOwner();
-
-        Operation *parentOp = block->getParentOp();
-        // TODO: determine whether we want to recurse backward into the other
-        // blocks of parentOp, which are not technically backward unless they
-        // flow into us. For now, just bail.
-        if (parentOp && backwardSlice.count(parentOp) == 0) {
-          if (parentOp->getNumRegions() != 1 &&
-              parentOp->getRegion(0).getBlocks().size() != 1) {
-            llvm::errs()
-                << "Error: Expected parentOp to have exactly one region and "
-                << "exactly one block, but found " << parentOp->getNumRegions()
-                << " regions and "
-                << (parentOp->getRegion(0).getBlocks().size()) << " blocks.\n";
-          };
-          matches(parentOp, backwardSlice, options, remainingDepth - 1);
-        }
-      } else {
-        llvm_unreachable("No definingOp and not a block argument\n");
-        return;
-      }
-    };
-
-    if (!options.omitUsesFromAbove) {
-      llvm::for_each(op->getRegions(), [&](Region &region) {
-        // Walk this region recursively to collect the regions that descend from
-        // this op's nested regions (inclusive).
-        SmallPtrSet<Region *, 4> descendents;
-        region.walk(
-            [&](Region *childRegion) { descendents.insert(childRegion); });
-        region.walk([&](Operation *op) {
-          for (OpOperand &operand : op->getOpOperands()) {
-            if (!descendents.contains(operand.get().getParentRegion()))
-              processValue(operand.get());
-          }
-        });
-      });
-    }
-
-    llvm::for_each(op->getOperands(), processValue);
-    backwardSlice.insert(op);
-    return true;
-  }
+  bool matches(Operation *rootOp, llvm::SetVector<Operation *> &backwardSlice,
+               query::QueryOptions &options, int64_t maxDepth);
 
 private:
   // The outer matcher (e.g., BackwardSliceMatcher) relies on the innerMatcher
   // to determine whether we want to traverse the DAG or not. For example, we
   // want to explore the DAG only if the top-level operation name is
   // "arith.addf".
-  mlir::query::matcher::DynMatcher innerMatcher;
-
+  query::matcher::DynMatcher innerMatcher;
   // maxDepth specifies the maximum depth that the matcher can traverse in the
   // DAG. For example, if maxDepth is 2, the matcher will explore the defining
   // operations of the top-level op up to 2 levels.
   int64_t maxDepth;
 };
-
-/// Fills `forwardSlice` with the computed forward slice (i.e. all
-/// the transitive uses of op)
-///
-///
-/// The implementation traverses the use chains in postorder traversal for
-/// efficiency reasons: if an operation is already in `forwardSlice`, no
-/// need to traverse its uses again. Since use-def chains form a DAG, this
-/// terminates.
-///
-/// Upon return to the root call, `forwardSlice` is filled with a
-/// postorder list of uses (i.e. a reverse topological order). To get a proper
-/// topological order, we just reverse the order in `forwardSlice` before
-/// returning.
-///
-/// Example starting from node 0
-/// ============================
-///
-///               0
-///    ___________|___________
-///    1       2      3      4
-///    |_______|      |______|
-///    |   |             |
-///    |   5             6
-///    |___|_____________|
-///      |               |
-///      7               8
-///      |_______________|
-///              |
-///              9
-///
-/// Assuming all local orders match the numbering order:
-/// 1. after getting back to the root getForwardSlice, `forwardSlice` may
-///    contain:
-///      {9, 7, 8, 5, 1, 2, 6, 3, 4}
-/// 2. reversing the result of 1. gives:
-///      {4, 3, 6, 2, 1, 5, 8, 7, 9}
-///
-class ForwardSliceMatcher {
-public:
-  ForwardSliceMatcher(mlir::query::matcher::DynMatcher &&innerMatcher,
-                      int64_t maxDepth)
-      : innerMatcher(std::move(innerMatcher)), maxDepth(maxDepth) {}
-
-  bool match(Operation *op, SetVector<Operation *> &forwardSlice,
-             mlir::query::QueryOptions &options) {
-    if (innerMatcher.match(op) &&
-        matches(op, forwardSlice, options, maxDepth)) {
-      if (!options.inclusive) {
-        // Don't insert the top level operation, we just queried on it and don't
-        // want it in the results.
-        forwardSlice.remove(op);
-      }
-      // Reverse to get back the actual topological order.
-      // std::reverse does not work out of the box on SetVector and I want an
-      // in-place swap based thing (the real std::reverse, not the LLVM
-      // adapter).
-      SmallVector<Operation *, 0> v(forwardSlice.takeVector());
-      forwardSlice.insert(v.rbegin(), v.rend());
-      return true;
-    }
-    return false;
-  }
-
-private:
-  bool matches(Operation *op, SetVector<Operation *> &forwardSlice,
-               mlir::query::QueryOptions &options, int64_t remainingDepth) {
-
-    // We need to check the current depth level;
-    // if we have reached level 0, we stop further traversing and insert
-    // the last user in def-use chain
-    if (remainingDepth == 0) {
-      forwardSlice.insert(op);
-      return true;
-    }
-
-    for (Region &region : op->getRegions())
-      for (Block &block : region)
-        for (Operation &blockOp : block)
-          if (forwardSlice.count(&blockOp) == 0)
-            matches(&blockOp, forwardSlice, options, remainingDepth - 1);
-    for (Value result : op->getResults()) {
-      for (Operation *userOp : result.getUsers())
-        // We omit traversing the same operations
-        if (forwardSlice.count(userOp) == 0)
-          matches(userOp, forwardSlice, options, remainingDepth - 1);
-    }
-
-    forwardSlice.insert(op);
-    return true;
-  }
-
-private:
-  // The outer matcher e.g (ForwardSliceMatcher) relies on the innerMatcher to
-  // determine whether we want to traverse the graph or not. E.g: we want to
-  // explore the DAG only if the top level operation name is "arith.addf"
-  mlir::query::matcher::DynMatcher innerMatcher;
-
-  // maxDepth specifies the maximum depth that the matcher can traverse the
-  // graph E.g: if maxDepth is 2, the matcher will explore the user
-  // operations of the top level op up to 2 levels
-  int64_t maxDepth;
-};
-
 } // namespace detail
 
 // Matches transitive defs of a top level operation up to 1 level
 inline detail::BackwardSliceMatcher
-m_DefinedBy(mlir::query::matcher::DynMatcher innerMatcher) {
+m_DefinedBy(query::matcher::DynMatcher innerMatcher) {
   return detail::BackwardSliceMatcher(std::move(innerMatcher), 1);
 }
 
 // Matches transitive defs of a top level operation up to N levels
 inline detail::BackwardSliceMatcher
-m_GetDefinitions(mlir::query::matcher::DynMatcher innerMatcher,
-                 int64_t maxDepth) {
+m_GetDefinitions(query::matcher::DynMatcher innerMatcher, int64_t maxDepth) {
   assert(maxDepth >= 0 && "maxDepth must be non-negative");
   return detail::BackwardSliceMatcher(std::move(innerMatcher), maxDepth);
 }
 
-// Matches uses of a top level operation up to 1 level
-inline detail::ForwardSliceMatcher
-m_UsedBy(mlir::query::matcher::DynMatcher innerMatcher) {
-  return detail::ForwardSliceMatcher(std::move(innerMatcher), 1);
-}
-
-// Matches uses of a top level operation up to N  levels
-inline detail::ForwardSliceMatcher
-m_GetUses(mlir::query::matcher::DynMatcher innerMatcher, int64_t maxDepth) {
-  assert(maxDepth >= 0 && "maxDepth must be non-negative");
-  return detail::ForwardSliceMatcher(std::move(innerMatcher), maxDepth);
-}
-
 /// Matches a constant foldable operation.
 inline detail::constant_op_matcher m_Constant() {
   return detail::constant_op_matcher();

mlir/include/mlir/Query/Matcher/MatchFinder.h

@@ -15,86 +15,41 @@
 #define MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
 
 #include "MatchersInternal.h"
+#include "mlir/Query/Query.h"
 #include "mlir/Query/QuerySession.h"
-#include "llvm/ADT/SetVector.h"
-#include "llvm/Support/SourceMgr.h"
-#include "llvm/Support/raw_ostream.h"
 
 namespace mlir::query::matcher {
 
+/// A class that provides utilities to find operations in a DAG
 class MatchFinder {
 
 public:
-  //
-  // getMatches walks the IR and prints operations as soon as it matches them
-  // if a matcher is to be further extracted into the function, then it does not
-  // print operations
-  //
-  static std::vector<Operation *>
-  getMatches(Operation *root, QueryOptions &options, DynMatcher matcher,
-             llvm::raw_ostream &os, QuerySession &qs) {
-    int matchCount = 0;
-    bool printMatchingOps = true;
-    // If matcher is to be extracted to a function, we don't want to print
-    // matching ops to sdout
-    if (matcher.hasFunctionName()) {
-      printMatchingOps = false;
-    }
-    std::vector<Operation *> matchedOps;
-    SetVector<Operation *> tempStorage;
-    os << "\n";
-    root->walk([&](Operation *subOp) {
-      if (matcher.match(subOp)) {
-        matchedOps.push_back(subOp);
-        if (printMatchingOps) {
-          os << "Match #" << ++matchCount << ":\n\n";
-          printMatch(os, qs, subOp, "root");
-        }
-      } else {
-        SmallVector<Operation *> printingOps;
-        if (matcher.match(subOp, tempStorage, options)) {
-          if (printMatchingOps) {
-            os << "Match #" << ++matchCount << ":\n\n";
-          }
-          SmallVector<Operation *> printingOps(tempStorage.takeVector());
-          for (auto op : printingOps) {
-            if (printMatchingOps) {
-              printMatch(os, qs, op, "root");
-            }
-            matchedOps.push_back(op);
-          }
-          printingOps.clear();
-        }
-      }
-    });
-    if (printMatchingOps) {
-      os << matchCount << (matchCount == 1 ? " match.\n\n" : " matches.\n\n");
-    }
-    return matchedOps;
-  }
-
-private:
-  // Overloaded version that doesn't print the binding
-  static void printMatch(llvm::raw_ostream &os, QuerySession &qs,
-                         mlir::Operation *op) {
-    auto fileLoc = op->getLoc()->dyn_cast<FileLineColLoc>();
-    SMLoc smloc = qs.getSourceManager().FindLocForLineAndColumn(
-        qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
+  /// A subclass which preserves the matching information
+  struct MatchResult {
+    MatchResult() = default;
+    MatchResult(Operation *rootOp, std::vector<Operation *> matchedOps);
 
-    llvm::SMDiagnostic diag =
-        qs.getSourceManager().GetMessage(smloc, llvm::SourceMgr::DK_Note,
+    /// Contains the root operation of the matching environment
+    Operation *rootOp = nullptr;
 
-                                         "");
-    diag.print("", os, true, false, true);
-  }
-  static void printMatch(llvm::raw_ostream &os, QuerySession &qs,
-                         mlir::Operation *op, const std::string &binding) {
-    auto fileLoc = op->getLoc()->findInstanceOf<FileLineColLoc>();
-    auto smloc = qs.getSourceManager().FindLocForLineAndColumn(
-        qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
-    qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,
-                                       "\"" + binding + "\" binds here");
-  }
+    /// Contains the matching enviroment. This allows the user to easily extract
+    /// the matched operations
+    std::vector<Operation *> matchedOps;
+  };
+  /// Traverses the DAG and collects the "rootOp" + "matching enviroment" for a
+  /// given Matcher
+  std::vector<MatchResult>
+  collectMatches(Operation *root, QueryOptions &options, DynMatcher matcher,
+                 llvm::raw_ostream &os, QuerySession &qs) const;
+  /// Prints the matched operation
+  void printMatch(llvm::raw_ostream &os, QuerySession &qs, Operation *op) const;
+  /// Labels the matched operation with the given binding (e.g., "root") and
+  /// prints it
+  void printMatch(llvm::raw_ostream &os, QuerySession &qs, Operation *op,
+                  const std::string &binding) const;
+  /// Flattens a vector of MatchResults into a vector of operations
+  std::vector<Operation *>
+  flattenMatchedOps(std::vector<MatchResult> &matches) const;
 };
 
 } // namespace mlir::query::matcher