fix

Author: wsmoses

test/lit_tests/dus_ds_while_simplify.mlir

@@ -324,7 +324,7 @@ module {
 // AUTOBATCH-NEXT:   %[[a3:.+]] = stablehlo.pad %[[a2]], %cst, low = [0, 0], high = [0, 6], interior = [0, 0] : (tensor<15x14xf32>, tensor<f32>) -> tensor<15x20xf32>
 // AUTOBATCH-NEXT:   %[[a4:.+]] = stablehlo.slice %arg1 [0:14, 0:1] : (tensor<15x3xf32>) -> tensor<14x1xf32>
 // AUTOBATCH-NEXT:   %[[a5:.+]] = stablehlo.reshape %[[a4]] : (tensor<14x1xf32>) -> tensor<14x1x1xf32>
-// AUTOBATCH-NEXT:   %[[a6:.+]] = stablehlo.broadcast_in_dim %[[a4]], dims = [0, 2, 1] : (tensor<14x1x1xf32>) -> tensor<14x5x3xf32>
+// AUTOBATCH-NEXT:   %[[a6:.+]] = stablehlo.broadcast_in_dim %[[a5]], dims = [0, 2, 1] : (tensor<14x1x1xf32>) -> tensor<14x5x3xf32> 
 // AUTOBATCH-NEXT:   %[[a7:.+]] = stablehlo.reduce(%[[a6]] init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<14x5x3xf32>, tensor<f32>) -> tensor<5x3xf32>
 // AUTOBATCH-NEXT:   return %[[a7]], %[[a3]] : tensor<5x3xf32>, tensor<15x20xf32>
 // AUTOBATCH-NEXT: }

test/lit_tests/linalg/syrk.mlir

@@ -102,18 +102,18 @@ module {
 // CUDA-NEXT: }
 
 // TPU: func.func @main2(%arg0: tensor<64x32xf32>, %arg1: tensor<64x64xf32>) -> tensor<64x64xf32> {
-// TPU-NEXT:     %cst = stablehlo.constant dense<3.000000e+00> : tensor<64x64xf32>
-// TPU-NEXT:     %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<64x64xf32>
+// TPU-NEXT:     %[[cst2:.+]] = stablehlo.constant dense<2.000000e+00> : tensor<64x64xf32>
+// TPU-NEXT:     %[[cst3:.+]] = stablehlo.constant dense<3.000000e+00> : tensor<64x64xf32>
 // TPU-NEXT:     %0 = stablehlo.iota dim = 0 : tensor<64x64xi32>
 // TPU-NEXT:     %1 = stablehlo.iota dim = 1 : tensor<64x64xi32>
 // TPU-NEXT:     %2 = stablehlo.compare  GT, %0, %1 : (tensor<64x64xi32>, tensor<64x64xi32>) -> tensor<64x64xi1>
 // TPU-NEXT:     %3 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<64x64xf32>) -> tensor<64x64xf32>
 // TPU-NEXT:     %4 = stablehlo.select %2, %arg1, %3 : tensor<64x64xi1>, tensor<64x64xf32>
 // TPU-NEXT:     %5 = stablehlo.dot_general %arg0, %arg0, contracting_dims = [1] x [1] : (tensor<64x32xf32>, tensor<64x32xf32>) -> tensor<64x64xf32>
-// TPU-NEXT:     %6 = stablehlo.multiply %cst_0, %5 : tensor<64x64xf32>
-// TPU-NEXT:     %7 = stablehlo.multiply %cst, %4 : tensor<64x64xf32>
-// TPU-NEXT:     %8 = stablehlo.add %6, %7 : tensor<64x64xf32>
-// TPU-NEXT:     return %8 : tensor<64x64xf32>
+// TPU-NEXT:     %[[a7:.+]] = stablehlo.multiply %[[cst3]], %4 : tensor<64x64xf32>
+// TPU-NEXT:     %[[a6:.+]] = stablehlo.multiply %[[cst2]], %5 : tensor<64x64xf32>
+// TPU-NEXT:     %[[a8:.+]] = stablehlo.add %[[a6]], %[[a7]] : tensor<64x64xf32>
+// TPU-NEXT:     return %[[a8]] : tensor<64x64xf32>
 // TPU-NEXT: }
 
 module {
@@ -157,8 +157,8 @@ module {
 // TPU-NEXT:     %[[cst2:.+]] = stablehlo.constant dense<2.000000e+00> : tensor<64x64xf32>
 // TPU-NEXT:     %[[cst3:.+]] = stablehlo.constant dense<3.000000e+00> : tensor<64x64xf32>
 // TPU-NEXT:     %[[a0:.+]] = stablehlo.dot_general %arg0, %arg0, contracting_dims = [1] x [1] : (tensor<64x32xf32>, tensor<64x32xf32>) -> tensor<64x64xf32>
-// TPU-NEXT:     %[[a1:.+]] = stablehlo.multiply %[[cst2]], %[[a0]] : tensor<64x64xf32>
 // TPU-NEXT:     %[[a2:.+]] = stablehlo.multiply %[[cst3]], %arg1 : tensor<64x64xf32>
+// TPU-NEXT:     %[[a1:.+]] = stablehlo.multiply %[[cst2]], %[[a0]] : tensor<64x64xf32>
 // TPU-NEXT:     %[[a3:.+]] = stablehlo.add %[[a1]], %[[a2]] : tensor<64x64xf32>
 // TPU-NEXT:     return %[[a3]] : tensor<64x64xf32>
 // TPU-NEXT: }