[mlir] multi-argument binding for top-level transform ops

[ftynse]

Summary

applyTransforms now takes an optional mapping to be associated with
trailing block arguments of the top-level transform op, in addition to
the payload root. This allows for more advanced forms of communication
between C++ code and the transform dialect interpreter, in particular
supplying operations without having to re-match them during
interpretation.

mlir/docs/Dialects/Transform.md

 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.
 
 A program transformation expressed using the Transform dialect can be
 programmatically triggered by calling:
 
 ```c++
-LogicalResult transform::applyTransforms(Operation *payloadRoot,
-                                         TransformOpInterface transform,
-                                         const TransformOptions &options);
+LogicalResult transform::applyTransforms(
+    Operation *payloadRoot,
+    ArrayRef<ArrayRef<PointerUnion<Operation *, Attribute>> extraMappings,
 ```
 
 that applies the transformations specified by the top-level `transform` to
-payload IR contained in `payloadRoot`.
+payload IR contained in `payloadRoot`. The payload root operation will be
+associated with the first argument of the entry block of the top-level transform
+op. This block may have additional arguments, handles or parameters. They will
+be associated with values provided as `extraMappings`. The call will report an
+error and return if the wrong number of mappings is provided.
 
 ## 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,

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

public:
 
   /// Returns true if the expensive checks are requested.
   bool getExpensiveChecksEnabled() const { return expensiveChecksEnabled; }
 
 private:
   bool expensiveChecksEnabled = true;
 };
 
+using Param = Attribute;
+using MappedValue = llvm::PointerUnion<Operation *, Param>;
+
 /// Entry point to the Transform dialect infrastructure. Applies the
 /// transformation specified by `transform` to payload IR contained in
 /// `payloadRoot`. The `transform` operation may contain other operations that
 /// will be executed following the internal logic of the operation. It must
 /// have the `PossibleTopLevelTransformOp` trait and not have any operands.
 /// This function internally keeps track of the transformation state.
 LogicalResult
 applyTransforms(Operation *payloadRoot, TransformOpInterface transform,
+                ArrayRef<ArrayRef<MappedValue>> extraMapping = {},
                 const TransformOptions &options = TransformOptions());
 
 /// 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
 /// which transformations apply are referred to as payload IR. Transform IR
 /// operates on values that can be associated either with a list of payload IR
 /// operations (such values are referred to as handles) or with a list of

 /// to create a `RegionScope` RAII object to create a new "stack frame" for
 /// values defined inside the region. The mappings from and to these values will
 /// be automatically dropped when the object goes out of scope, typically at the
 /// end of the `apply` function of the parent operation. If a region contains
 /// blocks with arguments, the client can map those arguments to payload IR ops
 /// using `mapBlockArguments`.
 class TransformState {
 public:
-  using Param = Attribute;
+  using Param = transform::Param;
 
 private:
   /// Mapping between a Value in the transform IR and the corresponding set of
   /// operations in the payload IR.
   using TransformOpMapping = DenseMap<Value, SmallVector<Operation *, 2>>;
 
   /// Mapping between a payload IR operation and the transform IR values it is
   /// associated with.
   using TransformOpReverseMapping =
       DenseMap<Operation *, SmallVector<Value, 2>>;
 
   /// Mapping between a Value in the transform IR and the corresponding list of
   /// parameters.
   using ParamMapping = DenseMap<Value, SmallVector<Param>>;
 
   /// The bidirectional mappings between transform IR values and payload IR
   /// operations, and the mapping between transform IR values and parameters.
   struct Mappings {
     TransformOpMapping direct;
     TransformOpReverseMapping reverse;
     ParamMapping params;
   };
 
-  friend LogicalResult applyTransforms(Operation *payloadRoot,
-                                       TransformOpInterface transform,
-                                       const TransformOptions &options);
+  friend LogicalResult applyTransforms(Operation *, TransformOpInterface,
+                                       ArrayRef<ArrayRef<MappedValue>>,
+                                       const TransformOptions &);
 
 public:
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
   Operation *getTopLevel() const;
 
+  /// Returns the number of extra mappings for the top-level operation.
+  size_t getNumTopLevelMappings() const { return topLevelMappedValues.size(); }
+
+  /// Returns the position-th extra mapping for the top-level operation.
+  ArrayRef<MappedValue> getTopLevelMapping(size_t position) const {
+    return topLevelMappedValues[position];
+  }
+
   /// 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;
 
   /// Returns the list of parameters that the given transform IR value
   /// corresponds to.
   ArrayRef<Attribute> getParams(Value value) const;
 
public:
   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);
   }
+  LogicalResult mapBlockArgument(BlockArgument argument,
+                                 ArrayRef<MappedValue> values);
 
   // Forward declarations to support limited visibility.
   class RegionScope;
 
   /// Creates a new region scope for the given region. The region is expected to
   /// be nested in the currently processed region.
   // Implementation note: this method is inline but implemented outside of the
   // class body to comply with visibility and full-declaration requirements.
