[mlir] GreedyPatternRewriter: Add ancestors to worklist

[springerm]

Summary

When adding an op to the worklist, also add its ancestors to the worklist. This allows for RewritePatterns to match an op a based on what is inside of the body of a.

This change fixes a problem that became apparent with vector.warp_execute_on_lane_0, but could probably be triggered with similar patterns. On a high level, the pattern extracts an op b with eligible = true from the body of an op a:

test.a {
  %0 = test.b() {eligible = true}
  yield %0
}

Afterwards:

%0 = test.b() {eligible = true}
test.a {
  yield %0
}

The pattern is an OpRewritePattern<OpA>. For some reason, test.a is not on the GreedyPatternRewriter's worklist. E.g., because no pattern could be applied and it was removed. Now, another pattern updates test.b, so that eligible is changed from false to true. The OpRewritePattern<OpA> could now be applied, but (without this revision) test.a is still not on the worklist.

Note: In the above example, an OpRewritePattern<OpB> could have been used instead of an OpRewritePattern<OpA>. With such a design, we can run into the same problem (when the eligible attr is on test.a and test.b is removed from the worklist because no patterns could be applied).

Note: This change uncovered an unrelated bug in TestSCFUtils.cpp that was triggered due to a change in the order in which ops are processed. A TODO is added to the broken code and test cases are adapted so that the bug is no longer triggered.

Depends On: D141365

[springerm]

Note: As an alternative to this change, it may also be possible to limit what users are allowed to do inside matchAndRewrite.

[ftynse]

Note: As an alternative to this change, it may also be possible to limit what users are allowed to do inside matchAndRewrite.

FWIW, we already have such limitations, e.g., don't bypass the rewriter. But since they are only English, many users just ignore them. I don't see an easy compiler-assisted way to enforce these: once the user gets ahold of an Operation *, they can do arbitrary mutation with the low-level API.

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

567–568

This isn't sound. You can easily end up adding parent operations that shouldn't be added, e.g. if you're in a function pass you shouldn't be adding the parent function/module/etc. This needs to gate on the upper limit of what can actually be considered.

[springerm]

address comments

[springerm]

@rriddle Can you take another look at this?

[rriddle]

Thanks for the ping, just coming out of vacation today. Looks reasonable to me, but curious about the cost of the check.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

313–322

Do you have any idea of the relative cost of this check? The pattern rewriter is fairly hot, just wondering if this has a visible impact on compile time.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

106

Can you remove this and make it an argument of the addToWorklist method? Seems to me this is only every local to the simplify method and adding "global" mutable state like does not seem ideal to me.
(also should be an ArrayRef I believe)

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

106

Actually I'm incorrect: this is called from many places, it may makes sense to have it here but please document the invariant a bit more:

/// This is set at the beginning of `simplify()` to the current scope the rewriter operates on.

I'm still bothered by this in the context of a class that:

1. has many public methods and simplify() isn't clearly the only entry point.
2. addToWorklist is a virtual method but now the default implementation is using a private field member here.

Seems like too many foot guns involved here...

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

313–322

We could also avoid repeating the tree-traversal of getParentOp hidden in findAncestorOpInRegion by having scope being a set and then do a single traversal but for each parent region querying the set.

That said it wouldn't be a win with a single region which may be the common case? (the code could handle it by branching on whether there is a single element in the set)

Also a nit: but Region::isAncestor(Region *other) seems a more suitable API than findAncestorOpInRegion here.

[springerm]

address comments

[springerm]

Added a dependent change. It is needed so that I could write the test case.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

106

I could prepare another cleanup that sets all fields already in the constructor. simplify would no longer take any parameters. Some of the fields can be const.

Then allow simplify to be called only once. This is how we use this class anyway. We create an instance, call simplify, then it is no longer used. Does that sound good?

313–322

I think it pays off even with just one region because we have now only a single while loop that traverses the ancestors. Before, it was two nested while loops. As for the loop structure, this is as efficient as it gets.

