18 changes: 13 additions & 5 deletions mlir/lib/Analysis/SliceAnalysis.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"

Expand Down Expand Up @@ -91,14 +93,13 @@ static void getBackwardSliceImpl(Operation *op,
if (options.filter && !options.filter(op))
return;

for (const auto &en : llvm::enumerate(op->getOperands())) {
auto operand = en.value();
if (auto *definingOp = operand.getDefiningOp()) {
auto processValue = [&](Value value) {
if (auto *definingOp = value.getDefiningOp()) {
if (backwardSlice->count(definingOp) == 0)
getBackwardSliceImpl(definingOp, backwardSlice, options);
} else if (auto blockArg = dyn_cast<BlockArgument>(operand)) {
} else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
if (options.omitBlockArguments)
continue;
return;

Block *block = blockArg.getOwner();
Operation *parentOp = block->getParentOp();
Expand All @@ -113,7 +114,14 @@ static void getBackwardSliceImpl(Operation *op,
} else {
llvm_unreachable("No definingOp and not a block argument.");
}
};

if (!options.omitUsesFromAbove) {
visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand *operand) {
processValue(operand->get());
});
}
llvm::for_each(op->getOperands(), processValue);

backwardSlice->insert(op);
}
Expand Down
28 changes: 27 additions & 1 deletion mlir/test/IR/slice.mlir
View file
Original file line number Diff line number Diff line change
@@ -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>
Expand Down Expand Up @@ -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
2 changes: 2 additions & 0 deletions mlir/test/lib/IR/TestSlicing.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down