[mlir] add readonly/consume annotations to transform named sequences

[ftynse]

Summary

Use the argument attribute mechanism for function-like operations to
annotate the arguments of named transform sequences as consuming or only
reading the handles passed as arguments. This makes it possible to
correctly specify handle invalidation for external named sequences by
requiring their declarations to always provide such annotations.
Additionally, these annotations remove the need to analyze the body of
a named sequence to understand its effects on the arguments. Make them
required for named sequences that are called from the same file, in
addition to external sequences.

Provide a convenience pass that infers annotations by analyzing bodies
of named sequences provided they are not called from the same file.

[springerm]

Do you also have to update TransformDialect::verifyOperationAttribute?

[ftynse]

Do you also have to update TransformDialect::verifyOperationAttribute?

It doesn't seem to get called, otherwise it would have emitted "unknown attribute" errors... Perhaps some calls are missing somewhere.

[ftynse]

Remove spurious change.

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

let extraClassDeclaration = [{
       constexpr const static llvm::StringLiteral
           kWithNamedSequenceAttrName = "transform.with_named_sequence";
 
       /// Names of the attribute attachable to an operation so it can be
       /// identified as root by the default interpreter pass.
       constexpr const static llvm::StringLiteral
           kTargetTagAttrName = "transform.target_tag";
 
+      /// Names of the attributes indicating whether an argument of an external
+      /// transform dialect symbol is consumed or only read.
+      constexpr const static llvm::StringLiteral
+          kArgConsumedAttrName = "transform.consumed";
+      constexpr const static llvm::StringLiteral
+          kArgReadOnlyAttrName = "transform.readonly";
+
       /// 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;
 
       /// Parses a type registered by this dialect or one of its extensions.
       ::mlir::Type parseType(::mlir::DialectAsmParser &parser) const override;
 
private:
 
       /// A map from type TypeID to its printing function. No need to do string
       /// lookups when the type is fully constructed.
       ::llvm::DenseMap<::mlir::TypeID, ExtensionTypePrintingHook>
       typePrintingHooks;
   }];
 }
 
