Added SetQuery, LetQuery, new implementation for matchers

Author: dbudii

mlir/include/mlir/IR/Matchers.h

@@ -59,15 +59,15 @@ struct NameOpMatcher {
   NameOpMatcher(StringRef name) : name(name) {}
   bool match(Operation *op) { return op->getName().getStringRef() == name; }
 
-  StringRef name;
+  std::string name;
 };
 
 /// The matcher that matches operations that have the specified attribute name.
 struct AttrOpMatcher {
   AttrOpMatcher(StringRef attrName) : attrName(attrName) {}
   bool match(Operation *op) { return op->hasAttr(attrName); }
 
-  StringRef attrName;
+  std::string attrName;
 };
 
 /// The matcher that matches operations that have the `ConstantLike` trait, and

mlir/include/mlir/Query/Matcher/ExtraMatchers.h

@@ -15,6 +15,9 @@
 
 #include "MatchFinder.h"
 #include "MatchersInternal.h"
+#include "mlir/IR/Region.h"
+#include "mlir/Query/Query.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 
@@ -24,80 +27,161 @@ namespace extramatcher {
 
 namespace detail {
 
-class DefinitionsMatcher {
+class BackwardSliceMatcher {
 public:
-  DefinitionsMatcher(matcher::DynMatcher &&InnerMatcher, unsigned Hops)
-      : InnerMatcher(std::move(InnerMatcher)), Hops(Hops) {}
+  BackwardSliceMatcher(matcher::DynMatcher &&innerMatcher, unsigned hops)
+      : innerMatcher(std::move(innerMatcher)), hops(hops) {}
 
 private:
-  bool matches(Operation *op, matcher::BoundOperationsGraphBuilder &Bound,
-               unsigned TempHops) {
-
-    llvm::DenseSet<mlir::Value> Ccache;
-    llvm::SmallVector<std::pair<Operation *, size_t>, 4> TempStorage;
-    TempStorage.push_back({op, TempHops});
-    while (!TempStorage.empty()) {
-      auto [CurrentOp, RemainingHops] = TempStorage.pop_back_val();
-
-      matcher::BoundOperationNode *CurrentNode =
-          Bound.addNode(CurrentOp, true, true);
-      if (RemainingHops == 0) {
-        continue;
-      }
+  bool matches(Operation *op, SetVector<Operation *> &backwardSlice,
+               QueryOptions &options, unsigned tempHops) {
 
-      for (auto Operand : CurrentOp->getOperands()) {
-        if (auto DefiningOp = Operand.getDefiningOp()) {
-          Bound.addEdge(CurrentOp, DefiningOp);
-          if (!Ccache.contains(Operand)) {
-            Ccache.insert(Operand);
-            TempStorage.emplace_back(DefiningOp, RemainingHops - 1);
-          }
-        } else if (auto BlockArg = Operand.dyn_cast<BlockArgument>()) {
-          auto *Block = BlockArg.getOwner();
+    bool validSlice = true;
+    if (op->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+      return false;
+    }
 
-          if (Block->isEntryBlock() &&
-              isa<FunctionOpInterface>(Block->getParentOp())) {
-            continue;
+    auto processValue = [&](Value value) {
+      if (tempHops == 0) {
+        return;
+      }
+      if (auto *definingOp = value.getDefiningOp()) {
+        if (backwardSlice.count(definingOp) == 0)
+          matches(definingOp, backwardSlice, options, tempHops - 1);
+      } else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
+        if (options.omitBlockArguments)
+          return;
+        Block *block = blockArg.getOwner();
+
+        Operation *parentOp = block->getParentOp();
+
+        if (parentOp && backwardSlice.count(parentOp) == 0) {
+          if (parentOp->getNumRegions() == 1 &&
+              parentOp->getRegion(0).getBlocks().size() == 1) {
+            validSlice = false;
+            return;
+          };
+          matches(parentOp, backwardSlice, options, tempHops - 1);
+        }
+      } else {
+        validSlice = false;
+        return;
+      }
+    };
+
+    if (!options.omitUsesFromAbove) {
+      llvm::for_each(op->getRegions(), [&](Region &region) {
+        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());
+            if (!validSlice)
+              return;
           }
+        });
+      });
+    }
 
-          Operation *ParentOp = BlockArg.getOwner()->getParentOp();
-          if (ParentOp) {
-            Bound.addEdge(CurrentOp, ParentOp);
-            if (!!Ccache.contains(BlockArg)) {
-              Ccache.insert(BlockArg);
-              TempStorage.emplace_back(ParentOp, RemainingHops - 1);
-            }
-          }
-        }
+    llvm::for_each(op->getOperands(), [&](Value operand) {
+      processValue(operand);
+      if (!validSlice)
+        return;
+    });
+    backwardSlice.insert(op);
+    if (!validSlice) {
+      return false;
+    }
+    return true;
+  }
+
+public:
+  bool match(Operation *op, SetVector<Operation *> &backwardSlice,
+             QueryOptions &options) {
+    if (innerMatcher.match(op) && matches(op, backwardSlice, options, hops)) {
+      if (!options.inclusive) {
+        backwardSlice.remove(op);
       }
+      return true;
     }
-    // We need at least 1 defining op
-    return Ccache.size() >= 2;
+    return false;
   }
 
+private:
+  matcher::DynMatcher innerMatcher;
+  unsigned hops;
+};
+
+class ForwardSliceMatcher {
 public:
-  bool match(Operation *op, matcher::BoundOperationsGraphBuilder &Bound) {
-    if (InnerMatcher.match(op) && matches(op, Bound, Hops)) {
+  ForwardSliceMatcher(matcher::DynMatcher &&innerMatcher, unsigned hops)
+      : innerMatcher(std::move(innerMatcher)), hops(hops) {}
+
+private:
+  bool matches(Operation *op, SetVector<Operation *> &forwardSlice,
+               QueryOptions &options, unsigned tempHops) {
+
+    if (tempHops == 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, tempHops - 1);
+    for (Value result : op->getResults()) {
+      for (Operation *userOp : result.getUsers())
+        if (forwardSlice.count(userOp) == 0)
+          matches(userOp, forwardSlice, options, tempHops - 1);
+    }
+
+    forwardSlice.insert(op);
+    return true;
+  }
+
+public:
+  bool match(Operation *op, SetVector<Operation *> &forwardSlice,
+             QueryOptions &options) {
+    if (innerMatcher.match(op) && matches(op, forwardSlice, options, hops)) {
+      if (!options.inclusive) {
+        forwardSlice.remove(op);
+      }
+      SmallVector<Operation *, 0> v(forwardSlice.takeVector());
+      forwardSlice.insert(v.rbegin(), v.rend());
       return true;
     }
     return false;
   }
 
 private:
-  matcher::DynMatcher InnerMatcher;
-  unsigned Hops;
+  matcher::DynMatcher innerMatcher;
+  unsigned hops;
 };
+
 } // namespace detail
 
-inline detail::DefinitionsMatcher
-definedBy(mlir::query::matcher::DynMatcher InnerMatcher) {
-  return detail::DefinitionsMatcher(std::move(InnerMatcher), 1);
+inline detail::BackwardSliceMatcher
+definedBy(mlir::query::matcher::DynMatcher innerMatcher) {
+  return detail::BackwardSliceMatcher(std::move(innerMatcher), 1);
+}
+
+inline detail::BackwardSliceMatcher
+getDefinitions(mlir::query::matcher::DynMatcher innerMatcher, unsigned hops) {
+  return detail::BackwardSliceMatcher(std::move(innerMatcher), hops);
+}
+
+inline detail::ForwardSliceMatcher
+usedBy(mlir::query::matcher::DynMatcher innerMatcher) {
+  return detail::ForwardSliceMatcher(std::move(innerMatcher), 1);
 }
 
-inline detail::DefinitionsMatcher
-getDefinitions(mlir::query::matcher::DynMatcher InnerMatcher, unsigned Hops) {
-  assert(Hops > 0 && "hops must be >= 1");
-  return detail::DefinitionsMatcher(std::move(InnerMatcher), Hops);
+inline detail::ForwardSliceMatcher
+getUses(mlir::query::matcher::DynMatcher innerMatcher, unsigned hops) {
+  return detail::ForwardSliceMatcher(std::move(innerMatcher), hops);
 }
 
 } // namespace extramatcher

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

@@ -23,18 +23,18 @@ namespace mlir::query::matcher {
 class MatchFinder {
 public:
   // Returns all operations that match the given matcher.
-  static BoundOperationsGraphBuilder getMatches(Operation *root,
-                                                DynMatcher matcher) {
-
-    BoundOperationsGraphBuilder Bound;
+  static SetVector<Operation *>
+  getMatches(Operation *root, QueryOptions &options, DynMatcher matcher) {
+    SetVector<Operation *> backwardSlice;
     root->walk([&](Operation *subOp) {
       if (matcher.match(subOp)) {
-        matcher::BoundOperationNode *currentNode = Bound.addNode(subOp);
-      } else if (matcher.match(subOp, Bound)) {
+        backwardSlice.insert(subOp);
+      } else {
+        matcher.match(subOp, backwardSlice, options);
         ////
       }
     });
-    return Bound;
+    return backwardSlice;
   }
 };
 

mlir/include/mlir/Query/Matcher/MatchersInternal.h

@@ -17,59 +17,13 @@
 #include <unordered_set>
 #include <vector>
 
-namespace mlir::query::matcher {
-
-struct BoundOperationNode {
-  Operation *op;
-  std::vector<BoundOperationNode *> Parents;
-  std::vector<BoundOperationNode *> Children;
-
-  bool IsRootNode;
-  bool DetailedPrinting;
-
-  BoundOperationNode(Operation *op, bool IsRootNode = false,
-                     bool DetailedPrinting = false)
-      : op(op), IsRootNode(IsRootNode), DetailedPrinting(DetailedPrinting) {}
-};
+namespace mlir {
+namespace query {
+struct QueryOptions;
+}
+} // namespace mlir
 
-class BoundOperationsGraphBuilder {
-public:
-  BoundOperationNode *addNode(Operation *op, bool IsRootNode = false,
-                              bool DetailedPrinting = false) {
-    auto It = Nodes.find(op);
-    if (It != Nodes.end()) {
-      return It->second.get();
-    }
-    auto Node =
-        std::make_unique<BoundOperationNode>(op, IsRootNode, DetailedPrinting);
-    BoundOperationNode *NodePtr = Node.get();
-    Nodes[op] = std::move(Node);
-    return NodePtr;
-  }
-
-  void addEdge(Operation *parentOp, Operation *childOp) {
-    BoundOperationNode *ParentNode = addNode(parentOp, false, false);
-    BoundOperationNode *ChildNode = addNode(childOp, false, false);
-
-    ParentNode->Children.push_back(ChildNode);
-    ChildNode->Parents.push_back(ParentNode);
-  }
-
-  BoundOperationNode *getNode(Operation *op) const {
-    auto It = Nodes.find(op);
-    return It != Nodes.end() ? It->second.get() : nullptr;
-  }
-
-  const llvm::MapVector<Operation *, std::unique_ptr<BoundOperationNode>> &
-  getNodes() const {
-    return Nodes;
-  }
-
-private:
-  llvm::MapVector<Operation *, std::unique_ptr<BoundOperationNode>> Nodes;
-};
-
-// Type traIt to detect if a matcher has a match(Operation*) method
+namespace mlir::query::matcher {
 template <typename T, typename = void>
 struct has_simple_match : std::false_type {};
 
@@ -78,15 +32,14 @@ struct has_simple_match<T, std::void_t<decltype(std::declval<T>().match(
                                std::declval<Operation *>()))>>
     : std::true_type {};
 
-// Type traIt to detect if a matcher has a match(Operation*,
-// BoundOperationsGraphBuilder&) method
 template <typename T, typename = void>
 struct has_bound_match : std::false_type {};
 
 template <typename T>
 struct has_bound_match<T, std::void_t<decltype(std::declval<T>().match(
                               std::declval<Operation *>(),
-                              std::declval<BoundOperationsGraphBuilder &>()))>>
+                              std::declval<SetVector<Operation *> &>(),
+                              std::declval<QueryOptions &>()))>>
     : std::true_type {};
 
 // Generic interface for matchers on an MLIR operation.
@@ -95,7 +48,8 @@ class MatcherInterface
 public:
   virtual ~MatcherInterface() = default;
   virtual bool match(Operation *op) = 0;
-  virtual bool match(Operation *op, BoundOperationsGraphBuilder &bound) = 0;
+  virtual bool match(Operation *op, SetVector<Operation *> &matchedOps,
+                     QueryOptions &options) = 0;
 };
 
 // MatcherFnImpl takes a matcher function object and implements
@@ -111,9 +65,10 @@ class MatcherFnImpl : public MatcherInterface {
     return false;
   }
 
-  bool match(Operation *op, BoundOperationsGraphBuilder &bound) override {
+  bool match(Operation *op, SetVector<Operation *> &matchedOps,
+             QueryOptions &options) override {
     if constexpr (has_bound_match<MatcherFn>::value)
-      return matcherFn.match(op, bound);
+      return matcherFn.match(op, matchedOps, options);
     return false;
   }
 
@@ -138,8 +93,9 @@ class DynMatcher {
   }
 
   bool match(Operation *op) const { return implementation->match(op); }
-  bool match(Operation *op, BoundOperationsGraphBuilder &bound) const {
-    return implementation->match(op, bound);
+  bool match(Operation *op, SetVector<Operation *> &matchedOps,
+             QueryOptions &options) const {
+    return implementation->match(op, matchedOps, options);
   }
 
   void setFunctionName(StringRef name) { functionName = name.str(); }

mlir/include/mlir/Query/Matcher/VariantValue.h

@@ -100,6 +100,8 @@ class VariantValue {
 
   // String representation of the type of the value.
   std::string getTypeAsString() const;
+  explicit operator bool() const { return hasValue(); }
+  bool hasValue() const { return type != ValueType::Nothing; }
 
 private:
   void reset();