[BACKPORT] Allow softmax type conversion to happen before or after elementwise (#1920)

Allow softmax type conversion to happen before or after elementwise ops in attention (#1911)

For the attention kernel, migraphx can decide to do softmax in higher precision. It would insert necessary convert operation to cast results back and forth around softmax ops.

But these casts could appear before or after any elementwise scaling ops that happen before softmax. Therefore this PR extends pattern matcher to allow such cases.

For example before, it only handled cases like :

gemm1  -> elementwise ops ->  convert(to_higher_precision) -> softmax -> convert(back_to_lower) -> gemm2

With this PR, it can also handle case where it is like following :

gemm1  -> convert(to_higher_precision) -> elementwise ops -> softmax -> convert(back_to_lower) -> gemm2

Author: umangyadav

mlir/include/mlir/Dialect/MIGraphX/Passes.h

@@ -22,6 +22,7 @@ namespace migraphx {
 
 #define GEN_PASS_DECL_MIGRAPHXREALIZEINT4PASS
 #define GEN_PASS_DECL_MIGRAPHXTRANSFORMPASS
+#define GEN_PASS_DECL_MIGRAPHXTOSASIMPLIFYPASS
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/MIGraphX/Passes.h.inc"
 

mlir/include/mlir/Dialect/MIGraphX/Passes.td

@@ -41,4 +41,13 @@ def MIGraphXRealizeInt4Pass : Pass<"migraphx-realize-int4", "func::FuncOp"> {
   }];
 }
 
+def MIGraphXTosaSimplifyPass : Pass<"migraphx-tosa-simplify", "func::FuncOp"> {
+  let summary = "Simplify TOSA operations after MIGraphX to TOSA conversion";
+  let description = [{
+    This pass simplifies TOSA operations after converting migraphx ops to tosa, optimizing
+    the representation of operations for better performance and clarity.
+  }];
+  let dependentDialects = ["tosa::TosaDialect", "func::FuncDialect"];
+}
+
 #endif // MLIR_DIALECT_MIGRAPHX_PASSES