Compute shortest depth in backwardSlice method

Relocate backwardSlice matcher to Query specific headers
Remove unncecessary code

Author: chios202

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

@@ -0,0 +1,85 @@
+//===- ExtraMatchers.h - Various common matchers --------------------------===//
+//
+// 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 provides matchers that depend on Query.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_MLIRQUERY_MATCHERS_EXTRAMATCHERS_H
+#define MLIR_TOOLS_MLIRQUERY_MATCHERS_EXTRAMATCHERS_H
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Query/Matcher/MatchersInternal.h"
+
+/// 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 9, assuming the matcher
+/// computes the slice for the first two depth levels
+/// ============================
+///    1       2      3      4
+///    |_______|      |______|
+///    |   |             |
+///    |   5             6
+///    |___|_____________|
+///      |               |
+///      7               8
+///      |_______________|
+///              |
+///              9
+///
+/// Assuming all local orders match the numbering order:
+///     {5, 7, 6, 8, 9}
+namespace mlir::query::matcher {
+class BackwardSliceMatcher {
+public:
+  explicit BackwardSliceMatcher(query::matcher::DynMatcher &&innerMatcher,
+                                int64_t maxDepth, bool inclusive,
+                                bool omitBlockArguments, bool omitUsesFromAbove)
+      : innerMatcher(std::move(innerMatcher)), maxDepth(maxDepth),
+        inclusive(inclusive), omitBlockArguments(omitBlockArguments),
+        omitUsesFromAbove(omitUsesFromAbove) {}
+  bool match(Operation *op, SetVector<Operation *> &backwardSlice) {
+    BackwardSliceOptions options;
+    return (innerMatcher.match(op) &&
+            matches(op, backwardSlice, options, maxDepth));
+  }
+
+private:
+  bool matches(Operation *rootOp, llvm::SetVector<Operation *> &backwardSlice,
+               BackwardSliceOptions &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".
+  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;
+
+  bool inclusive;
+  bool omitBlockArguments;
+  bool omitUsesFromAbove;
+};
+
+// Matches transitive defs of a top level operation up to N levels
+inline BackwardSliceMatcher
+m_GetDefinitions(query::matcher::DynMatcher innerMatcher, int64_t maxDepth,
+                 bool inclusive, bool omitBlockArguments,
+                 bool omitUsesFromAbove) {
+  assert(maxDepth >= 0 && "maxDepth must be non-negative");
+  return BackwardSliceMatcher(std::move(innerMatcher), maxDepth, inclusive,
+                              omitBlockArguments, omitUsesFromAbove);
+}
+} // namespace mlir::query::matcher
+
+#endif // MLIR_TOOLS_MLIRQUERY_MATCHERS_EXTRAMATCHERS_H

mlir/include/mlir/Query/Matcher/Marshallers.h

@@ -50,6 +50,36 @@ struct ArgTypeTraits<llvm::StringRef> {
   }
 };
 
