Update verifier for tf.While/tf.WhileRegion to not check for operand result type compatibility if op is shape invariant.

Shape invariant tf.While/tf.WhileRegion can have different result shapes from operand shapes.

PiperOrigin-RevId: 344929917
Change-Id: Ic160211075c0bb91548c71fec9f2d880ccd79723

Author: andyly

tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc

@@ -2534,21 +2534,20 @@ OpFoldResult VariableShapeOp::fold(ArrayRef<Attribute> operands) {
 
 static LogicalResult VerifyWhileTypes(Operation *op, TypeRange cond_input,
                                       TypeRange body_input,
-                                      TypeRange body_result) {
-  // Collect all the type lists for the op so that different pairs of type lists
-  // can be compared for the compatibility.
-  constexpr int kNumTypeLists = 5;
-  const std::array<TypeRangeWithDesc, kNumTypeLists> type_lists = {{
-      {op->getOperandTypes(), "input"},
+                                      TypeRange body_result,
+                                      bool shape_invariant) {
+  const TypeRangeWithDesc input_type = {op->getOperandTypes(), "input"};
+  const TypeRangeWithDesc result_type = {op->getResultTypes(), "result"};
+  constexpr int kNumRegionTypeLists = 3;
+  const std::array<TypeRangeWithDesc, kNumRegionTypeLists> region_types = {{
       {body_result, "body result"},
-      {op->getResultTypes(), "result"},
       {cond_input, "condition input"},
       {body_input, "body input"},
   }};
 
   // A pair of type lists should be cast compatible with each other if one is
   // converted to the another for a function call or assignment or there is a
-  // common source of inputs for both.  Therefore, the While op requires the
+  // common source of inputs for both. Therefore, the While op requires the
   // following pairs of type lists to be cast compatible for the tensor_cast
   // operation:
   //
@@ -2557,7 +2556,8 @@ static LogicalResult VerifyWhileTypes(Operation *op, TypeRange cond_input,
   // * Operands and body inputs to call the body function for the first
   //   iteration if the cond functions returns True or equivalent result.
   // * Operands and results to assign cond function arguments to op results if
-  //   the cond function returns False or equivalent result.
+  //   the cond function returns False or equivalent result. If the op is shape
+  //   invariant, this does not hold as shapes can differ.
   // * All three pairs using cond inputs, body inputs and results as operand is
   //   a common source for all three.
   // * Body result and cond inputs to call the cond function for the subsequent
@@ -2566,17 +2566,28 @@ static LogicalResult VerifyWhileTypes(Operation *op, TypeRange cond_input,
   //
   // Note that the operands and body results need not be compatible as they are
   // never converted from one to the another nor there is a common source
-  // tensors.  Compatibility requirement is not transitive.
-
-  for (int i = 0; i < kNumTypeLists; ++i) {
-    // Skip the first pair as the While op operands and body function results
-    // does not need to be compatible with each other.
-    for (int j = std::max(2, i + 1); j < kNumTypeLists; ++j) {
-      auto &a = type_lists[i];
-      auto &b = type_lists[j];
-      if (failed(VerifyTypeRangesAreCompatible(op, a, b))) return failure();
-    }
-  }
+  // tensors. Compatibility requirement is not transitive.
+
+  if (!shape_invariant &&
+      failed(VerifyTypeRangesAreCompatible(op, input_type, result_type)))
+    return failure();
+
+  // Skip the first pair as the While op operands and body function results does
+  // not need to be compatible with each other.
+  for (int i = 1; i < kNumRegionTypeLists; ++i)
+    if (failed(VerifyTypeRangesAreCompatible(op, input_type, region_types[i])))
+      return failure();
+
+  for (int i = 0; i < kNumRegionTypeLists; ++i)
+    if (failed(VerifyTypeRangesAreCompatible(op, result_type, region_types[i])))
+      return failure();
+
+  for (int i = 0; i < kNumRegionTypeLists; ++i)
+    for (int j = i + 1; j < kNumRegionTypeLists; ++j)
+      if (failed(VerifyTypeRangesAreCompatible(op, region_types[i],
+                                               region_types[j])))
+        return failure();
+
   return success();
 }
 
@@ -2601,7 +2612,8 @@ static LogicalResult Verify(WhileOp op) {
 
   if (failed(VerifyWhileTypes(op, /*cond_input=*/cond_fn_type.getInputs(),
                               /*body_input=*/body_fn_type.getInputs(),
-                              /*body_result=*/body_fn_type.getResults())))
+                              /*body_result=*/body_fn_type.getResults(),
+                              op.shape_invariant())))
     return failure();
   return success();
 }
@@ -2626,7 +2638,8 @@ static LogicalResult Verify(WhileRegionOp op) {
   Operation *body_yield = op.body().front().getTerminator();
   if (failed(VerifyWhileTypes(op, /*cond_input=*/op.cond().getArgumentTypes(),
                               /*body_input=*/op.body().getArgumentTypes(),
-                              /*body_result=*/body_yield->getOperandTypes())))
+                              /*body_result=*/body_yield->getOperandTypes(),
+                              op.shape_invariant())))
     return failure();
   return success();
 }

tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir

@@ -1849,7 +1849,7 @@ func private @testWhileBody(tensor<*xf32>) -> (tensor<*xi32>)
 // Test invalid 'While' operation
 func @testWhileResult(tensor<*xf32>) -> (tensor<*xf32>) {
 ^bb0(%arg0: tensor<*xf32>):
-  // expected-error @+1 {{'tf.While' op body result type tensor<*xi32> is incompatible with result type tensor<*xf32> at index 0}}
+  // expected-error @+1 {{'tf.While' op result type tensor<*xf32> is incompatible with body result type tensor<*xi32> at index 0}}
   %1 = "tf.While"(%arg0) {
     cond = @testWhileCond,
     body = @testWhileBody,
@@ -1933,6 +1933,19 @@ func @testWhileResult(tensor<*x!tf.resource<tensor<32xf32>>>) -> (tensor<!tf.res
   return %1 : tensor<!tf.resource>
 }
 
+// -----
+
+func private @cond(tensor<1x?x3xf32>) -> tensor<i1>
+func private @body(tensor<1x?x3xf32>) -> tensor<1x?x3xf32>
+
+// Test shape invariant 'While' operation verifier with different operand and
+// result shapes.
+// CHECK-LABEL: func @testShapeInvariantWhile
+func @testShapeInvariantWhile(%arg0: tensor<1x2x3xf32>) -> tensor<1x8x3xf32> {
+  %0 = "tf.While"(%arg0) {cond = @cond, body = @body, is_stateless = false, shape_invariant} : (tensor<1x2x3xf32>) -> tensor<1x8x3xf32>
+  return %0 : tensor<1x8x3xf32>
+}
+
 // -----
 // WhileRegion tests
 
@@ -2084,7 +2097,7 @@ func @testInvalidWhileRegion_I_BI_TypeMismatch(%arg0 : tensor<i32>) -> (tensor<i
 // -----
 
 func @testInvalidWhileRegion_O_BO_CountMismatch(%arg0 : tensor<i32>) -> (tensor<i32>) {
-  // expected-error @+1 {{'tf.WhileRegion' op body results (size = 2) should have the same number of values as results (size = 1)}}
+  // expected-error @+1 {{'tf.WhileRegion' op results (size = 1) should have the same number of values as body results (size = 2)}}
   %0 = "tf.WhileRegion"(%arg0) (
     {
      ^bb0(%carg: tensor<i32>):
@@ -2102,7 +2115,7 @@ func @testInvalidWhileRegion_O_BO_CountMismatch(%arg0 : tensor<i32>) -> (tensor<
 // -----
 
 func @testInvalidWhileRegionMismatch_O_BO_TypeMismatch(%arg0 : tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>) {
-  // expected-error @+1 {{'tf.WhileRegion' op body result type tensor<f32> is incompatible with result type tensor<i32> at index 0}}
+  // expected-error @+1 {{'tf.WhileRegion' op result type tensor<i32> is incompatible with body result type tensor<f32> at index 0}}
   %0 = "tf.WhileRegion"(%arg0) (
     {
      ^bb0(%carg: tensor<i32>):
@@ -2207,6 +2220,23 @@ func @testInvalidWhileRegionConditionOutputType(%arg : tensor<i32>) -> (tensor<i
   return %0 : tensor<i32>
 }
 
+// -----
+
+// Test shape invariant 'WhileRegion' operation verifier with different operand
+// and result shapes.
+// CHECK-LABEL: func @testShapeInvariantWhileRegion
+func @testShapeInvariantWhileRegion(%arg0: tensor<1x2x3xf32>) -> tensor<1x8x3xf32> {
+  %0 = "tf.WhileRegion"(%arg0) ( {
+  ^cond(%carg0: tensor<1x?x3xf32>):
+    %1 = "tf.SomeCondOp"(%carg0) : (tensor<1x?x3xf32>) -> tensor<i1>
+    "tf.Yield"(%1) : (tensor<i1>) -> ()
+  }, {
+  ^body(%barg0: tensor<1x?x3xf32>):
+    %1 = "tf.SomeBodyOp"(%barg0) : (tensor<1x?x3xf32>) -> tensor<1x?x3xf32>
+    "tf.Yield"(%1) : (tensor<1x?x3xf32>) -> ()
+  }) {is_stateless = false, shape_invariant} : (tensor<1x2x3xf32>) -> tensor<1x8x3xf32>
+  return %0 : tensor<1x8x3xf32>
+}
 
 // -----