[mehdi_amini]

LGTM, but let's wait for River to have confirmation.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

324–328

Maybe slightly simplify and reduces the bracing.

[mehdi_amini]

In D140304#4039256, @mehdi_amini wrote:

LGTM, but let's wait for River to have confirmation.

Did you get confirmation from River?

[springerm]

In D140304#4050960, @mehdi_amini wrote:

In D140304#4039256, @mehdi_amini wrote:

LGTM, but let's wait for River to have confirmation.

Did you get confirmation from River?

There was no further activity for multiple days, so I assume this is fine by him.

[mehdi_amini]

That isn't really what I intended with my conditional approval, so please get explicit approval today from River ASAP.

[springerm]

In D140304#4050998, @mehdi_amini wrote:

That isn't really what I intended with my conditional approval, so please get explicit approval today from River ASAP.

@rriddle Is this good or should I revert?

[rriddle]

Looks reasonable to me! (Sorry I was OOO for a bit)

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
                                       const GreedyRewriteConfig &config);
 
   /// Simplify the operations within the given regions.
   bool simplify(MutableArrayRef<Region> regions);
 
-  /// Add the given operation to the worklist.
+  /// Add the given operation to the worklist. Parent ops may or may not be
+  /// added to the worklist, depending on the type of rewrite driver. By
+  /// default, parent ops are added.
   virtual void addToWorklist(Operation *op);
 
   /// Pop the next operation from the worklist.
   Operation *popFromWorklist();
 
   /// If the specified operation is in the worklist, remove it.
   void removeFromWorklist(Operation *op);
 
   /// Notifies the driver that the specified operation may have been modified
   /// in-place.
   void finalizeRootUpdate(Operation *op) override;
 
 protected:
+  /// Add the given operation to the worklist.
+  void addSingleOpToWorklist(Operation *op);
+
   // Implement the hook for inserting operations, and make sure that newly
   // inserted ops are added to the worklist for processing.
   void notifyOperationInserted(Operation *op) override;
 
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further
   // simplifications.
protected:
   /// if they aren't the root of a pattern.
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
 private:
+  /// Only ops within this scope are simplified.
+  MutableArrayRef<Region> scope;
+
   /// Configuration information for how to simplify.
   GreedyRewriteConfig config;
 
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
     : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 }
 
 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
+  scope = regions;
+
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
 
   /// A utility function to log a process result for the given reason.
   auto logResult = [&](StringRef result, const llvm::Twine &msg = {}) {
     logger.unindent();
     logger.startLine() << "} -> " << result;
#endif
   } while (changed && (iteration++ < config.maxIterations ||
                        config.maxIterations == GreedyRewriteConfig::kNoLimit));
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
 void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
+  while (op) {
+    bool inScope = llvm::any_of(scope, [&](Region &r) {
+      return r.findAncestorOpInRegion(*op) != nullptr;
+    });
+    if (!inScope)
+      break;
+
+    addSingleOpToWorklist(op);
+    op = op->getParentOp();
+  }
+}
+
+void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *op) {
   // Check to see if the worklist already contains this op.
   if (worklistMap.count(op))
     return;
 
   worklistMap[op] = worklist.size();
   worklist.push_back(op);
 }
 
 Operation *GreedyPatternRewriteDriver::popFromWorklist() {
   auto *op = worklist.back();
   worklist.pop_back();

 // MultiOpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// This is a specialized GreedyPatternRewriteDriver to apply patterns and
 /// perform folding for a supplied set of ops. It repeatedly simplifies while
 /// restricting the rewrites to only the provided set of ops or optionally
-/// to those directly affected by it (result users or operand providers).
+/// to those directly affected by it (result users or operand providers). Parent
+/// ops are not considered.
 class MultiOpPatternRewriteDriver : public GreedyPatternRewriteDriver {
 public:
   explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
                                        const FrozenRewritePatternSet &patterns,
                                        bool strict)
       : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
         strictMode(strict) {}
 
   bool simplifyLocally(ArrayRef<Operation *> op);
 
   void addToWorklist(Operation *op) override {
     if (!strictMode || strictModeFilteredOps.contains(op))
-      GreedyPatternRewriteDriver::addToWorklist(op);
+      GreedyPatternRewriteDriver::addSingleOpToWorklist(op);
   }
 
 private:
   void notifyOperationInserted(Operation *op) override {
     GreedyPatternRewriteDriver::notifyOperationInserted(op);
     if (strictMode)
       strictModeFilteredOps.insert(op);
   }
 

mlir/test/Dialect/Math/expand-math.mlir

 // CHECK: %[[EXP2:.+]] = math.exp %[[DOUBLEDX]] : f32
 // CHECK: %[[DIVIDEND2:.+]] = arith.subf %[[EXP2]], %[[ONE]] : f32
 // CHECK: %[[DIVISOR2:.+]] = arith.addf %[[EXP2]], %[[ONE]] : f32
 // CHECK: %[[RES2:.+]] = arith.divf %[[DIVIDEND2]], %[[DIVISOR2]] : f32
 // CHECK: %[[COND:.+]] = arith.cmpf oge, %arg0, %[[ZERO]] : f32
 // CHECK: %[[RESULT:.+]] = arith.select %[[COND]], %[[RES1]], %[[RES2]] : f32
 // CHECK: return %[[RESULT]]
 
-// ----
+// -----
 
 // CHECK-LABEL: func @ctlz
 func.func @ctlz(%arg: i32) -> i32 {
-  // CHECK: %[[C0:.+]] = arith.constant 0 : i32
-  // CHECK: %[[C32:.+]] = arith.constant 32 : i32
-  // CHECK: %[[C1:.+]] = arith.constant 1 : i32
+  // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i32
+  // CHECK-DAG: %[[C32:.+]] = arith.constant 32 : i32
+  // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i32
   // CHECK: %[[WHILE:.+]]:3 = scf.while (%[[A1:.+]] = %arg0, %[[A2:.+]] = %[[C32]], %[[A3:.+]] = %[[C0]])
   // CHECK:   %[[CMP:.+]] = arith.cmpi ne, %[[A1]], %[[A3]]
   // CHECK:   scf.condition(%[[CMP]]) %[[A1]], %[[A2]], %[[A3]]
   // CHECK:   %[[SHR:.+]] = arith.shrui %[[A1]], %[[C1]]
   // CHECK:   %[[SUB:.+]] = arith.subi %[[A2]], %[[C1]]
   // CHECK:   scf.yield %[[SHR]], %[[SUB]], %[[A3]]
   %res = math.ctlz %arg : i32
 
   // CHECK: return %[[WHILE]]#1
   return %res : i32
 }

mlir/test/Dialect/SCF/loop-pipelining.mlir

 
 // -----
 
 // CHECK-LABEL: backedge_different_stage
 //  CHECK-SAME:   (%[[A:.*]]: memref<?xf32>) -> f32 {
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
-//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 1.000000e+00 : f32
+//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 2.000000e+00 : f32
 // Prologue:
 //       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
 //  CHECK-NEXT:   %[[ADD0:.*]] = arith.addf %[[L0]], %[[CSTF]] : f32
 //  CHECK-NEXT:   %[[L1:.*]] = memref.load %[[A]][%[[C1]]] : memref<?xf32>
 // Kernel:
 //  CHECK-NEXT:   %[[R:.*]]:3 = scf.for %[[IV:.*]] = %[[C0]] to %[[C2]]
 //  CHECK-SAME:     step %[[C1]] iter_args(%[[C:.*]] = %[[CSTF]],
 //  CHECK-SAME:     %[[ADDARG:.*]] = %[[ADD0]], %[[LARG:.*]] = %[[L1]]) -> (f32, f32, f32) {
-//  CHECK-NEXT:     %[[ADD1:.*]] = arith.addf %[[LARG]], %[[ADDARG]] : f32
+//  CHECK-NEXT:     %[[MUL0:.*]] = arith.mulf %[[ADDARG]], %[[CSTF]] : f32
+//  CHECK-NEXT:     %[[ADD1:.*]] = arith.addf %[[LARG]], %[[MUL0]] : f32
 //  CHECK-NEXT:     %[[IV2:.*]] = arith.addi %[[IV]], %[[C2]] : index
 //  CHECK-NEXT:     %[[L2:.*]] = memref.load %[[A]][%[[IV2]]] : memref<?xf32>
-//  CHECK-NEXT:     scf.yield %[[ADDARG]], %[[ADD1]], %[[L2]] : f32, f32, f32
+//  CHECK-NEXT:     scf.yield %[[MUL0]], %[[ADD1]], %[[L2]] : f32, f32, f32
 //  CHECK-NEXT:   }
 // Epilogue:
-//  CHECK-NEXT:   %[[ADD2:.*]] = arith.addf %[[R]]#2, %[[R]]#1 : f32
-//  CHECK-NEXT:   return %[[ADD2]] : f32
+//  CHECK-NEXT:   %[[MUL1:.*]] = arith.mulf %[[R]]#1, %[[CSTF]] : f32
+//  CHECK-NEXT:   %[[ADD2:.*]] = arith.addf %[[R]]#2, %[[MUL1]] : f32
+//  CHECK-NEXT:   %[[MUL2:.*]] = arith.mulf %[[ADD2]], %[[CSTF]] : f32
+//  CHECK-NEXT:   return %[[MUL2]] : f32
 func.func @backedge_different_stage(%A: memref<?xf32>) -> f32 {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c4 = arith.constant 4 : index
-  %cf = arith.constant 1.0 : f32
+  %cf = arith.constant 2.0 : f32
   %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf) -> (f32) {
     %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
     %A1_elem = arith.addf %A_elem, %arg0 { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
     %A2_elem = arith.mulf %cf, %A1_elem { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 0 } : f32
     scf.yield %A2_elem : f32
   }  { __test_pipelining_loop__ }
   return %r : f32
 }
 
 // -----
 
 // CHECK-LABEL: region_backedge_different_stage
 //  CHECK-SAME:   (%[[A:.*]]: memref<?xf32>) -> f32 {
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
-//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 1.000000e+00 : f32
+//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 2.000000e+00 : f32
 // Prologue:
 //       CHECK:   %[[L0:.*]] = scf.execute_region
 //  CHECK-NEXT:     memref.load %[[A]][%[[C0]]] : memref<?xf32>
 //       CHECK:   %[[ADD0:.*]] = scf.execute_region
 //  CHECK-NEXT:   arith.addf %[[L0]], %[[CSTF]] : f32
 //       CHECK:   %[[L1:.*]] = scf.execute_region
 //  CHECK-NEXT:     memref.load %[[A]][%[[C1]]] : memref<?xf32>
 // Kernel:
 //       CHECK:   %[[R:.*]]:3 = scf.for %[[IV:.*]] = %[[C0]] to %[[C2]]
 //  CHECK-SAME:     step %[[C1]] iter_args(%[[C:.*]] = %[[CSTF]],
 //  CHECK-SAME:     %[[ADDARG:.*]] = %[[ADD0]], %[[LARG:.*]] = %[[L1]]) -> (f32, f32, f32) {
+//       CHECK:     %[[MUL0:.*]] = arith.mulf %[[ADDARG]], %[[CSTF]] : f32
 //       CHECK:     %[[ADD1:.*]] = scf.execute_region
-//  CHECK-NEXT:       arith.addf %[[LARG]], %[[ADDARG]] : f32
+//  CHECK-NEXT:       arith.addf %[[LARG]], %[[MUL0]] : f32
 //       CHECK:     %[[IV2:.*]] = arith.addi %[[IV]], %[[C2]] : index
 //       CHECK:     %[[L2:.*]] = scf.execute_region
 //  CHECK-NEXT:       memref.load %[[A]][%[[IV2]]] : memref<?xf32>
-//       CHECK:     scf.yield %[[ADDARG]], %[[ADD1]], %[[L2]] : f32, f32, f32
+//       CHECK:     scf.yield %[[MUL0]], %[[ADD1]], %[[L2]] : f32, f32, f32
 //  CHECK-NEXT:   }
 // Epilogue:
+//       CHECK:   %[[MUL1:.*]] = arith.mulf %[[R]]#1, %[[CSTF]] : f32
 //       CHECK:   %[[ADD2:.*]] = scf.execute_region
-//  CHECK-NEXT:    arith.addf %[[R]]#2, %[[R]]#1 : f32
-//       CHECK:   return %[[ADD2]] : f32
+//  CHECK-NEXT:    arith.addf %[[R]]#2, %[[MUL1]] : f32
+//       CHECK:   %[[MUL2:.*]] = arith.mulf %[[ADD2]], %[[CSTF]] : f32
+//       CHECK:   return %[[MUL2]] : f32
 
 func.func @region_backedge_different_stage(%A: memref<?xf32>) -> f32 {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c4 = arith.constant 4 : index
-  %cf = arith.constant 1.0 : f32
+  %cf = arith.constant 2.0 : f32
   %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf) -> (f32) {
     %A_elem = scf.execute_region -> f32 {
       %A_elem1 = memref.load %A[%i0] : memref<?xf32>
       scf.yield %A_elem1 : f32
     } { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 }
     %A1_elem = scf.execute_region -> f32 {
       %inner = arith.addf %A_elem, %arg0 : f32
       scf.yield %inner : f32
     }  { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 }
     %A2_elem = arith.mulf %cf, %A1_elem { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 0 } : f32
     scf.yield %A2_elem : f32
   }  { __test_pipelining_loop__ }
   return %r : f32
 }
 
 
 // -----
 
 // CHECK-LABEL: backedge_same_stage
 //  CHECK-SAME:   (%[[A:.*]]: memref<?xf32>) -> f32 {
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
-//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 1.000000e+00 : f32
+//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 2.000000e+00 : f32
 // Prologue:
 //       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
 // Kernel:
 //  CHECK-NEXT:   %[[R:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C3]]
 //  CHECK-SAME:     step %[[C1]] iter_args(%[[C:.*]] = %[[CSTF]],
 //  CHECK-SAME:     %[[LARG:.*]] = %[[L0]]) -> (f32, f32) {
 //  CHECK-NEXT:     %[[ADD0:.*]] = arith.addf %[[LARG]], %[[C]] : f32
