[mlir] switch the transform loop extension to use types

[ftynse]

Summary

Add types to the Loop (SCF) extension of the transform dialect.

See https://discourse.llvm.org/t/rfc-type-system-for-the-transform-dialect/65702

[ftynse]

Rebase.

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.h

 //===- SCFTransformOps.h - SCF transformation ops ---------------*- 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_SCF_TRANSFORMOPS_SCFTRANSFORMOPS_H
 #define MLIR_DIALECT_SCF_TRANSFORMOPS_SCFTRANSFORMOPS_H
 
-#include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 namespace func {
 class FuncOp;
 } // namespace func
 namespace scf {
 class ForOp;

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td

 //===- SCFTransformOps.td - SCF (loop) transformation ops --*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SCF_TRANSFORM_OPS
 #define SCF_TRANSFORM_OPS
 
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
-include "mlir/Dialect/PDL/IR/PDLTypes.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
+def Transform_ScfForOp : Transform_ConcreteOpType<"scf.for">;
+
 def GetParentForOp : Op<Transform_Dialect, "loop.get_parent_for",
     [NavigationTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Gets a handle to the parent 'for' loop of the given operation";
   let description = [{
     Produces a handle to the n-th (default 1) parent `scf.for` loop for each
     Payload IR operation associated with the operand. Fails if such a loop
     cannot be found. The list of operations associated with the handle contains
     parent operations in the same order as the list associated with the operand,
     except for operations that are parents to more than one input which are only
     present once.
   }];
 
   let arguments =
-    (ins PDL_Operation:$target,
+    (ins TransformTypeInterface:$target,
          DefaultValuedAttr<ConfinedAttr<I64Attr, [IntPositive]>,
                            "1">:$num_loops);
-  let results = (outs PDL_Operation:$parent);
+  let results = (outs TransformTypeInterface:$parent);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat =
+    "$target attr-dict `:` functional-type(operands, results)";
 }
 
 def LoopOutlineOp : Op<Transform_Dialect, "loop.outline",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Outlines a loop into a named function";
   let description = [{
      Moves the loop into a separate function with the specified name and
      replaces the loop in the Payload IR with a call to that function. Takes
      care of forwarding values that are used in the loop as function arguments.
      If the operand is associated with more than one loop, each loop will be
      outlined into a separate function. The provided name is used as a _base_
      for forming actual function names following SymbolTable auto-renaming
      scheme to avoid duplicate symbols. Expects that all ops in the Payload IR
      have a SymbolTable ancestor (typically true because of the top-level
      module). Returns the handle to the list of outlined functions in the same
      order as the operand handle.
   }];
 
-  let arguments = (ins PDL_Operation:$target,
+  // Note that despite the name of the transform operation and related utility
+  // functions, the actual implementation does not require the operation to be
+  // a loop.
+  let arguments = (ins TransformTypeInterface:$target,
                    StrAttr:$func_name);
-  let results = (outs PDL_Operation:$transformed);
+  let results = (outs TransformTypeInterface:$transformed);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat =
+    "$target attr-dict `:` functional-type(operands, results)";
 }
 
 def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Peels the last iteration of the loop";
   let description = [{
      Updates the given loop so that its step evenly divides its range and puts
let description = [{
      Note that even though the Payload IR modification may be performed
      in-place, this operation consumes the operand handle and produces a new
      one.
 
      TODO: Return both the peeled loop and the remainder loop.
   }];
 
   let arguments =
-      (ins PDL_Operation:$target,
+      (ins Transform_ScfForOp:$target,
            DefaultValuedAttr<BoolAttr, "false">:$fail_if_already_divisible);
   // TODO: Return both the peeled loop and the remainder loop.
-  let results = (outs PDL_Operation:$transformed);
+  let results = (outs TransformTypeInterface:$transformed);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat =
+    "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::scf::ForOp target,
         ::llvm::SmallVector<::mlir::Operation *> &results,
         ::mlir::transform::TransformState &state);
   }];
 }
let description = [{
 
     This operation ignores non-scf::For ops and drops them in the return.
     If all the operations referred to by the `target` PDLOperation pipeline
     properly, the transform succeeds. Otherwise the transform silently fails.
     The return handle points to only the subset of successfully produced
     pipelined loops, which can be empty.
   }];
 
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins Transform_ScfForOp:$target,
                    DefaultValuedAttr<I64Attr, "1">:$iteration_interval,
                    DefaultValuedAttr<I64Attr, "10">:$read_latency);
