[mlir][Linalg] Make HoistRedundantTensorSubsetsOp not a FunctionalStyleTransformOpTrait

[nicolasvasilache]

Fix test

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

 }
 
 //===----------------------------------------------------------------------===//
 // HoistRedundantTensorSubsetsOp
 //===----------------------------------------------------------------------===//
 
 def HoistRedundantTensorSubsetsOp :
   Op<Transform_Dialect, "structured.hoist_redundant_tensor_subsets",
-    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformEachOpTrait, TransformOpInterface]> {
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
   let description = [{
     Hoists supported tensor subset extract/insert operation pairs out of 
     immediately enclosing loop iteratively, if the following conditions
     are true:
        1. The 2 ops access the same tensor subset.
        2. All operands are invariant under the enclosing loop.
     
     The supported subset extract/insert operation pairs currently comprise:
        - tensor.extract_slice / tensor.insert_slice
        - vector.transfer_read / vector.transfer_write on tensors
     
     Only scf.for loops are currently supported.
 
     When applied to:
        1. an scf.for loop, hoist out of this loop only.
        2. a non-loop op, apply hoisting to all the contained loop ops.
 
     #### Return modes:
 
-    The operation always succeeds and returns a handle to the transformed
-    function op.
+    The operation always succeeds and returns nothing.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
-  let results = (outs TransformHandleTypeInterface:$transformed);
+  let results = (outs);
 
-  let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
+  let assemblyFormat = [{
+    $target 
+    attr-dict 
+    `:` functional-type(operands, results)
+  }];
 
-  let builders = [
-    OpBuilder<(ins "Value":$target)>,
-  ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 //===----------------------------------------------------------------------===//
 // HoistRedundantTensorSubsetsOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::HoistRedundantTensorSubsetsOp::applyToOne(
     Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  IRRewriter rewriter(target->getContext());
+  TrackingListener listener(state, *this);
+  IRRewriter rewriter(target->getContext(), &listener);
   auto forOp = dyn_cast<scf::ForOp>(target);
   if (forOp) {
-    scf::ForOp newForOp =
-        linalg::hoistRedundantSubsetExtractInsert(rewriter, forOp);
-    results.push_back(newForOp);
     return DiagnosedSilenceableFailure::success();
   }
 
   // TODO: walking in some reverse / inside-out order would be more efficient
   // and would capture more cases.
   target->walk([&](scf::ForOp forOp) {
     hoistRedundantSubsetExtractInsert(rewriter, forOp);
   });
-  results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
+void transform::HoistRedundantTensorSubsetsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
 //===----------------------------------------------------------------------===//
 // InsertSliceToCopyOp
 //===----------------------------------------------------------------------===//
 template <typename OpTy>
 DiagnosedSilenceableFailure doit(RewriterBase &rewriter, OpTy target,
                                  transform::ApplyToEachResultList &results,
                                  transform::TransformState &state) {
   static_assert(llvm::is_one_of<OpTy, tensor::InsertSliceOp,

mlir/test/Dialect/Linalg/hoisting.mlir

return %0#0,  %0#1, %0#2, %0#3, %0#4,  %0#5 :
         tensor<?x?xf32>, tensor<?x?xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.hoist_redundant_tensor_subsets %0
-    : (!pdl.operation) -> !pdl.operation
+    : (!pdl.operation) -> ()
 }
 
 // -----
 
 // CHECK-LABEL: func @hoist_vector_transfer_pairs_disjoint_tensor(
 //  CHECK-SAME:   %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
 //  CHECK-SAME:   %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
 //  CHECK-SAME:   %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
// CHECK:                   vector.transfer_write %[[R]]#2, %[[TENSOR5]]{{.*}} : vector<3xf32>, tensor<?x?xf32>
   return %0#0,  %0#1, %0#2, %0#3 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.hoist_redundant_tensor_subsets %0
-    : (!pdl.operation) -> !pdl.operation
+    : (!pdl.operation) -> ()
 }
 
 // -----
 
 // CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor_and_slices
 //  CHECK-SAME:   %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
 //  CHECK-SAME:   %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
 //  CHECK-SAME:   %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
func.func @hoist_vector_transfer_pairs_tensor_and_slices(
   return %0#0, %0#1, %0#2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.hoist_redundant_tensor_subsets %0
-    : (!pdl.operation) -> !pdl.operation
+    : (!pdl.operation) -> ()
 }
 
 // -----
 
 // CHECK-LABEL: func @hoist_vector_transfer_write_pairs_disjoint_tensor(
 //  CHECK-SAME:   %[[T:.*]]: tensor<?x?xf32>,
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
func.func @hoist_vector_transfer_write_pairs_disjoint_tensor(
   return %1 : tensor<?x?xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.hoist_redundant_tensor_subsets %0
-    : (!pdl.operation) -> !pdl.operation
+    : (!pdl.operation) -> ()
 }
 
 // -----
 
 // CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor_and_slices_static_large_tensor
 //  CHECK-SAME:   %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<100x100xf32>,
 //  CHECK-SAME:   %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<200x200xf32>,
 //  CHECK-SAME:   %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<300x300xf32>
func.func @hoist_vector_transfer_pairs_tensor_and_slices_static_large_tensor(
   return %0#0, %0#1, %0#2 : tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.hoist_redundant_tensor_subsets %0
-    : (!pdl.operation) -> !pdl.operation
+    : (!pdl.operation) -> ()
 }
 
 // -----
 
 // CHECK-LABEL:  func.func @hoist_vector_transfer_read(
 // CHECK-DAG:      %[[C0:.+]] = arith.constant 0 : index
 // CHECK-DAG:      %[[C128:.+]] = arith.constant 128 : index
 // CHECK-DAG:      %[[C1024:.+]] = arith.constant 1024 : index