[Mlir-commits] [mlir] e8bfec2 - [mlir] [scf] Add RegionBranchOpInterface to scf.forall and scf.parallel op

[Matthias Springer]

Author: donald chen
Date: 2023-05-24T12:13:20+02:00
New Revision: e8bfec26a6dace4f7d8d21e74d1e8d05f9714a63

URL: https://github.com/llvm/llvm-project/commit/e8bfec26a6dace4f7d8d21e74d1e8d05f9714a63
DIFF: https://github.com/llvm/llvm-project/commit/e8bfec26a6dace4f7d8d21e74d1e8d05f9714a63.diff

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

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

Differential Revision: https://reviews.llvm.org/D151287

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
    mlir/lib/Dialect/SCF/IR/SCF.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 4173245931478..1eed4cec1e420 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -363,6 +363,7 @@ def ForallOp : SCF_Op<"forall", [
        AutomaticAllocationScope,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
+       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      ]> {
   let summary = "evaluate a block multiple times in parallel";
   let description = [{
@@ -814,6 +815,7 @@ def ParallelOp : SCF_Op<"parallel",
      AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveMemoryEffects,
+     DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      SingleBlockImplicitTerminator<"scf::YieldOp">]> {
   let summary = "parallel for operation";
   let description = [{

diff  --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 7048d1348ca89..acfeb0f1e205d 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1682,6 +1682,25 @@ void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
               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
 //===----------------------------------------------------------------------===//
@@ -2976,6 +2995,26 @@ void ParallelOp::getCanonicalizationPatterns(RewritePatternSet &results,
           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
 //===----------------------------------------------------------------------===//