[MLIR] Introduce scf.execute_region op

[bondhugula]

Summary

Introduce the execute_region op that is able to hold a region which it
executes exactly once. The op encapsulates a CFG within itself while
isolating it from the surrounding control flow. Proposal discussed here:
https://llvm.discourse.group/t/introduce-std-inlined-call-op-proposal/282

execute_region enables one to inline a function without lowering out all
other higher level control flow constructs (affine.for/if, scf.for/if)
to the flat list of blocks / CFG form. It thus allows the benefit of
transforms on higher level control flow ops available in the presence of
the inlined calls. The inlined calls continue to benefit from
propagation of SSA values across their top boundary. Functions won’t
have to remain outlined until later than desired. Abstractions like
affine execute_regions, lambdas with implicit captures could be lowered
to this without first lowering out structured loops/ifs or outlining.
But two potential early use cases are of: (1) an early inliner (which
can inline functions by introducing execute_region ops), (2) lowering of
an affine.execute_region, which cleanly maps to an scf.execute_region
when going from the affine dialect to the scf dialect.

Signed-off-by: Uday Bondhugula <[e-mail]>

[bondhugula]

Unintended duplicate

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

def ConditionOp : SCF_Op<"condition",
 
   let assemblyFormat =
       [{ `(` $condition `)` attr-dict ($args^ `:` type($args))? }];
 
   // Override the default verifier, everything is checked by traits.
   let verifier = ?;
 }
 
+//===----------------------------------------------------------------------===//
+// ExecuteRegionOp
+//===----------------------------------------------------------------------===//
+
+def ExecuteRegionOp : SCF_Op<"execute_region"> {
+  let summary = "operation that executes its region exactly once";
+  let description = [{
+    The `execute_region` operation executes the region held exactly once. The op
+    cannot have any operands, nor does its region have any arguments. All SSA
+    values that dominate the op can be accessed inside. The op's region can have
+    multiple blocks and the blocks can have terminators the same way as FuncOp.
+    The values returned from this op's region define the op's results.  The op
+    primarily provides control flow encapsulation and isolation from a parent
+    op's control flow restrictions if any; for example, it allows representation
+    of inlined calls in the inside of structured control flow ops with
+    restrictions like affine.for/if, scf.for/if ops, and thus the optimization
+    of IR in such a mixed form.
+
+    Example:
+
+    ```mlir
+    scf.for %i = 0 to 128 step %c1 {
+      %y = scf.execute_region -> i32 {
+        %x = load %A[%i] : memref<128xi32>
+        scf.yield %x : i32
+      }
+    }
+
+    affine.for %i = 0 to 100 {
+      "foo"() : () -> ()
+      %v = scf.execute_region -> i64 {
+        cond_br %cond, ^bb1, ^bb2
+
+      ^bb1:
+        %c1 = constant 1 : i64
+        br ^bb3(%c1 : i64)
+
+      ^bb2:
+        %c2 = constant 2 : i64
+        br ^bb3(%c2 : i64)
+
+      ^bb3(%x : i64):
+        scf.yield %x : i64
+      }
+      "bar"(%v) : (i64) -> ()
+    }
+    ```
+  }];
+
+  let results = (outs Variadic<AnyType>);
+
+  let regions = (region AnyRegion:$region);
+
+  // TODO: If the parent is a func like op (which would be the case if all other
+  // ops are from the std dialect), the inliner logic could be readily used to
+  // inline.
+  let hasCanonicalizer = 0;
+
+  // TODO: can fold if it returns a constant.
+  // TODO: Single block execute_region ops can be readily inlined irrespective
+  // of which op is a parent. Add a fold for this.
+  let hasFolder = 0;
+}
+
 def ForOp : SCF_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
