[mlir][transform] SequenceOp: Top-level operations can be used as matchers

[springerm]

Summary

As a convenience to the user, top-level sequence ops can optionally be used as matchers: the op type is specified by the type of the block argument.

This is similar to how pass pipeline targets can be specified on the command line (-pass-pipeline='builtin.module(func.func(...))).

Depends On: D153113

[ftynse]

This changes semantics a bit: previously, the first block argument of a sequence was guaranteed to the the payload root, which for example could be used to check if other handles are pointing to operations nested under that. This seems useful, although I am not particularly keen on special-casing the behavior for a specific transform type given that the type system is extensible.

One thing I have been pondering would be to update the syntax so we don't use an explicit block, which could make it clear at least in the syntax that we are applying the sequence to the "batch" of operations instead of one.

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

let description = [{
     errors and report them after attempting all transformations in the sequence.
 
     The entry block of this operation has a single argument that maps to either
     the operand if provided or the top-level container operation of the payload
     IR, typically the root operation of the pass interpreting the transform
     dialect. Operand omission is only allowed for sequences not contained in
     another sequence.
 
+    The type of the block argument must match the type of the operand. If the
+    sequence is a top-level transform (without an operand), it can be used for
+    matching operations if the specified type within the top-level container
+    payload IR (including the container op itself). E.g.:
+
+    ```mlir
+    transform.sequence failures(propagate) {
+    ^bb1(%arg1: !transform.any_op):
+      // %arg1 is mapped to the top-level container of the payload IR, which is
+      // typically a module
+    }
+
+    transform.sequence failures(propagate) {
+    ^bb1(%arg1: !transform.op<"func.func>"):
+      // %arg1 is mapped to all "func.func" ops within and including the
+      // top-level container of the payload IR. Nested operations that have the
+      // specified op type are not included.
+    }
+    ```
+
     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,
                        Variadic<Transform_AnyHandleOrParamType>:$extra_bindings);

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

if (op->getNumOperands() != 0) {
     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";
     }
 
+    // Top-level transforms can be used for matching. If no concrete operation
+    // type is specified, the block argument is mapped to the top-level op.
+    // Otherwise, it is mapped to all ops of the specified type within the
+    // top-level op (including the top-level op itself). Once an op is added as
+    // a target, its descendants are not explored any further.
+    BlockArgument bbArg = region.front().getArgument(0);
+    if (auto bbArgType = dyn_cast<transform::OperationType>(bbArg.getType())) {
+      state.getTopLevel()->walk<WalkOrder::PreOrder>([&](Operation *op) {
+        if (op->getName().getStringRef() == bbArgType.getOperationName()) {
+          targets.push_back(op);
+          return WalkResult::skip();
+        }
+        return WalkResult::advance();
+      });
+    } else {
       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()) {

mlir/test/Dialect/Tensor/fold-empty-op.mlir

 // RUN: mlir-opt -split-input-file -test-transform-dialect-interpreter %s | FileCheck %s
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.tensor.fold_tensor_empty
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
 // CHECK: #[[$MAP2:.+]] = affine_map<()[s0] -> (s0 * 28)>
 
 func.func @empty_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
   %0 = tensor.empty(%arg0) : tensor<6x5x?xf32>
   %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]]
func.func @rank_reducing_empty_tensor_extract(%sz : index, %idx : index) -> tensor<2xf32> {
   // CHECK-NOT: extract
   %r = tensor.extract_slice %a[%idx, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
   return %r: tensor<2xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.tensor.fold_tensor_empty
         {fold_single_use_only = true}
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 func.func @double_use_of_tensor_empty(%arg0: index, %arg1: index)
     -> (tensor<5x?x20xf32>, tensor<5x?x20xf32>)
 {
   %0 = tensor.empty(%arg0) : tensor<?x10x40xf32>
   %1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]
     : tensor<?x10x40xf32> to tensor<5x?x20xf32>
   %2 = tensor.extract_slice %0[1, 1, 1] [5, %arg1, 20] [1, 1, 1]
     : tensor<?x10x40xf32> to tensor<5x?x20xf32>
   return %1, %2 : tensor<5x?x20xf32>, tensor<5x?x20xf32>
 }
 // CHECK-LABEL: func @double_use_of_tensor_empty(
 //       CHECK:   tensor.empty{{.*}} : tensor<?x10x40xf32>
 //       CHECK:   tensor.extract_slice
 //       CHECK:   tensor.extract_slice

mlir/test/Dialect/Tensor/rewrite-as-constant.mlir

 // RUN: mlir-opt -split-input-file -test-transform-dialect-interpreter %s | FileCheck %s
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.tensor.rewrite_as_constant
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // CHECK-LABEL: func @tensor_generate_constant(
 //       CHECK:   %[[cst:.*]] = arith.constant dense<5.000000e+00> : tensor<2x3x5xf32>
 //       CHECK:   return %[[cst]]
 func.func @tensor_generate_constant() -> tensor<2x3x5xf32> {
   %cst = arith.constant 5.0 : f32
   %0 = tensor.generate {
     ^bb0(%arg0: index, %arg1: index, %arg2: index):
     tensor.yield %cst : f32
   } : tensor<2x3x5xf32>
   return %0 : tensor<2x3x5xf32>
 }

mlir/test/Dialect/Vector/vector-contract-matvec-transforms.mlir

func.func @redpar_vecmattrans2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
   %x = memref.load %arg1[] : memref<vector<2xf32>>
   %b = memref.load %arg2[] : memref<vector<2xf32>>
   %0 = vector.contract #redpar_vecmattrans_trait %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
   memref.store %0, %arg2[] : memref<vector<2xf32>>
   return
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-multi-reduction-lowering.mlir

func.func @vector_multi_reduction_parallel_middle(%arg0: vector<3x4x5xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
     return %0 : vector<4xf32>
 }
 
 // CHECK-LABEL: func @vector_multi_reduction_parallel_middle
 //  CHECK-SAME:   %[[INPUT:.+]]: vector<3x4x5xf32>, %[[ACC:.+]]: vector<4xf32>
 //       CHECK: vector.transpose %[[INPUT]], [1, 0, 2] : vector<3x4x5xf32> to vector<4x3x5xf32>
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_multi_reduction lowering_strategy = "innerreduction"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir

 func.func @vector_multi_reduction_to_scalar(%arg0: vector<2x3xf32>, %acc: f32) -> f32 {
   %0 = vector.multi_reduction <add>, %arg0, %acc [0, 1] : vector<2x3xf32> to f32
   return %0 : f32
 }
 // CHECK-LABEL: func @vector_multi_reduction_to_scalar
 //       CHECK:   return %{{.+}}
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_multi_reduction lowering_strategy = "innerparallel"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-transfer-drop-unit-dims-patterns.mlir

 }
 // CHECK-LABEL: func @transfer_write_and_vector_rank_reducing_to_0d
 //  CHECK-SAME:     %[[MEMREF:.+]]: memref<1x1x1x1x1xf32>, %[[VECTOR:.+]]: vector<1x1x1xf32>
 //       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[MEMREF]][0, 0, 0, 0, 0] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : memref<1x1x1x1x1xf32> to memref<f32>
 //       CHECK:   %[[SHCAST:.+]] = vector.shape_cast %[[VECTOR]] : vector<1x1x1xf32> to vector<f32>
 //       CHECK:   vector.transfer_write %[[SHCAST]], %[[SUBVIEW]]{{.*}} : vector<f32>, memref<f32>
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.rank_reducing_subview_patterns
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-transfer-full-partial-split-copy-transform.mlir

func.func @split_vector_transfer_read_strided_2d(
   // CHECK-SAME:   memref<?x8xf32, strided<[?, 1], offset: ?>>, vector<4x8xf32>
   %1 = vector.transfer_read %A[%i, %j], %f0 :
     memref<7x8xf32, strided<[?, 1], offset: ?>>, vector<4x8xf32>
 
   return %1 : vector<4x8xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "linalg-copy"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 func.func @split_vector_transfer_write_2d(%V: vector<4x8xf32>, %A: memref<?x8xf32>, %i: index, %j: index) {
   vector.transfer_write %V, %A[%i, %j] :
     vector<4x8xf32>, memref<?x8xf32>
   return

 // CHECK:             %[[DEST_VIEW:.*]] = memref.subview %[[DEST]][0, 0] [%[[VAL_20]], %[[VAL_21]]] [1, 1]
 // CHECK:             memref.copy %[[VAL_22]], %[[DEST_VIEW]]
 // CHECK-SAME:            : memref<?x?xf32, strided<[8, 1], offset: ?>> to memref<?x?xf32, strided{{.*}}>
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "linalg-copy"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 func.func @split_vector_transfer_write_strided_2d(
     %V: vector<4x8xf32>, %A: memref<7x8xf32, strided<[?, 1], offset: ?>>,
     %i: index, %j: index) {
   vector.transfer_write %V, %A[%i, %j] :

 // CHECK:             %[[DEST_VIEW:.*]] = memref.subview %[[DEST]][0, 0] [%[[VAL_20]], %[[VAL_21]]] [1, 1]
 // CHECK:             memref.copy %[[VAL_22]], %[[DEST_VIEW]]
 // CHECK-SAME:            : memref<?x?xf32, strided<[8, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>>
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "linalg-copy"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-transfer-full-partial-split.mlir

func.func @split_vector_transfer_read_strided_2d(
   %1 = vector.transfer_read %A[%i, %j], %f0 :
     memref<7x8xf32, strided<[?, 1], offset: ?>>, vector<4x8xf32>
 
   // CHECK: return %[[res]] : vector<4x8xf32>
   return %1 : vector<4x8xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 func.func @split_vector_transfer_write_2d(%V: vector<4x8xf32>, %A: memref<?x8xf32>, %i: index, %j: index) {
   vector.transfer_write %V, %A[%i, %j] :
     vector<4x8xf32>, memref<?x8xf32>
   return

 // CHECK-SAME:            %[[DEST]][%[[I]], %[[J]]]
 // CHECK-SAME:            : vector<4x8xf32>, memref<?x8xf32>
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 func.func @split_vector_transfer_write_strided_2d(
     %V: vector<4x8xf32>, %A: memref<7x8xf32, strided<[?, 1], offset: ?>>,
     %i: index, %j: index) {
   vector.transfer_write %V, %A[%i, %j] :

 // CHECK-SAME:            : memref<vector<4x8xf32>>
 // CHECK:             vector.transfer_write %[[VAL_20]], %[[DEST]][%[[I]], %[[J]]]
 // CHECK-SAME:            : vector<4x8xf32>, memref<7x8xf32, strided<[?, 1], offset: ?>>
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 func.func private @fake_side_effecting_fun(%0: vector<2x2xf32>) -> ()
 
 // CHECK-LABEL: transfer_read_within_async_execute
 func.func @transfer_read_within_async_execute(%A : memref<?x?xf32>) -> !async.token {
func.func @transfer_read_within_scf_for(%A : memref<?x?xf32>, %lb : index, %ub : index, %step : index) {
   scf.for %i = %lb to %ub step %step {
     %0 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x?xf32>, vector<2x2xf32>
     func.call @fake_side_effecting_fun(%0) : (vector<2x2xf32>) -> ()
   }
   return
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-transfer-tensor-slice-transforms.mlir

 // RUN: mlir-opt -split-input-file -test-transform-dialect-interpreter %s | FileCheck %s
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.fold_tensor_slice_into_transfer
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // CHECK-LABEL: func @transfer_read_of_extract_slice(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
 //   CHECK-DAG:   %[[c4:.*]] = arith.constant 4 : index
 //   CHECK-DAG:   %[[c8:.*]] = arith.constant 8 : index
 //       CHECK:   %[[add:.*]] = arith.addi %[[s1]], %[[c4]]
 //       CHECK:   %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[add]]], %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<5x6xf32>

mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir

func.func @transfer_broadcasting_complex(%mem : memref<10x20x30x8x8xf32>, %i : index) -> vector<3x2x4x5xf32> {
   %res = vector.transfer_read %mem[%i, %i, %i, %i, %i], %cf0 
     {in_bounds = [true, true, true, true], permutation_map = #broadcast_2d_in_4d} 
       : memref<10x20x30x8x8xf32>, vector<3x2x4x5xf32>
   return %res : vector<3x2x4x5xf32>
 }
 
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transfer max_transfer_rank = 99
     transform.apply_patterns.vector.transfer_permutation_patterns
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, 0, 0)>
 #map1 = affine_map<(d0, d1, d2, d3) -> (0, d1, 0, d0)>
 #map2 = affine_map<(d0, d1, d2, d3) -> (d3, d1, 0, 0)>
 #map3 = affine_map<(d0, d1) -> (d1, d0, 0, 0)>
func.func @transfer_write_broadcast_unit_dim(
   // CHECK: %[[NEW_VEC2:.*]] = vector.broadcast %{{.*}} : vector<8x16xf32> to vector<1x8x16xf32>
   // CHECK: %[[NEW_VEC3:.*]] = vector.transpose %[[NEW_VEC2]], [1, 2, 0] : vector<1x8x16xf32> to vector<8x16x1xf32>
   // CHECK: vector.transfer_write %[[NEW_VEC3]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] : vector<8x16x1xf32>, memref<?x?x?x?xf32>
 
   return %0 : tensor<?x?x?x?xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transfer max_transfer_rank = 99
     transform.apply_patterns.vector.transfer_permutation_patterns
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Dialect/Vector/vector-transpose-lowering.mlir

func.func @transpose1023_1x1x8x8xf32(%arg0: vector<1x1x8x8xf32>) -> vector<1x1x8x8xf32> {
   // Note the single 2-D extract/insert pair since 2 and 3 are not transposed!
   //      CHECK: vector.extract {{.*}}[0, 0] : vector<1x1x8x8xf32>
   // CHECK-NEXT: vector.insert {{.*}} [0, 0] : vector<8x8xf32> into vector<1x1x8x8xf32>
   %0 = vector.transpose %arg0, [1, 0, 2, 3] : vector<1x1x8x8xf32> to vector<1x1x8x8xf32>
   return %0 : vector<1x1x8x8xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose lowering_strategy = "eltwise"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 // CHECK-LABEL: func @transpose
 func.func @transpose(%arg0: vector<2x4xf32>) -> vector<4x2xf32> {
   //      CHECK: vector.shape_cast %{{.*}} : vector<2x4xf32> to vector<8xf32>
   //            0 4
   // 0 1 2 3    1 5
   // 4 5 6 7 -> 2 6
   //            3 7
   // CHECK-NEXT: vector.shuffle %{{.*}} [0, 4, 1, 5, 2, 6, 3, 7] : vector<8xf32>, vector<8xf32>
   // CHECK-NEXT: vector.shape_cast %{{.*}} : vector<8xf32> to vector<4x2xf32>
   %0 = vector.transpose %arg0, [1, 0] : vector<2x4xf32> to vector<4x2xf32>
   return %0 : vector<4x2xf32>
 }
 
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose lowering_strategy = "shuffle_1d"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 // CHECK-LABEL: func @transpose(
 func.func @transpose(%arg0: vector<2x4xf32>) -> vector<4x2xf32> {
   // CHECK:       vector.shape_cast {{.*}} : vector<2x4xf32> to vector<8xf32>
   // CHECK:       vector.flat_transpose %{{.*}} {columns = 2 : i32, rows = 4 : i32} : vector<8xf32> -> vector<8xf32>
   // CHECK:       vector.shape_cast {{.*}} : vector<8xf32> to vector<4x2xf32>
   %0 = vector.transpose %arg0, [1, 0] : vector<2x4xf32> to vector<4x2xf32>
   return %0 : vector<4x2xf32>
 }
 
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose lowering_strategy = "flat_transpose"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 // CHECK-LABEL: func @transpose4x8
 func.func @transpose4x8xf32(%arg0: vector<4x8xf32>) -> vector<8x4xf32> {
   //      CHECK: vector.extract {{.*}}[0]
   // CHECK-NEXT: vector.extract {{.*}}[1]
func.func @transpose210_1x8x8xf32(%arg0: vector<1x8x8xf32>) -> vector<8x8x1xf32> {
   // CHECK-NEXT: vector.insert {{.*}}[6]
   // CHECK-NEXT: vector.insert {{.*}}[7]
   // CHECK-NEXT: vector.shape_cast %{{.*}} : vector<8x8xf32> to vector<8x8x1xf32>
   %0 = vector.transpose %arg0, [2, 1, 0] : vector<1x8x8xf32> to vector<8x8x1xf32>
   return %0 : vector<8x8x1xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose avx2_lowering_strategy = true
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 func.func @transpose_shuffle16x16xf32(%arg0: vector<16x16xf32>) -> vector<16x16xf32> {
   // CHECK: vector.shuffle {{.*}} [0, 16, 1, 17, 4, 20, 5, 21, 8, 24, 9, 25, 12, 28, 13, 29] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [2, 18, 3, 19, 6, 22, 7, 23, 10, 26, 11, 27, 14, 30, 15, 31] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [0, 16, 1, 17, 4, 20, 5, 21, 8, 24, 9, 25, 12, 28, 13, 29] : vector<16xf32>, vector<16xf32>
func.func @transpose_shuffle16x16xf32(%arg0: vector<16x16xf32>) -> vector<16x16xf32> {
   // CHECK: vector.shuffle {{.*}} [4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
   %0 = vector.transpose %arg0, [1, 0] : vector<16x16xf32> to vector<16x16xf32>
   return %0 : vector<16x16xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose lowering_strategy = "shuffle_16x16"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }
 
 // -----
 
 // CHECK-LABEL: func @transpose021_shuffle16x16xf32
 func.func @transpose021_shuffle16x16xf32(%arg0: vector<1x16x16xf32>) -> vector<1x16x16xf32> {
   // CHECK: vector.shuffle {{.*}} [0, 16, 1, 17, 4, 20, 5, 21, 8, 24, 9, 25, 12, 28, 13, 29] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [2, 18, 3, 19, 6, 22, 7, 23, 10, 26, 11, 27, 14, 30, 15, 31] : vector<16xf32>, vector<16xf32>
func.func @transpose021_shuffle16x16xf32(%arg0: vector<1x16x16xf32>) -> vector<1x16x16xf32> {
   // CHECK: vector.shuffle {{.*}} [4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
   // CHECK: vector.shuffle {{.*}} [4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
   %0 = vector.transpose %arg0, [0, 2, 1] : vector<1x16x16xf32> to vector<1x16x16xf32>
   return %0 : vector<1x16x16xf32>
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose lowering_strategy = "shuffle_16x16"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }

mlir/test/Integration/Dialect/Vector/CPU/test-shuffle16x16.mlir

func.func @entry() {
   // CHECK-SAME:  ( 12, 28, 44, 60, 76, 92, 108, 124, 140, 156, 172, 188, 204, 220, 236, 252 ),
   // CHECK-SAME:  ( 13, 29, 45, 61, 77, 93, 109, 125, 141, 157, 173, 189, 205, 221, 237, 253 ),
   // CHECK-SAME:  ( 14, 30, 46, 62, 78, 94, 110, 126, 142, 158, 174, 190, 206, 222, 238, 254 ),
   // CHECK-SAME:  ( 15, 31, 47, 63, 79, 95, 111, 127, 143, 159, 175, 191, 207, 223, 239, 255 ) )
   return
 }
 
 transform.sequence failures(propagate) {
-^bb1(%module_op: !transform.any_op):
-  %func_op = transform.structured.match ops{["func.func"]} in %module_op
-      : (!transform.any_op) -> !transform.any_op
+^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.lower_transpose lowering_strategy = "shuffle_16x16"
-  } : !transform.any_op
+  } : !transform.op<"func.func">
 }