diff --git a/mlir/include/mlir/Analysis/SliceAnalysis.h b/mlir/include/mlir/Analysis/SliceAnalysis.h index 99279fdfe427c..a4f5d937cd51d 100644 --- a/mlir/include/mlir/Analysis/SliceAnalysis.h +++ b/mlir/include/mlir/Analysis/SliceAnalysis.h @@ -47,6 +47,11 @@ struct BackwardSliceOptions : public SliceOptions { /// backward slice computation traverses block arguments and asserts that the /// parent op has a single region with a single block. bool omitBlockArguments = false; + + /// When omitUsesFromAbove is true, the backward slice computation omits + /// traversing values that are captured from above. + /// TODO: this should default to `false` after users have been updated. + bool omitUsesFromAbove = true; }; using ForwardSliceOptions = SliceOptions; diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp index 2b1cf411ceeee..8803ba994b2c1 100644 --- a/mlir/lib/Analysis/SliceAnalysis.cpp +++ b/mlir/lib/Analysis/SliceAnalysis.cpp @@ -16,6 +16,7 @@ #include "mlir/IR/Operation.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Support/LLVM.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SetVector.h" #include "llvm/ADT/SmallPtrSet.h" @@ -91,14 +92,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(operand)) { + } else if (auto blockArg = dyn_cast(value)) { if (options.omitBlockArguments) - continue; + return; Block *block = blockArg.getOwner(); Operation *parentOp = block->getParentOp(); @@ -113,7 +113,24 @@ static void getBackwardSliceImpl(Operation *op, } else { llvm_unreachable("No definingOp and not a block argument."); } + }; + + if (!options.omitUsesFromAbove) { + llvm::for_each(op->getRegions(), [&](Region ®ion) { + // Walk this region recursively to collect the regions that descend from + // this op's nested regions (inclusive). + SmallPtrSet descendents; + region.walk( + [&](Region *childRegion) { descendents.insert(childRegion); }); + region.walk([&](Operation *op) { + for (OpOperand &operand : op->getOpOperands()) { + if (!descendents.contains(operand.get().getParentRegion())) + processValue(operand.get()); + } + }); + }); } + llvm::for_each(op->getOperands(), processValue); backwardSlice->insert(op); } diff --git a/mlir/test/IR/slice.mlir b/mlir/test/IR/slice.mlir index 0a32a0f231baf..87d446c8f415a 100644 --- a/mlir/test/IR/slice.mlir +++ b/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 @@ -33,3 +33,29 @@ func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) { // CHECK-DAG: %[[B:.+]] = memref.alloc(%[[ARG2]], %[[ARG1]]) : memref // CHECK-DAG: %[[C:.+]] = memref.alloc(%[[ARG0]], %[[ARG1]]) : memref // 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 diff --git a/mlir/test/lib/IR/TestSlicing.cpp b/mlir/test/lib/IR/TestSlicing.cpp index c3d0d151c6d24..e99d5976d6d9d 100644 --- a/mlir/test/lib/IR/TestSlicing.cpp +++ b/mlir/test/lib/IR/TestSlicing.cpp @@ -39,6 +39,8 @@ static LogicalResult createBackwardSliceFunction(Operation *op, SetVector 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);