Update jax-ml/jax to commit 6aaf234b1ee98800fcb80b9014a052a66dadbd8a

patches/xla.patch

@@ -1,9 +1,9 @@
 diff --git a/third_party/gpus/rocm/BUILD.tpl b/third_party/gpus/rocm/BUILD.tpl
-index 6e05471a99..c29d37acd3 100644
+index 09432ba..c2bf315 100644
 --- a/third_party/gpus/rocm/BUILD.tpl
 +++ b/third_party/gpus/rocm/BUILD.tpl
-@@ -468,6 +468,36 @@ cc_library(
-     data = [":system_libs_data"],
+@@ -570,6 +570,36 @@ rocm_lib_import(
+     deps = [":system_libs"],
  )
 
 +sh_binary(
@@ -40,18 +40,18 @@ index 6e05471a99..c29d37acd3 100644
      name = "toolchain_data",
      srcs = glob(
 diff --git a/third_party/rocm_device_libs/build_defs.bzl b/third_party/rocm_device_libs/build_defs.bzl
-index 5b9aee4f08..f55aabe5ad 100644
+index 49c0119..bd4ea0c 100644
 --- a/third_party/rocm_device_libs/build_defs.bzl
 +++ b/third_party/rocm_device_libs/build_defs.bzl
-@@ -111,6 +111,7 @@ def _bitcode_library_impl(ctx):
+@@ -124,6 +124,7 @@ def _bitcode_library_impl(ctx):
          inputs = [strip_out],
          outputs = [final_bc],
          arguments = [strip_out.path, "-o", final_bc.path],
 +        use_default_shell_env = True,
          progress_message = "Preparing final bitcode for {}".format(ctx.label.name),
          mnemonic = "BitcodeFinalize",
      )
-@@ -127,17 +128,17 @@ bitcode_library = rule(
+@@ -139,17 +140,17 @@ bitcode_library = rule(
          "hdrs": attr.label_list(allow_files = [".h"]),
          "file_specific_flags": attr.string_list_dict(),
          "_clang": attr.label(
@@ -73,10 +73,10 @@ index 5b9aee4f08..f55aabe5ad 100644
              cfg = "exec",
          ),
 diff --git a/third_party/rocm_device_libs/rocm_device_libs.BUILD b/third_party/rocm_device_libs/rocm_device_libs.BUILD
-index ee162b63ab..98a55618a4 100644
+index 1709a6b..8994c5f 100644
 --- a/third_party/rocm_device_libs/rocm_device_libs.BUILD
 +++ b/third_party/rocm_device_libs/rocm_device_libs.BUILD
-@@ -28,36 +28,12 @@ cc_binary(
+@@ -43,36 +43,12 @@ cc_binary(
      ],
  )
 
@@ -117,20 +117,3 @@ index ee162b63ab..98a55618a4 100644
 -    },
 +    actual = "@local_config_rocm//rocm:ocml"
  )
-
-diff --git a/xla/pjrt/distributed/topology_util.cc b/xla/pjrt/distributed/topology_util.cc
---- a/xla/pjrt/distributed/topology_util.cc
-+++ b/xla/pjrt/distributed/topology_util.cc
-@@ -307,8 +307,10 @@ absl::Status ExchangeTopologies(absl::string_view platform, int node_id,
-   if (num_nodes == 1) {
-     LocalTopologyProto* topology = global_topology->add_processes();
-     *topology = local_topology;
--    for (DeviceProto& device : *topology->mutable_devices()) {
--      device.set_global_device_id(device.local_device_ordinal());
-+    if (assign_global_device_ids) {
-+      for (int i = 0; i < topology->devices_size(); ++i) {
-+        topology->mutable_devices(i)->set_global_device_id(i);
-+      }
-     }
-     return absl::OkStatus();
-   }

src/enzyme_ad/jax/Passes/ConvertPolygeistToLLVM.cpp

@@ -4525,7 +4525,7 @@ populateCStyleGPUFuncLoweringPatterns(RewritePatternSet &patterns,
     } else if (gpuTarget == "rocm") {
       mlir::populateGpuToROCDLConversionPatterns(typeConverter, patterns,
                                                  mlir::gpu::amd::Runtime::HIP,
-                                                 amdgpu::Chipset());
+                                                 ROCDL::TargetInfo());
     }
   }
 }

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -14631,7 +14631,7 @@ struct CompareOpCanon final
                 ? stablehlo::makeTensor(rhs.resizeSplat(RankedTensorType::get(
                       {}, rhs.getType().getElementType())))
                 : stablehlo::constantOp(rhs),
-            op.getComparisonDirection(), ty));
+            op.getComparisonDirection(), op.getCompareType(), ty));
         if (isSplat)
           out = out.resizeSplat(op.getType());
 

