[mlir][SliceAnalysis] Add an options object to forward and backward slice.

[mravishankar]

Summary

Add an options object to allow control of the slice computation (for
both forward and backward slice). This makes the ABI stable, and also
allows avoiding an assert that makes the slice analysis unusable for
operations with multiple blocks.

[hanchung]

mlir/lib/Analysis/SliceAnalysis.cpp

107–111

Let's merge two if-statement into one if-statement.

[mravishankar]

Rebase

[nicolasvasilache]

It seems quite dangerous to just drop everything on the floor when we have != 1 regions / blocks.

Additionally, there are no tests in this PR.

Not sure what you are trying to achieve here but you may be in a "this needs a better algorithm" territory.

[mravishankar]

In D151520#4383339, @nicolasvasilache wrote:

It seems quite dangerous to just drop everything on the floor when we have != 1 regions / blocks.

I am not changing the behavior. I was getting an assert if the region had more than 1 block. Makes it impossible to use this with a func.func op with control flow.

Additionally, there are no tests in this PR.

Not sure what you are trying to achieve here but you may be in a "this needs a better algorithm" territory.

Better algorithm is definitely needed. But I am not sure what that looks like. I am just trying to use the existing slice function for a very trivial use case and being foiled by an assert. So, if having the assert is important I can maybe adapt the existing setup to add a SliceOptions object. It can assert by default, but opt-in to avoid the assert for people to do at their risk (I can add test. Didnt add cause I expected questions like above to be asked, so was waiting for more input to see where we go).

[mravishankar]

Add an options object to allow side stepping the assert.

[mravishankar]

Rebase

[mravishankar]

Add some comments.

[nicolasvasilache]

mlir/include/mlir/Analysis/SliceAnalysis.h

42

Given the behavior change I am seeing here, I would rename as omitBlockArguments and rephrase as:

When omitBlockArguments is true, the backward slice computation omits traversing any block arguments.
When omitBlockArguments is false, the backward slice computation traverses block arguments and asserts that the parent op has a single region with a single block.

[mravishankar]

Address comments.

[mravishankar]

Rebase

mlir/include/mlir/Analysis/SliceAnalysis.h

 
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class BlockArgument;
 class Operation;
 class Value;
 
+struct SliceOptions {
   /// Type of the condition to limit the propagation of transitive use-defs.
-/// This can be used in particular to limit the propagation to a given Scope or
-/// to avoid passing through certain types of operation in a configurable
+  /// This can be used in particular to limit the propagation to a given Scope
+  /// or to avoid passing through certain types of operation in a configurable
   /// manner.
-using TransitiveFilter = llvm::function_ref<bool(Operation *)>;
+  using TransitiveFilter = std::function<bool(Operation *)>;
+  TransitiveFilter filter = nullptr;
+
+  /// Include the top level op in the slice.
+  bool inclusive = false;
+};
+
+struct BackwardSliceOptions : public SliceOptions {
+  /// When omitBlockArguments is true, the backward slice computation omits
+  /// traversing any block arguments. When omitBlockArguments is false, the
+  /// backward slice computation traverses block arguments and asserts that the
+  /// parent op has a single region with a single block.
+  bool omitBlockArguments = false;
+};
+
+using ForwardSliceOptions = SliceOptions;
 
 /// Fills `forwardSlice` with the computed forward slice (i.e. all
 /// the transitive uses of op), **without** including that operation.
 ///
 /// This additionally takes a TransitiveFilter which acts as a frontier:
 /// when looking at uses transitively, an operation that does not pass the
 /// filter is never propagated through. This allows in particular to carve out
 /// the scope within a ForOp or the scope within an IfOp.

 /// Assuming all local orders match the numbering order:
 /// 1. after getting back to the root getForwardSlice, `forwardSlice` may
 ///    contain:
 ///      {9, 7, 8, 5, 1, 2, 6, 3, 4}
 /// 2. reversing the result of 1. gives:
 ///      {4, 3, 6, 2, 1, 5, 8, 7, 9}
 ///
 void getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
-                     TransitiveFilter filter = nullptr /* pass-through*/,
-                     bool inclusive = false);
+                     ForwardSliceOptions options = {});
 
 /// Value-rooted version of `getForwardSlice`. Return the union of all forward
 /// slices for the uses of the value `root`.
 void getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
