[Mlir-commits] [mlir] 4fdc019 - [mlir][SCF] Add `SingleBlock` op trait to "scf.while"

[Matthias Springer]

Author: Matthias Springer
Date: 2023-08-31T08:56:31+02:00
New Revision: 4fdc019a89105c956d4c79fb79ca1b416eaa72f2

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

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

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.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
    mlir/test/Dialect/SCF/invalid.mlir
    mlir/test/Transforms/inlining.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 477e1aefab0e2f..232e6b0bf4ed77 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -976,7 +976,7 @@ def ReduceReturnOp :
 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

diff  --git a/mlir/test/Dialect/SCF/invalid.mlir b/mlir/test/Dialect/SCF/invalid.mlir
index 0ff3eaadc8fecd..0cf587af42637c 100644
--- a/mlir/test/Dialect/SCF/invalid.mlir
+++ b/mlir/test/Dialect/SCF/invalid.mlir
@@ -478,11 +478,25 @@ func.func @while_empty_region() {
 // -----
 
 func.func @while_empty_block() {
-  // expected-error at +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:
+  ^bb0:
+    // expected-note @below{{terminator here}}
+    "test.foo"() : () -> ()
   } do {
-   ^bb0:
+  ^bb0:
+    "test.bar"() : () -> ()
   }
 }
 

diff  --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index a0b4e0082b5439..9544f1eb09170e 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -227,18 +227,40 @@ func.func @func_with_block_args_location_callee2(%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