adapt to XLA changes

deps/ReactantExtra/BUILD

@@ -1326,9 +1326,6 @@ cc_library(
         "@xla//xla/service:metrics_proto_cc_impl",
         "@xla//xla/service:custom_call_target_registry",
         "@xla//xla/service/cpu:cpu_compiler",
-        "@xla//xla/stream_executor/tpu:tpu_on_demand_compiler",
-        "@xla//xla/stream_executor/tpu:tpu_executor",
-        "@xla//xla/stream_executor/tpu:tpu_transfer_manager",
         "@xla//xla/service/cpu:cpu_transfer_manager",
         "@xla//xla/tsl/protobuf:protos_all_cc_impl",
         "@xla//xla/tsl/framework:allocator_registry_impl",

deps/ReactantExtra/WORKSPACE

@@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023"
 
 NSYNC_SHA256 = ""
 
-ENZYMEXLA_COMMIT = "41b5f3d3ea3ef334d196f965944854d7dcb9b51a"
+ENZYMEXLA_COMMIT = "36b215f2a65e5722ce2262a32ffee61cba295c3a"
 
 ENZYMEXLA_SHA256 = ""
 

src/CompileOptions.jl

@@ -193,6 +193,10 @@ 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. (Default: `false`).
+  - `loop_unswitch_threshold`: Threshold passed to the `loop_unswitch` pattern inside the
+    LICM pass group. A negative value disables the pattern entirely. (Default: `10`).
+  - `excluded_passes`: A list of pass base-names (e.g. `["loop_unswitch"]`) that should be
+    omitted from the generated pass pipeline. (Default: `String[]`).
   - `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.
@@ -241,6 +245,8 @@ struct CompileOptions
     disable_scatter_gather_optimization_passes::Bool
     disable_pad_optimization_passes::Bool
     disable_licm_optimization_passes::Bool
+    loop_unswitch_threshold::Int
+    excluded_passes::Vector{String}
     disable_reduce_slice_fusion_passes::Bool
     disable_slice_to_batch_passes::Bool
     disable_concat_to_batch_passes::Bool
@@ -276,6 +282,8 @@ function CompileOptions(;
     disable_scatter_gather_optimization_passes::Bool=false,
     disable_pad_optimization_passes::Bool=false,
     disable_licm_optimization_passes::Bool=false,
+    loop_unswitch_threshold::Int=10,
+    excluded_passes::Vector{String}=String[],
     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,
@@ -342,6 +350,8 @@ function CompileOptions(;
         disable_scatter_gather_optimization_passes,
         disable_pad_optimization_passes,
         disable_licm_optimization_passes,
+        loop_unswitch_threshold,
+        excluded_passes,
         disable_reduce_slice_fusion_passes,
         disable_slice_to_batch_passes,
         disable_concat_to_batch_passes,
@@ -397,6 +407,8 @@ 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.loop_unswitch_threshold,
+        compile_options.excluded_passes,
         compile_options.disable_reduce_slice_fusion_passes,
         compile_options.disable_slice_to_batch_passes,
         compile_options.disable_concat_to_batch_passes,
@@ -439,6 +451,8 @@ function __compile_options_with_updated_sync(compile_options::CompileOptions, sy
         compile_options.disable_scatter_gather_optimization_passes,
         compile_options.disable_pad_optimization_passes,
         compile_options.disable_licm_optimization_passes,
+        compile_options.loop_unswitch_threshold,
+        compile_options.excluded_passes,
         compile_options.disable_reduce_slice_fusion_passes,
         compile_options.disable_slice_to_batch_passes,
         compile_options.disable_concat_to_batch_passes,

src/compiler/OptimizationPasses.jl

@@ -43,6 +43,12 @@ function optimization_passes(
     raise_shlo_to_blas_lapack::Bool=true,
     self_to_convolution::Bool=false,
 )
+    # Build the excluded-passes pointer array. Julia Strings are null-terminated,
+    # so unsafe_convert(Cstring, s) gives a stable C pointer while the String is live.
+    excl_strs = compile_options.excluded_passes
+    excl_cstrs = Cstring[Base.unsafe_convert(Cstring, s) for s in excl_strs]
+    excl_ptr = isempty(excl_cstrs) ? Ptr{Cstring}(C_NULL) : pointer(excl_cstrs)
+
     options = MLIR.API.EnzymeXLATransformPassesOptions(
         compile_options.max_constant_threshold,          # max_constant_threshold
         WHILE_UNROLL_THRESHOLD[],                        # while_unroll_threshold
@@ -70,12 +76,15 @@ function optimization_passes(
         !compile_options.disable_concat_to_batch_passes,             # enable_concat_to_batch_passes
         !compile_options.disable_loop_raising_passes,                # enable_loop_raising_passes
         !compile_options.disable_licm_optimization_passes,           # enable_licm_optimization_passes
+        Int64(compile_options.loop_unswitch_threshold),              # loop_unswitch_threshold
         !compile_options.disable_pad_optimization_passes,            # enable_pad_optimization_passes
+        excl_ptr,                                                    # excluded_passes
+        Csize_t(length(excl_cstrs)),                                 # num_excluded_passes
     )
 
     main_passes_ptr = Ref{Cstring}()
     lower_passes_ptr = Ref{Cstring}()
-    GC.@preserve options begin
+    GC.@preserve options excl_strs excl_cstrs begin
         MLIR.API.enzymexlaGetTransformPassesList(
             Ref(options), main_passes_ptr, lower_passes_ptr
         )

src/mlir/libMLIR_h.jl

@@ -14112,7 +14112,10 @@ struct EnzymeXLATransformPassesOptions
     enable_concat_to_batch_passes::Bool
     enable_loop_raising_passes::Bool
     enable_licm_optimization_passes::Bool
+    loop_unswitch_threshold::Int64
     enable_pad_optimization_passes::Bool
+    excluded_passes::Ptr{Cstring}
+    num_excluded_passes::Csize_t
 end
 
 """