[mlir][transform] Introduce transform.sequence op

[ftynse]

Summary

Sequence is an important transform combination primitive that just indicates
transform ops being applied in a row. The simplest version requires fails
immediately if any transformation in the sequence fails. Introducing this
operation allows one to start placing transform IR within other IR.

Depends On D123135

[ftynse]

More cmake.

[Mogball]

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

22

nit

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

102

editor struggles :(

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

53

nit

mlir/test/Dialect/Transform/ops.mlir

13

nit

[ftynse]

Address review.

[ftynse]

Better documentation.

[Mogball]

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

75

Why does this need to be wrapped in this def?

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

25

nit

[rriddle]

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

176

Half written sentence?

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

116–119

The general strategy is to const_cast and call the non-const, instead of duplicating the method.

128

Mappins?

129

Why not just friend TransformState?

166

Can you document this field? and the one below?

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

28

tranfsormation

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

69

[ftynse]

Address review.

[ftynse]

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

75

Because the definition of regionStack is also wrapped, and data members of a class are ABI-breaking (or at least quickly become such).

128

Muffins :)

129

Just a preference for tighter visibility control. I can change if you prefer otherwise.

[ftynse]

CMake.

[tschuett]

opertaion -> operation in the summary.

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

-# The dialect does not have its own ops, so just generate the dialect files.
+# Generate the dialect files from the dialect .td.
+#
+# TODO: Make it possible to use XDialect instead of XOpsDialect in
+# add_mlir_dialect.
 set(LLVM_TARGET_DEFINITIONS TransformDialect.td)
 mlir_tablegen(TransformDialect.h.inc -gen-dialect-decls -dialect=transform)
 mlir_tablegen(TransformDialect.cpp.inc -gen-dialect-defs -dialect=transform)
 add_public_tablegen_target(MLIRTransformDialectIncGen)
 add_dependencies(mlir-headers MLIRTransformDialectIncGen)
 
+add_mlir_dialect(TransformOps transform)
+
 add_mlir_interface(TransformInterfaces)

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

let description = [{
 
     For the sake of reusability, transformations should be implemented as
     utility functions that are called from the interface methods of transform
     ops rather than having the methods directly act on the payload IR.
   }];
 
   let name = "transform";
   let cppNamespace = "::mlir::transform";
+  let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 
   let extraClassDeclaration = [{
     // Make addOperations available to the TransformDialectExtension class.
     private:
       using ::mlir::Dialect::addOperations;
 
       template <typename, typename...>
       friend class TransformDialectExtension;
   }];
 }
 
+// An op in the transform dialect. Ops defined for extensions
+class TransformDialectOp<string mnemonic, list<Trait> traits = []>
+    : Op<Transform_Dialect, mnemonic, traits>;
+
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT

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

class TransformState {
   /// Mapping between a Value in the transform IR and the corresponding set of
   /// operations in the payload IR.
   using TransformOpMapping = DenseMap<Value, SmallVector<Operation *>>;
 
   /// Mapping between a payload IR operation and the transform IR value it is
   /// currently associated with.
   using TransformOpReverseMapping = DenseMap<Operation *, Value>;
 
+  struct Mappings {
+    TransformOpMapping direct;
+    TransformOpReverseMapping reverse;
+  };
+
 public:
-  /// Creates a state for the transformation rooted at the given op.
-  explicit TransformState(Operation *root);
+  /// Creates a state for transform ops living in the given region. The parent
+  /// operation of the region. The second argument points to the root operation
+  /// in the payload IR beind transformed, which may or may not contain the
+  /// region with transform ops.
+  TransformState(Region &region, Operation *root);
 
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
   Operation *getTopLevel() const;
 
   /// Returns the list of ops that the given transform IR value corresponds to.
   /// This is helpful for transformations that apply to a particular handle.
   ArrayRef<Operation *> getPayloadOps(Value value) const;
 