+template <>
+struct ArgTypeTraits<int64_t> {
+  static bool hasCorrectType(const VariantValue &value) {
+    return value.isSigned();
+  }
+
+  static unsigned get(const VariantValue &value) { return value.getSigned(); }
+
+  static ArgKind getKind() { return ArgKind::Signed; }
+
+  static std::optional<std::string> getBestGuess(const VariantValue &) {
+    return std::nullopt;
+  }
+};
+
+template <>
+struct ArgTypeTraits<bool> {
+  static bool hasCorrectType(const VariantValue &value) {
+    return value.isBoolean();
+  }
+
+  static unsigned get(const VariantValue &value) { return value.getBoolean(); }
+
+  static ArgKind getKind() { return ArgKind::Boolean; }
+
+  static std::optional<std::string> getBestGuess(const VariantValue &) {
+    return std::nullopt;
+  }
+};
+
 template <>
 struct ArgTypeTraits<DynMatcher> {
 

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

@@ -7,33 +7,48 @@
 //===----------------------------------------------------------------------===//
 //
 // This file contains the MatchFinder class, which is used to find operations
-// that match a given matcher.
+// that match a given matcher and print them.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
 #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"
 
 namespace mlir::query::matcher {
 
-// MatchFinder is used to find all operations that match a given matcher.
+/// A class that provides utilities to find operations in a DAG
 class MatchFinder {
+
 public:
-  // Returns all operations that match the given matcher.
-  static std::vector<Operation *> getMatches(Operation *root,
-                                             DynMatcher matcher) {
-    std::vector<Operation *> matches;
-
-    // Simple match finding with walk.
-    root->walk([&](Operation *subOp) {
-      if (matcher.match(subOp))
-        matches.push_back(subOp);
-    });
-
-    return matches;
-  }
+  /// A subclass which preserves the matching information
+  struct MatchResult {
+    MatchResult() = default;
+    MatchResult(Operation *rootOp, std::vector<Operation *> matchedOps);
+
+    /// Contains the root operation of the matching environment
+    Operation *rootOp = nullptr;
+    /// 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,
+                                          DynMatcher matcher) 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

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

@@ -8,8 +8,9 @@
 //
 // Implements the base layer of the matcher framework.
 //
-// Matchers are methods that return a Matcher which provides a method
-// match(Operation *op)
+// Matchers are methods that return a Matcher which provides a method one of the
+// following methods: match(Operation *op), match(Operation *op,
+// SetVector<Operation *> &matchedOps)
 //
 // The matcher functions are defined in include/mlir/IR/Matchers.h.
 // This file contains the wrapper classes needed to construct matchers for
@@ -25,13 +26,39 @@
 
 namespace mlir::query::matcher {
 
+// Defaults to false if T has no match() method with the signature:
+// match(Operation* op).
+template <typename T, typename = void>
+struct has_simple_match : std::false_type {};
+
+// Specialized type trait that evaluates to true if T has a match() method
+// with the signature: match(Operation* op).
+template <typename T>
+struct has_simple_match<T, std::void_t<decltype(std::declval<T>().match(
+                               std::declval<Operation *>()))>>
+    : std::true_type {};
+
+// Defaults to false if T has no match() method with the signature:
+// match(Operation* op, SetVector<Operation*>&).
+template <typename T, typename = void>
+struct has_bound_match : std::false_type {};
+
+// Specialized type trait that evaluates to true if T has a match() method
+// with the signature: match(Operation* op, SetVector<Operation*>&).
+template <typename T>
+struct has_bound_match<T, std::void_t<decltype(std::declval<T>().match(
+                              std::declval<Operation *>(),
+                              std::declval<SetVector<Operation *> &>()))>>
+    : std::true_type {};
+
 // Generic interface for matchers on an MLIR operation.
 class MatcherInterface
     : public llvm::ThreadSafeRefCountedBase<MatcherInterface> {
 public:
   virtual ~MatcherInterface() = default;
 
   virtual bool match(Operation *op) = 0;
+  virtual bool match(Operation *op, SetVector<Operation *> &matchedOps) = 0;
 };
 
 // MatcherFnImpl takes a matcher function object and implements
@@ -40,14 +67,25 @@ template <typename MatcherFn>
 class MatcherFnImpl : public MatcherInterface {
 public:
   MatcherFnImpl(MatcherFn &matcherFn) : matcherFn(matcherFn) {}
-  bool match(Operation *op) override { return matcherFn.match(op); }
+
+  bool match(Operation *op) override {
+    if constexpr (has_simple_match<MatcherFn>::value)
+      return matcherFn.match(op);
+    return false;
+  }
+
+  bool match(Operation *op, SetVector<Operation *> &matchedOps) override {
+    if constexpr (has_bound_match<MatcherFn>::value)
+      return matcherFn.match(op, matchedOps);
+    return false;
+  }
 
 private:
   MatcherFn matcherFn;
 };
 
-// Matcher wraps a MatcherInterface implementation and provides a match()
-// method that redirects calls to the underlying implementation.
+// Matcher wraps a MatcherInterface implementation and provides match()
+// methods that redirect calls to the underlying implementation.
 class DynMatcher {
 public:
   // Takes ownership of the provided implementation pointer.
@@ -62,12 +100,13 @@ class DynMatcher {
   }
 
   bool match(Operation *op) const { return implementation->match(op); }
+  bool match(Operation *op, SetVector<Operation *> &matchedOps) const {
+    return implementation->match(op, matchedOps);
+  }
 
-  void setFunctionName(StringRef name) { functionName = name.str(); };
-
-  bool hasFunctionName() const { return !functionName.empty(); };
-
-  StringRef getFunctionName() const { return functionName; };
+  void setFunctionName(StringRef name) { functionName = name.str(); }
+  bool hasFunctionName() const { return !functionName.empty(); }
+  StringRef getFunctionName() const { return functionName; }
 
 private:
   llvm::IntrusiveRefCntPtr<MatcherInterface> implementation;

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

@@ -21,7 +21,7 @@
 namespace mlir::query::matcher {
 
 // All types that VariantValue can contain.
-enum class ArgKind { Matcher, String };
+enum class ArgKind { Boolean, Matcher, Signed, String };
 
 // A variant matcher object to abstract simple and complex matchers into a
 // single object type.
@@ -81,6 +81,8 @@ class VariantValue {
   // Specific constructors for each supported type.
   VariantValue(const llvm::StringRef string);
   VariantValue(const VariantMatcher &matcher);
+  VariantValue(int64_t signedValue);
+  VariantValue(bool setBoolean);
 
   // String value functions.
   bool isString() const;
@@ -92,21 +94,36 @@ class VariantValue {
   const VariantMatcher &getMatcher() const;
   void setMatcher(const VariantMatcher &matcher);
 
+  // Signed value functions.
+  bool isSigned() const;
+  int64_t getSigned() const;
+  void setSigned(int64_t signedValue);
+
+  // Boolean value functions.
+  bool isBoolean() const;
+  bool getBoolean() const;
+  void setBoolean(bool booleanValue);
   // 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();
 
   // All supported value types.
   enum class ValueType {
+    Boolean,
+    Matcher,
     Nothing,
+    Signed,
     String,
-    Matcher,
   };
 
   // All supported value types.
   union AllValues {
+    bool Boolean;
+    int64_t Signed;
     llvm::StringRef *String;
     VariantMatcher *Matcher;
   };

mlir/lib/Query/Matcher/CMakeLists.txt

@@ -1,4 +1,6 @@
 add_mlir_library(MLIRQueryMatcher
+  MatchFinder.cpp
+  ExtraMatchers.cpp
   Parser.cpp
   RegistryManager.cpp
   VariantValue.cpp