Introduced iterative bytecode execution.

[sfuniak]

Summary

This is commit 2 of 4 for the multi-root matching in PDL, discussed in https://llvm.discourse.group/t/rfc-multi-root-pdl-patterns-for-kernel-matching/4148 (topic flagged for review).

This commit implements the features needed for the execution of the new operations pdl_interp.get_accepting_ops, pdl_interp.choose_op:

1. The implementation of the generation and execution of the two ops.
2. The addition of Stack of bytecode positions within the ByteCodeExecutor. This is needed because in pdl_interp.choose_op, we iterate over the values returned by pdl_interp.get_accepting_ops until we reach finalize. When we reach finalize, we need to return back to the position marked in the stack.
3. The functionality to extend the lifetime of values that cross the nondeterministic choice. The existing bytecode generator allocates the values to memory positions by representing the liveness of values as a collection of disjoint intervals over the matcher positions. This is akin to register allocation, and substantially reduces the footprint of the bytecode executor. However, because with iterative operation pdl_interp.choose_op, execution "returns" back, so any values whose original liveness cross the nondeterminstic choice must have their lifetime executed until finalize.

Testing: pdl-bytecode.mlir test

[sfuniak]

Made PDL ByteCode work with the new pdl_interp.foreach region.

This update migrates the ByteCode generator and executor to support the new ForEach construct in the pdl_interp dialect. I have reverted the pdl_interp.finalize operation to not perform any looping, and rely on pdl_interp.continue to advance the execution of the current loop to the next iterate.

[rriddle]

Nice!

mlir/lib/Rewrite/ByteCode.cpp

314–315

Same comment about debug fields as in the header.

359–361

mlir/lib/Rewrite/ByteCode.h

34

nit: Drop llvm:: for ArrayRef (it's re-exported in mlir namespace)

186–187

This is only for LLVM_DEBUG? Can you move this to the bottom and wrap with NDEBUG? We should keep anything debug specific sectioned off from everything else.

187

nit: Drop std:: from size_t

[rriddle]

mlir/lib/Rewrite/ByteCode.cpp

457–460

That is... weird. Can you file an llvm bug and reference it here?

586–601

612

1516–1522

llvm::reverse(*block)?

1516–1522

This feels expensive, I'm also not sure this is correct. If we move the insertion point to the operands, we may end up creating a side effecting operation incorrectly. For example, what happens in the situation where:

%cst = some.constant
%alloc = some.alloc
...
%load = some.load
... <Current Insertion Point>

If the current insertion point is the same above, if we were about to create a some.store %cst in %alloc operation wouldn't this incorrectly create it above the load (thus leading to a miscompile)?

1754–1756

Why not grab the users from all of the values?

1762–1765

1774–1775

Yeah, the opRangeMemory should hold the actual storage. memory should hold the address to the data.

[sfuniak]

mlir/lib/Rewrite/ByteCode.cpp

1516–1522

I can see how this change could be problematic. It was mostly introduced, because the current insertion point is at the root of the pattern, and with downward edge traversal, there is no guarantee that all the inputs are defined at this point. However, now that the pdl.rewrite accepts an optional argument forcing the root, this is less of an issue, because the user can manually specify the desired root that's known to work. Furthermore, in our use case, we rely on graph region where the order does not matter anyway. I will revert this change.

[Mogball]

General question: what's the cost if I never use multi-root patterns?

[sfuniak]

In D108547#3060427, @Mogball wrote:

General question: what's the cost if I never use multi-root patterns?

If you do not use multi-root patterns, the increase in cost of PDL-to-PDLInterp lowering is negligible (the system will scan through the pattern to determine the candidate roots, which there will be exactly 1). There is no increase in the cost of the bytecode execution, because there will be no pdl_interp.foreach operations in the resulting PDLInterp IR.

[sfuniak]

mlir/lib/Rewrite/ByteCode.cpp

314–315

Actually, blockToAddr is always needed to resolve the forward references. It's moved here from
void Generator::generate(ModuleOp module).
What you probably refer to is addrToBlock below. I could guard that one, but as I indicated in my earlier comment, it may be better to determine at runtime whether this map is to be populated, rather than compile time.

mlir/lib/Rewrite/ByteCode.h

186–187

I am a little apprehensive about wrapping this with NDEBUG in a public header. Doing so would break the code if the users of this header file are compiled with the different setting of the NDEBUG flag as the ByteCode.cpp. From llvm/lib/Support/Debug.cpp:

Even though LLVM might be built with NDEBUG, define symbols that the code
built without NDEBUG can depend on via the llvm/Support/Debug.h header.

Furthermore, eventually, I would like to have a debugger-like functionality in ByteCode, so that you can execute the operations one by one. So in the long run, we may want to control at runtime whether or not to populate this map.

[rriddle]

mlir/lib/Rewrite/ByteCode.h

186–187

It's not a public header though? This can't be included outside fo lib/Rewrite. I'd also like to avoid any kind of debugging facilities being baked into release builds, that type of functionality should be reserved for compilation setups that want it.

[sfuniak]

mlir/lib/Rewrite/ByteCode.h

186–187

I suppose you are right, it's not a public header. We access the bytecode only through FrozenRewritePatternSet. I will do as you suggested.

[sfuniak]

Addressed review feedback and made further improvements:

1. Fixed the insertion point on pdl_interp.create_operation.
2. Fixed the implementation of pdl_interp.get_users to match its semantics specified in stacked diff #1.
3. Eliminated the extra map storing the block for each instruction address. Instead, we store the line number of the instruction directly in the bytecode when debug messages are turned on.
4. pdl_interp.foreach no longer consumes the range being iterated over. Instead, we store the index into the range and update this index every time we call pdl_interp.continue.
5. PDLByteCodeMutableState::opRangeMemory now stores an owning reference to each range.
6. Improved unit tests coverage.

[sfuniak]

I am sorry, I am not sure why my recent changes to PDLInterps files are showing up here. It's as if the base commit was not set right. I did update the stack diff #1. Any ideas?

[sfuniak]

Attempt to fix a rebase problem.

[sfuniak]

mlir/lib/Rewrite/ByteCode.cpp

457–460

I figured this out: IntervalMap has a user-defined destructor but not a user-defined copy constructor. The compiler-generated implicitly-declared copy constructor is wrong. We should just define a move constructor, which is sufficient for most practical purposes. I will submit a diff next week and reference it here.

[sfuniak]

Updated comment with the IntervalMap fix.

[sfuniak]

@rriddle: This is the last remaining diff. If you can re-review it whenever you are free, that would be greatly appreciated. Then we can land all 4 and iterate. Thank you!

[sfuniak]

Gentle ping, @rriddle @mehdi_amini

[rriddle]

mlir/lib/Rewrite/ByteCode.cpp

598

nit: Drop the trivial braces here.

699–704

The debug related functionality feels separable, can you split this out?

932

Could probably use llvm::SaveAndRestore here.

1098

1537–1538

nit: Please move comments to their own line.

1654

We don't really need the reference.

1786

Walking the user count is going to be O(N), what's the trade off here vs. not reserving? Are we banking on a small number of users in practice?

1799

typo: acceessing

1800–1802

Why do we need all of this? Can we augment/add an overload of executeGetOperandsResults for the case we are interested in?

[sfuniak]

mlir/lib/Rewrite/ByteCode.cpp

699–704

will do

932

I find llvm::SaveAndRestore a bit too magical and no less verbose. It certainly helps in cases when we set a variable to an arbitrary value, but in this case, we are increasing and decreasing level, which I think is cleaner as is. We do not have exceptions to guard against either. But you are the code owner, so the decision ultimately rests with you.

1786

Not reserving will require several allocs, one alloc every time we expand the storage. I figured walking the users upfront would be cheaper (though still O(N)). But on the second thought, the number of users will often be small, so maybe SmallVector will do.

1800–1802

We need the comparison, because we need to check if the extracted operand(s) at the specified position match(es) the given value(s). I will create an overload of executeGetOperandsResults as you suggested that does this.

[sfuniak]

Split out the debug functionality & further review feedback.

[Mogball]

You mentioned that the cost (when running the bytecode) for pattern sets with no multi-root patterns should be negligible, but wouldn't pattern compilation slow down even for pattern sets without any multi-root patterns? (Because they all have to be scanned).

mlir/lib/Rewrite/ByteCode.cpp

1544

You could just assert(kind == PDLValue::Kind::Operation given that the only other case in this switch is an abort.

1798

If a null value is read, you could just exit early (and set range to empty).

1804

Same here. If there are no values to read (either null pointer or empty range) you can just return.

1817

The implementation of get_users in the bytecode is a little bit complex, stemming from the four cases of whether an index is specified cross whether the operand is a value or value range. Would this be any better with two different ops? One with an index and one without. It'll get rid of the sentinel value too.

[sfuniak]

There are two places where we do an extra traversal: 1) when verifying the PDL pattern (checking for connectivity), and 2) when forming the predicate tree (detecting the roots). Both are O(E), where E is the total number of operands across all the operations in the PDL pattern -- hardly an expensive operation. Just to be sure, I did a barebones test, where I took 100 copies of test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir concatenated together, and ran mlir-opt -split-input-file -convert-pdl-to-pdl-interp > /dev/null on the resulting file. The runtime of my version and the baseline version of mlir-opt were the same.

mlir/lib/Rewrite/ByteCode.cpp

1544

That's true for now, but we are going to soon follow this with another diff where the iteration type is a ValueRange. I thought I'd anticipate the change.

1817

That's a good point, I will create a new op code and split this up into two execute functions.

[sfuniak]

Split GetUsers into GetUsersAll and GetUsersAt.

[Mogball]

In D108547#3139714, @sfuniak wrote:

There are two places where we do an extra traversal: 1) when verifying the PDL pattern (checking for connectivity), and 2) when forming the predicate tree (detecting the roots). Both are O(E), where E is the total number of operands across all the operations in the PDL pattern -- hardly an expensive operation. Just to be sure, I did a barebones test, where I took 100 copies of test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir concatenated together, and ran mlir-opt -split-input-file -convert-pdl-to-pdl-interp > /dev/null on the resulting file. The runtime of my version and the baseline version of mlir-opt were the same.