-  let results = (outs PDL_Operation:$transformed);
+  let results = (outs TransformTypeInterface:$transformed);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat =
+    "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::scf::ForOp target,
         ::llvm::SmallVector<::mlir::Operation *> &results,
         ::mlir::transform::TransformState &state);
   }];
 }
let description = [{
     This operation ignores non-scf::For ops and drops them in the return.
     If all the operations referred to by the `target` PDLOperation unroll
     properly, the transform succeeds. Otherwise the transform silently fails.
 
     Does not return handles as the operation may result in the loop being
     removed after a full unrolling.
   }];
 
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins Transform_ScfForOp:$target,
                        ConfinedAttr<I64Attr, [IntPositive]>:$factor);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat = "$target attr-dict `:` type($target)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::scf::ForOp target,
         ::llvm::SmallVector<::mlir::Operation *> &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // SCF_TRANSFORM_OPS

mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFTransformOps
   SCFTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SCF/TransformOps
 
   DEPENDS
   MLIRSCFTransformOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRFuncDialect
   MLIRIR
-  MLIRPDLDialect
   MLIRSCFDialect
   MLIRSCFTransforms
   MLIRSCFUtils
   MLIRTransformDialect
   MLIRVectorDialect
 )

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

 //===- SCFTransformOps.cpp - Implementation of SCF transformation ops -----===//
 //
 // 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/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Patterns.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 

 namespace {
 class SCFTransformDialectExtension
     : public transform::TransformDialectExtension<
           SCFTransformDialectExtension> {
 public:
   using Base::Base;
 
   void init() {
-    declareDependentDialect<pdl::PDLDialect>();
-
     declareGeneratedDialect<AffineDialect>();
     declareGeneratedDialect<func::FuncDialect>();
 
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
         >();
   }

mlir/python/mlir/dialects/_loop_transform_ops_ext.py

 #  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
 
 try:
   from ..ir import *
   from ._ods_common import get_op_result_or_value as _get_op_result_or_value
-  from ..dialects import pdl
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
 from typing import Optional, Union
 
 
 def _get_int64_attr(arg: Optional[Union[int, IntegerAttr]],
                     default_value: int = None):
   if isinstance(arg, IntegerAttr):
     return arg
 
   if arg is None:
     assert default_value is not None, "must provide default value"
     arg = default_value
 
   return IntegerAttr.get(IntegerType.get_signless(64), arg)
 
 
 class GetParentForOp:
   """Extension for GetParentForOp."""
 
   def __init__(self,
+               result_type: Type,
                target: Union[Operation, Value],
                *,
                num_loops: int = 1,
                ip=None,
                loc=None):
     super().__init__(
-        pdl.OperationType.get(),
+        result_type,
         _get_op_result_or_value(target),
         num_loops=_get_int64_attr(num_loops, default_value=1),
         ip=ip,
         loc=loc)
 
 
 class LoopOutlineOp:
   """Extension for LoopOutlineOp."""
 
   def __init__(self,
+               result_type: Type,
                target: Union[Operation, Value],
                *,
                func_name: Union[str, StringAttr],
                ip=None,
                loc=None):
     super().__init__(
-        pdl.OperationType.get(),
+        result_type,
         _get_op_result_or_value(target),
         func_name=(func_name if isinstance(func_name, StringAttr) else
                    StringAttr.get(func_name)),
         ip=ip,
         loc=loc)
 
 
 class LoopPeelOp:
   """Extension for LoopPeelOp."""
 
   def __init__(self,
+               result_type: Type,
                target: Union[Operation, Value],
                *,
                fail_if_already_divisible: Union[bool, BoolAttr] = False,
                ip=None,
                loc=None):
     super().__init__(
-        pdl.OperationType.get(),
+        result_type,
         _get_op_result_or_value(target),
         fail_if_already_divisible=(fail_if_already_divisible if isinstance(
             fail_if_already_divisible, BoolAttr) else
                                    BoolAttr.get(fail_if_already_divisible)),
         ip=ip,
         loc=loc)
 
 
 class LoopPipelineOp:
   """Extension for LoopPipelineOp."""
 
   def __init__(self,
+               result_type: Type,
                target: Union[Operation, Value],
                *,
                iteration_interval: Optional[Union[int, IntegerAttr]] = None,
                read_latency: Optional[Union[int, IntegerAttr]] = None,
                ip=None,
                loc=None):
     super().__init__(
-        pdl.OperationType.get(),
+        result_type,
         _get_op_result_or_value(target),
         iteration_interval=_get_int64_attr(iteration_interval, default_value=1),
         read_latency=_get_int64_attr(read_latency, default_value=10),
         ip=ip,
         loc=loc)
 
 
 class LoopUnrollOp:

mlir/test/Dialect/Linalg/transform-op-fuse.mlir

 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1
     %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]}
