[mlir][Transform] NFC - Pass TransformState as an argument to applyToOne methods

[nicolasvasilache]

Summary

This will allow implementing state-dependent behavior in the future.

[nicolasvasilache]

Rebase.

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

def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
   }];
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target);
+        ::mlir::linalg::LinalgOp target, TransformState &state);
   }];
 }
 
 def FuseOp : Op<Transform_Dialect, "structured.fuse",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let description = [{
     Tiles the operations pointed to by the target handle and fuses their
def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
   }];
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target);
+        ::mlir::linalg::LinalgOp target, TransformState &state);
   }];
 }
 
 def InterchangeOp : Op<Transform_Dialect, "structured.interchange",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
     TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Interchanges the iterators of the operations pointed to by the target handle
     using the iterator interchange attribute.
   }];
 
   let arguments =
     (ins PDL_Operation:$target,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$iterator_interchange);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target);
+        ::mlir::linalg::LinalgOp target, TransformState &state);
   }];
 }
 
 def PadOp : Op<Transform_Dialect, "structured.pad",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Pads the operations pointed to by the target handle using the options
(ins PDL_Operation:$target,
           "{}">:$transpose_paddings);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target);
+        ::mlir::linalg::LinalgOp target, TransformState &state);
   }];
 }
 
 def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Indicates that ops of a specific kind in the given function should be
     scalarized (i.e. their dynamic dimensions tiled by 1).
 
     This operation returns the tiled op but not the loops.
 
     We make this design choice because it is hard to know ahead of time the
     number of loops that will be produced (it depends on the number of dynamic
     dimensions after multiple transformations have been applied).
   }];
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$result);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target);
+        ::mlir::linalg::LinalgOp target, TransformState &state);
   }];
 }
 
 def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
        [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
         TransformEachOpTrait, TransformOpInterface]> {
   let description = [{
     Indicates that the given `target` op should be transformed with the 
def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
   let results = (outs PDL_Operation:$fill_op,
                       PDL_Operation:$split_linalg_op,
                       PDL_Operation:$combining_linalg_op);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
-        ::mlir::linalg::LinalgOp target);
+        ::mlir::linalg::LinalgOp target, TransformState &state);
   }];
 }
 
 def TileOp : Op<Transform_Dialect, "structured.tile",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface]> {
   let description = [{
     Indicates that the given `target` op should be tiled with the options
def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
 
   let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<BoolAttr, "false">:$vectorize_padding);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<Operation *> applyToOne(::mlir::Operation *target);
+    ::mlir::FailureOr<Operation *> applyToOne(
+      ::mlir::Operation *target, TransformState &state);
   }];
 }
 
 #endif // LINALG_TRANSFORM_OPS

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

def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
   let arguments =
       (ins PDL_Operation:$target,
            DefaultValuedAttr<BoolAttr, "false">:$fail_if_already_divisible);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(::mlir::scf::ForOp loop);
+    ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(
+      ::mlir::scf::ForOp loop, TransformState &state);
   }];
 }
  
 def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Applies software pipelining to the loop";
   let description = [{
def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
   let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<I64Attr, "1">:$iteration_interval,
                    DefaultValuedAttr<I64Attr, "10">:$read_latency);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(::mlir::scf::ForOp loop);
+    ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(
+      ::mlir::scf::ForOp loop, TransformState &state);
   }];
 }
 
 def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Unrolls the given loop with the given unroll factor";
   let description = [{
      Unrolls each loop associated with the given handle to have up to the given
      number of loop body copies per iteration. If the unroll factor is larger
      than the loop trip count, the latter is used as the unroll factor instead.
      Does not produce a new handle as the operation may result in the loop being
      removed after a full unrolling.
   }];
 
   let arguments = (ins PDL_Operation:$target,
                        Confined<I64Attr, [IntPositive]>:$factor);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::LogicalResult applyToOne(::mlir::scf::ForOp loop);
+    ::mlir::LogicalResult applyToOne(
+      ::mlir::scf::ForOp loop, TransformState &state);
   }];
 }
 
 #endif // SCF_TRANSFORM_OPS

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

LogicalResult mapBlockArguments(TransformState &state) {
     return mapBlockArguments(state, this->getOperation()->getRegion(0));
   }
 };
 
 /// Trait implementing the TransformOpInterface for operations applying a
 /// transformation to a single operation handle and producing one or multiple
 /// operation handles.
 /// The op must implement a method with one of the following signatures:
