[Mlir-commits] [mlir] [MLIR] Introduce support for early exits (PR #166688)

[Matthias Springer]

================
@@ -146,6 +147,182 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
 
   let hasVerifier = 1;
 }
+//===----------------------------------------------------------------------===//
+// LoopOp
+//===----------------------------------------------------------------------===//
+
+def LoopOp : SCF_Op<"loop",[
+    AutomaticAllocationScope,
+    OpAsmOpInterface,
+    RecursiveMemoryEffects,
+    PropagateControlFlowBreak,
+    TokenProducerTrait,
+    DeclareOpInterfaceMethods<RegionBranchOpInterface,
+        ["getEntrySuccessorOperands", "getSuccessorInputs"]>,
+    SingleBlockImplicitTerminator<"op_impl::LoopOpImplicitTerminatorType">,
+    HasBreakingControlFlowOpInterface,
+    HasNestedTerminator<["ContinueOp", "BreakOp"]>
+  ]> {
+  let summary = "Loop until a break operation";
+  let description = [{
+    The `loop` operation represents an unstructured infinite loop that executes
+    until a `break` is reached.
+
+    The loop consists of (1) a set of loop-carried values which are initialized by
+    `initValues` and updated by each iteration of the loop, and
+    (2) a region which represents the loop body.
+
+    The loop will execute the body of the loop until a `break` is dynamically executed.
+
+    Each control path of the loop must be terminated by:
+
+    - a `continue` that yields the next iteration's value for each loop carried variable.
+    - a `break` that terminates the loop and yields the final loop carried values.
+
+    As long as each loop iteration is terminated by one of these operations they may be combined with other control
+    flow operations to express different control flow patterns.
+
+    The loop operation produces one return value for each loop carried variable. The type of the `i`-th return
+    value is that of the `i`-th loop carried variable and its value is the final value of the
+    `i`-th loop carried variable.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$initValues);
+  let results = (outs Variadic<AnyType>:$resultValues);
+  let regions = (region SizedRegion<1>:$region);
+
+  let extraClassDeclaration = [{
+    /// Return the iteration values of the loop region.
+    Block::BlockArgListType getRegionIterValues() {
+      return getRegion().getArguments().drop_front();
+    }
+
+    /// Return the `index`-th region iteration value.
+    BlockArgument getRegionIterValue(unsigned index) {
+      return getRegionIterValues()[index];
+    }
+
+    /// Return the loop control token.
+    BlockArgument getControlToken() {
+      return getRegion().getArgument(0);
+    }
+
+    /// Returns the number of region arguments for loop-carried values.
+    unsigned getNumRegionIterValues() {
+      return getRegion().getNumArguments() - 1;
+    }
+
+    /// Returns the loop block body
+    Block *getBody() { return &getRegion().front(); }
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasRegionVerifier = 1;
+  let hasCanonicalizer = 1;
+}
+
+
+//===----------------------------------------------------------------------===//
+// BreakOp
+//===----------------------------------------------------------------------===//
+
+def BreakOp : SCF_Op<"break", [
+    Terminator, RegionExitTerminatorOpInterface,
+    DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+      ["getMutableSuccessorOperands"]>,
+    ParentOneOf<["IfOp", "LoopOp"]>
+  ]> {
+  let summary = "Break from loop";
+  let description = [{
+    The `break` operation is a terminator that exits one or more nested
+    regions and terminates the `scf.loop` that defined its control token.
+
+    The `break` may yield any number of operands; their types must match the
+    result types of the target `scf.loop`.
+
+    Example — break out of the immediately enclosing loop:
+    ```mlir
+    scf.loop token(%loop) -> i32 {
+      scf.break [%loop] %result : i32
+    }
+    ```
+
+    Example — break out of a loop through an enclosing `scf.if`:
+    ```mlir
+    scf.loop token(%loop) {
+      scf.if %cond {
+        scf.break [%loop]
+      }
+      scf.continue [%loop]
+    }
+    ```
+  }];
+
+
+  let arguments = (ins Token:$targetToken, Variadic<AnyType>:$args);
+  let assemblyFormat = [{
+    ` ` `[` $targetToken `]` ($args^ `:` type($args))? attr-dict
+  }];
+  let extraClassDeclaration = [{
+    /// RegionExitTerminatorOpInterface: resolve the token identifying the loop.
+    ::llvm::SmallVector<::mlir::Operation *>
+    getPotentialTargets();
+  }];
+  let hasVerifier = 1;
+}
+
+
+//===----------------------------------------------------------------------===//
+// ContinueOp
+//===----------------------------------------------------------------------===//
+
+def ContinueOp : SCF_Op<"continue", [
+    Terminator, RegionExitTerminatorOpInterface,
+    DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+      ["getMutableSuccessorOperands"]>, ParentOneOf<["IfOp", "LoopOp"]>
+  ]> {
+  let summary = "Continue to next loop iteration";
+  let description = [{
+    The `continue` operation is a terminator that re-enters a `scf.loop`
+    for its next iteration. The target loop is the one that defined the control
+    token operand.
+
+    The operands of `continue` become the loop-carried values (iter_args) for
+    the next iteration; their types must match the loop's iter_arg types.
+
+    Example — continue the immediately enclosing loop:
+    ```mlir
+    scf.loop token(%loop) iter_args(%i = %init) : i32 {
+      %next = arith.addi %i, %one : i32
+      scf.continue [%loop] %next : i32
+    }
+    ```
+
+    Example — continue an outer loop from inside a nested `scf.if`:
+    ```mlir
+    scf.loop token(%outer) iter_args(%counter = %init) : i64 {
+      scf.loop token(%inner) iter_args(%inner_arg = %counter) : i64 {
+        scf.if %restart_outer {
+          scf.continue [%outer] %inner_arg : i64
+        }
+        scf.continue [%inner] %inner_arg : i64
+      }
+      scf.continue [%outer] %counter : i64
+    }
+    ```
+  }];
+
+  let arguments = (ins Token:$targetToken, Variadic<AnyType>:$args);
+  let assemblyFormat = [{
+    ` ` `[` $targetToken `]` ($args^ `:` type($args))? attr-dict
+  }];
+  let extraClassDeclaration = [{
+    /// RegionExitTerminatorOpInterface: resolve the token identifying the loop.
+    ::llvm::SmallVector<::mlir::Operation *>
+    getPotentialTargets();
----------------
matthias-springer wrote:

nit: `DeclareOpInterfaceMethods`

https://github.com/llvm/llvm-project/pull/166688