[Mlir-commits] [mlir] [mlir] MLIR-QUERY slice-matchers implementation (PR #115670)

[Jacques Pienaar]

================
@@ -363,8 +366,267 @@ 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.
+///
+/// Example starting from node 8
+/// ============================
+///
+///    1       2      3      4
+///    |_______|      |______|
+///    |   |             |
+///    |   5             6
+///    |___|_____________|
+///      |               |
+///      7               8
+///      |_______________|
+///              |
+///              9
+///
+/// Assuming all local orders match the numbering order:
+///    {1, 2, 5, 3, 4, 6}
+///
+
+class BackwardSliceMatcher {
+public:
+  BackwardSliceMatcher(mlir::query::matcher::DynMatcher &&innerMatcher,
+                       int64_t maxDepth)
+      : innerMatcher(std::move(innerMatcher)), maxDepth(maxDepth) {}
+
+  bool match(Operation *op, SetVector<Operation *> &backwardSlice,
+             mlir::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");
----------------
jpienaar wrote:

How about returning at end of if, flattenting "else if" (as post return) and then cast'ing (which would assert) rather than dyn\_cast? For Value, these are the only two variants (its from an op or from a block arg), so the cast should suffice in both behavior and documenting the invariant.

https://github.com/llvm/llvm-project/pull/115670