-    transform.loop.peel %loops#0
+    %loop = transform.cast %loops#0 : !pdl.operation to !transform.op<"scf.for">
+    transform.loop.peel %loop : (!transform.op<"scf.for">) -> !pdl.operation
   }
 }
 
 // -----
 
 // CHECK-LABEL: func.func @interchange_reduction
 //  CHECK-SAME: (%[[INPUT:.+]]: tensor<12x7x25xf32>)
 func.func @interchange_reduction(%input: tensor<12x7x25xf32>) -> tensor<12x25xf32> {

mlir/test/Dialect/Linalg/transform-op-match.mlir

 // RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
 
 func.func @bar() {
   // expected-remark @below {{matched op name}}
   // expected-remark @below {{matched attr name}}
   %0 = arith.constant {my_attr} 0: i32
   // expected-remark @below {{matched op name}}
   %1 = arith.constant 1 : i32
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %match_name = transform.structured.match ops{["arith.constant"]} in %arg1
-    transform.test_print_remark_at_operand %match_name, "matched op name"
+    transform.test_print_remark_at_operand %match_name, "matched op name" : !pdl.operation
     transform.test_consume_operand %match_name
 
     %match_attr = transform.structured.match ops{["arith.constant"]} attributes{my_attr} in %arg1
-    transform.test_print_remark_at_operand %match_attr, "matched attr name"
+    transform.test_print_remark_at_operand %match_attr, "matched attr name" : !pdl.operation
     transform.test_consume_operand %match_attr
   }
 }
 
 // -----
 
 func.func @by_type() {
   %0 = arith.constant 0: i32
   // expected-remark @below {{matched op name}}
   %1 = arith.constant 1.0 : f32
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %match_name = transform.structured.match
       ops{["arith.constant"]} filter_result_type = f32 in %arg1
-    transform.test_print_remark_at_operand %match_name, "matched op name"
+    transform.test_print_remark_at_operand %match_name, "matched op name" : !pdl.operation
     transform.test_consume_operand %match_name
   }
 }
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>

 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %match_attr = transform.structured.match
         ops{["linalg.generic"]}
         attributes{iterator_types = ["parallel", "parallel", "parallel"]}
         in %arg1
-    transform.test_print_remark_at_operand %match_attr, "matched complex attr"
+    transform.test_print_remark_at_operand %match_attr, "matched complex attr" : !pdl.operation
     transform.test_consume_operand %match_attr
 
     %no_match = transform.structured.match
         attributes{iterator_types = ["parallel", "parallel", "reduction"]}
         in %arg1
   // expected-remark @below {{0}}
     transform.test_print_number_of_associated_payload_ir_ops %no_match
   }
 }

mlir/test/Dialect/MemRef/transform-ops.mlir

func.func @multi_buffer(%in: memref<16xf32>) {
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["memref.alloc"]} in %arg1
   %1 = transform.memref.multibuffer %0 {factor = 2 : i64}
   // Verify that the returned handle is usable.
-  transform.test_print_remark_at_operand %1, "transformed"
+  transform.test_print_remark_at_operand %1, "transformed" : !pdl.operation
 }

mlir/test/Dialect/SCF/transform-ops.mlir

 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["arith.addi"]} in %arg1
     // CHECK: = transform.loop.get_parent_for
