From 404fb1b537889f9809acc60a4682b2ccc311e07f Mon Sep 17 00:00:00 2001 From: jumerckx <31353884+jumerckx@users.noreply.github.com> Date: Fri, 3 Jul 2026 05:20:00 -0500 Subject: [PATCH 01/10] add loop_unswitch_threshold and excluded_passes options --- src/CompileOptions.jl | 14 ++++++++++++++ src/compiler/OptimizationPasses.jl | 11 ++++++++++- src/mlir/libMLIR_h.jl | 3 +++ 3 files changed, 27 insertions(+), 1 deletion(-) 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 """ From 5cfb06226e404b763e95ad28d33107cb54056d84 Mon Sep 17 00:00:00 2001 From: jumerckx <31353884+jumerckx@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:30:11 -0500 Subject: [PATCH 02/10] bump ENZYMEXLA_COMMIT --- deps/ReactantExtra/WORKSPACE | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index 360317c044..4f82d1d7d5 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "41b5f3d3ea3ef334d196f965944854d7dcb9b51a" +ENZYMEXLA_COMMIT = "06aa1d6756de89bbed41caa625c6ee21df041885" ENZYMEXLA_SHA256 = "" From c63413b6892b644a122e794f48198560e6d3c5fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Wed, 8 Jul 2026 16:28:05 -0500 Subject: [PATCH 03/10] remove legacy xla/stream_executor/tpu deps for ReactantExtra --- deps/ReactantExtra/BUILD | 3 --- 1 file changed, 3 deletions(-) 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", From c60639a1a88ee7ef9035261a4e3cb3fa4f990526 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 04:14:10 -0500 Subject: [PATCH 04/10] ssh into runner to debug --- .github/workflows/CommonCI.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml index dde0c648a9..6180745263 100644 --- a/.github/workflows/CommonCI.yml +++ b/.github/workflows/CommonCI.yml @@ -229,6 +229,8 @@ jobs: earlyoom -m 3 -s 100 -r 300 --prefer 'julia' & fi julia --project --color=yes --check-bounds=yes -e 'import Pkg; Pkg.test(;coverage=true, julia_args=["--check-bounds=yes", "--compiled-modules=yes", "--depwarn=yes"], test_args=`${{ env.runtest_test_args }}`, force_latest_compatible_version=false, allow_reresolve=${{ !(inputs.downgrade_testing) }})' + - name: "Setup tmate session" + uses: mxschmitt/action-tmate@v3 - name: "Upload MLIR modules" uses: actions/upload-artifact@v7 From 9965d671c4dacb1c8bad134105ab02eb063cac2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 04:38:16 -0500 Subject: [PATCH 05/10] put tmate job before failing step --- .github/workflows/CommonCI.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml index 6180745263..b568bfed10 100644 --- a/.github/workflows/CommonCI.yml +++ b/.github/workflows/CommonCI.yml @@ -219,6 +219,10 @@ jobs: "Reactant", "xla_runtime" => "IFRT"; force=true ) shell: julia --color=yes --project=. {0} + + - name: "Setup tmate session" + uses: mxschmitt/action-tmate@v3 + - name: "Run Tests (IFRT)" if: ${{ inputs.runtime == 'ifrt' || inputs.runtime == 'both' }} timeout-minutes: 120 @@ -229,8 +233,6 @@ jobs: earlyoom -m 3 -s 100 -r 300 --prefer 'julia' & fi julia --project --color=yes --check-bounds=yes -e 'import Pkg; Pkg.test(;coverage=true, julia_args=["--check-bounds=yes", "--compiled-modules=yes", "--depwarn=yes"], test_args=`${{ env.runtest_test_args }}`, force_latest_compatible_version=false, allow_reresolve=${{ !(inputs.downgrade_testing) }})' - - name: "Setup tmate session" - uses: mxschmitt/action-tmate@v3 - name: "Upload MLIR modules" uses: actions/upload-artifact@v7 From 74abb305b1cb9f925243ad8539db04ae4dadf8b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 06:29:16 -0500 Subject: [PATCH 06/10] temporarily use debug build --- .github/workflows/CommonCI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml index b568bfed10..fd55a9cc36 100644 --- a/.github/workflows/CommonCI.yml +++ b/.github/workflows/CommonCI.yml @@ -124,7 +124,7 @@ jobs: if: ${{ inputs.localjll || inputs.enzymejax_commit != '' }} run: | julia --color=yes --project=deps -e 'using Pkg; Pkg.instantiate()' - julia --color=yes --project=deps deps/build_local.jl --cc clang --push-cache + julia --color=yes --project=deps deps/build_local.jl --cc clang --push-cache --debug cp LocalPreferences.toml test/ # Compile Julia if assertions are enabled From 8bdd854c938407e7fd0334e6abec22017693a552 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 09:07:33 -0500 Subject: [PATCH 07/10] disable tmate and debug build --- .github/workflows/CommonCI.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml index fd55a9cc36..fb6eab4475 100644 --- a/.github/workflows/CommonCI.yml +++ b/.github/workflows/CommonCI.yml @@ -124,7 +124,7 @@ jobs: if: ${{ inputs.localjll || inputs.enzymejax_commit != '' }} run: | julia --color=yes --project=deps -e 'using Pkg; Pkg.instantiate()' - julia --color=yes --project=deps deps/build_local.jl --cc clang --push-cache --debug + julia --color=yes --project=deps deps/build_local.jl --cc clang --push-cache cp LocalPreferences.toml test/ # Compile Julia if assertions are enabled @@ -220,8 +220,8 @@ jobs: ) shell: julia --color=yes --project=. {0} - - name: "Setup tmate session" - uses: mxschmitt/action-tmate@v3 + # - name: "Setup tmate session" + # uses: mxschmitt/action-tmate@v3 - name: "Run Tests (IFRT)" if: ${{ inputs.runtime == 'ifrt' || inputs.runtime == 'both' }} From 71c93520020d7fe2dbb746b05cc712e071bdac9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 09:59:55 -0500 Subject: [PATCH 08/10] clean workflow --- .github/workflows/CommonCI.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml index fb6eab4475..5391672b2d 100644 --- a/.github/workflows/CommonCI.yml +++ b/.github/workflows/CommonCI.yml @@ -220,9 +220,6 @@ jobs: ) shell: julia --color=yes --project=. {0} - # - name: "Setup tmate session" - # uses: mxschmitt/action-tmate@v3 - - name: "Run Tests (IFRT)" if: ${{ inputs.runtime == 'ifrt' || inputs.runtime == 'both' }} timeout-minutes: 120 From c4a4f2deb202514c8a24535c5f7a71f19b5f408c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 22:52:18 +0200 Subject: [PATCH 09/10] remove line --- .github/workflows/CommonCI.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml index 5391672b2d..dde0c648a9 100644 --- a/.github/workflows/CommonCI.yml +++ b/.github/workflows/CommonCI.yml @@ -219,7 +219,6 @@ jobs: "Reactant", "xla_runtime" => "IFRT"; force=true ) shell: julia --color=yes --project=. {0} - - name: "Run Tests (IFRT)" if: ${{ inputs.runtime == 'ifrt' || inputs.runtime == 'both' }} timeout-minutes: 120 From 9648f7f8e39fcadecb59d0a4fcb0fa9259936b6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Thu, 9 Jul 2026 15:51:49 -0500 Subject: [PATCH 10/10] bump Enzyme-JAX commit --- deps/ReactantExtra/WORKSPACE | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index 4f82d1d7d5..41e7e0185d 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "06aa1d6756de89bbed41caa625c6ee21df041885" +ENZYMEXLA_COMMIT = "36b215f2a65e5722ce2262a32ffee61cba295c3a" ENZYMEXLA_SHA256 = ""