fix: enable batching except slice to batch passes (#1985)

* fix: enable batching by default

* feat: finegrained control of options

Author: avik-pal

docs/src/tutorials/raising.md

@@ -123,7 +123,7 @@ raising).
 
 ```@example raising_stablehlo
 @code_hlo compile_options = CompileOptions(;
-    disable_auto_batching_passes=true
+    disable_loop_raising_passes=true
 ) compute_attractive_force(positions_ra, masses_ra, 2.0f0)
 ```
 
@@ -133,7 +133,7 @@ tensor IR.
 
 ```@example raising_stablehlo
 hlo = @code_hlo compile_options=CompileOptions(;
-    disable_auto_batching_passes=false
+    disable_loop_raising_passes=false
 ) compute_attractive_force(positions_ra, masses_ra, 2.0f0)
 @assert !contains(repr(hlo), "stablehlo.while") #hide
 hlo
@@ -145,7 +145,7 @@ the values are identical.
 ```@example raising_stablehlo
 y_jl = compute_attractive_force(positions, masses, 2.0f0)
 y_ra = @jit compile_options=CompileOptions(;
-    disable_auto_batching_passes=false
+    disable_loop_raising_passes=false
 ) compute_attractive_force(positions_ra, masses_ra, 2.0f0)
 maximum(abs, Array(y_ra) .- y_jl)
 ```
@@ -154,7 +154,7 @@ Let's time the execution of the two versions.
 
 ```@example raising_stablehlo
 fn1 = @compile sync=true compile_options=CompileOptions(;
-    disable_auto_batching_passes=true
+    disable_loop_raising_passes=true
 ) compute_attractive_force(positions_ra, masses_ra, 2.0f0)
 fn2 = @compile sync=true compute_attractive_force(positions_ra, masses_ra, 2.0f0)
 ```

src/CompileOptions.jl

@@ -94,7 +94,7 @@ Fine-grained control over the compilation options for the Reactant compiler.
     the computation graph. If `:down`, they will be propagated down. Defaults to `:up`.
   - `max_constant_threshold`: If the number of elements in a constant is greater than this
     threshold (for a non-splatted constant), we will throw an error.
-  - `inline`: If `true`, all functions will be inlined. This is `true` by default.
+  - `inline`: If `true`, all functions will be inlined. (Default: `true`).
 
 ## Raising Options
 
@@ -107,7 +107,7 @@ Fine-grained control over the compilation options for the Reactant compiler.
 ## Dialect Specific Options
 
   - `legalize_chlo_to_stablehlo`: If `true`, `chlo` dialect ops will be converted to
-    `stablehlo` ops. This is `false` by default.
+    `stablehlo` ops. (Default: `false`).
 
 ## Backend Specific Options
 
@@ -153,13 +153,21 @@ Fine-grained control over the compilation options for the Reactant compiler.
     notice or deprecation cycle.
 
   - `disable_scatter_gather_optimization_passes`: Disables the scatter-gather
-    optimization passes. This is `false` by default.
+    optimization passes. (Default: `false`).
   - `disable_pad_optimization_passes`: Disables the pad optimization passes. This is
     `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.
