@llvm/pr-subscribers-mlir

Author: Fabian Mora (fabianmcg)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 2d15544e871b3..23c987dd777e2 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -633,6 +633,7 @@ def ForallOp : SCF_Op<"forall", [
 def InParallelOp : SCF_Op<"forall.in_parallel", [
        Pure,
        Terminator,
+       ReturnLike,
        DeclareOpInterfaceMethods<ParallelCombiningOpInterface>,
        HasParent<"ForallOp">,
       ] # GraphRegionNoTerminator.traits> {

@llvm/pr-subscribers-mlir-scf

Author: Fabian Mora (fabianmcg)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 2d15544e871b3..23c987dd777e2 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -633,6 +633,7 @@ def ForallOp : SCF_Op<"forall", [
 def InParallelOp : SCF_Op<"forall.in_parallel", [
        Pure,
        Terminator,
+       ReturnLike,
        DeclareOpInterfaceMethods<ParallelCombiningOpInterface>,
        HasParent<"ForallOp">,
       ] # GraphRegionNoTerminator.traits> {

Could you elaborate more on the reasoning here?

Sure.

Right now, scf.forall control-flow semantics are broken. scf.forall declares to be RegionBranchOpInterface but scf.in_parallel is not a RegionBranchTerminatorOpInterface, and this breaks things like dataflow. Furthermore, I'd argue this breaks the implicit contract of region control-flow, where each RegionBranchOpInterface should be matched by a RegionBranchTerminatorOpInterface.

There were 2 bad options to start fixing the control-flow semantics in this case:

1. Remove RegionBranchOpInterface from scf.forall. This sounds like a breaking change that might break semantics for everyone. Further, in the future we would re-introduce it once the whole semantics issue is fixed.
2. Start fixing scf.in_parallel, and try to break as few things as possible during the process to make it comply with RegionBranchTerminatorOpInterface.

I opted for 2. This patch fixes scf.in_parallel semantics for the buffer case (as the buffer case has no arguments) but not the tensor case, because as you mentioned this op doesn't forward its operands to the parent op.

The proper fix to this whole situation, (IMO) is to add a second region to scf.forall and do the in parallel computations there for tensors, then RegionBranchTerminatorOpInterface can always apply. But that's a large breaking change specially for testing, so I'm trying to figure out ways to keep the changes small, and this was my choice among a set of bad choices (I'm open to hear alternatives) as a first step, as it fixes the easy half of the problem (the buffer case).

Now, outside of this patch I would argue that control-flow semantics (outside blocks) in MLIR are somewhat broken, as this is not the first time I encounter a situation like this. But that's a larger conversation and out of scope for this patch.

The proper fix to this whole situation, (IMO) is to add a second region to scf.forall and do the in parallel computations there for tensors, then RegionBranchTerminatorOpInterface can always apply. But that's a large breaking change specially for testing, so I'm trying to figure out ways to keep the changes small, and this was my choice among a set of bad choices (I'm open to hear alternatives) as a first step, as it fixes the easy half of the problem (the buffer case).

This sounds like a good idea, but you'll have to forward values from one region to another. For testing, we can follow what I suggested on another thread and have a custom version of mlir-opt that accepts the previous syntax/structure and prints the new structure. Doing it with commented-out FileCheck lines is a bit tricky though.

Now, outside of this patch I would argue that control-flow semantics (outside blocks) in MLIR are somewhat broken, as this is not the first time I encounter a situation like this. But that's a larger conversation and out of scope for this patch.

This is something that can be raised to the forums. I know the modeling is quite limited, but I don't know whether is broken.

While the op does not have explicit operands, the op is actually implicitly capturing SSA value, which counts as operands to me.
However it does not just forward its operands to the parent, since it relies on its body to form the actual values to yield (in a non-trivial manner).

I don't quite see what's broken in terms of control-flow here? The fact that it is a parallel fan-out?
The RegionBanchOpInterface is limited in its ability to track the actual data-flow here but that's another issue.

The proper fix to this whole situation, (IMO) is to add a second region to scf.forall and do the in parallel computations there for tensors, then RegionBranchTerminatorOpInterface can always apply.

That wouldn't help alone: you would need to materialize the reduction in the same way that XLA/HLO does with its reduce: by taking two values, writing the combiner, and yield the new value (instead of the "magic" parallel_insert_slice model existing now)

I suggested on another thread and have a custom version of mlir-opt that accepts the previous syntax/structure and prints the new structure. Doing it with commented-out FileCheck lines is a bit tricky though.

Yes, that's why I mentioned clang's Rewriter class, it would allow to keep the comments. However, since then I think it's easier to parse the comments into ops.

I don't quite see what's broken in terms of control-flow here?

I think the semantics of the implementations are broken here, eg. forall having the interface but in_parallel not having the matching interface. In the same way, I don't think SSACFG region semantics are necessarily broken, but the implementation with RegionBranch* likely is.

However it does not just forward its operands to the parent, since it relies on its body to form the actual values to yield (in a non-trivial manner).

That's another thing I considered, extending the interface and other data structures to also work with ranges of ops. However, there were two things I learned while trying to do that approach:

1. MLIR relies heavily in control-flow operands to be actual op operands, eg. dataflow relies on getting the operand number.
2. Mutation as it's currently prescribed by the interface is impossible for non values. This is one of the things I think it's broken in the interface, IMO mutation should be handled by the interface and we should remove the notion of operands from the interface.

Consequently, the proper fix there it's an even bigger breaking change. I'd argue this needs to happen anyways, but I don't want to propose that until we know how generic control-flow looks in MLIR, as the new system could potentially replace all control-flow and fix things in one go.

That wouldn't help alone: you would need to materialize the reduction in the same way that XLA/HLO does with its reduce: by taking two values, writing the combiner, and yield the new value (instead of the "magic" parallel_insert_slice model existing now)

Yeah, a change in parallel_insert_slice is also needed to yield a value. However, I see it as a potentially easier change, with the correct handling happening in the parallel_result region of forall.