Thanks for checking!

Just to be clear, this concern doesn't preclude any of the changes you've suggested. I'm just wondering whether explicitly marking a pattern as single-rooted or multi-rooted would be useful (should the performance have been a concern).

mlir/lib/Rewrite/ByteCode.cpp

1544

Gotcha

1798

Unrelated, but this makes me think we need an AttrSizedSegments interface that hides a lot of this work for us.

1826–1840

For this one here, I think you can extract the loop and move it below the if/else of Value vs ValueRange. These functions look much cleaner now, thanks!

[Mogball]

mlir/lib/Rewrite/ByteCode.cpp

905–906

[sfuniak]

Simplified GetUsers.

[Mogball]

Looking pretty good to me.

mlir/lib/Rewrite/ByteCode.cpp

1546

[rriddle]

mlir/lib/Rewrite/ByteCode.cpp

138

or the first value in a range

Is this right?

1754–1756

Unresolved?

mlir/lib/Rewrite/ByteCode.h

31

I think we shouldn't define OpRange like this, given that it has a different ownership model than TypeRange/ValueRange.

[sfuniak]

mlir/lib/Rewrite/ByteCode.cpp

1754–1756

Grabbing only the first value in a range is consistent with the definition of pdl_interp.get_defining_op.

I believe, the present definition is correct. You will match the entirety of a value range, so we could take the users of any value. We could also take the intersection of users, but that's more costly and not really needed, because we follow up the users query with operand comparison.

mlir/lib/Rewrite/ByteCode.h

31

Okay... can I just rename it to OwningOpRange then? I still think that using an owning reference below in std::vector<OpRange> opRangeMemory; is the right thing to do (no matter what we choose to call it).

[sfuniak]

Fixed the semantics of get_users for value range and implemented pdl_interp.get_value.

[rriddle]

LGTM after resolving the comments in the parent and pulling in the changes to here. I'd really like to get this in tree and iterate from there.

mlir/lib/Rewrite/ByteCode.h

31

Yeah, OwningOpRange is fine. Only had problems with the name, just want to avoid a situation where we would forget that OpRange has different semantics.

[sfuniak]

Implemented pdl_interp.extract.

[sfuniak]

Minor cleanups.

mlir/include/mlir/IR/PatternMatch.h

public:
   explicit operator bool() const { return value; }
 
   /// Return the kind of this value.
   Kind getKind() const { return kind; }
 
   /// Print this value to the provided output stream.
   void print(raw_ostream &os) const;
 
+  /// Print the specified value kind to an output stream.
+  static void print(raw_ostream &os, Kind kind);
+
 private:
   /// Find the index of a given type in a range of other types.
   template <typename...>
   struct index_of_t;
   template <typename T, typename... R>
   struct index_of_t<T, T, R...> : std::integral_constant<size_t, 0> {};
   template <typename T, typename F, typename... R>
   struct index_of_t<T, F, R...>
private:
   Kind kind;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, PDLValue value) {
   value.print(os);
   return os;
 }
 
+inline raw_ostream &operator<<(raw_ostream &os, PDLValue::Kind kind) {
+  PDLValue::print(os, kind);
+  return os;
+}
+
 //===----------------------------------------------------------------------===//
 // PDLResultList
 
 /// The class represents a list of PDL results, returned by a native rewrite
 /// method. It provides the mechanism with which to pass PDLValues back to the
 /// PDL bytecode.
 class PDLResultList {
 public:

mlir/lib/IR/PatternMatch.cpp

case Kind::Value:
     os << cast<Value>();
     break;
   case Kind::ValueRange:
     llvm::interleaveComma(cast<ValueRange>(), os);
     break;
   }
 }
 
