[mlir] Drop reliance of SliceAnalysis on specific ops.

[nicolasvasilache]

Summary

SliceAnalysis originally was developed in the context of affine.for within mlfunc.
It predates the notion of region.
This revision updates it to not hardcode specific ops like scf::ForOp.
When rooted at an op, the behavior of the slice computation changes as it recurses into the regions of the op. This does not support gathering all values transitively depending on a loop induction variable anymore.
Additional variants rooted at a Value are added to also support the existing behavior.

[ftynse]

Please look at the compilation failure on Windows before submitting.

mlir/include/mlir/Analysis/SliceAnalysis.h

80–81

This may break gcc-5, similarly to the other patch with lambdas as default arguments that we debugged on Friday...

mlir/lib/Analysis/SliceAnalysis.cpp

85

FuncOp is still a specific op :) Should we stop at anything IsolatedFromAbove instead?

[nicolasvasilache]

mlir/include/mlir/Analysis/SliceAnalysis.h

80–81

Moved to llvm::funtion_ref + nullptr default, thanks!

[nicolasvasilache]

Addres comments + rebase.

mlir/include/mlir/Analysis/SliceAnalysis.h

 
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 
 class Operation;
+class Value;
 
 /// Type of the condition to limit the propagation of transitive use-defs.
 /// This can be used in particular to limit the propagation to a given Scope or
 /// to avoid passing through certain types of operation in a configurable
 /// manner.
 using TransitiveFilter = std::function<bool(Operation *)>;
 
 /// Fills `forwardSlice` with the computed forward slice (i.e. all

 /// 2. reversing the result of 1. gives:
 ///      {4, 3, 6, 2, 1, 5, 8, 7, 9}
 ///
 void getForwardSlice(
     Operation *op, llvm::SetVector<Operation *> *forwardSlice,
     TransitiveFilter filter = /* pass-through*/
     [](Operation *) { return true; });
 
+/// Value-rooted version of `getForwardSlice`. Return the union of all forward
+/// slices for the uses of the value `root`.
+void getForwardSlice(
+    Value root, llvm::SetVector<Operation *> *forwardSlice,
+    TransitiveFilter filter = /* pass-through*/
+    [](Operation *) { return true; });
+
 /// Fills `backwardSlice` with the computed backward slice (i.e.
 /// all the transitive defs of op), **without** including that operation.
 ///
 /// This additionally takes a TransitiveFilter which acts as a frontier:
 /// when looking at defs transitively, an operation that does not pass the
 /// filter is never propagated through. This allows in particular to carve out
 /// the scope within a ForInst or the scope within an IfInst.
 ///

 /// Assuming all local orders match the numbering order:
 ///    {1, 2, 5, 3, 4, 6}
 ///
 void getBackwardSlice(
     Operation *op, llvm::SetVector<Operation *> *backwardSlice,
     TransitiveFilter filter = /* pass-through*/
     [](Operation *) { return true; });
 
+/// Value-rooted version of `getBackwardSlice`. Return the union of all backward
+/// slices for the op defining or owning the value `root`.
+void getBackwardSlice(
+    Value root, llvm::SetVector<Operation *> *backwardSlice,
+    TransitiveFilter filter = /* pass-through*/
+    [](Operation *) { return true; });
+
 /// Iteratively computes backward slices and forward slices until
 /// a fixed point is reached. Returns an `llvm::SetVector<Operation *>` which
 /// **includes** the original operation.
 ///
 /// This allows building a slice (i.e. multi-root DAG where everything
 /// that is reachable from an Value in forward and backward direction is
 /// contained in the slice).
 /// This is the abstraction we need to materialize all the operations for

