[mlir] use side effects in the Transform dialect

[ftynse]

Summary

Currently, the sequence of Transform dialect operations only supports a single
use of each operand (verified by the transform.sequence operation). This was
originally motivated by the need to guard against accessing a payload IR
operation associated with a transform IR value after this operation has likely
been rewritten by a transformation. However, not all Transform dialect
operations rewrite payload IR, in particular the "navigation" operation such as
transform.pdl_match do not.

Introduce memory effects to the Transform dialect operations to describe their
effect on the payload IR and the mapping between payload IR opreations and
transform IR values. Use these effects to replace the single-use rule, allowing
repeated reads and disallowing use-after-free, where operations with the "free"
effect are considered to "consume" the transform IR value and rewrite the
corresponding payload IR operations). As an additional improvement, this
enables code motion transformation on the transform IR itself.

[Mogball]

Does allowing the reordering of transform operations really make sense if they do not declare all their pattern inputs and outputs? E.g.

alloc/write(%1) = transform.foo read(%0) 
alloc/write(%2) = transform.bar read(%0)

What if foo creates other ops for which bar optionally checks and which, if present, change what bar does? Reordering these two operations becomes invalid. Any transform op that corresponds to a pattern, for example, that is not purely functional (i.e. that does not declare all of its input ops and output ops) will need to be annotated with a generic MemoryRead and MemoryWrite.

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td

68

69

125

150

152

Can you add some transform IR code to go along with the explanation?

mlir/include/mlir/Dialect/Transform/IR/TransformEffects.td

39

42

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

385

392

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

111–113

[ftynse]

Address review.

[ftynse]

In D124181#3465833, @Mogball wrote:

Does allowing the reordering of transform operations really make sense if they do not declare all their pattern inputs and outputs? E.g.

alloc/write(%1) = transform.foo read(%0) 
alloc/write(%2) = transform.bar read(%0)

What if foo creates other ops for which bar optionally checks and which, if present, change what bar does? Reordering these two operations becomes invalid. Any transform op that corresponds to a pattern, for example, that is not purely functional (i.e. that does not declare all of its input ops and output ops) will need to be annotated with a generic MemoryRead and MemoryWrite.

Indeed, such ops will have to be annotated with read/write on the PayloadIRResource, which is the right mechanism to disallow reordering. With time, we can use the parameter field of the EffectInstance to qualify payload IR changes at a finer grain if desired.

This does allow permuting and deduplicating read-only ops as well as repeated reads from the same handle, which sounds quite useful.

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

111–113

Why? I expect clear not to free the dynamically allocated memory (if any) so it can be reused on future iterations if necessary, this would be impossible with a vector freshly created in each iteration.

[Mogball]

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

111–113

Oh, right you are

[rriddle]

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

183

Missing static here?

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.h

 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringMap.h"
 
