Improve depth limiting approach

Author: chios202

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

@@ -80,37 +80,33 @@ bool BackwardSliceMatcher<Matcher>::matches(
     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.
+  // Initializing the root op with a depth of 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)
+    // If the subOp hasn't been recorded in opDepths, it is deeper than
+    // maxDepth.
+    if (!opDepths.contains(subOp))
       return false;
     // Examine subOp's operands to compute depths of their defining operations.
     for (auto operand : subOp->getOperands()) {
+      int64_t newDepth = opDepths[subOp] + 1;
+      // If the newDepth is greater than maxDepth, further computation can be
+      // skipped.
+      if (newDepth > maxDepth)
+        continue;
+
       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)) {
+        // Registers the minimum depth
+        if (!opDepths.contains(definingOp) || newDepth < opDepths[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)) {
+
+        if (!opDepths.contains(parentOp) || newDepth < opDepths[parentOp])
           opDepths[parentOp] = newDepth;
-        } else {
-          opDepths[parentOp] = std::min(opDepths[parentOp], newDepth);
-        }
-        return !(opDepths[parentOp] > maxDepth);
       }
     }
     return true;
@@ -119,7 +115,7 @@ bool BackwardSliceMatcher<Matcher>::matches(
   return true;
 }
 
-// Matches transitive defs of a top-level operation up to N levels.
+/// 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,
@@ -130,6 +126,15 @@ m_GetDefinitions(Matcher innerMatcher, int64_t maxDepth, bool inclusive,
                                        omitUsesFromAbove);
 }
 
+/// Matches all transitive defs of a top-level operation up to N levels
+template <typename Matcher>
+inline BackwardSliceMatcher<Matcher> m_GetAllDefinitions(Matcher innerMatcher,
+                                                         int64_t maxDepth) {
+  assert(maxDepth >= 0 && "maxDepth must be non-negative");
+  return BackwardSliceMatcher<Matcher>(std::move(innerMatcher), maxDepth, true,
+                                       false, false);
+}
+
 } // namespace mlir::query::matcher
 
 #endif // MLIR_TOOLS_MLIRQUERY_MATCHERS_EXTRAMATCHERS_H

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

@@ -1,4 +1,4 @@
-// RUN: mlir-query %s -c "m getDefinitions(hasOpName(\"arith.addf\"),2,true,false,false)" | FileCheck %s
+// RUN: mlir-query %s -c "m getAllDefinitions(hasOpName(\"arith.addf\"),2)" | FileCheck %s
 
 #map = affine_map<(d0, d1) -> (d0, d1)>
 func.func @slice_use_from_above(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) {

mlir/tools/mlir-query/mlir-query.cpp

@@ -43,6 +43,9 @@ int main(int argc, char **argv) {
   matcherRegistry.registerMatcher(
       "getDefinitions",
       query::matcher::m_GetDefinitions<query::matcher::DynMatcher>);
+  matcherRegistry.registerMatcher(
+      "getAllDefinitions",
+      query::matcher::m_GetAllDefinitions<query::matcher::DynMatcher>);
   matcherRegistry.registerMatcher("hasOpAttrName",
                                   static_cast<HasOpAttrName *>(m_Attr));
   matcherRegistry.registerMatcher("hasOpName", static_cast<HasOpName *>(m_Op));