Fix castings with invalid assumptions in mlir/lite.

PiperOrigin-RevId: 344753711
Change-Id: I8f01a52929888a124387323a8f30cebf048de15d

Author: abattery

tensorflow/compiler/mlir/lite/ir/tfl_ops.cc

@@ -494,6 +494,8 @@ Attribute ConstFoldUnaryOp(Type result_type, Attribute operand,
   assert(IsF32ShapedType(result_type) || IsBF16ShapedType(result_type));
   auto result_shape_type = result_type.cast<ShapedType>();
 
+  if (!result_shape_type.hasStaticShape()) return {};
+
   if (auto dense_elements = operand.dyn_cast_or_null<DenseElementsAttr>()) {
     SmallVector<APFloat, 16> new_values;
     const int num_elements = result_shape_type.getNumElements();
@@ -1740,7 +1742,8 @@ static LogicalResult Verify(LSTMOp op) {
           op.forget_layer_norm_coefficients().getType().cast<ShapedType>();
       // If this lstm has layer normalization, this input value,
       // "forget_layer_norm_coefficients" should be a 1D tensor.
-      if (forget_layer_norm_coefficients.getRank() != 1 ||
+      if (!forget_layer_norm_coefficients.hasRank() ||
+          forget_layer_norm_coefficients.getRank() != 1 ||
           forget_layer_norm_coefficients.getDimSize(0) != n_cell)
         return op.emitOpError(
             "coefficient inputs have more than 2 dimensions or "

tensorflow/compiler/mlir/lite/tests/dilated-conv.mlir

@@ -283,6 +283,28 @@ func @testDilatedDepthWiseConvWithExpandSqueeze3(%arg0: tensor<1x128x128xf32>, %
   // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32>
 }
 
+func @testAvoidDilatedConvWithExpand(%arg0: tensor<*xf32>, %arg1: tensor<5x5x1x1xf32>, %arg2: tensor<128xf32>) -> tensor<1x128x128xf32> {
+  %cst = constant dense<[2, 2]> : tensor<2xi32>
+  %cst_0 = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
+  %cst_1 = constant dense<4> : tensor<2x2xi32>
+  %cst_2 = constant dense<0> : tensor<2x2xi32>
+  %0 = "tf.SpaceToBatchND"(%arg0, %cst, %cst_1) : (tensor<*xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x68x68xf32>
+  %1 = "tf.ExpandDims"(%0, %cst_0) : (tensor<4x68x68xf32>, tensor<i32>) -> tensor<4x68x68x1xf32>
+  %2 = "tf.Conv2D"(%1, %arg1) {padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<4x68x68x1xf32>, tensor<5x5x1x1xf32>) -> tensor<4x64x64x1xf32>
+  %3 = "tf.Squeeze"(%2) {squeeze_dims = [3]} : (tensor<4x64x64x1xf32>) -> tensor<4x64x64xf32>
+  %4 = "tf.BatchToSpaceND"(%3, %cst, %cst_2) : (tensor<4x64x64xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<1x128x128xf32>
+  %5 = "tf.BiasAdd"(%4, %arg2) : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32>
+  return %5 : tensor<1x128x128xf32>
+
+  // CHECK-LABEL: testAvoidDilatedConvWithExpand
+  // CHECK: "tf.SpaceToBatchND"
+  // CHECK: "tf.ExpandDims"
+  // CHECK: "tf.Conv2D"
+  // CHECK: "tf.Squeeze"
+  // CHECK: "tf.BatchToSpaceND"
+  // CHECK: "tf.BiasAdd"
+}
+
 func @testDilatedConvWithDifferentExpandSqueezeAxis(%arg0: tensor<1x128x128xf32>, %arg1: tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> {
   %cst = constant dense<[2, 2]> : tensor<2xi32>
   %cst_0 = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

tensorflow/compiler/mlir/lite/tests/ops.mlir

@@ -810,6 +810,16 @@ func @testLstmWithInvalidInputsRankMatch(%arg0: tensor<1x4xf32>, %arg1: tensor<4
   return %24 : tensor<1x4xf32>
 }
 
+// -----
+
+// Coefficient inputs of LSTM op have unknown rank.
+func @testLstmWithInvalidInputsRankMatch(%arg0: tensor<1x4xf32>, %arg1: tensor<4x4xf32>, %arg2: tensor<4x4xf32>, %arg3: tensor<4x4xf32>, %arg4: tensor<4x4xf32>, %arg5: tensor<4x4xf32>, %arg6: tensor<4x4xf32>, %arg7: tensor<4x4xf32>, %arg8: tensor<4x4xf32>, %arg9: tensor<4xf32>, %arg10: tensor<4xf32>, %arg11: tensor<4xf32>, %arg12: tensor<1x4xf32>, %arg13: tensor<4xf32>, %arg14: tensor<4xf32>, %arg15: tensor<4xf32>, %arg16: tensor<4x4xf32>, %arg17: tensor<4xf32>, %arg18: tensor<3xf32>, %arg19: tensor<3xf32>, %arg20: tensor<3xf32>, %arg21: tensor<*xf32>) -> tensor<1x4xf32> {
+  %cst0 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
+  %cst1 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
+  // expected-error @+1 {{'tfl.lstm' op coefficient inputs have more than 2 dimensions or don't match the dimension with input operand `input_to_output_weights`.}}
+  %24 = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %cst0, %cst1, %arg18, %arg19, %arg20, %arg21) ({}) {cell_clip = 0.000000e+00 : f32, fused_activation_function = "NONE", kernel_type = "FULL", proj_clip = 0.000000e+00 : f32} : (tensor<1x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<1x4xf32>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>, tensor<*xf32>) -> tensor<1x4xf32>
+  return %24 : tensor<1x4xf32>
+}
 
 // -----
 

tensorflow/compiler/mlir/lite/tests/optimize.mlir

@@ -343,6 +343,18 @@ func @fuseAddIntoFollowingFullyConnected(%arg0: tensor<4x2xf32>) -> tensor<4x2xf
 // CHECK-NEXT: return %[[fc]] : tensor<4x2xf32>
 }
 
+// CHECK-LABEL: @doNotFuseAddIntoFollowingFullyConnected
+func @doNotFuseAddIntoFollowingFullyConnected(%arg0: tensor<4x2xf32>, %arg1: tensor<*xf32>) -> tensor<4x2xf32> {
+  %cst1 = constant dense<1.5> : tensor<f32>
+  %0 = "tfl.add"(%arg0, %cst1) {fused_activation_function = "NONE"} : (tensor<4x2xf32>, tensor<f32>) -> tensor<4x2xf32>
+  %cst = constant dense<2.0> : tensor<2xf32>
+  %1 = "tfl.fully_connected"(%0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x2xf32>, tensor<*xf32>, tensor<2xf32>) -> tensor<4x2xf32>
+  return %1 : tensor<4x2xf32>
+
+// CHECK: "tfl.add"
+// CHECK: "tfl.fully_connected"
+}
+
 // CHECK-LABEL: @fuseMulIntoFollowingFullyConnected
 func @fuseMulIntoFollowingFullyConnected(%arg0: tensor<4x2xf32>) -> tensor<4x2xf32> {
   %cst2 = constant dense<1.5> : tensor<f32>

tensorflow/compiler/mlir/lite/transforms/dense_to_sparse.cc

@@ -175,6 +175,9 @@ InspectResult InspectWeight(
   } else if (auto cst = dyn_cast<QConstOp>(inst)) {
     attr = cst.value();
     type = cst.getType().cast<ShapedType>();
+  } else {
+    result.can_compress = false;
+    return result;
   }
 
   // Currently we only support compressing weights of ops:
@@ -222,6 +225,8 @@ std::vector<T> BuildSparsityParameterAttribute(
   } else if (auto cst = dyn_cast<QConstOp>(inst)) {
     attr = cst.value();
     type = cst.getType().cast<ShapedType>();
+  } else {
+    assert(false && "Expected a constant-like op");
   }
   const int dims_count = type.getRank();
   std::vector<int> shape(dims_count);

tensorflow/compiler/mlir/lite/transforms/dilated_conv.h

@@ -185,7 +185,12 @@ LogicalResult ConvertTFDilatedConvOp<Conv2dOpTy>::matchAndRewrite(
   llvm::Optional<ArrayAttr> dilations_attr = ExtractDilationsAttrFromBlockShape(
       stb_op.block_shape(), bts_op.block_shape(), rewriter);
   if (!dilations_attr.hasValue()) return failure();
-  op.setAttr("dilations", dilations_attr.getValue());
+
+  if (expand_op) {
+    if (stb_op.input().getType().dyn_cast<RankedTensorType>() == nullptr) {
+      return failure();
+    }
+  }
 
   // TODO(b/149936532): Check that the input width & height are multiples of
   // dilation rate.
@@ -234,6 +239,9 @@ LogicalResult ConvertTFDilatedConvOp<Conv2dOpTy>::matchAndRewrite(
     }
   }
 
+  // Set dilations
+  op.setAttr("dilations", dilations_attr.getValue());
+
   if (expand_op) {
     // If there is `expand_op`, we need to rewire the inputs to bypass the
     // `SpaceToBatch`, `BatchToSpace` and `Pad` op. E.g, turning

tensorflow/compiler/mlir/lite/transforms/optimize.cc

@@ -351,10 +351,10 @@ struct FuseFullyConnectedAndAdd : public OpRewritePattern<TFL::AddOp> {
         // to properly broadcast the scalar to `{num_channels}` shape.
 
         // Get the number of channels if possible.
-        auto filter_type = filter.getType().cast<ShapedType>();
+        auto filter_type = filter.getType().dyn_cast<RankedTensorType>();
         // Filter must be a `2D` tensor with `{num_channels, num_features}`
         // shape. The following check is rejecting unknown rank (-1).
-        if (filter_type.getRank() != 2) {
+        if (filter_type == nullptr || filter_type.getRank() != 2) {
           return failure();
         }
         int num_channels = filter_type.getShape()[0];