Integrate LLVM at llvm/llvm-project@355e0f94af5a

Updates LLVM usage to match
[355e0f94af5a](llvm/llvm-project@355e0f94af5a)

PiperOrigin-RevId: 834865231

Author: Google-ML-Automation

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -3112,8 +3112,7 @@ def _yielded_values(outs, avals):
   switch_op = scf_dialect.IndexSwitchOp(
       yielded_types,
       _as_index(_ensure_ir_value(index, index_aval.dtype)),
-      ir.DenseI64ArrayAttr.get(range(len(branches) - 1)),
-      num_caseRegions=len(branches) - 1,
+      range(len(branches) - 1),
   )
 
   # ``RegionSequence`` in MLIR does not support slicing, so the
@@ -3124,7 +3123,7 @@ def _yielded_values(outs, avals):
   regions = regions[1:] + regions[:1]
   treedef = None
   for branch, region in zip(branches, regions):
-    with ir.InsertionPoint(region.blocks.append()):
+    with ir.InsertionPoint(region.blocks[0]):
       outs = lower_jaxpr_to_mosaic_gpu(
           ctx.module_ctx, ctx.launch_ctx, branch.jaxpr, args, consts=branch.consts
       )

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -1332,7 +1332,7 @@ def _mgpu_arrive_expect_tx_op_lowering_rule(
   barrier = utils.DialectBarrierRef.from_barrier_memref(
       arrive_expect_tx_op.barrier
   )
-  nvvm.mbarrier_arrive_expect_tx_shared(barrier.get_ptr(), bytes)
+  nvvm.mbarrier_arrive_expect_tx(barrier.get_ptr(), bytes)
 
   return []
 
@@ -2155,15 +2155,14 @@ def _index_switch_op_lowering_rule(
       _infer_flat_result_types(switch_op, out_layouts),
       switch_op.arg,
       switch_op.cases,
-      len(switch_op.regions) - 1,
   )
 
   results_template: Sequence[_VectorTemplate | None] = []
   for region, new_region in zip(
       switch_op.regions, new_switch_op.regions, strict=True
   ):
     [block] = region.blocks
-    new_block = new_region.blocks.append()
+    new_block = new_region.blocks[0]
     results_template = _move_scf_block_to_block_with_flattened_arguments(
         ctx, block, new_block, scf.YieldOp, []
     )

jax/experimental/mosaic/gpu/launch_context.py

@@ -1157,8 +1157,10 @@ def async_copy(
 
       if arrive:
         arrive_predicate = utils.single_thread_predicate(utils.ThreadSubset.WARPGROUP)
-        nvvm.mbarrier_arrive_expect_tx_shared(
-            barrier_ptr, transfer_bytes, predicate=arrive_predicate,
+        nvvm.mbarrier_arrive_expect_tx(
+            barrier_ptr,
+            transfer_bytes,
+            predicate=arrive_predicate,
         )
 
       gmem_strides, _ = gmem_ref_ty.get_strides_and_offset()
@@ -1286,7 +1288,7 @@ def async_copy(
               arith.CmpIPredicate.eq, self.cluster_idx(collective), c(0, index),
           )
           arrive_predicate = arith.andi(predicate, first_block)
-          nvvm.mbarrier_arrive_expect_tx_shared(
+          nvvm.mbarrier_arrive_expect_tx(
               barrier_ptr, transfer_bytes, predicate=arrive_predicate
           )
         rank = len(slice_shape)
@@ -1307,7 +1309,7 @@ def async_copy(
         )
       else:
         if arrive:
-          nvvm.mbarrier_arrive_expect_tx_shared(
+          nvvm.mbarrier_arrive_expect_tx(
               barrier_ptr, transfer_bytes, predicate=predicate
           )
         if collective_size > 1:

jax/experimental/mosaic/gpu/utils.py

@@ -995,7 +995,7 @@ def wait_parity(self, parity, orders_tensor_core=False):
     i32 = ir.IntegerType.get_signless(32)
     ticks = arith.constant(i32, 10000000)
     parity = arith.extui(i32, parity)
-    nvvm.mbarrier_try_wait_parity_shared(self.get_ptr(), parity, ticks)
+    nvvm.mbarrier_try_wait_parity(self.get_ptr(), parity, ticks)
     if orders_tensor_core:
       llvm.inline_asm(
           ir.Type.parse("!llvm.void"),
@@ -1064,9 +1064,7 @@ def arrive_expect_tx(
     elif ir.IndexType.isinstance(bytes.type):
       i32 = ir.IntegerType.get_signless(32)
       bytes = arith.index_cast(i32, bytes)
-    nvvm.mbarrier_arrive_expect_tx_shared(
-        self.get_ptr(), bytes, predicate=predicate
-    )
+    nvvm.mbarrier_arrive_expect_tx(self.get_ptr(), bytes, predicate=predicate)
 
   def get_ptr(self):
     i64 = ir.IntegerType.get_signless(64)

jaxlib/mosaic/gpu/serde.cc

@@ -50,7 +50,9 @@ constexpr llvm::StringRef kVersionAttrName = "stable_mosaic_gpu.version";
 // lowering after 2025-11-13.
 // TODO(apaszke): Update the forward-compatible version to 5 in Mosaic GPU
 // lowering after 2025-12-07.
-constexpr int kVersion = 5;
+// TODO(apaszke): Update the forward-compatible version to 6 in Mosaic GPU
+// lowering after 2025-12-18.
+constexpr int kVersion = 6;
 
 using SerdeRuleType = jaxlib::mosaic::SerdeRuleType;
 
@@ -174,6 +176,38 @@ LogicalResult nvvm_mbarrier_init_shared_upgrade(Operation* op, int version,
   return success();
 }
 
+LogicalResult nvvm_mbarrier_try_wait_parity_shared_upgrade(Operation* op,
+                                                           int version,
+                                                           bool& erased) {
+  // https://github.com/llvm/llvm-project/commit/7eeae8e41d7827d84de12df7b5ecfab3058900cb
+  if (version < 6) {
+    mlir::OpBuilder b(op->getParentRegion());
+    b.setInsertionPointAfter(op);
+    mlir::NVVM::MBarrierTryWaitParityOp::create(
+        b, op->getLoc(), op->getOperand(0), op->getOperand(1),
+        op->getOperand(2));
+    op->erase();
+    erased = true;
+  }
+  return success();
+}
+
+LogicalResult nvvm_mbarrier_arrive_expect_tx_shared_upgrade(Operation* op,
+                                                            int version,
+                                                            bool& erased) {
+  // https://github.com/llvm/llvm-project/commit/7eeae8e41d7827d84de12df7b5ecfab3058900cb
+  if (version < 6) {
+    mlir::OpBuilder b(op->getParentRegion());
+    b.setInsertionPointAfter(op);
+    mlir::NVVM::MBarrierArriveExpectTxOp::create(
+        b, op->getLoc(), op->getOperand(0), op->getOperand(1),
+        op->getNumOperands() < 3 ? Value{} : op->getOperand(2));
+    op->erase();
+    erased = true;
+  }
+  return success();
+}
+
 const llvm::StringMap<SerdeRuleType>& upgrade_rules() {
   static auto rules = new llvm::StringMap<SerdeRuleType>{
       {::llvm::StringLiteral("vector.extractelement"),
@@ -185,6 +219,10 @@ const llvm::StringMap<SerdeRuleType>& upgrade_rules() {
       {::llvm::StringLiteral("vector.splat"), vector_splat_upgrade},
       {::llvm::StringLiteral("nvvm.mbarrier.init.shared"),
        nvvm_mbarrier_init_shared_upgrade},
+      {::llvm::StringLiteral("nvvm.mbarrier.try_wait.parity.shared"),
+       nvvm_mbarrier_try_wait_parity_shared_upgrade},
+      {::llvm::StringLiteral("nvvm.mbarrier.arrive.expect_tx.shared"),
+       nvvm_mbarrier_arrive_expect_tx_shared_upgrade},
   };
   return *rules;
 }

tests/mosaic/gpu_layout_inference_test.py

@@ -509,20 +509,19 @@ def test_infer_index_switch_op_layouts(
       index_switch = scf.IndexSwitchOp(
           [out_type, out_type, f32],
           condition,
-          ir.DenseI64ArrayAttr.get(range(3)),
-          num_caseRegions=2,
+          range(2),
       )
-      with ir.InsertionPoint(index_switch.caseRegions[0].blocks.append()):
+      with ir.InsertionPoint(index_switch.caseRegions[0].blocks[0]):
         out0, out1, dummy0 = undefs(out_type, out_type, f32)
         if out0_layout is not None:
           out0 = layout_cast(out0, out0_layout)
         yield0 = scf.YieldOp([out0, out1, dummy0])
-      with ir.InsertionPoint(index_switch.caseRegions[1].blocks.append()):
+      with ir.InsertionPoint(index_switch.caseRegions[1].blocks[0]):
         out2, out3, dummy1 = undefs(out_type, out_type, f32)
         if out3_layout is not None:
           out3 = layout_cast(out3, out3_layout)
         yield1 = scf.YieldOp([out2, out3, dummy1])
-      with ir.InsertionPoint(index_switch.defaultRegion.blocks.append()):
+      with ir.InsertionPoint(index_switch.defaultRegion.blocks[0]):
         out4, out5, dummy2 = undefs(out_type, out_type, f32)
         if out4_layout is not None:
           out4 = layout_cast(out4, out4_layout)