-// Base class for ops that belong to the tranfsorm dialect. Ops defined in
+// Base class for ops that belong to the transform dialect. Ops defined in
 // extensions of this dialect may also use this.
 class TransformDialectOp<string mnemonic, list<Trait> traits = []>
     : Op<Transform_Dialect, mnemonic, traits>;
 
 // Trait for operations that may be top-level operations in Transform IR.
 // Operations must have one single-block region and must be usable without
 // operands. See the C++ definition of the trait for more information.
 def PossibleTopLevelTransformOpTrait
     : NativeOpTrait<"PossibleTopLevelTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT

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

 /// Checks whether the transform op consumes the given handle.
 bool isHandleConsumed(Value handle, transform::TransformOpInterface transform);
 
 /// Populates `effects` with the memory effects indicating the access to payload
 /// IR resource.
 void modifiesPayload(SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
 void onlyReadsPayload(SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
 
+/// Populates `consumedArguments` with positions of `block` arguments that are
+/// consumed by the operations in the `block`.
+void getConsumedBlockArguments(
+    Block &block, llvm::SmallDenseSet<unsigned> &consumedArguments);
+
 /// Trait implementing the MemoryEffectOpInterface for operations that "consume"
 /// their operands and produce new results.
 template <typename OpTy>
 class FunctionalStyleTransformOpTrait
     : public OpTrait::TraitBase<OpTy, FunctionalStyleTransformOpTrait> {
 public:
   /// This op "consumes" the operands by reading and freeing then, "produces"
   /// the results by allocating and writing it and reads/writes the payload IR

mlir/include/mlir/Dialect/Transform/Transforms/Passes.h

-//===- CheckUses.h - Expensive transform value validity checks --*- C++ -*-===//
+//===- Passes.h - Transform dialect pass entry points -----------*- 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_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/Transform/Transforms/Passes.td

let description = [{
     `TransformMappingResource`. It does not currently perform an SCCP-style data
     flow analysis to prove that some branches are not taken, however, SCCP and
     other control flow simplifications can be performed on the transform IR
     prior to this pass provided that transform ops implement the relevant
     control flow interfaces.
   }];
 }
 
+def InferEffectsPass : Pass<"transform-infer-effects"> {
+  let summary = "infer transform side effects for symbols";
+  let description = [{
+    This pass analyzes the definitions of transform dialect callable symbol
+    operations, such as `transform.named_sequence`, and annotates the symbol
+    arguments with attributes indicating the side effects that the nested
+    operations have on them.
+  }];
+}
+
 #endif // MLIR_DIALECT_TRANSFORM_TRANSFORMS_PASSES

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

LogicalResult transform::TransformDialect::verifyOperationAttribute(
   }
   if (attribute.getName().getValue() == kTargetTagAttrName) {
     if (!attribute.getValue().isa<StringAttr>()) {
       return op->emitError()
              << attribute.getName() << " attribute must be a string";
     }
     return success();
   }
+  if (attribute.getName().getValue() == kArgConsumedAttrName ||
+      attribute.getName().getValue() == kArgReadOnlyAttrName) {
+    if (!attribute.getValue().isa<UnitAttr>()) {
+      return op->emitError()
+             << attribute.getName() << " must be a unit attribute";
+    }
+    return success();
+  }
   return emitError(op->getLoc())
          << "unknown attribute: " << attribute.getName();
 }
 
 #include "mlir/Dialect/Transform/IR/TransformDialectEnums.cpp.inc"

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

void transform::modifiesPayload(
   effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
 }
 
 void transform::onlyReadsPayload(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
 }
 
+void transform::getConsumedBlockArguments(
+    Block &block, llvm::SmallDenseSet<unsigned int> &consumedArguments) {
+  SmallVector<MemoryEffects::EffectInstance> effects;
+  for (Operation &nested : block) {
+    auto iface = dyn_cast<MemoryEffectOpInterface>(nested);
+    if (!iface)
+      continue;
+
+    effects.clear();
+    iface.getEffects(effects);
+    for (const MemoryEffects::EffectInstance &effect : effects) {
+      BlockArgument argument =
+          dyn_cast_or_null<BlockArgument>(effect.getValue());
+      if (!argument || argument.getOwner() != &block ||
+          !isa<MemoryEffects::Free>(effect.getEffect()) ||
+          effect.getResource() != transform::TransformMappingResource::get()) {
+        continue;
+      }
+      consumedArguments.insert(argument.getArgNumber());
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Utilities for TransformOpInterface.
 //===----------------------------------------------------------------------===//
 
 LogicalResult transform::detail::verifyTransformOpInterface(Operation *op) {
   auto iface = cast<MemoryEffectOpInterface>(op);
   SmallVector<MemoryEffects::EffectInstance> effects;
   iface.getEffects(effects);

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

remapArgumentEffects(Block &block, ValueRange operands,
   }
 }
 
 static DiagnosedSilenceableFailure
 verifyNamedSequenceOp(transform::NamedSequenceOp op);
 
 void transform::IncludeOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  // Always mark as modifying the payload.
+  // TODO: a mechanism to annotate effects on payload. Even when all handles are
+  // only read, the payload may still be modified, so we currently stay on the
+  // conservative side and always indicate modification. This may prevent some
+  // code reordering.
+  modifiesPayload(effects);
+
+  // Results are always produced.
+  producesHandle(getResults(), effects);
+
+  // Adds default effects to operands and results. This will be added if
+  // preconditions fail so the trait verifier doesn't complain about missing
+  // effects and the real precondition failure is reported later on.
+  auto defaultEffects = [&] { onlyReadsHandle(getOperands(), effects); };
+
   // Bail if the callee is unknown. This may run as part of the verification
   // process before we verified the validity of the callee or of this op.
   auto target =
       getOperation()->getAttrOfType<SymbolRefAttr>(getTargetAttrName());
   if (!target)
-    return;
+    return defaultEffects();
   auto callee = SymbolTable::lookupNearestSymbolFrom<NamedSequenceOp>(
       getOperation(), getTarget());
   if (!callee)
-    return;
+    return defaultEffects();
   DiagnosedSilenceableFailure earlyVerifierResult =
       verifyNamedSequenceOp(callee);
   if (!earlyVerifierResult.succeeded()) {
     (void)earlyVerifierResult.silence();
-    return;
+    return defaultEffects();
   }
 
-  // Carry over effects from the callee.
-  // TODO: external callees must provides attributes annotating the
-  // readonly/consume effects on operands.
-  if (!callee.isExternal())
-    remapArgumentEffects(callee.getBody().front(), getOperands(), effects);
-
-  // Proper effects.
-  onlyReadsHandle(getOperands(), effects);
-  producesHandle(getResults(), effects);
+  for (unsigned i = 0, e = getNumOperands(); i < e; ++i) {
+    if (callee.getArgAttr(i, TransformDialect::kArgConsumedAttrName))
+      consumesHandle(getOperand(i), effects);
+    else
+      onlyReadsHandle(getOperand(i), effects);
+  }
 }
 
 template <typename... Tys>
 static bool implementSameInterface(Type t1, Type t2) {
   return ((isa<Tys>(t1) && isa<Tys>(t2)) || ... || false);
 }
 
+/// Checks that the attributes of the named sequence operation have correct
+/// consumption effect annotations. If `alsoVerifyInternal`, checks for
+/// annotations being present even if they can be inferred from the body.
+static DiagnosedSilenceableFailure
+verifyNamedSequenceConsumeAnnotations(transform::NamedSequenceOp op,
+                                      bool alsoVerifyInternal = false) {
+  llvm::SmallDenseSet<unsigned> consumedArguments;
+  if (!op.isExternal()) {
+    transform::getConsumedBlockArguments(op.getBody().front(),
+                                         consumedArguments);
+  }
+  for (unsigned i = 0, e = op.getFunctionType().getNumInputs(); i < e; ++i) {
+    bool isConsumed =
+        op.getArgAttr(i, transform::TransformDialect::kArgConsumedAttrName) !=
+        nullptr;
+    bool isReadOnly =
+        op.getArgAttr(i, transform::TransformDialect::kArgReadOnlyAttrName) !=
+        nullptr;
+    if (isConsumed && isReadOnly) {
+      return op.emitSilenceableError()
+             << "argument #" << i << " cannot be both readonly and consumed";
+    }
+    if ((op.isExternal() || alsoVerifyInternal) && !isConsumed && !isReadOnly) {
+      return op.emitSilenceableError()
+             << "must provide consumed/readonly status for arguments of "
+                "external or called ops";
+    }
+    if (op.isExternal())
+      continue;
+
+    if (consumedArguments.contains(i) && !isConsumed && isReadOnly) {
+      return op.emitSilenceableError()
+             << "argument #" << i
+             << " is consumed in the body but is not marked as such";
+    }
+    if (!consumedArguments.contains(i) && isConsumed) {
+      Diagnostic warning(op->getLoc(), DiagnosticSeverity::Warning);
+      warning << "argument #" << i
+              << " is not consumed in the body but is marked as consumed";
+      return DiagnosedSilenceableFailure::silenceableFailure(
+          std::move(warning));
+    }
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
 LogicalResult
 transform::IncludeOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   // Access through indirection and do additional checking because this may be
   // running before the main op verifier.
   auto targetAttr = getOperation()->getAttrOfType<SymbolRefAttr>("target");
   if (!targetAttr)
     return emitOpError() << "expects a 'target' symbol reference attribute";
 
if (!implementSameInterface<TransformHandleTypeInterface,
                                 TransformParamTypeInterface>(resultType,
                                                              funcType)) {
       return emitOpError() << "type of result #" << i
                            << " must implement the same transform dialect "
                               "interface as the corresponding callee result";
     }
   }
 
-  return success();
+  return verifyNamedSequenceConsumeAnnotations(target,
+                                               /*alsoVerifyInternal=*/true)
+      .checkAndReport();
 }
 
 //===----------------------------------------------------------------------===//
 // MergeHandlesOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::MergeHandlesOp::apply(transform::TransformResults &results,
if (auto parent = op->getParentOfType<transform::TransformOpInterface>()) {
     DiagnosedSilenceableFailure diag =
         emitSilenceableFailure(op)
         << "cannot be defined inside another transform op";
     diag.attachNote(parent.getLoc()) << "ancestor transform op";
     return diag;
   }
 
   if (op.isExternal() || op.getBody().empty())
-    return DiagnosedSilenceableFailure::success();
+    return verifyNamedSequenceConsumeAnnotations(op);
 
   if (op.getBody().front().empty())
     return emitSilenceableFailure(op) << "expected a non-empty body block";
 
   Operation *terminator = &op.getBody().front().back();
   if (!isa<transform::YieldOp>(terminator)) {
     DiagnosedSilenceableFailure diag = emitSilenceableFailure(op)
                                        << "expected '"
for (auto [i, operandType, resultType] :
     if (operandType == resultType)
       continue;
     return emitSilenceableFailure(terminator)
            << "the type of the terminator operand #" << i
            << " must match the type of the corresponding parent op result ("
            << operandType << " vs " << resultType << ")";
   }
 
-  return DiagnosedSilenceableFailure::success();
+  return verifyNamedSequenceConsumeAnnotations(op);
 }
 
 LogicalResult transform::NamedSequenceOp::verify() {
   // Actual verification happens in a separate function for reusability.
   return verifyNamedSequenceOp(*this).checkAndReport();
 }
 
 //===----------------------------------------------------------------------===//

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

 add_mlir_dialect_library(MLIRTransformDialectTransforms
   CheckUses.cpp
+  InferEffects.cpp
   TransformInterpreterPassBase.cpp
 
   DEPENDS
   MLIRTransformDialectTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRTransformDialect
   MLIRIR
   MLIRPass
 )

mlir/lib/Dialect/Transform/Transforms/InferEffects.cpp

+//===- InferEffects.cpp - Infer memory effects for named symbols ----------===//
+//
+// 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/Transforms/Passes.h"
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/FunctionInterfaces.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/DenseSet.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace transform {
+#define GEN_PASS_DEF_INFEREFFECTSPASS
+#include "mlir/Dialect/Transform/Transforms/Passes.h.inc"
+} // namespace transform
+} // namespace mlir
+
+static LogicalResult inferSideEffectAnnotations(Operation *op) {
+  if (!isa<transform::TransformOpInterface>(op))
+    return success();
+
+  auto func = dyn_cast<FunctionOpInterface>(op);
+  if (!func || func.isExternal())
+    return success();
+
+  if (!func.getFunctionBody().hasOneBlock()) {
+    return op->emitError()
+           << "only single-block operations are currently supported";
+  }
+
+  // Note that there can't be an inclusion of an unannotated symbol because it
+  // wouldn't have passed the verifier, so recursion isn't necessary here.
+  llvm::SmallDenseSet<unsigned> consumedArguments;
+  transform::getConsumedBlockArguments(func.getFunctionBody().front(),
+                                       consumedArguments);
+
+  for (unsigned i = 0, e = func.getNumArguments(); i < e; ++i) {
+    func.setArgAttr(i,
+                    consumedArguments.contains(i)
+                        ? transform::TransformDialect::kArgConsumedAttrName
+                        : transform::TransformDialect::kArgReadOnlyAttrName,
+                    UnitAttr::get(op->getContext()));
+  }
+  return success();
+}
+
+namespace {
+class InferEffectsPass
+    : public transform::impl::InferEffectsPassBase<InferEffectsPass> {
+public:
+  void runOnOperation() override {
+    WalkResult result = getOperation()->walk([](Operation *op) {
+      return failed(inferSideEffectAnnotations(op)) ? WalkResult::interrupt()
+                                                    : WalkResult::advance();
+    });
+    if (result.wasInterrupted())
+      return signalPassFailure();
+  }
+};
+} // namespace

mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp

 //===- TransformInterpreterPassBase.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
 //
 //===----------------------------------------------------------------------===//
 //
 // Base class with shared implementation for transform dialect interpreter
 // passes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/FileUtilities.h"
saveReproToTempFile(llvm::dbgs(), target, transform, passName,
                         debugPayloadRootTag, debugTransformRootTag,
                         transformLibraryFileName, binaryName);
   });
 }
 
 /// Replaces external symbols in `block` with their (non-external) definitions
 /// from the given module.
 static LogicalResult defineDeclaredSymbols(Block &block, ModuleOp definitions) {
+  MLIRContext &ctx = *definitions->getContext();
+  auto consumedName =
+      StringAttr::get(&ctx, transform::TransformDialect::kArgConsumedAttrName);
+  auto readOnlyName =
+      StringAttr::get(&ctx, transform::TransformDialect::kArgReadOnlyAttrName);
+
   for (Operation &op : llvm::make_early_inc_range(block)) {
     LLVM_DEBUG(DBGS() << op << "\n");
     auto symbol = dyn_cast<SymbolOpInterface>(op);
     if (!symbol)
       continue;
     if (symbol->getNumRegions() == 1 && !symbol->getRegion(0).empty())
       continue;
 
for (Operation &op : llvm::make_early_inc_range(block)) {
 
     LLVM_DEBUG(llvm::dbgs() << "found @" << externalSymbol << "\n");
     if (symbolFunc.getFunctionType() != externalSymbolFunc.getFunctionType()) {
       return symbolFunc.emitError()
              << "external definition has a mismatching signature ("
              << externalSymbolFunc.getFunctionType() << ")";
     }
 
+    for (unsigned i = 0, e = symbolFunc.getNumArguments(); i < e; ++i) {
+      bool isExternalConsumed =
+          externalSymbolFunc.getArgAttr(i, consumedName) != nullptr;
+      bool isExternalReadonly =
+          externalSymbolFunc.getArgAttr(i, readOnlyName) != nullptr;
+      bool isConsumed = symbolFunc.getArgAttr(i, consumedName) != nullptr;
+      bool isReadonly = symbolFunc.getArgAttr(i, readOnlyName) != nullptr;
+      if (!isExternalConsumed && !isExternalReadonly) {
+        if (isConsumed)
+          externalSymbolFunc.setArgAttr(i, consumedName, UnitAttr::get(&ctx));
+        else if (isReadonly)
+          externalSymbolFunc.setArgAttr(i, readOnlyName, UnitAttr::get(&ctx));
+        continue;
+      }
+
+      if ((isExternalConsumed && !isConsumed) ||
+          (isExternalReadonly && !isReadonly)) {
+        return symbolFunc.emitError()
+               << "external definition has mismatching consumption annotations "
+                  "for argument #"
+               << i;
+      }
+    }
+
     OpBuilder builder(&op);
     builder.setInsertionPoint(&op);
     builder.clone(*externalSymbol);
     symbol->erase();
   }
 
   return success();
 }

mlir/test/Dialect/Transform/infer-effects.mlir

+// RUN: mlir-opt %s --transform-infer-effects | FileCheck %s
+
+module attributes { transform.with_named_sequence } {
+  // CHECK-LABEL: @infer
+  // CHECK-SAME: %{{.*}}: !transform.any_op {transform.consumed}
+  // CHECK-SAME: %{{.*}}: !transform.any_op {transform.readonly}
+  // CHECK-SAME: %{{.*}}: !transform.param<i32> {transform.readonly}
+  transform.named_sequence @infer(%op: !transform.any_op, %other: !transform.any_op, %param: !transform.param<i32>) {
+    transform.test_consume_operand %op : !transform.any_op
+    transform.test_print_remark_at_operand %other, "" : !transform.any_op
+    transform.yield
+  }
+}

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

 module attributes { transform.with_named_sequence} {
   func.func private @foo()
 
   transform.named_sequence @nested(%arg0: !transform.any_op) -> !transform.op<"builtin.module"> {
     // expected-error @below {{the type of the terminator operand #0 must match the type of the corresponding parent op result}}
     transform.yield %arg0 : !transform.any_op
   }
 }
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  // expected-error @below {{must provide consumed/readonly status for arguments of external or called ops}}
+  transform.named_sequence @foo(%op: !transform.any_op )
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  // expected-error @below {{argument #0 cannot be both readonly and consumed}}
+  transform.named_sequence @foo(%op: !transform.any_op { transform.readonly, transform.consumed } )
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  // expected-error @below {{must provide consumed/readonly status for arguments of external or called ops}}
+  transform.named_sequence @foo(%op: !transform.any_op) {
+    transform.test_print_remark_at_operand %op, "message" : !transform.any_op
+    transform.yield
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  // expected-error @below {{argument #0 cannot be both readonly and consumed}}
+  transform.named_sequence @foo(%op: !transform.any_op {transform.readonly, transform.consumed}) {
+    transform.test_print_remark_at_operand %op, "message" : !transform.any_op
+    transform.yield
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  // expected-warning @below {{argument #0 is not consumed in the body but is marked as consume}}
+  transform.named_sequence @foo(%op: !transform.any_op {transform.consumed}) {
+    transform.test_print_remark_at_operand %op, "message" : !transform.any_op
+    transform.yield
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  // expected-error @below {{argument #0 is consumed in the body but is not marked as such}}
+  transform.named_sequence @foo(%op: !transform.any_op {transform.readonly}) {
+    transform.test_consume_operand %op : !transform.any_op
+    transform.yield
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    transform.yield
+  }
+}
+
+// -----
+
+// Checking that consumptions annotations are used correctly in invocation checks.
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @foo(%op: !transform.any_op { transform.consumed } )
+
+  // expected-error @below {{'transform.sequence' op block argument #0 has more than one potential consumer}}
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-note @below {{used here as operand #0}}
+    transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    // expected-note @below {{used here as operand #0}}
+    transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    transform.yield
+  }
+}

mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-and-schedule.mlir

 // RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-symbol-decl.mlir transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
 // RUN:             --verify-diagnostics
 
 // RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-symbol-decl.mlir transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-symbol-decl.mlir transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
 // RUN:             --verify-diagnostics
 
 // The external transform script has a declaration to the named sequence @foo,
 // the definition of which is provided in another file. Repeated application
 // of the same pass should not be a problem. Note that the same diagnostic
 // produced twice at the same location only needs to be matched once.
 
 // expected-remark @below {{message}}
+// expected-remark @below {{unannotated}}
 module {}

mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-invalid.mlir

 // RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter)" \
 // RUN:             --verify-diagnostics --split-input-file
 
 // The definition of the @foo named sequence is provided in another file. It
 // will be included because of the pass option.
 
 module attributes {transform.with_named_sequence} {
   // expected-error @below {{external definition has a mismatching signature}}
-  transform.named_sequence private @foo(!transform.op<"builtin.module">)
+  transform.named_sequence private @foo(!transform.op<"builtin.module"> {transform.readonly})
 
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.op<"builtin.module">):
     include @foo failures(propagate) (%arg0) : (!transform.op<"builtin.module">) -> ()
   }
 }
 
 // -----
 
 module attributes {transform.with_named_sequence} {
   transform.named_sequence private @undefined_sequence()
 
   transform.sequence failures(suppress) {
   ^bb0(%arg0: !transform.any_op):
     // expected-error @below {{unresolved external named sequence}}
     include @undefined_sequence failures(suppress) () : () -> ()
   }
 }
+
+// -----
+
+module attributes {transform.with_named_sequence} {
+  // expected-error @below {{external definition has mismatching consumption annotations for argument #0}}
+  transform.named_sequence private @consuming(%arg0: !transform.any_op {transform.readonly})
+
+  transform.sequence failures(suppress) {
+  ^bb0(%arg0: !transform.any_op):
+    include @consuming failures(suppress) (%arg0) : (!transform.any_op) -> ()
+  }
+}

mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir

 // RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
-// RUN:             --verify-diagnostics | FileCheck %s
+// RUN:             --verify-diagnostics --split-input-file | FileCheck %s
 
 // RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter)" \
-// RUN:             --verify-diagnostics | FileCheck %s
+// RUN:             --verify-diagnostics --split-input-file | FileCheck %s
 
 // RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
-// RUN:             --verify-diagnostics | FileCheck %s
+// RUN:             --verify-diagnostics --split-input-file | FileCheck %s
 
 // The definition of the @foo named sequence is provided in another file. It
 // will be included because of the pass option. Repeated application of the
 // same pass, with or without the library option, should not be a problem.
 // Note that the same diagnostic produced twice at the same location only
 // needs to be matched once.
 
 // expected-remark @below {{message}}
+// expected-remark @below {{unannotated}}
 module attributes {transform.with_named_sequence} {
   // CHECK: transform.named_sequence @foo
   // CHECK: test_print_remark_at_operand %{{.*}}, "message"
-  transform.named_sequence private @foo(!transform.any_op)
+  transform.named_sequence private @foo(!transform.any_op {transform.readonly})
+
+  // CHECK: transform.named_sequence @unannotated
+  // CHECK: test_print_remark_at_operand %{{.*}}, "unannotated"
+  transform.named_sequence private @unannotated(!transform.any_op {transform.readonly})
 
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
     include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+    include @unannotated failures(propagate) (%arg0) : (!transform.any_op) -> ()
   }
 }