   /// Applies the transformation specified by the given transform op and updates
   /// the state accordingly.
   LogicalResult applyTransform(TransformOpInterface transform);
 
+  LogicalResult mapBlockArguments(BlockArgument argument,
+                                  ArrayRef<Operation *> operations) {
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+    assert(argument.getParentRegion() == regionStack.back() &&
+           "mapping block arguments from a region other than the active one");
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
+    return setPayloadOps(argument, operations);
+  }
+
+  class RegionScope;
+  inline RegionScope make_region_scope(Region &region);
+
+  class RegionScope {
+  public:
+    ~RegionScope() {
+      state.mappings.erase(region);
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+      state.regionStack.pop_back();
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
+    }
+
+  private:
+    RegionScope(TransformState &state, Region &region)
+        : state(state), region(&region) {
+      auto res = state.mappings.try_emplace(this->region);
+      assert(res.second && "the region scope is already present");
+      (void)res;
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+      assert(state.regionStack.back()->isProperAncestor(&region) &&
+             "scope started at a non-nested region");
+      state.regionStack.push_back(&region);
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
+    }
+
+    TransformState &state;
+    Region *region;
+
+    friend RegionScope TransformState::make_region_scope(Region &);
+  };
+  friend class RegionScope;
+
 private:
   /// Identifier for storing top-level value in the `operations` mapping.
   static constexpr Value kTopLevelValue = Value();
 
+  const Mappings &getMapping(Value value) const {
+    auto it = mappings.find(value.getParentRegion());
+    assert(it != mappings.end() &&
+           "trying to find a mapping for a value from an unmapped region");
+    return it->second;
+  }
+
+  Mappings &getMapping(Value value) {
+    auto it = mappings.find(value.getParentRegion());
+    assert(it != mappings.end() &&
+           "trying to find a mapping for a value from an unmapped region");
+    return it->second;
+  }
+
+  const Mappings &getMapping(Operation *operation) const {
+    auto it = mappings.find(operation->getParentRegion());
+    assert(it != mappings.end() &&
+           "trying to find a mapping for an operation from an unmapped region");
+    return it->second;
+  }
+
   /// Sets the payload IR ops associated with the given transform IR value.
   /// Fails if this would result in multiple transform IR values with uses
   /// corresponding to the same payload IR ops. For example, a hypothetical
   /// "find function by name" transform op would (indirectly) call this
   /// function for its result. Having two such calls in a row with for different
   /// values, e.g. coming from different ops:
   ///
   ///   %0 = transform.find_func_by_name { name = "myfunc" }
private:
   /// The callback function is called once per associated operation and is
   /// expected to return the modified operation or nullptr. In the latter case,
   /// the corresponding operation is no longer associated with the transform IR
   /// value.
   void updatePayloadOps(Value value,
                         function_ref<Operation *(Operation *)> callback);
 
   /// The mapping between payload IR values and transform IR ops.
-  TransformOpMapping operationMapping;
-  TransformOpReverseMapping reverseMapping;
+  llvm::SmallDenseMap<Region *, Mappings> mappings;
+
+  Operation *topLevel;
+
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+  SmallVector<Region *> regionStack;
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 };
 
 /// Local mapping between values defined by a specific op implementing the
 /// TransformOpInterface and the payload IR ops they correspond to.
 class TransformResults {
   friend class TransformState;
 
 public:
private:
   /// op has results. Each entry is a reference to a contiguous segment in
   /// the `operations` list that contains the pointers to operations. This
   /// allows for operations to be stored contiguously without nested vectors and
   /// for different segments to be set in any order.
   SmallVector<ArrayRef<Operation *>, 2> segments;
   SmallVector<Operation *> operations;
 };
 
+TransformState::RegionScope TransformState::make_region_scope(Region &region) {
+  return RegionScope(*this, region);
+}
+
 } // 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.h

+//===- TransformDialect.h - Transform dialect operations --------*- 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_TRANSFORMOPS_H
+#define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
+
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformOps.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_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/Dialect/PDL/IR/PDLTypes.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+
+def SequenceOp : TransformDialectOp<"sequence",
+    [DeclareOpInterfaceMethods<TransformOpInterface>, OpAsmOpInterface,
+     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 tranfsormation
+    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.
+
+    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 arguments = (ins Optional<PDL_Operation>:$root);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat =
+    "($root^)? attr-dict-with-keyword regions (`:` type($results)^)?";
+
+  let extraClassDeclaration = [{
+    /// Allow the dialect prefix to be omitted.
+    static StringRef getDefaultDialect() { return "transform"; }
+
+    Block *getBodyBlock() {
+      return &getBody().front();
+    }
+  }];
+
+  let hasVerifier = 1;
+}
+
+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 assemblyFormat = "operands attr-dict (`:` type($operands)^)?";
+
+  let builders = [
+    OpBuilder<(ins), [{
+      return build($_builder, $_state, ::mlir::ValueRange());
+    }]>
+  ];
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS
 No newline at end of file

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/TransformDialect.cpp.inc"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
 
 using namespace mlir;
 
