[mlir] [scf] Add RegionBranchOpInterface to scf.forall and scf.parallel op

[cxy-1993]

Summary

Add RegionBranchOpIntefface to scf.forall and scf.parallel op to make analysis trace through subregions.

[cxy-1993]

As discussed on https://discourse.llvm.org/t/why-scf-forall-op-doesnt-have-regionbranchop-interface/70789/4, these two operation has special terminators.

Operation forall's terminator is InParallelOp without input operands.
Operation parallel's termnator is yield op without input operands.

So Interface "RegionBranchTerminatorOpInterface" don't need any more, return empty operands will keep “value propagation” invalid.

[cxy-1993]

Update coding style.

[springerm]

scf.parallel can have an scf.reduce terminator.
E.g.:

%init = arith.constant 0.0 : f32
scf.parallel (%iv) = (%lb) to (%ub) step (%step) init (%init) -> f32 {
  %elem_to_reduce = load %buffer[%iv] : memref<100xf32>
  scf.reduce(%elem_to_reduce) : f32 {
    ^bb0(%lhs : f32, %rhs: f32):
      %res = arith.addf %lhs, %rhs : f32
      scf.reduce.return %res : f32
  }
}

[cxy-1993]

In D151287#4367113, @springerm wrote:

scf.parallel can have an scf.reduce terminator.
E.g.:

%init = arith.constant 0.0 : f32
scf.parallel (%iv) = (%lb) to (%ub) step (%step) init (%init) -> f32 {
  %elem_to_reduce = load %buffer[%iv] : memref<100xf32>
  scf.reduce(%elem_to_reduce) : f32 {
    ^bb0(%lhs : f32, %rhs: f32):
      %res = arith.addf %lhs, %rhs : f32
      scf.reduce.return %res : f32
  }
}

scf.parallel has a implicit scf.yield terminator(https://github.com/llvm/llvm-project/blob/cf1ef4161006e8119761b3a137423c23436bcf33/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td#L812).
And scf.reduce don't have terminator trait(https://github.com/llvm/llvm-project/blob/cf1ef4161006e8119761b3a137423c23436bcf33/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td#L900).

The actual result of scf.parallel is generated by scf.reduce, but terminator of scf.parallel yield empty. So we can prevent value propagation without registe RegionBranchTerminatorOpInterface to terminator.

[cxy-1993]

Could you please help me push this patch to master @springerm

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

 // ForallOp
 //===----------------------------------------------------------------------===//
 
 def ForallOp : SCF_Op<"forall", [
        AttrSizedOperandSegments,
        AutomaticAllocationScope,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
+       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      ]> {
   let summary = "evaluate a block multiple times in parallel";
   let description = [{
     `scf.forall` is a target-independent multi-dimensional parallel
     region application operation. It has exactly one block that represents the
     parallel body and it takes index operands that specify lower bounds, upper
     bounds and steps.
 

 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 def ParallelOp : SCF_Op<"parallel",
     [AutomaticAllocationScope,
      AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveMemoryEffects,
+     DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      SingleBlockImplicitTerminator<"scf::YieldOp">]> {
   let summary = "parallel for operation";
   let description = [{
     The "scf.parallel" operation represents a loop nest taking 4 groups of SSA
     values as operands that represent the lower bounds, upper bounds, steps and
     initial values, respectively. The operation defines a variadic number of
     SSA values for its induction variables. It has one region capturing the
     loop body. The induction variables are represented as an argument of this

mlir/lib/Dialect/SCF/IR/SCF.cpp

 
 void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
   results.add<DimOfForallOp, FoldTensorCastOfOutputIntoForallOp,
               ForallOpControlOperandsFolder,
               ForallOpSingleOrZeroIterationDimsFolder>(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(std::optional<unsigned> index,
+                                ArrayRef<Attribute> operands,
+                                SmallVectorImpl<RegionSuccessor> &regions) {
+  // If the predecessor is ForallOp, branch into the body with empty arguments.
+  if (!index) {
+    regions.push_back(RegionSuccessor(&getRegion()));
+    return;
+  }
+
+  // Otherwise, the loop should branch back to the parent operation.
+  assert(*index == 0 && "expected loop region");
+  regions.push_back(RegionSuccessor());
+}
+
 //===----------------------------------------------------------------------===//
 // InParallelOp
 //===----------------------------------------------------------------------===//
 
 // Build a InParallelOp with mixed static and dynamic entries.
 void InParallelOp::build(OpBuilder &b, OperationState &result) {
   OpBuilder::InsertionGuard g(b);
   Region *bodyRegion = result.addRegion();

 
 void ParallelOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results
       .add<ParallelOpSingleOrZeroIterationDimsFolder, MergeNestedParallelLoops>(
           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 ParallelOp::getSuccessorRegions(std::optional<unsigned> index,
+                                ArrayRef<Attribute> operands,
+                                SmallVectorImpl<RegionSuccessor> &regions) {
+  // If the predecessor is ParallelOp, branch into the body with empty
+  // arguments.
+  if (!index) {
+    regions.push_back(RegionSuccessor(&getRegion()));
+    return;
+  }
+
+  assert(*index == 0 && "expected loop region");
+  // Otherwise, the loop should branch back to the parent operation.
+  regions.push_back(RegionSuccessor());
+}
+
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
 void ReduceOp::build(
     OpBuilder &builder, OperationState &result, Value operand,
     function_ref<void(OpBuilder &, Location, Value, Value)> bodyBuilderFn) {
   auto type = operand.getType();