private:
   /// Identifier for storing top-level value in the `operations` mapping.
   static constexpr Value kTopLevelValue = Value();
 
   /// Creates a state for transform ops living in the given region. The second
   /// argument points to the root operation in the payload IR being transformed,
   /// which may or may not contain the region with transform ops. Additional
   /// options can be provided through the trailing configuration object.
   TransformState(Region *region, Operation *payloadRoot,
+                 ArrayRef<ArrayRef<MappedValue>> extraMappings = {},
                  const TransformOptions &options = TransformOptions());
 
   /// Returns the mappings frame for the reigon in which the value is defined.
   const Mappings &getMapping(Value value) const {
     return const_cast<TransformState *>(this)->getMapping(value);
   }
   Mappings &getMapping(Value value) {
     auto it = mappings.find(value.getParentRegion());
private:
 
   /// Extensions attached to the TransformState, identified by the TypeID of
   /// their type. Only one extension of any given type is allowed.
   DenseMap<TypeID, std::unique_ptr<Extension>> extensions;
 
   /// The top-level operation that contains all payload IR, typically a module.
   Operation *topLevel;
 
+  /// Storage for extra mapped values (payload operations or parameters) to be
+  /// associated with additional entry block arguments of the top-level
+  /// transform operation. Each entry in `topLevelMappedValues` is a reference
+  /// to a contiguous block in `topLevelMappedValueStorage`.
+  // TODO: turn this into a proper named data structure, there are several more
+  // below.
+  SmallVector<ArrayRef<MappedValue>> topLevelMappedValues;
+  SmallVector<MappedValue> topLevelMappedValueStorage;
+
   /// Additional options controlling the transformation state behavior.
   TransformOptions options;
 
   /// The mapping from invalidated handles to the error-reporting functions that
   /// describe when the handles were invalidated. Calling such a function emits
   /// a user-visible diagnostic with an additional note pointing to the given
   /// location.
   DenseMap<Value, std::function<void(Location)>> invalidatedHandles;

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

 namespace mlir {
 namespace transform {
 enum class FailurePropagationMode : uint32_t;
 class FailurePropagationModeAttr;
 
 /// A builder function that populates the body of a SequenceOp.
 using SequenceBodyBuilderFn = ::llvm::function_ref<void(
     ::mlir::OpBuilder &, ::mlir::Location, ::mlir::BlockArgument)>;
+using SequenceBodyBuilderArgsFn =
+    ::llvm::function_ref<void(::mlir::OpBuilder &, ::mlir::Location,
+                              ::mlir::BlockArgument, ::mlir::ValueRange)>;
 } // namespace transform
 } // namespace mlir
 
 #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

 
 def SequenceOp : TransformDialectOp<"sequence",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getSuccessorEntryOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      OpAsmOpInterface, PossibleTopLevelTransformOpTrait,
-     SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
+     SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">,
+     AttrSizedOperandSegments]> {
   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.
     The behavior of the operation when a nested transformation produces a
     silenceable error is controlled by the `failure_propagation_mode` attribute.
     When set to `propagate`, the failure of any nested transformation in the
let description = [{
     another sequence.
 
     The body of the sequence terminates with an implicit or explicit
     `transform.yield` op. The operands of the terminator are returned as the
     results of the sequence op.
   }];
 
   let arguments = (ins FailurePropagationMode:$failure_propagation_mode,
-                       Optional<TransformHandleTypeInterface>:$root);
+                       Optional<TransformHandleTypeInterface>:$root,
+                       Variadic<Transform_AnyHandleOrParamType>:$extra_bindings);
   let results = (outs Variadic<TransformHandleTypeInterface>:$results);
   let regions = (region SizedRegion<1>:$body);
 
   let assemblyFormat =
-    "($root^ `:` type($root))? (`->` type($results)^)? `failures` `(` "
+    "custom<SequenceOpOperands>($root, type($root), $extra_bindings, type($extra_bindings))"
+    " (`->` type($results)^)? `failures` `(` "
     "$failure_propagation_mode `)` attr-dict-with-keyword regions";
 
   let builders = [
     // Build a sequence with a root.
     OpBuilder<(ins
         "::mlir::TypeRange":$resultTypes,
         "::mlir::transform::FailurePropagationMode":$failure_propagation_mode,
         "::mlir::Value":$root, "SequenceBodyBuilderFn":$bodyBuilder)>,
 
-    // Build a sequence without a root but a certain bbArg type.
+    // Build a sequence with a root and additional arguments.
+    OpBuilder<(ins
+        "::mlir::TypeRange":$resultTypes,
+        "::mlir::transform::FailurePropagationMode":$failure_propagation_mode,
+        "::mlir::Value":$root, "::mlir::ValueRange":$extraBindings,
+        "SequenceBodyBuilderArgsFn":$bodyBuilder)>,
+
+    // Build a top-level sequence (no root).
+    OpBuilder<(ins
+        "::mlir::TypeRange":$resultTypes,
+        "::mlir::transform::FailurePropagationMode":$failure_propagation_mode,
+        "::mlir::Type":$bbArgType, "SequenceBodyBuilderFn":$bodyBuilder)>,
+
+    // Build a top-level sequence (no root) with extra arguments.
     OpBuilder<(ins
         "::mlir::TypeRange":$resultTypes,
         "::mlir::transform::FailurePropagationMode":$failure_propagation_mode,
-        "::mlir::Type":$bbArgType, "SequenceBodyBuilderFn":$bodyBuilder)>
+        "::mlir::Type":$bbArgType, "::mlir::TypeRange":$extraBindingTypes,
+        "SequenceBodyBuilderArgsFn":$bodyBuilder)>
   ];
 
   let extraClassDeclaration = [{
     /// Allow the dialect prefix to be omitted.
     static StringRef getDefaultDialect() { return "transform"; }
   }];
 
   let hasVerifier = 1;

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

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TransformState
 //===----------------------------------------------------------------------===//
 
 constexpr const Value transform::TransformState::kTopLevelValue;
 
-transform::TransformState::TransformState(Region *region,
-                                          Operation *payloadRoot,
-                                          const TransformOptions &options)
+transform::TransformState::TransformState(
+    Region *region, Operation *payloadRoot,
+    ArrayRef<ArrayRef<MappedValue>> extraMappings,
     : topLevel(payloadRoot), options(options) {
+  topLevelMappedValues.reserve(extraMappings.size());
+  for (ArrayRef<MappedValue> mapping : extraMappings) {
+    size_t start = topLevelMappedValueStorage.size();
+    llvm::append_range(topLevelMappedValueStorage, mapping);
+    topLevelMappedValues.push_back(
+        ArrayRef<MappedValue>(topLevelMappedValueStorage)
+            .slice(start, mapping.size()));
+  }
+
   auto result = mappings.try_emplace(region);
   assert(result.second && "the region scope is already present");
   (void)result;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   regionStack.push_back(region);
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 }
 
if (iterator != mapping.reverse.end()) {
       found = true;
     }
   }
 
   return success(found);
 }
 
 LogicalResult
