[mlir] Introduce wrapInZeroTripCheck in LoopLikeOpInterface

[pzread]

Adds a new method wrapInZeroTripCheck to LoopLikeOpInterface, to create zero-trip-check around the loop

The purpose is to let loop ops (e.g. scf.while, scf.for) implement their own transformations to add zero-trip-check. The zero-trip-check creates a guard (e.g. scf.if) around the loop and the condition will be true only if the loop body will run at least once. An example usage is to combine with mlir::moveLoopInvariantCode to hoist resource-intense loop invariants into that guard region so they will only run once but also not run when the loop body won't run at all.

The implementation for scf.while can be found in the follow-up change: #80349

[llvmbot]

@llvm/pr-subscribers-mlir-scf
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Jerry Wu (pzread)

Changes

Adds a new method replaceWithZeroTripCheck to LoopLikeOpInterface, to create zero-trip-check around the loop

The purpose is to let loop ops (e.g. scf.while, scf.for) implement their own transformations to add zero-trip-check. The zero-trip-check creates a guard (e.g. scf.if) around the loop and the condition will be true only if the loop body will run at least once. An example usage is to hoist resource-intense loop invariants into this guard region so they will only run once but also not run when the loop body won't run at all.

The implementation of scf.while can be found in the follow-up change: #80349

---

Full diff: https://github.com/llvm/llvm-project/pull/80331.diff

1 Files Affected:

• (modified) mlir/include/mlir/Interfaces/LoopLikeInterface.td (+22)

diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index e2ac85a3f7725..77409cb3a8274 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -220,6 +220,28 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
       /*defaultImplementation=*/[{
         return ::mlir::failure();
       }]
+    >,
+    InterfaceMethod<[{
+        Add a zero-trip-check around the loop to check if the loop body is ever
+        run and return the new loop inside the check. The loop body is moved
+        over to the new loop. Returns "failure" if the loop doesn't support
+        this transformation.
+
+        After the transformation, the ops inserted to the parent region of the
+        loop are guaranteed to be run only if the loop body runs at least one
+        iteration.
+
+        Note: Ops in the loop body might be rearranged because of loop rotating
+        to maintain the semantic. Terminators might be removed/added during this
+        transformation.
+      }],
+      /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
+      /*methodName=*/"replaceWithZeroTripCheck",
+      /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::failure();
+      }]
     >
   ];
 

[dcaballe]

You can also mention that loop conditions might be checked redundantly so they are expected to not have side effects and that the utility won't check that but just apply the transformation.

[pzread]

Added a few words to mention that.

[joker-eph]

Can you add some tests for this?

[pzread]

Can you add some tests for this?

Done.

[matthias-springer]

Instead of creating a "check" (it is not really clear what the "check" is; it could be an scf.if, it could be a new basic block with a conditional branch, etc.), this method could also return just the condition. (Which evaluates to true if the loop has at least one iteration.) That would keep the interface a bit simpler. What do you think?

[pzread]

I think the tricky part of adding zero-trip-check is to do loop rotation when the loop condition has side-effects. So if we simply return the loop condition, callers will need to implement loop rotation by themselves for each loop op, which can be complicated.

If the kind of "check" op is a concern, maybe we can add a callback and let callers to create the check op in the callback by themselves?

[dcaballe]

Would it make sense to start with scf.if and generalize later as needed?

I think it's worth clarifying that loop rotation will only be possible for scf.while ops for now, that is, turning regular a "while-do" into a "do-while". That kind of rotation won't be possible at this level of abstraction for scf.for, hence the comment about redundant first iteration check. I think the rotation responsibility belongs to the interface itself.

If we want to make it extra configurable, we could separate the zero-trip-check method from the rotation method...

[matthias-springer]

This makes sense, I did not think of side effects.

The lambda sounds like a good idea, but I'm not sure what exactly it should look like. (Should it return a Block * into which the loop is moved?) We can generalize later if needed.

I think it would also work for scf.for. The check condition is lb < ub.

[pzread]

I think whether the loop rotation is needed (or possible) is the implementation details of each type of loop. scf.while needs it because the before block is in the loop and can contain ops with side-effects. scf.for is simpler as the lb < ub should have no side-effect and lightweight, which can be run twice (for scf.if and the for loop)

[joker-eph]

Please make it a pass: we're avoiding C++ unit-tests as much as possible.

I would implement the interface for scf.while and scf.for and have the pass just walk the IR, find all LoopLikeOpInterface, and call replaceWithZeroTripCheck.

[pzread]

Done. PTAL. I prefer to keep this change focusing on the interface definition so I only add a scf.while test with no transformation. The actual transformation and tests will be added in the follow-up change #80349

[pzread]

Kindly ping : )

[dcaballe]

Would it make sense to start with scf.if and generalize later as needed?

I think it's worth clarifying that loop rotation will only be possible for scf.while ops for now, that is, turning regular a "while-do" into a "do-while". That kind of rotation won't be possible at this level of abstraction for scf.for, hence the comment about redundant first iteration check. I think the rotation responsibility belongs to the interface itself.

If we want to make it extra configurable, we could separate the zero-trip-check method from the rotation method...

[MaheshRavishankar]