-///   - FailureOr<convertible-to-Operation*> applyToOne(OpTy)
-///   - FailureOr<SmallVector<convertible-to-Operation*>> applyToOne(OpTy)
-///   - LogicalResult applyToOne(OpTy)
+///   - FailureOr<convertible-to-Operation*> applyToOne(OpTy, state)
+///   - FailureOr<SmallVector<convertible-to-Operation*>>applyToOne(OpTy, state)
+///   - LogicalResult applyToOne(OpTy, state)
 /// to perform a transformation that is applied in turn to all payload IR
 /// operations that correspond to the handle of the transform IR operation.
 /// In the functions above, OpTy is either Operation * or a concrete payload IR
 /// Op class that the transformation is applied to (NOT the class of the
 /// transform IR op). The op is expected to have a single operand.
 template <typename OpTy>
 class TransformEachOpTrait
     : public OpTrait::TraitBase<OpTy, TransformEachOpTrait> {
using TransformOpType = typename llvm::function_traits<
       decltype(&OpTy::applyToOne)>::template arg_t<0>;
   ArrayRef<Operation *> targets =
       state.getPayloadOps(this->getOperation()->getOperand(0));
   SmallVector<SmallVector<Operation *>, 1> results;
   // In the multi-result case, collect the number of results each transform
   // produced.
   DiagnosedSilenceableFailure result = detail::applyTransformToEach(
       targets, results, [&](TransformOpType specificOp) {
-        return static_cast<OpTy *>(this)->applyToOne(specificOp);
+        return static_cast<OpTy *>(this)->applyToOne(specificOp, state);
       });
   if (!result.succeeded())
     return result;
   for (const SmallVector<Operation *> &oneTargetResults : results) {
     if (OpTy::template hasTrait<OpTrait::ZeroResults>())
       continue;
     if (OpTy::template hasTrait<OpTrait::OneResult>()) {
       transformResults.set(

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

if (failed(result))
     return failure();
   return cast<LinalgOp>(result->getOperation());
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
-FailureOr<LinalgOp> transform::DecomposeOp::applyToOne(LinalgOp target) {
+FailureOr<LinalgOp> transform::DecomposeOp::applyToOne(LinalgOp target,
+                                                       TransformState &state) {
   FailureOr<LinalgOp> windowed =
       tryApply<DownscaleSizeOneWindowed2DConvolution>(target);
   if (succeeded(windowed))
     return windowed;
 
   FailureOr<LinalgOp> depthwise =
       tryApply<DownscaleDepthwiseConv2DNhwcHwcOp>(target);
   if (succeeded(depthwise))
LogicalResult transform::FuseOp::verify() {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GeneralizeOp
 //===----------------------------------------------------------------------===//
 
-FailureOr<LinalgOp> transform::GeneralizeOp::applyToOne(LinalgOp target) {
+FailureOr<LinalgOp> transform::GeneralizeOp::applyToOne(LinalgOp target,
+                                                        TransformState &state) {
   // Exit early if no transformation is needed.
   if (isa<GenericOp>(target))
     return target;
 
   FailureOr<LinalgOp> generic = tryApply<LinalgGeneralizationPattern>(target);
   if (succeeded(generic))
     return generic;
 
   return reportUnknownTransformError(target);
 }
 
 //===----------------------------------------------------------------------===//
 // InterchangeOp
 //===----------------------------------------------------------------------===//
 
-FailureOr<LinalgOp> transform::InterchangeOp::applyToOne(LinalgOp target) {
+FailureOr<LinalgOp>
+transform::InterchangeOp::applyToOne(LinalgOp target, TransformState &state) {
   SmallVector<unsigned> interchangeVector =
       extractUIntArray(getIteratorInterchange());
   // Exit early if no transformation is needed.
   if (interchangeVector.empty())
     return target;
 
   auto genericTarget = dyn_cast<GenericOp>(target.getOperation());
   if (!genericTarget) {
LogicalResult transform::InterchangeOp::verify() {
   }
   return success();
 }
 
 //===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//
 
-FailureOr<LinalgOp> transform::PadOp::applyToOne(LinalgOp target) {
+FailureOr<LinalgOp> transform::PadOp::applyToOne(LinalgOp target,
+                                                 TransformState &state) {
   // Convert the integer packing flags to booleans.
   SmallVector<bool> packPaddings;
   for (int64_t packPadding : extractI64Array(getPackPaddings()))
     packPaddings.push_back(static_cast<bool>(packPadding));
 
   // Convert the padding values to attributes.
   SmallVector<Attribute> paddingValues;
   for (auto const &it :
LogicalResult transform::PadOp::verify() {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ScalarizeOp
 //===----------------------------------------------------------------------===//
 
-FailureOr<LinalgOp> transform::ScalarizeOp::applyToOne(LinalgOp target) {
+FailureOr<LinalgOp> transform::ScalarizeOp::applyToOne(LinalgOp target,
+                                                       TransformState &state) {
   LinalgTilingOptions tilingOptions;
   tilingOptions.scalarizeDynamicDims();
   // Tiling with "scalarize_dyn_dims" actually sets the same lambda as the tile
   // sizes and asserts that it is not already set.
   SmallVector<int64_t> emptyTileSizes;
   LinalgTilingPattern pattern(getContext(), tilingOptions);
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<TiledLinalgOp> result =
       pattern.returningMatchAndRewrite(target, rewriter);
   if (failed(result))
     return failure();
 
   return result->op;
 }
 
 //===----------------------------------------------------------------------===//
 // SplitReductionOp
 //===----------------------------------------------------------------------===//
 
 FailureOr<SmallVector<Operation *>>
-transform::SplitReductionOp::applyToOne(LinalgOp target) {
+transform::SplitReductionOp::applyToOne(LinalgOp target,
+                                        TransformState &state) {
   ControlSplitReductionFn splitFn = [&](LinalgOp) {
     return std::pair<int64_t, unsigned>(getSplitFactor(),
                                         getInsertSplitDimension());
   };
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<SplitReductionResult> splitResult =
       splitReduction(rewriter, target, splitFn);
void TileOp::print(OpAsmPrinter &p) {
   p << getTarget();
   p.printOptionalAttrDict((*this)->getAttrs());
 }
 
 //===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
-FailureOr<Operation *> VectorizeOp::applyToOne(Operation *target) {
+FailureOr<Operation *> VectorizeOp::applyToOne(Operation *target,
+                                               TransformState &state) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     InFlightDiagnostic diag = emitOpError()
                               << "applies only to isolated-from-above targets";
     diag.attachNote(target->getLoc()) << "non-isolated target";
     return diag;
   }
 
   MLIRContext *ctx = getContext();

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

transform::LoopOutlineOp::apply(transform::TransformResults &results,
   results.set(getTransformed().cast<OpResult>(), transformed);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // LoopPeelOp
 //===----------------------------------------------------------------------===//
 
-FailureOr<scf::ForOp> transform::LoopPeelOp::applyToOne(scf::ForOp loop) {
+FailureOr<scf::ForOp> transform::LoopPeelOp::applyToOne(scf::ForOp loop,
+                                                        TransformState &state) {
   scf::ForOp result;
   IRRewriter rewriter(loop->getContext());
   LogicalResult status =
       scf::peelAndCanonicalizeForLoop(rewriter, loop, result);
   if (failed(status)) {
     if (getFailIfAlreadyDivisible())
       return reportUnknownTransformError(loop);
     return loop;
loopScheduling(scf::ForOp forOp,
   for (const auto &it : wrappedSchedule) {
     for (Operation *op : it.second) {
       unsigned cycle = opCycles[op];
       schedule.push_back(std::make_pair(op, cycle / iterationInterval));
     }
   }
 }
 
-FailureOr<scf::ForOp> transform::LoopPipelineOp::applyToOne(scf::ForOp loop) {
+FailureOr<scf::ForOp>
+transform::LoopPipelineOp::applyToOne(scf::ForOp loop, TransformState &state) {
   scf::PipeliningOption options;
   options.getScheduleFn =
       [this](scf::ForOp forOp,
              std::vector<std::pair<Operation *, unsigned>> &schedule) mutable {
         loopScheduling(forOp, schedule, getIterationInterval(),
                        getReadLatency());
       };
 
   scf::ForLoopPipeliningPattern pattern(options, loop->getContext());
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(loop);
   FailureOr<scf::ForOp> patternResult =
       pattern.returningMatchAndRewrite(loop, rewriter);
   if (failed(patternResult))
     return reportUnknownTransformError(loop);
   return patternResult;
 }
 
 //===----------------------------------------------------------------------===//
 // LoopUnrollOp
 //===----------------------------------------------------------------------===//
 
-LogicalResult transform::LoopUnrollOp::applyToOne(scf::ForOp loop) {
+LogicalResult transform::LoopUnrollOp::applyToOne(scf::ForOp loop,
+                                                  TransformState &state) {
   if (failed(loopUnrollByFactor(loop, getFactor())))
     return reportUnknownTransformError(loop);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//

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

for (Operation *op : state.getPayloadOps(getTarget()))
     op->erase();
 
   if (getFailAfterErase())
     return emitSilenceableError() << "silencable error";
   return DiagnosedSilenceableFailure::success();
 }
 
 FailureOr<SmallVector<Operation *>>
-mlir::test::TestWrongNumberOfResultsOp::applyToOne(Operation *) {
+mlir::test::TestWrongNumberOfResultsOp::applyToOne(
+    Operation *, transform::TransformState &state) {
   return SmallVector<Operation *>{};
 }
 
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension

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

def TestWrongNumberOfResultsOp
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$a,
                       PDL_Operation:$b,
                       PDL_Operation:$c);
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
-        ::mlir::Operation *target);
+        ::mlir::Operation *target, transform::TransformState &state);
   }];
 }
 
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD