Integrate LLVM at llvm/llvm-project@917d1f20aecf

Updates LLVM usage to match
[917d1f20aecf](llvm/llvm-project@917d1f20aecf)

PiperOrigin-RevId: 823542980

Author: Google-ML-Automation

jax/_src/lib/mlir/dialects/__init__.py

@@ -17,6 +17,7 @@
 from typing import Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
+  from jaxlib.mlir.dialects import _gpu_ops_gen as _gpu_ops_gen
   from jaxlib.mlir.dialects import arith as arith
   from jaxlib.mlir.dialects import builtin as builtin
   from jaxlib.mlir.dialects import cf as cf
@@ -35,6 +36,7 @@
 else:
   from jax._src import lazy_loader as _lazy
   __getattr__, __dir__, __all__ = _lazy.attach("jaxlib.mlir.dialects", [
+      "_gpu_ops_gen",
       "arith",
       "builtin",
       "cf",

jax/experimental/mosaic/gpu/core.py

@@ -32,13 +32,14 @@
 from jax._src import core as jax_core
 from jax._src import dtypes
 from jax._src import lib
-from jax._src import sharding_impls
 from jax._src import mesh as mesh_lib
+from jax._src import sharding_impls
 from jax._src import util as jax_util
 from jax._src.interpreters import mlir
 from jax._src.lib import mosaic_gpu_dialect as dialect
 from jaxlib.mlir import ir
 from jaxlib.mlir import passmanager
+from jaxlib.mlir.dialects import _gpu_ops_gen
 from jaxlib.mlir.dialects import arith
 from jaxlib.mlir.dialects import builtin
 from jaxlib.mlir.dialects import func
@@ -579,9 +580,14 @@ def _launch(
         )
   else:
     cluster_kwargs = {}
-  launch_op = gpu.LaunchOp(
-      token.type, [token], *grid_vals, *block_vals,
-      dynamicSharedMemorySize=c(smem_bytes, i32), **cluster_kwargs)
+  launch_op = _gpu_ops_gen.LaunchOp(
+      token.type,
+      [token],
+      *grid_vals,
+      *block_vals,
+      dynamicSharedMemorySize=c(smem_bytes, i32),
+      **cluster_kwargs,
+  )
   launch_op.body.blocks.append(*([index] * (12 + 2 * len(cluster_kwargs))))  # Append an empty block
   with ir.InsertionPoint(launch_op.body.blocks[0]):
     dynamic_smem = gpu.dynamic_shared_memory(

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -25,6 +25,7 @@
 from jax._src.interpreters import mlir as mlir_interpreter
 from jax._src.lib import mosaic_gpu_dialect as mgpu
 from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import _gpu_ops_gen
 from jax._src.lib.mlir.dialects import arith
 from jax._src.lib.mlir.dialects import builtin
 from jax._src.lib.mlir.dialects import func
@@ -2062,7 +2063,7 @@ def _index_switch_op_lowering_rule(
 
 
 @_register_lowering(func.FuncOp)
-@_register_lowering(gpu.LaunchOp)
+@_register_lowering(_gpu_ops_gen.LaunchOp)
 def _traverse_op_lowering_rule(
     ctx: LoweringContext, op: ir.OpView
 ) -> MlirLoweringRuleResult:

jax/experimental/mosaic/gpu/launch_context.py

@@ -23,6 +23,7 @@
 
 from jax._src.lib import mosaic_gpu_dialect as mgpu_dialect
 from jaxlib.mlir import ir
+from jaxlib.mlir.dialects import _gpu_ops_gen
 from jaxlib.mlir.dialects import arith
 from jaxlib.mlir.dialects import builtin
 from jaxlib.mlir.dialects import func
@@ -32,9 +33,9 @@
 from jaxlib.mlir.dialects import nvvm
 import numpy as np
 
+from . import fragmented_array as fa
 from . import profiler
 from . import utils
-from . import fragmented_array as fa
 
 TMA_DESCRIPTOR_BYTES = 128
 TMA_DESCRIPTOR_ALIGNMENT = 64
@@ -304,7 +305,7 @@ class Scratch:
              : (!llvm.array<256 x i8>) -> !llvm.ptr
 
   """
-  def __init__(self, gpu_launch_op: gpu.LaunchOp):
+  def __init__(self, gpu_launch_op: _gpu_ops_gen.LaunchOp):
     self.next_offset: int = 0
     self.host_init: list[Callable[[ir.Value], None]] = []
     self._ops_created = False