[mlir][Linalg][Transform] Avoid FunctionalStyleTransformOpTrait where unnecesseary to improve usability

Differential Revision: https://reviews.llvm.org/D146305

Author: nicolasvasilache

mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td

@@ -17,8 +17,7 @@ include "mlir/IR/OpBase.td"
 
 def MapNestedForallToThreads :
   Op<Transform_Dialect, "gpu.map_nested_forall_to_threads",
-    [FunctionalStyleTransformOpTrait,
-     MemoryEffectsOpInterface,
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformEachOpTrait,
      TransformOpInterface]> {
   let description = [{
@@ -72,9 +71,7 @@ def MapNestedForallToThreads :
       scf.forall operations with mappings other than gpu.thread are
       ignored.
 
-      The returned handle points to the same LaunchOp operand, consuming it and
-      producing a new SSA value to satisfy chaining and linearity of the IR
-      properties.
+      This operation returns nothing.
 
       #### Example:
 
@@ -111,18 +108,19 @@ def MapNestedForallToThreads :
       ```
     }];
 
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins TransformHandleTypeInterface:$target,
                    DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$block_dims,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$warp_dims,
                    DefaultValuedAttr<BoolAttr, "true">:$sync_after_distribute);
-  let results = (outs PDL_Operation:$result);
+  let results = (outs);
 
   let assemblyFormat = [{
     $target
     `block_dims` `=` $block_dims
     (`warp_dims` `=` $warp_dims^)?
     (`sync_after_distribute` `=` $sync_after_distribute^)?
     attr-dict
+    `:` functional-type(operands, results)
   }];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(

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

@@ -1651,11 +1651,13 @@ def TileToScfForOp : Op<Transform_Dialect, "structured.tile_to_scf_for",
 //===----------------------------------------------------------------------===//
 
 def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
-    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformEachOpTrait, TransformOpInterface]> {
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
   let description = [{
     Indicates that the given `target` op all the ops it contains should be
     vectorized with the configuration specified by the attributes of this op.
+
     This vectorization only handles structured ops that operate on shaped types
     and does not vectorize loops or straight-line. Internally, it applies a
     set of rewrite patterns, some of which enable vectorization and some of
@@ -1685,24 +1687,22 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
 
     This operation produces `definiteFailure` if vectorization fails for any
     reason.
-    The operation always returns the handle to the target op that is expected
-    to be isolated from above.
+    This operation returns nothing.
   }];
 
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins TransformHandleTypeInterface:$target,
                    UnitAttr:$vectorize_padding,
                    UnitAttr:$vectorize_nd_extract,
                    UnitAttr:$disable_multi_reduction_to_contract_patterns,
                    UnitAttr:$disable_transfer_permutation_map_lowering_patterns);
-  let results = (outs PDL_Operation:$transformed);
+  let results = (outs);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat = [{
+    $target 
+    attr-dict 
+    `:` functional-type(operands, results)
+  }];
 
-  let builders = [
-    OpBuilder<(ins "Value":$target,
-               CArg<"bool", "false">:$vectorizePadding,
-               CArg<"bool", "false">:$vectorizeNDExtract)>,
-  ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
@@ -1711,6 +1711,10 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// MaskedVectorizeOp
+//===----------------------------------------------------------------------===//
+
 def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformOpInterface]> {
@@ -1765,8 +1769,9 @@ def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
 
 def HoistRedundantVectorTransfersOp :
   Op<Transform_Dialect, "structured.hoist_redundant_vector_transfers",
-    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformEachOpTrait, TransformOpInterface]> {
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
   let description = [{
     Hoist vector.transfer_read / vector.transfer_write pairs out of immediately
     enclosing scf::ForOp iteratively, if the following conditions are true:
@@ -1782,18 +1787,17 @@ def HoistRedundantVectorTransfersOp :
 
     #### 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 assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
+  let results = (outs);
+  let assemblyFormat = [{
+    $target 
+    attr-dict 
+    `:` functional-type(operands, results)
+  }];
 
-  let builders = [
-    OpBuilder<(ins "Value":$target)>,
-  ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::func::FuncOp target,
@@ -1884,8 +1888,9 @@ def ConvertConv2DToImg2ColOp : Op<Transform_Dialect,
 
 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
@@ -1905,18 +1910,18 @@ def HoistRedundantTensorSubsetsOp :
 
     #### 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/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

@@ -17,8 +17,9 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def LowerVectorsOp : Op<Transform_Dialect, "vector.lower_vectors",
-    [DeclareOpInterfaceMethods<TransformOpInterface>,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
   let description = [{
     Indicates that the vector operations nested under the isolated from above op
     `target` should be lowered to finer-grained vector primitives.
@@ -27,10 +28,14 @@ def LowerVectorsOp : Op<Transform_Dialect, "vector.lower_vectors",
 
     This is usally a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
+
+    #### Return modes:
+
+    The operation returns nothing.
   }];
 
   // TODO: evolve this to proper enums.
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins TransformHandleTypeInterface:$target,
      DefaultValuedAttr<VectorContractLoweringAttr,
        "vector::VectorContractLowering::OuterProduct">:$contraction_lowering,
      DefaultValuedAttr<VectorMultiReductionLoweringAttr,
@@ -43,7 +48,7 @@ def LowerVectorsOp : Op<Transform_Dialect, "vector.lower_vectors",
      DefaultValuedAttr<BoolAttr, "false">:$transpose_avx2_lowering,
      DefaultValuedAttr<BoolAttr, "true">:$unroll_vector_transfers
   );
-  let results = (outs PDL_Operation:$results);
+  let results = (outs);
 
   let builders = [
     OpBuilder<(ins "Type":$resultType, "Value":$target,
@@ -66,6 +71,14 @@ def LowerVectorsOp : Op<Transform_Dialect, "vector.lower_vectors",
       | `transpose_lowering` `=` $transpose_lowering
     )
     attr-dict
+    `:` functional-type(operands, results)
+  }];
+  
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
   }];
 }
 

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

@@ -848,6 +848,12 @@ DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForallToThreadsImpl(
   return DiagnosedSilenceableFailure::success();
 }
 
+void transform::MapNestedForallToThreads::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getTarget(), effects);
+  modifiesPayload(effects);
+}
+
 DiagnosedSilenceableFailure transform::MapNestedForallToThreads::applyToOne(
     Operation *target, ApplyToEachResultList &results, TransformState &state) {
   LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
@@ -880,7 +886,6 @@ DiagnosedSilenceableFailure transform::MapNestedForallToThreads::applyToOne(
       mapNestedForallToThreadsImpl(rewriter, transformOp, gpuLaunch, blockDims,
                                    getWarpDims(), getSyncAfterDistribute());
 
-  results.push_back(gpuLaunch.getOperation());
   return diag;
 }
 

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

@@ -1786,7 +1786,7 @@ LogicalResult transform::PadOp::verify() {
 }
 
 //===---------------------------------------------------------------------===//
-// HoistPadOp
+// PadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
@@ -2977,21 +2977,6 @@ void transform::TileToScfForOp::getEffects(
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
-void transform::VectorizeOp::build(OpBuilder &builder, OperationState &result,
-                                   Value target, bool vectorizePadding,
-                                   bool vectorizeExtract) {
-  result.addOperands(target);
-  if (vectorizePadding) {
-    result.addAttribute(VectorizeOp::getVectorizePaddingAttrName(result.name),
-                        builder.getUnitAttr());
-  }
-  if (vectorizeExtract) {
-    result.addAttribute(VectorizeOp::getVectorizeNdExtractAttrName(result.name),
-                        builder.getUnitAttr());
-  }
-  result.addTypes(pdl::OperationType::get(builder.getContext()));
-}
-
 namespace {
 /// This is an helper only to call vectorize via a pattern inside of
 /// VectorizeOp::applyToOne.
@@ -3050,10 +3035,15 @@ transform::VectorizeOp::applyToOne(Operation *target,
   if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns))))
     return emitDefaultDefiniteFailure(target);
 
-  results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
+void transform::VectorizeOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
 //===----------------------------------------------------------------------===//
 // MaskedVectorizeOp
 //===----------------------------------------------------------------------===//
@@ -3133,22 +3123,6 @@ SmallVector<OpFoldResult> MaskedVectorizeOp::getMixedVectorSizes() {
   return getMixedValues(getStaticVectorSizes(), getVectorSizes(), b);
 }
 
-//===----------------------------------------------------------------------===//
-// HoistRedundantVectorTransfersOp
-//===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure
-transform::HoistRedundantVectorTransfersOp::applyToOne(
-    func::FuncOp target, transform::ApplyToEachResultList &results,
-    transform::TransformState &state) {
-  // WARNING: This hoisting does not model parallelism and is generally
-  // incorrect when used on distributed loops with memref semantics!
-  // TODO: obsolete and should be retired.
-  linalg::hoistRedundantVectorTransfers(target);
-  results.push_back(target);
-  return DiagnosedSilenceableFailure::success();
-}
-
 //===----------------------------------------------------------------------===//
 // ConvertConv2DToImg2ColOp.
 //===----------------------------------------------------------------------===//
@@ -3193,9 +3167,7 @@ transform::HoistRedundantTensorSubsetsOp::applyToOne(
   IRRewriter rewriter(target->getContext());
   auto forOp = dyn_cast<scf::ForOp>(target);
   if (forOp) {
-    scf::ForOp newForOp =
-        linalg::hoistRedundantSubsetExtractInsert(rewriter, forOp);
-    results.push_back(newForOp);
+    linalg::hoistRedundantSubsetExtractInsert(rewriter, forOp);
     return DiagnosedSilenceableFailure::success();
   }
 
@@ -3204,10 +3176,36 @@ transform::HoistRedundantTensorSubsetsOp::applyToOne(
   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);
+}
+
+//===----------------------------------------------------------------------===//
+// HoistRedundantVectorTransfersOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::HoistRedundantVectorTransfersOp::applyToOne(
+    func::FuncOp target, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  // WARNING: This hoisting does not model parallelism and is generally
+  // incorrect when used on distributed loops with memref semantics!
+  // TODO: obsolete and should be retired.
+  linalg::hoistRedundantVectorTransfers(target);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::HoistRedundantVectorTransfersOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//