Eqsat ci

.bazelrc

@@ -1,5 +1,6 @@
 build --announce_rc
 
+query --experimental_repo_remote_exec
 build --experimental_repo_remote_exec
 build --cxxopt=-std=c++17 --host_cxxopt=-std=c++17
 build --cxxopt=-w --host_cxxopt=-w
@@ -11,6 +12,7 @@ build --define framework_shared_object=true
 build --define tsl_protobuf_header_only=true
 build --define=use_fast_cpp_protos=true
 build --define=allow_oversize_protos=true
+build --check_visibility=false
 
 build -c opt
 
@@ -27,4 +29,22 @@ build:cuda --@local_config_cuda//:enable_cuda
 build:cuda --@xla//xla/python:jax_cuda_pip_rpaths=true
 # 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_cuda//:cuda_compiler=nvcc
+
+build:cudaci --repo_env TF_NEED_CUDA=1
+build:cudaci --repo_env TF_NVCC_CLANG=0
+build:cudaci --repo_env CUDA_NVCC=1
+build:cudaci --repo_env CC="/usr/bin/gcc"
+build:cudaci --repo_env TF_NCCL_USE_STUB=1
+build:cudaci --repo_env=HERMETIC_CUDA_VERSION="12.6.1"
+build:cudaci --repo_env=HERMETIC_CUDNN_VERSION="9.4.0"
+# "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:cudaci --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
+build:cudaci --crosstool_top="@local_config_cuda//crosstool:toolchain"
+build:cudaci --@local_config_cuda//:enable_cuda
+build:cudaci --@xla//xla/python:jax_cuda_pip_rpaths=true
+# Default hermetic CUDA and CUDNN versions.
+build:cudaci --@local_config_cuda//cuda:include_cuda_libs=true
+build:cudaci --@local_config_cuda//:cuda_compiler=nvcc
+

.buildkite/gpu_pipeline.yml

@@ -1,37 +1,97 @@
 steps:
   - label: "CUDA"
     agents:
-      queue: "juliagpu"
+      queue: "benchmark"
+      gpu: "rtx4070"
       cuda: "*"
     if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
+    timeout_in_minutes: 180
     commands: |
+      pwd
+      env
       echo "--- Setup :python: Dependencies"
       mkdir -p .local/bin
       export PATH="`pwd`/.local/bin:`pwd`/conda/bin:\$PATH"
       echo "openssl md5 | cut -d' ' -f2" > .local/bin/md5
       chmod +x .local/bin/md5
 
+      # No one tells us what to do
+      unset NV_LIBCUBLAS_VERSION
+      unset NVIDIA_VISIBLE_DEVICES
+      unset NV_NVML_DEV_VERSION
+      unset NV_LIBNCCL_DEV_PACKAGE
+      unset NV_LIBNCCL_DEV_PACKAGE_VERSION
+      unset NVIDIA_REQUIRE_CUDA
+      unset NV_LIBCUBLAS_DEV_PACKAGE
+      unset NV_NVTX_VERSION
+
       curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-amd64
 
       mv bazel* .local/bin/bazel
       chmod +x .local/bin/bazel
+      export PATH="`pwd`/.local/bin:\$PATH"
 
       export RUSTUP_HOME=`pwd`/.local/.rustup
-      export CARGO_HOME=`pwd`/.local/.cargo bash 
+      export CARGO_HOME=`pwd`/.local/.cargo bash
+
       curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
       export PATH="`pwd`/.local/.cargo/bin:\$PATH"
       rustup default stable
       cargo install cxxbridge-cmd
 
       mkdir -p .baztmp
+      touch src/enzyme_ad/jax/deps/tensat/cargo-bazel-lock.json
 
       echo "--- :python: Test"
 