-void transform::TransformDialect::initialize() {}
+#include "mlir/Dialect/Transform/IR/TransformDialect.cpp.inc"
+
+void transform::TransformDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
+      >();
+}

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

 //===- TransformInterfaces.cpp - Transform Dialect Interfaces -------------===//
 //
 // 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/TransformInterfaces.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Operation.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TransformState
 //===----------------------------------------------------------------------===//
 
 constexpr const Value transform::TransformState::kTopLevelValue;
 
-transform::TransformState::TransformState(Operation *root) {
-  operationMapping[kTopLevelValue].push_back(root);
+transform::TransformState::TransformState(Region &region, Operation *root)
+    : topLevel(root) {
+  auto res = mappings.try_emplace(&region);
+  assert(res.second && "the region scope is already present");
+  (void)res;
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+  regionStack.push_back(&region);
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 }
 
-Operation *transform::TransformState::getTopLevel() const {
-  return operationMapping.lookup(kTopLevelValue).front();
-}
+Operation *transform::TransformState::getTopLevel() const { return topLevel; }
 
 ArrayRef<Operation *>
 transform::TransformState::getPayloadOps(Value value) const {
+  const TransformOpMapping &operationMapping = getMapping(value).direct;
   auto iter = operationMapping.find(value);
   assert(iter != operationMapping.end() && "unknown handle");
   return iter->getSecond();
 }
 
 LogicalResult
 transform::TransformState::setPayloadOps(Value value,
                                          ArrayRef<Operation *> targets) {
   assert(value != kTopLevelValue &&
          "attempting to reset the transformation root");
 
   if (value.use_empty())
     return success();
 
   // Setting new payload for the value without cleaning it first is a misuse of
   // the API, assert here.
   SmallVector<Operation *> storedTargets(targets.begin(), targets.end());
+  Mappings &mappings = getMapping(value);
   bool inserted =
-      operationMapping.insert({value, std::move(storedTargets)}).second;
+      mappings.direct.insert({value, std::move(storedTargets)}).second;
   assert(inserted && "value is already associated with another list");
   (void)inserted;
 
   // Having multiple handles to the same operation is an error in the transform
   // expressed using the dialect and may be constructed by valid API calls from
   // valid IR. Emit an error here.
   for (Operation *op : targets) {
-    auto insertionResult = reverseMapping.insert({op, value});
+    auto insertionResult = mappings.reverse.insert({op, value});
     if (!insertionResult.second) {
       InFlightDiagnostic diag = op->emitError()
                                 << "operation tracked by two handles";
       diag.attachNote(value.getLoc()) << "handle";
       diag.attachNote(insertionResult.first->second.getLoc()) << "handle";
       return diag;
     }
   }
 
   return success();
 }
 
 void transform::TransformState::removePayloadOps(Value value) {
-  for (Operation *op : operationMapping[value])
-    reverseMapping.erase(op);
-  operationMapping.erase(value);
+  Mappings &mappings = getMapping(value);
+  for (Operation *op : mappings.direct[value])
+    mappings.reverse.erase(op);
+  mappings.direct.erase(value);
 }
 
 void transform::TransformState::updatePayloadOps(
     Value value, function_ref<Operation *(Operation *)> callback) {
-  auto it = operationMapping.find(value);
-  assert(it != operationMapping.end() && "unknown handle");
+  auto it = getMapping(value).direct.find(value);
+  assert(it != getMapping(value).direct.end() && "unknown handle");
   SmallVector<Operation *> &association = it->getSecond();
   SmallVector<Operation *> updated;
   updated.reserve(association.size());
 
   for (Operation *op : association)
     if (Operation *updatedOp = callback(op))
       updated.push_back(updatedOp);
 
   std::swap(association, updated);
 }
 
 LogicalResult
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   transform::TransformResults results(transform->getNumResults());
   if (failed(transform.apply(results, *this)))
     return failure();
 
   for (Value target : transform->getOperands())
     removePayloadOps(target);
 
-  for (auto &en : llvm::enumerate(transform->getResults()))
+  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();
+  }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TransformResults
 //===----------------------------------------------------------------------===//
 

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

