[mlir][SCF] Add `SingleBlock` op trait to "scf.while"

[springerm]

Summary

This trait is needed so that unstructured control flow is not inlined into "scf.while" ops.

Note: The two regions of "scf.while" are already defined as SizedRegion<1>. SingleBlock can be queried from C++, SizedRegion<n> not.

Fixes #64976.

[javedabsar]

LGTM other than the nitpick comment.

mlir/test/Transforms/inlining.mlir

255

Guess the 'func_with_block_args_location' is used here for its 'cf.br' rather than loc part.
Maybe create a new caller with better name for the purpose here?

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

 
 //===----------------------------------------------------------------------===//
 // WhileOp
 //===----------------------------------------------------------------------===//
 
 def WhileOp : SCF_Op<"while",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getEntrySuccessorOperands"]>,
-     RecursiveMemoryEffects]> {
+     RecursiveMemoryEffects, SingleBlock]> {
   let summary = "a generic 'while' loop";
   let description = [{
     This operation represents a generic "while"/"do-while" loop that keeps
     iterating as long as a condition is satisfied. There is no restriction on
     the complexity of the condition. It consists of two regions (with single
     block each): "before" region and "after" region. The names of regions
     indicates whether they execute before or after the condition check.
     Therefore, if the main loop payload is located in the "before" region, the

mlir/test/Dialect/SCF/invalid.mlir

func.func @while_empty_region() {
   scf.while : () -> () {
   } do {
   }
 }
 
 // -----
 
 func.func @while_empty_block() {
-  // expected-error@+1 {{expects the 'before' region to terminate with 'scf.condition'}}
+  // expected-error @below {{expects a non-empty block}}
   scf.while : () -> () {
-   ^bb0:
   } do {
-   ^bb0:
   }
 }
 
 // -----
 
+func.func @while_invalid_terminator() {
+  // expected-error @below {{expects the 'before' region to terminate with 'scf.condition'}}
+  scf.while : () -> () {
+  ^bb0:
+    // expected-note @below{{terminator here}}
+    "test.foo"() : () -> ()
+  } do {
+  ^bb0:
+    "test.bar"() : () -> ()
+  }
+}
+
+// -----
+
 func.func @while_cross_region_type_mismatch() {
   %true = arith.constant true
   // expected-error@+1 {{'scf.while' op  region control flow edge from Region #0 to Region #1: source has 0 operands, but target successor needs 1}}
   scf.while : () -> () {
     scf.condition(%true)
   } do {
   ^bb0(%arg0: i32):
     scf.yield

mlir/test/Transforms/inlining.mlir

 }
 
 // CHECK-LABEL: func @func_with_block_args_location_callee2
 func.func @func_with_block_args_location_callee2(%arg0 : i32) {
   call @func_with_block_args_location(%arg0) : (i32) -> ()
   return
 }
 
-// CHECK-LABEL: func @func_with_block_args_location_callee3
-func.func @func_with_block_args_location_callee3(%arg0 : i32) {
+func.func @func_with_multiple_blocks(%arg0 : i32) {
+  cf.br ^bb1(%arg0 : i32)
+^bb1(%x : i32):
+  "test.foo" (%x) : (i32) -> () loc("bar")
+  return
+}
+
+// CHECK-LABEL: func @func_with_multiple_blocks_callee1
+func.func @func_with_multiple_blocks_callee1(%arg0 : i32) {
   "test.dummy_op"() ({
     // Call cannot be inlined because "test.dummy" may not support unstructured
     // control flow in its body.
-    // CHECK: call @func_with_block_args_location
-    call @func_with_block_args_location(%arg0) : (i32) -> ()
+    // CHECK: call @func_with_multiple_blocks
+    call @func_with_multiple_blocks(%arg0) : (i32) -> ()
     "test.terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
+// CHECK-LABEL: func @func_with_multiple_blocks_callee2
+func.func @func_with_multiple_blocks_callee2(%arg0 : i32, %c : i1) {
+  %0 = scf.while (%arg1 = %arg0) : (i32) -> (i32) {
+    // Call cannot be inlined because scf.while does not support unstructured
+    // control flow in its body.
+    // CHECK: call @func_with_multiple_blocks
+    func.call @func_with_multiple_blocks(%arg0) : (i32) -> ()
+    scf.condition(%c) %arg1 : i32
+  } do {
+  ^bb0(%arg1: i32):
+    scf.yield %arg1 : i32
+  }
+  return
+}
+
 // Check that we can handle argument and result attributes.
 test.conversion_func_op @handle_attr_callee_fn_multi_arg(%arg0 : i16, %arg1 : i16 {"test.handle_argument"}) -> (i16 {"test.handle_result"}, i16) {
   %0 = arith.addi %arg0, %arg1 : i16
   %1 = arith.subi %arg0, %arg1 : i16
   "test.return"(%0, %1) : (i16, i16) -> ()
 }
 test.conversion_func_op @handle_attr_callee_fn(%arg0 : i32 {"test.handle_argument"}) -> (i32 {"test.handle_result"}) {
   "test.return"(%arg0) : (i32) -> ()