+transform::TransformState::mapBlockArgument(BlockArgument argument,
+                                            ArrayRef<MappedValue> values) {
+  if (argument.getType().isa<TransformHandleTypeInterface>()) {
+    SmallVector<Operation *> operations;
+    operations.reserve(values.size());
+    for (MappedValue value : values) {
+      if (auto *op = value.dyn_cast<Operation *>()) {
+        operations.push_back(op);
+        continue;
+      }
+      return emitError(argument.getLoc())
+             << "wrong kind of value provided for top-level operation handle";
+    }
+    return setPayloadOps(argument, operations);
+  }
+
+  assert(argument.getType().isa<TransformParamTypeInterface>() &&
+         "unsupported kind of block argument");
+  SmallVector<Param> parameters;
+  parameters.reserve(values.size());
+  for (MappedValue value : values) {
+    if (auto attr = value.dyn_cast<Attribute>()) {
+      parameters.push_back(attr);
+      continue;
+    }
+    return emitError(argument.getLoc())
+           << "wrong kind of value provided for top-level parameter";
+  }
+  return setParams(argument, parameters);
+}
+
+LogicalResult
 transform::TransformState::setPayloadOps(Value value,
                                          ArrayRef<Operation *> targets) {
   assert(value != kTopLevelValue &&
          "attempting to reset the transformation root");
   assert(!value.getType().isa<TransformParamTypeInterface>() &&
          "cannot associate payload ops with a value of parameter type");
 
   for (Operation *target : targets) {

 
 //===----------------------------------------------------------------------===//
 // Utilities for PossibleTopLevelTransformOpTrait.
 //===----------------------------------------------------------------------===//
 
 LogicalResult transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
     TransformState &state, Operation *op, Region &region) {
   SmallVector<Operation *> targets;
-  if (op->getNumOperands() != 0)
+  SmallVector<SmallVector<MappedValue>> extraMappings;
+  if (op->getNumOperands() != 0) {
     llvm::append_range(targets, state.getPayloadOps(op->getOperand(0)));
-  else
+    for (Value operand : op->getOperands().drop_front()) {
+      SmallVector<MappedValue> &mapped = extraMappings.emplace_back();
+      if (operand.getType().isa<TransformHandleTypeInterface>()) {
+        llvm::append_range(mapped, state.getPayloadOps(operand));
+      } else {
+        assert(operand.getType().isa<TransformParamTypeInterface>() &&
+               "unsupported kind of transform dialect value");
+        llvm::append_range(mapped, state.getParams(operand));
+      }
+    }
+  } else {
+    if (state.getNumTopLevelMappings() !=
+        region.front().getNumArguments() - 1) {
+      return emitError(op->getLoc())
+             << "operation expects " << region.front().getNumArguments() - 1
+             << " extra value bindings, but " << state.getNumTopLevelMappings()
+             << " were provided to the interpreter";
+    }
+
     targets.push_back(state.getTopLevel());
+    for (unsigned i = 0, e = state.getNumTopLevelMappings(); i < e; ++i)
+      extraMappings.push_back(llvm::to_vector(state.getTopLevelMapping(i)));
+  }
+
+  if (failed(state.mapBlockArguments(region.front().getArgument(0), targets)))
+    return failure();
+
+  for (BlockArgument argument : region.front().getArguments().drop_front()) {
+    if (failed(state.mapBlockArgument(
+            argument, extraMappings[argument.getArgNumber() - 1])))
+      return failure();
+  }
 
-  return state.mapBlockArguments(region.front().getArgument(0), targets);
+  return success();
 }
 
 LogicalResult
 transform::detail::verifyPossibleTopLevelTransformOpTrait(Operation *op) {
   // Attaching this trait without the interface is a misuse of the API, but it
   // cannot be caught via a static_assert because interface registration is
   // dynamic.
   assert(isa<TransformOpInterface>(op) &&
          "should implement TransformOpInterface to have "
          "PossibleTopLevelTransformOpTrait");
 
   if (op->getNumRegions() < 1)
     return op->emitOpError() << "expects at least one region";
 
   Region *bodyRegion = &op->getRegion(0);
   if (!llvm::hasNItems(*bodyRegion, 1))
     return op->emitOpError() << "expects a single-block region";
 
   Block *body = &bodyRegion->front();
-  if (body->getNumArguments() != 1 ||
-      !body->getArgumentTypes()[0].isa<TransformHandleTypeInterface>()) {
+  if (body->getNumArguments() == 0) {
     return op->emitOpError()
-           << "expects the entry block to have one argument "
-              "of type implementing TransformHandleTypeInterface";
+           << "expects the entry block to have at least one argument";
+  }
+  if (!body->getArgument(0).getType().isa<TransformHandleTypeInterface>()) {
+    return op->emitOpError()
+           << "expects the first entry block argument to be of type "
+              "implementing TransformHandleTypeInterface";
+  }
+  BlockArgument arg = body->getArgument(0);
+  if (op->getNumOperands() != 0) {
+    if (arg.getType() != op->getOperand(0).getType()) {
+      return op->emitOpError()
+             << "expects the type of the block argument to match "
+                "the type of the operand";
+    }
+  }
+  for (BlockArgument arg : body->getArguments().drop_front()) {
+    if (arg.getType()
+            .isa<TransformHandleTypeInterface, TransformParamTypeInterface>())
+      continue;
+
+    InFlightDiagnostic diag =
+        op->emitOpError()
+        << "expects trailing entry block arguments to be of type implementing "
+           "TransformHandleTypeInterface or TransformParamTypeInterface";
+    diag.attachNote() << "argument #" << arg.getArgNumber() << " does not";
+    return diag;
   }
 
   if (auto *parent =
           op->getParentWithTrait<PossibleTopLevelTransformOpTrait>()) {
-    if (op->getNumOperands() == 0) {
+    if (op->getNumOperands() != body->getNumArguments()) {
       InFlightDiagnostic diag =
           op->emitOpError()
-          << "expects the root operation to be provided for a nested op";
+          << "expects operands to be provided for a nested op";
       diag.attachNote(parent->getLoc())
           << "nested in another possible top-level op";
       return diag;
     }
   }
 
   return success();
 }
LogicalResult transform::detail::verifyTransformOpInterface(Operation *op) {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Entry point.
 //===----------------------------------------------------------------------===//
 
-LogicalResult transform::applyTransforms(Operation *payloadRoot,
-                                         TransformOpInterface transform,
-                                         const TransformOptions &options) {
+LogicalResult
+transform::applyTransforms(Operation *payloadRoot,
+                           ArrayRef<ArrayRef<MappedValue>> extraMapping,
 #ifndef NDEBUG
   if (!transform->hasTrait<PossibleTopLevelTransformOpTrait>() ||
       transform->getNumOperands() != 0) {
     transform->emitError()
         << "expected transform to start at the top-level transform op";
     llvm::report_fatal_error("could not run transforms",
                              /*gen_crash_diag=*/false);
   }
 #endif // NDEBUG
 
-  TransformState state(transform->getParentRegion(), payloadRoot, options);
+  TransformState state(transform->getParentRegion(), payloadRoot, extraMapping,
+                       options);
   return state.applyTransform(transform).checkAndReport();
 }
 
 //===----------------------------------------------------------------------===//
 // Generated interface implementation.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.cpp.inc"

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

 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
 
 using namespace mlir;
 
+static ParseResult parseSequenceOpOperands(
+    OpAsmParser &parser, Optional<OpAsmParser::UnresolvedOperand> &root,
+    Type &rootType,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &extraBindings,
+    SmallVectorImpl<Type> &extraBindingTypes);
+static void printSequenceOpOperands(OpAsmPrinter &printer, Operation *op,
+                                    Value root, Type rootType,
+                                    ValueRange extraBindings,
+                                    TypeRange extraBindingTypes);
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // PatternApplicatorExtension
 //===----------------------------------------------------------------------===//
 
 namespace {
transform::SequenceOp::apply(transform::TransformResults &results,
   }
 
   // Forward the operation mapping for values yielded from the sequence to the
   // values produced by the sequence op.
   forwardTerminatorOperands(getBodyBlock(), state, results);
   return DiagnosedSilenceableFailure::success();
 }
 
+static ParseResult parseSequenceOpOperands(
+    OpAsmParser &parser, Optional<OpAsmParser::UnresolvedOperand> &root,
+    Type &rootType,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &extraBindings,
+    SmallVectorImpl<Type> &extraBindingTypes) {
+  OpAsmParser::UnresolvedOperand rootOperand;
+  OptionalParseResult hasRoot = parser.parseOptionalOperand(rootOperand);
+  if (!hasRoot.has_value()) {
+    root = std::nullopt;
+    return success();
+  }
+  if (failed(hasRoot.value()))
+    return failure();
+  root = rootOperand;
+
+  if (succeeded(parser.parseOptionalComma())) {
+    if (failed(parser.parseOperandList(extraBindings)))
+      return failure();
+  }
+  if (failed(parser.parseColon()))
+    return failure();
+
+  // The paren is truly optional.
+  (void)parser.parseOptionalLParen();
+
+  if (failed(parser.parseType(rootType))) {
+    return failure();
+  }
+
+  if (!extraBindings.empty()) {
+    if (parser.parseComma() || parser.parseTypeList(extraBindingTypes))
+      return failure();
+  }
+
+  if (extraBindingTypes.size() != extraBindings.size()) {
+    return parser.emitError(parser.getNameLoc(),
+                            "expected types to be provided for all operands");
+  }
+
+  // The paren is truly optional.
+  (void)parser.parseOptionalRParen();
+  return success();
+}
+
+static void printSequenceOpOperands(OpAsmPrinter &printer, Operation *op,
+                                    Value root, Type rootType,
+                                    ValueRange extraBindings,
+                                    TypeRange extraBindingTypes) {
+  if (!root)
+    return;
+
+  printer << root;
+  bool hasExtras = !extraBindings.empty();
+  if (hasExtras) {
+    printer << ", ";
+    printer.printOperands(extraBindings);
+  }
+
+  printer << " : ";
+  if (hasExtras)
+    printer << "(";
+
+  printer << rootType;
+  if (hasExtras) {
+    printer << ", ";
+    llvm::interleaveComma(extraBindingTypes, printer.getStream());
+    printer << ")";
+  }
+}
+
 /// 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 iface = dyn_cast<transform::TransformOpInterface>(use.getOwner());
   if (!iface)
     return true;
 
diag.attachNote(use.getOwner()->getLoc())
         << "used here as operand #" << use.getOperandNumber();
     return diag;
   }
 
   return success();
 }
 
 LogicalResult transform::SequenceOp::verify() {
-  assert(getBodyBlock()->getNumArguments() == 1 &&
-         "the number of arguments must have been verified to be 1 by "
+  assert(getBodyBlock()->getNumArguments() >= 1 &&
+         "the number of arguments must have been verified to be more than 1 by "
          "PossibleTopLevelTransformOpTrait");
 
-  BlockArgument arg = getBodyBlock()->getArgument(0);
-  if (getRoot()) {
-    if (arg.getType() != getRoot().getType()) {
-      return emitOpError() << "expects the type of the block argument to match "
-                              "the type of the operand";
-    }
+  if (!getRoot() && !getExtraBindings().empty()) {
+    return emitOpError()
+           << "does not expect extra operands when used as top-level";
   }
 
-  // Check if the block argument has more than one consuming use.
-  if (failed(checkDoubleConsume(
-          arg, [this]() { return (emitOpError() << "block argument #0"); }))) {
+  // Check if a block argument has more than one consuming use.
+  for (BlockArgument arg : getBodyBlock()->getArguments()) {
+    if (failed(checkDoubleConsume(arg, [this, arg]() {
+          return (emitOpError() << "block argument #" << arg.getArgNumber());
+        }))) {
       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";
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(
+/// Appends to `effects` the memory effect instances on `target` with the same
+/// resource and effect as the ones the operation `iface` having on `source`.
+static void
+remapEffects(MemoryEffectOpInterface iface, BlockArgument source, Value target,
              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);
-  }
+  SmallVector<MemoryEffects::EffectInstance> nestedEffects;
+  iface.getEffectsOnValue(source, nestedEffects);
+  for (const auto &effect : nestedEffects)
+    effects.emplace_back(effect.getEffect(), target, effect.getResource());
 
+void transform::SequenceOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getRoot(), effects);
+  onlyReadsHandle(getExtraBindings(), effects);
+  producesHandle(getResults(), effects);
+
   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);
-      }
-
+      auto iface = cast<MemoryEffectOpInterface>(&op);
       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);
-    }
+    auto iface = cast<MemoryEffectOpInterface>(&op);
 
-    SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
-    iface.getEffectsOnValue(getBodyBlock()->getArgument(0), nestedEffects);
-    for (const auto &effect : nestedEffects)
-      effects.emplace_back(effect.getEffect(), getRoot(), effect.getResource());
+    remapEffects(iface, getBodyBlock()->getArgument(0), getRoot(), effects);
+    for (auto [source, target] : llvm::zip(
+             getBodyBlock()->getArguments().drop_front(), getExtraBindings())) {
+      remapEffects(iface, source, target, effects);
+    }
   }
 }
 
 OperandRange transform::SequenceOp::getSuccessorEntryOperands(
     std::optional<unsigned> index) {
   assert(index && *index == 0 && "unexpected region index");
-  if (getOperation()->getNumOperands() == 1)
+  if (getOperation()->getNumOperands() > 0)
     return getOperation()->getOperands();
   return OperandRange(getOperation()->operand_end(),
                       getOperation()->operand_end());
 }
 
 void transform::SequenceOp::getSuccessorRegions(
     std::optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {

 }
 
 void transform::SequenceOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
   (void)operands;
   bounds.emplace_back(1, 1);
 }
 
+template <typename FnTy>
+static void buildSequenceBody(OpBuilder &builder, OperationState &state,
+                              Type bbArgType, TypeRange extraBindingTypes,
+                              FnTy bodyBuilder) {
+  SmallVector<Type> types;
+  types.reserve(1 + extraBindingTypes.size());
+  types.push_back(bbArgType);
+  llvm::append_range(types, extraBindingTypes);
+
+  OpBuilder::InsertionGuard guard(builder);
+  Region *region = state.regions.back().get();
+  Block *bodyBlock = builder.createBlock(region, region->begin(),
+                                         extraBindingTypes, {state.location});
+
+  // Populate body.
+  builder.setInsertionPointToStart(bodyBlock);
+  if constexpr (llvm::function_traits<FnTy>::num_args == 3) {
+    bodyBuilder(builder, state.location, bodyBlock->getArgument(0));
+  } else {
+    bodyBuilder(builder, state.location, bodyBlock->getArgument(0),
+                bodyBlock->getArguments().drop_front());
+  }
+}
+
 void transform::SequenceOp::build(OpBuilder &builder, OperationState &state,
                                   TypeRange resultTypes,
                                   FailurePropagationMode failurePropagationMode,
                                   Value root,
                                   SequenceBodyBuilderFn bodyBuilder) {
-  build(builder, state, resultTypes, failurePropagationMode, root);
-  Region *region = state.regions.back().get();
+  build(builder, state, resultTypes, failurePropagationMode, root,
+        /*extraBindings=*/ValueRange());
   Type bbArgType = root.getType();
-  OpBuilder::InsertionGuard guard(builder);
-  Block *bodyBlock = builder.createBlock(
-      region, region->begin(), TypeRange{bbArgType}, {state.location});
+  buildSequenceBody(builder, state, bbArgType,
+                    /*extraBindingTypes=*/TypeRange(), bodyBuilder);
+}
 
-  // Populate body.
-  builder.setInsertionPointToStart(bodyBlock);
-  bodyBuilder(builder, state.location, bodyBlock->getArgument(0));
+void transform::SequenceOp::build(OpBuilder &builder, OperationState &state,
+                                  TypeRange resultTypes,
+                                  FailurePropagationMode failurePropagationMode,
+                                  Value root, ValueRange extraBindings,
+                                  SequenceBodyBuilderArgsFn bodyBuilder) {
+  build(builder, state, resultTypes, failurePropagationMode, root,
+        extraBindings);
+  buildSequenceBody(builder, state, root.getType(), extraBindings.getTypes(),
+                    bodyBuilder);
 }
 
 void transform::SequenceOp::build(OpBuilder &builder, OperationState &state,
                                   TypeRange resultTypes,
                                   FailurePropagationMode failurePropagationMode,
                                   Type bbArgType,
                                   SequenceBodyBuilderFn bodyBuilder) {
-  build(builder, state, resultTypes, failurePropagationMode, /*root=*/Value());
-  Region *region = state.regions.back().get();
-  OpBuilder::InsertionGuard guard(builder);
-  Block *bodyBlock = builder.createBlock(
-      region, region->begin(), TypeRange{bbArgType}, {state.location});
+  build(builder, state, resultTypes, failurePropagationMode, /*root=*/Value(),
+        /*extraBindings=*/ValueRange());
+  buildSequenceBody(builder, state, bbArgType,
+                    /*extraBindingTypes=*/TypeRange(), bodyBuilder);
+}
 
-  // Populate body.
-  builder.setInsertionPointToStart(bodyBlock);
-  bodyBuilder(builder, state.location, bodyBlock->getArgument(0));
+void transform::SequenceOp::build(OpBuilder &builder, OperationState &state,
+                                  TypeRange resultTypes,
+                                  FailurePropagationMode failurePropagationMode,
+                                  Type bbArgType, TypeRange extraBindingTypes,
+                                  SequenceBodyBuilderArgsFn bodyBuilder) {
+  build(builder, state, resultTypes, failurePropagationMode, /*root=*/Value(),
+        /*extraBindings=*/ValueRange());
+  buildSequenceBody(builder, state, bbArgType, extraBindingTypes, bodyBuilder);
 }
 
 //===----------------------------------------------------------------------===//
 // WithPDLPatternsOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::WithPDLPatternsOp::apply(transform::TransformResults &results,

mlir/python/mlir/dialects/_transform_ops_ext.py

super().__init__(
         [_get_op_result_or_value(h) for h in handles],
         loc=loc,
         ip=ip)
 
 
 class SequenceOp:
 
   def __init__(self, failure_propagation_mode, results: Sequence[Type],
-               target: Union[Operation, Value, Type]):
+               target: Union[Operation, Value, Type],
+               extra_bindings: Optional[Union[Sequence[Value], Sequence[Type],
+                                              Operation, OpView]] = None):
     root = _get_op_result_or_value(target) if isinstance(
         target, (Operation, Value)) else None
     root_type = root.type if not isinstance(target, Type) else target
     if not isinstance(failure_propagation_mode, Attribute):
       failure_propagation_mode_attr = IntegerAttr.get(
           IntegerType.get_signless(32), failure_propagation_mode._as_int())
     else:
       failure_propagation_mode = failure_propagation_mode
+
+    if extra_bindings is None:
+      extra_bindings = []
+    if isinstance(extra_bindings, (Operation, OpView)):
+      extra_bindings = _get_op_results_or_values(extra_bindings)
+
+    extra_binding_types = []
+    if len(extra_bindings) != 0:
+      if isinstance(extra_bindings[0], Type):
+        extra_binding_types = extra_bindings
+        extra_bindings = []
+      else:
+        extra_binding_types = [v.type for v in extra_bindings]
+
     super().__init__(results_=results,
                      failure_propagation_mode=failure_propagation_mode_attr,
-                     root=root)
-    self.regions[0].blocks.append(root_type)
+                     root=root,
+                     extra_bindings=extra_bindings)
+    self.regions[0].blocks.append(*tuple([root_type] + extra_binding_types))
 
   @property
   def body(self) -> Block:
     return self.regions[0].blocks[0]
 
   @property
   def bodyTarget(self) -> Value:
     return self.body.arguments[0]
 
