[mlir] add producer fusion to structured transform ops

[ftynse]

Summary

This relies on the existing TileAndFuse pattern for tensor-based structured
ops. It complements pure tiling, from which some utilities are generalized.

Depends On D127300

[ftynse]

Rebase

[ftynse]

Rerebase.

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

def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
         ::mlir::linalg::LinalgOp target);
   }];
 }
 
+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
+    producers greedily using the options provided as attributes.
+  }];
+
+  let arguments =
+    (ins PDL_Operation:$target,
+         DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,
+         DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_interchange);
+  let results = (outs PDL_Operation:$transformed,
+                      Variadic<PDL_Operation>:$loops);
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
 def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Transforms a named structued operation into the generic form with the
     explicit attached region. The operand handle must point to a list of
     structured operations, it is consumed by the transformation and is not
     expected to be used afterwards. The resulting handle points to the list
def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
         ::mlir::linalg::LinalgOp target);
   }];
 }
 
 def TileOp : Op<Transform_Dialect, "structured.tile",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
-        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+        FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface]> {
   let description = [{
     Indicates that the given `target` op should be tiled with the options
     provided as attributes. This transform generates a loop nest with a smaller
     ("tiled") target operation in its body. Currently limited to LinalgOps.
 
     `sizes` are the tile sizes. A tile size of `0` indicates that the
     respective dimension should not be tiled. No loop will be generated for such
     dimensions. If all tile sizes are `0`, this transform is effectively a

mlir/include/mlir/Dialect/SCF/Patterns.h

 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 namespace scf {
 /// Generate a pipelined version of the scf.for loop based on the schedule given
 /// as option. This applies the mechanical transformation of changing the loop
 /// and generating the prologue/epilogue for the pipelining and doesn't make any
 /// decision regarding the schedule.
-/// Based on the options the loop is split into several stages.
+/// Based on the option the loop is split into several stages.
 /// The transformation assumes that the scheduling given by user is valid.
 /// For example if we break a loop into 3 stages named S0, S1, S2 we would
-/// generate the following code with the number in parenthesis as the iteration
+/// generate the following code with the number in parenthesis the iteration
 /// index:
 /// S0(0)                        // Prologue
 /// S0(1) S1(0)                  // Prologue
 /// scf.for %I = %C0 to %N - 2 {
 ///  S0(I+2) S1(I+1) S2(I)       // Pipelined kernel
 /// }
 /// S1(N) S2(N-1)                // Epilogue
 /// S2(N)                        // Epilogue

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

 
 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/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
-def GetParentForOp : Op<Transform_Dialect, "loop.get_parent_for",
+def GetParentLoopOp : Op<Transform_Dialect, "loop.get_parent_loop",
     [NavigationTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
-  let summary = "Gets a handle to the parent 'for' loop of the given operation";
+  let summary = "Gets a handle to the parent 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,
          DefaultValuedAttr<Confined<I64Attr, [IntPositive]>,
                            "1">:$num_loops);
   let results = (outs PDL_Operation:$parent);
 
   let assemblyFormat = "$target attr-dict";
 }
 
 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.
+     replaces the loop in the Payload IR with the call to that function. Takes
+     care or 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,
                    StrAttr:$func_name);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 }
 
 def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
-  let summary = "Peels the last iteration of the loop";
+  let summary = "Peels of the last iterations of the loop";
   let description = [{
-     Updates the given loop so that its step evenly divides its range and puts
-     the remaining iteration into a separate loop or a conditional. Note that
-     even though the Payload IR modification may be performed in-place, this
+     Updates the given loop so that its step evenly divides range and puts the
+     remaining iterations into a separate loop or a conditional. Note that even
+     though the Payload IR modification may be performed in-place, this
      operation consumes the operand handle and produces a new one. Applies to
      each loop associated with the operand handle individually. The results
      follow the same order as the operand.
-
-     Note: If it can be proven statically that the step already evenly divides
-     the range, this op is a no-op. In the absence of sufficient static
-     information, this op may peel a loop, even if the step always divides the
-     range evenly at runtime.
   }];
 
-  let arguments =
-      (ins PDL_Operation:$target,
-           DefaultValuedAttr<BoolAttr, "false">:$fail_if_already_divisible);
+  let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(::mlir::scf::ForOp loop);
   }];
 }

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

 /// cannot modify the Payload IR without reading it first). This is intended
 /// to disallow reordering of Transform IR operations that mutate the Payload IR
 /// while still allowing the reordering of those that only access it.
 struct PayloadIRResource
     : public SideEffects::Resource::Base<PayloadIRResource> {
   StringRef getName() override { return "transform.payload_ir"; }
 };
 
