Refine result type of std.tensor_cast op during shape inference.

Prakalp Srivastava · tensorflower-gardener · commit 1ce189fed323 · 2020-04-28T15:54:35.000-07:00
std.tensor_cast operand and result are cast compatible, so it is safe to refine the result type of std.tensor_cast to be same as operand.

PiperOrigin-RevId: 308913809
Change-Id: I414ad8f32d15e864faddff75e5bf2eaa4bb95262
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir b/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir
@@ -291,6 +291,15 @@ func @multiple_blocks_one_return(%arg0: tensor<?xf32>) -> tensor<*xf32> {
     return %0 : tensor<?x?x?xf32>
   }
 
+  // Tests that tensor_cast result shapes are refined.
+  // CHECK-LABEL: func @tensor_cast_refine
+  func @tensor_cast_refine(%arg0: tensor<4xi32>) -> (tensor<*xi32>) {
+    // CHECK: tensor_cast
+    // CHECK-SAME: tensor<4xi32> to tensor<4xi32>
+    %0 = tensor_cast %arg0 : tensor<4xi32> to tensor<*xi32>
+    return %0 : tensor<*xi32>
+  }
+
   // CHECK-LABEL: func @fold_cast
   func @fold_cast(%arg0: tensor<*xf32>) -> tensor<*xf32> {
     // CHECK-NOT: Cast
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc b/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc
@@ -185,6 +185,10 @@ bool InferShapeForNonTFDialectOperation(Operation* op, Dialect* tf_dialect) {
         iter_sink.getOperands().drop_front().take_front(), iter_source,
         tf_dialect);
   }
+  if (auto tensor_cast = dyn_cast<mlir::TensorCastOp>(op)) {
+    return InferShapeForPassThroughOps(
+        tensor_cast.getOperation()->getOperands(), op, tf_dialect);
+  }
   return false;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -185,6 +185,10 @@ bool InferShapeForNonTFDialectOperation(Operation* op, Dialect* tf_dialect) {`
`185`	`185`	`iter_sink.getOperands().drop_front().take_front(), iter_source,`
`186`	`186`	`tf_dialect);`
`187`	`187`	`}`
	`188`	`+ if (auto tensor_cast = dyn_cast<mlir::TensorCastOp>(op)) {`
	`189`	`+ return InferShapeForPassThroughOps(`
	`190`	`+ tensor_cast.getOperation()->getOperands(), op, tf_dialect);`
	`191`	`+ }`
`188`	`192`	`return false;`
`189`	`193`	`}`
`190`	`194`