[mlir] Support operations with multiple results in slicing

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

Differential Revision: https://reviews.llvm.org/D83627

2 files changed, 27 insertions, 20 deletions

diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp
index a09fcf4bea06..8f5f87ba620e 100644
--- a/mlir/lib/Analysis/SliceAnalysis.cpp
+++ b/mlir/lib/Analysis/SliceAnalysis.cpp
@@ -41,24 +41,23 @@ static void getForwardSliceImpl(Operation *op,
   }
 
   if (auto forOp = dyn_cast<AffineForOp>(op)) {
-    for (auto *ownerOp : forOp.getInductionVar().getUsers())
-      if (forwardSlice->count(ownerOp) == 0)
-        getForwardSliceImpl(ownerOp, forwardSlice, filter);
+    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 (auto *ownerOp : forOp.getInductionVar().getUsers())
-      if (forwardSlice->count(ownerOp) == 0)
-        getForwardSliceImpl(ownerOp, forwardSlice, filter);
-    for (auto result : forOp.getResults())
-      for (auto *ownerOp : result.getUsers())
-        if (forwardSlice->count(ownerOp) == 0)
-          getForwardSliceImpl(ownerOp, forwardSlice, filter);
+    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");
-    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 (Value result : op->getResults()) {
+      for (Operation *userOp : result.getUsers())
+        if (forwardSlice->count(userOp) == 0)
+          getForwardSliceImpl(userOp, forwardSlice, filter);
     }
   }
 
@@ -172,12 +171,9 @@ struct DFSState {
 } // 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 (Value result : current->getResults()) {
+    for (Operation *op : result.getUsers())
       DFSPostorder(op, state);
-    }
   }
   bool inserted;
   using IterTy = decltype(state->seen.begin());
diff --git a/mlir/test/Dialect/Affine/slicing-utils.mlir b/mlir/test/Dialect/Affine/slicing-utils.mlir
index 5cc0c3ddcdfb..e11a66b0d0eb 100644
--- a/mlir/test/Dialect/Affine/slicing-utils.mlir
+++ b/mlir/test/Dialect/Affine/slicing-utils.mlir
@@ -274,6 +274,17 @@ func @slicing_test_function_argument(%arg0: 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