+  @property
+  def bodyExtraArgs(self) -> BlockArgumentList:
+    return self.body.arguments[1:]
+
 
 class WithPDLPatternsOp:
 
   def __init__(self,
                target: Union[Operation, Value, Type],
                *,
                loc=None,
                ip=None):

mlir/test/Dialect/Transform/multi-arg-top-level-ops.mlir

+// RUN: mlir-opt %s --pass-pipeline='builtin.module(test-transform-dialect-interpreter{bind-first-extra-to-ops=func.func bind-second-extra-to-ops=func.return})' \
+// RUN:             --split-input-file --verify-diagnostics
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_op):
+  transform.test_print_remark_at_operand %arg1, "first extra" : !transform.any_op
+  transform.test_print_remark_at_operand %arg2, "second extra" : !transform.any_op
+}
+
+// expected-remark @below {{first extra}}
+func.func @foo() {
+  // expected-remark @below {{second extra}}
+  return
+}
+
+// expected-remark @below {{first extra}}
+func.func @bar(%arg0: i1) {
+  cf.cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  // expected-remark @below {{second extra}}
+  return
+^bb2:
+  // expected-remark @below {{second extra}}
+  return
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.param<i64>):
+  // expected-error @above {{wrong kind of value provided for top-level parameter}}
+}
+
+func.func @foo() {
+  return
+}
+
+// -----
+
+// expected-error @below {{operation expects 1 extra value bindings, but 2 were provided to the interpreter}}
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op):
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_op):
+  transform.sequence %arg0, %arg1, %arg2 : !transform.any_op, !transform.any_op, !transform.any_op failures(propagate) {
+  ^bb0(%arg3: !transform.any_op, %arg4: !transform.any_op, %arg5: !transform.any_op):
+    transform.test_print_remark_at_operand %arg4, "first extra" : !transform.any_op
+    transform.test_print_remark_at_operand %arg5, "second extra" : !transform.any_op
+  }
+}
+
+// expected-remark @below {{first extra}}
+func.func @foo() {
+  // expected-remark @below {{second extra}}
+  return
+}
+
+// expected-remark @below {{first extra}}
+func.func @bar(%arg0: i1) {
+  cf.cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  // expected-remark @below {{second extra}}
+  return
+^bb2:
+  // expected-remark @below {{second extra}}
+  return
+}