-    %1 = transform.loop.get_parent_for %0
-    %2 = transform.loop.get_parent_for %0 { num_loops = 2 }
-    %3 = transform.loop.get_parent_for %0 { num_loops = 3 }
-    transform.test_print_remark_at_operand %1, "third loop"
-    transform.test_print_remark_at_operand %2, "second loop"
-    transform.test_print_remark_at_operand %3, "first loop"
+    %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
+    %2 = transform.loop.get_parent_for %0 { num_loops = 2 } : (!pdl.operation) -> !transform.op<"scf.for">
+    %3 = transform.loop.get_parent_for %0 { num_loops = 3 } : (!pdl.operation) -> !transform.op<"scf.for">
+    transform.test_print_remark_at_operand %1, "third loop" : !transform.op<"scf.for">
+    transform.test_print_remark_at_operand %2, "second loop" : !transform.op<"scf.for">
+    transform.test_print_remark_at_operand %3, "first loop" : !transform.op<"scf.for">
   }
 }
 
 // -----
 
 func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
   // expected-note @below {{target op}}
   arith.addi %arg0, %arg1 : index
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["arith.addi"]} in %arg1
     // expected-error @below {{could not find an 'scf.for' parent}}
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
   }
 }
 
 // -----
 
 // Outlined functions:
 //
 // CHECK: func @foo(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}})
func.func @loop_outline_op(%arg0: index, %arg1: index, %arg2: index) {
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["arith.addi"]} in %arg1
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_for %0  : (!pdl.operation) -> !transform.op<"scf.for">
     // CHECK: = transform.loop.outline %{{.*}}
-    transform.loop.outline %1 {func_name = "foo"}
+    transform.loop.outline %1 {func_name = "foo"} : (!transform.op<"scf.for">) -> !pdl.operation
   }
 }
 
 // -----
 
 func.func private @cond() -> i1
 func.func private @body()
 

 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["scf.while"]} in %arg1
     // expected-error @below {{failed to outline}}
-    transform.loop.outline %0 {func_name = "foo"}
+    transform.loop.outline %0 {func_name = "foo"} : (!pdl.operation) -> !pdl.operation
   }
 }
 
 // -----
 
 // CHECK-LABEL: @loop_peel_op
 func.func @loop_peel_op() {
   // CHECK: %[[C0:.+]] = arith.constant 0
func.func @loop_peel_op() {
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["arith.addi"]} in %arg1
-    %1 = transform.loop.get_parent_for %0
-    transform.loop.peel %1
+    %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
+    transform.loop.peel %1 : (!transform.op<"scf.for">) -> !pdl.operation
   }
 }
 
 // -----
 
 func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["arith.addf"]} in %arg1
-    %1 = transform.loop.get_parent_for %0
-    %2 = transform.loop.pipeline %1
+    %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
+    %2 = transform.loop.pipeline %1 : (!transform.op<"scf.for">) -> !pdl.operation
     // Verify that the returned handle is usable.
-    transform.test_print_remark_at_operand %2, "transformed"
+    transform.test_print_remark_at_operand %2, "transformed" : !pdl.operation
   }
 }
 
 // -----
 
 // CHECK-LABEL: @loop_unroll_op
 func.func @loop_unroll_op() {
   %c0 = arith.constant 0 : index
   %c42 = arith.constant 42 : index
   %c5 = arith.constant 5 : index
   // CHECK: scf.for %[[I:.+]] =
   scf.for %i = %c0 to %c42 step %c5 {
     // CHECK-COUNT-4: arith.addi %[[I]]
     arith.addi %i, %i : index
   }
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["arith.addi"]} in %arg1
-    %1 = transform.loop.get_parent_for %0
-    transform.loop.unroll %1 { factor = 4 }
+    %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
+    transform.loop.unroll %1 { factor = 4 } : !transform.op<"scf.for">
   }
 }
 

mlir/test/Dialect/Transform/expensive-checks.mlir

^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     // expected-note @below {{handle to invalidated ops}}
     %0 = pdl_match @return in %arg1 : (!pdl.operation) -> !pdl.operation
     %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
     // expected-note @below {{invalidated by this transform op that consumes its operand #0}}
     test_consume_operand %1
     // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