-#include "mlir/Dialect/Transform/IR/TransformDialect.h.inc"
-
 namespace mlir {
 namespace transform {
-
 #ifndef NDEBUG
 namespace detail {
 /// Asserts that the operations provided as template arguments implement the
-/// TransformOpInterface. This must be a dynamic assertion since interface
-/// implementations may be registered at runtime.
+/// TransformOpInterface and MemoryEffectsOpInterface. This must be a dynamic
+/// assertion since interface implementations may be registered at runtime.
 template <typename OpTy>
 static inline void checkImplementsTransformInterface(MLIRContext *context) {
   // Since the operation is being inserted into the Transform dialect and the
   // dialect does not implement the interface fallback, only check for the op
   // itself having the interface implementation.
   RegisteredOperationName opName =
       *RegisteredOperationName::lookup(OpTy::getOperationName(), context);
-  assert(opName.hasInterface<TransformOpInterface>() &&
+  assert((opName.hasInterface<TransformOpInterface>() ||
+          opName.hasTrait<OpTrait::IsTerminator>()) &&
+         "non-terminator ops injected into the transform dialect must "
+         "implement TransformOpInterface");
+  assert(opName.hasInterface<MemoryEffectOpInterface>() &&
          "ops injected into the transform dialect must implement "
-         "TransformOpInterface");
+         "MemoryEffectsOpInterface");
 }
 } // namespace detail
 #endif // NDEBUG
+} // namespace transform
+} // namespace mlir
+
+#include "mlir/Dialect/Transform/IR/TransformDialect.h.inc"
+
+namespace mlir {
+namespace transform {
 
 /// Base class for extensions of the Transform dialect that supports injecting
 /// operations into the Transform dialect at load time. Concrete extensions are
 /// expected to derive this class and register operations in the constructor.
 /// They can be registered with the DialectRegistry and automatically applied
 /// to the Transform dialect when it is loaded.
 template <typename DerivedTy, typename... ExtraDialects>
 class TransformDialectExtension
for (const DialectLoader &loader : dialectLoaders)
       loader(context);
     for (const Initializer &init : opInitializers)
       init(transformDialect);
     transformDialect->mergeInPDLMatchHooks(std::move(pdlMatchConstraintFns));
   }
 
 protected:
   /// Injects the operations into the Transform dialect. The operations must
-  /// implement the TransformOpInterface and the implementation must be already
-  /// available when the operation is injected.
+  /// implement the TransformOpInterface and MemoryEffectsOpInterface, and the
+  /// implementations must be already available when the operation is injected.
   template <typename... OpTys>
   void registerTransformOps() {
     opInitializers.push_back([](TransformDialect *transformDialect) {
-      transformDialect->addOperations<OpTys...>();
-
-#ifndef NDEBUG
-      (void)std::initializer_list<int>{
-          (detail::checkImplementsTransformInterface<OpTys>(
-               transformDialect->getContext()),
-           0)...};
-#endif // NDEBUG
+      transformDialect->addOperationsChecked<OpTys...>();
     });
   }
 
   /// Declares that this Transform dialect extension depends on the dialect
   /// provided as template parameter. When the Transform dialect is loaded,
   /// dependent dialects will be loaded as well. This is intended for dialects
   /// that contain attributes and types used in creation and canonicalization of
   /// the injected operations.

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td

let description = [{
     ```
 
     The values defined by operations in this dialect correspond to (groups of)
     operations in the payload IR. In the example above, `%0` corresponds to the
     set of loops found in the payload IR that satisfy the condition, and `%1`
     correspond to groups of outer and inner loops, respectively, produced by
     the tiling transformation.
 
+    Overall, Transform IR ops are expected to be contained in a single top-level
+    op. Such top-level ops specify how to apply the transformations described
+    by the operations they contain, e.g., `transform.sequence` executes
+    transformations one by one and fails if any of them fails. Such ops are
+    expected to have the `PossibleTopLevelTransformOpTrait` and may be used
+    without arguments.
+
+    ## Dialect Extension Mechanism
+
     This dialect is designed to be extensible, that is, clients of this dialect
     are allowed to inject additional operations into this dialect using the
     `TransformDialectExtension` mechanism. This allows the dialect to avoid a
     dependency on the implementation of the transformation as well as to avoid
     introducing dialect-specific transform dialects. In the example above,
     the operations may have been injected by a notional `loop` dialect rather
     than defined in this dialect, hence the common prefix.
 
     It is recommended to prefix injected operations with one or several
     dot-separated words that indicate which extension adds them. For
     dialect-specific transformations, the prefix is naturally the name of the
     dialect, e.g., `transform.affine.reschedule`. For dialect-agnostic
     transformations (typically implemented using interfaces), the prefix may
     be derived from the interface name or from a common concept, e.g.,
     `transform.loop.tile` may apply to any loop-like operation that implements
     `TileableOpInterface`. The C++ classes for the dialect extension should
     include the prefix in their name, e.g., `AffineTransformDialectExtension` or
     `LoopTransformDialectExtension` in the cases above. Unprefixed operation
     names are reserved for ops defined directly in the Transform dialect.
 
-    Overall, Transform IR ops are expected to be contained in a single top-level
-    op. Such top-level ops specifie how to apply the transformations described
-    by operations they contain, e.g., `transform.sequence` executes
-    transformations one by one and fails if any of them fails. Such ops are
-    expected to have the `PossibleTopLevelTransformOpTrait` and may be used
-    without arguments.
+    Operations injected into the dialect must:
+
+      * Implement the `TransformOpInterface` to execute the corresponding
+        transformation on the payload IR.
+
+      * Implement the `MemoryEffectsOpInterface` to annotate the effects of
+        the transform IR operation on the payload IR as well as on the mapping
+        between transform IR values and payload IR operations. See below for
+        the description of available effects.
+
+    The presence of interface implementations is checked at runtime when the
+    dialect is loaded to allow for those implementations to be supplied by
+    separate dialect extensions if desired.
+
+    ## Side Effects
+
+    The Transform dialect relies on MLIR side effect modelling to enable
+    optimization of the transform IR. More specifically, it provides several
+    side effect resource objects and expects operations to describe their
+    effects on these resources.
+
+      * `TransformMappingResource` - side effect resource corresponding to the
+        mapping between transform IR values and payload IR operations.
+        
+        - An `Allocate` effect from this resource means creating a new mapping
+          entry, it is always accompanied by a `Write` effet.
+          
+        - A `Read` effect from this resource means accessing the mapping.
+        
+        - A `Free` effect on this resource indicates the removal of the mapping
+          entry, typically after a transformation that modifies the payload IR
+          operations associated with one of the transform IR operation's
+          operands. It is always accompanied by a `Read` effect.
+
+      * `PayloadIRResource` - side effect resource corresponding to the payload
+        IR itself.
+
+        - A `Read` effect from this resource means accessing the payload IR.
+
+        - A `Write` effect on this resource means mutating the payload IR. It is
+          almost always accompanied by a `Read`.
+
+    The typical flow of values in the transform IR is as follows. Most
+    operations produce new transform IR values and immediately associate them
+    with a list of payload IR operations. This corresponds to `Allocate` and
+    `Write` effects on the `TransformMappingResource`, and often requires at
+    least a `Read` effect on the `PayloadIRResource`. Transform operations that
+    only inspect the payload IR to produce new handles are usually limited to
+    these effects on their operands. Transform operations that mutate the
+    payload IR are thought to _consume_ the handles provided as operands, that
+    is have the `Read` and `Free` effects on them. As with the usual memory
+    effects, using a value after it was freed is incorrect. In case of the
+    transform IR, this value is likely associated with payload IR operations
+    that were modified or even removed by the transformation, so it is
+    meaningless to refer to them. When further transformations are desired, the
+    transform operations can return _new_ handles that can be read or consumed
+    by subsequent operations.
 
     ## Intended Use and Integrations
 
     The transformation control infrastructure provided by this dialect is
     positioned roughly between rewrite patterns and passes. A transformation
     that is executed by a transform operation is likely to be sufficiently
     complex to require at least a set of patterns to be implemented. It is also
     expected to be more focused than a pass: a pass typically applies identical
def Transform_Dialect : Dialect {
 
   let extraClassDeclaration = [{
       /// Returns the named PDL constraint functions available in the dialect
       /// as a map from their name to the function.
       const ::llvm::StringMap<::mlir::PDLConstraintFunction> &
       getPDLConstraintHooks() const;
 
     private:
-      // Make addOperations available to the TransformDialectExtension class.
-      using ::mlir::Dialect::addOperations;
+      /// Registers operations specified as template parameters with this
+      /// dialect. Checks that they implement the required interfaces.
+      template <typename... OpTys>
+      void addOperationsChecked() {
+        addOperations<OpTys...>();
+
+        #ifndef NDEBUG
+        (void)std::initializer_list<int>{
+          (detail::checkImplementsTransformInterface<OpTys>(getContext()),
+           0)...};
+        #endif // NDEBUG
+      }
 
       template <typename, typename...>
       friend class TransformDialectExtension;
 
       /// Takes ownership of the named PDL constraint function from the given
       /// map and makes them available for use by the operations in the dialect.
       void mergeInPDLMatchHooks(
           ::llvm::StringMap<::mlir::PDLConstraintFunction> &&constraintFns);

mlir/include/mlir/Dialect/Transform/IR/TransformEffects.td

+
+//===- TransformEffect.td - Transform side effects ---------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines side effects and associated resources for operations in the
+// Transform dialect and extensions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_EFFECTS_TD
+#define MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_EFFECTS_TD
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Effects on the mapping between Transform IR values and Payload IR ops.
+//===----------------------------------------------------------------------===//
+
+// Side effect resource corresponding to the mapping between transform IR values
+// and Payload IR operations.
+def TransformMappingResource
+    : Resource<"::mlir::transform::TransformMappingResource">;
+
+// Describes the creation of a new entry in the transform mapping. Should be
+// accompanied by the Write effect as the entry is immediately initialized by
+// any reasonable transform operation.
+def TransformMappingAlloc : MemAlloc<TransformMappingResource>;
+
+// Describes the removal of an entry in the transform mapping. Typically
+// accompanied by the Read effect.
+def TransformMappingFree : MemFree<TransformMappingResource>;
+
+// Describes the access to the mapping. Read-only accesses can be reordered.
+def TransformMappingRead : MemRead<TransformMappingResource>;
+
+// Describes a modification of an existing entry in the mapping. It is rarely
+// used alone, and is mostly accompanied by the Allocate effect.
+def TransformMappingWrite : MemWrite<TransformMappingResource>;
+
+//===----------------------------------------------------------------------===//
+// Effects on Payload IR.
+//===----------------------------------------------------------------------===//
+
+// Side effect resource corresponding to the Payload IR itself.
+def PayloadIRResource : Resource<"::mlir::transform::PayloadIRResource">;
+
+// Corresponds to the read-only access to the Payload IR through some operation
+// handles in the Transform IR.
+def PayloadIRRead : MemRead<PayloadIRResource>;
+
+// Corresponds to the mutation of the Payload IR through an operation handle in
+// the Transform IR. Should be accompanied by the Read effect for most transform
+// operations (only a complete overwrite of the root op of the Payload IR is a
+// write-only modification).
+def PayloadIRWrite : MemWrite<PayloadIRResource>;
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_EFFECTS_TD

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

 //===- TransformInterfaces.h - Transform Dialect Interfaces -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
 
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
 namespace mlir {
 namespace transform {
 
 class TransformOpInterface;
 
 /// The state maintained across applications of various ops implementing the
 /// TransformOpInterface. The operations implementing this interface and the
 /// surrounding structure are referred to as transform IR. The operations to
public:
   /// state is expected to be already scoped at the region of this operation.
   /// Returns failure if the mapping failed, e.g., the value is already mapped.
   LogicalResult mapBlockArguments(TransformState &state) {
     return detail::mapPossibleTopLevelTransformOpBlockArguments(
         state, this->getOperation());
   }
 };
 
+/// Side effect resource corresponding to the mapping between Transform IR
+/// values and Payload IR operations. An Allocate effect from this resource
+/// means creating a new mapping entry, it is always accompanied by a Write
+/// effet. A Read effect from this resource means accessing the mapping. A Free
+/// effect on this resource indicates the removal of the mapping entry,
+/// typically after a transformation that modifies the Payload IR operations
+/// associated with one of the Transform IR operation's operands. It is always
+/// accompanied by a Read effect. Read-after-Free and double-Free are not
+/// allowed (they would be problematic with "regular" memory effects too) as
+/// they indicate an attempt to access Payload IR operations that have been
+/// modified, potentially erased, by the previous tranfsormations.
+// TODO: consider custom effects if these are not enabling generic passes such
+// as CSE/DCE to work.
+struct TransformMappingResource
+    : public SideEffects::Resource::Base<TransformMappingResource> {
+  StringRef getName() override { return "transform.mapping"; }
+};
+
+/// Side effect resource corresponding to the Payload IR itself. Only Read and
+/// Write effects are expected on this resource, with Write always accompanied
+/// by a Read (short of fully replacing the top-level Payload IR operation, one
+/// cannot modify the Payload IR without reading it first). This is intended
+/// to disallow reordering of Transform IR operations that mutate the Payload IR
+/// while still allowing the reordering of those that only access it.
+struct PayloadIRResource
+    : public SideEffects::Resource::Base<PayloadIRResource> {
+  StringRef getName() override { return "transform.payload_ir"; }
+};
+
 } // namespace transform
 } // namespace mlir
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h.inc"
 
 #endif // DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

 //===- TransformOps.td - Transform dialect operations ------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS
 
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 
 def PDLMatchOp : TransformDialectOp<"pdl_match",
     [DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Finds ops that match the named PDL pattern";
   let description = [{
     Find Payload IR ops nested within the Payload IR op associated with the
     operand that match the PDL pattern identified by its name. The pattern is
     expected to be defined in the closest surrounding `WithPDLPatternsOp`.
 
     Produces a Transform IR value associated with the list of Payload IR ops
     that matched the pattern. The order of results in the list is that of the
     Operation::walk, clients are advised not to rely on a specific order though.
     If the operand is assocaited with multiple Payload IR ops, finds matching
     ops nested within each of those and produces a single list containing all
     of the matched ops.
 
     The tranfsormation is considered successful regardless of whether some
     Payload IR ops actually matched the pattern and only fails if the pattern
     could not be looked up or compiled.
   }];
 
-  let arguments = (ins PDL_Operation:$root, SymbolRefAttr:$pattern_name);
-  let results = (outs PDL_Operation:$matched);
+  let arguments = (ins
+    Arg<PDL_Operation, "Payload IR scope to match within",
+        [TransformMappingRead, PayloadIRRead]>:$root, 
+    SymbolRefAttr:$pattern_name);
+  let results = (outs
+    Res<PDL_Operation, "Handle to the matched Payload IR ops",
+        [TransformMappingAlloc, TransformMappingWrite]>:$matched);
 
   let assemblyFormat = "$pattern_name `in` $root attr-dict";
 }
 
 def SequenceOp : TransformDialectOp<"sequence",
-    [DeclareOpInterfaceMethods<TransformOpInterface>, OpAsmOpInterface,
-     PossibleTopLevelTransformOpTrait,
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     OpAsmOpInterface, PossibleTopLevelTransformOpTrait,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
   let summary = "Contains a sequence of other transform ops to apply";
   let description = [{
     The transformations indicated by the sequence are applied in order of their
     appearance. Each value produced by a transformation within the sequence
     corresponds to an operation or a group of operations in the payload IR.
-    Each value may be used at most once by another transformation operation as
-    the transformation is likely to replace the transformed operation with
-    another operation or a group thereof. In such cases, the transformation
-    operation is expected to produce a new value to denote the newly produced
-    operations that can be transformed further. During application, if any
-    transformation in the sequence fails, the entire sequence fails immediately
-    leaving the payload IR in potentially invalid state, i.e., this operation
-    offers no transformation rollback capabilities.
+    During application, if any transformation in the sequence fails, the entire
+    sequence fails immediately leaving the payload IR in potentially invalid
+    state, i.e., this operation offers no transformation rollback capabilities.
 
     The entry block of this operation has a single argument that maps to either
     the operand if provided or the top-level container operation of the payload
     IR, typically the root operation of the pass interpreting the transform
     dialect. Operand omission is only allowed for sequences not contained in
     another sequence.
   }];
 
let extraClassDeclaration = [{
     static StringRef getDefaultDialect() { return "transform"; }
   }];
 
   let hasVerifier = 1;
 }
 
 def WithPDLPatternsOp : TransformDialectOp<"with_pdl_patterns",
     [DeclareOpInterfaceMethods<TransformOpInterface>, NoTerminator,
-     OpAsmOpInterface, PossibleTopLevelTransformOpTrait, SymbolTable]> {
+     OpAsmOpInterface, PossibleTopLevelTransformOpTrait, RecursiveSideEffects,
+     SymbolTable]> {
   let summary = "Contains PDL patterns available for use in transforms";
   let description = [{
     This op contains a set of named PDL patterns that are available for the
     Transform dialect operations to be used for pattern matching. For example,
     PDLMatchOp can be used to produce a Transform IR value associated with all
     Payload IR operations that match the pattern as follows:
 
     ```mlir
let description = [{
 
     This op is expected to contain `pdl.pattern` operations and exactly one
     another Transform dialect operation that gets executed with all patterns
     available. This op is a possible top-level Transform IR op, the argument of
     its entry block corresponds to either the root op of the payload IR or the
     ops associated with its operand when provided.
   }];
 
-  let arguments = (ins Optional<PDL_Operation>:$root);
+  let arguments = (ins
+    Arg<Optional<PDL_Operation>, "Root operation of the Payload IR",
+        [TransformMappingRead]>:$root);
   let regions = (region SizedRegion<1>:$body);
   let assemblyFormat = "($root^)? attr-dict-with-keyword regions";
 
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     /// Allow the dialect prefix to be omitted.
     static StringRef getDefaultDialect() { return "transform"; }
   }];
 }
 
 def YieldOp : TransformDialectOp<"yield", [Terminator]> {
   let summary = "Yields operation handles from a transform IR region";
   let description = [{
     This terminator operation yields operation handles from regions of the
     transform IR ops back to the containing op. It is not itself associated with
     any transformation on the payload IR and is used for flow purposes only.
   }];
 
-  let arguments = (ins Variadic<AnyType>:$operands);
+  let arguments = (ins 
+    Arg<Variadic<AnyType>, "Opration handles yielded back to the parent",
+        [TransformMappingRead]>:$operands);
   let assemblyFormat = "operands attr-dict (`:` type($operands)^)?";
 
   let builders = [
     OpBuilder<(ins), [{
       return build($_builder, $_state, ::mlir::ValueRange());
     }]>
   ];
 }
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS

mlir/lib/Dialect/Transform/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRTransformDialect
   TransformDialect.cpp
   TransformInterfaces.cpp
   TransformOps.cpp
 
   DEPENDS
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPDL
   MLIRPDLInterp
   MLIRRewrite
+  MLIRSideEffectInterfaces
   )

mlir/lib/Dialect/Transform/IR/TransformDialect.cpp

 //===- TransformDialect.cpp - Transform Dialect Definition ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
 
 using namespace mlir;
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.cpp.inc"
 
 void transform::TransformDialect::initialize() {
-  addOperations<
+  // Using the checked version to enable the same assertions as for the ops from
+  // extensions.
+  addOperationsChecked<
 #define GET_OP_LIST
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
       >();
 }
 
 void transform::TransformDialect::mergeInPDLMatchHooks(
     llvm::StringMap<PDLConstraintFunction> &&constraintFns) {
   // Steal the constraint functions form the given map.
   for (auto &it : constraintFns)
     pdlMatchHooks.registerConstraintFunction(it.getKey(), std::move(it.second));
 }
 
 const llvm::StringMap<PDLConstraintFunction> &
 transform::TransformDialect::getPDLConstraintHooks() const {
   return pdlMatchHooks.getConstraintFunctions();
 }

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

 }
 
 LogicalResult
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   transform::TransformResults results(transform->getNumResults());
   if (failed(transform.apply(results, *this)))
     return failure();
 