mlir/test/Dialect/Transform/multi-arg-top-level-params.mlir

+// RUN: mlir-opt %s --pass-pipeline='builtin.module(test-transform-dialect-interpreter{bind-first-extra-to-params=1,2,3 bind-second-extra-to-params=42,45})' \
+// RUN:          --split-input-file --verify-diagnostics
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation, %arg1: !transform.param<i64>, %arg2: !transform.param<i64>):
+  // expected-remark @below {{1 : i64, 2 : i64, 3 : i64}}
+  transform.test_print_param %arg1 : !transform.param<i64>
+  // expected-remark @below {{42 : i64, 45 : i64}}
+  transform.test_print_param %arg2 : !transform.param<i64>
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation, %arg1: !transform.any_op, %arg2: !transform.param<i64>):
+  // expected-error @above {{wrong kind of value provided for top-level operation handle}}
+}
+
+// -----
+
+// expected-error @below {{operation expects 3 extra value bindings, but 2 were provided to the interpreter}}
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation, %arg1: !transform.param<i64>, %arg2: !transform.param<i64>, %arg3: !transform.param<i64>):
+}

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

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
-// expected-error @below {{expects the entry block to have one argument of type implementing TransformHandleTypeInterface}}
+// expected-error @below {{expects the entry block to have at least one argument}}
 transform.sequence failures(propagate) {
 }
 
 // -----
 
