[mlir] Support operations with multiple results in slicing

[ThomasRaoux]

Summary

Right now slicing would assert if an operation with multiple results is in the slice.

[bondhugula]

mlir/lib/Analysis/SliceAnalysis.cpp

57

Nit: this shouldn't really be named ownerOp but a userOp. It's an owner op just for the OpOperand which isn't really in the picture.

173–174

Missing space - clang-format please.

174

Nit: Please avoid auto or use a better name (agree that this existed before). auto -> OpOperand or u -> opOperand.

174

You could just do an r.getUsers(). You only need the user ops.

[ThomasRaoux]

Address review comments.

[ThomasRaoux]

mlir/lib/Analysis/SliceAnalysis.cpp

57

I'm not really changing this part of the code and the same naming is used several times so I would rather keep it consistent and send a NFC patch to change all the uses if you think it should be changed.

[bondhugula]

mlir/lib/Analysis/SliceAnalysis.cpp

57

If you are already touching a line, an NFC change to that same line (for example something like naming that improves readability) is completely fine to go along. We don't need a separate NFC patch (unless there are additional surrounding lines of code being modified).

[nicolasvasilache]

mlir/lib/Analysis/SliceAnalysis.cpp

57

+1 for s/ownerOp/userOp/g in this file.
You can just fold it as part of this revision as the changes are small.

174–175

nit - trivial braces here and below.

[mehdi_amini]

mlir/lib/Analysis/SliceAnalysis.cpp

174

The comment about auto haven't been addressed: in general please spell the type when using auto does not improve readability (Value -> auto does not seem better to me, just like Operation * -> auto *)

[ThomasRaoux]

mlir/lib/Analysis/SliceAnalysis.cpp

174

Replace auto.

mlir/lib/Analysis/SliceAnalysis.cpp

static void getForwardSliceImpl(Operation *op,
   // 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)) {
     return;
   }
 
   if (auto forOp = dyn_cast<AffineForOp>(op)) {
-    for (auto *ownerOp : forOp.getInductionVar().getUsers())
-      if (forwardSlice->count(ownerOp) == 0)
-        getForwardSliceImpl(ownerOp, forwardSlice, filter);
+    for (auto *userOp : forOp.getInductionVar().getUsers())
+      if (forwardSlice->count(userOp) == 0)
+        getForwardSliceImpl(userOp, forwardSlice, filter);
   } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-    for (auto *ownerOp : forOp.getInductionVar().getUsers())
-      if (forwardSlice->count(ownerOp) == 0)
-        getForwardSliceImpl(ownerOp, forwardSlice, filter);
+    for (auto *userOp : forOp.getInductionVar().getUsers())
+      if (forwardSlice->count(userOp) == 0)
+        getForwardSliceImpl(userOp, forwardSlice, filter);
     for (auto result : forOp.getResults())
-      for (auto *ownerOp : result.getUsers())
-        if (forwardSlice->count(ownerOp) == 0)
-          getForwardSliceImpl(ownerOp, forwardSlice, filter);
+      for (auto *userOp : result.getUsers())
+        if (forwardSlice->count(userOp) == 0)
+          getForwardSliceImpl(userOp, forwardSlice, filter);
   } else {
     assert(op->getNumRegions() == 0 && "unexpected generic op with regions");
-    assert(op->getNumResults() <= 1 && "unexpected multiple results");
-    if (op->getNumResults() > 0) {
-      for (auto *ownerOp : op->getResult(0).getUsers())
-        if (forwardSlice->count(ownerOp) == 0)
-          getForwardSliceImpl(ownerOp, forwardSlice, filter);
+    for (auto result : op->getResults()) {
+      for (auto *userOp : result.getUsers())
+        if (forwardSlice->count(userOp) == 0)
+          getForwardSliceImpl(userOp, forwardSlice, filter);
     }
   }
 
   forwardSlice->insert(op);
 }
 
 void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
                            TransitiveFilter filter) {
struct DFSState {
   DFSState(const SetVector<Operation *> &set)
       : toSort(set), topologicalCounts(), seen() {}
   const SetVector<Operation *> &toSort;
   SmallVector<Operation *, 16> topologicalCounts;
   DenseSet<Operation *> seen;
 };
 } // namespace
 
 static void DFSPostorder(Operation *current, DFSState *state) {
-  assert(current->getNumResults() <= 1 && "NYI: multi-result");
-  if (current->getNumResults() > 0) {
-    for (auto &u : current->getResult(0).getUses()) {
-      auto *op = u.getOwner();
+  for (auto result : current->getResults()) {
+    for (auto *op : result.getUsers())
       DFSPostorder(op, state);
-    }
-  }
   bool inserted;
   using IterTy = decltype(state->seen.begin());
   IterTy iter;
   std::tie(iter, inserted) = state->seen.insert(current);
   if (inserted) {
     if (state->toSort.count(current) > 0) {
       state->topologicalCounts.push_back(current);
     }

mlir/test/Dialect/Affine/slicing-utils.mlir

 // BWD-LABEL: slicing_test_function_argument
 // FWDBWD-LABEL: slicing_test_function_argument
 func @slicing_test_function_argument(%arg0: index) -> index {
   // BWD: matched: {{.*}} (index, index) -> index backward static slice:
   %0 = "slicing-test-op"(%arg0, %arg0): (index, index) -> index
   return %0 : index
 }
 
+// FWD-LABEL: slicing_test_multiple_return
+// BWD-LABEL: slicing_test_multiple_return
+// FWDBWD-LABEL: slicing_test_multiple_return
+func @slicing_test_multiple_return(%arg0: index) -> (index, index) {
+  // BWD: matched: {{.*}} (index, index) -> (index, index) backward static slice:
+  // FWD: matched: %{{.*}}:2 = "slicing-test-op"(%arg0, %arg0) : (index, index) -> (index, index) forward static slice:
+  // FWD: return %{{.*}}#0, %{{.*}}#1 : index, index
+  %0:2 = "slicing-test-op"(%arg0, %arg0): (index, index) -> (index, index)
+  return %0#0, %0#1 : index, index
+}
+
 // This test dumps 2 sets of outputs: first the test outputs themselves followed
 // by the module. These labels isolate the test outputs from the module dump.
 // FWD-LABEL: slicing_test
 // BWD-LABEL: slicing_test
 // FWDBWD-LABEL: slicing_test
 // FWD-LABEL: slicing_test_2
 // BWD-LABEL: slicing_test_2
 // FWDBWD-LABEL: slicing_test_2
 // FWD-LABEL: slicing_test_3
 // BWD-LABEL: slicing_test_3
 // FWDBWD-LABEL: slicing_test_3
 // FWD-LABEL: slicing_test_function_argument
 // BWD-LABEL: slicing_test_function_argument
 // FWDBWD-LABEL: slicing_test_function_argument