+    optimization passes. (Default: `false`).
+  - `disable_reduce_slice_fusion_passes`: Disables fusion of slice elementwise and reduce
+    operations. (Default `false`).
+  - `disable_slice_to_batch_passes`: Disables the slice to batch fusion optimization passes.
+    (Default: `true`). _(Note that this is generally an expensive pass to run)_
+  - `disable_concat_to_batch_passes`: Disables concatenate to batch fusion passes.
+    (Default: `false`).
+  - `disable_loop_raising_passes`: Disables raising passes for `stablehlo.while`.
+    (Default: `false`).
+  - `disable_structured_tensors_passes`: Disables structured tensors detection and
+    propagation passes. (Default `false`).
 """
 struct CompileOptions
     optimization_passes::Union{Symbol,String}
@@ -188,7 +196,11 @@ struct CompileOptions
     disable_scatter_gather_optimization_passes::Bool
     disable_pad_optimization_passes::Bool
     disable_licm_optimization_passes::Bool
-    disable_auto_batching_passes::Bool
+    disable_reduce_slice_fusion_passes::Bool
+    disable_slice_to_batch_passes::Bool
+    disable_concat_to_batch_passes::Bool
+    disable_loop_raising_passes::Bool
+    disable_structured_tensors_passes::Bool
 end
 
 function CompileOptions(;
@@ -212,7 +224,11 @@ 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=true,
+    disable_reduce_slice_fusion_passes::Bool=false,
+    disable_slice_to_batch_passes::Bool=true, # expensive + introduces all-to-all in GB25
+    disable_concat_to_batch_passes::Bool=false,
+    disable_loop_raising_passes::Bool=false,
+    disable_structured_tensors_passes::Bool=false,
 )
     optimization_passes isa Bool &&
         (optimization_passes = ifelse(optimization_passes, :all, :none))
@@ -261,7 +277,11 @@ function CompileOptions(;
         disable_scatter_gather_optimization_passes,
         disable_pad_optimization_passes,
         disable_licm_optimization_passes,
-        disable_auto_batching_passes,
+        disable_reduce_slice_fusion_passes,
+        disable_slice_to_batch_passes,
+        disable_concat_to_batch_passes,
+        disable_loop_raising_passes,
+        disable_structured_tensors_passes,
     )
 end
 
@@ -303,7 +323,11 @@ function __compile_options_with_reversed_propagation(compile_options::CompileOpt
         compile_options.disable_scatter_gather_optimization_passes,
         compile_options.disable_pad_optimization_passes,
         compile_options.disable_licm_optimization_passes,
-        compile_options.disable_auto_batching_passes,
+        compile_options.disable_reduce_slice_fusion_passes,
+        compile_options.disable_slice_to_batch_passes,
+        compile_options.disable_concat_to_batch_passes,
+        compile_options.disable_loop_raising_passes,
+        compile_options.disable_structured_tensors_passes,
     )
 end
 

src/Compiler.jl

@@ -926,18 +926,49 @@ function optimization_passes(
     if !is_sharded
         # these passes don't have optimized sharding implementations
         if raise_shlo_to_blas_lapack
-            append!(transform_passes_list, ["dot_general_to_syrk"])
+            if !compile_options.disable_structured_tensors_passes
+                append!(transform_passes_list, ["dot_general_to_syrk"])
+            end
         end
     end
 
-    if !compile_options.disable_auto_batching_passes
+    if !compile_options.disable_slice_to_batch_passes
+        append!(
+            transform_passes_list,
+            [
+                "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",
+                "convolution_slice_to_batch",
+            ],
+        )
+    end
+
+    if !compile_options.disable_reduce_slice_fusion_passes
         append!(
             transform_passes_list,
             [
                 "add_reduce_slice_fusion",
                 "mul_reduce_slice_fusion",
                 "min_reduce_slice_fusion",
                 "max_reduce_slice_fusion",
+                "and_reduce_slice_fusion",
+                "xor_reduce_slice_fusion",
+                "or_reduce_slice_fusion",
+            ],
+        )
+    end
+
+    if !compile_options.disable_concat_to_batch_passes
+        append!(
+            transform_passes_list,
+            [
                 "concat_insert_dim_dot_general",
                 "concat_insert_dim_gather",
                 "concat_insert_dim_iota",
@@ -946,21 +977,14 @@ function optimization_passes(
                 "concat_insert_dim_reduce_window",
                 "concat_insert_dim_elementwise",
                 "concat_insert_dim_convolution",
-                "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",
-                "convolution_slice_to_batch",
-                "greedy_while_loop_batch_fission",
             ],
         )
     end
 
+    if !compile_options.disable_loop_raising_passes
+        append!(transform_passes_list, ["greedy_while_loop_batch_fission"])
+    end
+
     if !compile_options.disable_licm_optimization_passes
         append!(
             transform_passes_list,

test/batching.jl

@@ -35,23 +35,23 @@ function run_auto_batching_tests(f::F, args...) where {F}
     @testset "$(nameof(F))" begin
         @testset "Correctness" begin
             res1 = @jit f(args...)
-            res2 = @jit compile_options = CompileOptions(;
-                disable_auto_batching_passes=true
-            ) f(args...)
+            res2 = @jit compile_options = CompileOptions(; disable_loop_raising_passes=true) f(
+                args...
+            )
             @test res1 ≈ res2
         end
 
         @testset "No while loops" begin
             hlo = repr(
                 @code_hlo compile_options = CompileOptions(;
-                    disable_auto_batching_passes=true
+                    disable_loop_raising_passes=true
                 ) f(args...)
             )
             @test occursin("stablehlo.while", hlo)
 
             hlo = repr(
                 @code_hlo compile_options = CompileOptions(;
-                    disable_auto_batching_passes=false
+                    disable_loop_raising_passes=false
                 ) f(args...)
             )
             @test !occursin("stablehlo.while", hlo)
@@ -119,11 +119,11 @@ end
     input1 = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float32, 10))
     input2 = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float32, 10))
 
-    hlo = @code_hlo compile_options = CompileOptions(; disable_auto_batching_passes=true) mctr(
+    hlo = @code_hlo compile_options = CompileOptions(; disable_loop_raising_passes=true) mctr(
         map_with_scalar_indexing, 1:8, input1, input2
     )
     @test contains(repr(hlo), "stablehlo.while")
-    hlo = @code_hlo compile_options = CompileOptions(; disable_auto_batching_passes=false) mctr(
+    hlo = @code_hlo compile_options = CompileOptions(; disable_loop_raising_passes=false) mctr(
         map_with_scalar_indexing, 1:8, input1, input2
     )
     @test !contains(repr(hlo), "stablehlo.while")