+// expected-error @below {{expects the first entry block argument to be of type implementing TransformHandleTypeInterface}}
+transform.sequence failures(propagate) {
+^bb0(%rag0: i64):
+}
+
+// -----
+
 // expected-note @below {{nested in another possible top-level op}}
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
-  // expected-error @below {{expects the root operation to be provided for a nested op}}
+  // expected-error @below {{expects operands to be provided for a nested op}}
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
   }
 }
 
 // -----
 
+// expected-error @below {{'transform.sequence' op expects trailing entry block arguments to be of type implementing TransformHandleTypeInterface or TransformParamTypeInterface}}
+// expected-note @below {{argument #1 does not}}
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: i64):
+}
+
+// -----
+
 // expected-error @below {{expected children ops to implement TransformOpInterface}}
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-note @below {{op without interface}}
   arith.constant 42.0 : f32
 }
 
 // -----
^bb0(%arg0: !transform.any_op):
   transform.sequence %arg0: !transform.any_op failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     transform.yield
   }
 }
 
 // -----
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_op):
+  // expected-error @below {{expected types to be provided for all operands}}
+  transform.sequence %arg0, %arg1, %arg2 : (!transform.any_op, !transform.any_op) failures(propagate) {
+  ^bb0(%arg3: !transform.any_op, %arg4: !transform.any_op, %arg5: !transform.any_op):
+  }
+}
+
+// -----
+
+%0 = "test.generate_something"() : () -> !transform.any_op
+// expected-error @below {{does not expect extra operands when used as top-level}}
+"transform.sequence"(%0) ({
+^bb0(%arg0: !transform.any_op):
+  "transform.yield"() : () -> ()
+}) {failure_propagation_mode = 1 : i32, operand_segment_sizes = array<i32: 0, 1>} : (!transform.any_op) -> ()
+
+// -----
+
 // expected-note @below {{nested in another possible top-level op}}
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
-  // expected-error @below {{expects the root operation to be provided for a nested op}}
+  // expected-error @below {{expects operands to be provided for a nested op}}
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
   }
 }
 
 // -----
 
 // expected-error @below {{expects only one non-pattern op in its body}}