-  for (Value target : transform->getOperands())
+  // Remove the mapping for the operand if it is consumed by the operation. This
+  // allows us to catch use-after-free with assertions later on.
+  auto memEffectInterface =
+      cast<MemoryEffectOpInterface>(transform.getOperation());
+  SmallVector<MemoryEffects::EffectInstance, 2> effects;
+  for (Value target : transform->getOperands()) {
+    effects.clear();
+    memEffectInterface.getEffectsOnValue(target, effects);
+    if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
+          return isa<transform::TransformMappingResource>(
+                     effect.getResource()) &&
+                 isa<MemoryEffects::Free>(effect.getEffect());
+        })) {
       removePayloadOps(target);
+    }
+  }
 
   for (auto &en : llvm::enumerate(transform->getResults())) {
     assert(en.value().getDefiningOp() == transform.getOperation() &&
            "payload IR association for a value other than the result of the "
            "current transform op");
     if (failed(setPayloadOps(en.value(), results.get(en.index()))))
       return failure();
   }

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

for (const auto &pair :
     Value terminatorOperand = std::get<0>(pair);
     OpResult result = std::get<1>(pair);
     results.set(result, state.getPayloadOps(terminatorOperand));
   }
 
   return success();
 }
 
+/// Returns `true` if the given op operand may be consuming the handle value in
+/// the Transform IR. That is, if it may have a Free effect on it.
+static bool isValueUsePotentialConsumer(OpOperand &use) {
+  // Conservatively assume the effect being present in absence of the interface.
+  auto memEffectInterface = dyn_cast<MemoryEffectOpInterface>(use.getOwner());
+  if (!memEffectInterface)
+    return true;
+
+  SmallVector<MemoryEffects::EffectInstance, 2> effects;
+  memEffectInterface.getEffectsOnValue(use.get(), effects);
+  return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
+    return isa<transform::TransformMappingResource>(effect.getResource()) &&
+           isa<MemoryEffects::Free>(effect.getEffect());
+  });
+}
+
+LogicalResult
+checkDoubleConsume(Value value,
+                   function_ref<InFlightDiagnostic()> reportError) {
+  OpOperand *potentialConsumer = nullptr;
+  for (OpOperand &use : value.getUses()) {
+    if (!isValueUsePotentialConsumer(use))
+      continue;
+
+    if (!potentialConsumer) {
+      potentialConsumer = &use;
+      continue;
+    }
+
+    InFlightDiagnostic diag = reportError()
+                              << " has more than one potential consumer";
+    diag.attachNote(potentialConsumer->getOwner()->getLoc())
+        << "used here as operand #" << potentialConsumer->getOperandNumber();
+    diag.attachNote(use.getOwner()->getLoc())
+        << "used here as operand #" << use.getOperandNumber();
+    return diag;
+  }
+
+  return success();
+}
+
 LogicalResult transform::SequenceOp::verify() {
+  // Check if the block argument has more than one consuming use.
+  for (BlockArgument argument : getBodyBlock()->getArguments()) {
+    auto report = [&]() {
+      return (emitOpError() << "block argument #" << argument.getArgNumber());
+    };
+    if (failed(checkDoubleConsume(argument, report)))
+      return failure();
+  }
+
+  // Check properties of the nested operations they cannot check themselves.
   for (Operation &child : *getBodyBlock()) {
     if (!isa<TransformOpInterface>(child) &&
         &child != &getBodyBlock()->back()) {
       InFlightDiagnostic diag =
           emitOpError()
           << "expected children ops to implement TransformOpInterface";
       diag.attachNote(child.getLoc()) << "op without interface";
       return diag;
     }
 
     for (OpResult result : child.getResults()) {
-      if (llvm::hasNItemsOrLess(result.getUses(), 1))
-        continue;
-      InFlightDiagnostic diag = child.emitError()
-                                << "result #" << result.getResultNumber()
-                                << " has more than one use";
-      for (OpOperand &use : result.getUses()) {
-        diag.attachNote(use.getOwner()->getLoc())
-            << "used here as operand #" << use.getOperandNumber();
-      }
-      return diag;
+      auto report = [&]() {
+        return (child.emitError() << "result #" << result.getResultNumber());
+      };
+      if (failed(checkDoubleConsume(result, report)))
+        return failure();
     }
   }
 
   if (getBodyBlock()->getTerminator()->getOperandTypes() !=
       getOperation()->getResultTypes()) {
     InFlightDiagnostic diag = emitOpError()
                               << "expects the types of the terminator operands "
                                  "to match the types of the result";
     diag.attachNote(getBodyBlock()->getTerminator()->getLoc()) << "terminator";
     return diag;
   }
   return success();
 }
 