let extraClassDeclaration = [{
     ConditionOp getConditionOp();
     Block::BlockArgListType getAfterArguments();
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def YieldOp : SCF_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
-                               ParentOneOf<["IfOp, ForOp", "ParallelOp",
-                                            "WhileOp"]>]> {
+                               ParentOneOf<["ExecuteRegionOp, ForOp",
+                                            "IfOp, ParallelOp, WhileOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "scf.yield" yields an SSA value from the SCF dialect op region and
     terminates the regions. The semantics of how the values are yielded is
     defined by the parent operation.
     If "scf.yield" has any operands, the operands must match the parent
     operation's results.
     If the parent operation defines no values, then the "scf.yield" may be

mlir/lib/Dialect/SCF/SCF.cpp

 }
 
 /// Default callback for IfOp builders. Inserts a yield without arguments.
 void mlir::scf::buildTerminatedBody(OpBuilder &builder, Location loc) {
   builder.create<scf::YieldOp>(loc);
 }
 
 //===----------------------------------------------------------------------===//
+// ExecuteRegionOp
+//===----------------------------------------------------------------------===//
+
+///
+/// (ssa-id `=`)? `execute_region` `->` function-result-type `{`
+///    block+
+/// `}`
+///
+/// Example:
+///   scf.execute_region -> i32 {
+///     %idx = load %rI[%i] : memref<128xi32>
+///     return %idx : i32
+///   }
+///
+static ParseResult parseExecuteRegionOp(OpAsmParser &parser,
+                                        OperationState &result) {
+  if (parser.parseOptionalArrowTypeList(result.types))
+    return failure();
+
+  // Introduce the body region and parse it.
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}) ||
+      parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  return success();
+}
+
+static void print(OpAsmPrinter &p, ExecuteRegionOp op) {
+  p << ExecuteRegionOp::getOperationName();
+  if (op.getNumResults() > 0)
+    p << " -> " << op.getResultTypes();
+
+  p.printRegion(op.region(),
+                /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/true);
+
+  p.printOptionalAttrDict(op->getAttrs());
+}
+
+static LogicalResult verify(ExecuteRegionOp op) {
+  if (op.region().empty())
+    return op.emitOpError("region needs to have at least one block");
+  if (op.region().front().getNumArguments() > 0)
+    return op.emitOpError("region cannot have any arguments");
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
                   Value ub, Value step, ValueRange iterArgs,
                   BodyBuilderFn bodyBuilder) {
   result.addOperands({lb, ub, step});
   result.addOperands(iterArgs);
static ParseResult parseForOp(OpAsmParser &parser, OperationState &result) {
   regionArgs.push_back(inductionVariable);
 
   if (succeeded(parser.parseOptionalKeyword("iter_args"))) {
     // Parse assignment list and results type list.
     if (parser.parseAssignmentList(regionArgs, operands) ||
         parser.parseArrowTypeList(result.types))
       return failure();
     // Resolve input operands.
-    for (auto operand_type : llvm::zip(operands, result.types))
-      if (parser.resolveOperand(std::get<0>(operand_type),
-                                std::get<1>(operand_type), result.operands))
+    for (auto operandType : llvm::zip(operands, result.types))
+      if (parser.resolveOperand(std::get<0>(operandType),
+                                std::get<1>(operandType), result.operands))
         return failure();
   }
   // Induction variable.
   argTypes.push_back(indexType);
   // Loop carried variables
   argTypes.append(result.types.begin(), result.types.end());
   // Parse the body region.
   Region *body = result.addRegion();

 
 Region &ForOp::getLoopBody() { return region(); }
 
 bool ForOp::isDefinedOutsideOfLoop(Value value) {
   return !region().isAncestor(value.getParentRegion());
 }
 
 LogicalResult ForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
-  for (auto op : ops)
+  for (auto *op : ops)
     op->moveBefore(*this);
   return success();
 }
 
 ForOp mlir::scf::getForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg)
     return ForOp();

 
 Region &ParallelOp::getLoopBody() { return region(); }
 
 bool ParallelOp::isDefinedOutsideOfLoop(Value value) {
   return !region().isAncestor(value.getParentRegion());
 }
 
 LogicalResult ParallelOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
-  for (auto op : ops)
+  for (auto *op : ops)
     op->moveBefore(*this);
   return success();
 }
 
 ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg)
     return ParallelOp();

mlir/test/Dialect/SCF/canonicalize.mlir

func @combineIfs4(%arg0 : i1, %arg2: i64) {
   }
   return
 }
 
 // CHECK-NEXT:     scf.if %arg0 {
 // CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
 // CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
 // CHECK-NEXT:     }