^bb1(%arg1: !pdl.operation):
   ^bb2(%arg2: !pdl.operation):
     // expected-note @below {{terminator}}
     transform.yield
   }
 }
 
 // -----
 
-// expected-error @below {{expects the entry block to have one argument of type implementing TransformHandleTypeInterface}}
+// expected-error @below {{expects the entry block to have at least one argument}}
 transform.alternatives {
 ^bb0:
   transform.yield
 }
 
 // -----
 
 transform.sequence failures(propagate) {

mlir/test/Dialect/Transform/ops.mlir

^bb0(%arg0: !pdl.operation):
   transform.sequence %0 : !pdl.operation failures(propagate) {
   ^bb2(%arg2: !pdl.operation):
   }
   transform.sequence %0 : !pdl.operation failures(propagate) {
   ^bb3(%arg3: !pdl.operation):
   }
 }
 
+// CHECK: transform.sequence failures(propagate)
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_op):
+  // CHECK: sequence %{{.*}}, %{{.*}}, %{{.*}} : (!transform.any_op, !transform.any_op, !transform.any_op) failures(propagate)
+  transform.sequence %arg0, %arg1, %arg2 : !transform.any_op, !transform.any_op, !transform.any_op failures(propagate) {
+  ^bb0(%arg3: !transform.any_op, %arg4: !transform.any_op, %arg5: !transform.any_op):
+  }
+}
+
+// CHECK: transform.sequence failures(propagate)
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_op):
+  // CHECK: sequence %{{.*}}, %{{.*}}, %{{.*}} : (!transform.any_op, !transform.any_op, !transform.any_op) failures(propagate)
+  transform.sequence %arg0, %arg1, %arg2 : (!transform.any_op, !transform.any_op, !transform.any_op) failures(propagate) {
+  ^bb0(%arg3: !transform.any_op, %arg4: !transform.any_op, %arg5: !transform.any_op):
+  }
+}
+
 // CHECK: transform.sequence
 // CHECK: foreach
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   transform.foreach %arg0 : !pdl.operation {
   ^bb1(%arg1: !pdl.operation):
   }
 }

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

 //===- TestTransformDialectInterpreter.cpp --------------------------------===//
 //
 // 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 a test pass that interprets Transform dialect operations in
 // the module.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple pass that applies transform dialect ops directly contained in a
 /// module.
public:
   StringRef getArgument() const override {
     return "test-transform-dialect-interpreter";
   }
 
   StringRef getDescription() const override {
     return "apply transform dialect operations one by one";
   }
 
