From 31225841216e735b677a257ccfd9e79a7c27bf62 Mon Sep 17 00:00:00 2001
From: Mahesh Ravishankar <ravishankarm@google.com>
Date: Fri, 26 Aug 2022 03:45:16 +0000
Subject: [PATCH] Change `softmax` test to use `maxf`.

The e2e softmax test uses `cmpf` -> `select` for max operations. Use
`maxf`instead. This allows the op to be vectorized. The TOSA to Linalg
lowering has been recently updated to do the same (and this test was
derived from using an older TOSA to Linalg lowering).
---
 tests/e2e/regression/softmax.mlir | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tests/e2e/regression/softmax.mlir b/tests/e2e/regression/softmax.mlir
index 7e4b5db25e23..5e51258d6d07 100644
--- a/tests/e2e/regression/softmax.mlir
+++ b/tests/e2e/regression/softmax.mlir
@@ -23,9 +23,8 @@ func.func @softmax() {
   %2 = linalg.fill ins(%cst_1 : f32) outs(%1 : tensor<12x128xf32>) -> tensor<12x128xf32>
   %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%0 : tensor<12x128x128xf32>) outs(%2 : tensor<12x128xf32>) {
   ^bb0(%arg0: f32, %arg1: f32):
-    %11 = arith.cmpf ogt, %arg0, %arg1 : f32
-    %12 = arith.select %11, %arg0, %arg1 : f32
-    linalg.yield %12 : f32
+    %11 = arith.maxf %arg0, %arg1 : f32
+    linalg.yield %11 : f32
   } -> tensor<12x128xf32>
   %4 = linalg.init_tensor [12, 128, 128] : tensor<12x128x128xf32>
   %5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0, %3 : tensor<12x128x128xf32>, tensor<12x128xf32>) outs(%4 : tensor<12x128x128xf32>) {
@@ -74,9 +73,8 @@ func.func @softmax_dynamic() {
   %2 = linalg.fill ins(%cst_1 : f32) outs(%1 : tensor<?x?xf32>) -> tensor<?x?xf32>
   %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%cst_3 : tensor<?x?x?xf32>) outs(%2 : tensor<?x?xf32>) {
   ^bb0(%arg0: f32, %arg1: f32):
-    %11 = arith.cmpf ogt, %arg0, %arg1 : f32
-    %12 = arith.select %11, %arg0, %arg1 : f32
-    linalg.yield %12 : f32
+    %11 = arith.maxf %arg0, %arg1 : f32
+    linalg.yield %11 : f32
   } -> tensor<?x?xf32>
   %4 = linalg.init_tensor [%dim_0, %dim_1, %dim_2] : tensor<?x?x?xf32>
   %5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%cst_3, %3 : tensor<?x?x?xf32>, tensor<?x?xf32>) outs(%4 : tensor<?x?x?xf32>) {