Update JAX & MLIR dependency chain to v0.4.28 (#931)

**Context:**
We target at 0.4.28 instead of 0.4.30 because many bufferization passes
are removed after [this llvm
commit](llvm/llvm-project#93535).

**Description of the Change:**
***Mandatory Updates:***
* Cmake
- Remove `MhloShapeOpsToStandard`
[link](tensorflow/mlir-hlo@57d2124)
  - Add `StablehloPasses`
  - Add `MhloQuantToIntConversion` (This will be removed after 0.4.29)
  -  `EnzymeStatic-18` => `EnzymeStatic-19`
- `RunnerUtils.h` requires `Float16Bits.h`
[link](llvm/llvm-project@7bc6c4a)
* LLVM
- `setDataLayout` must happen before code generation or they will use
the default one.
[link](https://discord.com/channels/636084430946959380/636732535434510338/1265407221324451871)
* MLIR
  - `updateRootInPlace` => `modifyOpInPlace`
  -  `startRootUpdate` => `startOpModification`
  -  `finalizeRootUpdate` => `finalizeOpModification`
- The order of transformed mlir expressions is different and required
fine-tuning for `CHECK-DAG`s.
- Using rewriter's method to release MeasurementOps (Caused by LLVM
commit-`b840d2968391dd610b792a65133a1edc1bcc397c`).
[link](llvm/llvm-project@b840d29)
- Allow `replaceTerminatorWithUnconditionalJumpToSuccessBlock` to accept
`LLVM:br`. (New LLVM/MLIR will reuse `LLVM::UnreachableOp` and use
`LLVM::br` to reach it.)
* Frontend
- `jax.linear_util` => `jax.extend.linear_util`
[link](https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-24-feb-6-2024)
  - `jax_ctx.module_context.replace` => `jax_ctx.replace`
- `gensym(jaxprs, suffix)` => `gensym(suffix)`
[link](jax-ml/jax@67df647)
- Move `name_stack` out of mlir.ModuleContext (functions like
`lower_jaxpr_to_fun` is taking `name_stack` now).
[link](jax-ml/jax#19856).
- Pass `LoweringRuleContext.ModuleContext` instead of
`LoweringRuleContext` to `jaxpr_subcomp`
- Patch new `_sin_lowering` and `_cos_lowering` with
`_nary_lower_hlo(sine/cosine)`.
  - Variable names (like `%0` => `%cst`) and orders in FileCheck.

***Deprecations/Warnings:***
* MLIR
- `x.cast<T>()` => `mlir::cast<T>(x)`
[link](https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443)
- `x.dyn_cast<T>()` => `mlir::dyn_cast<T>(x)`
[link](https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443)
- `x.isa<T>()` => `mlir::isa<T>(x)`
[link](https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443)
* Enzyme
  - Add `DCMAKE_POLICY_DEFAULT_CMP0116`

**Related GitHub Issues:**
#863
[sc-67111]

---------

Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: erick-xanadu <110487834+erick-xanadu@users.noreply.github.com>

Author: 3 people

.dep-versions

@@ -1,7 +1,7 @@
 # Always update the version check in catalyst.__init__ when changing the JAX version.
-jax=0.4.23
-mhlo=4611968a5f6818e6bdfb82217b9e836e0400bba9
-llvm=cd9a641613eddf25d4b25eaa96b2c393d401d42c
+jax=0.4.28
+mhlo=89a891c986650c33df76885f5620e0a92150d90f
+llvm=3a8316216807d64a586b971f51695e23883331f7
 enzyme=v0.0.130
 
 # Always remove custom PL/LQ versions before release.

doc/changelog.md

@@ -170,6 +170,8 @@
   [(#945)](https://github.com/PennyLaneAI/catalyst/pull/945)
   [(#962)](https://github.com/PennyLaneAI/catalyst/pull/962)
 
+* Update JAX to `v0.4.28`. [(#931)](https://github.com/PennyLaneAI/catalyst/pull/931)
+
 <h3>Breaking changes</h3>
 
 * Return values of qjit-compiled functions that were previously `numpy.ndarray` are now of type

frontend/catalyst/__init__.py

@@ -23,7 +23,7 @@
 
 import jaxlib as _jaxlib
 
-_jaxlib_version = "0.4.23"
+_jaxlib_version = "0.4.28"
 if _jaxlib.__version__ != _jaxlib_version:
     import warnings
 

frontend/catalyst/compiler.py

@@ -165,7 +165,6 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
         "func.func(chlo-legalize-to-hlo)",
         "stablehlo-legalize-to-hlo",
         "func.func(mhlo-legalize-control-flow)",
-        "func.func(hlo-legalize-shapeops-to-standard)",
         "func.func(hlo-legalize-to-linalg)",
         "func.func(mhlo-legalize-to-std)",
         "func.func(hlo-legalize-sort)",

frontend/catalyst/jax_extras/lowering.py

@@ -125,7 +125,6 @@ def custom_lower_jaxpr_to_module(
         backend_or_name=None,
         platforms=[platform],
         axis_context=axis_context,
-        name_stack=name_stack,
         keepalives=keepalives,
         channel_iterator=channel_iter,
         host_callbacks=host_callbacks,
@@ -149,6 +148,7 @@ def custom_lower_jaxpr_to_module(
             replicated_args=replicated_args,
             arg_shardings=arg_shardings,
             result_shardings=result_shardings,
+            name_stack=name_stack,
         )
 
         for op in ctx.module.body.operations:

frontend/catalyst/jax_extras/patches.py

@@ -18,19 +18,23 @@
 from __future__ import annotations
 
 import jax
+from jax._src.lax.lax import _nary_lower_hlo
 from jax._src.lax.slicing import (
     _gather_shape_computation,
     _is_sorted,
     _no_duplicate_dims,
     _rank,
     _sorted_dims_in_range,
 )
+from jax._src.lib.mlir.dialects import hlo
 from jax.core import AbstractValue, Tracer, concrete_aval
 
 __all__ = (
     "get_aval2",
     "_no_clean_up_dead_vars",
     "_gather_shape_rule_dynamic",
+    "_sin_lowering2",
+    "_cos_lowering2",
 )
 
 
@@ -180,3 +184,13 @@ def _gather_shape_rule_dynamic(
             )
 
     return _gather_shape_computation(indices, dimension_numbers, slice_sizes)
+
+
+def _sin_lowering2(ctx, x):
+    """Use hlo.sine lowering instead of the new sin lowering from jax 0.4.28"""
+    return _nary_lower_hlo(hlo.sine, ctx, x)
+
+
+def _cos_lowering2(ctx, x):
+    """Use hlo.cosine lowering instead of the new cosine lowering from jax 0.4.28"""
+    return _nary_lower_hlo(hlo.cosine, ctx, x)

frontend/catalyst/jax_extras/tracing.py

@@ -44,7 +44,7 @@
     trace_to_jaxpr_dynamic2,
 )
 from jax._src.lax.control_flow import _initial_style_jaxpr
-from jax._src.lax.lax import _abstractify
+from jax._src.lax.lax import _abstractify, cos_p, sin_p
 from jax._src.lax.slicing import (
     _argnum_weak_type,
     _gather_dtype_rule,
@@ -80,14 +80,14 @@
     new_jaxpr_eqn,
     thread_local_state,
 )
+from jax.extend.linear_util import transformation_with_aux, wrap_init
 from jax.interpreters.partial_eval import (
     DynamicJaxprTrace,
     DynamicJaxprTracer,
     convert_constvars_jaxpr,
     make_jaxpr_effects,
 )
 from jax.lax import convert_element_type
-from jax.linear_util import transformation_with_aux, wrap_init
 from jax.tree_util import (
     PyTreeDef,
     tree_flatten,
@@ -97,7 +97,12 @@
 )
 from jaxlib.xla_extension import PyTreeRegistry
 
-from catalyst.jax_extras.patches import _gather_shape_rule_dynamic, get_aval2
+from catalyst.jax_extras.patches import (
+    _cos_lowering2,
+    _gather_shape_rule_dynamic,
+    _sin_lowering2,
+    get_aval2,
+)
 from catalyst.logging import debug_logger
 from catalyst.tracing.type_signatures import verify_static_argnums_type
 from catalyst.utils.patching import Patcher
@@ -288,7 +293,7 @@ def __init__(self, boxid: int, e: JaxprEqn):
 def jaxpr_pad_consts(jaxprs: List[Jaxpr]) -> List[ClosedJaxpr]:
     """Align the constants of Jaxpr programs. Return the list of corresponding programs accepting
     the same constants."""
-    newvar = gensym(jaxprs, suffix="_")
+    newvar = gensym("_")
 
     # List of constant variables of all jaxprs, preprended with '_'
     all_mangled_constvars: List[List[Var]] = []
@@ -519,6 +524,10 @@ def abstractify(args, kwargs):
     )
     register_lowering(gather2_p, _gather_lower)
 
+    # TBD
+    register_lowering(sin_p, _sin_lowering2)
+    register_lowering(cos_p, _cos_lowering2)
+
     primitive_batchers2 = jax._src.interpreters.batching.primitive_batchers.copy()
     for primitive in jax._src.interpreters.batching.primitive_batchers.keys():
         if primitive.name == "gather":
@@ -532,6 +541,8 @@ def make_jaxpr_f(*args, **kwargs):
             (jax._src.interpreters.partial_eval, "get_aval", get_aval2),
             (jax._src.lax.slicing, "gather_p", gather2_p),
             (jax._src.interpreters.batching, "primitive_batchers", primitive_batchers2),
+            (jax._src.lax.lax, "_sin_lowering", _sin_lowering2),
+            (jax._src.lax.lax, "_cos_lowering", _cos_lowering2),
         ), ExitStack():
             f = wrap_init(fun)
             if static_argnums:

frontend/catalyst/jax_primitives.py

@@ -330,8 +330,8 @@ def _python_callback_lowering(
     fwd_jaxpr = custom_grad._fwd_jaxpr
     rev_jaxpr = custom_grad._bwd_jaxpr
     ctx = jax_ctx.module_context
-    mlir_fwd = _func_def_lowering(ctx, call_jaxpr=fwd_jaxpr, fn=fwd)
-    mlir_rev = _func_def_lowering(ctx, call_jaxpr=rev_jaxpr, fn=rev)
+    mlir_fwd = _func_def_lowering(ctx, call_jaxpr=fwd_jaxpr, fn=fwd, name_stack=jax_ctx.name_stack)
+    mlir_rev = _func_def_lowering(ctx, call_jaxpr=rev_jaxpr, fn=rev, name_stack=jax_ctx.name_stack)
     sym_fwd = mlir_fwd.sym_name.value + ".fwd"
 
     argc = len(args)
@@ -397,11 +397,11 @@ def _func_def_impl(ctx, *args, call_jaxpr, fn, call=True):  # pragma: no cover
     raise NotImplementedError()
 
 
-def _func_def_lowering(ctx, fn, call_jaxpr) -> str:
+def _func_def_lowering(ctx, fn, call_jaxpr, name_stack) -> str:
     """Create a func::FuncOp from JAXPR."""
     if isinstance(call_jaxpr, core.Jaxpr):
         call_jaxpr = core.ClosedJaxpr(call_jaxpr, ())
-    func_op = mlir.lower_jaxpr_to_fun(ctx, fn.__name__, call_jaxpr, tuple())
+    func_op = mlir.lower_jaxpr_to_fun(ctx, fn.__name__, call_jaxpr, tuple(), name_stack=name_stack)
 
     if isinstance(fn, qml.QNode):
         func_op.attributes["qnode"] = ir.UnitAttr.get()
@@ -443,7 +443,7 @@ def _func_lowering(ctx, *args, call_jaxpr, fn, call=True):
     if fn in mlir_fn_cache:
         func_op = mlir_fn_cache[fn]
     else:
-        func_op = _func_def_lowering(ctx.module_context, fn, call_jaxpr)
+        func_op = _func_def_lowering(ctx.module_context, fn, call_jaxpr, name_stack=ctx.name_stack)
         mlir_fn_cache[fn] = func_op
 
     symbol_name = func_op.name.value
@@ -1585,14 +1585,13 @@ def emit_branches(preds, branch_jaxprs, ip):
 
             # if block
             source_info_util.extend_name_stack("if")
-            if_ctx = jax_ctx.module_context.replace(
-                name_stack=jax_ctx.module_context.name_stack.extend("if")
-            )
+            if_ctx = jax_ctx.replace(name_stack=jax_ctx.name_stack.extend("if"))
             with ir.InsertionPoint(if_block):
                 # recursively generate the mlir for the if block
                 out = mlir.jaxpr_subcomp(
-                    if_ctx,
+                    if_ctx.module_context,
                     true_jaxpr.jaxpr,
+                    if_ctx.name_stack,
                     mlir.TokenSet(),
                     [mlir.ir_constants(c) for c in true_jaxpr.consts],
                     *([a] for a in flat_args_plus_consts),  # fn expects [a1], [a2], [a3] format
@@ -1603,17 +1602,16 @@ def emit_branches(preds, branch_jaxprs, ip):
 
             # else block
             source_info_util.extend_name_stack("else")
-            else_ctx = jax_ctx.module_context.replace(
-                name_stack=jax_ctx.module_context.name_stack.extend("else")
-            )
+            else_ctx = jax_ctx.replace(name_stack=jax_ctx.name_stack.extend("else"))
             else_block = if_op_scf.else_block
             if len(preds) == 1:
                 # Base case: reached the otherwise block
                 otherwise_jaxpr = branch_jaxprs[-1]
                 with ir.InsertionPoint(else_block):
                     out = mlir.jaxpr_subcomp(
-                        else_ctx,
+                        else_ctx.module_context,
                         otherwise_jaxpr.jaxpr,
+                        else_ctx.name_stack,
                         mlir.TokenSet(),
                         [mlir.ir_constants(c) for c in otherwise_jaxpr.consts],
                         *([a] for a in flat_args_plus_consts),
@@ -1695,15 +1693,16 @@ def _while_loop_lowering(
 
     # cond block
     cond_block = while_op_scf.regions[0].blocks.append(*loop_carry_types)
-    name_stack = jax_ctx.module_context.name_stack.extend("while")
-    cond_ctx = jax_ctx.module_context.replace(name_stack=name_stack.extend("cond"))
+    name_stack = jax_ctx.name_stack.extend("while")
+    cond_ctx = jax_ctx.replace(name_stack=name_stack.extend("cond"))
     with ir.InsertionPoint(cond_block):
         cond_args = [cond_block.arguments[i] for i in range(len(loop_carry_types))]
 
         # recursively generate the mlir for the while cond
         ((pred,),), _ = mlir.jaxpr_subcomp(
-            cond_ctx,
+            cond_ctx.module_context,
             cond_jaxpr.jaxpr,
+            cond_ctx.name_stack,
             mlir.TokenSet(),
             [mlir.ir_constants(c) for c in cond_jaxpr.consts],
             *([a] for a in (cond_consts + cond_args)),  # fn expects [a1], [a2], [a3] format
@@ -1715,14 +1714,15 @@ def _while_loop_lowering(
 
     # body block
     body_block = while_op_scf.regions[1].blocks.append(*loop_carry_types)
-    body_ctx = jax_ctx.module_context.replace(name_stack=name_stack.extend("body"))
+    body_ctx = jax_ctx.replace(name_stack=name_stack.extend("body"))
     with ir.InsertionPoint(body_block):
         body_args = [body_block.arguments[i] for i in range(len(loop_carry_types))]
 
         # recursively generate the mlir for the while body
         out, _ = mlir.jaxpr_subcomp(
-            body_ctx,
+            body_ctx.module_context,
             body_jaxpr.jaxpr,
+            body_ctx.name_stack,
             mlir.TokenSet(),
             [mlir.ir_constants(c) for c in cond_jaxpr.consts],
             *([a] for a in (body_consts + body_args)),  # fn expects [a1], [a2], [a3] format
@@ -1831,9 +1831,9 @@ def _cast_to_index(p):
 
     for_op_scf = ForOp(lower_bound, upper_bound, step, iter_args=loop_args)
 
-    name_stack = jax_ctx.module_context.name_stack.extend("for")
+    name_stack = jax_ctx.name_stack.extend("for")
     body_block = for_op_scf.body
-    body_ctx = jax_ctx.module_context.replace(name_stack=name_stack.extend("body"))
+    body_ctx = jax_ctx.replace(name_stack=name_stack.extend("body"))
 
     with ir.InsertionPoint(body_block):
         body_args = list(body_block.arguments)
@@ -1858,8 +1858,9 @@ def _cast_to_index(p):
 
         # Recursively generate the mlir for the loop body
         out, _ = mlir.jaxpr_subcomp(
-            body_ctx,
+            body_ctx.module_context,
             body_jaxpr.jaxpr,
+            body_ctx.name_stack,
             mlir.TokenSet(),
             [mlir.ir_constants(c) for c in body_jaxpr.consts],
             *body_args,
@@ -1935,10 +1936,9 @@ def _adjoint_lowering(
     with ir.InsertionPoint(adjoint_block):
         source_info_util.extend_name_stack("adjoint")
         out, _ = mlir.jaxpr_subcomp(
-            jax_ctx.module_context.replace(
-                name_stack=jax_ctx.module_context.name_stack.extend("adjoint")
-            ),
+            jax_ctx.module_context,
             jaxpr.jaxpr,
+            jax_ctx.name_stack.extend("adjoint"),
             mlir.TokenSet(),
             [mlir.ir_constants(c) for c in jaxpr.consts],
             *([a] for a in chain(consts, cargs, adjoint_block.arguments)),  # [3]