-      # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --test_output=errors //test/...
-      # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no //test:bench_vs_xla
-      CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no --config=cuda //test:llama
+      export CUDA_DIR=`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvcc_cu12/site-packages/nvidia/cuda_nvcc
+      export XLA_FLAGS=--xla_gpu_cuda_data_dir=\$CUDA_DIR
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cusolver_cu12/site-packages/nvidia/cusolver:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cudnn_cu12/site-packages/nvidia/cudnn/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/test.runfiles/pypi_nvidia_cublas_cu12/site-packages/nvidia/cublas/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_cupti_cu12/site-packages/nvidia/cuda_cupti/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_runtime_cu12/site-packages/nvidia/cuda_runtime/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvrtc_cu12/site-packages/nvidia/cuda_nvrtc/lib:\$LD_LIBRARY_PATH"
+      export PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvcc_cu12/site-packages/nvidia/cuda_nvcc/bin:\$PATH"
+      export PATH="`pwd`/bazel-bin/test/llama.runfiles/python_x86_64-unknown-linux-gnu/bin:\$PATH"
+      export TF_CPP_MIN_LOG_LEVEL=0
+
+      BAZEL_BUILD_FLAGS=()
+      BAZEL_BUILD_FLAGS+=(--define=no_aws_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_gcp_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_hdfs_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_kafka_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=no_ignite_support=true)
+      BAZEL_BUILD_FLAGS+=(--define=grpc_no_ares=true)
+
+      BAZEL_BUILD_FLAGS+=(--define=llvm_enable_zlib=false)
+
+      BAZEL_BUILD_FLAGS+=(--verbose_failures)
+      BAZEL_BUILD_FLAGS+=(--cxxopt=-std=c++17 --host_cxxopt=-std=c++17)
+      BAZEL_BUILD_FLAGS+=(--cxxopt=-DTCP_USER_TIMEOUT=0)
+      BAZEL_BUILD_FLAGS+=(--check_visibility=false)
+      BAZEL_BUILD_FLAGS+=(--experimental_cc_shared_library)
+
+      mkdir .eqsat-tmp
+      export TMPDIR="`pwd`/.eqsat-tmp"
+      export TMP=$TMPDIR
+      export TEMP=$TMPDIR
+      BAZEL_BUILD_FLAGS+=(--action_env=TMP=\$TMPDIR --action_env=TEMP=\$TMPDIR --action_env=TMPDIR=\$TMPDIR --sandbox_tmpfs_path=\$TMPDIR)
+      BAZEL_BUILD_FLAGS+=(--action_env=EQSAT_PLATFORM=gpu)
+      BAZEL_BUILD_FLAGS+=(--config=cudaci)
+
+      export EQSAT_PLATFORM=gpu
+
+      CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp run --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL //builddeps:requirements.update || echo "no req update"
+      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --test_output=errors //test/... || echo "fail1"
+      find `pwd`/bazel-bin/test/llama.runfiles > finds.txt
+      # HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:bench_vs_xla || echo "fail2"
+      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s \${BAZEL_BUILD_FLAGS[@]} //test:llama || echo "fail3"
+      HERMETIC_PYTHON_VERSION="3.12" bazel-bin/test/llama
       cat bazel-out/*/testlogs/test/llama/test.log
     artifact_paths:
+      - "finds.txt"
+      - "finds2.txt"
       - "bazel-out/*/testlogs/test/llama/test.log"
       - "bazel-out/*/testlogs/test/llama/bench_vs_xla.log"
+      - "bazel-bin/test/llama.runfiles/results_*.csv"

.buildkite/pipeline.yml

@@ -8,7 +8,7 @@ steps:
         os:
           - macos
         python:
-          - "3.12"
+          - "3.11"
     agents:
       queue: "juliaecosystem"
       os: "{{matrix.os}}"
@@ -55,6 +55,7 @@ steps:
       cargo install cxxbridge-cmd
 
       mkdir -p .baztmp
+      touch src/enzyme_ad/jax/deps/tensat/cargo-bazel-lock.json
       # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --test_output=errors //test/...
       # CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no //test:bench_vs_xla
       CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --check_visibility=false --cache_test_results=no //test:llama

.gitignore

@@ -1,3 +1,4 @@
+*.pyc
 *.swp
 *.swo
 bazel-*

WORKSPACE

@@ -26,10 +26,9 @@ http_archive(
 )
 
 load("@jax//third_party/xla:workspace.bzl", "XLA_COMMIT", "XLA_SHA256")
-XLA_COMMIT = "7d4f8d1e8a91e67a713ac69796a22f343d292327"
 http_archive(
     name = "xla",
-    #sha256 = XLA_SHA256,
+    sha256 = XLA_SHA256,
     strip_prefix = "xla-" + XLA_COMMIT,
     urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
     patch_cmds = XLA_PATCHES,

build.sh

@@ -23,6 +23,6 @@ export CUDA_HOME=$HOME/miniconda3/
 export PATH=$PATH:$CUDA_HOME/bin
 export CUDACXX=$CUDA_HOME/bin/nvcc
 BAZEL_BUILD_FLAGS+=(--config=cuda)
-HERMETIC_PYTHON_VERSION=3.12 bazel build ${BAZEL_BUILD_FLAGS[@]} :wheel
-pip install bazel-bin/enzyme_ad-0.0.8-py312-none-manylinux2014_x86_64.whl --no-deps --force-reinstall
+HERMETIC_PYTHON_VERSION=3.11 bazel build ${BAZEL_BUILD_FLAGS[@]} :wheel
+pip install bazel-bin/enzyme_ad-0.0.8-py311-none-manylinux2014_x86_64.whl --no-deps --force-reinstall
 

builddeps/BUILD

@@ -10,6 +10,9 @@ compile_pip_requirements(
         "--build-isolation",
         "--rebuild",
     ],
+    extra_deps = [
+        # "@pypi_wheel//:pkg"
+    ],
     requirements_in = "requirements.in",
     requirements_txt = REQUIREMENTS,
     generate_hashes = True,

builddeps/requirements.in

@@ -3,6 +3,7 @@
 #
 -r test-requirements.txt
 
-jax >= 0.4.21
-jaxlib >= 0.4.21
+jax
+jaxlib
 absl_py >= 2.0.0
+ortools >= 9.10