Repository: google/jax
Pull request #22019 (2024-06-21 01:23:52)

[Pallas] Implement block-invariant sampling.
[Pallas] Implement block-invariant sampling.



Patch:
diff --git a/jax/BUILD b/jax/BUILD
index 3105256eb..b3d85106b 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -205,6 +205,7 @@ py_library_providing_imports_info(
         "_src/ad_checkpoint.py",
         "_src/api.py",
         "_src/array.py",
+        "_src/blocked_sampler.py",
         "_src/callback.py",
         "_src/checkify.py",
         "_src/custom_batching.py",
diff --git a/jax/_src/blocked_sampler.py b/jax/_src/blocked_sampler.py
new file mode 100644
index 000000000..66d0f53ab
--- /dev/null
+++ b/jax/_src/blocked_sampler.py
@@ -0,0 +1,164 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Protocol, Sequence
+import jax
+from jax._src import random
+from jax._src.typing import Array, ArrayLike
+from jax import numpy as jnp
+
+NdKeyList = Any
+Shape = random.Shape
+
+class SampleFn(Protocol):
+  def __call__(self, key: random.KeyArrayLike, *args, shape: Shape,
+               **kwargs) -> Array:
+    ...
+
+
+def _compute_scalar_index(iteration_index: Sequence[int],
+                          total_size: Shape,
+                          block_size: Shape,
+                          block_index: Sequence[int]) -> int:
+  ndims = len(iteration_index)
+  dim_size = 1
+  total_idx = 0
+  for i in range(ndims-1, -1, -1):
+    dim_idx = block_index[i] + iteration_index[i] * block_size[i]
+    total_idx += dim_idx * dim_size
+    dim_size *= total_size[i]
+  return total_idx
+
+
+def blocked_fold_in(
+  global_key: random.KeyArrayLike,
+  total_size: Shape,
+  block_size: Shape,
+  tile_size: Shape,
+  block_index: Sequence[ArrayLike],
+  ) -> NdKeyList:
+  """Computes a grid of keys for block-invariant sampling.
+
+  Suppose we wished to construct a 16x512 array of random numbers, using
+  block sizes of 16x128 and 16x256. We could select an tile size of 8x128
+  (which divides both 16x128 and 16x256) and divide the total array in tiles as:
+  ---------------------------------
+  | 8x128 | 8x128 | 8x128 | 8x128 |
+  ---------------------------------
+  | 8x128 | 8x128 | 8x128 | 8x128 |
+  ---------------------------------
+
+  We generate a key for each tile as:
+    tile_key = fold_in(global_key, tile_idx)
+
+  Where the tile_idx is the row-major raveled index of each element:
+  -----------------
+  | 0 | 1 | 2 | 3 |
+  -----------------
+  | 4 | 5 | 6 | 7 |
+  -----------------
+
+  We then compute and return the keys required to sample the tiles that make
+  up the current block (specified via `block_index`).
+  With a 16x256 block size, each block requires 4 (2x2) tile keys:
+  ---------------
+  | 0, 1 | 2, 3 |
+  | 4, 5 | 6, 7 |
+  ---------------
+  Therefore, we return a grid of 2x2 keys for each block (2 blocks total).
+
+  With a 16x128 block size, each block requires 2 (2x1) tile keys:
+  -----------------
+  | 0 | 1 | 2 | 3 |
+  | 4 | 5 | 6 | 7 |
+  -----------------
+  Therefore, we return a grid of 2x1 keys for each block (4 blocks total).
+
+  Args:
+    global_key: The global key shared between all blocks.
+    total_size: The shape of the array being generated.
+    block_size: The shape of an individual block.
+    tile_size: The shape of a `tile`, which is the smallest unit at
+      which samples are generated. This should be selected to be a divisor
+      of all block sizes one needs to be invariant to.
+    block_index: The index denoting which block to generate keys for.
+
+  Returns:
+    An N-dimensional nested list of keys required to sample the tiles
+    corresponding to the block specified by `block_index`.
+  """
+  size_in_blocks = tuple(
+      _shape // _element for _shape, _element in zip(block_size, tile_size))
+
+  def _keygen_loop(axis, prefix):
+    if axis == len(size_in_blocks):
+      subtile_key = jax.random.fold_in(
+          global_key, _compute_scalar_index(
+              block_index, total_size, size_in_blocks, prefix))
+      return subtile_key
+    else:
+      keys = []
+      for i in range(size_in_blocks[axis]):
+        keys.append(_keygen_loop(axis+1, prefix+(i,)))
+      return keys
+  return _keygen_loop(0, tuple())
+
+
+def sample_block(
+    sampler_fn: SampleFn,
+    keys: NdKeyList,
+    block_size: Shape,
+    tile_size: Shape,
+    *args,
+    **kwargs
+  ) -> jax.Array:
+  """Draws random samples for a single block.
+
+  This function is intended to be used in conjunction with `blocked_fold_in`:
+  ```
+  key_list = blocked_fold_in(global_key, total_size, block_size, tile_size,
+                             block_index)
+  samples = sample_block(jax.random.uniform, key_list, block_size, tile_size)
+  ```
+
+  Args:
+    sampler_fn: A random sampling function, e.g. jax.random.uniform.
+    keys: A grid of keys generated by `blocked_fold_in`.
+    block_size: The shape of an individual block.
+    tile_size: The shape of a `tile`, which is the smallest unit at
+      which samples are generated. This should be selected to be a divisor
+      of all block sizes one needs to be invariant to.
+    args: varargs for sampler_fn.
+    kwargs: kwargs for sampler_fn.
+
+  Returns:
+    An array of random samples drawn using sampler_fn.
+  """
+  size_in_tiles = tuple(
+      _shape // _element for _shape, _element in zip(block_size, tile_size))
+  def _nested_index(arr: jax.Array, idx: Sequence[int]) -> jax.Array:
+    if len(idx) == 1:
+      return arr[idx[0]]
+    return _nested_index(arr[idx[0]], idx[1:])
+
+  def _sample_loop(axis: int, prefix: tuple[int, ...]) -> jax.Array:
+    if axis == len(size_in_tiles):
+      return sampler_fn(_nested_index(keys, prefix), *args,
+                        shape=tile_size, **kwargs)
+    else:
+      samples = []
+      for i in range(size_in_tiles[axis]):
+        samples.append(_sample_loop(axis+1, prefix+(i,)))
+      return jnp.concatenate(samples, axis=axis)
+  return _sample_loop(0, tuple())
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index 375691aea..853a97a86 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -1005,14 +1005,15 @@ def _prng_key_load_lowering_rule(ctx: LoweringRuleContext, *args_flat, args_tree
   assert isinstance(aval_out.dtype, prng.KeyTy)
   ref_block_shape = aval_out.dtype._impl.key_shape
 
-  if len(ref_block_shape) != 1:
-    raise NotImplementedError("Seed key_data must be 1D.")
-  if ref_block_shape[0] != 1:
-    raise NotImplementedError("Seed key_data of shape != (1,) not supported.")
+  if len(ref_block_shape) != 2:
+    raise NotImplementedError("Seed key_data must be 2D.")
+  if tuple(ref_block_shape) != (1, 1):
+    raise NotImplementedError(
+      f"Seed key_data of shape != (1, 1) not supported. Got: {ref_block_shape}")
 
   load_ops = []
   for i in range(ref_block_shape[0]):
-    idx = NDIndexer(indices=(i,), shape=ref_block_shape,
+    idx = NDIndexer(indices=(0, i), shape=ref_block_shape,
                     int_indexer_shape=tuple())
     starts, _, _, _, _ = _indexer_to_start_size_stride(
         idx,
@@ -2478,7 +2479,8 @@ def _prng_seed_lowering_rule(ctx: LoweringRuleContext, *seeds):
   # takes in a list of integers as input.
   all_integers = all(isinstance(seed.type, ir.IntegerType) for seed in seeds)
   if not all_integers:
-    raise ValueError("All seed data must be integers.")
+    seed_types = [seed.type for seed in seeds]
+    raise ValueError(f"All seed data must be scalar integers. Got {seed_types}")
   return tpu.PRNGSeed32Op(seeds).results
 lowering_rules[tpu_primitives.prng_seed_p] = _prng_seed_lowering_rule
 
@@ -2531,12 +2533,14 @@ def random_wrap_lowering(ctx, key_data, *, impl):
     # If the key data lives in vregs, need to unpack it to sregs.
     key_data_list = []
     key_data_shape = key_data.type.shape
-    if len(key_data_shape) != 1:
-      raise NotImplementedError("Seed key_data must be 1D.")
-    if key_data_shape[0] != 1:
-      raise NotImplementedError("key_data with shape != (1,) not supported.")
-    for i in range(key_data_shape[0]):
-      key_data_list.append(vector.ExtractOp(key_data, [], [i]))
+    if len(key_data_shape) != 2:
+      raise NotImplementedError("Seed key_data must be 2D.")
+    if tuple(key_data_shape) != (1, 1):
+      raise NotImplementedError(
+        "Seed key_data of shape != (1, 1) not supported. "
+        f"Got: {key_data_shape}")
+    for i in range(key_data_shape[1]):
+      key_data_list.append(vector.ExtractOp(key_data, [], [0, i]))
     return KeyScalarBundle(scalars=key_data_list)
   if isinstance(key_data, KeyScalarBundle):
     return key_data
diff --git a/jax/_src/pallas/mosaic/random.py b/jax/_src/pallas/mosaic/random.py
index 9e07d125a..295196bc9 100644
--- a/jax/_src/pallas/mosaic/random.py
+++ b/jax/_src/pallas/mosaic/random.py
@@ -11,20 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable
+from typing import Callable, Optional
 
 import jax
 import numpy as np
 from jax import numpy as jnp
 from jax import random as jax_api_random
+from jax._src import blocked_sampler
 from jax._src import typing
 from jax._src.pallas.mosaic.primitives import prng_seed
 from jax._src.pallas.mosaic.primitives import prng_random_bits
+from jax._src.pallas import primitives
 from jax._src import prng as jax_prng
 
 
 Shape = jax_prng.Shape
-SampleFnType = Any
+SampleFnType = blocked_sampler.SampleFn
 KeylessSampleFnType = Callable[..., jax.Array]
 
 set_seed = prng_seed
@@ -32,6 +34,7 @@ set_seed = prng_seed
 FOLD_IN_ROUNDS = 128
 SUPPORTED_CONVERSION_KEYS = ["rbg", "unsafe_rbg", "pallas_tpu"]
 
+
 def to_pallas_key(key: jax_prng.PRNGKeyArray) -> jax_prng.PRNGKeyArray:
   """Helper function for converting non-Pallas PRNG keys into Pallas keys."""
 
@@ -79,15 +82,16 @@ def _fold_in(key: jax_prng.PRNGKeyArray, data: typing.Array):
 
   mixed = key_bits[..., FOLD_IN_ROUNDS-1] ^ data_bits[..., FOLD_IN_ROUNDS-1]
   assert mixed.shape == key_shape
-  impl: jax_prng.PRNGSpec = jax.random.key_impl(key)  # type: ignore
-  return jax.random.wrap_key_data(mixed, impl=impl)
+  return jax.random.wrap_key_data(mixed, impl="pallas_tpu")
 
 def _split(key: typing.Array, shape: Shape):
   del key, shape
   raise NotImplementedError()
 
 tpu_key_impl = jax_prng.PRNGImpl(
-   key_shape=(1,),
+   # Pallas currently only supports 2D+ windows, so set the key_shape
+   # to be 2D to have better compatibility with setting BlockSpecs.
+   key_shape=(1, 1),
    seed=_seed_func,
    split=_split,
    random_bits=_random_bits,
@@ -152,11 +156,63 @@ def _make_stateful_sampler(sampler: SampleFnType) -> KeylessSampleFnType:
     placeholder_key = jax_api_random.key(0, impl=tpu_internal_stateful_impl)
     return sampler(placeholder_key, *args, **kwargs)
   # Remove key argument from docstring.
-  doc_lines = filter(
-      lambda line: "key:" not in line, sampler.__doc__.split("\n"))
-  new_sampler.__doc__ = "\n".join(doc_lines)
+  if sampler.__doc__:
+    doc_lines = filter(
+        lambda line: "key:" not in line, sampler.__doc__.split("\n"))
+    new_sampler.__doc__ = "\n".join(doc_lines)
   return new_sampler
 
-bits = _make_stateful_sampler(jax_api_random.bits)
-uniform = _make_stateful_sampler(jax_api_random.uniform)
-bernoulli = _make_stateful_sampler(jax_api_random.bernoulli)
+bits = _make_stateful_sampler(jax_api_random.bits)  # type: ignore
+uniform = _make_stateful_sampler(jax_api_random.uniform)  # type: ignore
+bernoulli = _make_stateful_sampler(jax_api_random.bernoulli)  # type: ignore
+
+
+def sample_block(sampler_fn: SampleFnType,
+                 global_key: jax_prng.PRNGKeyArray,
+                 block_size: Shape,
+                 tile_size: Shape,
+                 total_size: Shape,
+                 block_index: Optional[tuple[typing.ArrayLike, ...]] = None,
+                 **kwargs) -> jax.Array:
+  """Samples a block of random values with invariance guarantees.
+
+  `sample_block` allows the sampling of identical blocks of random values
+  across kernels with different block shapes and iteration orders. Each call
+  to `sample_block` returns a `block_size`-shaped array of random samples
+  corresponding to the `block_index`.
+
+  `tile_size` should be chosen such that it is a divisor to all block sizes
+  one needs to be invariant to. The larger the `tile_size`, the more
+  efficient the sampling process wil be and therefore the best choice is
+  typically the greatest common divisor between all possible block sizes.
+
+  Args:
+    sampler_fn: A sampling function that consumes a key and returns
+      random samples.
+    global_key: The global key to use for sampling.
+    block_size: The shape of an individual block.
+    tile_size: The shape of a `tile`, which is the smallest unit at
+      which samples are generated. This should be selected to be a divisor
+      of all block sizes one needs to be invariant to.
+    total_size: The total size of the array to sample.
+    block_index: The index denoting which block to generate keys for. Defaults
+      to the program_id for each block axis.
+    **kwargs: Additional arguments to pass to the sampler_fn.
+
+  Returns:
+    A `block_size` shaped array of samples for the current block corresponding
+    to `block_index`.
+  """
+  if len(block_size) != len(tile_size):
+    raise ValueError(f"block_size ({len(block_size)}) and tile_size "
+                     f"({len(tile_size)}) must have the same length.")
+
+  if block_index is None:
+    num_axes = len(block_size)
+    block_index = tuple(
+      primitives.program_id(axis) for axis in range(num_axes))
+
+  keys = blocked_sampler.blocked_fold_in(
+      global_key, total_size, block_size, tile_size, block_index)
+  return blocked_sampler.sample_block(
+      sampler_fn, keys, block_size, tile_size, **kwargs)


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index c420a3ba9..4a3c05f1c 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1087,6 +1087,11 @@ jax_test(
     main = "third_party/scipy/line_search_test.py",
 )
 
+jax_test(
+    name = "blocked_sampler_test",
+    srcs = ["blocked_sampler_test.py"],
+)
+
 py_test(
     name = "tree_util_test",
     srcs = ["tree_util_test.py"],
diff --git a/tests/blocked_sampler_test.py b/tests/blocked_sampler_test.py
new file mode 100644
index 000000000..1f8f2b645
--- /dev/null
+++ b/tests/blocked_sampler_test.py
@@ -0,0 +1,90 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax import numpy as jnp
+from jax._src import blocked_sampler
+from jax._src import config
+from jax._src import test_util as jtu
+import numpy as np
+
+
+config.parse_flags_with_absl()
+
+
+def call_kernel(
+  kernel,
+  grid: tuple[int, int],
+  transpose_grid: bool,
+  *args
+  ):
+  """Calls a kernel over a grid and concatenates results to a single array."""
+  if transpose_grid:
+    grid = (grid[1], grid[0])
+  m, n = grid
+  return jnp.concatenate([
+      jnp.concatenate([
+          kernel(i, j, *args) for j in range(n)], axis=1)
+      for i in range(m)], axis=0)
+
+
+def uniform_kernel(i: int, j: int, total_size, block_size, tile_size):
+  """Uniform random sampling kernel function."""
+  global_key = jax.random.key(0)
+  keys = blocked_sampler.blocked_fold_in(global_key,
+                                         total_size=total_size,
+                                         block_size=block_size,
+                                         tile_size=tile_size,
+                                         block_index=(i, j))
+  return blocked_sampler.sample_block(jax.random.uniform,
+                                      keys,
+                                      block_size=block_size,
+                                      tile_size=tile_size,
+                                      minval=0.0, maxval=1.0)
+
+
+class BlockedSamplerTest(jtu.JaxTestCase):
+
+  @parameterized.named_parameters(
+    dict(testcase_name='8x128_vs_16x256', total_size=(32, 256),
+         block_size_a=(8, 128), block_size_b=(16, 256),
+         tile_size=(8, 128), transpose_grid=False),
+    dict(testcase_name='transpose_8x128_vs_16x256', total_size=(32, 256),
+         block_size_a=(8, 128), block_size_b=(16, 256),
+         tile_size=(8, 128), transpose_grid=True),
+    dict(testcase_name='8x128_vs_32x128', total_size=(32, 128),
+         block_size_a=(8, 128), block_size_b=(32, 128),
+         tile_size=(8, 128), transpose_grid=False),
+    dict(testcase_name='16x256_vs_32x128', total_size=(32, 256),
+         block_size_a=(16, 256), block_size_b=(32, 128),
+         tile_size=(8, 128), transpose_grid=False),
+  )
+  def test_block_shape_invariance(self, total_size, block_size_a,
+                          block_size_b, tile_size, transpose_grid):
+    grid_a = tuple(_tot // _blk for _tot, _blk in zip(total_size, block_size_a))
+    result_a = call_kernel(
+        uniform_kernel, grid_a, transpose_grid,
+        total_size, block_size_a, tile_size)
+
+    grid_b = tuple(_tot // _blk for _tot, _blk in zip(total_size, block_size_b))
+    result_b = call_kernel(
+        uniform_kernel, grid_b, transpose_grid,
+        total_size, block_size_b, tile_size)
+    np.testing.assert_array_equal(result_a, result_b)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/pallas/tpu/pallas_random_test.py b/tests/pallas/tpu/pallas_random_test.py
index 0a6c22dea..64892b707 100644
--- a/tests/pallas/tpu/pallas_random_test.py
+++ b/tests/pallas/tpu/pallas_random_test.py
@@ -18,7 +18,6 @@ from absl.testing import parameterized
 import jax
 from jax import random as jax_random
 from jax._src import test_util as jtu
-from jax._src.pallas.mosaic import core as tpu_core
 from jax._src.pallas.mosaic import random as plrandom
 from jax.experimental import pallas as pl
 from jax.experimental.pallas import tpu as pltpu
@@ -100,7 +99,7 @@ class PRNGTest(jtu.JaxTestCase):
     o_shape = jax.ShapeDtypeStruct((8, 128), jnp.float32)
     result = pl.pallas_call(
         body,
-        in_specs=[pl.BlockSpec(memory_space=tpu_core.TPUMemorySpace.SMEM)],
+        in_specs=[pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM)],
         out_shape=o_shape,
     )(key)
     self.assertGreaterEqual(jnp.min(result), 0)
@@ -118,7 +117,7 @@ class PRNGTest(jtu.JaxTestCase):
     o_shape = jax.ShapeDtypeStruct((8, 128), jnp.float32)
     result = pl.pallas_call(
         body,
-        in_specs=[pl.BlockSpec(memory_space=tpu_core.TPUMemorySpace.SMEM)],
+        in_specs=[pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM)],
         out_shape=o_shape,
     )(key)
     self.assertGreaterEqual(jnp.min(result), 0)
@@ -142,7 +141,7 @@ class PRNGTest(jtu.JaxTestCase):
     o_shape = jax.ShapeDtypeStruct((2, 8, 128), jnp.float32)
     result = pl.pallas_call(
         body,
-        in_specs=[pl.BlockSpec(memory_space=tpu_core.TPUMemorySpace.SMEM)],
+        in_specs=[pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM)],
         out_shape=o_shape,
     )(key)
     result_a = result[0]
@@ -150,5 +149,55 @@ class PRNGTest(jtu.JaxTestCase):
     np.testing.assert_array_compare(np.not_equal, result_a, result_b)
 
 
+class BlockInvarianceTest(parameterized.TestCase):
+
+  def setUp(self):
+    if not jtu.test_device_matches(["tpu"]):
+      self.skipTest("Need TPU devices")
+    super().setUp()
+
+  def test_block_invariance(self):
+
+    def make_kernel_body(index_map):
+      def body(key_ref, o_ref):
+        key = key_ref[0, 0]
+        samples = plrandom.sample_block(
+            jax.random.uniform,
+            key,
+            block_size=o_ref[...].shape,
+            tile_size=(16, 128),
+            total_size=(64, 512),
+            block_index=index_map(pl.program_id(0), pl.program_id(1)),
+            minval=0.0,
+            maxval=1.0)
+        o_ref[...] = samples
+      return body
+
+    global_key = jax_random.key(0, impl="pallas_tpu")
+    o_shape = jnp.ones((64, 512), dtype=jnp.float32)
+    key_spec = pl.BlockSpec(lambda i, j: (0, 0),
+                            block_shape=(1, 1),
+                            memory_space=pltpu.TPUMemorySpace.SMEM)
+    out_spec = pl.BlockSpec(lambda i, j: (i, j), block_shape=(16, 128))
+    result_16x128 = pl.pallas_call(
+        make_kernel_body(index_map=lambda i, j: (i, j)),
+        out_shape=o_shape,
+        in_specs=[key_spec],
+        out_specs=out_spec,
+        grid=(4, 4),
+    )(global_key)
+
+    out_spec = pl.BlockSpec(lambda i, j: (j, i), block_shape=(32, 256))
+    result_32x256 = pl.pallas_call(
+        make_kernel_body(index_map=lambda i, j: (j, i)),
+        in_specs=[key_spec],
+        out_shape=o_shape,
+        out_specs=out_spec,
+        grid=(2, 2),
+    )(global_key)
+    np.testing.assert_array_equal(result_16x128, result_32x256)
+
+
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())