[SCF][Transform] Add transform.loop.fuse_sibling

[Groverkss]

Summary

This patch adds a new transform operation transform.loop.fuse_sibling,
which given two loops, fuses them, assuming that they are independent.
The transform operation itself performs very basic checks to ensure
IR legality, and leaves the responsibility of ensuring independence on the user.

[springerm]

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

319–320

Mention which kind of "loop" ops are supported. Also mention that the bounds must match and that a silenceable failure is produced otherwise. Also mention that both input handles are mapped to exactly one op and what happens if they are not.

324–325

Add the following: "This operation consumes the loop and sibling handles and produces the fused_loop handle."

328–329

I would rename these to indicate which loop is fused into which one. Maybe something like target and source.

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

377–379

What happens if the result of one loop is used as an "shared_outs" operand of the other loop? Also, could there be issues where the resulting IR violates dominance because some operands of ops inside the loop are not defined after the other loop. I think it should be possible to write a relatively simple check with getUsedValuesDefinedAbove.

382–383

I would make this and the isOpSibiling check a silenceable error. When there is an obvious error like "op type does not match", indicating there's something inherently wrong with the transform dialect script, I usually make it a definite error.

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

974–975

Can you give this function a RewriterBase & and use that rewriter to modify the IR. That would make it easier to reuse this helper function.

1025–1033

These should then also go through the rewriter.

[nicolasvasilache]

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

320

The legality of loop fusion is a tricky topic, we can't just approximate with "neither is an ancestor of another".

You should at least restrict the transform to only work in the case of loops with no recursive side effects (i.e. on tensors only and without function calls like print).
And this should also be well documented.

[ftynse]

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

319–320

Can this description be more extensive? It should mention that only forall loops can be fused, and which preconditions should hold for the loops. Also mention that there is no real legality check performed (see below).

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

331–332

The bounds being the same is not a sufficient criterion for the fusion to be legal. Consider

scf.forall %i in %bounds
  store %i, %mem[%i]

scf.forall %i in %bounds
  load %mem[%bounds - %i]

that would trivially be incorrect if two loops are fused.

I'm not necessary opposed to having a "dangerous" transform exposed as an op, but it must be very clearly documented as such.

335–340

isa followed by cast, even at a distance, is an anti-pattern. Call dyn_cast and check the resulting variables for being non-null instead.

369–372

This should be documented in the op description.

379

Nit: don't start error messages with a capital letter.

389

Nit: it's possible to construct an ArrayRef temporary from a single element, no need to use a vector here.

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

977

Pass in RewriterBase & as an argument to this function and use the rewriter available in the apply method.

980