mlir/test/Dialect/Transform/test-interpreter-external-symbol-def.mlir

 // RUN: mlir-opt %s
 
 module attributes {transform.with_named_sequence} {
-  transform.named_sequence @foo(%arg0: !transform.any_op) {
+  transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) {
     transform.test_print_remark_at_operand %arg0, "message" : !transform.any_op
     transform.yield
   }
+
+  transform.named_sequence @consuming(%arg0: !transform.any_op {transform.consumed}) {
+    transform.test_consume_operand %arg0 : !transform.any_op
+    transform.yield
+  }
+
+  transform.named_sequence @unannotated(%arg0: !transform.any_op) {
+    transform.test_print_remark_at_operand %arg0, "unannotated" : !transform.any_op
+    transform.yield
+  }
 }

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

^bb1(%arg1: !transform.any_op):
   %op = transform.get_defining_op %bbarg : (!transform.any_value) -> !transform.any_op
   transform.test_print_remark_at_operand %op, "matched" : !transform.any_op
 }
 
 // -----
 
 module @named_inclusion attributes { transform.with_named_sequence } {
 
-  transform.named_sequence @foo(%arg0: !transform.any_op) -> () {
+  transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) -> () {
     // expected-remark @below {{applying transformation "a"}}
     transform.test_transform_op "a"
     transform.yield
   }
 
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
     include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
   }
 }
 
 // -----
 
 module @named_inclusion_in_named attributes { transform.with_named_sequence } {
 
-  transform.named_sequence @foo(%arg0: !transform.any_op) -> () {
+  transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) -> () {
     // expected-remark @below {{applying transformation "a"}}
     transform.test_transform_op "a"
     transform.yield
   }
 
-  transform.named_sequence @bar(%arg0: !transform.any_op) -> () {
+  transform.named_sequence @bar(%arg0: !transform.any_op {transform.readonly}) -> () {
     // expected-remark @below {{applying transformation "b"}}
     transform.test_transform_op "b"
     transform.include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
     transform.yield
   }
 
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
     transform.include @bar failures(suppress) (%arg0) : (!transform.any_op) -> ()
   }
 }
 
 // -----
 
 // expected-remark @below {{operation}}
 module @named_operands attributes { transform.with_named_sequence } {
 
-  transform.named_sequence @foo(%arg0: !transform.any_op, %arg1: !transform.any_value) -> () {
+  transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly},
+                                %arg1: !transform.any_value {transform.readonly}) -> () {
     transform.test_print_remark_at_operand %arg0, "operation" : !transform.any_op
     transform.test_print_remark_at_operand_value %arg1, "value" : !transform.any_value
     transform.yield
   }
 
   transform.sequence failures(propagate) {
   // expected-remark @below {{value}}
   // expected-note @below {{value handle points to a block argument #0 in block #0 in region #0}}
   ^bb0(%arg0: !transform.any_op):
     %0 = transform.test_produce_value_handle_to_self_operand %arg0 : (!transform.any_op) -> !transform.any_value
     include @foo failures(propagate) (%arg0, %0) : (!transform.any_op, !transform.any_value) -> ()
   }
 }
 
 // -----
 
 // expected-remark @below {{operation}}
 module @named_return attributes { transform.with_named_sequence } {
 
   // expected-remark @below {{value}}
   // expected-note @below {{value handle points to a block argument #0 in block #0 in region #0}}
-  transform.named_sequence @foo(%arg0: !transform.any_op) -> (!transform.any_op, !transform.any_value) {
+  transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) -> (!transform.any_op, !transform.any_value) {
     %0 = transform.test_produce_value_handle_to_self_operand %arg0 : (!transform.any_op) -> !transform.any_value
     transform.yield %arg0, %0 : !transform.any_op, !transform.any_value
   }
 
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
     %0:2 = include @foo failures(propagate) (%arg0) : (!transform.any_op) -> (!transform.any_op, !transform.any_value)
     transform.test_print_remark_at_operand %0#0, "operation" : !transform.any_op
     transform.test_print_remark_at_operand_value %0#1, "value" : !transform.any_value
   }
 }