[mlir][SCF] Fix region branch op interfaces for scf.forall and its terminator

[matthias-springer]

scf.forall does not completely implement the RegionBranchOpInterface: scf.forall.in_parallel does not implement the RegionBranchTerminatorOpInterface.

Incomplete interface implementation is a problem for transformations that try to understand the control flow by querying the RegionBranchOpInterface.

Detailed explanation of what is wrong with the current implementation.

• There is exactly one region branch point: "parent". in_parallel is not a region branch point because it does not implement the RegionBranchTerminatorOpInterface. (Clarified in [mlir][Interfaces][NFC] Document that RegionBranchTerminatorOpInterface is mandatory #174978.)
• ForallOp::getSuccessorRegions(parent) returns one region successors: the region of the scf.forall op.
• Since there is no region branch point in the region, there is no way to leave the region. This means: once you enter the region, you are stuck in it indefinitely. (It is unspecified what happens once you are in the region, but we can be sure that you cannot leave it.)

This commit adds the RegionBranchTerminatorOpInterface (via ReturnLike) to scf.forall.in_parallel to indicate the after leaving the region, the control flow returns to the parent. (Note: Only block terminators in directly nested regions can be region branch terminators, so in_parallel is the only possible op. I.e., parallel_insert_slice cannot be a region branch terminator.)

This commit also removes all successor operands / inputs from the implementation. The number of successor operands and successor inputs must match, but scf.forall.in_parallel has no operands. Therefore, the region must also have 0 successor inputs. Therefore, the scf.forall op must also have 0 successor operands.

This commit also adds a missing control flow edge from "parent" to "parent": in case of 0 threads, the region is not entered.

Depends on #174978.

[llvmbot]

@llvm/pr-subscribers-mlir-scf

Author: Matthias Springer (matthias-springer)

Changes

scf.forall does not completely implement the RegionBranchOpInterface. scf.forall.in_parallel does not implement the RegionBranchTerminatorOpInterface and it is unspecified which values are forwarded from the in_parallel op to the successor inputs.

Incomplete interface implementation is a problem for transformations that try to understand the data flow by querying the RegionBranchOpInterface, such as the patterns in #174094. Such transformations may be produce incorrect IR.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/SCF/IR/SCFOps.td (-1)
• (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (-17)
• (modified) mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir (+1-1)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 8bdf3e0b566ef..8c6c5e56e3645 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -326,7 +326,6 @@ def ForallOp : SCF_Op<"forall", [
            "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
-       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        DestinationStyleOpInterface,
        HasParallelRegion
      ]> {
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 8803a6d136f7a..0c6d539be0db4 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2013,23 +2013,6 @@ void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
               ForallOpReplaceConstantInductionVar>(context);
 }
 
-/// Given the region at `index`, or the parent operation if `index` is None,
-/// return the successor regions. These are the regions that may be selected
-/// during the flow of control. `operands` is a set of optional attributes that
-/// correspond to a constant value for each operand, or null if that operand is
-/// not a constant.
-void ForallOp::getSuccessorRegions(RegionBranchPoint point,
-                                   SmallVectorImpl<RegionSuccessor> &regions) {
-  // In accordance with the semantics of forall, its body is executed in
-  // parallel by multiple threads. We should not expect to branch back into
-  // the forall body after the region's execution is complete.
-  if (point.isParent())
-    regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
-  else
-    regions.push_back(
-        RegionSuccessor(getOperation(), getOperation()->getResults()));
-}
-
 //===----------------------------------------------------------------------===//
 // InParallelOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir b/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir
index 171a35fdeafb9..5c277f31dd02e 100644
--- a/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir
+++ b/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir
@@ -346,7 +346,7 @@ func.func @affine_loop_no_use_iv() {
 // CHECK-LABEL: test_tag: forall:
 // CHECK-NEXT: operand #0: live
 // CHECK-NEXT: region: #0:
-// CHECK-NEXT:   argument: #0: live
+// CHECK-NEXT:   argument: #0: not live
 
 func.func @forall_no_use_iv_has_side_effect_op(%idx1: index, %idx2: index) {
   scf.parallel (%i) = (%idx1) to (%idx2) step (%idx2) {

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

scf.forall does not completely implement the RegionBranchOpInterface. scf.forall.in_parallel does not implement the RegionBranchTerminatorOpInterface and it is unspecified which values are forwarded from the in_parallel op to the successor inputs.

Incomplete interface implementation is a problem for transformations that try to understand the data flow by querying the RegionBranchOpInterface, such as the patterns in #174094. Such transformations may be produce incorrect IR.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/SCF/IR/SCFOps.td (-1)
• (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (-17)
• (modified) mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir (+1-1)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 8bdf3e0b566ef..8c6c5e56e3645 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -326,7 +326,6 @@ def ForallOp : SCF_Op<"forall", [
            "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
-       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        DestinationStyleOpInterface,
        HasParallelRegion
      ]> {
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 8803a6d136f7a..0c6d539be0db4 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2013,23 +2013,6 @@ void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
               ForallOpReplaceConstantInductionVar>(context);
 }
 
-/// Given the region at `index`, or the parent operation if `index` is None,
-/// return the successor regions. These are the regions that may be selected
-/// during the flow of control. `operands` is a set of optional attributes that
-/// correspond to a constant value for each operand, or null if that operand is
-/// not a constant.
-void ForallOp::getSuccessorRegions(RegionBranchPoint point,
-                                   SmallVectorImpl<RegionSuccessor> &regions) {
-  // In accordance with the semantics of forall, its body is executed in
-  // parallel by multiple threads. We should not expect to branch back into
-  // the forall body after the region's execution is complete.
-  if (point.isParent())
-    regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
-  else
-    regions.push_back(
-        RegionSuccessor(getOperation(), getOperation()->getResults()));
-}
-
 //===----------------------------------------------------------------------===//
 // InParallelOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir b/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir
index 171a35fdeafb9..5c277f31dd02e 100644
--- a/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir
+++ b/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir
@@ -346,7 +346,7 @@ func.func @affine_loop_no_use_iv() {
 // CHECK-LABEL: test_tag: forall:
 // CHECK-NEXT: operand #0: live
 // CHECK-NEXT: region: #0:
-// CHECK-NEXT:   argument: #0: live
+// CHECK-NEXT:   argument: #0: not live
 
 func.func @forall_no_use_iv_has_side_effect_op(%idx1: index, %idx2: index) {
   scf.parallel (%i) = (%idx1) to (%idx2) step (%idx2) {

[matthias-springer]

Related PR: #147491

RegionBranchOpInterface has always been fishy on scf.forall. Why does this op implement the interface at all? What do we gain from it? @cxy-1993

[ftynse]

I am not sure about the "interface implementation completeness": do we require somewhere that terminators in operations implementing RegionBranchOpInterface must implement RegionBranchTerminatorOpInterface (if so, I'd perhaps add an assertion somehow)? In my reading, the current implementation says as much as we can say using this interface about control flow: it flows from before the op to the region and then to after the op, no values are being forwarded as-is anywhere (hence no need for RegionBranchTerminatorOpInterface). IMO, short of the requirement connecting the two interfaces, the logic using them should be robust to the terminator not implementing the interface since nothing guaranteed it. We can also have RegionBranchTerminatorOpInterface on scf.in_parallel and have it return an empty range of operands as forwarded to results.

I'm also not opposed to the change, but I think we need to nail it down this time.

[joker-eph]

Outside of the requirement (or not) for RegionBranchTerminatorOpInterface, the support for RegionBranchOpInterface for the parallel loop has always been a bit weird, because the op does more than "control flow" with a special logic to "merge" results.

[matthias-springer]

do we require somewhere that terminators in operations implementing RegionBranchOpInterface must implement RegionBranchTerminatorOpInterface

Not all terminators are necessarily RegionBranchTerminatorOpInterface. You could also have a BranchOpInterface terminator. (If there's just a single block in the region and the terminator is a branching terminator for unstructured control flow, we have an infinite loop. You could potentially also have an op which implements both interfaces.) Long story short -- we cannot simplify verify that every terminator implements the RegionBranchTerminatorOpInterface.

If you have a transformation that analyzes RegionBranchTerminatorOpInterface and RegionBranchOpInterface, we must be able to query both interfaces. Otherwise, the transformation may "miss" some control flow edges. That's the case in the patterns that I'm adding in #174094. (For the moment, this is not a problem because those patterns are enabled selectively via a whitelist. But eventually, I'd like to enable them globally.)

I see two options:

1. Require region branch terminators to implement the RegionBranchTerminatorOpInterface. This can be checked in the verifier by checking all region branch points. (Similar to how we check that scf.for has an scf.yield terminator.)
2. Document that implementing the RegionBranchTerminatorOpInterface is optional. Add a helper function RegionBranchOpInterface::modelsCompleteDataflow (which performs the same check as in Option 1), so transformations can opt out.

it flows from before the op to the region and then to after the op, no values are being forwarded as-is anywhere (hence no need for RegionBranchTerminatorOpInterface).

For an op that does forward any values, we could still require that the RegionBranchTerminatorOpInterface interface is implemented, and getMutableSuccessorOperands would return an an empty list.

the parallel loop has always been a bit weird, because the op does more than "control flow" with a special logic to "merge" results

Yes, that the problem that I wanted to address first. It's unclear to me why the op even implements the interface. If the answer is "we need it" for some reason, then Option 1 above is not feasible with the current design. @cxy-1993 please take a look when you get a chance.

[matthias-springer]

I found additional code that misbehaves when terminators model region control flow but do not implement the RegionBranchTerminatorOpInterface:

/// Return `true` if region `r` is reachable from region `begin` according to
/// the RegionBranchOpInterface (by taking a branch).
static bool isRegionReachable(Region *begin, Region *r);

The above function may return "false". To be fair, I wrote that function (and the places where it is used) some time ago.

[ftynse]

Long story short -- we cannot simplify verify that every terminator implements the RegionBranchTerminatorOpInterface.

I never said that all terminators should implement it, only

terminators in operations implementing RegionBranchOpInterface

which is exactly

1. Require region branch terminators to implement the RegionBranchTerminatorOpInterface. This can be checked in the verifier by checking all region branch points.

Currently we don't seem to require it. I have written such transformations and in a conservative way: dyn_cast the terminator to the interface and assume the worst case (any terminator operand may flow to any block argument / parent result) if not implemented. I don't think not implementing RegionBranchTerminatorOpInterface makes us miss any control flow edge as these are not specified by this interface. They are specified by the RegionBranchOpInterface::getSuccessorRegion, the terminator interface only saying which operands are forwarded along those edges. I see a getSuccessorRegions function on the terminator interface, but its default implementation defers to the region interface on the parent and we don't have non-default implementation upstream. IMO it should be turned into a non-method and just an extra declaration on the interface.

[matthias-springer]

Long story short -- we cannot simplify verify that every terminator implements the RegionBranchTerminatorOpInterface.

I never said that all terminators should implement it, only

terminators in operations implementing RegionBranchOpInterface

I think we're still talking about different things. There are legitimate cases where even a terminator inside of a RegionBranchOpInterface cannot implement the RegionBranchTerminatorOpInterface (case 3 below) -- in the presence of unstructured control flow, a region may have multiple blocks. There are 3 cases:

1. A block terminator is conceptually a region branch terminator and implements the RegionBranchTerminatorOpInterface.
2. A block terminator is conceptually a region branch terminator but does not implement the RegionBranchTerminatorOpInterface. (E.g., because we say that it's optional to implement that interface.)
3. A block terminator is conceptually not a region branch terminator and does not implement the RegionBranchTerminatorOpInterface. (E.g., because it a unstructured control flow branching op.)

I'd like to see if we can rule out case 2 to simplify the design.

Here's a problem with case 2 that I hinted to above: While the RegionBranchOpInterface itself captures the entire region control flow, the current API does not allow you to query the region successor in case 2. RegionBranchOpInterface::getSuccessorRegions takes a RegionBranchPoint parameter, but to construct a RegionBranchPoint you need a RegionBranchTerminatorOpInterface op. You cannot construct a RegionBranchPoint from just an Operation *.

If we relax this API, I agree that we could probably safely allow case 2 and even support conservative value propagation as you described above. But that means that we also have to check/update all existing analyses/transformations that query control flow from the RegionBranchOpInterface. They all implicitly assume that case 2 does not exist (due to the API limitation). Given that all analyses/transformations already operate under that assumption, is there any harm in forbidding case 2? I'm going to update this PR accordingly so we can see what that design looks like in detail.

edit: Even the RegionBranchOpInterface::getRegionSuccessor(RegionBranchPoint, SmallVector<RegionSuccessor> &) interface method cannot model control flow for case 2.

I see a getSuccessorRegions function on the terminator interface, but its default implementation defers to the region interface on the parent and we don't have non-default implementation upstream. IMO it should be turned into a non-method and just an extra declaration on the interface.

I just looked into this. It's a bit different from RegionBranchOpInterface::getSuccessorRegions: it takes an additional ArrayRef<Attribute> operands parameter. This RegionBranchTerminatorOpInterface::getSuccessorRegions interface method is similar to RegionBranchOpInterface::getEntrySuccessorRegions, and while these interface methods may not be used much, at least the design is consistent.

[matthias-springer]

The discussion here sidetracked a bit. I created a new PR to make RegionBranchTerminatorOpInterface mandatory. Let's continue the discussion there and focus only on scf.forall on this PR.

[matthias-springer]

After clarifying the semantics of RegionBranchOpInterface #174978, I believe it is safe to implement the interface on scf.forall, as along as scf.forall.in_parallel is also implemented on the terminator. Details in the commit message.

Basically what @ftynse said above:

We can also have RegionBranchTerminatorOpInterface on scf.in_parallel and have it return an empty range of operands as forwarded to results.

As @joker-eph mentioned above, the main issue seems to be with the special merge logic of the op. This commit drops that part from the interface implementation.

[matthias-springer]

Note: This is a shortcut for RegionBranchTerminatorOpInterface.

[matthias-springer]

note: This now falls back to the default implementation of getEntrySuccessorOperands: an empty operand range.

[matthias-springer]

I am going to add a test once #174094 has been merged. (With #174094, there's no good way to test this.)

I added a test case.

[joker-eph]

Can we test the 0 thread case?

[matthias-springer]

The getSuccessorRegions implementation doesn't actually check the number of iterations of the loop. I.e., doesn't matter if 0 or more iterations. It just (conservatively) populates both region successors. (This could be improved across the entire SCF dialect. E.g., the implementation for scf.for also doesn't look at the number of iterations...)

I updated the test case to use an SSA value for the number of threads. In that case, we must populate both region successors because we do not statically know the branching behavior. Is that good enough? (That test case used to produce incorrect results before.)

[matthias-springer]

This is ready to be merged. Could you take another look?

[llvm-ci]

LLVM Buildbot has detected a new failure on builder ppc64le-mlir-rhel-clang running on ppc64le-mlir-rhel-test while building mlir at step 6 "test-build-check-mlir-build-only-check-mlir".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/129/builds/36454

Here is the relevant piece of the build log for the reference

Step 6 (test-build-check-mlir-build-only-check-mlir) failure: 1200 seconds without output running [b'ninja', b'check-mlir'], attempting to kill
...
PASS: MLIR :: mlir-translate/split-markers.mlir (3801 of 3812)
PASS: MLIR :: Pass/crash-recovery.mlir (3802 of 3812)
PASS: MLIR :: Pass/remarks.mlir (3803 of 3812)
PASS: MLIR :: mlir-reduce/dce-test.mlir (3804 of 3812)
PASS: MLIR :: mlir-opt/split-markers.mlir (3805 of 3812)
PASS: MLIR :: Transforms/location-snapshot.mlir (3806 of 3812)
PASS: MLIR :: mlir-tblgen/op-error.td (3807 of 3812)
PASS: MLIR :: mlir-runner/utils.mlir (3808 of 3812)
PASS: MLIR :: Pass/pipeline-parsing.mlir (3809 of 3812)
PASS: MLIR :: Pass/pipeline-options-parsing.mlir (3810 of 3812)
command timed out: 1200 seconds without output running [b'ninja', b'check-mlir'], attempting to kill
process killed by signal 9
program finished with exit code -1
elapsedTime=1982.809037