Cleanup impl & add test

Author: IanWood1

mlir/lib/Analysis/SliceAnalysis.cpp

@@ -92,7 +92,13 @@ static void getBackwardSliceImpl(Operation *op,
   if (options.filter && !options.filter(op))
     return;
 
-  for (const auto &en : llvm::enumerate(op->getOperands())) {
+  auto operands = op->getOperands();
+  SetVector<Value> valuesToFollow(operands.begin(), operands.end());
+  if (!options.omitUsesFromAbove) {
+    getUsedValuesDefinedAbove(op->getRegions(), valuesToFollow);
+  }
+
+  for (const auto &en : llvm::enumerate(valuesToFollow)) {
     auto operand = en.value();
     if (auto *definingOp = operand.getDefiningOp()) {
       if (backwardSlice->count(definingOp) == 0)
@@ -116,19 +122,6 @@ static void getBackwardSliceImpl(Operation *op,
     }
   }
 
-  // Visit values that are defined above.
-  if (!options.omitUsesFromAbove) {
-    visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand *operand) {
-      if (Operation *definingOp = operand->get().getDefiningOp()) {
-        getBackwardSliceImpl(definingOp, backwardSlice, options);
-        return;
-      }
-      Operation *bbAargOwner =
-          cast<BlockArgument>(operand->get()).getOwner()->getParentOp();
-      getBackwardSliceImpl(bbAargOwner, backwardSlice, options);
-    });
-  }
-
   backwardSlice->insert(op);
 }
 

mlir/test/IR/slice.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -slice-analysis-test %s | FileCheck %s
+// RUN: mlir-opt -slice-analysis-test -split-input-file %s | FileCheck %s
 
 func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
   %a = memref.alloc(%arg0, %arg2) : memref<?x?xf32>
@@ -33,3 +33,29 @@ func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
 //   CHECK-DAG:   %[[B:.+]] = memref.alloc(%[[ARG2]], %[[ARG1]]) : memref<?x?xf32>
 //   CHECK-DAG:   %[[C:.+]] = memref.alloc(%[[ARG0]], %[[ARG1]]) : memref<?x?xf32>
 //       CHECK:   return
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @slice_use_from_above(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) {
+  %0 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    %2 = arith.addf %in, %in : f32
+    linalg.yield %2 : f32
+  } -> tensor<5x5xf32>
+  %collapsed = tensor.collapse_shape %0 [[0, 1]] : tensor<5x5xf32> into tensor<25xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    %c2 = arith.constant 2 : index
+    %extracted = tensor.extract %collapsed[%c2] : tensor<25xf32>
+    %2 = arith.addf %extracted, %extracted : f32
+    linalg.yield %2 : f32
+  } -> tensor<5x5xf32>
+  return
+}
+
+// CHECK-LABEL: func @slice_use_from_above__backward_slice__0
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor 
+//       CHECK:   %[[A:.+]] = linalg.generic {{.*}} ins(%[[ARG0]]
+//       CHECK:   %[[B:.+]] = tensor.collapse_shape %[[A]]
+//       CHECK:   return

mlir/test/lib/IR/TestSlicing.cpp

@@ -39,6 +39,8 @@ static LogicalResult createBackwardSliceFunction(Operation *op,
   SetVector<Operation *> slice;
   BackwardSliceOptions options;
   options.omitBlockArguments = omitBlockArguments;
+  // TODO: Make this default.
+  options.omitUsesFromAbove = false;
   getBackwardSlice(op, &slice, options);
   for (Operation *slicedOp : slice)
     builder.clone(*slicedOp, mapper);