feat: enable new auto-batching passes

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.164"
+version = "0.2.165"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -103,7 +103,7 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.16"
-Reactant_jll = "0.0.240"
+Reactant_jll = "0.0.241"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/CompileOptions.jl

@@ -157,6 +157,8 @@ Fine-grained control over the compilation options for the Reactant compiler.
     `false` by default.
   - `disable_licm_optimization_passes`: Disables the Loop Invariant Code Motion (LICM)
     optimization passes. This is `false` by default.
+  - `disable_auto_batching_passes`: Disables the auto-batching optimization passes. This
+    is `false` by default.
 """
 struct CompileOptions
     optimization_passes::Union{Symbol,String}
@@ -185,6 +187,7 @@ struct CompileOptions
     disable_scatter_gather_optimization_passes::Bool
     disable_pad_optimization_passes::Bool
     disable_licm_optimization_passes::Bool
+    disable_auto_batching_passes::Bool
 end
 
 function CompileOptions(;
@@ -208,6 +211,7 @@ function CompileOptions(;
     disable_scatter_gather_optimization_passes::Bool=false,
     disable_pad_optimization_passes::Bool=false,
     disable_licm_optimization_passes::Bool=false,
+    disable_auto_batching_passes::Bool=false,
 )
     optimization_passes isa Bool &&
         (optimization_passes = ifelse(optimization_passes, :all, :none))
@@ -256,6 +260,7 @@ function CompileOptions(;
         disable_scatter_gather_optimization_passes,
         disable_pad_optimization_passes,
         disable_licm_optimization_passes,
+        disable_auto_batching_passes,
     )
 end
 

src/Compiler.jl

@@ -903,14 +903,36 @@ function optimization_passes(
         "self_add_to_convolution_like($(Int(backend == "tpu")))",
         "self_mul_to_convolution_like($(Int(backend == "tpu")))",
         "subtract_multiply_const_to_add_mul_const",
-        "concat_insert_dim_dot_general",
-        "concat_insert_dim_gather",
-        "concat_insert_dim_iota",
-        "concat_insert_dim_reduce",
-        "concat_insert_dim_sort",
-        "concat_insert_dim_reduce_window",
     ]
 
+    if !compile_options.disable_auto_batching_passes
+        append!(
+            transform_passes_list,
+            [
+                "trivial_reduce_window_to_reduce_op",
+                "add_reduce_slice_fusion",
+                "mul_reduce_slice_fusion",
+                "min_reduce_slice_fusion",
+                "max_reduce_slice_fusion",
+                "concat_insert_dim_dot_general",
+                "concat_insert_dim_gather",
+                "concat_insert_dim_iota",
+                "concat_insert_dim_reduce",
+                "concat_insert_dim_sort",
+                "concat_insert_dim_reduce_window",
+                "dot_general_slice_to_batch",
+                "gather_slice_to_batch",
+                "iota_slice_to_batch",
+                "reduce_slice_to_batch",
+                "sort_slice_to_batch",
+                "transpose_slice_to_batch",
+                "broadcastindim_slice_to_batch",
+                "reducewindow_slice_to_batch",
+                "elementwise_slice_to_batch",
+            ],
+        )
+    end
+
     if !compile_options.disable_licm_optimization_passes
         append!(
             transform_passes_list,