feat: use the generalized passes from upcoming enzymejax (#1850)

* feat: use the generalized passes from upcoming enzymejax [skip ci]

* chore: bump jll version

* test: update

* feat: new passes

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

src/CompileOptions.jl

@@ -12,10 +12,11 @@ communication.
     wrap_comm::Int = 0
     extend_comm::Int = 0
     dus_to_pad_manual_comp_comm::Int = 0 # 2
-    dus_to_pad_comm::Int = 1
+    dus_to_pad_comm::Int = 0
     concat_two_operands_comm::Int = 0
     concat_to_pad_comm::Int = 1
-    extend_to_pad_comm::Int = 1
+    extend_to_pad_comm::Int = 0
+    extend_to_pad_comm2::Int = 1
     wrap_to_pad_comm::Int = 1
 end
 

src/Compiler.jl

@@ -759,6 +759,9 @@ function optimization_passes(
         "noop_slice<16>",
         "noop_reverse<16>",
         "slice_slice<16>",
+        "dynamic_slice_slice<16>",
+        "slice_dynamic_slice<16>",
+        "dynamic_slice_dynamic_slice<16>",
         "shift_right_logical_simplify<16>",
         "slice_simplify<16>",
         "convert_simplify<16>",
@@ -816,20 +819,7 @@ function optimization_passes(
         "dus_dus",
         "dus_dus_concat",
         "abs_positive_simplify",
-        "transpose_unary_transpose_abs",
-        "transpose_unary_transpose_neg",
-        "transpose_unary_transpose_sqrt",
-        "transpose_unary_transpose_rsqrt",
-        "transpose_unary_transpose_ceil",
-        "transpose_unary_transpose_convert",
-        "transpose_unary_transpose_cosine",
-        "transpose_unary_transpose_exp",
-        "transpose_unary_transpose_expm1",
-        "transpose_unary_transpose_log",
-        "transpose_unary_transpose_log1p",
-        "transpose_unary_transpose_sign",
-        "transpose_unary_transpose_sine",
-        "transpose_unary_transpose_tanh",
+        "transpose_elementwise_transpose",
         "select_comp_iota_const_simplify<1>",
         "sign_abs_simplify<1>",
         "broadcastindim_is_reshape",
@@ -1147,6 +1137,9 @@ function optimization_passes(
                 "concat_appending_reshape",
                 "slice_reshape",
                 "slice_reshape_slice<1>",
+                "dynamic_slice_reshape_slice<1>",
+                "slice_reshape_dynamic_slice<1>",
+                "dynamic_slice_reshape_dynamic_slice<1>",
                 "slice_reshape_concat<1>",
                 "slice_reshape_elementwise<1>",
                 "slice_reshape_dot_general<1>",
@@ -1193,17 +1186,7 @@ function optimization_passes(
             transform_passes_list,
             [
                 "reorder_elementwise_and_shape_op<16>",
-                "binary_op_transpose_simplify_add",
-                "binary_op_transpose_simplify_sub",
-                "binary_op_transpose_simplify_mul",
-                "binary_op_transpose_simplify_div",
-                "binary_op_transpose_simplify_min",
-                "binary_op_transpose_simplify_max",
-                "binary_op_transpose_simplify_pow",
-                "binary_op_transpose_simplify_rem",
-                "binary_op_transpose_simplify_or",
-                "binary_op_transpose_simplify_and",
-                "binary_op_transpose_simplify_xor",
+                "elementwise_all_transpose_operands_simplify",
                 "slice_transpose",
                 "einsum_transpose<1>",
                 "slice_reshape_transpose<1>",

test/optimize_comm.jl

@@ -100,10 +100,8 @@ if length(addressable_devices) ≥ 8
 
         hlo = repr(@code_xla shardy_passes = :to_mhlo_shardings dus2(rx, ry))
         @test !contains(hlo, "all-to-all")
-        @test !contains(hlo, "all-gather") broken =
-            Reactant.XLA.REACTANT_XLA_RUNTIME == "PJRT"
-        @test contains(hlo, "collective-permute") broken =
-            Reactant.XLA.REACTANT_XLA_RUNTIME == "PJRT"
+        @test !contains(hlo, "all-gather")
+        @test contains(hlo, "collective-permute")
 
         dus2(x, y)
         @jit shardy_passes = :to_mhlo_shardings dus2(rx, ry)