Try to fix macos (#69)

* Try to fix macos

* Try later baselisk

* Add caching

* try hidden file

Author: wsmoses

.buildkite/pipeline.yml

@@ -12,6 +12,17 @@ steps:
       os: "{{matrix.os}}"
       arch: "{{matrix.arch}}"
     if: build.tag == null
+    plugins:
+      - cache#v0.6.0:
+          manifest: .buildkite/pipeline.yml
+          path: .local
+          restore: file
+          save: file
+      - cache#v0.6.0:
+          manifest: WORKSPACE
+          path: .baztmp
+          restore: file
+          save: file
     commands: |
       mkdir -p .local/bin
       export PATH="`pwd`/.local/bin:`pwd`/conda/bin:\$PATH"
@@ -26,17 +37,17 @@ steps:
           sed -i.bak 's~targets = \[.*\]~targets = \[\"X86\", \"AMDGPU\"]~g' WORKSPACE
           curl -fLO https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-MacOSX-{{matrix.arch}}.sh
         fi
-        curl -fLO "https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-darwin"
+        curl -fLO "https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-darwin"
         mv bazelisk-darwin .local/bin/bazel
         chmod +x .local/bin/bazel
         chmod +x Miniconda*.sh
         ./Miniconda*.sh -b -p `pwd`/conda
         rm Miniconda*.sh
       elif [ "{{matrix.os}}" == "linux" ]; then
         if [ "{{matrix.arch}}" == "aarch64" ]; then
-          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-arm64
+          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-arm64
         else
-          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64
+          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-amd64
         fi
         mv bazel* .local/bin/bazel
         chmod +x .local/bin/bazel
@@ -58,12 +69,13 @@ steps:
       # conda install -c conda-forge cxx-compiler -y
       python -m ensurepip --upgrade
       python -m pip install --user numpy wheel
-      mkdir baztmp
-      bazel --output_user_root=`pwd`/baztmp build :enzyme_ad
+      mkdir -p .baztmp
+      rm -f bazel-bin/*.whl
+      bazel --output_user_root=`pwd`/.baztmp build :enzyme_ad
       cp bazel-bin/*.whl .
       python -m pip install *.whl
       python -m pip install "jax[cpu]"
-      bazel --output_user_root=`pwd`/baztmp test --test_output=errors ...
+      bazel --output_user_root=`pwd`/.baztmp test --test_output=errors ...
     artifact_paths:
       - "*.whl"
 

.buildkite/secure_pipeline.yml

@@ -33,17 +33,17 @@ steps:
           sed -i.bak 's~targets = \[.*\]~targets = \[\"X86\", \"AMDGPU\"]~g' WORKSPACE
           curl -fLO https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-MacOSX-{{matrix.arch}}.sh
         fi
-        curl -fLO "https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-darwin"
+        curl -fLO "https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-darwin"
         mv bazelisk-darwin .local/bin/bazel
         chmod +x .local/bin/bazel
         chmod +x Miniconda*.sh
         ./Miniconda*.sh -b -p `pwd`/conda
         rm Miniconda*.sh
       elif [ "{{matrix.os}}" == "linux" ]; then      
         if [ "{{matrix.arch}}" == "aarch64" ]; then
-          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-arm64
+          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-arm64
         else
-          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64
+          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-amd64
         fi
         mv bazel* .local/bin/bazel
         chmod +x .local/bin/bazel

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -3677,10 +3677,11 @@ struct PadDotGeneral : public OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
       SmallVector<int64_t> sliceLow, sliceHigh, sliceStride;
       for (auto &&[pos, size] :
            llvm::enumerate(otherArg.getType().cast<TensorType>().getShape())) {
-        auto it = llvm::find_if(otherDimsToSlice, [&](auto &tup) {
-          return std::get<0>(tup) == pos;
-        });
-        if (it == otherDimsToSlice.end()) {
+        std::optional<std::tuple<int64_t, int64_t, int64_t, int64_t>> it;
+        for (auto tup : otherDimsToSlice)
+          if (std::get<0>(tup) == pos)
+            it = tup;
+        if (!it) {
           sliceLow.push_back(0);
           sliceHigh.push_back(size);
           sliceStride.push_back(1);