[mlir] fixes to transform::SequenceOp

[ftynse]

Summary

Harden the verifier to check that the block argument type matches the
operand type, when present. This was overlooked when transform dialect
types were introduced.

Fix the builders to preserve the insertion point before creating the
block, otherwise the insertion point is updated to be within the block
by createBlock and never reset to be after the sequence op itself,
leading all following operations to be created in the unexpected to
the caller place.

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

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 "
+         "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";
+    }
+  }
+
   // Check if the block argument has more than one consuming use.
-  for (BlockArgument argument : getBodyBlock()->getArguments()) {
-    auto report = [&]() {
-      return (emitOpError() << "block argument #" << argument.getArgNumber());
-    };
-    if (failed(checkDoubleConsume(argument, report)))
-      return failure();
+  if (failed(checkDoubleConsume(
+          arg, [this]() { return (emitOpError() << "block argument #0"); }))) {
   }
 
   // Check properties of the nested operations they cannot check themselves.
   for (Operation &child : *getBodyBlock()) {
     if (!isa<TransformOpInterface>(child) &&
         &child != &getBodyBlock()->back()) {
       InFlightDiagnostic diag =
           emitOpError()

 
 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();
-  auto bbArgType = root.getType();
+  Type bbArgType = root.getType();
+  OpBuilder::InsertionGuard guard(builder);
   Block *bodyBlock = builder.createBlock(
       region, region->begin(), TypeRange{bbArgType}, {state.location});
 
   // Populate body.
-  OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToStart(bodyBlock);
   bodyBuilder(builder, state.location, bodyBlock->getArgument(0));
 }
 
 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});
 
   // Populate body.
-  OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToStart(bodyBlock);
   bodyBuilder(builder, state.location, bodyBlock->getArgument(0));
 }
 
 //===----------------------------------------------------------------------===//
 // WithPDLPatternsOp
 //===----------------------------------------------------------------------===//
 

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

 %0 = transform.sequence -> !pdl.operation failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-note @below {{terminator}}
   transform.yield
 }
 
 // -----
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{expects the type of the block argument to match the type of the operand}}
+  transform.sequence %arg0: !transform.any_op failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    transform.yield
+  }
+}
+
+// -----
+
 // 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}}
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
   }
 }