Another take at enabling Pallas GPU tests on x64

Note that for_loop_p no longer assumes that the loop index is an int32.

Closes jax-ml#18847

Author: superbobry

jax/_src/core.py

@@ -2076,14 +2076,16 @@ def raise_to_shaped(aval: AbstractValue, weak_type=None):
     if handler: return handler(aval, weak_type)
   raise TypeError(type(aval))
 
-raise_to_shaped_mappings : dict[type, Callable] = {
-  AbstractToken: lambda aval, _: aval,
-  Bot: lambda aval, _: aval,
-  UnshapedArray: lambda aval, _: aval,
-  ShapedArray: lambda aval, weak_type: ShapedArray(
-      aval.shape, aval.dtype, weak_type, aval.named_shape),
-  DConcreteArray: lambda aval, weak_type: DShapedArray(
-      aval.shape, aval.dtype, weak_type),
+raise_to_shaped_mappings: dict[type, Callable] = {
+    AbstractToken: lambda aval, _: aval,
+    Bot: lambda aval, _: aval,
+    UnshapedArray: lambda aval, _: aval,
+    ShapedArray: lambda aval, weak_type: ShapedArray(
+        aval.shape, aval.dtype, weak_type, aval.named_shape
+    ),
+    DConcreteArray: lambda aval, weak_type: DShapedArray(
+        aval.shape, aval.dtype, weak_type
+    ),
 }
 
 ### Operations on shapes and dimension sizes.

jax/_src/lax/control_flow/for_loop.py

@@ -132,7 +132,7 @@ def wrapped_body(i, refs):
   nsteps, = nsteps
   flat_state, state_tree = tree_flatten(init_state)
   state_avals = map(state_utils.val_to_ref_aval, flat_state)
-  idx_aval = core.ShapedArray((), jnp.dtype("int32"))
+  idx_aval = core.ShapedArray((), dtypes.canonicalize_dtype(jnp.int64))
   jaxpr, consts, out_tree = _trace_to_jaxpr_with_refs(
       body, state_tree, [idx_aval, *state_avals])
   if out_tree != tree_structure(None):
@@ -251,7 +251,7 @@ def body(i, state):
 
 def _for_impl_unrolled(body, nsteps, unroll, *args):
   remainder = nsteps % unroll
-  i = jnp.int32(0)
+  i = jnp.astype(0, dtypes.canonicalize_dtype(jnp.int64))
   state = list(args)
 
   for _ in range(remainder):
@@ -748,15 +748,15 @@ def discharged_for_loop(nsteps, body, init_state, *, reverse: bool = False):
   """
   flat_state, state_tree = tree_flatten(init_state)
   state_avals = map(state_utils.val_to_ref_aval, flat_state)
-  idx_aval = core.ShapedArray((), jnp.dtype("int32"))
+  idx_aval = core.ShapedArray((), dtypes.canonicalize_dtype(jnp.int64))
   jaxpr, consts, out_tree = _trace_to_jaxpr_with_refs(
       body, state_tree, [idx_aval, *state_avals])
   if out_tree != tree_structure(None):
     raise Exception("`body` should not return anything.")
   discharged_jaxpr, discharged_consts = discharge_state(jaxpr, consts)
 
   def fori_body(i, carry):
-    i = jnp.int32(i)
+    i = jnp.astype(i, dtypes.canonicalize_dtype(jnp.int64))
     if reverse:
       i = nsteps - i - 1
     out_flat = core.eval_jaxpr(discharged_jaxpr, discharged_consts,

jax/_src/pallas/pallas_call.py

@@ -274,7 +274,7 @@ def body(carry):
           len(blocks),
           len(scratch_values),
       )
-      blocks = jax.core.eval_jaxpr(discharged_jaxpr, discharged_consts, *scalars,
+      blocks = jax_core.eval_jaxpr(discharged_jaxpr, discharged_consts, *scalars,
                                    *blocks, *scratch)
     blocks = blocks[grid_mapping.num_index_operands:]
     blocks, out_scratch = split_list(blocks, [num_inout])
@@ -787,7 +787,7 @@ def checked_kernel_fn(*args):
     # errors before other arguments.
     jaxpr_args = [*input_error_vals, *scalars, *inputs, *outputs, *scratch]
     assert len(checked_jaxpr.jaxpr.invars) == len(jaxpr_args)
-    result_flat = jax.core.eval_jaxpr(
+    result_flat = jax_core.eval_jaxpr(
         checked_jaxpr.jaxpr, checked_jaxpr.consts, *jaxpr_args)
     output_errors, _ = split_list(result_flat, [num_err_vals])
     # Store new errors back in the error refs.

jax/_src/pallas/primitives.py

@@ -26,6 +26,7 @@
 from jax import tree_util
 from jax._src import ad_util
 from jax._src import core as jax_core
+from jax._src import dtypes
 from jax._src import effects
 from jax._src import pretty_printer as pp
 from jax._src import state
@@ -359,7 +360,12 @@ def _load_discharge_rule(in_avals, out_avals, *args_flat, args_tree, **_):
     # of bounds, it will instead move the start_index backwards so the slice
     # will fit in memory.
     ref = _pad_values_to_avoid_dynamic_slice_oob_shift(ref, slice_sizes)
-    out_ones = lax.dynamic_slice(ref, slice_starts, slice_sizes=slice_sizes)
+    idx_dtype = dtypes.canonicalize_dtype(jnp.int64)
+    out_ones = lax.dynamic_slice(
+      ref,
+      [jnp.astype(s, idx_dtype) for s in slice_starts],
+      slice_sizes=slice_sizes,
+    )
     out_indexer = tuple(0 if scalar else slice(None) for scalar in scalar_dims)
     out = out_ones[out_indexer]
   elif all(not isinstance(s, Slice) for s in idx.indices):

jax/_src/pallas/triton/BUILD

@@ -52,6 +52,7 @@ pytype_strict_library(
         "//jax",
         "//jax:ad_util",
         "//jax:api_util",
+        "//jax:config",
         "//jax:core",
         "//jax:mlir",
         "//jax:partial_eval",

jax/_src/pallas/triton/lowering.py

@@ -29,6 +29,7 @@
 from jax._src import ad_checkpoint
 from jax._src import ad_util
 from jax._src import api_util
+from jax._src import config
 from jax._src import core as jax_core
 from jax._src import custom_derivatives
 from jax._src import linear_util as lu
@@ -2263,9 +2264,10 @@ def _for_lowering_rule(
   del which_linear
   if reverse or unroll != 1:
     raise NotImplementedError
-  lower_bound = _i32_constant(0)
-  upper_bound = _i32_constant(nsteps)
-  step = _i32_constant(1)
+  _i_constant = _i64_constant if config.enable_x64.value else _i32_constant
+  lower_bound = _i_constant(0)
+  upper_bound = _i_constant(nsteps)
+  step = _i_constant(1)
   init_args = map(_ensure_ir_value, args, ctx.avals_in)
   # Partially discharge state from jaxpr for non-pointers
   should_discharge = [

jax/_src/state/primitives.py

@@ -190,7 +190,7 @@ def _swap_abstract_eval(ref_aval: AbstractRef,
                        f"Expected shape: {expected_out_shape}. "
                        f"Value shape: {val_aval.shape}. "
                        f"Indices: {indexers}. ")
-    if ref_aval.dtype != val_aval.dtype:
+    if ref_aval.dtype != val_aval.dtype and not val_aval.weak_type:
       raise ValueError("Invalid dtype for `swap`. "
                        f"Ref dtype: {ref_aval.dtype}. "
                        f"Value dtype: {val_aval.dtype}. ")

jax/_src/state/types.py

@@ -132,6 +132,12 @@ class AbstractRef(core.AbstractValue):
   def __init__(self, inner_aval: core.AbstractValue):
     self.inner_aval = inner_aval
 
+  @property
+  def weak_type(self) -> bool:
+    if not hasattr(self.inner_aval, "weak_type"):
+      raise AttributeError
+    return self.inner_aval.weak_type
+
   def update(self, inner_aval=None):
     if inner_aval is None:
       return AbstractRef(self.inner_aval)

tests/pallas/BUILD

@@ -41,8 +41,6 @@ jax_test(
     disable_configs = [
         "gpu",
         "gpu_x32",
-        "gpu_a100",
-        "gpu_h100",
         "gpu_p100",
         "gpu_p100_x32",
     ],