feat: enable auto-batching passes (#1799)

* feat: enable auto-batching passes

* chore: bump reactant version

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>", "Mosè Giordano <mose@gnu.org>"]
-version = "0.2.172"
+version = "0.2.173"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Compiler.jl

@@ -937,6 +937,7 @@ function optimization_passes(
                 "broadcastindim_slice_to_batch",
                 "reducewindow_slice_to_batch",
                 "elementwise_slice_to_batch",
+                "greedy_while_loop_batch_fission",
             ],
         )
     end

test/batching.jl

@@ -29,3 +29,59 @@ end
 
     @test @jit(f3(A_ra, 1)) ≈ A .+ 1
 end
+
+# Auto-Batching
+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...)
+            @test res1 ≈ res2
+        end
+
+        @testset "No while loops" begin
+            hlo = repr(
+                @code_hlo compile_options = CompileOptions(;
+                    disable_auto_batching_passes=true
+                ) f(args...)
+            )
+            @test occursin("stablehlo.while", hlo)
+
+            hlo = repr(@code_hlo f(args...))
+            @test !occursin("stablehlo.while", hlo)
+        end
+    end
+end
+
+function looped_reduction(y, x)
+    z = copy(y)
+    @trace for i in 1:size(x, 2)
+        z[:, i, :] = dropdims(sum(abs2, x[:, i, :, :]; dims=3); dims=3)
+    end
+    return z
+end
+
+@testset "Loop of Reduces => Single Reduction" begin
+    x = Reactant.to_rarray(rand(Float32, 3, 256, 5, 7))
+    y = Reactant.to_rarray(rand(Float32, 3, 260, 5))
+
+    run_auto_batching_tests(looped_reduction, y, x)
+end
+
+function naive_batched_matmul(x, y)
+    @assert size(x, 3) == size(y, 3)
+    z = similar(x, size(x, 1), size(y, 2), size(x, 3))
+    @trace for i in 1:size(x, 3)
+        z[:, :, i] = x[:, :, i] * y[:, :, i]
+    end
+    return z
+end
+
+@testset "Naive Batched Matmul => Single Dot General" begin
+    x = Reactant.to_rarray(rand(Float32, 3, 256, 5))
+    y = Reactant.to_rarray(rand(Float32, 256, 7, 5))
+
+    run_auto_batching_tests(naive_batched_matmul, x, y)
+end