+//===- TransformDialect.cpp - Transform dialect operations ----------------===//
+//
+// 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/TransformOps.h"
+#include "mlir/IR/Builders.h"
+
+#include "mlir/IR/OpImplementation.h"
+
+using namespace mlir;
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
+
+LogicalResult transform::SequenceOp::apply(transform::TransformResults &results,
+                                           transform::TransformState &state) {
+  SmallVector<Operation *> targets;
+  if (getRoot())
+    llvm::append_range(targets, state.getPayloadOps(getRoot()));
+  else
+    targets.push_back(state.getTopLevel());
+
+  // Map the entry block argument to the list of operations.
+  auto scope = state.make_region_scope(*getBodyBlock()->getParent());
+  if (failed(state.mapBlockArguments(getBodyBlock()->getArgument(0), targets)))
+    return failure();
+
+  // Apply the sequenced ops one by one.
+  for (Operation &transform : getBodyBlock()->without_terminator())
+    if (failed(state.applyTransform(cast<TransformOpInterface>(transform))))
+      return failure();
+
+  // Forward the operation mapping for values yielded from the sequence to the
+  // values produced by the sequence op.
+  for (const auto &pair :
+       llvm::zip(getBodyBlock()->getTerminator()->getOperands(),
+                 getOperation()->getOpResults())) {
+    Value terminatorOperand = std::get<0>(pair);
+    OpResult result = std::get<1>(pair);
+    results.set(result, state.getPayloadOps(terminatorOperand));
+  }
+
+  return success();
+}
+
+LogicalResult transform::SequenceOp::verify() {
+  if (getBodyBlock()->getNumArguments() != 1 ||
+      !getBodyBlock()->getArgumentTypes()[0].isa<pdl::OperationType>()) {
+    return emitOpError()
+           << "expected the entry block to have one argument of type "
+           << pdl::OperationType::get(getContext());
+  }
+
+  if (auto parent = getOperation()->getParentOfType<transform::SequenceOp>()) {
+    if (!getRoot()) {
+      InFlightDiagnostic diag =
+          emitOpError()
+          << "expected the root operation to be provided for a nested sequence";
+      diag.attachNote(parent.getLoc()) << "nested in another sequence";
+      return diag;
+    }
+  }
+
+  for (Operation &child : getBodyBlock()->getOperations()) {
+    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;
+    }
+  }
+
+  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();
+}

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

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error @below {{expected the entry block to have one argument of type '!pdl.operation'}}
+transform.sequence {
+}
+
+// -----
+
+// expected-note @below {{nested in another sequence}}
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{expected the root operation to be provided for a nested sequence}}
+  transform.sequence {
+  ^bb1(%arg1: !pdl.operation):
+  }
+}
+
+// -----
+
+// expected-error @below {{expected children ops to implement TransformOpInterface}}
+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

mlir/test/Dialect/Transform/ops.mlir

+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK: transform.sequence
+// CHECK: ^{{.+}}(%{{.+}}: !pdl.operation):
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // CHECK: sequence %{{.+}}
+  // CHECK: ^{{.+}}(%{{.+}}: !pdl.operation):
+  sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+  }
+}
 No newline at end of file

mlir/test/Dialect/Transform/test-interpreter.mlir

 // expected-error @below {{operation tracked by two handles}}
 %0 = transform.test_produce_param_or_forward_operand 42
 // expected-note @below {{handle}}
 %1 = transform.test_produce_param_or_forward_operand from %0
 // expected-note @below {{handle}}
 %2 = transform.test_produce_param_or_forward_operand from %0
 transform.test_consume_operand_if_matches_param_or_fail %1[42]
 transform.test_consume_operand_if_matches_param_or_fail %2[42]
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  sequence %arg0 {
+  ^bb0(%arg1: !pdl.operation):
+    // expected-remark @below {{applying transformation "a"}}
+    test_transform_op "a"
+    // expected-remark @below {{applying transformation "b"}}
+    test_transform_op "b"
+    // expected-remark @below {{applying transformation "c"}}
+    test_transform_op "c"
+  }
+  // expected-remark @below {{applying transformation "d"}}
+  test_transform_op "d"
+  // expected-remark @below {{applying transformation "e"}}
+  test_transform_op "e"
+}
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  %0 = test_produce_param_or_forward_operand 42
+  sequence %0 {
+  ^bb0(%arg1: !pdl.operation):
+    // expected-remark @below {{succeeded}}
+    test_consume_operand_if_matches_param_or_fail %arg1[42]
+  }
+}
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  %0 = sequence %arg0 {
+  ^bb0(%arg1: !pdl.operation):
+    %1 = test_produce_param_or_forward_operand 42
+    yield %1 : !pdl.operation
+  } : !pdl.operation
+  // expected-remark @below {{succeeded}}
+  test_consume_operand_if_matches_param_or_fail %0[42]
+}

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

