Fix failing TFRT lit tests.

Mahesh Ravishankar · tensorflower-gardener · commit c8dafc9f9ae6 · 2021-12-14T13:38:10.000-08:00
PiperOrigin-RevId: 416380991
Change-Id: I8bc650f81f1adcdf972923419275e77086b92f4e
diff --git a/tensorflow/compiler/mlir/tfrt/tests/jit/tf_cpurt_vectorize_tiled_ops.mlir b/tensorflow/compiler/mlir/tfrt/tests/jit/tf_cpurt_vectorize_tiled_ops.mlir
@@ -1,4 +1,4 @@
-// RUN: tf-tfrt-opt %s --tf-cpurt-vectorize-tiled-ops --split-input-file |\
+// RUN: tf-tfrt-opt %s --vectorize-tiled-ops --split-input-file |\
 // RUN: FileCheck %s
 
 #map0 = affine_map<(d0) -> (d0)>
diff --git a/tensorflow/compiler/mlir/tfrt/tests/jit/vectorize_tiled_ops.mlir b/tensorflow/compiler/mlir/tfrt/tests/jit/vectorize_tiled_ops.mlir
@@ -34,26 +34,21 @@ func @tiled_add(%A: tensor<8xf32>, %B: tensor<8xf32>,
 // CHECK-LABEL: func @tiled_add
 
 // CHECK-DAG:  %[[CST:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
 
-// CHECK: linalg.tiled_loop
+// CHECK: linalg.tiled_loop (%[[IV0:arg[0-9]]])
 // CHECK-SAME: ins (%[[A:arg[0-9]]] = %{{arg[0-9]}}: tensor<8xf32>,
 // CHECK-SAME:      %[[B:arg[0-9]]] = %{{arg[0-9]}}: tensor<8xf32>
 // CHECK-SAME: outs (%[[C:arg[0-9]]] = %{{arg[0-9]}}: tensor<8xf32>)
 
-// CHECK-NEXT: %[[SUB_A:.*]] = tensor.extract_slice %[[A]]
-// CHECK-NEXT: %[[SUB_B:.*]] = tensor.extract_slice %[[B]]
-// CHECK-NEXT: %[[SUB_C:.*]] = tensor.extract_slice %[[C]]
-
-// CHECK-NEXT: %[[LHS:.*]] = vector.transfer_read %1[%[[C0]]], %[[CST]]
-// CHECK-SAME:   {in_bounds = [true]} : tensor<2xf32>, vector<2xf32>
-// CHECK-NEXT: %[[RHS:.*]] = vector.transfer_read %2[%[[C0]]], %[[CST]]
-// CHECK-SAME:   {in_bounds = [true]} : tensor<2xf32>, vector<2xf32>
+// CHECK-NEXT: %[[LHS:.*]] = vector.transfer_read %[[A]][%[[IV0]]], %[[CST]]
+// CHECK-SAME:   {in_bounds = [true]} : tensor<8xf32>, vector<2xf32>
+// CHECK-NEXT: %[[RHS:.*]] = vector.transfer_read %[[B]][%[[IV0]]], %[[CST]]
+// CHECK-SAME:   {in_bounds = [true]} : tensor<8xf32>, vector<2xf32>
 
 // CHECK-NEXT: %[[SUM:.*]] = arith.addf %[[LHS]], %[[RHS]] : vector<2xf32>
 
-// CHECK-NEXT: %{{.*}} = vector.transfer_write %[[SUM]], %[[SUB_C]][%[[C0]]]
-// CHECK-SAME:   {in_bounds = [true]} : vector<2xf32>, tensor<2xf32>
+// CHECK-NEXT: %{{.*}} = vector.transfer_write %[[SUM]], %[[C]][%[[IV0]]]
+// CHECK-SAME:   {in_bounds = [true]} : vector<2xf32>, tensor<8xf32>
 
 // -----
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-// RUN: tf-tfrt-opt %s --tf-cpurt-vectorize-tiled-ops --split-input-file \|\`
	`1`	`+// RUN: tf-tfrt-opt %s --vectorize-tiled-ops --split-input-file \|\`
`2`	`2`	`// RUN: FileCheck %s`
`3`	`3`
`4`	`4`	`#map0 = affine_map<(d0) -> (d0)>`