-/// Trait implementing the MemoryEffectOpInterface for single-operand zero- or
-/// single-result operations that "consume" their operand and produce a new
-/// result.
+/// Trait implementing the MemoryEffectOpInterface for single-operand operations
+/// that "consume" their operand and produce a new result.
 template <typename OpTy>
 class FunctionalStyleTransformOpTrait
     : public OpTrait::TraitBase<OpTy, FunctionalStyleTransformOpTrait> {
 public:
   /// This op "consumes" the operand by reading and freeing it, "produces" the
-  /// result by allocating and writing it and reads/writes the payload IR in the
-  /// process.
+  /// results by allocating and writing it and reads/writes the payload IR in
+  /// the process.
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
     effects.emplace_back(MemoryEffects::Read::get(),
                          this->getOperation()->getOperand(0),
                          TransformMappingResource::get());
     effects.emplace_back(MemoryEffects::Free::get(),
                          this->getOperation()->getOperand(0),
                          TransformMappingResource::get());
-    if (this->getOperation()->getNumResults() == 1) {
-      effects.emplace_back(MemoryEffects::Allocate::get(),
-                           this->getOperation()->getResult(0),
+    for (Value result : this->getOperation()->getResults()) {
+      effects.emplace_back(MemoryEffects::Allocate::get(), result,
                            TransformMappingResource::get());
-      effects.emplace_back(MemoryEffects::Write::get(),
-                           this->getOperation()->getResult(0),
+      effects.emplace_back(MemoryEffects::Write::get(), result,
                            TransformMappingResource::get());
     }
     effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
     effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
   }
 
   /// Checks that the op matches the expectations of this trait.
   static LogicalResult verifyTrait(Operation *op) {
     static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
                   "expected single-operand op");
-    static_assert(OpTy::template hasTrait<OpTrait::ZeroResults>() ||
-                      OpTy::template hasTrait<OpTrait::OneResult>(),
-                  "expected zero- or single-result op");
     if (!op->getName().getInterface<MemoryEffectOpInterface>()) {
       op->emitError()
           << "FunctionalStyleTransformOpTrait should only be attached to ops "
              "that implement MemoryEffectOpInterface";
     }
     return success();
   }
 };

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

InterfaceMethod<
       /*arguments=*/(ins
           "::mlir::transform::TransformResults &":$transformResults,
           "::mlir::transform::TransformState &":$state
     )>,
   ];
 
   let extraSharedClassDeclaration = [{
     /// Emits a generic transform error for the current transform operation
-    /// targeting the given Payload IR operation and returns failure. Should
+    /// targeting the given Payload IR opreation and returns failure. Should
     /// be only used as a last resort when the transformation itself provides
     /// no further indication as to the reason of the failure.
     ::mlir::LogicalResult reportUnknownTransformError(
         ::mlir::Operation *target) {
       ::mlir::InFlightDiagnostic diag = $_op->emitError() << "failed to apply";
       diag.attachNote(target->getLoc()) << "attempted to apply to this op";
       return diag;
     }

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

FailureOr<LinalgOp> depthwise =
       tryApply<DownscaleDepthwiseConv2DNhwcHwcOp>(target);
   if (succeeded(depthwise))
     return depthwise;
 
   return reportUnknownTransformError(target);
 }
 
 //===----------------------------------------------------------------------===//
+// FuseOp
+//===----------------------------------------------------------------------===//
+
+/// Apply a tiling transformation to all payload ops and store both the
+/// tiled operation as well as the created tile loops.
+static LogicalResult
+applyTilingToAll(Operation *transformOp, Value target,
+                 ArrayRef<int64_t> tileSizes,
+                 transform::TransformResults &transformResults,
+                 transform::TransformState &state,
+                 function_ref<FailureOr<TiledLinalgOp>(LinalgOp)> applyFn) {
+  // Number of loops: Number of tiles sizes that are not zero.
+  size_t numLoops = tileSizes.size() - llvm::count(tileSizes, 0);
+  // All payload ops. These should all be LinalgOps for now.
+  ArrayRef<Operation *> payloadOps = state.getPayloadOps(target);
+
+  SmallVector<Operation *> tiledLinalgOps;
+  SmallVector<SmallVector<Operation *>> loopOps(numLoops);
+  for (unsigned int i = 0; i < numLoops; ++i)
+    loopOps[i].reserve(payloadOps.size());
+
+  for (Operation *target : payloadOps) {
+    auto linalgOp = dyn_cast<linalg::LinalgOp>(target);
+    if (!linalgOp)
+      return transformOp->emitError("only LinalgOps are supported");
+
+    FailureOr<TiledLinalgOp> tiled = applyFn(linalgOp);
+    if (failed(tiled))
+      return failure();
+
+    tiledLinalgOps.push_back(tiled->op);
+    if (tiled->loops.size() != numLoops)
+      // Not enough loops were generated. This usually means that the input size
+      // was smaller than the tiling size.
+      // TODO: LinalgTilingPattern should return failure().
+      return failure();
+    for (unsigned int i = 0; i < numLoops; ++i)
+      loopOps[i].push_back(tiled->loops[i]);
+  }
+
+  transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);
+  for (unsigned int i = 0; i < numLoops; ++i)
+    transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);
+  return success();
+}
+
+/// Parse a tiling-like operation that returns the tiled op as well as the
+/// created tile loops. The function counts the non-zero tile sizes to compute
+/// the number of results.
+static ParseResult parseTileLikeOp(OpAsmParser &parser, OperationState &result,
+                                   StringRef sizesAttrName) {
+  OpAsmParser::UnresolvedOperand targetOperand;
+  SMLoc opLoc = parser.getCurrentLocation();
+  if (parser.parseOperand(targetOperand) ||
+      parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+  Attribute sizesAttr = result.attributes.get(sizesAttrName);
+  if (!sizesAttr)
+    return parser.emitError(opLoc)
+           << "expected '" << sizesAttrName << "' attribute";
+  auto sizesArrayAttr = sizesAttr.dyn_cast<ArrayAttr>();
+  if (!sizesArrayAttr)
+    return parser.emitError(opLoc)
+           << "'" << sizesAttrName << "' attribute must be an array";
+  Type pdlOpType = parser.getBuilder().getType<pdl::OperationType>();
+  size_t numExpectedLoops =
+      sizesArrayAttr.size() - llvm::count(extractI64Array(sizesArrayAttr), 0);
+  result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOpType));
+  if (parser.resolveOperand(targetOperand, pdlOpType, result.operands))
+    return failure();
+  return success();
+}
+
+LogicalResult
+transform::FuseOp::apply(mlir::transform::TransformResults &transformResults,
+                         mlir::transform::TransformState &state) {
+  LinalgTilingAndFusionOptions fusionOptions;
+  fusionOptions.tileSizes = extractI64Array(getTileSizes());
+  fusionOptions.tileInterchange = extractI64Array(getTileInterchange());
+
+  return applyTilingToAll(
+      getOperation(), getTarget(), fusionOptions.tileSizes, transformResults,
+      state, [&](LinalgOp linalgOp) -> FailureOr<TiledLinalgOp> {
+        LinalgTileAndFuseTensorOpsPattern pattern(getContext(), fusionOptions);
+        SimpleRewriter rewriter(getContext());
+        rewriter.setInsertionPoint(linalgOp);
+        FailureOr<TileLoopNest> tileLoopNest =
+            pattern.returningMatchAndRewrite(linalgOp, rewriter);
+        if (failed(tileLoopNest))
+          return failure();
+
+        TiledLinalgOp tiledLinalgOp;
+        tiledLinalgOp.op = tileLoopNest->getRootOp();
+        tiledLinalgOp.loops = {tileLoopNest->getLoopOps().begin(),
+                               tileLoopNest->getLoopOps().end()};
+        return tiledLinalgOp;
+      });
+}
+
+ParseResult transform::FuseOp::parse(OpAsmParser &parser,
+                                     OperationState &result) {
+  return parseTileLikeOp(
+      parser, result,
+      transform::FuseOp::getTileSizesAttrName(result.name).getValue());
+}
+
+void transform::FuseOp::print(OpAsmPrinter &p) {
+  p << ' ';
+  p << getTarget();
+  p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+LogicalResult transform::FuseOp::verify() {
+  SmallVector<int64_t> permutation = extractI64Array(getTileInterchange());
+  auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, permutation.size()));
+  if (!std::is_permutation(sequence.begin(), sequence.end(),
+                           permutation.begin(), permutation.end())) {
+    return emitOpError() << "expects interchange to be a permutation, found "
+                         << getTileInterchange();
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // GeneralizeOp
 //===----------------------------------------------------------------------===//
 
 FailureOr<LinalgOp> transform::GeneralizeOp::applyToOne(LinalgOp target) {
   // Exit early if no transformation is needed.
   if (isa<GenericOp>(target))
     return target;
 
FailureOr<LinalgOp> transform::ScalarizeOp::applyToOne(LinalgOp target) {
 
   return result->op;
 }
 
 //===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
 
-/// Apply a tiling transformation to all payload ops and store both the
-/// tiled operation as well as the created tile loops.
-static LogicalResult
-applyTilingToAll(Operation *transformOp, Value target,
-                 ArrayRef<int64_t> tileSizes,
-                 transform::TransformResults &transformResults,
-                 transform::TransformState &state,
-                 function_ref<FailureOr<TiledLinalgOp>(LinalgOp)> applyFn) {
-  // Number of loops: Number of tiles sizes that are not zero.
-  size_t numLoops = tileSizes.size() - llvm::count(tileSizes, 0);
-  // All payload ops. These should all be LinalgOps for now.
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(target);
-
-  SmallVector<Operation *> tiledLinalgOps;
-  SmallVector<SmallVector<Operation *>> loopOps(numLoops);
-  for (unsigned int i = 0; i < numLoops; ++i)
-    loopOps[i].reserve(payloadOps.size());
-
-  for (Operation *target : payloadOps) {
-    auto linalgOp = dyn_cast<linalg::LinalgOp>(target);
-    if (!linalgOp)
-      return transformOp->emitError("only LinalgOps are supported");
-
-    FailureOr<TiledLinalgOp> tiled = applyFn(linalgOp);
-    if (failed(tiled))
-      return failure();
-
-    tiledLinalgOps.push_back(tiled->op);
-    if (tiled->loops.size() != numLoops)
-      // Not enough loops were generated. This usually means that the input size
-      // was smaller than the tiling size.
-      // TODO: LinalgTilingPattern should return failure().
-      return failure();
-    for (unsigned int i = 0; i < numLoops; ++i)
-      loopOps[i].push_back(tiled->loops[i]);
-  }
-
-  transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);
-  for (unsigned int i = 0; i < numLoops; ++i)
-    transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);
-  return success();
-}
-
 LogicalResult transform::TileOp::apply(TransformResults &transformResults,
                                        TransformState &state) {
   LinalgTilingOptions tilingOptions;
   SmallVector<int64_t> tileSizes = extractI64Array(getSizes());
 
   if (!tileSizes.empty())
     tilingOptions.setTileSizes(tileSizes);
   tilingOptions.setInterchange(extractUIntArray(getInterchange()));
   LinalgTilingPattern pattern(getContext(), tilingOptions);
 
   return applyTilingToAll(getOperation(), getTarget(), tileSizes,
                           transformResults, state, [&](LinalgOp linalgOp) {
                             SimpleRewriter rewriter(linalgOp.getContext());
                             return pattern.returningMatchAndRewrite(linalgOp,
                                                                     rewriter);
                           });
 }
 
 ParseResult transform::TileOp::parse(OpAsmParser &parser,
                                      OperationState &result) {
-  StringRef sizesAttrName = TileOp::getSizesAttrName(result.name).getValue();
-  OpAsmParser::UnresolvedOperand targetOperand;
-  SMLoc opLoc = parser.getCurrentLocation();
-  if (parser.parseOperand(targetOperand) ||
-      parser.parseOptionalAttrDict(result.attributes))
-    return failure();
-  Attribute sizesAttr = result.attributes.get(sizesAttrName);
-  if (!sizesAttr)
-    return parser.emitError(opLoc)
-           << "expected '" << sizesAttrName << "' attribute";
-  auto sizesArrayAttr = sizesAttr.dyn_cast<ArrayAttr>();
-  if (!sizesArrayAttr)
-    return parser.emitError(opLoc)
-           << "'" << sizesAttrName << "' attribute must be an array";
-  Type pdlOpType = parser.getBuilder().getType<pdl::OperationType>();
-  size_t numExpectedLoops =
-      sizesArrayAttr.size() - llvm::count(extractI64Array(sizesArrayAttr), 0);
-  result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOpType));
-  if (parser.resolveOperand(targetOperand, pdlOpType, result.operands))
-    return failure();
-  return success();
+  return parseTileLikeOp(parser, result,
+                         TileOp::getSizesAttrName(result.name).getValue());
 }
 
 void TileOp::print(OpAsmPrinter &p) {
   p << ' ';
   p << getTarget();
   p.printOptionalAttrDict((*this)->getAttrs());
 }
 