+//  CHECK-NEXT:     %[[MUL0:.*]] = arith.mulf %[[ADD0]], %[[CSTF]] : f32
 //  CHECK-NEXT:     %[[IV1:.*]] = arith.addi %[[IV]], %[[C1]] : index
 //  CHECK-NEXT:     %[[L2:.*]] = memref.load %[[A]][%[[IV1]]] : memref<?xf32>
-//  CHECK-NEXT:     scf.yield %[[ADD0]], %[[L2]] : f32, f32
+//  CHECK-NEXT:     scf.yield %[[MUL0]], %[[L2]] : f32, f32
 //  CHECK-NEXT:   }
 // Epilogue:
 //  CHECK-NEXT:   %[[ADD1:.*]] = arith.addf %[[R]]#1, %[[R]]#0 : f32
-//  CHECK-NEXT:   return %[[ADD1]] : f32
+//  CHECK-NEXT:   %[[MUL1:.*]] = arith.mulf %[[ADD1]], %[[CSTF]] : f32
+//  CHECK-NEXT:   return %[[MUL1]] : f32
 func.func @backedge_same_stage(%A: memref<?xf32>) -> f32 {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c4 = arith.constant 4 : index
-  %cf = arith.constant 1.0 : f32
+  %cf = arith.constant 2.0 : f32
   %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf) -> (f32) {
     %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
     %A1_elem = arith.addf %A_elem, %arg0 { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
     %A2_elem = arith.mulf %cf, %A1_elem { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
     scf.yield %A2_elem : f32
   }  { __test_pipelining_loop__ }
   return %r : f32
 }
 
 // -----
 
-// CHECK: @pipeline_op_with_region(%[[ARG0:.+]]: memref<?xf32>, %[[ARG1:.+]]: memref<?xf32>, %[[ARG2:.+]]: memref<?xf32>) {
+// CHECK: @pipeline_op_with_region(%[[ARG0:.+]]: memref<?xf32>, %[[ARG1:.+]]: memref<?xf32>, %[[ARG2:.+]]: memref<?xf32>, %[[CF:.*]]: f32) {
 // CHECK:   %[[C0:.+]] = arith.constant 0 :
 // CHECK:   %[[C3:.+]] = arith.constant 3 :
 // CHECK:   %[[C1:.+]] = arith.constant 1 :
 // CHECK:   %[[APRO:.+]] = memref.alloc() :
 // CHECK:   %[[BPRO:.+]] = memref.alloc() :
 // CHECK:   %[[ASV0:.+]] = memref.subview %[[ARG0]][%[[C0]]] [8] [1] :
 // CHECK:   %[[BSV0:.+]] = memref.subview %[[ARG1]][%[[C0]]] [8] [1] :
 
indexing_maps = [
       #map1,
       #map2,
       #map1
     ],
   iterator_types = ["parallel"],
   __test_pipelining_stage__ = 1,
   __test_pipelining_op_order__ = 2
 }