+void transform::SequenceOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  auto *mappingResource = TransformMappingResource::get();
+  effects.emplace_back(MemoryEffects::Read::get(), getRoot(), mappingResource);
+
+  for (Value result : getResults()) {
+    effects.emplace_back(MemoryEffects::Allocate::get(), result,
+                         mappingResource);
+    effects.emplace_back(MemoryEffects::Write::get(), result, mappingResource);
+  }
+
+  if (!getRoot()) {
+    for (Operation &op : *getBodyBlock()) {
+      auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
+      if (!iface) {
+        // TODO: fill all possible effects; or require ops to actually implement
+        // the memory effect interface always
+        assert(false);
+      }
+
+      SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
+      iface.getEffects(effects);
+    }
+    return;
+  }
+
+  // Carry over all effects on the argument of the entry block as those on the
+  // operand, this is the same value just remapped.
+  for (Operation &op : *getBodyBlock()) {
+    auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
+    if (!iface) {
+      // TODO: fill all possible effects; or require ops to actually implement
+      // the memory effect interface always
+      assert(false);
+    }
+
+    SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
+    iface.getEffectsOnValue(getBodyBlock()->getArgument(0), nestedEffects);
+    for (const auto &effect : nestedEffects)
+      effects.emplace_back(effect.getEffect(), getRoot(), effect.getResource());
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // WithPDLPatternsOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 transform::WithPDLPatternsOp::apply(transform::TransformResults &results,
                                     transform::TransformState &state) {
   OwningOpRef<ModuleOp> pdlModuleOp =

mlir/test/Dialect/Transform/ops-invalid.mlir

 transform.sequence {
 ^bb0(%arg0: !pdl.operation):
   // expected-note @below {{op without interface}}
   arith.constant 42.0 : f32
 }
 
 // -----
 
-transform.sequence {
-^bb0(%arg0: !pdl.operation):
-  // expected-error @below {{result #0 has more than one use}}
-  %0 = transform.sequence %arg0 {
-  ^bb1(%arg1: !pdl.operation):
-  } : !pdl.operation
-  // expected-note @below {{used here as operand #0}}
-  transform.sequence %0 {
-  ^bb2(%arg2: !pdl.operation):
-  }
-  // expected-note @below {{used here as operand #0}}
-  transform.sequence %0 {
-  ^bb3(%arg3: !pdl.operation):
-  }
-}
-
-// -----
-
 // expected-error @below {{expects the types of the terminator operands to match the types of the resul}}
 %0 = transform.sequence {
 ^bb0(%arg0: !pdl.operation):
   // expected-note @below {{terminator}}
   transform.yield
 } : !pdl.operation
 
 // -----

 
 // expected-error @below {{expects a single-block region}}
 "transform.test_transform_unrestricted_op_no_interface"() ({
 ^bb0(%arg0: !pdl.operation):
   "test.potential_terminator"() : () -> ()
 ^bb1:
   "test.potential_terminator"() : () -> ()
 }) : () -> ()
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{result #0 has more than one potential consumer}}
+  %0 = test_produce_param_or_forward_operand 42
+  // expected-note @below {{used here as operand #0}}
+  test_consume_operand_if_matches_param_or_fail %0[42]
+  // expected-note @below {{used here as operand #0}}
+  test_consume_operand_if_matches_param_or_fail %0[42]
+}
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{result #0 has more than one potential consumer}}
+  %0 = test_produce_param_or_forward_operand 42
+  // expected-note @below {{used here as operand #0}}
+  test_consume_operand_if_matches_param_or_fail %0[42]
+  // expected-note @below {{used here as operand #0}}
+  transform.sequence %0 {
+  ^bb1(%arg1: !pdl.operation):
+    test_consume_operand_if_matches_param_or_fail %arg1[42]
+  }
+}
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{result #0 has more than one potential consumer}}
+  %0 = test_produce_param_or_forward_operand 42
+  // expected-note @below {{used here as operand #0}}
+  test_consume_operand_if_matches_param_or_fail %0[42]
+  transform.sequence %0 {
+  ^bb1(%arg1: !pdl.operation):
+    // expected-note @below {{used here as operand #0}}
+    test_consume_operand_if_matches_param_or_fail %0[42]
+  }
+}
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{result #0 has more than one potential consumer}}
+  %0 = test_produce_param_or_forward_operand 42
+  // expected-note @below {{used here as operand #0}}
+  test_consume_operand_if_matches_param_or_fail %0[42]
+  // expected-note @below {{used here as operand #0}}
+  transform.sequence %0 {
+  ^bb1(%arg1: !pdl.operation):
+    transform.sequence %arg1 {
+    ^bb2(%arg2: !pdl.operation):
+      test_consume_operand_if_matches_param_or_fail %arg2[42]
+    }
+  }
+}

mlir/test/Dialect/Transform/ops.mlir

 // CHECK: ^{{.+}}(%[[ARG:.+]]: !pdl.operation):
 transform.sequence {
 ^bb0(%arg0: !pdl.operation):
   // CHECK: with_pdl_patterns %[[ARG]]
   with_pdl_patterns %arg0 {
   ^bb1(%arg1: !pdl.operation):
   }
 }
+
+// Using the same value multiple times without consuming it is fine.
+// CHECK: transform.sequence
+// CHECK: %[[V:.+]] = sequence
+// CHECK: sequence %[[V]]
+// CHECK: sequence %[[V]]
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    yield %arg1 : !pdl.operation
+  } : !pdl.operation
+  transform.sequence %0 {
+  ^bb2(%arg2: !pdl.operation):
+  }
+  transform.sequence %0 {
+  ^bb3(%arg3: !pdl.operation):
+  }
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

 
 using namespace mlir;
 
 namespace {
 /// Simple transform op defined outside of the dialect. Just emits a remark when
 /// applied. This op is defined in C++ to test that C++ definitions also work
 /// for op injection into the Transform dialect.
 class TestTransformOp
-    : public Op<TestTransformOp, transform::TransformOpInterface::Trait> {
+    : public Op<TestTransformOp, transform::TransformOpInterface::Trait,
+                MemoryEffectOpInterface::Trait> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTransformOp)
 
   using Op::Op;
 
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static constexpr llvm::StringLiteral getOperationName() {
if (result.getValue().succeeded())
       state.addAttribute("message", message);
     return result.getValue();
   }
 
   void print(OpAsmPrinter &printer) {
     if (getMessage())
       printer << " " << getMessage();
   }
+
+  // No side effects.
+  void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 };
 
 /// A test op to exercise the verifier of the PossibleTopLevelTransformOpTrait
 /// in cases where it is attached to ops that do not comply with the trait
 /// requirements. This op cannot be defined in ODS because ODS generates strict
 /// verifiers that overalp with those in the trait and run earlier.
 class TestTransformUnrestrictedOpNoInterface
     : public Op<TestTransformUnrestrictedOpNoInterface,
                 transform::PossibleTopLevelTransformOpTrait,