-                     TransitiveFilter filter = nullptr /* pass-through*/,
-                     bool inclusive = false);
+                     ForwardSliceOptions options = {});
 
 /// Fills `backwardSlice` with the computed backward slice (i.e.
 /// all the transitive defs of op), **without** including that operation.
 ///
 /// This additionally takes a TransitiveFilter which acts as a frontier:
 /// when looking at defs transitively, an operation that does not pass the
 /// filter is never propagated through. This allows in particular to carve out
 /// the scope within a ForOp or the scope within an IfOp.

 ///      |_______________|
 ///              |
 ///              9
 ///
 /// Assuming all local orders match the numbering order:
 ///    {1, 2, 5, 3, 4, 6}
 ///
 void getBackwardSlice(Operation *op, SetVector<Operation *> *backwardSlice,
-                      TransitiveFilter filter = nullptr /* pass-through*/,
-                      bool inclusive = false);
+                      BackwardSliceOptions options = {});
 
 /// Value-rooted version of `getBackwardSlice`. Return the union of all backward
 /// slices for the op defining or owning the value `root`.
 void getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice,
-                      TransitiveFilter filter = nullptr /* pass-through*/,
-                      bool inclusive = false);
+                      BackwardSliceOptions options = {});
 
 /// Iteratively computes backward slices and forward slices until
 /// a fixed point is reached. Returns an `SetVector<Operation *>` which
 /// **includes** the original operation.
 ///
 /// This allows building a slice (i.e. multi-root DAG where everything
 /// that is reachable from an Value in forward and backward direction is
 /// contained in the slice).

 ///   5. for any {x} |= defs,    uses |= uses(x); because all go through op
 ///   6. for any {x} |= uses, uses(x) |= uses;    because all go through op
 ///
 /// Intuitively, we should be able to recurse like:
 ///   preorder(defs) - op - postorder(uses)
 /// and keep things ordered but this is still hand-wavy and not worth the
 /// trouble for now: punt to a simple worklist-based solution.
 ///
-SetVector<Operation *>
-getSlice(Operation *op,
-         TransitiveFilter backwardFilter = nullptr /* pass-through*/,
-         TransitiveFilter forwardFilter = nullptr /* pass-through*/,
-         bool inclusive = false);
+SetVector<Operation *> getSlice(Operation *op,
+                                BackwardSliceOptions backwardSliceOptions = {},
+                                ForwardSliceOptions forwardSliceOptions = {});
 
 /// Multi-root DAG topological sort.
 /// Performs a topological sort of the Operation in the `toSort` SetVector.
 /// Returns a topologically sorted SetVector.
 SetVector<Operation *> topologicalSort(const SetVector<Operation *> &toSort);
 
 /// Utility to match a generic reduction given a list of iteration-carried
 /// arguments, `iterCarriedArgs` and the position of the potential reduction

mlir/lib/Analysis/SliceAnalysis.cpp

 #include "llvm/ADT/SmallPtrSet.h"
 
 ///
 /// Implements Analysis functions specific to slicing in Function.
 ///
 
 using namespace mlir;
 
