From 4385eb303924462c646ec911526383f9e703cd2e Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Fri, 17 Oct 2025 01:08:12 -0500 Subject: [PATCH 1/4] Try enzyme-jax bump --- deps/ReactantExtra/WORKSPACE | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index 8df2ae5292..6185f5c88e 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -1,11 +1,9 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" - NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "996b3e5ccbe5024d69d9a365e5f4db79f41dbc6c" - +ENZYMEXLA_COMMIT = "fc38e3809323c70f425dfb3c61ac55d1abee595a" ENZYMEXLA_SHA256 = "" http_archive( @@ -73,19 +71,9 @@ LLVM_TARGETS = [ # ], # ) -http_archive( - name = "rules_ml_toolchain", - patch_cmds = [ - """ - sed -i.bak0 "/D_FORTIFY_SOURCE/d" cc/features/BUILD third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl -""", - ], - sha256 = "e7e44c4e349a1c1f31398bd2257c51432e73ea0e7e24cce67090b68b0b50007e", - strip_prefix = "rules_ml_toolchain-55dcd0a52c7e0f9eec9927a32512229c09ac3b3e", - urls = [ - "https://github.com/google-ml-infra/rules_ml_toolchain/archive/55dcd0a52c7e0f9eec9927a32512229c09ac3b3e.tar.gz", - ], -) +load("@enzyme_ad//third_party/ml_toolchain:workspace.bzl", ml_toolchain_workspace = "repo") + +ml_toolchain_workspace() load("@enzyme_ad//third_party/jax:workspace.bzl", jax_workspace = "repo") From 3eba91927f6e3dd1b4ca5364219fac5f0a34fefd Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Fri, 17 Oct 2025 01:21:07 -0500 Subject: [PATCH 2/4] fix --- deps/ReactantExtra/WORKSPACE | 2 +- src/Ops.jl | 2 +- test/control_flow.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index 6185f5c88e..1743e1dbf2 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -3,7 +3,7 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "fc38e3809323c70f425dfb3c61ac55d1abee595a" +ENZYMEXLA_COMMIT = "96bcf384c5dce3341b347554c32bd02b2defc0eb" ENZYMEXLA_SHA256 = "" http_archive( diff --git a/src/Ops.jl b/src/Ops.jl index 39c9790dea..252d87dc84 100644 --- a/src/Ops.jl +++ b/src/Ops.jl @@ -2149,7 +2149,7 @@ end ) if !mincut - MLIR.IR.attr!(while_op, "enzymexla.disable_min_cut", MLIR.IR.UnitAttribute()) + MLIR.IR.attr!(while_op, "enzyme.disable_mincut", MLIR.IR.UnitAttribute()) end if checkpointing diff --git a/test/control_flow.jl b/test/control_flow.jl index 4657f588e9..4442ebe805 100644 --- a/test/control_flow.jl +++ b/test/control_flow.jl @@ -706,7 +706,7 @@ end @test for_no_track_numbers_ra(x_ra, n_ra) == for_no_track_numbers(x, n) ir = sprint(show, @code_hlo optimize = "enzyme-batch" for_no_track_numbers(x_ra, n_ra)) - @test contains(ir, "enzymexla.disable_min_cut") + @test contains(ir, "enzyme.disable_mincut") @test contains(ir, "enzymexla.enable_checkpointing") end From 771ea730ea31d964272ba2ef8d8079df4e89b72a Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Fri, 17 Oct 2025 04:47:42 -0500 Subject: [PATCH 3/4] fix --- deps/ReactantExtra/.bazelrc | 44 ++++++++++++++++++------------------ deps/ReactantExtra/API.cpp | 9 ++++---- deps/ReactantExtra/WORKSPACE | 8 +++++-- 3 files changed, 33 insertions(+), 28 deletions(-) diff --git a/deps/ReactantExtra/.bazelrc b/deps/ReactantExtra/.bazelrc index 0358af1af6..cb98372303 100644 --- a/deps/ReactantExtra/.bazelrc +++ b/deps/ReactantExtra/.bazelrc @@ -29,32 +29,32 @@ build --repo_env=RULES_PYTHON_ENABLE_PYSTAR=0 build -c opt -build:cuda --repo_env TF_NEED_CUDA=1 -build:cuda --repo_env TF_NVCC_CLANG=1 -build:cuda --repo_env TF_NCCL_USE_STUB=1 -build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.8.1" -build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.8.0" -build:cuda --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" +common:cuda --repo_env TF_NEED_CUDA=1 +common:cuda --repo_env TF_NVCC_CLANG=1 +common:cuda --repo_env TF_NCCL_USE_STUB=1 +common:cuda --repo_env=HERMETIC_CUDA_VERSION="12.8.1" +common:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.8.0" +common:cuda --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" # "sm" means we emit only cubin, which is forward compatible within a GPU generation. # "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations. -build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,compute_90" -build:cuda --crosstool_top="@local_config_cuda//crosstool:toolchain" -build:cuda --@local_config_cuda//:enable_cuda +common:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,compute_90" +common:cuda --crosstool_top="@local_config_cuda//crosstool:toolchain" +common:cuda --@local_config_cuda//:enable_cuda # Default hermetic CUDA and CUDNN versions. -build:cuda --@local_config_cuda//cuda:include_cuda_libs=true -build:cuda --@local_config_cuda//:cuda_compiler=nvcc -# build:cuda --@local_config_nvshmem//:override_include_nvshmem_libs=true -# build:cuda --@local_config_nvshmem//cuda:include_nvshmem_libs=true +common:cuda --@local_config_cuda//cuda:include_cuda_libs=true +common:cuda --@local_config_cuda//:cuda_compiler=nvcc +# common:cuda --@local_config_nvshmem//:override_include_nvshmem_libs=true +# common:cuda --@local_config_nvshmem//cuda:include_nvshmem_libs=true -build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain -build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true -build:rocm --repo_env TF_NEED_ROCM=1 -build:rocm --action_env TF_ROCM_AMDGPU_TARGETS="gfx900,gfx906,gfx908,gfx90a,gfx940,gfx941,gfx942,gfx1030,gfx1100,gfx1200,gfx1201" +common:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain +common:rocm --define=using_rocm=true --define=using_rocm_hipcc=true +common:rocm --repo_env TF_NEED_ROCM=1 +common:rocm --action_env TF_ROCM_AMDGPU_TARGETS="gfx900,gfx906,gfx908,gfx90a,gfx940,gfx941,gfx942,gfx1030,gfx1100,gfx1200,gfx1201" # Build with hipcc for ROCm and clang for the host. -build:rocm --action_env=TF_ROCM_CLANG="1" -build:rocm --action_env=CLANG_COMPILER_PATH="/usr/lib/llvm-18/bin/clang" -build:rocm --copt=-Wno-gnu-offsetof-extensions -build:rocm --copt=-Qunused-arguments -build:rocm --action_env=TF_HIPCC_CLANG="1" +common:rocm --action_env=TF_ROCM_CLANG="1" +common:rocm --action_env=CLANG_COMPILER_PATH="/usr/lib/llvm-18/bin/clang" +common:rocm --copt=-Wno-gnu-offsetof-extensions +common:rocm --copt=-Qunused-arguments +common:rocm --action_env=TF_HIPCC_CLANG="1" diff --git a/deps/ReactantExtra/API.cpp b/deps/ReactantExtra/API.cpp index 80d43090a5..3687056788 100644 --- a/deps/ReactantExtra/API.cpp +++ b/deps/ReactantExtra/API.cpp @@ -1626,8 +1626,9 @@ REACTANT_ABI HeldIfrtArray *ifrt_client_assemble_array_from_single_shards( REACTANT_ABI HeldIfrtArray * ifrt_pjrt_array_create(ifrt::PjRtClient *client, HeldValue> *buffer) { - return reactant::capture(tsl::RCReference( - MyValueOrThrow(xla::ifrt::PjRtArray::Create(client, buffer->obj())))); + return reactant::capture( + tsl::RCReference(MyValueOrThrow(xla::ifrt::PjRtArray::Create( + client, buffer->obj(), /*has_custom_layout*/ false)))); } // we might me interested in the `Compiler::Compile` method variant that accepts @@ -2373,7 +2374,7 @@ REACTANT_ABI bool hlo_sharding_check_eq(xla::HloSharding *hloSharding, #pragma endregion -typedef ifrt::Future<> IfRtFutureType; +typedef tsl::Future<> IfRtFutureType; REACTANT_ABI void ifrt_free_future(IfRtFutureType *Future) { delete Future; } @@ -2600,7 +2601,7 @@ REACTANT_ABI void ifrt_loaded_executable_execute( // there is only 1 status and is valid because we set `options.fill_status = // true` *futures = true; - *status = new FutureType(result.status); + *status = new IfRtFutureType(result.status); for (int i = 0; i < num_results; i++) { op_results[i] = reactant::capture(result.outputs[i]); diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index 1743e1dbf2..bb561f7815 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -3,7 +3,7 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "96bcf384c5dce3341b347554c32bd02b2defc0eb" +ENZYMEXLA_COMMIT = "f46d7016fbc50dc87125bbdfddf71ea6782073fd" ENZYMEXLA_SHA256 = "" http_archive( @@ -34,7 +34,11 @@ LLVM_TARGETS = [ "AArch64", "X86", "ARM", -] +] +#+ [ +# "PowerPC", +# "SystemZ" +#] # Uncomment these lines to use a custom LLVM commit # LLVM_COMMIT = "b39c5cb6977f35ad727d86b2dd6232099734ffd3" From e193eba386bc0a7af45ca65f30d6aeb29f9f0741 Mon Sep 17 00:00:00 2001 From: William Moses Date: Fri, 17 Oct 2025 19:02:38 +0800 Subject: [PATCH 4/4] Update ENZYMEXLA_COMMIT hash in WORKSPACE --- deps/ReactantExtra/WORKSPACE | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE index bb561f7815..66adb0ab58 100644 --- a/deps/ReactantExtra/WORKSPACE +++ b/deps/ReactantExtra/WORKSPACE @@ -3,7 +3,7 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023" NSYNC_SHA256 = "" -ENZYMEXLA_COMMIT = "f46d7016fbc50dc87125bbdfddf71ea6782073fd" +ENZYMEXLA_COMMIT = "6137ac98e710adf6f4e953bf441db4e25b2db40f" ENZYMEXLA_SHA256 = "" http_archive(