-func.func @pipeline_op_with_region(%A: memref<?xf32>, %B: memref<?xf32>, %result: memref<?xf32>) {
+func.func @pipeline_op_with_region(%A: memref<?xf32>, %B: memref<?xf32>, %result: memref<?xf32>, %cf: f32) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c4 = arith.constant 4 : index
-  %cf = arith.constant 1.0 : f32
   %a_buf = memref.alloc() : memref<2x8xf32>
   %b_buf = memref.alloc() : memref<2x8xf32>
   scf.for %i0 = %c0 to %c4 step %c1 {
     %A_view = memref.subview %A[%i0][8][1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 3 } : memref<?xf32> to memref<8xf32, #map>
     %B_view = memref.subview %B[%i0][8][1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 4 } : memref<?xf32> to memref<8xf32, #map>
     %buf_idx = affine.apply  affine_map<(d0)->(d0 mod 2)> (%i0)[] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 5 }
     %a_buf_view = memref.subview %a_buf[%buf_idx,0][1,8][1,1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 6 } : memref<2x8xf32> to memref<8xf32, #map>
     %b_buf_view = memref.subview %b_buf[%buf_idx,0][1,8][1,1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 7 } : memref<2x8xf32> to memref<8xf32, #map>

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

getSchedule(scf::ForOp forOp,
 
     schedule.resize(forOp.getBody()->getOperations().size() - 1);
     forOp.walk([&schedule](Operation *op) {
       auto attrStage =
           op->getAttrOfType<IntegerAttr>(kTestPipeliningStageMarker);
       auto attrCycle =
           op->getAttrOfType<IntegerAttr>(kTestPipeliningOpOrderMarker);
       if (attrCycle && attrStage) {
+        // TODO: Index can be out-of-bounds if ops of the loop body disappear
+        // due to folding.
         schedule[attrCycle.getInt()] =
             std::make_pair(op, unsigned(attrStage.getInt()));
       }
     });
   }
 
   /// Helper to generate "predicated" version of `op`. For simplicity we just
   /// wrap the operation in a scf.ifOp operation.