-    test_print_remark_at_operand %0, "remark"
+    test_print_remark_at_operand %0, "remark" : !pdl.operation
   }
 }
 
 // -----
 
 func.func @func1() {
   // expected-note @below {{repeated target op}}
   return
^bb0(%arg0: !pdl.operation):
 
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @func in %arg1 : (!pdl.operation) -> !pdl.operation
     %1 = pdl_match @return in %arg1 : (!pdl.operation) -> !pdl.operation
     %2 = replicate num(%0) %1 : !pdl.operation, !pdl.operation
     // expected-error @below {{a handle passed as operand #0 and consumed by this operation points to a payload operation more than once}}
     test_consume_operand %2
-    test_print_remark_at_operand %0, "remark"
+    test_print_remark_at_operand %0, "remark" : !pdl.operation
   }
 }
 
 
 // -----
 
 // expected-note @below {{ancestor payload op}}
 // expected-note @below {{nested payload op}}

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

 
 // -----
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = pdl_match @some in %arg1 : (!pdl.operation) -> !pdl.operation
-    test_print_remark_at_operand %0, "matched"
+    test_print_remark_at_operand %0, "matched" : !pdl.operation
   }
 
   pdl.pattern @some : benefit(1) {
     %0 = pdl.operation "test.some_op"
     pdl.rewrite %0 with "transform.dialect"
   }
 
   pdl.pattern @other : benefit(1) {
pdl.pattern @const : benefit(1) {
     pdl.rewrite %0 with "transform.dialect"
   }
 
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %f = pdl_match @const in %arg1 : (!pdl.operation) -> !pdl.operation
     // CHECK: %{{.+}} = get_closest_isolated_parent %{{.+}}
     %m = get_closest_isolated_parent %f : (!pdl.operation) -> !pdl.operation
-    test_print_remark_at_operand %m, "parent function"
+    test_print_remark_at_operand %m, "parent function" : !pdl.operation
   }
 }
 
 // -----
 
 func.func @foo() {
   %0 = arith.constant 0 : i32
   return
^bb1(%arg1: !pdl.operation):
     transform.alternatives %1 : !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
       %2 = transform.pdl_match @match_call in %arg2 : (!pdl.operation) -> !pdl.operation
       // expected-remark @below {{applying}}
       transform.test_emit_remark_and_erase_operand %2, "applying" {fail_after_erase}
     }, {
     ^bb2(%arg2: !pdl.operation):
       %2 = transform.pdl_match @match_call in %arg2 : (!pdl.operation) -> !pdl.operation
-      transform.test_print_remark_at_operand %2, "still here"
+      transform.test_print_remark_at_operand %2, "still here" : !pdl.operation
       // This alternative succeeds.
     }, {
     ^bb2(%arg2: !pdl.operation):
       // This alternative is never run, so we must not have a remark here.
       %2 = transform.pdl_match @match_call in %arg2 : (!pdl.operation) -> !pdl.operation
       transform.test_emit_remark_and_erase_operand %2, "should not happen" {fail_after_erase}
     }
   }
pdl.pattern @match_const : benefit(1) {
     %2 = pdl.operation "arith.constant"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
     pdl.rewrite %2 with "transform.dialect"
   }
 
 
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.pdl_match @match_const in %arg1 : (!pdl.operation) -> !pdl.operation
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !pdl.operation
     // expected-error @below {{only isolated-from-above ops can be alternative scopes}}
     alternatives %1 : !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
     }
   }
 }
 // -----
 
pdl.pattern @subi : benefit(1) {
     pdl.rewrite %2 with "transform.dialect"
   }
 
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = pdl_match @addi in %arg1 : (!pdl.operation) -> !pdl.operation
     %1 = pdl_match @subi in %arg1 : (!pdl.operation) -> !pdl.operation
     %2 = merge_handles %0, %1 : !pdl.operation
-    test_print_remark_at_operand %2, "matched"
+    test_print_remark_at_operand %2, "matched" : !pdl.operation
   }
 }
 
 // -----
 
 func.func @foo() {
   "op" () { target_me } : () -> ()
   // expected-note @below {{when applied to this op}}
^bb0(%arg0: !pdl.operation):
 
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %f = pdl_match @const in %arg1 : (!pdl.operation) -> !pdl.operation
     transform.foreach %f : !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
       // expected-remark @below {{1}}
       transform.test_print_number_of_associated_payload_ir_ops %arg2
-      transform.test_print_remark_at_operand %arg2, "transform applied"
+      transform.test_print_remark_at_operand %arg2, "transform applied" : !pdl.operation
     }
   }
 }
 
 // -----
 
 func.func @bar() {
   scf.execute_region {
^bb1(%arg1: !pdl.operation):
     %results = transform.foreach %f : !pdl.operation -> !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
       %g = transform.pdl_match @const in %arg2 : (!pdl.operation) -> !pdl.operation
       transform.yield %g : !pdl.operation
     }
 
     // expected-remark @below {{3}}
     transform.test_print_number_of_associated_payload_ir_ops %results
-    transform.test_print_remark_at_operand %results, "transform applied"
+    transform.test_print_remark_at_operand %results, "transform applied" : !pdl.operation
   }
 }
 
 // -----
 
 func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
   // expected-remark @below {{found muli}}
   %0 = arith.muli %arg0, %arg1 : index  
   arith.addi %0, %arg1 : index  
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %addi = transform.structured.match ops{["arith.addi"]} in %arg1
   %muli = get_producer_of_operand %addi[0] : (!pdl.operation) -> !pdl.operation