-static void getForwardSliceImpl(Operation *op,
-                                SetVector<Operation *> *forwardSlice,
-                                TransitiveFilter filter) {
+static void
+getForwardSliceImpl(Operation *op, SetVector<Operation *> *forwardSlice,
+                    SliceOptions::TransitiveFilter filter = nullptr) {
   if (!op)
     return;
 
   // Evaluate whether we should keep this use.
   // This is useful in particular to implement scoping; i.e. return the
   // transitive forwardSlice in the current scope.
   if (filter && !filter(op))
     return;
 
   for (Region &region : op->getRegions())
     for (Block &block : region)
       for (Operation &blockOp : block)
         if (forwardSlice->count(&blockOp) == 0)
           getForwardSliceImpl(&blockOp, forwardSlice, filter);
   for (Value result : op->getResults()) {
     for (Operation *userOp : result.getUsers())
       if (forwardSlice->count(userOp) == 0)
         getForwardSliceImpl(userOp, forwardSlice, filter);
   }
 
   forwardSlice->insert(op);
 }
 
 void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
-                           TransitiveFilter filter, bool inclusive) {
-  getForwardSliceImpl(op, forwardSlice, filter);
-  if (!inclusive) {
+                           ForwardSliceOptions options) {
+  getForwardSliceImpl(op, forwardSlice, options.filter);
+  if (!options.inclusive) {
     // Don't insert the top level operation, we just queried on it and don't
     // want it in the results.
     forwardSlice->remove(op);
   }
 
   // Reverse to get back the actual topological order.
   // std::reverse does not work out of the box on SetVector and I want an
   // in-place swap based thing (the real std::reverse, not the LLVM adapter).
   std::vector<Operation *> v(forwardSlice->takeVector());
   forwardSlice->insert(v.rbegin(), v.rend());
 }
 
 void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
-                           TransitiveFilter filter, bool inclusive) {
+                           SliceOptions options) {
   for (Operation *user : root.getUsers())
-    getForwardSliceImpl(user, forwardSlice, filter);
+    getForwardSliceImpl(user, forwardSlice, options.filter);
 
   // Reverse to get back the actual topological order.
   // std::reverse does not work out of the box on SetVector and I want an
   // in-place swap based thing (the real std::reverse, not the LLVM adapter).
   std::vector<Operation *> v(forwardSlice->takeVector());
   forwardSlice->insert(v.rbegin(), v.rend());
 }
 
 static void getBackwardSliceImpl(Operation *op,
                                  SetVector<Operation *> *backwardSlice,
-                                 TransitiveFilter filter) {
+                                 BackwardSliceOptions options) {
   if (!op || op->hasTrait<OpTrait::IsIsolatedFromAbove>())
     return;
 
   // Evaluate whether we should keep this def.
   // This is useful in particular to implement scoping; i.e. return the
   // transitive backwardSlice in the current scope.
-  if (filter && !filter(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()) {
       if (backwardSlice->count(definingOp) == 0)
-        getBackwardSliceImpl(definingOp, backwardSlice, filter);
+        getBackwardSliceImpl(definingOp, backwardSlice, options);
     } else if (auto blockArg = dyn_cast<BlockArgument>(operand)) {
+      if (options.omitBlockArguments)
+        continue;
+
       Block *block = blockArg.getOwner();
       Operation *parentOp = block->getParentOp();
       // TODO: determine whether we want to recurse backward into the other
       // blocks of parentOp, which are not technically backward unless they flow
       // into us. For now, just bail.
       if (parentOp && backwardSlice->count(parentOp) == 0) {
         assert(parentOp->getNumRegions() == 1 &&
                parentOp->getRegion(0).getBlocks().size() == 1);
-        getBackwardSliceImpl(parentOp, backwardSlice, filter);
+        getBackwardSliceImpl(parentOp, backwardSlice, options);
       }
     } else {
       llvm_unreachable("No definingOp and not a block argument.");
     }
   }
 
   backwardSlice->insert(op);
 }
 
 void mlir::getBackwardSlice(Operation *op,
                             SetVector<Operation *> *backwardSlice,
-                            TransitiveFilter filter, bool inclusive) {
-  getBackwardSliceImpl(op, backwardSlice, filter);
+                            BackwardSliceOptions options) {
+  getBackwardSliceImpl(op, backwardSlice, options);
 
-  if (!inclusive) {
+  if (!options.inclusive) {
     // Don't insert the top level operation, we just queried on it and don't
     // want it in the results.
     backwardSlice->remove(op);
   }
 }
 
 void mlir::getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice,
-                            TransitiveFilter filter, bool inclusive) {
+                            BackwardSliceOptions options) {
   if (Operation *definingOp = root.getDefiningOp()) {
-    getBackwardSlice(definingOp, backwardSlice, filter, inclusive);
+    getBackwardSlice(definingOp, backwardSlice, options);
     return;
   }
   Operation *bbAargOwner = cast<BlockArgument>(root).getOwner()->getParentOp();
-  getBackwardSlice(bbAargOwner, backwardSlice, filter, inclusive);
+  getBackwardSlice(bbAargOwner, backwardSlice, options);
 }
 
 SetVector<Operation *> mlir::getSlice(Operation *op,
-                                      TransitiveFilter backwardFilter,
-                                      TransitiveFilter forwardFilter,
-                                      bool inclusive) {
+                                      BackwardSliceOptions backwardSliceOptions,
+                                      ForwardSliceOptions forwardSliceOptions) {
   SetVector<Operation *> slice;
   slice.insert(op);
 
   unsigned currentIndex = 0;
   SetVector<Operation *> backwardSlice;
   SetVector<Operation *> forwardSlice;
   while (currentIndex != slice.size()) {
     auto *currentOp = (slice)[currentIndex];
     // Compute and insert the backwardSlice starting from currentOp.
     backwardSlice.clear();
-    getBackwardSlice(currentOp, &backwardSlice, backwardFilter, inclusive);
+    getBackwardSlice(currentOp, &backwardSlice, backwardSliceOptions);
     slice.insert(backwardSlice.begin(), backwardSlice.end());
 
     // Compute and insert the forwardSlice starting from currentOp.
     forwardSlice.clear();
-    getForwardSlice(currentOp, &forwardSlice, forwardFilter, inclusive);
+    getForwardSlice(currentOp, &forwardSlice, forwardSliceOptions);
     slice.insert(forwardSlice.begin(), forwardSlice.end());
     ++currentIndex;
   }
   return topologicalSort(slice);
 }
 
 namespace {
 /// DFS post-order implementation that maintains a global count to work across

 
 /// Returns true if `value` (transitively) depends on iteration-carried values
 /// of the given `ancestorOp`.
 static bool dependsOnCarriedVals(Value value,
                                  ArrayRef<BlockArgument> iterCarriedArgs,
                                  Operation *ancestorOp) {
   // Compute the backward slice of the value.
   SetVector<Operation *> slice;
-  getBackwardSlice(value, &slice,
-                   [&](Operation *op) { return !ancestorOp->isAncestor(op); });
+  BackwardSliceOptions sliceOptions;
+  sliceOptions.filter = [&](Operation *op) {
+    return !ancestorOp->isAncestor(op);
+  };
+  getBackwardSlice(value, &slice, sliceOptions);
 
   // Check that none of the operands of the operations in the backward slice are
   // loop iteration arguments, and neither is the value itself.
   SmallPtrSet<Value, 8> iterCarriedValSet(iterCarriedArgs.begin(),
                                           iterCarriedArgs.end());
   if (iterCarriedValSet.contains(value))
     return true;
 

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

static bool supportsMMaMatrixType(Operation *op, bool useNvGpu) {
   if (auto unsignedExtend = dyn_cast<arith::ExtUIOp>(op))
     return integerExtendSupportsMMAMatrixType<arith::ExtUIOp>(unsignedExtend);
   return elementwiseSupportsMMAMatrixType(op);
 }
 
 /// Return an unsorted slice handling scf.for region differently than
 /// `getSlice`. In scf.for we only want to include as part of the slice elements
 /// that are part of the use/def chain.
-static SetVector<Operation *> getSliceContract(Operation *op,
-                                               TransitiveFilter backwardFilter,
-                                               TransitiveFilter forwardFilter) {
+static SetVector<Operation *>
+getSliceContract(Operation *op, BackwardSliceOptions backwardSliceOptions,
+                 ForwardSliceOptions forwardSliceOptions) {
   SetVector<Operation *> slice;
   slice.insert(op);
   unsigned currentIndex = 0;
   SetVector<Operation *> backwardSlice;
   SetVector<Operation *> forwardSlice;
   while (currentIndex != slice.size()) {
     auto *currentOp = (slice)[currentIndex];
     // Compute and insert the backwardSlice starting from currentOp.
     backwardSlice.clear();
-    getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
+    getBackwardSlice(currentOp, &backwardSlice, backwardSliceOptions);
     slice.insert(backwardSlice.begin(), backwardSlice.end());
 
     // Compute and insert the forwardSlice starting from currentOp.
     forwardSlice.clear();
     // Special case for ForOp, we don't want to include the whole region but
     // only the value using the region arguments.
     // TODO: We should refine this to only care about the region arguments being
     // converted to matrix type.
     if (auto forOp = dyn_cast<scf::ForOp>(currentOp)) {
       for (Value forOpResult : forOp.getResults())
-        getForwardSlice(forOpResult, &forwardSlice, forwardFilter);
+        getForwardSlice(forOpResult, &forwardSlice, forwardSliceOptions);
       for (BlockArgument &arg : forOp.getRegionIterArgs())
-        getForwardSlice(arg, &forwardSlice, forwardFilter);
+        getForwardSlice(arg, &forwardSlice, forwardSliceOptions);
     } else {
-      getForwardSlice(currentOp, &forwardSlice, forwardFilter);
+      getForwardSlice(currentOp, &forwardSlice, forwardSliceOptions);
     }
     slice.insert(forwardSlice.begin(), forwardSlice.end());
     ++currentIndex;
   }
   return slice;
 }
 
 // Analyze slice of operations based on convert op to figure out if the whole
 // slice can be converted to MMA operations.
 static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
                                              bool useNvGpu) {
   auto hasVectorDest = [](Operation *op) {
     return llvm::any_of(op->getResultTypes(),
                         [](Type t) { return isa<VectorType>(t); });
   };
+  BackwardSliceOptions backwardSliceOptions;
+  backwardSliceOptions.filter = hasVectorDest;
+
   auto hasVectorSrc = [](Operation *op) {
     return llvm::any_of(op->getOperandTypes(),
                         [](Type t) { return isa<VectorType>(t); });
   };
+  ForwardSliceOptions forwardSliceOptions;
+  forwardSliceOptions.filter = hasVectorSrc;
+
   SetVector<Operation *> opToConvert;
   op->walk([&](vector::ContractionOp contract) {
     if (opToConvert.contains(contract.getOperation()))
       return;
     SetVector<Operation *> dependentOps =
-        getSliceContract(contract, hasVectorDest, hasVectorSrc);
+        getSliceContract(contract, backwardSliceOptions, forwardSliceOptions);
     // If any instruction cannot use MMA matrix type drop the whole
     // chain. MMA matrix are stored in an opaque type so they cannot be used
     // by all operations.
     if (llvm::any_of(dependentOps, [useNvGpu](Operation *op) {
           return !supportsMMaMatrixType(op, useNvGpu);
         }))
       return;
     opToConvert.insert(dependentOps.begin(), dependentOps.end());

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

 // Get all the ops in the backwards slice starting from `padOp` and that
 // are dominated by the outermost enclosing loop.
 // This also requires tracking ops defining values used in the region but
 // defined above.
 static void computeBackwardSlice(tensor::PadOp padOp,
                                  scf::ForOp outermostEnclosingForOp,
                                  SetVector<Operation *> &backwardSlice) {
   DominanceInfo domInfo(outermostEnclosingForOp);
-  auto filter = [&](Operation *op) {
+  BackwardSliceOptions sliceOptions;
+  sliceOptions.filter = [&](Operation *op) {
     return domInfo.dominates(outermostEnclosingForOp, op) &&
            !padOp->isProperAncestor(op);
   };
+  sliceOptions.inclusive = true;
+
   // First, add the ops required to compute the region to the backwardSlice.
   SetVector<Value> valuesDefinedAbove;
   getUsedValuesDefinedAbove(padOp.getRegion(), padOp.getRegion(),
                             valuesDefinedAbove);
   for (Value v : valuesDefinedAbove) {
-    getBackwardSlice(v, &backwardSlice, filter, /*inclusive=*/true);
+    getBackwardSlice(v, &backwardSlice, sliceOptions);
   }
   // Then, add the backward slice from padOp itself.
-  getBackwardSlice(padOp.getOperation(), &backwardSlice, filter,
-                   /*inclusive=*/true);
+  getBackwardSlice(padOp.getOperation(), &backwardSlice, sliceOptions);
 }
 
 //===----------------------------------------------------------------------===//
 // HoistPaddingAnalysis Implementation.
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Analysis class to support tensor::PadOp hoisting across multiple enclosing

mlir/lib/Dialect/SCF/Utils/Utils.cpp

 
 // Hoist the ops within `outer` that appear before `inner`.
 // Such ops include the ops that have been introduced by parametric tiling.
 // Ops that come from triangular loops (i.e. that belong to the program slice
 // rooted at `outer`) and ops that have side effects cannot be hoisted.
 // Return failure when any op fails to hoist.
 static LogicalResult hoistOpsBetween(scf::ForOp outer, scf::ForOp inner) {
   SetVector<Operation *> forwardSlice;
-  getForwardSlice(
-      outer.getInductionVar(), &forwardSlice,
-      [&inner](Operation *op) { return op != inner.getOperation(); });
+  ForwardSliceOptions options;
+  options.filter = [&inner](Operation *op) {
+    return op != inner.getOperation();
+  };
+  getForwardSlice(outer.getInductionVar(), &forwardSlice, options);
   LogicalResult status = success();
   SmallVector<Operation *, 8> toHoist;
   for (auto &op : outer.getBody()->without_terminator()) {
     // Stop when encountering the inner loop.
     if (&op == inner.getOperation())
       break;
     // Skip over non-hoistable ops.
     if (forwardSlice.count(&op) > 0) {

mlir/test/IR/slice_multiple_blocks.mlir

+// RUN: mlir-opt --pass-pipeline="builtin.module(slice-analysis-test{omit-block-arguments=true})" %s | FileCheck %s
+
+func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %a = memref.alloc(%arg0, %arg2) : memref<?x?xf32>
+  %b = memref.alloc(%arg2, %arg1) : memref<?x?xf32>
+  cf.br ^bb1
+^bb1() :
+  %c = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
+  %d = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
+  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
+               outs(%c : memref<?x?xf32>)
+  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
+               outs(%d : memref<?x?xf32>)
+  memref.dealloc %c : memref<?x?xf32>
+  memref.dealloc %b : memref<?x?xf32>
+  memref.dealloc %a : memref<?x?xf32>
+  memref.dealloc %d : memref<?x?xf32>
+  return
+}
+// CHECK-LABEL: func @slicing_linalg_op__backward_slice__0
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
+//   CHECK-DAG:   %[[A:.+]] = memref.alloc(%[[ARG0]], %[[ARG2]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[B:.+]] = memref.alloc(%[[ARG2]], %[[ARG1]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[C:.+]] = memref.alloc(%[[ARG0]], %[[ARG1]]) : memref<?x?xf32>
+//       CHECK:   return
+
+// CHECK-LABEL: func @slicing_linalg_op__backward_slice__1
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
+//   CHECK-DAG:   %[[A:.+]] = memref.alloc(%[[ARG0]], %[[ARG2]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[B:.+]] = memref.alloc(%[[ARG2]], %[[ARG1]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[C:.+]] = memref.alloc(%[[ARG0]], %[[ARG1]]) : memref<?x?xf32>
+//       CHECK:   return

mlir/test/lib/IR/TestSlicing.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 
 /// Create a function with the same signature as the parent function of `op`
 /// with name being the function name and a `suffix`.
 static LogicalResult createBackwardSliceFunction(Operation *op,
-                                                 StringRef suffix) {
+                                                 StringRef suffix,
+                                                 bool omitBlockArguments) {
   func::FuncOp parentFuncOp = op->getParentOfType<func::FuncOp>();
   OpBuilder builder(parentFuncOp);
   Location loc = op->getLoc();
   std::string clonedFuncOpName = parentFuncOp.getName().str() + suffix.str();
   func::FuncOp clonedFuncOp = builder.create<func::FuncOp>(
       loc, clonedFuncOpName, parentFuncOp.getFunctionType());
   IRMapping mapper;
   builder.setInsertionPointToEnd(clonedFuncOp.addEntryBlock());
   for (const auto &arg : enumerate(parentFuncOp.getArguments()))
     mapper.map(arg.value(), clonedFuncOp.getArgument(arg.index()));
   SetVector<Operation *> slice;
-  getBackwardSlice(op, &slice);
+  BackwardSliceOptions options;
+  options.omitBlockArguments = omitBlockArguments;
+  getBackwardSlice(op, &slice, options);
   for (Operation *slicedOp : slice)
     builder.clone(*slicedOp, mapper);
   builder.create<func::ReturnOp>(loc);
   return success();
 }
 
 namespace {
 /// Pass to test slice generated from slice analysis.
 struct SliceAnalysisTestPass
     : public PassWrapper<SliceAnalysisTestPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SliceAnalysisTestPass)
 
   StringRef getArgument() const final { return "slice-analysis-test"; }
   StringRef getDescription() const final {
     return "Test Slice analysis functionality.";
   }
+
+  Option<bool> omitBlockArguments{
+      *this, "omit-block-arguments",
+      llvm::cl::desc("Test Slice analysis with multiple blocks but slice "
+                     "omiting block arguments"),
+      llvm::cl::init(true)};
+
   void runOnOperation() override;
   SliceAnalysisTestPass() = default;
   SliceAnalysisTestPass(const SliceAnalysisTestPass &) {}
 };
 } // namespace
 
 void SliceAnalysisTestPass::runOnOperation() {
   ModuleOp module = getOperation();
   auto funcOps = module.getOps<func::FuncOp>();
   unsigned opNum = 0;
   for (auto funcOp : funcOps) {
-    if (!llvm::hasSingleElement(funcOp.getBody())) {
-      funcOp->emitOpError("Does not support functions with multiple blocks");
-      signalPassFailure();
-      return;
-    }
     // TODO: For now this is just looking for Linalg ops. It can be generalized
     // to look for other ops using flags.
     funcOp.walk([&](Operation *op) {
       if (!isa<linalg::LinalgOp>(op))
         return WalkResult::advance();
       std::string append =
           std::string("__backward_slice__") + std::to_string(opNum);
-      (void)createBackwardSliceFunction(op, append);
+      (void)createBackwardSliceFunction(op, append, omitBlockArguments);
       opNum++;
       return WalkResult::advance();
     });
   }
 }
 
 namespace mlir {
 void registerSliceAnalysisTestPass() {
   PassRegistration<SliceAnalysisTestPass>();
 }
 } // namespace mlir