Late to the party here, but I think these transformations dont belong in the interface. Looking at the LoopLikeOpInterface it is actually doing some very heavy weight transformations that not all loop ops that implement this interface can handle. I tried to do this too in an earlier PR, and it was rightly pointed out that those kind of transformations dont go into the interface definition.

Instead of adding oto the interface definition, maybe these are just transformation methods that are written on the interface. My read of interfaces are that they are a wrapper around the core operation to have a unified way to interpret operations. Adding transformations of this kind to an interface seem odd to me.

I am not pointing out something really off in this specific PR, but it is contributing to an aspect of the interface that I wanted to flag.

[pzread]

Late to the party here, but I think these transformations dont belong in the interface. Looking at the LoopLikeOpInterface it is actually doing some very heavy weight transformations that not all loop ops that implement this interface can handle. I tried to do this too in an earlier PR, and it was rightly pointed out that those kind of transformations dont go into the interface definition.

Instead of adding oto the interface definition, maybe these are just transformation methods that are written on the interface. My read of interfaces are that they are a wrapper around the core operation to have a unified way to interpret operations. Adding transformations of this kind to an interface seem odd to me.

I am not pointing out something really off in this specific PR, but it is contributing to an aspect of the interface that I wanted to flag.

My initial thought is to create a common interface on loop ops so users can combine with utilities like moveLoopInvariantCode to hoist loop invariants and be op-agnostic:

llvm-project/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp

Lines 106 to 115 in 7880b2c

	size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
	return moveLoopInvariantCode(
	loopLike.getLoopRegions(),
	[&](Value value, Region *) {
	return loopLike.isDefinedOutsideOfLoop(value);
	},
	[&](Operation *op, Region *) {
	return isMemoryEffectFree(op) && isSpeculatable(op);
	},
	[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });

I also have the same question when writing this as an interface. Do you have an example where you think this transformation should live?

[pzread]

Late to the party here, but I think these transformations dont belong in the interface. Looking at the LoopLikeOpInterface it is actually doing some very heavy weight transformations that not all loop ops that implement this interface can handle. I tried to do this too in an earlier PR, and it was rightly pointed out that those kind of transformations dont go into the interface definition.
Instead of adding oto the interface definition, maybe these are just transformation methods that are written on the interface. My read of interfaces are that they are a wrapper around the core operation to have a unified way to interpret operations. Adding transformations of this kind to an interface seem odd to me.
I am not pointing out something really off in this specific PR, but it is contributing to an aspect of the interface that I wanted to flag.

My initial thought is to create a common interface on loop ops so users can combine with utilities like moveLoopInvariantCode to hoist loop invariants and be op-agnostic:

llvm-project/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp

Lines 106 to 115 in 7880b2c

	size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
	return moveLoopInvariantCode(
	loopLike.getLoopRegions(),
	[&](Value value, Region *) {
	return loopLike.isDefinedOutsideOfLoop(value);
	},
	[&](Operation *op, Region *) {
	return isMemoryEffectFree(op) && isSpeculatable(op);
	},
	[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });

I also have the same question when writing this as an interface. Do you have an example where you think this transformation should live?

IMO, I can probably write this as a transformation helper like this one:

llvm-project/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

Lines 990 to 993 in 8b0f47b

	FailureOr<scf::SCFTileAndFuseResult>
	mlir::scf::tileConsumerAndFuseProducersUsingSCF(
	RewriterBase &rewriter, TilingInterface consumer,
	const scf::SCFTileAndFuseOptions &options) {

[pzread]

As @MaheshRavishankar suggested, create #81050 to directly add a transformation for scf while loop.

[MaheshRavishankar]

Late to the party here, but I think these transformations dont belong in the interface. Looking at the LoopLikeOpInterface it is actually doing some very heavy weight transformations that not all loop ops that implement this interface can handle. I tried to do this too in an earlier PR, and it was rightly pointed out that those kind of transformations dont go into the interface definition.
Instead of adding oto the interface definition, maybe these are just transformation methods that are written on the interface. My read of interfaces are that they are a wrapper around the core operation to have a unified way to interpret operations. Adding transformations of this kind to an interface seem odd to me.
I am not pointing out something really off in this specific PR, but it is contributing to an aspect of the interface that I wanted to flag.

My initial thought is to create a common interface on loop ops so users can combine with utilities like moveLoopInvariantCode to hoist loop invariants and be op-agnostic:

llvm-project/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp

Lines 106 to 115 in 7880b2c

	size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
	return moveLoopInvariantCode(
	loopLike.getLoopRegions(),
	[&](Value value, Region *) {
	return loopLike.isDefinedOutsideOfLoop(value);
	},
	[&](Operation *op, Region *) {
	return isMemoryEffectFree(op) && isSpeculatable(op);
	},
	[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });

I also have the same question when writing this as an interface. Do you have an example where you think this transformation should live?

IMO, I can probably write this as a transformation helper like this one:

llvm-project/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

Lines 990 to 993 in 8b0f47b

	FailureOr<scf::SCFTileAndFuseResult>
	mlir::scf::tileConsumerAndFuseProducersUsingSCF(
	RewriterBase &rewriter, TilingInterface consumer,
	const scf::SCFTileAndFuseOptions &options) {

This is what I was thinking of. Thanks for addressing.