-  transform.test_print_remark_at_operand %muli, "found muli"
+  transform.test_print_remark_at_operand %muli, "found muli" : !pdl.operation
 }
 
 // -----
 
 func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
   // expected-note @below {{target op}}
   %0 = arith.muli %arg0, %arg1 : index  
   return

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

def TestConsumeOperandIfMatchesParamOrFail
   let assemblyFormat = "$operand `[` $parameter `]` attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestPrintRemarkAtOperandOp
   : Op<Transform_Dialect, "test_print_remark_at_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins 
-    Arg<PDL_Operation, "",
+    Arg<TransformTypeInterface, "",
         [TransformMappingRead, PayloadIRRead]>:$operand,
     StrAttr:$message);
-  let assemblyFormat = "$operand `,` $message attr-dict";
+  let assemblyFormat =
+    "$operand `,` $message attr-dict `:` type($operand)";
   let cppNamespace = "::mlir::test";
 }
 
 def TestAddTestExtensionOp
   : Op<Transform_Dialect, "test_add_test_extension",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         NoSideEffect]> {
   let arguments = (ins StrAttr:$message);

mlir/test/python/dialects/transform_loop_ext.py

with InsertionPoint(module.body):
       print("\nTEST:", f.__name__)
       f()
     print(module)
   return f
 
 
 @run
 def getParentLoop():
-  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], pdl.OperationType.get())
   with InsertionPoint(sequence.body):
-    loop.GetParentForOp(sequence.bodyTarget, num_loops=2)
+    loop.GetParentForOp(transform.OperationType.get("scf.for"), sequence.bodyTarget, num_loops=2)
     transform.YieldOp()
   # CHECK-LABEL: TEST: getParentLoop
   # CHECK: = transform.loop.get_parent_for %
   # CHECK: num_loops = 2
 
 
 @run
 def loopOutline():
-  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], transform.OperationType.get("scf.for"))
   with InsertionPoint(sequence.body):
-    loop.LoopOutlineOp(sequence.bodyTarget, func_name="foo")
+    loop.LoopOutlineOp(pdl.OperationType.get(), sequence.bodyTarget, func_name="foo")
     transform.YieldOp()
   # CHECK-LABEL: TEST: loopOutline
   # CHECK: = transform.loop.outline %
   # CHECK: func_name = "foo"
 
 
 @run
 def loopPeel():
-  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], transform.OperationType.get("scf.for"))
   with InsertionPoint(sequence.body):
-    loop.LoopPeelOp(sequence.bodyTarget)
+    loop.LoopPeelOp(pdl.OperationType.get(), sequence.bodyTarget)
     transform.YieldOp()
   # CHECK-LABEL: TEST: loopPeel
   # CHECK: = transform.loop.peel %
 
 
 @run
 def loopPipeline():
-  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], transform.OperationType.get("scf.for"))
   with InsertionPoint(sequence.body):
-    loop.LoopPipelineOp(sequence.bodyTarget, iteration_interval=3)
+    loop.LoopPipelineOp(pdl.OperationType.get(), sequence.bodyTarget, iteration_interval=3)
     transform.YieldOp()
   # CHECK-LABEL: TEST: loopPipeline
   # CHECK: = transform.loop.pipeline %
   # CHECK-DAG: iteration_interval = 3
   # CHECK-DAG: read_latency = 10
 
 
 @run
 def loopUnroll():
-  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], transform.OperationType.get("scf.for"))
   with InsertionPoint(sequence.body):
     loop.LoopUnrollOp(sequence.bodyTarget, factor=42)
     transform.YieldOp()
   # CHECK-LABEL: TEST: loopUnroll
   # CHECK: transform.loop.unroll %
   # CHECK: factor = 42