[MLIR][Transform] Introduce loop.coalesce transform op.

[kaitingwang]

Summary

This patch made a minor refactor of LoopCoalescing.cpp's walkLoops templated method and placed it in Affine's LoopUtils.cpp/h.
This method is also renamed as coalescePerfectlyNestedLoops method. This minor change enables this method to be invoked
by both the original LoopCoalescing pass as well as the newly introduced loop.coalesce transform op.

The loop.coalesce transform op has the ability to coalesce affine, and scf loop nests, leveraging existing LoopCoalescing
mechanism. I have created it inside the SCFTransformOps.td instead of AffineTransformOps.td as it feels to be similar
in spirit as the loop.unroll op that can handle both scf and affine loops. Please let me know if you feel that this op
should be moved into AffineTransformOps.td instead.

The testcase added illustrates loop.coalesce transform op working for scf, affine loops (inner, outer) as well as
coalesced loop can be further unrolled (achieving composibility).

[ftynse]

mlir/include/mlir/Dialect/Affine/LoopUtils.h

46–48

This won't work without explicitly listing available specializations as extern. Or moving the body of the function template to the header.

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

12 ↗

(On Diff #487110)

Please don't add back the dependency on PDL, we don't want it in most cases.

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

12

Ditto.

195

197–198

202–203

Does the transform fail if the coalescing didn't happen? It can still return the original loop so the transformation continues if the silenceable failures are suppressed.

206

There is absolutely no reason why this should be a PDL_Operation and not a TransformHandleTypeInterface.

209

Nit: can we just use functional-type for a more homogeneous syntax?

mlir/lib/Dialect/Affine/Transforms/LoopCoalescing.cpp

37–42

You can request a PreOrder walk via a template parameter, and skip descending into nested operations once the first loop is found.

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

2848

Nit: please don't specify the explicit number of stack elements unless you have a strong reason to pick a specific number. Here and below.

2899–2900

This should be in the header.

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

264

It would be better if this emitted a silenceable failure when the input was not a supported loop.

mlir/test/Dialect/SCF/transform-op-coalesce.mlir

2

There doesn't seem to be a need to allow unregistered dialects in this test.

8

Please use sigils (%, #, etc.) outside the pattern match. Some very old tests may be doing otherwise, but they are not the example to follow.

[kaitingwang]

Much appreciated the detail review comments from @ftynse ! I went over each of them carefully and hopefully not missing anything. Please let me know if further improvements are needed! Thank you once again.

[ftynse]

Thanks for iterating! I have another batch of comments and this should be good to go.

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

195

This doesn't seem to have been addressed. If you are looking at reviews in the email, consider opening the web version. It contains suggested changes to the code that are not being sent over email.

197–198

Ditto.

mlir/lib/Dialect/Affine/Transforms/LoopCoalescing.cpp

39–41

Could this check the result of coalescePerfectlyNestedLoops and, if it is a failure, emit an error and signal pass failure? I see we were not doing this before because coalescePerfectlyNestedLoops was not propagating errors, but now it does, so sounds like a usability improvement. Good for a separate diff, too.

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

268

It is invalid to create a freestanding Note diagnostic, it must always be attached to a diagnostic with larger severity. This code can use emitSilenceableError() instead of manually creating the diagnostic.

Please add a test exercising this failure mode. We usually have tests for all user-visible error messages, which let us catch situations like the one described above.

269

Nit: start diagnostic messages from a lower-case letter.

mlir/test/Dialect/SCF/transform-op-coalesce.mlir

8

Sorry, I wasn't clear here. I mean something like %[[IV0:.*]] so the % symbol is present in the CHECK line, but is not part of the regular expression.

62–68

There are still some % inside the regular expression here, they should go outside, like in %[[IV2:[a-zA-Z0-9]+]].

[kaitingwang]

Thanks @ftynse ! Indeed, I didn't realize the review e-mail doesn't cover all the comments. I also managed to spectacularly misunderstand your comments on the sigils being 'outside' the pattern match. Hopefully, this time catches it all.

I didn't make the runOnOperation for LoopCoalescing.cpp return signalPassFailure() as you suggested because test/Dialect/Affine/loop-coalescing.mlir has a case (the last case in the test) where no coalescing takes place yet expecting the pass to return 0. I wasn't sure if I should modified this assumed behavior and this testcase. Let me know what you think and I'd be happy to modify it further.

[kaitingwang]

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

195

Indeed I missed it. Your wording sounds much better!

mlir/test/Dialect/SCF/transform-op-coalesce.mlir

8

Your comment was actually quite clear...Not sure what I was thinking the first time. :) Thanks for following up.

[ftynse]

Thanks!

I didn't make the runOnOperation for LoopCoalescing.cpp return signalPassFailure() as you suggested because test/Dialect/Affine/loop-coalescing.mlir has a case (the last case in the test) where no coalescing takes place yet expecting the pass to return 0. I wasn't sure if I should modified this assumed behavior and this testcase. Let me know what you think and I'd be happy to modify it further.

I see. A comment above the error-suppressing line along the lines of "this pass should succeed even if no coalescing happened" would be helpful if somebody looks at this later.

mlir/include/mlir/Dialect/Affine/LoopUtils.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_LOOPUTILS_H
 #define MLIR_DIALECT_AFFINE_LOOPUTILS_H
 
 #include "mlir/IR/Block.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/RegionUtils.h"
 #include <optional>
 
 namespace mlir {
 class AffineForOp;
 class AffineMap;
 class LoopLikeOpInterface;
 struct MemRefRegion;
 class OpBuilder;
 class Value;
 class ValueRange;
 
 namespace func {
 class FuncOp;
 } // namespace func
 
 namespace scf {
 class ForOp;
 class ParallelOp;
 } // namespace scf
 
 /// Unrolls this for operation completely if the trip count is known to be
 /// constant. Returns failure otherwise.
 LogicalResult loopUnrollFull(AffineForOp forOp);
 
 /// Unrolls this for operation by the specified unroll factor. Returns failure
 /// if the loop cannot be unrolled either due to restrictions or due to invalid
 /// unroll factors. Requires positive loop bounds and step. If specified,
 /// annotates the Ops in each unrolled iteration by applying `annotateFn`.
 /// When `cleanUpUnroll` is true, we can ensure the cleanup loop is unrolled
 /// regardless of the unroll factor.
 LogicalResult loopUnrollByFactor(
     AffineForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr,
     bool cleanUpUnroll = false);
 

 ///      // full_tile
 ///    else
 ///      // partial tile
 ///
 LogicalResult
 separateFullTiles(MutableArrayRef<AffineForOp> nest,
                   SmallVectorImpl<AffineForOp> *fullTileNest = nullptr);
 
+/// Walk either an scf.for or an affine.for to find a band to coalesce.
+template <typename LoopOpTy>
+LogicalResult coalescePerfectlyNestedLoops(LoopOpTy op) {
+  LogicalResult result(failure());
+  SmallVector<LoopOpTy> loops;
+  getPerfectlyNestedLoops(loops, op);
+
+  // Look for a band of loops that can be coalesced, i.e. perfectly nested
+  // loops with bounds defined above some loop.
+  // 1. For each loop, find above which parent loop its operands are
+  // defined.
+  SmallVector<unsigned, 4> operandsDefinedAbove(loops.size());
+  for (unsigned i = 0, e = loops.size(); i < e; ++i) {
+    operandsDefinedAbove[i] = i;
+    for (unsigned j = 0; j < i; ++j) {
+      if (areValuesDefinedAbove(loops[i].getOperands(), loops[j].getRegion())) {
+        operandsDefinedAbove[i] = j;
+        break;
+      }
+    }
+  }
+
+  // 2. Identify bands of loops such that the operands of all of them are
+  // defined above the first loop in the band.  Traverse the nest bottom-up
+  // so that modifications don't invalidate the inner loops.
+  for (unsigned end = loops.size(); end > 0; --end) {
+    unsigned start = 0;
+    for (; start < end - 1; ++start) {
+      auto maxPos =
+          *std::max_element(std::next(operandsDefinedAbove.begin(), start),
+                            std::next(operandsDefinedAbove.begin(), end));
+      if (maxPos > start)
+        continue;
+      assert(maxPos == start &&
+             "expected loop bounds to be known at the start of the band");
+      auto band = llvm::makeMutableArrayRef(loops.data() + start, end - start);
+      if (succeeded(coalesceLoops(band)))
+        result = success();
+      break;
+    }
+    // If a band was found and transformed, keep looking at the loops above
+    // the outermost transformed loop.
+    if (start != end - 1)
+      end = start + 1;
+  }
+  return result;
+}
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_LOOPUTILS_H

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

 //===- SCFTransformOps.td - SCF (loop) transformation ops --*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SCF_TRANSFORM_OPS
 #define SCF_TRANSFORM_OPS
 
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def Transform_ScfForOp : Transform_ConcreteOpType<"scf.for">;
 
def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def LoopCoalesceOp : Op<Transform_Dialect, "loop.coalesce", [
+  FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+  TransformOpInterface, TransformEachOpTrait]> {
+  let summary = "Coalesces the perfect loop nest enclosed by a given loop";
+  let description = [{
+    Given a perfect loop nest identified by the outermost loop,
+    perform loop coalescing in a bottom-up one-by-one manner.
+
+    #### Return modes
+
+    The return handle points to the coalesced loop if coalescing happens, or
+    the given input loop if coalescing does not happen.
+  }];
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs TransformHandleTypeInterface:$transformed);
+
+  let assemblyFormat =
+      "$target attr-dict `:` functional-type($target, $transformed)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // SCF_TRANSFORM_OPS

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

 std::optional<std::pair<AffineExpr, AffineExpr>>
 getSCFMinMaxExpr(Value value, SmallVectorImpl<Value> &dims,
                  SmallVectorImpl<Value> &symbols,
                  llvm::function_ref<bool(Operation *)> loopFilter = nullptr);
 
 /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
 /// `loops` contains a list of perfectly nested loops with bounds and steps
 /// independent of any loop induction variable involved in the nest.
-void coalesceLoops(MutableArrayRef<scf::ForOp> loops);
+LogicalResult coalesceLoops(MutableArrayRef<scf::ForOp> loops);
 
 /// Take the ParallelLoop and for each set of dimension indices, combine them
 /// into a single dimension. combinedDimensions must contain each index into
 /// loops exactly once.
 void collapseParallelLoops(scf::ParallelOp loops,
                            ArrayRef<std::vector<unsigned>> combinedDimensions);
 
 /// Promotes the loop body of a scf::ForOp to its containing block if the loop

mlir/lib/Dialect/Affine/Transforms/LoopCoalescing.cpp

 #define DEBUG_TYPE PASS_NAME
 
 using namespace mlir;
 
 namespace {
 struct LoopCoalescingPass
     : public impl::LoopCoalescingBase<LoopCoalescingPass> {
 
-  /// Walk either an scf.for or an affine.for to find a band to coalesce.
-  template <typename LoopOpTy>
-  static void walkLoop(LoopOpTy op) {
-    // Ignore nested loops.
-    if (op->template getParentOfType<LoopOpTy>())
-      return;
-
-    SmallVector<LoopOpTy, 4> loops;
-    getPerfectlyNestedLoops(loops, op);
-    LLVM_DEBUG(llvm::dbgs()
-               << "found a perfect nest of depth " << loops.size() << '\n');
-
-    // Look for a band of loops that can be coalesced, i.e. perfectly nested
-    // loops with bounds defined above some loop.
-    // 1. For each loop, find above which parent loop its operands are
-    // defined.
-    SmallVector<unsigned, 4> operandsDefinedAbove(loops.size());
-    for (unsigned i = 0, e = loops.size(); i < e; ++i) {
-      operandsDefinedAbove[i] = i;
-      for (unsigned j = 0; j < i; ++j) {
-        if (areValuesDefinedAbove(loops[i].getOperands(),
-                                  loops[j].getRegion())) {
-          operandsDefinedAbove[i] = j;
-          break;
-        }
-      }
-      LLVM_DEBUG(llvm::dbgs()
-                 << "  bounds of loop " << i << " are known above depth "
-                 << operandsDefinedAbove[i] << '\n');
-    }
-
-    // 2. Identify bands of loops such that the operands of all of them are
-    // defined above the first loop in the band.  Traverse the nest bottom-up
-    // so that modifications don't invalidate the inner loops.
-    for (unsigned end = loops.size(); end > 0; --end) {
-      unsigned start = 0;
-      for (; start < end - 1; ++start) {
-        auto maxPos =
-            *std::max_element(std::next(operandsDefinedAbove.begin(), start),
-                              std::next(operandsDefinedAbove.begin(), end));
-        if (maxPos > start)
-          continue;
-
-        assert(maxPos == start &&
-               "expected loop bounds to be known at the start of the band");
-        LLVM_DEBUG(llvm::dbgs() << "  found coalesceable band from " << start
-                                << " to " << end << '\n');
-
-        auto band = llvm::MutableArrayRef(loops.data() + start, end - start);
-        (void)coalesceLoops(band);
-        break;
-      }
-      // If a band was found and transformed, keep looking at the loops above
-      // the outermost transformed loop.
-      if (start != end - 1)
-        end = start + 1;
-    }
-  }
-
   void runOnOperation() override {
     func::FuncOp func = getOperation();
-    func.walk([&](Operation *op) {
+    func.walk<WalkOrder::PreOrder>([](Operation *op) {
       if (auto scfForOp = dyn_cast<scf::ForOp>(op))
-        walkLoop(scfForOp);
+        (void)coalescePerfectlyNestedLoops(scfForOp);
       else if (auto affineForOp = dyn_cast<AffineForOp>(op))
-        walkLoop(affineForOp);
+        (void)coalescePerfectlyNestedLoops(affineForOp);
     });
   }
 };
 
 } // namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLoopCoalescingPass() {
   return std::make_unique<LoopCoalescingPass>();
 }

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
 
 #define DEBUG_TYPE "loop-utils"
 
mlir::separateFullTiles(MutableArrayRef<AffineForOp> inputNest,
   elseBlock->getOperations().splice(std::prev(elseBlock->end()),
                                     firstLoop->getBlock()->getOperations(),
                                     Block::iterator(firstLoop));
 
   if (fullTileNest)
     *fullTileNest = std::move(fullTileLoops);
 
   return success();
 }

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

transform::LoopUnrollOp::applyToOne(Operation *op,
                                     transform::TransformState &state) {
   LogicalResult result(failure());
   if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op))
     result = loopUnrollByFactor(scfFor, getFactor());
   else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
     result = loopUnrollByFactor(affineFor, getFactor());
 
   if (failed(result)) {
-    Diagnostic diag(op->getLoc(), DiagnosticSeverity::Note);
-    diag << "Op failed to unroll";
-    return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                       << "failed to unroll";
+    return diag;
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
+// LoopCoalesceOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::LoopCoalesceOp::applyToOne(Operation *op,
+                                      transform::ApplyToEachResultList &results,
+                                      transform::TransformState &state) {
+  LogicalResult result(failure());
+  if (scf::ForOp scfForOp = dyn_cast<scf::ForOp>(op))
+    result = coalescePerfectlyNestedLoops(scfForOp);
+  else if (AffineForOp affineForOp = dyn_cast<AffineForOp>(op))
+    result = coalescePerfectlyNestedLoops(affineForOp);
+
+  results.push_back(op);
+  if (failed(result)) {
+    DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                       << "failed to coalesce";
+    return diag;
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 class SCFTransformDialectExtension
     : public transform::TransformDialectExtension<
           SCFTransformDialectExtension> {
 public:
   using Base::Base;
 
   void init() {
     declareGeneratedDialect<AffineDialect>();
     declareGeneratedDialect<func::FuncDialect>();
 
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
 
 void mlir::scf::registerTransformDialectExtension(DialectRegistry &registry) {
   registry.addExtensions<SCFTransformDialectExtension>();
 }

mlir/lib/Dialect/SCF/Utils/Utils.cpp

auto loopPieces = normalizeLoop(builder, innerBuilder, loop.getLoc(),
                                   loop.getLowerBound(), loop.getUpperBound(),
                                   loop.getStep(), loop.getInductionVar());
 
   loop.setLowerBound(loopPieces.lowerBound);
   loop.setUpperBound(loopPieces.upperBound);
   loop.setStep(loopPieces.step);
 }
 
-void mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
+LogicalResult mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
   if (loops.size() < 2)
-    return;
+    return failure();
 
   scf::ForOp innermost = loops.back();
   scf::ForOp outermost = loops.front();
 
   // 1. Make sure all loops iterate from 0 to upperBound with step 1.  This
   // allows the following code to assume upperBound is the number of iterations.
   for (auto loop : loops)
     normalizeLoop(loop, outermost, innermost);
LogicalResult mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
   // 4. Move the operations from the innermost just above the second-outermost
   // loop, delete the extra terminator and the second-outermost loop.
   scf::ForOp second = loops[1];
   innermost.getBody()->back().erase();
   outermost.getBody()->getOperations().splice(
       Block::iterator(second.getOperation()),
       innermost.getBody()->getOperations());
   second.erase();
+  return success();
 }
 
 void mlir::collapseParallelLoops(
     scf::ParallelOp loops, ArrayRef<std::vector<unsigned>> combinedDimensions) {
   OpBuilder outsideBuilder(loops);
   Location loc = loops.getLoc();
 
   // Presort combined dimensions.

mlir/test/Dialect/SCF/transform-op-coalesce.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
+
+func.func @coalesce_inner() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+
+  // CHECK: scf.for %[[IV0:.+]]
+  // CHECK:   scf.for %[[IV1:.+]]
+  // CHECK:     scf.for %[[IV2:.+]]
+  // CHECK-NOT:   scf.for %[[IV3:.+]]
+  scf.for %i = %c0 to %c10 step %c1 {
+    scf.for %j = %c0 to %c10 step %c1 {
+      scf.for %k = %i to %j step %c1 {
+        // Inner loop must have been removed.
+        scf.for %l = %i to %j step %c1 {
+          arith.addi %i, %j : index
+        }
+      } {coalesce}
+    }
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1
+  %1 = transform.cast %0 : !pdl.operation to !transform.op<"scf.for">
+  %2 = transform.loop.coalesce %1: (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
+}
+
+// -----
+
+func.func @coalesce_outer(%arg1: memref<64x64xf32, 1>, %arg2: memref<64x64xf32, 1>, %arg3: memref<64x64xf32, 1>) attributes {} {
+  // CHECK: affine.for %[[IV1:.+]] = 0 to %[[UB:.+]] {
+  // CHECK-NOT: affine.for %[[IV2:.+]]
+  affine.for %arg4 = 0 to 64 {
+    affine.for %arg5 = 0 to 64 {
+      // CHECK: %[[IDX0:.+]] = affine.apply #[[MAP0:.+]](%[[IV1]])[%{{.+}}]
+      // CHECK: %[[IDX1:.+]] = affine.apply #[[MAP1:.+]](%[[IV1]])[%{{.+}}]
+      // CHECK-NEXT: %{{.+}} = affine.load %{{.+}}[%[[IDX1]], %[[IDX0]]] : memref<64x64xf32, 1>
+      %0 = affine.load %arg1[%arg4, %arg5] : memref<64x64xf32, 1>
+      %1 = affine.load %arg2[%arg4, %arg5] : memref<64x64xf32, 1>
+      %2 = arith.addf %0, %1 : f32
+      affine.store %2, %arg3[%arg4, %arg5] : memref<64x64xf32, 1>
+    }
+  } {coalesce}
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["affine.for"]} attributes {coalesce} in %arg1
+  %1 = transform.cast %0 : !pdl.operation to !transform.op<"affine.for">
+  %2 = transform.loop.coalesce %1 : (!transform.op<"affine.for">) -> (!transform.op<"affine.for">)
+}
+
+// -----
+
+func.func @coalesce_and_unroll(%arg1: memref<64x64xf32, 1>, %arg2: memref<64x64xf32, 1>, %arg3: memref<64x64xf32, 1>) attributes {} {
+  // CHECK: scf.for %[[IV1:.+]] =
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+
+  scf.for %arg4 = %c0 to %c64 step %c1 {
+    // CHECK-NOT: scf.for
+    scf.for %arg5 = %c0 to %c64 step %c1 {
+      // CHECK: %[[IDX0:.+]] = arith.remsi %[[IV1]]
+      // CHECK: %[[IDX1:.+]] = arith.divsi %[[IV1]]
+      // CHECK-NEXT: %{{.+}} = memref.load %{{.+}}[%[[IDX1]], %[[IDX0]]] : memref<64x64xf32, 1>
+      %0 = memref.load %arg1[%arg4, %arg5] : memref<64x64xf32, 1>
+      %1 = memref.load %arg2[%arg4, %arg5] : memref<64x64xf32, 1>
+      %2 = arith.addf %0, %1 : f32
+      // CHECK: memref.store
+      // CHECK: memref.store
+      // CHECK: memref.store
+      // Residual loop must have a single store.
+      // CHECK: memref.store
+      memref.store %2, %arg3[%arg4, %arg5] : memref<64x64xf32, 1>
+    }
+  } {coalesce}
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1
+  %1 = transform.cast %0 : !pdl.operation to !transform.op<"scf.for">
+  %2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
+  transform.loop.unroll %2 {factor = 3} : !transform.op<"scf.for">
+}

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

+// RUN: mlir-opt %s -test-transform-dialect-interpreter --split-input-file --verify-diagnostics
+
+#map0 = affine_map<(d0) -> (d0 * 110)>
+#map1 = affine_map<(d0) -> (696, d0 * 110 + 110)>
+func.func @test_loops_do_not_get_coalesced() {
+  affine.for %i = 0 to 7 {
+    affine.for %j = #map0(%i) to min #map1(%i) {
+    }
+  } {coalesce}
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["affine.for"]} attributes {coalesce} in %arg1
+  %1 = transform.cast %0 : !pdl.operation to !transform.op<"affine.for">
+  // expected-error @below {{failed to coalesce}}
+  %2 = transform.loop.coalesce %1: (!transform.op<"affine.for">) -> (!transform.op<"affine.for">)
+}
+
+// -----
+
+func.func @test_loops_do_not_get_unrolled() {
+  affine.for %i = 0 to 7 {
+    arith.addi %i, %i : index
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+  %1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+  // expected-error @below {{failed to unroll}}
+  transform.loop.unroll %1 { factor = 8 } : !transform.op<"affine.for">
+}
+
+// -----
+
+func.func private @cond() -> i1
+func.func private @body()
+
+func.func @loop_outline_op_multi_region() {
+  // expected-note @below {{target op}}
+  scf.while : () -> () {
+    %0 = func.call @cond() : () -> i1
+    scf.condition(%0)
+  } do {
+  ^bb0:
+    func.call @body() : () -> ()
+    scf.yield
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["scf.while"]} in %arg1
+  // expected-error @below {{failed to outline}}
+  transform.loop.outline %0 {func_name = "foo"} : (!pdl.operation) -> !pdl.operation
+}

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

^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["arith.addi"]} in %arg1
   %1 = transform.loop.get_parent_for %0  : (!pdl.operation) -> !transform.op<"scf.for">
   // CHECK: = transform.loop.outline %{{.*}}
   transform.loop.outline %1 {func_name = "foo"} : (!transform.op<"scf.for">) -> !pdl.operation
 }
 
 // -----
 
-func.func private @cond() -> i1
-func.func private @body()
-
-func.func @loop_outline_op_multi_region() {
-  // expected-note @below {{target op}}
-  scf.while : () -> () {
-    %0 = func.call @cond() : () -> i1
-    scf.condition(%0)
-  } do {
-  ^bb0:
-    func.call @body() : () -> ()
-    scf.yield
-  }
-  return
-}
-
-transform.sequence failures(propagate) {
-^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["scf.while"]} in %arg1
-  // expected-error @below {{failed to outline}}
-  transform.loop.outline %0 {func_name = "foo"} : (!pdl.operation) -> !pdl.operation
-}
-
-// -----
-
 // CHECK-LABEL: @loop_peel_op
 func.func @loop_peel_op() {
   // CHECK: %[[C0:.+]] = arith.constant 0
   // CHECK: %[[C42:.+]] = arith.constant 42
   // CHECK: %[[C5:.+]] = arith.constant 5
   // CHECK: %[[C40:.+]] = arith.constant 40
   // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C40]] step %[[C5]]
   // CHECK:   arith.addi