diff --git a/deps/ReactantExtra/BUILD b/deps/ReactantExtra/BUILD index 7d4ba563c5..d7c28cb7c2 100644 --- a/deps/ReactantExtra/BUILD +++ b/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", diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index 360317c044..41e7e0185d 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "41b5f3d3ea3ef334d196f965944854d7dcb9b51a" +ENZYMEXLA_COMMIT = "36b215f2a65e5722ce2262a32ffee61cba295c3a" ENZYMEXLA_SHA256 = "" diff --git a/src/CompileOptions.jl b/src/CompileOptions.jl index a4806185eb..045b732c5b 100644 --- a/src/CompileOptions.jl +++ b/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, diff --git a/src/compiler/OptimizationPasses.jl b/src/compiler/OptimizationPasses.jl index 3f7a275466..bd4f7abce1 100644 --- a/src/compiler/OptimizationPasses.jl +++ b/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 ) diff --git a/src/mlir/libMLIR_h.jl b/src/mlir/libMLIR_h.jl index 88c56b79f4..aff7a0cee5 100755 --- a/src/mlir/libMLIR_h.jl +++ b/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 """