+  ArrayRef<transform::MappedValue>
+  findOperationsByName(Operation *root, StringRef name,
+                       SmallVectorImpl<transform::MappedValue> &storage) {
+    size_t start = storage.size();
+    root->walk([&](Operation *op) {
+      if (op->getName().getStringRef() == name) {
+        storage.push_back(op);
+      }
+    });
+    return ArrayRef(storage).drop_front(start);
+  }
+
+  ArrayRef<transform::MappedValue>
+  createParameterMapping(MLIRContext &context, ArrayRef<int> values,
+                         SmallVectorImpl<transform::MappedValue> &storage) {
+    size_t start = storage.size();
+    llvm::append_range(storage, llvm::map_range(values, [&](int v) {
+                         Builder b(&context);
+                         return transform::MappedValue(b.getI64IntegerAttr(v));
+                       }));
+    return ArrayRef(storage).drop_front(start);
+  }
+
   void runOnOperation() override {
+    if (!bindFirstExtraToOps.empty() && !bindFirstExtraToParams.empty()) {
+      emitError(UnknownLoc::get(&getContext()))
+          << "cannot bind the first extra top-level argument to both "
+             "operations and parameters";
+      return signalPassFailure();
+    }
+    if (!bindSecondExtraToOps.empty() && !bindSecondExtraToParams.empty()) {
+      emitError(UnknownLoc::get(&getContext()))
+          << "cannot bind the second extra top-level argument to both "
+             "operations and parameters";
+      return signalPassFailure();
+    }
+    if ((!bindSecondExtraToOps.empty() || !bindSecondExtraToParams.empty()) &&
+        bindFirstExtraToOps.empty() && bindFirstExtraToParams.empty()) {
+      emitError(UnknownLoc::get(&getContext()))
+          << "cannot bind the second extra top-level argument without binding "
+             "the first";
+      return signalPassFailure();
+    }
+
+    SmallVector<transform::MappedValue> extraMappingStorage;
+    SmallVector<ArrayRef<transform::MappedValue>> extraMapping;
+    if (!bindFirstExtraToOps.empty()) {
+      extraMapping.push_back(findOperationsByName(
+          getOperation(), bindFirstExtraToOps.getValue(), extraMappingStorage));
+    } else if (!bindFirstExtraToParams.empty()) {
+      extraMapping.push_back(createParameterMapping(
+          getContext(), bindFirstExtraToParams, extraMappingStorage));
+    }
+    if (!bindSecondExtraToOps.empty()) {
+      extraMapping.push_back(findOperationsByName(
+          getOperation(), bindSecondExtraToOps, extraMappingStorage));
+    } else if (!bindSecondExtraToParams.empty()) {
+      extraMapping.push_back(createParameterMapping(
+          getContext(), bindSecondExtraToParams, extraMappingStorage));
+    }
+
     ModuleOp module = getOperation();
     for (auto op :
          module.getBody()->getOps<transform::TransformOpInterface>()) {
       if (failed(transform::applyTransforms(
-              module, op,
+              module, op, extraMapping,
               transform::TransformOptions().enableExpensiveChecks(
                   enableExpensiveChecks))))
         return signalPassFailure();
     }
   }
 
   Option<bool> enableExpensiveChecks{
       *this, "enable-expensive-checks", llvm::cl::init(false),
       llvm::cl::desc("perform expensive checks to better report errors in the "
                      "transform IR")};
+
+  Option<std::string> bindFirstExtraToOps{
+      *this, "bind-first-extra-to-ops",
+      llvm::cl::desc("bind the first extra argument of the top-level op to "
+                     "payload operations of the given kind")};
+  ListOption<int> bindFirstExtraToParams{
+      *this, "bind-first-extra-to-params",
+      llvm::cl::desc("bind the first extra argument of the top-level op to "
+                     "the given integer parameters")};
+
+  Option<std::string> bindSecondExtraToOps{
+      *this, "bind-second-extra-to-ops",
+      llvm::cl::desc("bind the second extra argument of the top-level op to "
+                     "payload operations of the given kind")};
+  ListOption<int> bindSecondExtraToParams{
+      *this, "bind-second-extra-to-params",
+      llvm::cl::desc("bind the second extra argument of the top-level op to "
+                     "the given integer parameters")};
 };
 
 struct TestTransformDialectEraseSchedulePass
     : public PassWrapper<TestTransformDialectEraseSchedulePass,
                          OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestTransformDialectEraseSchedulePass)
 

mlir/test/python/dialects/transform.py

def testNestedSequenceOp():
   # CHECK:     ^{{.*}}(%[[ARG2:.+]]: !pdl.operation):
   # CHECK:       yield %[[ARG2]] : !pdl.operation
   # CHECK:     }
   # CHECK:   }
   # CHECK: }
 
 
 @run
+def testSequenceOpWithExtras():
+  sequence = transform.SequenceOp(
+      transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get(),
+      [transform.AnyOpType.get(),
+       transform.OperationType.get("foo.bar")])
+  with InsertionPoint(sequence.body):
+    transform.YieldOp()
+  # CHECK-LABEL: TEST: testSequenceOpWithExtras
+  # CHECK: transform.sequence failures(propagate)
+  # CHECK: ^{{.*}}(%{{.*}}: !transform.any_op, %{{.*}}: !transform.any_op, %{{.*}}: !transform.op<"foo.bar">):
+
+
+@run
+def testNestedSequenceOpWithExtras():
+  sequence = transform.SequenceOp(
+      transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get(),
+      [transform.AnyOpType.get(),
+       transform.OperationType.get("foo.bar")])
+  with InsertionPoint(sequence.body):
+    nested = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], sequence.bodyTarget,
+                                  sequence.bodyExtraArgs)
+    with InsertionPoint(nested.body):
+      transform.YieldOp()
+    transform.YieldOp()
+  # CHECK-LABEL: TEST: testNestedSequenceOpWithExtras
+  # CHECK: transform.sequence failures(propagate)
+  # CHECK: ^{{.*}}(%[[ARG0:.*]]: !transform.any_op, %[[ARG1:.*]]: !transform.any_op, %[[ARG2:.*]]: !transform.op<"foo.bar">):
+  # CHECK:   sequence %[[ARG0]], %[[ARG1]], %[[ARG2]] : (!transform.any_op, !transform.any_op, !transform.op<"foo.bar">)
+
+
+@run
 def testTransformPDLOps():
   withPdl = transform.WithPDLPatternsOp(pdl.OperationType.get())
   with InsertionPoint(withPdl.body):
     sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
                                     [pdl.OperationType.get()],
                                     withPdl.bodyTarget)
     with InsertionPoint(sequence.body):
       match = transform.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "pdl_matcher")