+
+// -----
+
+// CHECK-LABEL: func @propagate_into_execute_region
+func @propagate_into_execute_region() {
+  %cond = constant 0 : i1
+  affine.for %i = 0 to 100 {
+    "test.foo"() : () -> ()
+    %v = scf.execute_region -> i64 {
+      cond_br %cond, ^bb1, ^bb2
+
+    ^bb1:
+      %c1 = constant 1 : i64
+      br ^bb3(%c1 : i64)
+
+    ^bb2:
+      %c2 = constant 2 : i64
+      br ^bb3(%c2 : i64)
+
+    ^bb3(%x : i64):
+      scf.yield %x : i64
+    }
+    "test.bar"(%v) : (i64) -> ()
+    // CHECK:      %[[C2:.*]] = constant 2 : i64
+    // CHECK:        scf.execute_region -> i64 {
+    // CHECK-NEXT:     scf.yield %[[C2]] : i64
+    // CHECK-NEXT:   }
+  }
+  return
+}

mlir/test/Dialect/SCF/invalid.mlir

func @parallel_invalid_yield(
   }
   return
 }
 
 // -----
 
 func @yield_invalid_parent_op() {
   "my.op"() ({
-   // expected-error@+1 {{'scf.yield' op expects parent op to be one of 'scf.if, scf.for, scf.parallel, scf.while'}}
+   // expected-error@+1 {{'scf.yield' op expects parent op to be one of 'scf.execute_region, scf.for, scf.if, scf.parallel, scf.while'}}
    scf.yield
   }) : () -> ()
   return
 }
 
 // -----
 
 func @while_parser_type_mismatch() {
func @while_bad_terminator() {
   // expected-error@+1 {{expects the 'after' region to terminate with 'scf.yield'}}
   scf.while : () -> () {
     scf.condition(%true)
   } do {
     // expected-note@+1 {{terminator here}}
     "some.other_terminator"() : () -> ()
   }
 }
+
+// -----
+
+func @execute_region() {
+  // expected-error @+1 {{region cannot have any arguments}}
+  "scf.execute_region"() ({
+  ^bb0(%i : i32):
+    scf.yield
+  }) : () -> ()
+  return
+}

mlir/test/Dialect/SCF/ops.mlir

scf.while : () -> () {
     scf.condition(%true)
   // CHECK: } do {
   } do {
     // CHECK: scf.yield
     scf.yield
   }
   return
 }
+
+// CHECK-LABEL: func @execute_region
+func @execute_region() -> i64 {
+  // CHECK:      scf.execute_region -> i64 {
+  // CHECK-NEXT:   constant
+  // CHECK-NEXT:   scf.yield
+  // CHECK-NEXT: }
+  %res = scf.execute_region -> i64 {
+    %c1 = constant 1 : i64
+    scf.yield %c1 : i64
+  }
+
+  // CHECK:       scf.execute_region {
+  // CHECK-NEXT:    br ^bb1
+  // CHECK-NEXT:  ^bb1:
+  // CHECK-NEXT:    scf.yield
+  // CHECK-NEXT:  }
+  "scf.execute_region"() ({
+  ^bb0:
+    br ^bb1
+  ^bb1:
+    scf.yield
+  }) : () -> ()
+  return %res : i64
+}

mlir/utils/vim/syntax/mlir.vim

 syn match mlirOps /\<affine\.dma_start\>/
 syn match mlirOps /\<affine\.dma_wait\>/
 syn match mlirOps /\<affine\.for\>/
 syn match mlirOps /\<affine\.if\>/
 syn match mlirOps /\<affine\.load\>/
 syn match mlirOps /\<affine\.parallel\>/
 syn match mlirOps /\<affine\.prefetch\>/
 syn match mlirOps /\<affine\.store\>/
-syn match mlirOps /\<loop\.for\>/
-syn match mlirOps /\<loop\.if\>/
+syn match mlirOps /\<scf\.execute_region\>/
+syn match mlirOps /\<scf\.for\>/
+syn match mlirOps /\<scf\.if\>/
+syn match mlirOps /\<scf\.yield\>/
 
 " TODO: dialect name prefixed ops (llvm or std).
 
 " Keywords.
 syn keyword mlirKeyword
       \ affine_map
       \ affine_set
       \ dense