+void PDLValue::print(raw_ostream &os, Kind kind) {
+  switch (kind) {
+  case Kind::Attribute:
+    os << "Attribute";
+    break;
+  case Kind::Operation:
+    os << "Operation";
+    break;
+  case Kind::Type:
+    os << "Type";
+    break;
+  case Kind::TypeRange:
+    os << "TypeRange";
+    break;
+  case Kind::Value:
+    os << "Value";
+    break;
+  case Kind::ValueRange:
+    os << "ValueRange";
+    break;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // PDLPatternModule
 //===----------------------------------------------------------------------===//
 
 void PDLPatternModule::mergeIn(PDLPatternModule &&other) {
   // Ignore the other module if it has no patterns.
   if (!other.pdlModule)
     return;

mlir/lib/Rewrite/ByteCode.h

 
 namespace detail {
 class PDLByteCode;
 
 /// Use generic bytecode types. ByteCodeField refers to the actual bytecode
 /// entries. ByteCodeAddr refers to size of indices into the bytecode.
 using ByteCodeField = uint16_t;
 using ByteCodeAddr = uint32_t;
+using OwningOpRange = llvm::OwningArrayRef<Operation *>;
 
 //===----------------------------------------------------------------------===//
 // PDLByteCodePattern
 //===----------------------------------------------------------------------===//
 
 /// All of the data pertaining to a specific pattern within the bytecode.
 class PDLByteCodePattern : public Pattern {
 public:
   static PDLByteCodePattern create(pdl_interp::RecordMatchOp matchOp,
                                    ByteCodeAddr rewriterAddr);
 
private:
   /// Allow access to data fields.
   friend class PDLByteCode;
 
   /// The mutable block of memory used during the matching and rewriting phases
   /// of the bytecode.
   std::vector<const void *> memory;
 
   /// A mutable block of memory used during the matching and rewriting phase of
+  /// the bytecode to store ranges of operations. These are always stored by
+  /// owning references, because at no point in the execution of the byte code
+  /// we get an indexed range (view) of operations.
+  std::vector<OwningOpRange> opRangeMemory;
+
+  /// A mutable block of memory used during the matching and rewriting phase of
   /// the bytecode to store ranges of types.
   std::vector<TypeRange> typeRangeMemory;
   /// A set of type ranges that have been allocated by the byte code interpreter
   /// to provide a guaranteed lifetime.
   std::vector<llvm::OwningArrayRef<Type>> allocatedTypeRangeMemory;
 
   /// A mutable block of memory used during the matching and rewriting phase of
   /// the bytecode to store ranges of values.
   std::vector<ValueRange> valueRangeMemory;
   /// A set of value ranges that have been allocated by the byte code
   /// interpreter to provide a guaranteed lifetime.
   std::vector<llvm::OwningArrayRef<Value>> allocatedValueRangeMemory;
 
+  /// The current index of ranges being iterated over for each level of nesting.
+  /// These are always maintained at 0 for the loops that are not active, so we
+  /// do not need to have a separate initialization phase for each loop.
+  std::vector<unsigned> loopIndex;
+
   /// The up-to-date benefits of the patterns held by the bytecode. The order
   /// of this array corresponds 1-1 with the array of patterns in `PDLByteCode`.
   std::vector<PatternBenefit> currentPatternBenefits;
 };
 
 //===----------------------------------------------------------------------===//
 // PDLByteCode
 //===----------------------------------------------------------------------===//
void executeByteCode(const ByteCodeField *inst, PatternRewriter &rewriter,
                        SmallVectorImpl<MatchResult> *matches) const;
 
   /// A vector containing pointers to uniqued data. The storage is intentionally
   /// opaque such that we can store a wide range of data types. The types of
   /// data stored here include:
   ///  * Attribute, OperationName, Type
   std::vector<const void *> uniquedData;
 
   /// A vector containing the generated bytecode for the matcher.
   SmallVector<ByteCodeField, 64> matcherByteCode;
 
   /// A vector containing the generated bytecode for all of the rewriters.
   SmallVector<ByteCodeField, 64> rewriterByteCode;
 
   /// The set of patterns contained within the bytecode.
   SmallVector<PDLByteCodePattern, 32> patterns;
 
   /// A set of user defined functions invoked via PDL.
   std::vector<PDLConstraintFunction> constraintFunctions;
   std::vector<PDLRewriteFunction> rewriteFunctions;
 
   /// The maximum memory index used by a value.
   ByteCodeField maxValueMemoryIndex = 0;
 
   /// The maximum number of different types of ranges.
+  ByteCodeField maxOpRangeCount = 0;
   ByteCodeField maxTypeRangeCount = 0;
   ByteCodeField maxValueRangeCount = 0;
+
+  /// The maximum number of nested loops.
+  ByteCodeField maxLoopLevel = 0;
 };
 
 } // end namespace detail
 } // end namespace mlir
 
 #endif // MLIR_REWRITE_BYTECODE_H_

mlir/lib/Rewrite/ByteCode.cpp

enum OpCode : ByteCodeField {
   /// Compare the operand count of an operation with a constant.
   CheckOperandCount,
   /// Compare the name of an operation with a constant.
   CheckOperationName,
   /// Compare the result count of an operation with a constant.
   CheckResultCount,
   /// Compare a range of types to a constant range of types.
   CheckTypes,
+  /// Continue to the next iteration of a loop.
+  Continue,
   /// Create an operation.
   CreateOperation,
   /// Create a range of types.
   CreateTypes,
   /// Erase an operation.
   EraseOp,
+  /// Extract the op from a range at the specified index.
+  ExtractOp,
+  /// Extract the type from a range at the specified index.
+  ExtractType,
+  /// Extract the value from a range at the specified index.
+  ExtractValue,
   /// Terminate a matcher or rewrite sequence.
   Finalize,
+  /// Iterate over a range of values.
+  ForEach,
   /// Get a specific attribute of an operation.
   GetAttribute,
   /// Get the type of an attribute.
   GetAttributeType,
   /// Get the defining operation of a value.
   GetDefiningOp,
   /// Get a specific operand of an operation.
   GetOperand0,
   GetOperand1,
   GetOperand2,
   GetOperand3,
   GetOperandN,
   /// Get a specific operand group of an operation.
   GetOperands,
   /// Get a specific result of an operation.
   GetResult0,
   GetResult1,
   GetResult2,
   GetResult3,
   GetResultN,
   /// Get a specific result group of an operation.
   GetResults,
+  /// Get the users of a value or a range of values.
+  GetUsers,
   /// Get the type of a value.
   GetValueType,
   /// Get the types of a value range.
   GetValueRangeTypes,
   /// Check if a generic value is not null.
   IsNotNull,
   /// Record a successful pattern match.
   RecordMatch,

 //===----------------------------------------------------------------------===//
 // ByteCode Generation
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Generator
 
 namespace {
+struct ByteCodeLiveRange;
 struct ByteCodeWriter;
 
+/// Check if the given class `T` can be converted to an opaque pointer.
+template <typename T, typename... Args>
+using has_pointer_traits = decltype(std::declval<T>().getAsOpaquePointer());
+
 /// This class represents the main generator for the pattern bytecode.
 class Generator {
 public:
   Generator(MLIRContext *ctx, std::vector<const void *> &uniquedData,
             SmallVectorImpl<ByteCodeField> &matcherByteCode,
             SmallVectorImpl<ByteCodeField> &rewriterByteCode,
             SmallVectorImpl<PDLByteCodePattern> &patterns,
             ByteCodeField &maxValueMemoryIndex,
+            ByteCodeField &maxOpRangeMemoryIndex,
             ByteCodeField &maxTypeRangeMemoryIndex,
             ByteCodeField &maxValueRangeMemoryIndex,
+            ByteCodeField &maxLoopLevel,
             llvm::StringMap<PDLConstraintFunction> &constraintFns,
             llvm::StringMap<PDLRewriteFunction> &rewriteFns)
       : ctx(ctx), uniquedData(uniquedData), matcherByteCode(matcherByteCode),
         rewriterByteCode(rewriterByteCode), patterns(patterns),
         maxValueMemoryIndex(maxValueMemoryIndex),
+        maxOpRangeMemoryIndex(maxOpRangeMemoryIndex),
         maxTypeRangeMemoryIndex(maxTypeRangeMemoryIndex),
-        maxValueRangeMemoryIndex(maxValueRangeMemoryIndex) {
+        maxValueRangeMemoryIndex(maxValueRangeMemoryIndex),
+        maxLoopLevel(maxLoopLevel) {
     for (auto it : llvm::enumerate(constraintFns))
       constraintToMemIndex.try_emplace(it.value().first(), it.index());
     for (auto it : llvm::enumerate(rewriteFns))
       externalRewriterToMemIndex.try_emplace(it.value().first(), it.index());
   }
 
   /// Generate the bytecode for the given PDL interpreter module.
   void generate(ModuleOp module);
public:
   }
 
 private:
   /// Allocate memory indices for the results of operations within the matcher
   /// and rewriters.
   void allocateMemoryIndices(FuncOp matcherFunc, ModuleOp rewriterModule);
 
   /// Generate the bytecode for the given operation.
+  void generate(Region *region, ByteCodeWriter &writer);
   void generate(Operation *op, ByteCodeWriter &writer);
   void generate(pdl_interp::ApplyConstraintOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::ApplyRewriteOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::AreEqualOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::BranchOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CheckAttributeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CheckOperandCountOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CheckOperationNameOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CheckResultCountOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CheckTypeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CheckTypesOp op, ByteCodeWriter &writer);
+  void generate(pdl_interp::ContinueOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CreateAttributeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CreateOperationOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CreateTypeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::CreateTypesOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::EraseOp op, ByteCodeWriter &writer);
+  void generate(pdl_interp::ExtractOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::FinalizeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetAttributeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetAttributeTypeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetDefiningOpOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetOperandOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetOperandsOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetResultOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetResultsOp op, ByteCodeWriter &writer);
+  void generate(pdl_interp::GetUsersOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::GetValueTypeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::InferredTypesOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::IsNotNullOp op, ByteCodeWriter &writer);
+  void generate(pdl_interp::ForEachOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::RecordMatchOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::ReplaceOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::SwitchAttributeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::SwitchTypeOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::SwitchTypesOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::SwitchOperandCountOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::SwitchOperationNameOp op, ByteCodeWriter &writer);
   void generate(pdl_interp::SwitchResultCountOp op, ByteCodeWriter &writer);
private:
   /// Mapping from rewriter function name to the bytecode address of the
   /// rewriter function in byte.
   llvm::StringMap<ByteCodeAddr> rewriterToAddr;
 
   /// Mapping from a uniqued storage object to its memory index within
   /// `uniquedData`.
   DenseMap<const void *, ByteCodeField> uniquedDataToMemIndex;
 
+  /// The current level of the foreach loop.
+  ByteCodeField curLoopLevel = 0;
+
   /// The current MLIR context.
   MLIRContext *ctx;
 
+  /// Mapping from block to its address.
+  DenseMap<Block *, ByteCodeAddr> blockToAddr;
+
   /// Data of the ByteCode class to be populated.
   std::vector<const void *> &uniquedData;
   SmallVectorImpl<ByteCodeField> &matcherByteCode;
   SmallVectorImpl<ByteCodeField> &rewriterByteCode;
   SmallVectorImpl<PDLByteCodePattern> &patterns;
   ByteCodeField &maxValueMemoryIndex;
+  ByteCodeField &maxOpRangeMemoryIndex;
   ByteCodeField &maxTypeRangeMemoryIndex;
   ByteCodeField &maxValueRangeMemoryIndex;
+  ByteCodeField &maxLoopLevel;
 };
 
 /// This class provides utilities for writing a bytecode stream.
 struct ByteCodeWriter {
   ByteCodeWriter(SmallVectorImpl<ByteCodeField> &bytecode, Generator &generator)
       : bytecode(bytecode), generator(generator) {}
 
   /// Append a field to the bytecode.
   void append(ByteCodeField field) { bytecode.push_back(field); }
   void append(OpCode opCode) { bytecode.push_back(opCode); }
 
   /// Append an address to the bytecode.
   void append(ByteCodeAddr field) {
     static_assert((sizeof(ByteCodeAddr) / sizeof(ByteCodeField)) == 2,
                   "unexpected ByteCode address size");
 
     ByteCodeField fieldParts[2];
     std::memcpy(fieldParts, &field, sizeof(ByteCodeAddr));
     bytecode.append({fieldParts[0], fieldParts[1]});
   }
 
-  /// Append a successor range to the bytecode, the exact address will need to
+  /// Append a single successor to the bytecode, the exact address will need to
   /// be resolved later.
-  void append(SuccessorRange successors) {
-    // Add back references to the any successors so that the address can be
-    // resolved later.
-    for (Block *successor : successors) {
-      unresolvedSuccessorRefs[successor].push_back(bytecode.size());
-      append(ByteCodeAddr(0));
-    }
+  void append(Block *successor) {
+    // Add back a reference to the successor so that the address can be resolved
+    // later.
+
+  /// Append a successor range to the bytecode, the exact address will need to
+  /// be resolved later.
+  void append(SuccessorRange successors) {
+    for (Block *successor : successors)
+      append(successor);
   }
 
   /// Append a range of values that will be read as generic PDLValues.
   void appendPDLValueList(OperandRange values) {
     bytecode.push_back(values.size());
     for (Value value : values)
       appendPDLValue(value);
   }
 
   /// Append a value as a PDLValue.
   void appendPDLValue(Value value) {
     appendPDLValueKind(value);
     append(value);
   }
 
   /// Append the PDLValue::Kind of the given value.
-  void appendPDLValueKind(Value value) {
-    // Append the type of the value in addition to the value itself.
+  void appendPDLValueKind(Value value) { appendPDLValueKind(value.getType()); }
+
+  /// Append the PDLValue::Kind of the given type.
+  void appendPDLValueKind(Type type) {
     PDLValue::Kind kind =
-        TypeSwitch<Type, PDLValue::Kind>(value.getType())
+        TypeSwitch<Type, PDLValue::Kind>(type)
             .Case<pdl::AttributeType>(
                 [](Type) { return PDLValue::Kind::Attribute; })
             .Case<pdl::OperationType>(
                 [](Type) { return PDLValue::Kind::Operation; })
             .Case<pdl::RangeType>([](pdl::RangeType rangeTy) {
               if (rangeTy.getElementType().isa<pdl::TypeType>())
                 return PDLValue::Kind::TypeRange;
               return PDLValue::Kind::ValueRange;
             })
             .Case<pdl::TypeType>([](Type) { return PDLValue::Kind::Type; })
             .Case<pdl::ValueType>([](Type) { return PDLValue::Kind::Value; });
     bytecode.push_back(static_cast<ByteCodeField>(kind));
   }
 
-  /// Check if the given class `T` has an iterator type.
-  template <typename T, typename... Args>
-  using has_pointer_traits = decltype(std::declval<T>().getAsOpaquePointer());
-
   /// Append a value that will be stored in a memory slot and not inline within
   /// the bytecode.
   template <typename T>
   std::enable_if_t<llvm::is_detected<has_pointer_traits, T>::value ||
                    std::is_pointer<T>::value>
   append(T value) {
     bytecode.push_back(generator.getMemIndex(value));
   }
struct ByteCodeWriter {
 
   /// The main generator producing PDL.
   Generator &generator;
 };
 
 /// This class represents a live range of PDL Interpreter values, containing
 /// information about when values are live within a match/rewrite.
 struct ByteCodeLiveRange {
-  using Set = llvm::IntervalMap<ByteCodeField, char, 16>;
+  using Set = llvm::IntervalMap<uint64_t, char, 16>;
   using Allocator = Set::Allocator;
 
-  ByteCodeLiveRange(Allocator &alloc) : liveness(alloc) {}
+  ByteCodeLiveRange(Allocator &alloc) : liveness(new Set(alloc)) {}
 
   /// Union this live range with the one provided.
   void unionWith(const ByteCodeLiveRange &rhs) {
-    for (auto it = rhs.liveness.begin(), e = rhs.liveness.end(); it != e; ++it)
-      liveness.insert(it.start(), it.stop(), /*dummyValue*/ 0);
+    for (auto it = rhs.liveness->begin(), e = rhs.liveness->end(); it != e;
+         ++it)
+      liveness->insert(it.start(), it.stop(), /*dummyValue*/ 0);
   }
 
   /// Returns true if this range overlaps with the one provided.
   bool overlaps(const ByteCodeLiveRange &rhs) const {
-    return llvm::IntervalMapOverlaps<Set, Set>(liveness, rhs.liveness).valid();
+    return llvm::IntervalMapOverlaps<Set, Set>(*liveness, *rhs.liveness)
+        .valid();
   }
 
   /// A map representing the ranges of the match/rewrite that a value is live in
   /// the interpreter.
-  llvm::IntervalMap<ByteCodeField, char, 16> liveness;
+  ///
+  /// We use std::unique_ptr here, because IntervalMap does not provide a
+  /// correct copy or move constructor. We can eliminate the pointer once
+  /// https://reviews.llvm.org/D113240 lands.
+  std::unique_ptr<llvm::IntervalMap<uint64_t, char, 16>> liveness;
+
+  /// The operation range storage index for this range.
+  Optional<unsigned> opRangeIndex;
 
   /// The type range storage index for this range.
   Optional<unsigned> typeRangeIndex;
 
   /// The value range storage index for this range.
   Optional<unsigned> valueRangeIndex;
 };
 } // end anonymous namespace
for (FuncOp rewriterFunc : rewriterModule.getOps<FuncOp>()) {
     rewriterToAddr.try_emplace(rewriterFunc.getName(), rewriterByteCode.size());
     for (Operation &op : rewriterFunc.getOps())
       generate(&op, rewriterByteCodeWriter);
   }
   assert(rewriterByteCodeWriter.unresolvedSuccessorRefs.empty() &&
          "unexpected branches in rewriter function");
 
   // Generate code for the matcher function.
-  DenseMap<Block *, ByteCodeAddr> blockToAddr;
-  llvm::ReversePostOrderTraversal<Region *> rpot(&matcherFunc.getBody());
   ByteCodeWriter matcherByteCodeWriter(matcherByteCode, *this);
-  for (Block *block : rpot) {
-    // Keep track of where this block begins within the matcher function.
-    blockToAddr.try_emplace(block, matcherByteCode.size());
-    for (Operation &op : *block)
-      generate(&op, matcherByteCodeWriter);
-  }
+  generate(&matcherFunc.getBody(), matcherByteCodeWriter);
 
   // Resolve successor references in the matcher.
   for (auto &it : matcherByteCodeWriter.unresolvedSuccessorRefs) {
     ByteCodeAddr addr = blockToAddr[it.first];
     for (unsigned offsetToFix : it.second)
       std::memcpy(&matcherByteCode[offsetToFix], &addr, sizeof(ByteCodeAddr));
   }
 }
void Generator::allocateMemoryIndices(FuncOp matcherFunc,
   }
 
   // The matcher function uses a more sophisticated numbering that tries to
   // minimize the number of memory indices assigned. This is done by determining
   // a live range of the values within the matcher, then the allocation is just
   // finding the minimal number of overlapping live ranges. This is essentially
   // a simplified form of register allocation where we don't necessarily have a
   // limited number of registers, but we still want to minimize the number used.
-  DenseMap<Operation *, ByteCodeField> opToIndex;
+  DenseMap<Operation *, unsigned> opToIndex;
   matcherFunc.getBody().walk([&](Operation *op) {
     opToIndex.insert(std::make_pair(op, opToIndex.size()));
   });
 
   // Liveness info for each of the defs within the matcher.
   ByteCodeLiveRange::Allocator allocator;
   DenseMap<Value, ByteCodeLiveRange> valueDefRanges;
 
   // Assign the root operation being matched to slot 0.
   BlockArgument rootOpArg = matcherFunc.getArgument(0);
   valueToMemIndex[rootOpArg] = 0;
 
   // Walk each of the blocks, computing the def interval that the value is used.
   Liveness matcherLiveness(matcherFunc);
-  for (Block &block : matcherFunc.getBody()) {
-    const LivenessBlockInfo *info = matcherLiveness.getLiveness(&block);
+  matcherFunc->walk([&](Block *block) {
+    const LivenessBlockInfo *info = matcherLiveness.getLiveness(block);
     assert(info && "expected liveness info for block");
     auto processValue = [&](Value value, Operation *firstUseOrDef) {
       // We don't need to process the root op argument, this value is always
       // assigned to the first memory slot.
       if (value == rootOpArg)
         return;
 
       // Set indices for the range of this block that the value is used.
       auto defRangeIt = valueDefRanges.try_emplace(value, allocator).first;
-      defRangeIt->second.liveness.insert(
+      defRangeIt->second.liveness->insert(
           opToIndex[firstUseOrDef],
           opToIndex[info->getEndOperation(value, firstUseOrDef)],
           /*dummyValue*/ 0);
 
       // Check to see if this value is a range type.
       if (auto rangeTy = value.getType().dyn_cast<pdl::RangeType>()) {
         Type eleType = rangeTy.getElementType();
-        if (eleType.isa<pdl::TypeType>())
+        if (eleType.isa<pdl::OperationType>())
+          defRangeIt->second.opRangeIndex = 0;
+        else if (eleType.isa<pdl::TypeType>())
           defRangeIt->second.typeRangeIndex = 0;
         else if (eleType.isa<pdl::ValueType>())
           defRangeIt->second.valueRangeIndex = 0;
       }
     };
 
     // Process the live-ins of this block.
-    for (Value liveIn : info->in())
-      processValue(liveIn, &block.front());
+    for (Value liveIn : info->in()) {
+      // Only process the value if it has been defined in the current region.
+      // Other values that span across pdl_interp.foreach will be added higher
+      // up. This ensures that the we keep them alive for the entire duration
+      // of the loop.
+      if (liveIn.getParentRegion() == block->getParent())
+        processValue(liveIn, &block->front());
+    }
+
+    // Process the block arguments for the entry block (those are not live-in).
+    if (block->isEntryBlock()) {
+      for (Value argument : block->getArguments())
+        processValue(argument, &block->front());
+    }
 
     // Process any new defs within this block.
-    for (Operation &op : block)
+    for (Operation &op : *block)
       for (Value result : op.getResults())
         processValue(result, &op);
-  }
+  });
 
   // Greedily allocate memory slots using the computed def live ranges.
   std::vector<ByteCodeLiveRange> allocatedIndices;
-  ByteCodeField numIndices = 1, numTypeRanges = 0, numValueRanges = 0;
+
+  // The number of memory indices currently allocated (and its next value).
+  // Recall that the root gets allocated memory index 0.
+  ByteCodeField numIndices = 1;
+
+  // The number of memory ranges of various types (and their next values).
+  ByteCodeField numOpRanges = 0, numTypeRanges = 0, numValueRanges = 0;
+
   for (auto &defIt : valueDefRanges) {
     ByteCodeField &memIndex = valueToMemIndex[defIt.first];
     ByteCodeLiveRange &defRange = defIt.second;
 
     // Try to allocate to an existing index.
     for (auto existingIndexIt : llvm::enumerate(allocatedIndices)) {
       ByteCodeLiveRange &existingRange = existingIndexIt.value();
       if (!defRange.overlaps(existingRange)) {
         existingRange.unionWith(defRange);
         memIndex = existingIndexIt.index() + 1;
 
-        if (defRange.typeRangeIndex) {
+        if (defRange.opRangeIndex) {
+          if (!existingRange.opRangeIndex)
+            existingRange.opRangeIndex = numOpRanges++;
+          valueToRangeIndex[defIt.first] = *existingRange.opRangeIndex;
+        } else if (defRange.typeRangeIndex) {
           if (!existingRange.typeRangeIndex)
             existingRange.typeRangeIndex = numTypeRanges++;
           valueToRangeIndex[defIt.first] = *existingRange.typeRangeIndex;
         } else if (defRange.valueRangeIndex) {
           if (!existingRange.valueRangeIndex)
             existingRange.valueRangeIndex = numValueRanges++;
           valueToRangeIndex[defIt.first] = *existingRange.valueRangeIndex;
         }
         break;
       }
     }
 
     // If no existing index could be used, add a new one.
     if (memIndex == 0) {
       allocatedIndices.emplace_back(allocator);
       ByteCodeLiveRange &newRange = allocatedIndices.back();
       newRange.unionWith(defRange);
 
-      // Allocate an index for type/value ranges.
-      if (defRange.typeRangeIndex) {
+      // Allocate an index for op/type/value ranges.
+      if (defRange.opRangeIndex) {
+        newRange.opRangeIndex = numOpRanges;
+        valueToRangeIndex[defIt.first] = numOpRanges++;
+      } else if (defRange.typeRangeIndex) {
         newRange.typeRangeIndex = numTypeRanges;
         valueToRangeIndex[defIt.first] = numTypeRanges++;
       } else if (defRange.valueRangeIndex) {
         newRange.valueRangeIndex = numValueRanges;
         valueToRangeIndex[defIt.first] = numValueRanges++;
       }
 
       memIndex = allocatedIndices.size();
       ++numIndices;
     }
   }
 
+  // Print the index usage and ensure that we did not run out of index space.
+  LLVM_DEBUG({
+    llvm::dbgs() << "Allocated " << allocatedIndices.size() << " indices "
+                 << "(down from initial " << valueDefRanges.size() << ").\n";
+  });
+  assert(allocatedIndices.size() <= std::numeric_limits<ByteCodeField>::max() &&
+         "Ran out of memory for allocated indices");
+
   // Update the max number of indices.
   if (numIndices > maxValueMemoryIndex)
     maxValueMemoryIndex = numIndices;
+  if (numOpRanges > maxOpRangeMemoryIndex)
+    maxOpRangeMemoryIndex = numOpRanges;
   if (numTypeRanges > maxTypeRangeMemoryIndex)
     maxTypeRangeMemoryIndex = numTypeRanges;
   if (numValueRanges > maxValueRangeMemoryIndex)
     maxValueRangeMemoryIndex = numValueRanges;
 }
 
+void Generator::generate(Region *region, ByteCodeWriter &writer) {
+  llvm::ReversePostOrderTraversal<Region *> rpot(region);
+  for (Block *block : rpot) {
+    // Keep track of where this block begins within the matcher function.
+    blockToAddr.try_emplace(block, matcherByteCode.size());
+    for (Operation &op : *block)
+      generate(&op, writer);
+  }
+}
+
 void Generator::generate(Operation *op, ByteCodeWriter &writer) {
   TypeSwitch<Operation *>(op)
       .Case<pdl_interp::ApplyConstraintOp, pdl_interp::ApplyRewriteOp,
             pdl_interp::AreEqualOp, pdl_interp::BranchOp,
             pdl_interp::CheckAttributeOp, pdl_interp::CheckOperandCountOp,
             pdl_interp::CheckOperationNameOp, pdl_interp::CheckResultCountOp,
             pdl_interp::CheckTypeOp, pdl_interp::CheckTypesOp,
-            pdl_interp::CreateAttributeOp, pdl_interp::CreateOperationOp,
-            pdl_interp::CreateTypeOp, pdl_interp::CreateTypesOp,
-            pdl_interp::EraseOp, pdl_interp::FinalizeOp,
-            pdl_interp::GetAttributeOp, pdl_interp::GetAttributeTypeOp,
-            pdl_interp::GetDefiningOpOp, pdl_interp::GetOperandOp,
-            pdl_interp::GetOperandsOp, pdl_interp::GetResultOp,
-            pdl_interp::GetResultsOp, pdl_interp::GetValueTypeOp,
+            pdl_interp::ContinueOp, pdl_interp::CreateAttributeOp,
+            pdl_interp::CreateOperationOp, pdl_interp::CreateTypeOp,
+            pdl_interp::CreateTypesOp, pdl_interp::EraseOp,
+            pdl_interp::ExtractOp, pdl_interp::FinalizeOp,
+            pdl_interp::ForEachOp, pdl_interp::GetAttributeOp,
+            pdl_interp::GetAttributeTypeOp, pdl_interp::GetDefiningOpOp,
+            pdl_interp::GetOperandOp, pdl_interp::GetOperandsOp,
+            pdl_interp::GetResultOp, pdl_interp::GetResultsOp,
+            pdl_interp::GetUsersOp, pdl_interp::GetValueTypeOp,
             pdl_interp::InferredTypesOp, pdl_interp::IsNotNullOp,
             pdl_interp::RecordMatchOp, pdl_interp::ReplaceOp,
             pdl_interp::SwitchAttributeOp, pdl_interp::SwitchTypeOp,
             pdl_interp::SwitchTypesOp, pdl_interp::SwitchOperandCountOp,
             pdl_interp::SwitchOperationNameOp, pdl_interp::SwitchResultCountOp>(
           [&](auto interpOp) { this->generate(interpOp, writer); })
       .Default([](Operation *) {
         llvm_unreachable("unknown `pdl_interp` operation");
writer.append(OpCode::CheckResultCount, op.operation(), op.count(),
                 op.getSuccessors());
 }
 void Generator::generate(pdl_interp::CheckTypeOp op, ByteCodeWriter &writer) {
   writer.append(OpCode::AreEqual, op.value(), op.type(), op.getSuccessors());
 }
 void Generator::generate(pdl_interp::CheckTypesOp op, ByteCodeWriter &writer) {
   writer.append(OpCode::CheckTypes, op.value(), op.types(), op.getSuccessors());
 }
+void Generator::generate(pdl_interp::ContinueOp op, ByteCodeWriter &writer) {
+  assert(curLoopLevel > 0 && "encountered pdl_interp.continue at top level");
+  writer.append(OpCode::Continue, ByteCodeField(curLoopLevel - 1));
+}
 void Generator::generate(pdl_interp::CreateAttributeOp op,
                          ByteCodeWriter &writer) {
   // Simply repoint the memory index of the result to the constant.
   getMemIndex(op.attribute()) = getMemIndex(op.value());
 }
 void Generator::generate(pdl_interp::CreateOperationOp op,
                          ByteCodeWriter &writer) {
   writer.append(OpCode::CreateOperation, op.operation(),

 }
 void Generator::generate(pdl_interp::CreateTypesOp op, ByteCodeWriter &writer) {
   writer.append(OpCode::CreateTypes, op.result(),
                 getRangeStorageIndex(op.result()), op.value());
 }
 void Generator::generate(pdl_interp::EraseOp op, ByteCodeWriter &writer) {
   writer.append(OpCode::EraseOp, op.operation());
 }
+void Generator::generate(pdl_interp::ExtractOp op, ByteCodeWriter &writer) {
+  OpCode opCode =
+      TypeSwitch<Type, OpCode>(op.result().getType())
+          .Case([](pdl::AttributeType) {
+            assert(false &&
+                   "attribute ranges are not supported in bytecode yet");
+            return OpCode();
+          })
+          .Case([](pdl::OperationType) { return OpCode::ExtractOp; })
+          .Case([](pdl::ValueType) { return OpCode::ExtractValue; })
+          .Case([](pdl::TypeType) { return OpCode::ExtractType; });
+  writer.append(opCode, op.range(), op.index(), op.result());
+}
 void Generator::generate(pdl_interp::FinalizeOp op, ByteCodeWriter &writer) {
   writer.append(OpCode::Finalize);
 }
 void Generator::generate(pdl_interp::GetAttributeOp op,
                          ByteCodeWriter &writer) {
   writer.append(OpCode::GetAttribute, op.attribute(), op.operation(),
                 op.nameAttr());
 }
writer.append(OpCode::GetResults,
                 index.getValueOr(std::numeric_limits<uint32_t>::max()),
                 op.operation());
   if (result.getType().isa<pdl::RangeType>())
     writer.append(getRangeStorageIndex(result));
   else
     writer.append(std::numeric_limits<ByteCodeField>::max());
   writer.append(result);
 }
+void Generator::generate(pdl_interp::GetUsersOp op, ByteCodeWriter &writer) {
+  Value operations = op.operations();
+  ByteCodeField rangeIndex = getRangeStorageIndex(operations);
+  writer.append(OpCode::GetUsers, operations, rangeIndex);
+  writer.appendPDLValue(op.value());
+}
 void Generator::generate(pdl_interp::GetValueTypeOp op,
                          ByteCodeWriter &writer) {
   if (op.getType().isa<pdl::RangeType>()) {
     Value result = op.result();
     writer.append(OpCode::GetValueRangeTypes, result,
                   getRangeStorageIndex(result), op.value());
   } else {
     writer.append(OpCode::GetValueType, op.result(), op.value());
   }
 }
 
 void Generator::generate(pdl_interp::InferredTypesOp op,
                          ByteCodeWriter &writer) {
   // InferType maps to a null type as a marker for inferring result types.
   getMemIndex(op.type()) = getMemIndex(Type());
 }
 void Generator::generate(pdl_interp::IsNotNullOp op, ByteCodeWriter &writer) {
   writer.append(OpCode::IsNotNull, op.value(), op.getSuccessors());
 }
+void Generator::generate(pdl_interp::ForEachOp op, ByteCodeWriter &writer) {
+  BlockArgument arg = op.getLoopVariable();
+  writer.append(OpCode::ForEach, getRangeStorageIndex(op.values()), arg);
+  writer.appendPDLValueKind(arg.getType());
+  writer.append(curLoopLevel, op.successor());
+  ++curLoopLevel;
+  if (curLoopLevel > maxLoopLevel)
+    maxLoopLevel = curLoopLevel;
+  generate(&op.region(), writer);
+  --curLoopLevel;
+}
 void Generator::generate(pdl_interp::RecordMatchOp op, ByteCodeWriter &writer) {
   ByteCodeField patternIndex = patterns.size();
   patterns.emplace_back(PDLByteCodePattern::create(
       op, rewriterToAddr[op.rewriter().getLeafReference().getValue()]));
   writer.append(OpCode::RecordMatch, patternIndex,
                 SuccessorRange(op.getOperation()), op.matchedOps());
   writer.appendPDLValueList(op.inputs());
 }

 // PDLByteCode
 //===----------------------------------------------------------------------===//
 
 PDLByteCode::PDLByteCode(ModuleOp module,
                          llvm::StringMap<PDLConstraintFunction> constraintFns,
                          llvm::StringMap<PDLRewriteFunction> rewriteFns) {
   Generator generator(module.getContext(), uniquedData, matcherByteCode,
                       rewriterByteCode, patterns, maxValueMemoryIndex,
-                      maxTypeRangeCount, maxValueRangeCount, constraintFns,
-                      rewriteFns);
+                      maxOpRangeCount, maxTypeRangeCount, maxValueRangeCount,
+                      maxLoopLevel, constraintFns, rewriteFns);
   generator.generate(module);
 
   // Initialize the external functions.
   for (auto &it : constraintFns)
     constraintFunctions.push_back(std::move(it.second));
   for (auto &it : rewriteFns)
     rewriteFunctions.push_back(std::move(it.second));
 }
 
 /// Initialize the given state such that it can be used to execute the current
 /// bytecode.
 void PDLByteCode::initializeMutableState(PDLByteCodeMutableState &state) const {
   state.memory.resize(maxValueMemoryIndex, nullptr);
+  state.opRangeMemory.resize(maxOpRangeCount);
   state.typeRangeMemory.resize(maxTypeRangeCount, TypeRange());
   state.valueRangeMemory.resize(maxValueRangeCount, ValueRange());
+  state.loopIndex.resize(maxLoopLevel, 0);
   state.currentPatternBenefits.reserve(patterns.size());
   for (const PDLByteCodePattern &pattern : patterns)
     state.currentPatternBenefits.push_back(pattern.getBenefit());
 }
 
 //===----------------------------------------------------------------------===//
 // ByteCode Execution
 
 namespace {
 /// This class provides support for executing a bytecode stream.
 class ByteCodeExecutor {
 public:
   ByteCodeExecutor(
       const ByteCodeField *curCodeIt, MutableArrayRef<const void *> memory,
+      MutableArrayRef<llvm::OwningArrayRef<Operation *>> opRangeMemory,
       MutableArrayRef<TypeRange> typeRangeMemory,
       std::vector<llvm::OwningArrayRef<Type>> &allocatedTypeRangeMemory,
       MutableArrayRef<ValueRange> valueRangeMemory,
       std::vector<llvm::OwningArrayRef<Value>> &allocatedValueRangeMemory,
-      ArrayRef<const void *> uniquedMemory, ArrayRef<ByteCodeField> code,
+      MutableArrayRef<unsigned> loopIndex, ArrayRef<const void *> uniquedMemory,
+      ArrayRef<ByteCodeField> code,
       ArrayRef<PatternBenefit> currentPatternBenefits,
       ArrayRef<PDLByteCodePattern> patterns,
       ArrayRef<PDLConstraintFunction> constraintFunctions,
       ArrayRef<PDLRewriteFunction> rewriteFunctions)
-      : curCodeIt(curCodeIt), memory(memory), typeRangeMemory(typeRangeMemory),
+      : curCodeIt(curCodeIt), memory(memory), opRangeMemory(opRangeMemory),
+        typeRangeMemory(typeRangeMemory),
         allocatedTypeRangeMemory(allocatedTypeRangeMemory),
         valueRangeMemory(valueRangeMemory),
         allocatedValueRangeMemory(allocatedValueRangeMemory),
-        uniquedMemory(uniquedMemory), code(code),
+        loopIndex(loopIndex), uniquedMemory(uniquedMemory), code(code),
         currentPatternBenefits(currentPatternBenefits), patterns(patterns),
         constraintFunctions(constraintFunctions),
         rewriteFunctions(rewriteFunctions) {}
 
   /// Start executing the code at the current bytecode index. `matches` is an
   /// optional field provided when this function is executed in a matching
   /// context.
   void execute(PatternRewriter &rewriter,
                SmallVectorImpl<PDLByteCode::MatchResult> *matches = nullptr,
                Optional<Location> mainRewriteLoc = {});
 
 private:
   /// Internal implementation of executing each of the bytecode commands.
   void executeApplyConstraint(PatternRewriter &rewriter);
   void executeApplyRewrite(PatternRewriter &rewriter);
   void executeAreEqual();
   void executeAreRangesEqual();
   void executeBranch();
   void executeCheckOperandCount();
   void executeCheckOperationName();
   void executeCheckResultCount();
   void executeCheckTypes();
+  void executeContinue();
   void executeCreateOperation(PatternRewriter &rewriter,
                               Location mainRewriteLoc);
   void executeCreateTypes();
   void executeEraseOp(PatternRewriter &rewriter);
+  template <typename T, typename Range, PDLValue::Kind kind>
+  void executeExtract();
+  void executeFinalize();
+  void executeForEach();
   void executeGetAttribute();
   void executeGetAttributeType();
   void executeGetDefiningOp();
   void executeGetOperand(unsigned index);
   void executeGetOperands();
   void executeGetResult(unsigned index);
   void executeGetResults();
+  void executeGetUsers();
   void executeGetValueType();
   void executeGetValueRangeTypes();
   void executeIsNotNull();
   void executeRecordMatch(PatternRewriter &rewriter,
                           SmallVectorImpl<PDLByteCode::MatchResult> &matches);
   void executeReplaceOp(PatternRewriter &rewriter);
   void executeSwitchAttribute();
   void executeSwitchOperandCount();
   void executeSwitchOperationName();
   void executeSwitchResultCount();
   void executeSwitchType();
   void executeSwitchTypes();
 
+  /// Pushes a code iterator to the stack.
+  void pushCodeIt(const ByteCodeField *it) { resumeCodeIt.push_back(it); }
+
+  /// Pops a code iterator from the stack, returning true on success.
+  void popCodeIt() {
+    assert(!resumeCodeIt.empty() && "attempt to pop code off empty stack");
+    curCodeIt = resumeCodeIt.back();
+    resumeCodeIt.pop_back();
+  }
+
   /// Read a value from the bytecode buffer, optionally skipping a certain
   /// number of prefix values. These methods always update the buffer to point
   /// to the next field after the read data.
   template <typename T = ByteCodeField>
   T read(size_t skipN = 0) {
     curCodeIt += skipN;
     return readImpl<T>();
   }
void handleSwitch(const T &value, RangeT &&cases, Comparator cmp = {}) {
     // Check to see if the attribute value is within the case list. Jump to
     // the correct successor index based on the result.
     for (auto it = cases.begin(), e = cases.end(); it != e; ++it)
       if (cmp(*it, value))
         return selectJump(size_t((it - cases.begin()) + 1));
     selectJump(size_t(0));
   }
 
+  /// Store a pointer to memory.
+  void storeToMemory(unsigned index, const void *value) {
+    memory[index] = value;
+  }
+
+  /// Store a value to memory as an opaque pointer.
+  template <typename T>
+  std::enable_if_t<llvm::is_detected<has_pointer_traits, T>::value>
+  storeToMemory(unsigned index, T value) {
+    memory[index] = value.getAsOpaquePointer();
+  }
+
   /// Internal implementation of reading various data types from the bytecode
   /// stream.
   template <typename T>
   const void *readFromMemory() {
     size_t index = *curCodeIt++;
 
     // If this type is an SSA value, it can only be stored in non-const memory.
     if (llvm::is_one_of<T, Operation *, TypeRange *, ValueRange *,
private:
   template <typename T>
   std::enable_if_t<std::is_same<T, PDLValue::Kind>::value, T> readImpl() {
     return static_cast<PDLValue::Kind>(readImpl<ByteCodeField>());
   }
 
   /// The underlying bytecode buffer.
   const ByteCodeField *curCodeIt;
 
+  /// The stack of bytecode positions at which to resume operation.
+  SmallVector<const ByteCodeField *> resumeCodeIt;
+
   /// The current execution memory.
   MutableArrayRef<const void *> memory;
+  MutableArrayRef<OwningOpRange> opRangeMemory;
   MutableArrayRef<TypeRange> typeRangeMemory;
   std::vector<llvm::OwningArrayRef<Type>> &allocatedTypeRangeMemory;
   MutableArrayRef<ValueRange> valueRangeMemory;
   std::vector<llvm::OwningArrayRef<Value>> &allocatedValueRangeMemory;
 
+  /// The current loop indices.
+  MutableArrayRef<unsigned> loopIndex;
+
   /// References to ByteCode data necessary for execution.
   ArrayRef<const void *> uniquedMemory;
   ArrayRef<ByteCodeField> code;
   ArrayRef<PatternBenefit> currentPatternBenefits;
   ArrayRef<PDLByteCodePattern> patterns;
   ArrayRef<PDLConstraintFunction> constraintFunctions;
   ArrayRef<PDLRewriteFunction> rewriteFunctions;
 };
void ByteCodeExecutor::executeCheckTypes() {
   LLVM_DEBUG(llvm::dbgs() << "Executing AreEqual:\n");
   TypeRange *lhs = read<TypeRange *>();
   Attribute rhs = read<Attribute>();
   LLVM_DEBUG(llvm::dbgs() << "  * " << lhs << " == " << rhs << "\n\n");
 
   selectJump(*lhs == rhs.cast<ArrayAttr>().getAsValueRange<TypeAttr>());
 }
 
+void ByteCodeExecutor::executeContinue() {
+  ByteCodeField level = read();
+  LLVM_DEBUG(llvm::dbgs() << "Executing Continue\n"
+                          << "  * Level: " << level << "\n");
+  ++loopIndex[level];
+  popCodeIt();
+}
+
 void ByteCodeExecutor::executeCreateTypes() {
   LLVM_DEBUG(llvm::dbgs() << "Executing CreateTypes:\n");
   unsigned memIndex = read();
   unsigned rangeIndex = read();
   ArrayAttr typesAttr = read<Attribute>().cast<ArrayAttr>();
 
   LLVM_DEBUG(llvm::dbgs() << "  * Types: " << typesAttr << "\n\n");
 
void ByteCodeExecutor::executeCreateOperation(PatternRewriter &rewriter,
 
   Operation *resultOp = rewriter.createOperation(state);
   memory[memIndex] = resultOp;
 
   LLVM_DEBUG({
     llvm::dbgs() << "  * Attributes: "
                  << state.attributes.getDictionary(state.getContext())
                  << "\n  * Operands: ";
     llvm::interleaveComma(state.operands, llvm::dbgs());
     llvm::dbgs() << "\n  * Result Types: ";
     llvm::interleaveComma(state.types, llvm::dbgs());
     llvm::dbgs() << "\n  * Result: " << *resultOp << "\n";
   });
 }
 
 void ByteCodeExecutor::executeEraseOp(PatternRewriter &rewriter) {
   LLVM_DEBUG(llvm::dbgs() << "Executing EraseOp:\n");
   Operation *op = read<Operation *>();
 
   LLVM_DEBUG(llvm::dbgs() << "  * Operation: " << *op << "\n");
   rewriter.eraseOp(op);
 }
 
+template <typename T, typename Range, PDLValue::Kind kind>
+void ByteCodeExecutor::executeExtract() {
+  LLVM_DEBUG(llvm::dbgs() << "Executing Extract" << kind << ":\n");
+  Range *range = read<Range *>();
+  unsigned index = read<uint32_t>();
+  unsigned memIndex = read();
+
+  if (!range) {
+    memory[memIndex] = nullptr;
+    return;
+  }
+
+  T result = index < range->size() ? (*range)[index] : T();
+  LLVM_DEBUG(llvm::dbgs() << "  * " << kind << "s(" << range->size() << ")\n"
+                          << "  * Index: " << index << "\n"
+                          << "  * Result: " << result << "\n");
+  storeToMemory(memIndex, result);
+}
+
+void ByteCodeExecutor::executeFinalize() {
+  LLVM_DEBUG(llvm::dbgs() << "Executing Finalize\n");
+}
+
+void ByteCodeExecutor::executeForEach() {
+  LLVM_DEBUG(llvm::dbgs() << "Executing ForEach:\n");
+  // Subtract 1 for the op code.
+  const ByteCodeField *it = curCodeIt - 1;
+  unsigned rangeIndex = read();
+  unsigned memIndex = read();
+  const void *value = nullptr;
+
+  switch (read<PDLValue::Kind>()) {
+  case PDLValue::Kind::Operation: {
+    unsigned &index = loopIndex[read()];
+    ArrayRef<Operation *> array = opRangeMemory[rangeIndex];
+    assert(index <= array.size() && "iterated past the end");
+    if (index < array.size()) {
+      LLVM_DEBUG(llvm::dbgs() << "  * Result: " << array[index] << "\n");
+      value = array[index];
+      break;
+    }
+
+    LLVM_DEBUG(llvm::dbgs() << "  * Done\n");
+    index = 0;
+    selectJump(size_t(0));
+    return;
+  }
+  default:
+    llvm_unreachable("unexpected `ForEach` value kind");
+  }
+
+  // Store the iterate value and the stack address.
+  memory[memIndex] = value;
+  pushCodeIt(it);
+
+  // Skip over the successor (we will enter the body of the loop).
+  read<ByteCodeAddr>();
+}
+
 void ByteCodeExecutor::executeGetAttribute() {
   LLVM_DEBUG(llvm::dbgs() << "Executing GetAttribute:\n");
   unsigned memIndex = read();
   Operation *op = read<Operation *>();
   StringAttr attrName = read<StringAttr>();
   Attribute attr = op->getAttr(attrName);
 
   LLVM_DEBUG(llvm::dbgs() << "  * Operation: " << *op << "\n"

 
 /// This function is the internal implementation of `GetResults` and
 /// `GetOperands` that provides support for extracting a value range from the
 /// given operation.
 template <template <typename> class AttrSizedSegmentsT, typename RangeT>
 static void *
 executeGetOperandsResults(RangeT values, Operation *op, unsigned index,
                           ByteCodeField rangeIndex, StringRef attrSizedSegments,
-                          MutableArrayRef<ValueRange> &valueRangeMemory) {
+                          MutableArrayRef<ValueRange> valueRangeMemory) {
   // Check for the sentinel index that signals that all values should be
   // returned.
   if (index == std::numeric_limits<uint32_t>::max()) {
     LLVM_DEBUG(llvm::dbgs() << "  * Getting all values\n");
     // `values` is already the full value range.
 
     // Otherwise, check to see if this operation uses AttrSizedSegments.
   } else if (op->hasTrait<AttrSizedSegmentsT>()) {
void ByteCodeExecutor::executeGetResults() {
   void *result = executeGetOperandsResults<OpTrait::AttrSizedResultSegments>(
       op->getResults(), op, index, rangeIndex, "result_segment_sizes",
       valueRangeMemory);
   if (!result)
     LLVM_DEBUG(llvm::dbgs() << "  * Invalid result range\n");
   memory[read()] = result;
 }
 
+void ByteCodeExecutor::executeGetUsers() {
+  LLVM_DEBUG(llvm::dbgs() << "Executing GetUsers:\n");
+  unsigned memIndex = read();
+  unsigned rangeIndex = read();
+  OwningOpRange &range = opRangeMemory[rangeIndex];
+  memory[memIndex] = &range;
+
+  range = OwningOpRange();
+  if (read<PDLValue::Kind>() == PDLValue::Kind::Value) {
+    // Read the value.
+    Value value = read<Value>();
+    if (!value)
+      return;
+    LLVM_DEBUG(llvm::dbgs() << "  * Value: " << value << "\n");
+
+    // Extract the users of a single value.
+    range = OwningOpRange(std::distance(value.user_begin(), value.user_end()));
+    llvm::copy(value.getUsers(), range.begin());
+  } else {
+    // Read a range of values.
+    ValueRange *values = read<ValueRange *>();
+    if (!values)
+      return;
+    LLVM_DEBUG({
+      llvm::dbgs() << "  * Values (" << values->size() << "): ";
+      llvm::interleaveComma(*values, llvm::dbgs());
+      llvm::dbgs() << "\n";
+    });
+
+    // Extract all the users of a range of values.
+    SmallVector<Operation *> users;
+    for (Value value : *values) {
+      users.append(value.user_begin(), value.user_end());
+    }
+    range = OwningOpRange(users.size());
+    llvm::copy(users, range.begin());
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "  * Result: " << range.size() << " operations\n");
+}
+
 void ByteCodeExecutor::executeGetValueType() {
   LLVM_DEBUG(llvm::dbgs() << "Executing GetValueType:\n");
   unsigned memIndex = read();
   Value value = read<Value>();
   Type type = value ? value.getType() : Type();
 
   LLVM_DEBUG(llvm::dbgs() << "  * Value: " << value << "\n"
                           << "  * Result: " << type << "\n");
   memory[memIndex] = type.getAsOpaquePointer();
 }
 
 void ByteCodeExecutor::executeGetValueRangeTypes() {
   LLVM_DEBUG(llvm::dbgs() << "Executing GetValueRangeTypes:\n");
   unsigned memIndex = read();
   unsigned rangeIndex = read();
   ValueRange *values = read<ValueRange *>();
   if (!values) {
     LLVM_DEBUG(llvm::dbgs() << "  * Values: <NULL>\n\n");
     memory[memIndex] = nullptr;
     return;
   }
 
   LLVM_DEBUG({
     llvm::dbgs() << "  * Values (" << values->size() << "): ";
     llvm::interleaveComma(*values, llvm::dbgs());
     llvm::dbgs() << "\n  * Result: ";
     llvm::interleaveComma(values->getType(), llvm::dbgs());
     llvm::dbgs() << "\n";
   });
   typeRangeMemory[rangeIndex] = values->getType();
   memory[memIndex] = &typeRangeMemory[rangeIndex];
 }
 
 void ByteCodeExecutor::executeIsNotNull() {
   LLVM_DEBUG(llvm::dbgs() << "Executing IsNotNull:\n");
   const void *value = read<const void *>();
 
   LLVM_DEBUG(llvm::dbgs() << "  * Value: " << value << "\n");
   selectJump(value != nullptr);
 }
 
 void ByteCodeExecutor::executeRecordMatch(
     PatternRewriter &rewriter,
     SmallVectorImpl<PDLByteCode::MatchResult> &matches) {
   LLVM_DEBUG(llvm::dbgs() << "Executing RecordMatch:\n");
   unsigned patternIndex = read();
   PatternBenefit benefit = currentPatternBenefits[patternIndex];
   const ByteCodeField *dest = &code[read<ByteCodeAddr>()];
 
   // If the benefit of the pattern is impossible, skip the processing of the
   // rest of the pattern.
   if (benefit.isImpossibleToMatch()) {
     LLVM_DEBUG(llvm::dbgs() << "  * Benefit: Impossible To Match\n");
     curCodeIt = dest;
     return;
   }
 
   // Create a fused location containing the locations of each of the
   // operations used in the match. This will be used as the location for
   // created operations during the rewrite that don't already have an
   // explicit location set.
   unsigned numMatchLocs = read();
   SmallVector<Location, 4> matchLocs;
   matchLocs.reserve(numMatchLocs);
   for (unsigned i = 0; i != numMatchLocs; ++i)
     matchLocs.push_back(read<Operation *>()->getLoc());
case CheckOperationName:
       executeCheckOperationName();
       break;
     case CheckResultCount:
       executeCheckResultCount();
       break;
     case CheckTypes:
       executeCheckTypes();
       break;
+    case Continue:
+      executeContinue();
+      break;
     case CreateOperation:
       executeCreateOperation(rewriter, *mainRewriteLoc);
       break;
     case CreateTypes:
       executeCreateTypes();
       break;
     case EraseOp:
       executeEraseOp(rewriter);
       break;
+    case ExtractOp:
+      executeExtract<Operation *, OwningOpRange, PDLValue::Kind::Operation>();
+      break;
+    case ExtractType:
+      executeExtract<Type, TypeRange, PDLValue::Kind::Type>();
+      break;
+    case ExtractValue:
+      executeExtract<Value, ValueRange, PDLValue::Kind::Value>();
+      break;
     case Finalize:
-      LLVM_DEBUG(llvm::dbgs() << "Executing Finalize\n\n");
+      executeFinalize();
+      LLVM_DEBUG(llvm::dbgs() << "\n");
       return;
+    case ForEach:
+      executeForEach();
+      break;
     case GetAttribute:
       executeGetAttribute();
       break;
     case GetAttributeType:
       executeGetAttributeType();
       break;
     case GetDefiningOp:
       executeGetDefiningOp();
while (true) {
     }
     case GetResultN:
       LLVM_DEBUG(llvm::dbgs() << "Executing GetResultN:\n");
       executeGetResult(read<uint32_t>());
       break;
     case GetResults:
       executeGetResults();
       break;
+    case GetUsers:
+      executeGetUsers();
+      break;
     case GetValueType:
       executeGetValueType();
       break;
     case GetValueRangeTypes:
       executeGetValueRangeTypes();
       break;
     case IsNotNull:
       executeIsNotNull();

 void PDLByteCode::match(Operation *op, PatternRewriter &rewriter,
                         SmallVectorImpl<MatchResult> &matches,
                         PDLByteCodeMutableState &state) const {
   // The first memory slot is always the root operation.
   state.memory[0] = op;
 
   // The matcher function always starts at code address 0.
   ByteCodeExecutor executor(
-      matcherByteCode.data(), state.memory, state.typeRangeMemory,
-      state.allocatedTypeRangeMemory, state.valueRangeMemory,
-      state.allocatedValueRangeMemory, uniquedData, matcherByteCode,
-      state.currentPatternBenefits, patterns, constraintFunctions,
-      rewriteFunctions);
+      matcherByteCode.data(), state.memory, state.opRangeMemory,
+      state.typeRangeMemory, state.allocatedTypeRangeMemory,
+      state.valueRangeMemory, state.allocatedValueRangeMemory, state.loopIndex,
+      uniquedData, matcherByteCode, state.currentPatternBenefits, patterns,
+      constraintFunctions, rewriteFunctions);
   executor.execute(rewriter, &matches);
 
   // Order the found matches by benefit.
   std::stable_sort(matches.begin(), matches.end(),
                    [](const MatchResult &lhs, const MatchResult &rhs) {
                      return lhs.benefit > rhs.benefit;
                    });
 }
 
 /// Run the rewriter of the given pattern on the root operation `op`.
 void PDLByteCode::rewrite(PatternRewriter &rewriter, const MatchResult &match,
                           PDLByteCodeMutableState &state) const {
   // The arguments of the rewrite function are stored at the start of the
   // memory buffer.
   llvm::copy(match.values, state.memory.begin());
 
   ByteCodeExecutor executor(
       &rewriterByteCode[match.pattern->getRewriterAddr()], state.memory,
-      state.typeRangeMemory, state.allocatedTypeRangeMemory,
-      state.valueRangeMemory, state.allocatedValueRangeMemory, uniquedData,
+      state.opRangeMemory, state.typeRangeMemory,
+      state.allocatedTypeRangeMemory, state.valueRangeMemory,
+      state.allocatedValueRangeMemory, state.loopIndex, uniquedData,
       rewriterByteCode, state.currentPatternBenefits, patterns,
       constraintFunctions, rewriteFunctions);
   executor.execute(rewriter, /*matches=*/nullptr, match.location);
 }

mlir/test/Rewrite/pdl-bytecode.mlir

 module @ir attributes { test.check_types_1 } {
   "test.op"() : () -> (i32, i64)
   "test.op"() : () -> i32
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::ContinueOp
+//===----------------------------------------------------------------------===//
+
+// Fully tested within the tests for other operations.
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::CreateAttributeOp
 //===----------------------------------------------------------------------===//
 
 // Fully tested within the tests for other operations.
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::CreateOperationOp
 //===----------------------------------------------------------------------===//

 
 //===----------------------------------------------------------------------===//
 // pdl_interp::EraseOp
 //===----------------------------------------------------------------------===//
 
 // Fully tested within the tests for other operations.
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::ExtractOp
+//===----------------------------------------------------------------------===//
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    %val = pdl_interp.get_result 0 of %root
+    %ops = pdl_interp.get_users of %val : !pdl.value
+    %op1 = pdl_interp.extract 1 of %ops : !pdl.operation
+    pdl_interp.is_not_null %op1 : !pdl.operation -> ^success, ^end
+  ^success:
+    pdl_interp.record_match @rewriters::@success(%op1 : !pdl.operation) : benefit(1), loc([%root]) -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.extract_op
+// CHECK: "test.success"
+// CHECK: %[[OPERAND:.*]] = "test.op"
+// CHECK: "test.op"(%[[OPERAND]])
+module @ir attributes { test.extract_op } {
+  %operand = "test.op"() : () -> i32
+  "test.op"(%operand) : (i32) -> (i32)
+  "test.op"(%operand, %operand) : (i32, i32) -> (i32)
+}
+
+// -----
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    %vals = pdl_interp.get_results of %root : !pdl.range<value>
+    %types = pdl_interp.get_value_type of %vals : !pdl.range<type>
+    %type1 = pdl_interp.extract 1 of %types : !pdl.type
+    pdl_interp.is_not_null %type1 : !pdl.type -> ^success, ^end
+  ^success:
+    pdl_interp.record_match @rewriters::@success(%root : !pdl.operation) : benefit(1), loc([%root]) -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.extract_type
+// CHECK: %[[OPERAND:.*]] = "test.op"
+// CHECK: "test.success"
+// CHECK: "test.op"(%[[OPERAND]])
+module @ir attributes { test.extract_type } {
+  %operand = "test.op"() : () -> i32
+  "test.op"(%operand) : (i32) -> (i32, i32)
+  "test.op"(%operand) : (i32) -> (i32)
+}
+
+// -----
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    %vals = pdl_interp.get_results of %root : !pdl.range<value>
+    %val1 = pdl_interp.extract 1 of %vals : !pdl.value
+    pdl_interp.is_not_null %val1 : !pdl.value -> ^success, ^end
+  ^success:
+    pdl_interp.record_match @rewriters::@success(%root : !pdl.operation) : benefit(1), loc([%root]) -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.extract_value
+// CHECK: %[[OPERAND:.*]] = "test.op"
+// CHECK: "test.success"
+// CHECK: "test.op"(%[[OPERAND]])
+module @ir attributes { test.extract_value } {
+  %operand = "test.op"() : () -> i32
+  "test.op"(%operand) : (i32) -> (i32, i32)
+  "test.op"(%operand) : (i32) -> (i32)
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::FinalizeOp
 //===----------------------------------------------------------------------===//
 
 // Fully tested within the tests for other operations.
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::ForEachOp
+//===----------------------------------------------------------------------===//
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    %val1 = pdl_interp.get_result 0 of %root
+    %ops1 = pdl_interp.get_users of %val1 : !pdl.value
+    pdl_interp.foreach %op1 : !pdl.operation in %ops1 {
+      %val2 = pdl_interp.get_result 0 of %op1
+      %ops2 = pdl_interp.get_users of %val2 : !pdl.value
+      pdl_interp.foreach %op2 : !pdl.operation in %ops2 {
+        pdl_interp.record_match @rewriters::@success(%op2 : !pdl.operation) : benefit(1), loc([%root]) -> ^cont
+      ^cont:
+        pdl_interp.continue
+      } -> ^cont
+    ^cont:
+      pdl_interp.continue
+    } -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.foreach
+// CHECK: "test.success"
+// CHECK: "test.success"
+// CHECK: "test.success"
+// CHECK: "test.success"
+// CHECK: %[[ROOT:.*]] = "test.op"
+// CHECK: %[[VALA:.*]] = "test.op"(%[[ROOT]])
+// CHECK: %[[VALB:.*]] = "test.op"(%[[ROOT]])
+module @ir attributes { test.foreach } {
+  %root = "test.op"() : () -> i32
+  %valA = "test.op"(%root) : (i32) -> (i32)
+  "test.op"(%valA) : (i32) -> (i32)
+  "test.op"(%valA) : (i32) -> (i32)
+  %valB = "test.op"(%root) : (i32) -> (i32)
+  "test.op"(%valB) : (i32) -> (i32)
+  "test.op"(%valB) : (i32) -> (i32)
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// pdl_interp::GetUsersOp
+//===----------------------------------------------------------------------===//
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    %val = pdl_interp.get_result 0 of %root
+    %ops = pdl_interp.get_users of %val : !pdl.value
+    pdl_interp.foreach %op : !pdl.operation in %ops {
+      pdl_interp.record_match @rewriters::@success(%op : !pdl.operation) : benefit(1), loc([%root]) -> ^cont
+    ^cont:
+      pdl_interp.continue
+    } -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.get_users_of_value
+// CHECK: "test.success"
+// CHECK: "test.success"
+// CHECK: %[[OPERAND:.*]] = "test.op"
+module @ir attributes { test.get_users_of_value } {
+  %operand = "test.op"() : () -> i32
+  "test.op"(%operand) : (i32) -> (i32)
+  "test.op"(%operand, %operand) : (i32, i32) -> (i32)
+}
+
+// -----
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    pdl_interp.check_result_count of %root is at_least 2 -> ^next, ^end
+  ^next:
+    %vals = pdl_interp.get_results of %root : !pdl.range<value>
+    %ops = pdl_interp.get_users of %vals : !pdl.range<value>
+    pdl_interp.foreach %op : !pdl.operation in %ops {
+      pdl_interp.record_match @rewriters::@success(%op : !pdl.operation) : benefit(1), loc([%root]) -> ^cont
+    ^cont:
+      pdl_interp.continue
+    } -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.get_all_users_of_range
+// CHECK: "test.success"
+// CHECK: "test.success"
+// CHECK: %[[OPERANDS:.*]]:2 = "test.op"
+module @ir attributes { test.get_all_users_of_range } {
+  %operands:2 = "test.op"() : () -> (i32, i32)
+  "test.op"(%operands#0) : (i32) -> (i32)
+  "test.op"(%operands#1) : (i32) -> (i32)
+}
+
+// -----
+
+module @patterns {
+  func @matcher(%root : !pdl.operation) {
+    pdl_interp.check_result_count of %root is at_least 2 -> ^next, ^end
+  ^next:
+    %vals = pdl_interp.get_results of %root : !pdl.range<value>
+    %val = pdl_interp.extract 0 of %vals : !pdl.value
+    %ops = pdl_interp.get_users of %val : !pdl.value
+    pdl_interp.foreach %op : !pdl.operation in %ops {
+      pdl_interp.record_match @rewriters::@success(%op : !pdl.operation) : benefit(1), loc([%root]) -> ^cont
+    ^cont:
+      pdl_interp.continue
+    } -> ^end
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    func @success(%matched : !pdl.operation) {
+      %op = pdl_interp.create_operation "test.success"
+      pdl_interp.erase %matched
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.get_first_users_of_range
+// CHECK: "test.success"
+// CHECK: %[[OPERANDS:.*]]:2 = "test.op"
+// CHECK: "test.op"
+module @ir attributes { test.get_first_users_of_range } {
+  %operands:2 = "test.op"() : () -> (i32, i32)
+  "test.op"(%operands#0) : (i32) -> (i32)
+  "test.op"(%operands#1) : (i32) -> (i32)
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::GetAttributeOp
 //===----------------------------------------------------------------------===//
 
 // Fully tested within the tests for other operations.
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::GetAttributeTypeOp
 //===----------------------------------------------------------------------===//