-void TileOp::getEffects(
-    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
-        &effects) {
-  // `target` arg is consumed and can no longer be used.
-  effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
-                       TransformMappingResource::get());
-  effects.emplace_back(MemoryEffects::Free::get(), getTarget(),
-                       TransformMappingResource::get());
-
-  for (Value r : getResults()) {
-    effects.emplace_back(MemoryEffects::Write::get(), r,
-                         TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Allocate::get(), r,
-                         TransformMappingResource::get());
-  }
-
-  effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
-  effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
-}
-
 //===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
 FailureOr<Operation *> VectorizeOp::applyToOne(Operation *target) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     InFlightDiagnostic diag = emitOpError()
                               << "applies only to isolated-from-above targets";

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

 /// A simple pattern rewriter that implements no special logic.
 class SimpleRewriter : public PatternRewriter {
 public:
   SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// GetParentForOp
+// GetParentLoopOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
-transform::GetParentForOp::apply(transform::TransformResults &results,
+transform::GetParentLoopOp::apply(transform::TransformResults &results,
                                   transform::TransformState &state) {
   SetVector<Operation *> parents;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     scf::ForOp loop;
     Operation *current = target;
     for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
       loop = current->getParentOfType<scf::ForOp>();
       if (!loop) {
         InFlightDiagnostic diag = emitError() << "could not find an '"

 // LoopPeelOp
 //===----------------------------------------------------------------------===//
 
 FailureOr<scf::ForOp> transform::LoopPeelOp::applyToOne(scf::ForOp loop) {
   scf::ForOp result;
   IRRewriter rewriter(loop->getContext());
   LogicalResult status =
       scf::peelAndCanonicalizeForLoop(rewriter, loop, result);
-  if (failed(status)) {
-    if (getFailIfAlreadyDivisible())
-      return reportUnknownTransformError(loop);
+  if (failed(status))
     return loop;
-  }
   return result;
 }
 
 //===----------------------------------------------------------------------===//
 // LoopPipelineOp
 //===----------------------------------------------------------------------===//
 
 /// Callback for PipeliningOption. Populates `schedule` with the mapping from an

mlir/python/mlir/dialects/_loop_transform_ops_ext.py

def _get_int64_attr(arg: Optional[Union[int, IntegerAttr]],
 
   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."""
+class GetParentLoopOp:
+  """Extension for GetParentLoopOp."""
 
   def __init__(self,
                target: Union[Operation, Value],
                *,
                num_loops: int = 1,
                ip=None,
                loc=None):
     super().__init__(
super().__init__(
                    StringAttr.get(func_name)),
         ip=ip,
         loc=loc)
 
 
 class LoopPeelOp:
   """Extension for LoopPeelOp."""
 
-  def __init__(self,
-               target: Union[Operation, Value],
-               *,
-               fail_if_already_divisible: Union[bool, BoolAttr] = False,
-               ip=None,
-               loc=None):
+  def __init__(self, target: Union[Operation, Value], *, ip=None, loc=None):
     super().__init__(
         pdl.OperationType.get(),
         _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,

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

+// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @fuse_unary
+func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+
+  //     CHECK:   scf.for
+  //     CHECK:     scf.for
+  //     CHECK:       linalg.elemwise_unary
+  //     CHECK:       linalg.elemwise_binary
+  %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = linalg.elemwise_binary ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = pdl.operation "linalg.elemwise_binary"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]}
+  }
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fuse_unary
+func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+
+  //     CHECK:   scf.for
+  //     CHECK:     scf.for
+  //     CHECK:       linalg.elemwise_unary
+  //     CHECK:       linalg.elemwise_binary
+  //     CHECK:   scf.for
+  //     CHECK:     scf.for
+  //     CHECK:       linalg.elemwise_unary
+  //     CHECK:       linalg.elemwise_binary
+  %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = linalg.elemwise_binary ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = pdl.operation "linalg.elemwise_binary"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]}
+    transform.loop.peel %loops#0
+  }
+}

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

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
 
-// CHECK-LABEL: @get_parent_for_op
-func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
+// CHECK-LABEL: @get_parent_loop_op
+func.func @get_parent_loop_op(%arg0: index, %arg1: index, %arg2: index) {
   // expected-remark @below {{first loop}}
   scf.for %i = %arg0 to %arg1 step %arg2 {
     // expected-remark @below {{second loop}}
     scf.for %j = %arg0 to %arg1 step %arg2 {
       // expected-remark @below {{third loop}}
       scf.for %k = %arg0 to %arg1 step %arg2 {
         arith.addi %i, %j : index
       }
pdl.pattern @match_addi : benefit(1) {
     %results = types
     %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     rewrite %op with "transform.dialect"
   }
 
   sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_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 }
+    // CHECK: = transform.loop.get_parent_loop
+    %1 = transform.loop.get_parent_loop %0
+    %2 = transform.loop.get_parent_loop %0 { num_loops = 2 }
+    %3 = transform.loop.get_parent_loop %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"
   }
 }
 
 // -----
 
-func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
+func.func @get_parent_loop_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):
   pdl.pattern @match_addi : benefit(1) {
     %args = operands
     %results = types
     %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     rewrite %op with "transform.dialect"
   }
 
   sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_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_loop %0
   }
 }
 
 // -----
 
 // Outlined functions:
 //
 // CHECK: func @foo(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}})
pdl.pattern @match_addi : benefit(1) {
     %results = types
     %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     rewrite %op with "transform.dialect"
   }
 
   sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_addi in %arg1
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_loop %0
     // CHECK: = transform.loop.outline %{{.*}}
     transform.loop.outline %1 {func_name = "foo"}
   }
 }
 
 // -----
 
 func.func private @cond() -> i1
pdl.pattern @match_addi : benefit(1) {
     %results = types
     %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     rewrite %op with "transform.dialect"
   }
 
   sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_addi in %arg1
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_loop %0
     transform.loop.peel %1
   }
 }
 
 // -----
 
 func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
   %c0 = arith.constant 0 : index
pdl.pattern @match_addf : benefit(1) {
     %results = types
     %op = operation "arith.addf"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     rewrite %op with "transform.dialect"
   }
 
   sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_addf in %arg1
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_loop %0
     %2 = transform.loop.pipeline %1
     // Verify that the returned handle is usable.
     transform.test_print_remark_at_operand %2, "transformed"
   }
 }
 
 // -----
 
pdl.pattern @match_addi : benefit(1) {
     %results = types
     %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     rewrite %op with "transform.dialect"
   }
 
   sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_addi in %arg1
-    %1 = transform.loop.get_parent_for %0
+    %1 = transform.loop.get_parent_loop %0
     transform.loop.unroll %1 { factor = 4 }
   }
 }
 

mlir/test/python/dialects/transform_loop_ext.py

with Context(), Location.unknown():
     print(module)
   return f
 
 
 @run
 def getParentLoop():
   sequence = transform.SequenceOp()
   with InsertionPoint(sequence.body):
-    loop.GetParentForOp(sequence.bodyTarget, num_loops=2)
+    loop.GetParentLoopOp(sequence.bodyTarget, num_loops=2)
     transform.YieldOp()
   # CHECK-LABEL: TEST: getParentLoop
-  # CHECK: = transform.loop.get_parent_for %
+  # CHECK: = transform.loop.get_parent_loop %
   # CHECK: num_loops = 2
 
 
 @run
 def loopOutline():
   sequence = transform.SequenceOp()
   with InsertionPoint(sequence.body):
     loop.LoopOutlineOp(sequence.bodyTarget, func_name="foo")