mlir/lib/Analysis/SliceAnalysis.cpp

 
 using namespace mlir;
 
 using llvm::SetVector;
 
 static void getForwardSliceImpl(Operation *op,
                                 SetVector<Operation *> *forwardSlice,
                                 TransitiveFilter filter) {
-  if (!op) {
+  if (!op)
     return;
-  }
 
   // Evaluate whether we should keep this use.
   // This is useful in particular to implement scoping; i.e. return the
   // transitive forwardSlice in the current scope.
-  if (!filter(op)) {
+  if (!filter(op))
     return;
-  }
 
-  if (auto forOp = dyn_cast<AffineForOp>(op)) {
-    for (Operation *userOp : forOp.getInductionVar().getUsers())
-      if (forwardSlice->count(userOp) == 0)
-        getForwardSliceImpl(userOp, forwardSlice, filter);
-  } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-    for (Operation *userOp : forOp.getInductionVar().getUsers())
-      if (forwardSlice->count(userOp) == 0)
-        getForwardSliceImpl(userOp, forwardSlice, filter);
-    for (Value result : forOp.getResults())
-      for (Operation *userOp : result.getUsers())
-        if (forwardSlice->count(userOp) == 0)
-          getForwardSliceImpl(userOp, forwardSlice, filter);
-  } else {
-    assert(op->getNumRegions() == 0 && "unexpected generic op with regions");
-    for (Value result : op->getResults()) {
-      for (Operation *userOp : result.getUsers())
-        if (forwardSlice->count(userOp) == 0)
-          getForwardSliceImpl(userOp, forwardSlice, filter);
-    }
-  }
+  for (Region &region : op->getRegions())
+    for (Block &block : region)
+      for (Operation &op : block)
+        if (forwardSlice->count(&op) == 0)
+          getForwardSliceImpl(&op, forwardSlice, filter);
 
   forwardSlice->insert(op);
 }
 
 void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
                            TransitiveFilter filter) {
   getForwardSliceImpl(op, forwardSlice, filter);
   // Don't insert the top level operation, we just queried on it and don't
   // want it in the results.
   forwardSlice->remove(op);
 
   // Reverse to get back the actual topological order.
   // std::reverse does not work out of the box on SetVector and I want an
   // in-place swap based thing (the real std::reverse, not the LLVM adapter).
   std::vector<Operation *> v(forwardSlice->takeVector());
   forwardSlice->insert(v.rbegin(), v.rend());
 }
 
+void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
+                           TransitiveFilter filter) {
+  for (Operation *user : root.getUsers())
+    getForwardSliceImpl(user, forwardSlice, filter);
+
+  // Reverse to get back the actual topological order.
+  // std::reverse does not work out of the box on SetVector and I want an
+  // in-place swap based thing (the real std::reverse, not the LLVM adapter).
+  std::vector<Operation *> v(forwardSlice->takeVector());
+  forwardSlice->insert(v.rbegin(), v.rend());
+}
+
 static void getBackwardSliceImpl(Operation *op,
                                  SetVector<Operation *> *backwardSlice,
                                  TransitiveFilter filter) {
-  if (!op)
+  if (!op || isa<FuncOp>(op))
     return;
 
-  assert((op->getNumRegions() == 0 ||
-          isa<AffineForOp, scf::ForOp, linalg::LinalgOp, linalg::PadTensorOp>(
-              op)) &&
-         "unexpected generic op with regions");
-
   // Evaluate whether we should keep this def.
   // This is useful in particular to implement scoping; i.e. return the
-  // transitive forwardSlice in the current scope.
-  if (!filter(op)) {
+  // transitive backwardSlice in the current scope.
+  if (!filter(op))
     return;
-  }
 
   for (auto en : llvm::enumerate(op->getOperands())) {
     auto operand = en.value();
-    if (auto blockArg = operand.dyn_cast<BlockArgument>()) {
-      if (auto affIv = getForInductionVarOwner(operand)) {
-        auto *affOp = affIv.getOperation();
-        if (backwardSlice->count(affOp) == 0)
-          getBackwardSliceImpl(affOp, backwardSlice, filter);
-      } else if (auto loopIv = scf::getForInductionVarOwner(operand)) {
-        auto *loopOp = loopIv.getOperation();
-        if (backwardSlice->count(loopOp) == 0)
-          getBackwardSliceImpl(loopOp, backwardSlice, filter);
-      } else if (blockArg.getOwner() !=
-                 &op->getParentOfType<FuncOp>().getBody().front()) {
-        op->emitError("unsupported CF for operand ") << en.index();
-        llvm_unreachable("Unsupported control flow");
-      }
-      continue;
-    }
-    auto *op = operand.getDefiningOp();
-    if (backwardSlice->count(op) == 0) {
-      getBackwardSliceImpl(op, backwardSlice, filter);
+    if (auto *definingOp = operand.getDefiningOp()) {
+      if (backwardSlice->count(definingOp) == 0)
+        getBackwardSliceImpl(definingOp, backwardSlice, filter);
+    } else if (auto blockArg = operand.dyn_cast<BlockArgument>()) {
+      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.
+      assert(parentOp->getNumRegions() == 1 &&
+             parentOp->getRegion(0).getBlocks().size() == 1);
+      if (backwardSlice->count(parentOp) == 0)
+        getBackwardSliceImpl(parentOp, backwardSlice, filter);
+    } else {
+      llvm_unreachable("No definingOp and not a block argument.");
     }
   }
 
   backwardSlice->insert(op);
 }
 
 void mlir::getBackwardSlice(Operation *op,
                             SetVector<Operation *> *backwardSlice,
                             TransitiveFilter filter) {
   getBackwardSliceImpl(op, backwardSlice, filter);
 
   // Don't insert the top level operation, we just queried on it and don't
   // want it in the results.
   backwardSlice->remove(op);
 }
 
+void mlir::getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice,
+                            TransitiveFilter filter) {
+  if (Operation *definingOp = root.getDefiningOp()) {
+    getBackwardSlice(definingOp, backwardSlice, filter);
+    return;
+  }
+  Operation *bbAargOwner = root.cast<BlockArgument>().getOwner()->getParentOp();
+  getBackwardSlice(bbAargOwner, backwardSlice, filter);
+}
+
 SetVector<Operation *> mlir::getSlice(Operation *op,
                                       TransitiveFilter backwardFilter,
                                       TransitiveFilter forwardFilter) {
   SetVector<Operation *> slice;
   slice.insert(op);
 
   unsigned currentIndex = 0;
   SetVector<Operation *> backwardSlice;

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

func.walk([&](vector::TransferReadOp transferRead) {
               cast<LoopLikeOpInterface>(loop.getOperation()))))
         llvm_unreachable(
             "Unexpected failure to move invariant code out of loop");
 
       LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
                         << "\n");
 
       llvm::SetVector<Operation *> forwardSlice;
-      getForwardSlice(transferRead, &forwardSlice);
+      getForwardSlice(transferRead.getOperation(), &forwardSlice);
 
       // Look for the last TransferWriteOp in the forwardSlice of
       // `transferRead` that operates on the same memref.
       vector::TransferWriteOp transferWrite;
       for (auto *sliceOp : llvm::reverse(forwardSlice)) {
         auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
         if (!candidateWrite || candidateWrite.source() != transferRead.source())
           continue;
nextEnclosingForOp =
         nextEnclosingForOp->getParentOfType<LoopLikeOpInterface>();
   }
   if (!outermostEnclosingForOp)
     return failure();
 
   // Get the backwards slice from `padTensorOp` that is dominated by the
   // outermost enclosing loop.
   DominanceInfo domInfo(outermostEnclosingForOp);
-  getBackwardSlice(padTensorOp, &backwardSlice, [&](Operation *op) {
+  getBackwardSlice(padTensorOp.getOperation(), &backwardSlice,
+                   [&](Operation *op) {
                      return domInfo.dominates(outermostEnclosingForOp, op);
                    });
 
   // Bail on any op with a region that is not a LoopLikeInterface or a LinalgOp.
   if (llvm::any_of(backwardSlice, [](Operation *op) {
         return op->getNumRegions() > 0 && !isa<LoopLikeOpInterface>(op) &&
                !isa<LinalgOp>(op);
       }))
     return failure();
 

mlir/lib/Transforms/Utils/LoopUtils.cpp

 
 // Hoist the ops within `outer` that appear before `inner`.
 // Such ops include the ops that have been introduced by parametric tiling.
 // Ops that come from triangular loops (i.e. that belong to the program slice
 // rooted at `outer`) and ops that have side effects cannot be hoisted.
 // Return failure when any op fails to hoist.
 static LogicalResult hoistOpsBetween(scf::ForOp outer, scf::ForOp inner) {
   SetVector<Operation *> forwardSlice;
-  getForwardSlice(outer.getOperation(), &forwardSlice, [&inner](Operation *op) {
-    return op != inner.getOperation();
-  });
+  getForwardSlice(
+      outer.getInductionVar(), &forwardSlice,
+      [&inner](Operation *op) { return op != inner.getOperation(); });
   LogicalResult status = success();
   SmallVector<Operation *, 8> toHoist;
   for (auto &op : outer.getBody()->without_terminator()) {
     // Stop when encountering the inner loop.
     if (&op == inner.getOperation())
       break;
     // Skip over non-hoistable ops.
     if (forwardSlice.count(&op) > 0) {
       status = failure();
       continue;
     }
-    // Skip scf::ForOp, these are not considered a failure.
-    if (op.getNumRegions() > 0)
+    // Skip intermediate scf::ForOp, these are not considered a failure.
+    if (isa<scf::ForOp>(op))
       continue;
     // Skip other ops with regions.
     if (op.getNumRegions() > 0) {
       status = failure();
       continue;
     }
     // Skip if op has side effects.
     // TODO: loads to immutable memory regions are ok.