[mlir] LoopLikeOpInterface vs. RegionBranchOpInterface

LoopLikeOpInterface ostensibly exists to enable LICM. However, a limitation of the interface is that it only supports one loop body, which prohibits LICM on something like a while op that has two “loop” regions. I thought about extending the interface to support multiple loop bodies, but I wonder if RegionBranchOpInterface makes LoopLikeOpInterface obsolete?

Should the functionality of LoopLikeOpInterface be moved to RegionBranchOpInterface and LICM modified to work on it?

What types of API would you suggest moving to RegionBranchOpInterface? I would be somewhat worried (from an initial perspective) about conflating RegionBranchOpInterface with loops. I’m definitely +1 on having RegionBranchOpInterface be the anchor for general control flow things, but some things surrounding loops aren’t very general (e.g induction variables).

– River

These are:

  1. isDefinedOutsideOfLoop(Value) which returns whether a value is defined outside the loop (usually by checking ancestor). This can be redefined as isDefinedOutsideRegion(Value, Region *). Control-flow sink can make use of this too.
  2. getLoopBody. This is sort of replaced by getInvocationBounds. LICM will operate on regions that are known to execute more than once.
  3. moveOutOfLoop will be redefined as moveOutOfRegion (and maybe also add a moveIntoRegion).

This is what I mean by LoopLikeOpInterface exists mostly to enable LICM/hoisting/etc. RegionBranchOpInterface provides more in terms of induction variables than LoopLikeOpInterface with region successor operands and such.

This is what I mean by LoopLikeOpInterface exists mostly to enable LICM/hoisting/etc.

It’s not meant just for that. Take a look at the latest codebase. It’s useful to get a lower bound, step, etc. from a LoopLikeOpInterface where possible – which allows transformation and analysis utilities to function in a generic way on “countable” loops. Ops implementing RegionBranchOpInterface aren’t and need not be countable loops by structure or design.

I see I’m looking a codebase a week out of date…

The new interface functions added to LoopLikeOpInterface seem too specific to me. The interface doesn’t represent “loop-like” operations in general anymore. It represents specifically “countable” loops with a lower bound, upper bound, and step. And the description of the interface has not been updated to reflect this change.

To me, this strengthens to case to split out the part of the API in LookLikeOpInterface that is used to implement LICM (as per the interface’s description) versus the “countable loop op interface” functions that were added.

There are multiple ways to look at this. Not every single op needs to be supporting every single interface method (it can return llvm::None on an Optional) to be participating in the OpInterface. In this case, both scf.for and affine.for support returning step and lower bound for example, and I see this as pretty general and widely useful. It enables a whole bunch of very basic simplifications without having to depend on the SCF dialect or the affine dialect. Two examples:

<abc>.for %i = .... {
  // This will fold to true when %i owner op dyn_cast<LoopLikeOpInterface>
  // (..).getLowerBound()` gives you a value <= 0.
  = arith.cmpi ge %i, %c0 : ...
  // This will fold if loop like op interface's lowerBound + step is a multiple of 4.
  = arith.remi %i, %c4

There are many such simplifications/folding hooks that I can think of that could now generically use LoopLikeOpInterface – sure, scf::ParallelOp or AffineParallelOp don’t provide any value for these interface methods (they return None), but so would even scf.for or affine.for when that information isn’t available (for e.g., when you have a max lower bound for affine.for).

It represents specifically “countable” loops with a lower bound, upper bound, and step

Lower bound and stride aren’t tied to countable loops: a loop need not be countable for it to have a lower bound or step. So, any for op that exists in the codebase or an out-of-tree dialect is now able to now make use of LoopLikeOpInterface to benefit from simplifications like the ones above (post @ThomasRaoux’s revision). Furthermore, countable loops themselves cover a pretty large surface area of loops.

If we’re moving in the direction of building features on top of LoopLikeOpInterface instead of introducing new interfaces, I’m not necessarily against it. But it seems weird that while loops, which are ostensibly “loop-like” operations, cannot implement the interface effectively because it has more than one loop body. But with the addition of the new interface methods, it doesn’t make sense for LoopLikeOpInterface to have multiple loop bodies anymore.

That said, I would like to

  1. Change the name of LoopLikeOpInterface to be a little more specific. It doesn’t generalize to all loops anymore (primarily while loops are excluded and some other more exotic loop types). Maybe not “CountableLoopLikeOpInterface”. Suggest a name?
  2. Extend LICM to work on RegionBranchOpInterface as well.

I think both of those seem somewhat reasonable to me (without bike shedding yet on naming). If anything, just updating the docs for LoopLikeOpInterface to properly document which loops it is intended for would be useful. LICM using RegionBranchOpInterface would need to be careful about when it actually hoists though, given that the contract is not entirely the same with loops.

– River

I can’t think of anything better than LoopLikeOpInterface here without making it awkward. Improving the doc with enough information would be sufficient I feel. We wouldn’t want to change the name on the more general/common thing to accommodate the more exotic ones – instead, the outliers could perhaps be more appropriately named.

Coming back to the while op (if that was the concrete trigger for this discussion), what hurdles do you see in getting it to work with LICM? LICM for while loops would be good to have.

I will go the route of implementing LICM for RegionBranchOpInterface, which should cover while loops that implement the interface.

1 Like