Update grammar

Make BackwardSlice matcher more generic
Capture values in tests

Author: chios202

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

@@ -12,16 +12,15 @@
 
 #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
+/// A matcher encapsulating `getBackwardSlice` method from SliceAnalysis.h.
 /// Additionally, it limits the slice computation to a certain depth level using
-/// a custom filter
+/// a custom filter.
 ///
-/// Example starting from node 9, assuming the matcher
-/// computes the slice for the first two depth levels
+/// Example: starting from node 9, assuming the matcher
+/// computes the slice for the first two depth levels:
 /// ============================
 ///    1       2      3      4
 ///    |_______|      |______|
@@ -37,18 +36,23 @@
 /// Assuming all local orders match the numbering order:
 ///     {5, 7, 6, 8, 9}
 namespace mlir::query::matcher {
+
+template <typename Matcher>
 class BackwardSliceMatcher {
 public:
-  explicit BackwardSliceMatcher(query::matcher::DynMatcher &&innerMatcher,
-                                int64_t maxDepth, bool inclusive,
-                                bool omitBlockArguments, bool omitUsesFromAbove)
+  BackwardSliceMatcher(Matcher 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) {
+
+  bool match(Operation *rootOp, SetVector<Operation *> &backwardSlice) {
     BackwardSliceOptions options;
-    return (innerMatcher.match(op) &&
-            matches(op, backwardSlice, options, maxDepth));
+    options.inclusive = inclusive;
+    options.omitUsesFromAbove = omitUsesFromAbove;
+    options.omitBlockArguments = omitBlockArguments;
+    return (innerMatcher.match(rootOp) &&
+            matches(rootOp, backwardSlice, options, maxDepth));
   }
 
 private:
@@ -57,29 +61,75 @@ class BackwardSliceMatcher {
 
 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
+  // to determine whether we want to traverse the IR or not. For example, we
+  // want to explore the IR only if the top-level operation name is
+  // `"arith.addf"`.
+  Matcher innerMatcher;
+  // `maxDepth` specifies the maximum depth that the matcher can traverse the
+  // IR. 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) {
+template <typename Matcher>
+bool BackwardSliceMatcher<Matcher>::matches(
+    Operation *rootOp, llvm::SetVector<Operation *> &backwardSlice,
+    BackwardSliceOptions &options, int64_t maxDepth) {
+  backwardSlice.clear();
+  llvm::DenseMap<Operation *, int64_t> opDepths;
+  // The starting point is the root op; therefore, we set its depth to 0.
+  opDepths[rootOp] = 0;
+  options.filter = [&](Operation *subOp) {
+    // If the subOp's depth exceeds maxDepth, we stop further slicing for this
+    // branch.
+    if (opDepths[subOp] > maxDepth)
+      return false;
+    // Examine subOp's operands to compute depths of their defining operations.
+    for (auto operand : subOp->getOperands()) {
+      if (auto definingOp = operand.getDefiningOp()) {
+        // Set the defining operation's depth to one level greater than
+        // subOp's depth.
+        int64_t newDepth = opDepths[subOp] + 1;
+        if (!opDepths.contains(definingOp)) {
+          opDepths[definingOp] = newDepth;
+        } else {
+          opDepths[definingOp] = std::min(opDepths[definingOp], newDepth);
+        }
+        return !(opDepths[subOp] > maxDepth);
+      } else {
+        auto blockArgument = cast<BlockArgument>(operand);
+        Operation *parentOp = blockArgument.getOwner()->getParentOp();
+        if (!parentOp)
+          continue;
+        int64_t newDepth = opDepths[subOp] + 1;
+        if (!opDepths.contains(parentOp)) {
+          opDepths[parentOp] = newDepth;
+        } else {
+          opDepths[parentOp] = std::min(opDepths[parentOp], newDepth);
+        }
+        return !(opDepths[parentOp] > maxDepth);
+      }
+    }
+    return true;
+  };
+  getBackwardSlice(rootOp, &backwardSlice, options);
+  return true;
+}
+
+// Matches transitive defs of a top-level operation up to N levels.
+template <typename Matcher>
+inline BackwardSliceMatcher<Matcher>
+m_GetDefinitions(Matcher 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);
+  return BackwardSliceMatcher<Matcher>(std::move(innerMatcher), maxDepth,
+                                       inclusive, omitBlockArguments,
+                                       omitUsesFromAbove);
 }
+
 } // namespace mlir::query::matcher
 
 #endif // MLIR_TOOLS_MLIRQUERY_MATCHERS_EXTRAMATCHERS_H

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

@@ -21,32 +21,35 @@
 
 namespace mlir::query::matcher {
 
-/// A class that provides utilities to find operations in a DAG
+/// A class that provides utilities to find operations in the IR.
 class MatchFinder {
 
 public:
-  /// A subclass which preserves the matching information
+  /// A subclass which preserves the matching information. Each instance
+  /// contains the `rootOp` along with the matching environment.
   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
+    /// Contains the matching environment.
     std::vector<Operation *> matchedOps;
   };
-  /// Traverses the DAG and collects the "rootOp" + "matching enviroment" for
-  /// a given Matcher
+
+  /// Traverses the IR and returns a vector of `MatchResult` for each match of
+  /// the `matcher`.
   std::vector<MatchResult> collectMatches(Operation *root,
                                           DynMatcher matcher) const;
-  /// Prints the matched operation
+
+  /// 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
+
+  /// 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
+
+  /// Flattens a vector of `MatchResult` into a vector of operations.
   std::vector<Operation *>
   flattenMatchedOps(std::vector<MatchResult> &matches) const;
 };

mlir/lib/Query/Matcher/CMakeLists.txt

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

mlir/lib/Query/Matcher/ExtraMatchers.cpp

@@ -16,8 +16,11 @@
 // provided to the parser.
 //
 // The grammar for the supported expressions is as follows:
-// <Expression>        := <StringLiteral> | <MatcherExpression>
+// <Expression>        := <Literal> | <MatcherExpression>
+// <Literal>           := <StringLiteral> | <NumericLiteral> | <BooleanLiteral>
 // <StringLiteral>     := "quoted string"
+// <BooleanLiteral>    := "true" | "false"
+// <NumericLiteral>    := [0-9]+
 // <MatcherExpression> := <MatcherName>(<ArgumentList>)
 // <MatcherName>       := [a-zA-Z]+
 // <ArgumentList>      := <Expression> | <Expression>,<ArgumentList>

mlir/lib/Query/Matcher/Parser.h

@@ -166,6 +166,7 @@ QueryRef QueryParser::doParse() {
 
   case ParsedQueryKind::Quit:
     return endQuery(new QuitQuery);
+
   case ParsedQueryKind::Match: {
     if (completionPos) {
       return completeMatcherExpression();

mlir/lib/Query/QueryParser.cpp

@@ -26,7 +26,7 @@ func.func @slice_use_from_above(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>)
 
 // CHECK: Match #2:
 
-// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %0 {{\[\[.*\]\]}} : tensor<5x5xf32> into tensor<25xf32>
+// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %[[LINALG]] {{\[\[.*\]\]}} : tensor<5x5xf32> into tensor<25xf32>
 // CHECK: %[[C2:.*]] = arith.constant {{.*}} : index
 // CHECK: %[[EXTRACTED:.*]] = tensor.extract %[[COLLAPSED]][%[[C2]]] : tensor<25xf32>
 // CHECK: %[[ADDF2:.*]] = arith.addf %[[EXTRACTED]], %[[EXTRACTED]] : f32  

mlir/test/mlir-query/complex-test.mlir

@@ -40,8 +40,9 @@ int main(int argc, char **argv) {
   query::matcher::Registry matcherRegistry;
 
   // Matchers registered in alphabetical order for consistency:
-  matcherRegistry.registerMatcher("getDefinitions",
-                                  query::matcher::m_GetDefinitions);
+  matcherRegistry.registerMatcher(
+      "getDefinitions",
+      query::matcher::m_GetDefinitions<query::matcher::DynMatcher>);
   matcherRegistry.registerMatcher("hasOpAttrName",
                                   static_cast<HasOpAttrName *>(m_Attr));
   matcherRegistry.registerMatcher("hasOpName", static_cast<HasOpName *>(m_Op));