diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp --- a/mlir/lib/Analysis/SliceAnalysis.cpp +++ b/mlir/lib/Analysis/SliceAnalysis.cpp @@ -41,24 +41,23 @@ } if (auto forOp = dyn_cast(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(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 @@ } // 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 --- a/mlir/test/Dialect/Affine/slicing-utils.mlir +++ b/mlir/test/Dialect/Affine/slicing-utils.mlir @@ -274,6 +274,17 @@ 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