feat: new passes (#1417)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.138"
+version = "0.2.139"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -90,7 +90,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.210"
+Reactant_jll = "0.0.212"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -887,6 +887,8 @@ function optimization_passes(
         "log_simplify",
         "neg_mul_const_simplify",
         "neg_div_const_simplify",
+        "reshape_deletions_broadcast_in_dim_simplify",
+        "reshape_insertions_broadcast_in_dim_simplify",
     ]
 
     if !compile_options.disable_scatter_gather_optimization_passes
@@ -1092,6 +1094,8 @@ function optimization_passes(
                 "transpose_batch_norm_inference",
                 "transpose_batch_norm_grad",
                 "transpose_if",
+                "transpose_fft",
+                "transpose_reshape",
             ],
         )
         if AGGRESSIVE_PROPAGATION[]
@@ -1128,10 +1132,23 @@ function optimization_passes(
     if compile_options.no_nan
         append!(
             transform_passes_list,
-            ["no_nan", "no_nan_self_sub_simplify", "no_nan_add_sub_simplify(1)"],
+            [
+                "no_nan",
+                "no_nan_self_sub_simplify",
+                "no_nan_add_sub_simplify(1)",
+                "no_nan_mul_simplify(1)",
+                "no_nan_div_simplify(1)",
+            ],
         )
     else
-        push!(transform_passes_list, "no_nan_add_sub_simplify(0)")
+        append!(
+            transform_passes_list,
+            [
+                "no_nan_add_sub_simplify(0)",
+                "no_nan_mul_simplify(0)",
+                "no_nan_div_simplify(0)",
+            ],
+        )
     end
 
     if compile_options.all_finite