-                transform::TransformOpInterface::Trait> {
+                transform::TransformOpInterface::Trait,
+                MemoryEffectOpInterface::Trait> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestTransformUnrestrictedOpNoInterface)
 
   using Op::Op;
 
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static constexpr llvm::StringLiteral getOperationName() {
     return llvm::StringLiteral(
         "transform.test_transform_unrestricted_op_no_interface");
   }
 
   LogicalResult apply(transform::TransformResults &results,
                       transform::TransformState &state) {
     return success();
   }
+
+  // No side effects.
+  void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 };
 } // namespace
 
 LogicalResult mlir::test::TestProduceParamOrForwardOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   if (getOperation()->getNumOperands() != 0) {
     results.set(getResult().cast<OpResult>(),
                 getOperation()->getOperand(0).getDefiningOp());

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTTRANSFORMDIALECTEXTENSION_TD
 #define MLIR_TESTTRANSFORMDIALECTEXTENSION_TD
 
 include "mlir/IR/OpBase.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 
 def TestProduceParamOrForwardOperandOp
   : Op<Transform_Dialect, "test_produce_param_or_forward_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
-  let arguments = (ins Optional<PDL_Operation>:$operand,
-                       OptionalAttr<I64Attr>:$parameter);
-  let results = (outs PDL_Operation:$res);
+  let arguments = (ins 
+     Arg<Optional<PDL_Operation>, "", [TransformMappingRead]>:$operand,
+  let results = (outs 
+     Res<PDL_Operation, "",
+         [TransformMappingAlloc, TransformMappingWrite]>:$res);
   let assemblyFormat = "(`from` $operand^)? ($parameter^)? attr-dict";
   let cppNamespace = "::mlir::test";
   let hasVerifier = 1;
 }
 
 def TestConsumeOperandIfMatchesParamOrFail
   : Op<Transform_Dialect, "test_consume_operand_if_matches_param_or_fail",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
-  let arguments = (ins PDL_Operation:$operand, I64Attr:$parameter);
+  let arguments = (ins 
+    Arg<PDL_Operation, "",
+        [TransformMappingWrite, TransformMappingFree]>:$operand,
+    I64Attr:$parameter);
   let assemblyFormat = "$operand `[` $parameter `]` attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestPrintRemarkAtOperandOp
   : Op<Transform_Dialect, "test_print_remark_at_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
-  let arguments = (ins PDL_Operation:$operand, StrAttr:$message);
+  let arguments = (ins 
+    Arg<PDL_Operation, "",
+        [TransformMappingRead, PayloadIRRead]>:$operand,
+    StrAttr:$message);
   let assemblyFormat = "$operand `,` $message attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 )
 
 td_library(
     name = "TransformDialectTdFiles",
     srcs = glob(["include/mlir/Dialect/Transform/IR/*.td"]),
     deps = [
         ":OpBaseTdFiles",
         ":PDLDialectTdFiles",
+        ":SideEffectInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "TransformDialectInterfacesIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
cc_library(
     name = "TransformDialect",
     srcs = glob(["lib/Dialect/Transform/IR/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Transform/IR/*.h"]),
     deps = [
         ":IR",
         ":PDLDialect",
         ":PDLInterpDialect",
         ":Rewrite",
+        ":SideEffectInterfaces",
         ":Support",
         ":TransformDialectIncGen",
         ":TransformDialectInterfacesIncGen",
         ":TransformOpsIncGen",
         "//llvm:Support",
     ],
 )