src/external/polly/include/polly/CodeGen/IRBuilder.h

@@ -123,7 +123,7 @@ class IRInserter final : public llvm::IRBuilderDefaultInserter {
   IRInserter(ScopAnnotator &A) : Annotator(&A) {}
 
   void InsertHelper(llvm::Instruction *I, const llvm::Twine &Name,
-                    llvm::BasicBlock::iterator InsertPt) const override {
+                    llvm::BasicBlock::iterator InsertPt) const {
     llvm::IRBuilderDefaultInserter::InsertHelper(I, Name, InsertPt);
     if (Annotator)
       Annotator->annotate(I);

test/lit_tests/raising/raiseimmerse3.mlir

@@ -27,15 +27,19 @@ module {
 }
 
 // CHECK:  func.func private @"##call__Z25gpu__mask_immersed_field_16CompilerMetadataI10StaticSizeI13_180__85__20_E12DynamicCheckvv7NDRangeILi3ES0_I11_12__6__20_ES0_I11_16__16__1_EvvEE11OffsetArrayI7Float64Li3E13CuTracedArrayIS9_Li3ELi1E13_194__99__34_EE5TupleI6Center4FaceSE_E20ImmersedBoundaryGridIS9_8Periodic14RightConnected7Bounded28OrthogonalSphericalShellGridIS9_SI_SJ_SK_28StaticVerticalDiscretizationIS8_IS9_Li1ESA_IS9_Li1ELi1E5_35__EES8_IS9_Li1ESA_IS9_Li1ELi1E5_34__EESO_SQ_ES8_IS9_Li2ESA_IS9_Li2ELi1E9_194__99_EE8TripolarI5Int64SV_SV_EvE16GridFittedBottomI5FieldISE_SE_vvvvS8_IS9_Li3ESA_IS9_Li3ELi1E12_194__99__1_EES9_vvvE23CenterImmersedConditionEvvvES9_#1335$par75"(%arg0: memref<34x99x194xf64, 1>, %arg1: memref<34xf64, 1>, %arg2: memref<1x99x194xf64, 1>, %arg3: i1) {
-// CHECK-NEXT:    %c0 = arith.constant 0 : index
+// CHECK-NEXT:    %false = arith.constant false
+// CHECK-NEXT:    %true = arith.constant true
 // CHECK-NEXT:    %cst = arith.constant 0.000000e+00 : f64
 // CHECK-NEXT:    %cst_0 = arith.constant 2.000000e+00 : f64
 // CHECK-NEXT:    affine.parallel (%arg4) = (0) to (85) {
-// CHECK-NEXT:      %0 = arith.cmpi eq, %arg4, %c0 : index
+// CHECK-NEXT:      %0 = affine.if #set(%arg4) -> i1 {
+// CHECK-NEXT:        affine.yield %true : i1
+// CHECK-NEXT:      } else {
+// CHECK-NEXT:        affine.yield %false : i1
+// CHECK-NEXT:      }
 // CHECK-NEXT:      %1 = arith.ori %0, %arg3 : i1
-// CHECK-NEXT:      %2 = arith.ori %arg3, %1 : i1
-// CHECK-NEXT:      %3 = arith.select %2, %cst, %cst_0 : f64
-// CHECK-NEXT:      affine.store %3, %arg0[7, %arg4 + 7, 7] : memref<34x99x194xf64, 1>
+// CHECK-NEXT:      %2 = arith.select %1, %cst, %cst_0 : f64
+// CHECK-NEXT:      affine.store %2, %arg0[7, %arg4 + 7, 7] : memref<34x99x194xf64, 1>
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return
 // CHECK-NEXT:  }

workspace.bzl

@@ -1,4 +1,4 @@
-JAX_COMMIT = "fbf6588d55c5662ccf98558015d5b2c4c4a7bf6f"
+JAX_COMMIT = "6aaf234b1ee98800fcb80b9014a052a66dadbd8a"
 JAX_SHA256 = ""
 
 ENZYME_COMMIT = "56b02e7e94265e7af38fe95e9fbf9818ef1d09c0"
@@ -70,10 +70,10 @@ echo "     InterpreterFallback* fallback) {" >> third_party/stablehlo/temporary.
 echo "+#if (defined(_WIN32) || defined(__CYGWIN__))" >> third_party/stablehlo/temporary.patch
 echo "+  llvm::report_fatal_error(\\"Op not supported on windows due to std::future\\");" >> third_party/stablehlo/temporary.patch
 echo "+#else" >> third_party/stablehlo/temporary.patch
-echo "   llvm::DefaultThreadPool threadPool;" >> third_party/stablehlo/temporary.patch
-echo "   llvm::ThreadPoolTaskGroup taskGroup(threadPool);" >> third_party/stablehlo/temporary.patch
+echo "   uint32_t numReplicas = programs.size();" >> third_party/stablehlo/temporary.patch
+echo "   uint32_t numPartitions = programs[0].size();" >> third_party/stablehlo/temporary.patch
 echo " " >> third_party/stablehlo/temporary.patch
-echo "@@ -211,6 +214,7 @@ SmallVector<InterpreterValue> evalRunParallelOp(" >> third_party/stablehlo/temporary.patch
+echo "@@ -223,6 +226,7 @@ SmallVector<InterpreterValue> evalRunParallelOp(" >> third_party/stablehlo/temporary.patch
 echo "   for (auto& output : taskOutputs) results.append(output);" >> third_party/stablehlo/temporary.patch
 echo "   // TODO(#1725): Figure out how to test the outfeed queue." >> third_party/stablehlo/temporary.patch
 echo "   return results;" >> third_party/stablehlo/temporary.patch
@@ -105,7 +105,7 @@ echo " llvm::Error evalPrintOp(PrintOp& op, InterpreterValue operand) {" >> thir
     # place copies and spills above the exec restore of an if/else join,
     # miscompiling kernels (llvm/llvm-project#222368). Drop once XLA's LLVM
     # includes it.
-    sed -i.bak0 "s/llvm:generated.patch\\\",/llvm:generated.patch\\\", \\\"\\/\\/:patches\\/llvm_amdgpu_bb_prolog.patch\\\", \\\"\\/\\/:patches\\/llvm_mlir_import_inrange_width.patch\\\", \\\"\\/\\/:patches\\/llvm_orc_unw_revert.patch\\\", \\\"\\/\\/:patches\\/llvm_affine_parallel_signless_minmax.patch\\\",/g" third_party/llvm/workspace.bzl
+    sed -i.bak0 "s/llvm:generated.patch\\\",/llvm:generated.patch\\\", \\\"\\/\\/:patches\\/llvm_amdgpu_bb_prolog.patch\\\", \\\"\\/\\/:patches\\/llvm_affine_parallel_signless_minmax.patch\\\",/g" third_party/llvm/workspace.bzl
     """,
     """
     sed -i.bak0 "s/tf_http_archive/http_archive/g" third_party/llvm/workspace.bzl