[Mlir-commits] [mlir] [mlir][tosa] Check for isolated regions in `tosa.while_loop` (PR #144865)

[Jack Frankland]

================
@@ -227,15 +227,79 @@ func.func @test_error_i32_unsigned_output(%arg0: tensor<1xi8>) -> tensor<1xi32>
 }
 
 // -----
-// CHECK-LABEL: cond_if_simplified_form
-func.func @test_cond_if_simplified_form(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
-  // expected-error at +1 {{'tosa.cond_if' op the current simplified form is not strictly conformant to the spec, please use the generic format}}
+
+func.func @test_cond_if_not_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
+    // expected-error at +1 {{'tosa.cond_if' op is not conformant to the TOSA specification. It requires the then/else regions are isolated from above.}}
+    %0 = "tosa.cond_if"(%arg2) ({
+      ^bb0():
+        tosa.yield %arg0 : tensor<f32>
+      },  {
+      ^bb0():
+        tosa.yield %arg1 : tensor<f32>
+      }) : (tensor<i1>) -> tensor<f32>
+    return %0 : tensor<f32>
+}
+
+// -----
+
+func.func @test_cond_if_simplified_form_not_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
+  // expected-error at +1 {{'tosa.cond_if' op is not conformant to the TOSA specification. It requires the then/else regions are isolated from above.}}
   %0 = tosa.cond_if %arg2 -> (tensor<f32>) {
-    %1 = tosa.add %arg0, %arg1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
-    tosa.yield %1 : tensor<f32>
+    tosa.yield %arg0 : tensor<f32>
   } else {
-    %1 = tosa.sub %arg0, %arg1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
-    tosa.yield %1 : tensor<f32>
+    tosa.yield %arg1 : tensor<f32>
   }
   return %0 : tensor<f32>
 }
+
+// -----
+
+// COM: Check isolated cond_if's are valid
+func.func @test_cond_if_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
+      tosa.yield %arg3 : tensor<f32>
+    },  {
+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
+      tosa.yield %arg4 : tensor<f32>
+    }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// -----
+
+func.func @test_while_loop_not_isolated_from_above(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<f32>) {
+  %0 = "tosa.const"() {values = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // expected-error at +1 {{'tosa.while_loop' op is not conformant to the TOSA specification. It requires the cond/body regions are isolated from above.}}
+  %1 = "tosa.while_loop"(%0) ({
+  ^bb0(%arg3: tensor<i32>):
+    %2 = "tosa.greater_equal"(%arg3, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    %3 = "tosa.logical_not"(%2) : (tensor<i1>) -> tensor<i1>
+    tosa.yield %3 : tensor<i1>
+  },  {
+  ^bb0(%arg3: tensor<i32>):
----------------
FranklandJack wrote:

Here it looks like the body graph is isolated from above, and it is only the condition graph that isn't. Should we add another test to check the error is triggered when the body graph isn't isolated from above?

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