Fix integer narrowing bug in BatchMatMul unrolling logic.

The implicit cast of shape values (of type int64_t) to 32-bit int results in the conversion of any ShapedType::kDynamic values (int64::min()) to 0. PiperOrigin-RevId: 492876058
tensorflow · Dec 5, 2022 · 3a5eb5f · 3a5eb5f
1 parent ee3c46d
commit 3a5eb5f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/unroll-batch-matmul.mlir b/tensorflow/compiler/mlir/tensorflow/tests/unroll-batch-matmul.mlir
@@ -448,6 +448,17 @@ func.func @batchMatMulV2MatrixAdjXY(%arg0: tensor<5x4xf32>, %arg1: tensor<6x5xf3
   // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32>
 }
 
+// -----
+
+func.func @batchMatMulV2DynamicSize(%arg0: tensor<?x?xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x4xf32> {
+  %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<?x?xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  func.return %0 : tensor<?x4xf32>
+
+  // CHECK-LABEL: batchMatMulV2DynamicSize
+  // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  // CHECK: return %[[MATMUL_1]] : tensor<?x4xf32>
+}
+
 // -----
 // ==== V3 tests ====
 

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/unroll_batch_matmul.cc b/tensorflow/compiler/mlir/tensorflow/transforms/unroll_batch_matmul.cc
@@ -198,8 +198,8 @@ LogicalResult ConvertTFBatchMatMulOp<BatchMatMulOpType>::matchAndRewrite(
     std::swap(rhs_shape[rhs_dims - 1], rhs_shape[rhs_dims - 2]);
   }
 
-  const int rows = lhs_shape[lhs_dims - 2];
-  const int cols = rhs_shape[rhs_dims - 1];
+  const int64_t rows = lhs_shape[lhs_dims - 2];
+  const int64_t cols = rhs_shape[rhs_dims - 1];
 
   if (lhs_shape[lhs_dims - 1] != rhs_shape[rhs_dims - 2]) {
     // Input dimensions must be compatible for multiplication.