Nit: expand auto unless the type is obvious from context (e.g., there's a cast on the RHS).

982

Nit: don't specify the number of vector stack elements unless there is a very good reason to.

983–985

Something like llvm::to_vector(llvm::concat(range1, range2)) would probably work here.

1025–1033

This must use proper rewriting API, e.g. replaceOp and eraseOp. Low-level Operation API will result to catastrophic memory problems when used with proper rewriters.

[ftynse]

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

327–329

This will accept loops from, e.g., different functions. A better, and faster, condition is to check that both operations belong to the same block.

Also, depending on the understanding of "sibling", this may accept as siblings two loops that have other operations (including other loops!) between them, which may not be legal.

[Groverkss]

• Address comments

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

320

Right. I don't really want to make this operation check for the legality of loop fusion other than something that may produce invalid IR as that is not easy, but I still want this operation to be useable in those cases. I have changed the description to reflect this better.

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

327–329

I updated the documentation to describe what this operation is meant to do better.

331–332

Added documentation to make it clear that it is the responsibility of the user to consider these.

377–379

If the shared_outs use happens, then it's a dependence. This is what I wanted to check with this operation but wrongly assumed that isAncestor does this.

What I want to check is if any result produced by the first loop is used by the second loop.

For the dominance problem, I added a check that all values used by target must be defined before source.

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

983–985

For some reason, to_vector refuses to work with this. Tried: llvm::to_vector(llvm::concat<Value>(range1, range2))

But that fails with:

error: taking the address of a temporary object of type 'mlir::Value' [-Waddress-of-temporary]

Using the old implementation for now.

[ftynse]

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

316

Nit: I'd just say "assuming the fusion is legal". Consumer/producer fusion is also fine here IMO.

326

s/matching/mapping

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

189–190

Nit: in my understanding, "siblings" doesn't imply independence and vice versa, these are two separate conditions that must be met. Two adjacent loops can depend on each other, e.g. in producer/consumer way or have the first loop compute the trip count of the second loop.

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

328

335

340

Not all operands necessarily have a defining op, they may be block arguments, at which point this will be dereferencing a null pointer and likely crash.

Furthermore, I suspect this condition one of the conditions here is either incorrect or unnecessary. I'm confused by the use of "source" and "target" here, so I'll just my own example:

%results = scf.forall {  // first_loop
}
// something in-between
scf.forall %operands {  // second_loop
}

If the body of first_loop is merged into second_loop, then the condition that has to be ensured is that all uses of %results are after second_loop. Since second_loop already follows first_loop, all operands of first_loop are already visible at second_loop if the input IR is valid, no need to check for that.
If the body of the second_loop is merged into the first_loop, then the condition is that operands of second_loop are visible at first_loop, but the results are already okay.

Also note that dominance is more complex than just order of operations in blocks. Operands may be defined in surrounding regions, or in dominating blocks that transfer control to the given block. Results can be used in nested regions, or dominated blocks to which the control is transferred. It is possible to assume (and actually check) that the current block belongs to a single-block region, which is the most common case in presence of structured control flow. When done carefully, this entire test will conservatively return failure. However, I am encouraging you to consider at least nested regions, otherwise the utility of this transform is quite limited.

358

Nit: I would rename this function to make it better correspond to what it actually does, e.g., isForallWithIdenticalConfiguration. Otherwise somebody will eventually try using it for actual fusion legality tests.

359–360

Nit: consider https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code. Here and below.

391

Nit: please follow https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements here and below.

403

I may sound obsessed with the term, but "not siblings" doesn't describe the actual check that has taken place, which also includes dominance checks. Couldn't you just return DiagnosedSilenceableFailure from the helper function and put a more descriptive error message there?

[Groverkss]

• Add better dominance checking startergy, add more tests, address comments

[Groverkss]

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

189–190

My understanding of siblings is that they are two loops that do not behave in a producer/consumer fashion. A result/side-effect produced by the first loop, should not be consumed by the other loop. Isn't this the same as being independent? Or am I missing something?

I'm happy to change the name if you feel this is not the right word.

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

340

Regarding the unnecessary conditions, I now check the ordering and only do the required checks.

Regarding dominance, I switched to using DominanceInfo, will that be enough? Also, I added some more test cases which check for dominance failures.

While adding these tests, I also noticed that not only do the operands need to be checked, but we also need to check if any value used by the second_loop should be visible by the first loop. So, I added another check for this.

[Groverkss]

• Add comment to explain context of isOpSibling

[ftynse]

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

189–190

In my understanding, sibling operations are merely operations that have a common parent. Nothing more. This is a natural extension of the parent/child terminology used extensively for operations nested in regions of other operations: two "child" operations that have a common parent are siblings. It doesn't imply anything about data flow between them. In non-polyhedral loop transformation terminology, "sibling fusion" is sometimes used in opposition to "producer/consumer fusion", but this is a very niche usage compared to the basic IR structure usage that is more pervasive.

I could live with something like doIndependentSiblingFusion so it is clearer that "sibling" qualifies the kind of fusion and not the relation between operations, but a more verbose and less ambiguous name is welcome.

[Groverkss]

• Change method name

[Groverkss]

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

189–190

Changed it to fuseIndependentSiblingForallLoops.

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

let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::scf::IfOp ifOp,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def LoopFuseSibling : Op<Transform_Dialect, "loop.fuse_sibling",
+  [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+   DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let summary = "Fuse a loop into another loop, assuming they are independent.";
+
+  let description = [{
+    Fuses the `target` loop into the `source` loop assuming they are
+    independent of each other. It is the responsibility of the user to ensure
+    that the given two loops are independent of each other, this operation will
+    not performa any legality checks and will simply fuse the two given loops.
+
+    Currently, the only fusion supported is when both `target` and `source`
+    are `scf.forall` operations. For `scf.forall` fusion, the bounds and the
+    matching must match, otherwise a silencable failure is produced.
+
+    The input handles `target` and `source` must map to exactly one operation,
+    a definite failure is produced otherwise.
+
+    #### Return modes
+
+    This operation consumes the `target` and `source` handles and produces the
+    `fused_loop` handle, which points to the fused loop.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       TransformHandleTypeInterface:$source);
+  let results = (outs TransformHandleTypeInterface:$fused_loop);
+  let assemblyFormat = "$target `into` $source attr-dict "
+                       " `:` functional-type(operands, results)";
+
+  let builders = [
+    OpBuilder<(ins "Value":$loop, "Value":$fused_loop)>
+  ];
+}
+
 #endif // SCF_TRANSFORM_OPS

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

 
 /// Get perfectly nested sequence of loops starting at root of loop nest
 /// (the first op being another AffineFor, and the second op - a terminator).
 /// A loop is perfectly nested iff: the first op in the loop's body is another
 /// AffineForOp, and the second op is a terminator).
 void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
                              scf::ForOp root);
 
+/// Given two scf.forall loops, `target` and `source`, fuses `target` into
+/// `source`. Assumes that the given loops are "siblings", i.e. they are
+/// independent of each other.
+///
+/// This function does not perform any legality checks and simply fuses the
+/// loops. The caller is responsible for ensuring that the loops are legal to
+/// fuse.
+scf::ForallOp fuseSiblingForallLoops(scf::ForallOp target, scf::ForallOp source,
+                                     RewriterBase &rewriter);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_

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

 
 void transform::TakeAssumedBranchOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getTarget(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
+// LoopFuseSibling
+//===----------------------------------------------------------------------===//
+
+/// Check if `target` and `source` are siblings.
+///
+/// This is a simple check that just checks if both operations are in the same
+/// block and some checks to ensure that the fused IR does not voilate
+/// dominance.
+static bool isOpSibling(Operation *target, Operation *source) {
+  // Check if both operations are in the same block.
+  if (target->getBlock() != source->getBlock())
+    return false;
+
+  // Check if fusion will voilate dominance. We check that every operand of
+  // `target` dominates `source` and every result of `target` is dominated by
+  // `source`.
+  for (Value operand : target->getOperands()) {
+    // Operand should be strictly before `source` in the block.
+    if (!operand.getDefiningOp()->isBeforeInBlock(source))
+      return false;
+  }
+  for (Operation *user : target->getUsers()) {
+    // User should be strictly after `source` in the block.
+    if (!source->isBeforeInBlock(user))
+      return false;
+  }
+
+  return true;
+}
+
+/// Check if `target` can be fused into `source`.
+///
+/// This is a simple check that just checks if both loops have same
+/// bounds, steps and mapping. This check does not ensure that the side effects
+/// of `target` are independent of `source` or vice-versa. It is the
+/// responsibility of the caller to ensure that.
+static bool isFusionLegal(Operation *target, Operation *source) {
+  if (auto targetOp = dyn_cast<scf::ForallOp>(target)) {
+    if (auto sourceOp = dyn_cast<scf::ForallOp>(source)) {
+      return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
+             targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
+             targetOp.getMixedStep() == sourceOp.getMixedStep() &&
+             targetOp.getMapping() == sourceOp.getMapping();
+    }
+    return false;
+  }
+  // TODO: Add fusion for more operations. Currently, we handle only scf.forall.
+  return false;
+}
+
+/// Fuse `target` into `source` assuming they are siblings.
+static Operation *fuseSiblings(Operation *target, Operation *source,
+                               RewriterBase &rewriter) {
+  if (auto targetOp = dyn_cast<scf::ForallOp>(target)) {
+    if (auto sourceOp = dyn_cast<scf::ForallOp>(source))
+      return fuseSiblingForallLoops(targetOp, sourceOp, rewriter);
+    return nullptr;
+  }
+  // TODO: Add fusion for more operations. Currently, we handle only scf.forall.
+  return nullptr;
+}
+
+DiagnosedSilenceableFailure
+transform::LoopFuseSibling::apply(transform::TransformRewriter &rewriter,
+                                  transform::TransformResults &results,
+                                  transform::TransformState &state) {
+  auto targetOps = state.getPayloadOps(getTarget());
+  auto sourceOps = state.getPayloadOps(getSource());
+
+  if (!llvm::hasSingleElement(targetOps) || !llvm::hasSingleElement(sourceOps))
+    return emitDefiniteFailure()
+           << "requires exactly one target handle (got "
+           << llvm::range_size(targetOps) << ") and exactly one "
+           << "source handle (got " << llvm::range_size(sourceOps) << ")";
+
+  Operation *target = *targetOps.begin();
+  Operation *source = *sourceOps.begin();
+
+  // Check if the target and source are siblings.
+  if (!isOpSibling(target, source))
+    return emitSilenceableFailure(target->getLoc())
+           << "operations are not siblings";
+
+  // Check if the target can be fused into source.
+  if (!isFusionLegal(target, source))
+    return emitSilenceableFailure(target->getLoc())
+           << "operations cannot be fused";
+
+  Operation *fusedLoop = fuseSiblings(target, source, rewriter);
+  assert(fusedLoop && "failed to fuse operations");
+
+  results.set(cast<OpResult>(getFusedLoop()), {fusedLoop});
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 class SCFTransformDialectExtension
     : public transform::TransformDialectExtension<
           SCFTransformDialectExtension> {
 public:

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

TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
 
   // TODO: for now we just ignore the result of band isolation.
   // In the future, mapping decisions may be impacted by the ability to
   // isolate perfectly nested bands.
   (void)tryIsolateBands(tileLoops);
 
   return tileLoops;
 }
+
+scf::ForallOp mlir::fuseSiblingForallLoops(scf::ForallOp target,
+                                           scf::ForallOp source,
+                                           RewriterBase &rewriter) {
+  unsigned numTargetOuts = target.getNumResults();
+  unsigned numSourceOuts = source.getNumResults();
+
+  OperandRange targetOuts = target.getOutputs();
+  OperandRange sourceOuts = source.getOutputs();
+
+  // Create fused shared_outs.
+  SmallVector<Value> fusedOuts;
+  fusedOuts.reserve(numTargetOuts + numSourceOuts);
+  fusedOuts.append(targetOuts.begin(), targetOuts.end());
+  fusedOuts.append(sourceOuts.begin(), sourceOuts.end());
+
+  // Create a new scf::forall op after the source loop.
+  rewriter.setInsertionPointAfter(source);
+  scf::ForallOp fusedLoop = rewriter.create<scf::ForallOp>(
+      source.getLoc(), source.getMixedLowerBound(), source.getMixedUpperBound(),
+      source.getMixedStep(), fusedOuts, source.getMapping());
+
+  // Map control operands.
+  IRMapping fusedMapping;
+  fusedMapping.map(target.getInductionVars(), fusedLoop.getInductionVars());
+  fusedMapping.map(source.getInductionVars(), fusedLoop.getInductionVars());
+
+  // Map shared outs.
+  fusedMapping.map(target.getOutputBlockArguments(),
+                   fusedLoop.getOutputBlockArguments().slice(0, numTargetOuts));
+  fusedMapping.map(
+      source.getOutputBlockArguments(),
+      fusedLoop.getOutputBlockArguments().slice(numTargetOuts, numSourceOuts));
+
+  // Append everything except the terminator into the fused operation.
+  rewriter.setInsertionPointToStart(fusedLoop.getBody());
+  for (Operation &op : target.getLoopBody().begin()->without_terminator())
+    rewriter.clone(op, fusedMapping);
+  for (Operation &op : source.getLoopBody().begin()->without_terminator())
+    rewriter.clone(op, fusedMapping);
+
+  // Fuse the old terminator in_parallel ops into the new one.
+  scf::InParallelOp targetTerm = target.getTerminator();
+  scf::InParallelOp sourceTerm = source.getTerminator();
+  scf::InParallelOp fusedTerm = fusedLoop.getTerminator();
+
+  rewriter.setInsertionPointToStart(fusedTerm.getBody());
+  for (Operation &op : targetTerm.getYieldingOps())
+    rewriter.clone(op, fusedMapping);
+  for (Operation &op : sourceTerm.getYieldingOps())
+    rewriter.clone(op, fusedMapping);
+
+  // Replace all uses of the old loops with the fused loop.
+  rewriter.replaceAllUsesWith(target.getResults(),
+                              fusedLoop.getResults().slice(0, numTargetOuts));
+  rewriter.replaceAllUsesWith(
+      source.getResults(),
+      fusedLoop.getResults().slice(numTargetOuts, numSourceOuts));
+
+  // Erase the old loops.
+  rewriter.eraseOp(target);
+  rewriter.eraseOp(source);
+
+  return fusedLoop;
+}

mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter --cse --canonicalize -split-input-file -verify-diagnostics | FileCheck %s
+
+func.func @test(%A : tensor<128x128xf32>, %B1 : tensor<128x128xf32>, %B2 : tensor<128x128xf32>) -> (tensor<128x128xf32>, tensor<128x128xf32>) {
+  %zero = arith.constant 0.0 : f32
+  %out_alloc = tensor.empty() : tensor<128x128xf32>
+  %out = linalg.fill ins(%zero : f32) outs(%out_alloc : tensor<128x128xf32>) -> tensor<128x128xf32>
+
+  // CHECK: scf.forall ([[I:%.*]]) in (4) shared_outs([[S1:%.*]] = [[IN1:%.*]], [[S2:%.*]] = [[IN2:%.*]]) -> (tensor<128x128xf32>, tensor<128x128xf32>) {
+  // CHECK:   [[T:%.*]] = affine.apply
+  // CHECK:   tensor.extract_slice [[S1]][[[T]], 0] [32, 128] [1, 1]
+  // CHECK:   [[OUT1:%.*]] = linalg.matmul
+  // CHECK:   tensor.extract_slice [[S2]][[[T]], 0] [32, 128] [1, 1]
+  // CHECK:   [[OUT2:%.*]] = linalg.matmul
+  // CHECK:   scf.forall.in_parallel {
+  // CHECK:     tensor.parallel_insert_slice [[OUT1]] into [[S1]][[[T]], 0] [32, 128] [1, 1]
+  // CHECK:     tensor.parallel_insert_slice [[OUT2]] into [[S2]][[[T]], 0] [32, 128] [1, 1]
+  // CHECK:   }
+  // CHECK: }
+  %out1 = linalg.matmul ins(%A, %B1 : tensor<128x128xf32>, tensor<128x128xf32>) outs(%out : tensor<128x128xf32>) -> tensor<128x128xf32>
+  %out2 = linalg.matmul ins(%A, %B2 : tensor<128x128xf32>, tensor<128x128xf32>) outs(%out : tensor<128x128xf32>) -> tensor<128x128xf32>
+
+  func.return %out1, %out2 : tensor<128x128xf32>, tensor<128x128xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%variant_op : !transform.any_op):
+  %matched = transform.structured.match ops{["linalg.matmul"]} in %variant_op : (!transform.any_op) -> (!transform.any_op)
+
+  %mm1, %mm2 = transform.split_handle %matched : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+
+  %loop1, %tiled_mm1 = transform.structured.tile_to_forall_op %mm1 tile_sizes [32] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+  %loop2, %tiled_mm2 = transform.structured.tile_to_forall_op %mm2 tile_sizes [32] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+
+  %fused_loop = transform.loop.fuse_sibling %loop1 into %loop2 : (!transform.any_op, !transform.any_op) -> !transform.any_op
+}