public:
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static constexpr llvm::StringLiteral getOperationName() {
     return llvm::StringLiteral("transform.test_transform_op");
   }
 
   LogicalResult apply(transform::TransformResults &results,
                       transform::TransformState &state) {
-    emitRemark() << "applying transformation";
+    InFlightDiagnostic remark = emitRemark() << "applying transformation";
+    if (Attribute message = getMessage())
+      remark << " " << message;
+
     return success();
   }
 
+  Attribute getMessage() { return getOperation()->getAttr("message"); }
+
   static ParseResult parse(OpAsmParser &parser, OperationState &state) {
+    StringAttr message;
+    OptionalParseResult result = parser.parseOptionalAttribute(message);
+    if (!result.hasValue())
       return success();
+
+    if (result.getValue().succeeded())
+      state.addAttribute("message", message);
+    return result.getValue();
   }
 
-  void print(OpAsmPrinter &printer) {}
+  void print(OpAsmPrinter &printer) {
+    if (getMessage())
+      printer << " " << getMessage();
+  }
 };
 } // namespace
 
 LogicalResult mlir::test::TestProduceParamOrForwardOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   if (getOperation()->getNumOperands() != 0) {
-    results.set(getResult().cast<OpResult>(), getOperand(0).getDefiningOp());
+    results.set(getResult().cast<OpResult>(),
+                getOperation()->getOperand(0).getDefiningOp());
   } else {
     results.set(getResult().cast<OpResult>(),
-                reinterpret_cast<Operation *>(*parameter()));
+                reinterpret_cast<Operation *>(*getParameter()));
   }
   return success();
 }
 
 LogicalResult mlir::test::TestProduceParamOrForwardOperandOp::verify() {
-  if (parameter().hasValue() ^ (getNumOperands() != 1))
+  if (getParameter().hasValue() ^ (getNumOperands() != 1))
     return emitOpError() << "expects either a parameter or an operand";
   return success();
 }
 
 LogicalResult mlir::test::TestConsumeOperandIfMatchesParamOrFail::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   ArrayRef<Operation *> payload = state.getPayloadOps(getOperand());
   assert(payload.size() == 1 && "expected a single target op");
   auto value = reinterpret_cast<intptr_t>(payload[0]);
-  if (value != parameter()) {
+  if (value != static_cast<intptr_t>(getParameter())) {
     return emitOpError() << "expected the operand to be associated with "
-                         << parameter() << " got " << value;
+                         << getParameter() << " got " << value;
   }
 
   emitRemark() << "succeeded";
   return success();
 }
 
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and

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

public:
   }
 
   StringRef getDescription() const override {
     return "apply transform dialect operations one by one";
   }
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
-    transform::TransformState state(module);
+    transform::TransformState state(module.getBodyRegion(), module);
     for (auto op :
          module.getBody()->getOps<transform::TransformOpInterface>()) {
       if (failed(state.applyTransform(op)))
         return signalPassFailure();
     }
   }
 };
 } // namespace

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

cc_library(
     ],
 )
 
 td_library(
     name = "TransformDialectTdFiles",
     srcs = glob(["include/mlir/Dialect/Transform/IR/*.td"]),
     deps = [
         ":OpBaseTdFiles",
+        ":PDLDialectTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "TransformDialectInterfacesIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
tbl_outs = [
             "include/mlir/Dialect/Transform/IR/TransformDialect.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Transform/IR/TransformDialect.td",
     deps = [":TransformDialectTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "TransformOpsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/mlir/Dialect/Transform/IR/TransformOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/mlir/Dialect/Transform/IR/TransformOps.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Transform/IR/TransformOps.td",
+    deps = [":TransformDialectTdFiles"],
+)
+
 cc_library(
     name = "TransformDialect",
     srcs = glob(["lib/Dialect/Transform/IR/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Transform/IR/*.h"]),
     deps = [
         ":IR",
+        ":PDLDialect",
         ":Support",
         ":TransformDialectIncGen",
         ":TransformDialectInterfacesIncGen",
+        ":TransformOpsIncGen",
         "//llvm:Support",
     ],
 )
 
 td_library(
     name = "ComplexOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/Complex/IR/ComplexBase.td",