Fix segfault caused by passing in a block argument into DefinedByConv2D constraint check.

PiperOrigin-RevId: 278115186
Change-Id: I8b40a2804cf7e84aacfd993855e751c75462276b

Author: andyly

tensorflow/compiler/mlir/tensorflow/tests/optimize-arg-operand-constraint.mlir

@@ -0,0 +1,13 @@
+// RUN: tf-opt -tf-optimize %s | FileCheck %s --dump-input=fail
+
+// Check passing an argument into DefinedByConv2D constraint does not crash.
+
+// CHECK-LABEL: func @main
+func @main(%arg0: tensor<1xf32>) -> tensor<1xf32>
+attributes  {tf.entry_function = {inputs = "input", outputs = "output_node"}} {
+  %0 = constant dense<2.000000e+00> : tensor<f32>
+  %1 = constant dense<1.000000e+00> : tensor<f32>
+  %2 = "tf.AddV2"(%arg0, %1) {T = "tfdtype$DT_FLOAT", device = "", name = "StatefulPartitionedCall/add"} : (tensor<1xf32>, tensor<f32>) -> tensor<1xf32>
+  %3 = "tf.Mul"(%2, %0) {T = "tfdtype$DT_FLOAT", device = "", name = "output_node"} : (tensor<1xf32>, tensor<f32>) -> tensor<1xf32>
+  return %3 : tensor<1xf32>
+}

tensorflow/compiler/mlir/tensorflow/transforms/optimize.td

@@ -21,7 +21,7 @@ def BroadcastableElements :
     Constraint<CPred<"TFL::IsBroadcastableElementsAttrs($0, $1)">>;
 def F32ElementsAttr : ElementsAttrBase<
     CPred<"$_self.cast<ElementsAttr>().getType().getElementType().isF32()">, "float constant tensor">;
-def DefinedByConv2D : Constraint<CPred<"llvm::isa<mlir::TF::Conv2DOp>($0->getDefiningOp())">>;
+def DefinedByConv2D : Constraint<CPred<"llvm::isa_and_nonnull<mlir::TF::Conv2DOp>($0->getDefiningOp())">>;
 
 // If we see a Conv2D op followed by Mul, then multiply the filter
 // with the value in Mul.