# Copyright 2025 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. import functools import itertools from absl.testing import absltest from absl.testing import parameterized import jax from jax import lax from jax._src import config from jax._src import state from jax._src import test_util as jtu from jax._src.pallas import einshape as einshape_lib from jax._src.pallas import primitives as pallas_primitives from jax._src.pallas.fuser import block_spec as block_spec_lib from jax._src.pallas.fuser import custom_fusion_lib from jax.experimental import checkify from jax.experimental import pallas as pl import jax.numpy as jnp import numpy as np jax.config.parse_flags_with_absl() class PullBlockSpecTest(jtu.JaxTestCase): def setUp(self): super().setUp() if config.enable_x64.value: self.skipTest('x64 not supported') def test_identity(self): def f(x): return x in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) # We should have no new values or scalar prefetch values. self.assertEmpty(scalar_prefetch_values) self.assertEmpty(value_block_specs) self.assertEqual(in_block_spec.block_shape, (128, 128)) self.assertEqual(in_block_spec.index_map(0, 1, 2), (0, 1)) x = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, new_values, x), x, ) def test_pull_block_spec_with_none_in_pytree(self): def f(vals, g): x, none_val = vals assert none_val is None return g + x in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) out_spec = pl.BlockSpec((128, 128), lambda i, j: (i, j)) kernel_fn, (in_specs, g_spec), _ = block_spec_lib.pull_block_spec( f, out_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )((in_type, None), in_type) self.assertEqual(in_specs[0].block_shape, (128, 128)) self.assertIsNone(in_specs[1]) self.assertEqual(g_spec.block_shape, (128, 128)) x = np.ones((128, 128), dtype=np.float32) g = 2.0 * np.ones((128, 128), dtype=np.float32) out = kernel_fn((0, 0), (), (x, None), g) np.testing.assert_array_equal(out, x + g) @parameterized.named_parameters( ('transpose', 'ab->ba', (32, 32), (16, 16), {}, (16, 16)), ( 'bcnt_to_bn_ct', 'bcnt->(bn)(ct)', (2, 4, 8, 16), (8, 32), {}, (1, 2, 8, 16), ), ( 'split_dims', '(ab)c->abc', (32, 64), (2, 8, 16), {'a': 4, 'b': 8}, (16, 16), ), ('merge_dims', 'abc->(ab)c', (4, 8, 64), (16, 16), {}, (2, 8, 16)), ( 'split_transpose_merge', 'a(bc)->c(ba)', (16, 32), (8, 16), {'b': 4, 'c': 8}, (16, 8), ), ( 'multi_split_merge', '(ab)(cd)->(ac)(bd)', (8, 128), (2, 64), {'a': 2, 'b': 4, 'c': 8, 'd': 16}, (4, 32), ), ) def test_einshape_in_input_fusion( self, equation, in_shape, out_block_shape, sizes, expected_in_block_shape ): def f(x): return einshape_lib.einshape(equation, x, **sizes) in_type = jax.ShapeDtypeStruct(in_shape, jnp.float32) f2, new_values, _ = block_spec_lib.get_fusion_values(f, in_type) out_eval = jax.eval_shape( lambda v: einshape_lib.einshape(equation, v, **sizes), in_type ) grid = tuple(s // bs for s, bs in zip(out_eval.shape, out_block_shape)) block_spec = pl.BlockSpec( block_shape=out_block_shape, index_map=lambda *pids: pids ) _, (_, in_block_spec), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=len(grid), scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) self.assertEqual(in_block_spec.block_shape, expected_in_block_shape) def test_einshape_in_input_fusion_non_contiguous_raises(self): def f(x): return einshape_lib.einshape('(ab)c->abc', x, a=4, b=8) in_type = jax.ShapeDtypeStruct((32, 64), jnp.float32) f2, new_values, _ = block_spec_lib.get_fusion_values(f, in_type) out_block_shape = (2, 4, 16) block_spec = pl.BlockSpec( block_shape=out_block_shape, index_map=lambda *pids: pids ) with self.assertRaisesRegex( NotImplementedError, 'SplitDims slice .* is non-contiguous' ): block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) def test_einshape_in_input_fusion_unblocked_split(self): def f(x): return einshape_lib.einshape('(ab)c->abc', x, a=4, b=8) in_type = jax.ShapeDtypeStruct((32, 64), jnp.float32) f2, new_values, _ = block_spec_lib.get_fusion_values(f, in_type) out_block_shape = (None, None, 16) block_spec = pl.BlockSpec( block_shape=out_block_shape, index_map=lambda *pids: pids ) _, (_, in_block_spec), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) self.assertEqual(in_block_spec.block_shape, (None, 16)) def test_einshape_in_input_fusion_partial_unblocked_split(self): def f(x): return einshape_lib.einshape('(ab)c->abc', x, a=4, b=8) in_type = jax.ShapeDtypeStruct((32, 64), jnp.float32) f2, new_values, _ = block_spec_lib.get_fusion_values(f, in_type) out_block_shape = (None, 8, 16) block_spec = pl.BlockSpec( block_shape=out_block_shape, index_map=lambda *pids: pids ) _, (_, in_block_spec), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) self.assertEqual(in_block_spec.block_shape, (32, 16)) def test_einshape_in_input_fusion_invalid_merge_block_size(self): def f(x): return einshape_lib.einshape('abc->(ab)c', x) in_type = jax.ShapeDtypeStruct((4, 8, 64), jnp.float32) f2, new_values, _ = block_spec_lib.get_fusion_values(f, in_type) out_block_shape = (64, 16) block_spec = pl.BlockSpec( block_shape=out_block_shape, index_map=lambda *pids: pids ) with self.assertRaises(NotImplementedError): block_spec_lib.pull_block_spec( f2, block_spec, grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) def test_const(self): x = np.ones((512, 512), dtype=np.float32) def f(): return x f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) self.assertEmpty(scalar_prefetch_values) self.assertEqual(value_block_specs[0].block_shape, (128, 128)) self.assertEqual(value_block_specs[0].index_map(0, 1, 2), (0, 1)) x_block = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn( (0, 0, 0), scalar_prefetch_values, (np.ones((128, 128), dtype=np.float32),), ), x_block, ) @parameterized.parameters([ jnp.exp, jnp.tanh, functools.partial(pallas_primitives.multiple_of, values=(32,)), ]) def test_elementwise(self, fn): def f(x): return fn(x) in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler( 0 ), )(new_values, in_type) ) self.assertEmpty(value_block_specs) self.assertEqual(in_block_spec.block_shape, (128, 128)) self.assertEqual(in_block_spec.index_map(0, 1, 2), (0, 1)) x = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (), x), fn(x), ) @parameterized.parameters([jnp.exp, jnp.tanh]) def test_elementwise_bias(self, fn): b = np.ones((512, 512), dtype=np.float32) def f(x): return fn(x) + b in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertLen(value_block_specs, 1) self.assertEqual(value_block_specs[0].block_shape, (128, 128)) self.assertEqual(value_block_specs[0].index_map(0, 1, 2), (0, 1)) self.assertEqual(in_block_spec.block_shape, (128, 128)) self.assertEqual(in_block_spec.index_map(0, 1, 2), (0, 1)) x = np.ones((128, 128), dtype=np.float32) b = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (b,), x), fn(x) + b, ) def test_custom_fusion(self): @custom_fusion_lib.custom_fusion def fn(x, y): return x + y fn.def_pull_block_spec(lambda bss: (bss[0], bss[0])) fn.def_push_block_spec(lambda bss: (bss[0],)) fn.def_eval_rule(lambda _, x, y: (fn(x, y),)) in_type = ( jax.ShapeDtypeStruct((512, 512), jnp.float32), jax.ShapeDtypeStruct((512, 512), jnp.float32), ) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( fn, *in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, *in_type) ) self.assertEmpty(value_block_specs) self.assertLen(in_block_specs, 2) x_block_spec, y_block_spec = in_block_specs self.assertEqual(x_block_spec.block_shape, (128, 128)) self.assertEqual( x_block_spec.index_map(0, 1, 2), block_spec.index_map(0, 1, 2) ) self.assertEqual(y_block_spec.block_shape, (128, 128)) self.assertEqual( y_block_spec.index_map(0, 1, 2), block_spec.index_map(0, 1, 2) ) x = np.ones((128, 128), dtype=np.float32) y = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, new_values, x, y), x + y ) def test_custom_fusion_prefetch(self): @custom_fusion_lib.custom_fusion def fn(x): return x * 2.0 fn.def_pull_block_spec(lambda bss: bss) fn.def_push_block_spec(lambda bss: bss) @fn.def_eval_rule def eval_rule(ctx, x): self.assertIsNotNone(ctx.scalar_prefetch) sp_handler = block_spec_lib.make_scalar_prefetch_handler(0) prefetch_val = sp_handler(*ctx.scalar_prefetch) return (fn(x) + prefetch_val,) in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( fn, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k, *sp: (i, j)) kernel_fn, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec(f2, block_spec, grid_len=3)( new_values, in_type ) ) self.assertEmpty(value_block_specs) self.assertEqual(in_block_spec.block_shape, (128, 128)) x = np.ones((128, 128), dtype=np.float32) prefetch_val = np.float32(5.0) scalar_prefetch_inputs = (prefetch_val,) result = kernel_fn((0, 0, 0), scalar_prefetch_inputs, new_values, x) np.testing.assert_array_equal(result, x * 2.0 + 5.0) @parameterized.product( fn=[ lax.mul, lax.add, lax.sub, lax.div, lax.max, lax.min, lax.lt, lax.eq, lax.gt, ], ) def test_binop(self, fn): in_type = ( jax.ShapeDtypeStruct((512, 512), jnp.float32), jax.ShapeDtypeStruct((512, 512), jnp.float32), ) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( fn, *in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, *in_type) ) self.assertEmpty(value_block_specs) self.assertLen(in_block_specs, 2) x_block_spec, y_block_spec = in_block_specs self.assertEqual(x_block_spec.block_shape, (128, 128)) self.assertEqual( x_block_spec.index_map(0, 1, 2), block_spec.index_map(0, 1, 2) ) self.assertEqual(y_block_spec.block_shape, (128, 128)) self.assertEqual( y_block_spec.index_map(0, 1, 2), block_spec.index_map(0, 1, 2) ) x = np.ones((128, 128), dtype=np.float32) y = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, new_values, x, y), fn(x, y), ) @parameterized.parameters(None, pl.Squeezed()) def test_binop_bcast_mapped_dim(self, squeezed): in_type = ( jax.ShapeDtypeStruct((128, 512), jnp.float32), jax.ShapeDtypeStruct((1, 512), jnp.float32), ) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( lax.add, *in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((squeezed, 128), lambda i, j: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, *in_type) ) self.assertEmpty(value_block_specs) self.assertLen(in_block_specs, 2) x_block_spec, y_block_spec = in_block_specs self.assertEqual(x_block_spec.block_shape, (squeezed, 128)) self.assertEqual(x_block_spec.index_map(2, 3), (2, 3)) self.assertEqual(y_block_spec.block_shape, (squeezed, 128)) self.assertEqual(y_block_spec.index_map(2, 3), (0, 3)) x = y = np.ones((128,), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0), scalar_prefetch_values, new_values, x, y), x + y ) def test_slice(self): x = jax.random.normal(jax.random.key(0), (4, 512, 512), dtype=np.float32) def f(): return jax.lax.slice(x, (1, 0, 0), (2, 512, 512)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((1, 128, 128), lambda i, j, k: (0, i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(0), )(new_values) self.assertLen(value_block_specs, 1) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (1, 128, 128)) self.assertEqual(x_block_spec.index_map(4, 2, 3), (1, 4, 2)) x = np.ones((1, 128, 128), dtype=np.float32) # Slice doesn't change value after pulling block spec. np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (x,)), x ) def test_slice_non_multiple_block_size(self): x = jax.random.normal(jax.random.key(0), (1, 12, 128), dtype=np.float32) def f(): return jax.lax.slice(x, (0, 0, 0), (1, 12, 128)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec((1, 8, 128), lambda i, j: (0, i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(0), )(new_values) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (1, 8, 128)) self.assertEqual(x_block_spec.index_map(1, 0), (0, 1, 0)) tile = np.ones((1, 8, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0), scalar_prefetch_values, (tile,)), tile ) def test_dynamic_slice_non_multiple_block_size(self): x = jnp.zeros((1, 12, 128), dtype=jnp.float32) i, j, k = jnp.int32(0), jnp.int32(0), jnp.int32(0) def f(): return jax.lax.dynamic_slice(x, (i, j, k), (1, 12, 128)) f2, new_values, sp = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec((1, 8, 128), lambda a, b, *_: (a, b, 0)) _, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) x_bs = value_block_specs[0] self.assertEqual(x_bs.block_shape, (1, 8, 128)) sp_mem = jax.tree.map(lambda v: v[None], sp) self.assertEqual(x_bs.index_map(0, 1, *sp_mem), (0, 1, 0)) def test_partial_slice_non_multiple_block_size_fails(self): x = jnp.zeros((1, 20, 128), dtype=jnp.float32) block_spec = pl.BlockSpec((1, 8, 128), lambda i, j: (0, i, j)) # Static slice fails at pull_block_spec call time def f_static(): return jax.lax.slice(x, (0, 0, 0), (1, 12, 128)) f2, new_values, _ = block_spec_lib.get_fusion_values(f_static) with self.assertRaisesRegex(ValueError, 'not a multiple of block_size'): block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) # Dynamic slice with literal 0 start def f_dynamic(): return jax.lax.dynamic_slice(x, (0, 0, 0), (1, 12, 128)) f2_dyn, new_vals_dyn, _ = block_spec_lib.get_fusion_values(f_dynamic) _, (val_bs,), _ = block_spec_lib.pull_block_spec( f2_dyn, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_vals_dyn) with self.assertRaisesRegex(ValueError, 'not a multiple of block_size'): val_bs[0].index_map(0, 0) def test_slice_block_size_larger_than_full_size(self): x = jnp.zeros((1, 4, 128), dtype=jnp.float32) def f(): return jax.lax.slice(x, (0, 0, 0), (1, 4, 128)) f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec((1, 8, 128), lambda i, j: (0, i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertEqual(value_block_specs[0].block_shape, (1, 8, 128)) tile = np.ones((1, 8, 128), dtype=np.float32) np.testing.assert_array_equal(kernel_fn((0, 0), (), (tile,)), tile) def test_squeeze(self): x = jax.random.normal(jax.random.key(0), (1, 512, 512), dtype=np.float32) def f(): return jnp.squeeze(x, axis=0) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (None, 128, 128)) self.assertEqual(x_block_spec.index_map(4, 2, 3), (0, 4, 2)) x = np.ones((128, 128), dtype=np.float32) # Squeeze doesn't change value after pulling block spec and the squeezed # dimensions are removed. np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (x,)), x ) def test_dynamic_slice_only(self): x = jax.random.normal(jax.random.key(0), (3, 4, 512, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) j = jnp.array(2, dtype=jnp.int32) def f(): return jax.lax.dynamic_slice(x, (i, j, 0, 0), (1, 1, 512, 512)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertLen(scalar_prefetch_values, 2) block_spec = pl.BlockSpec( (1, 1, 128, 128), lambda i, j, k, *_: (0, 0, i, j) ) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch_values = jax.tree.map( lambda x: x[None], scalar_prefetch_values ) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (1, 1, 128, 128)) self.assertEqual( x_block_spec.index_map(0, 1, 2, *scalar_prefetch_values), (1, 2, 0, 1) ) x = np.ones((1, 1, 128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (x,)), x ) def test_dynamic_slice_squeeze(self): x = jax.random.normal(jax.random.key(0), (3, 4, 512, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) j = jnp.array(2, dtype=jnp.int32) def f(): return jnp.squeeze( jax.lax.dynamic_slice(x, (i, j, 0, 0), (1, 1, 512, 512)), axis=(0, 1) ) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertLen(scalar_prefetch_values, 2) scalar_prefetch_values = jax.tree.map( lambda x: x[None], scalar_prefetch_values ) block_spec = pl.BlockSpec((128, 128), lambda i, j, k, *_: (i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (None, None, 128, 128)) self.assertEqual( x_block_spec.index_map(0, 1, 2, *scalar_prefetch_values), (1, 2, 0, 1) ) x = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (x,)), x ) def test_strict_mode_compares_index_maps_reading_prefetched_scalars(self): x = jax.random.normal(jax.random.key(0), (3, 4, 512, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) j = jnp.array(2, dtype=jnp.int32) def f(): # `x` is sliced twice, so pulling compares the two index maps it arrives # at, and both of them read the scalars `i` and `j`. a = jax.lax.dynamic_slice(x, (i, j, 0, 0), (1, 1, 512, 512)) b = jax.lax.dynamic_slice(x, (i, j, 0, 0), (1, 1, 512, 512)) return a * b f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec( (1, 1, 128, 128), lambda a, b, c, *_: (0, 0, a, b) ) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3 )(new_values) self.assertLen(kernel_fn.scalar_prefetch, 2) scalar_prefetch = jax.tree.map(lambda v: v[None], kernel_fn.scalar_prefetch) self.assertEqual( value_block_specs[0].index_map(0, 1, 2, *scalar_prefetch), (1, 2, 0, 1) ) block = np.full((1, 1, 128, 128), 3.0, dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch, (block,)), block * block ) def test_strict_mode_rejects_differing_index_maps_reading_scalars(self): x = jax.random.normal(jax.random.key(0), (3, 4, 512, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) j = jnp.array(2, dtype=jnp.int32) def f(): # As above, but the two slices disagree, so the comparison must say so. a = jax.lax.dynamic_slice(x, (i, j, 0, 0), (1, 1, 512, 512)) b = jax.lax.dynamic_slice(x, (j, i, 0, 0), (1, 1, 512, 512)) return a * b f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec( (1, 1, 128, 128), lambda a, b, c, *_: (0, 0, a, b) ) with self.assertRaisesRegex(ValueError, "Fusion contains a DAG"): block_spec_lib.pull_block_spec(f2, block_spec, grid_len=3)(new_values) def test_dynamic_slice_derived_scalar_prefetch(self): x = jax.random.normal(jax.random.key(0), (3, 4, 512, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) j = jnp.array(2, dtype=jnp.int32) def f(): return jax.lax.dynamic_slice(x, (i, j, 0, 0), (1, 1, 512, 512)) f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec( (1, 1, 128, 128), lambda i, j, k, *_: (0, 0, i, j) ) # No handler: the fusion is handed only its own scalars, which the pull # derives and exposes as kernel_fn.scalar_prefetch. kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3 )(new_values) self.assertLen(kernel_fn.scalar_prefetch, 2) self.assertEqual(kernel_fn.scalar_prefetch[0], 1) self.assertEqual(kernel_fn.scalar_prefetch[1], 2) scalar_prefetch = jax.tree.map(lambda x: x[None], kernel_fn.scalar_prefetch) self.assertEqual( value_block_specs[0].index_map(0, 1, 2, *scalar_prefetch), (1, 2, 0, 1) ) x = np.ones((1, 1, 128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch, (x,)), x ) def test_scalar_prefetch_used_in_body_is_read_from_smem_dedup(self): x = jnp.zeros((512, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) def f(): return jax.lax.dynamic_slice(x, (i, 0), (128, 512)) + i.astype(x.dtype) f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec((128, 128), lambda i, j, *_: (0, j)) kernel_fn, _, _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2 )(new_values) self.assertLen(kernel_fn.scalar_prefetch, 1) # The body sees the staged value, not the one the fusion closed over. staged = (jnp.array([7], dtype=jnp.int32),) x = np.zeros((128, 128), dtype=np.float32) np.testing.assert_array_equal(kernel_fn((0, 0), staged, (x,)), x + 7) def test_scalar_prefetch_derived_by_binop_used_in_body(self): x = jnp.zeros(8, dtype=jnp.int32) a = jnp.int32(1) def f(): i = a + 1 return jax.lax.dynamic_slice(x, (i,), (1,)) + i _, _, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(scalar_prefetch_values, 1) @parameterized.named_parameters( ('broadcast_in_dim', lambda x, i: jax.lax.broadcast_in_dim(i, (1,), ())), ('dynamic_slice', lambda x, i: jax.lax.dynamic_slice(x, (i,), (1,))), ( 'dynamic_update_slice', lambda x, i: jax.lax.dynamic_update_slice(x, x, (i,)), ), ) def test_scalar_prefetch_derived_by_unsupported_op_used_in_body(self, op): x = jnp.zeros(1, dtype=jnp.int32) i = jnp.int32(0) def f(): j = op(x, i) return jax.lax.dynamic_slice(x, (j[0],), (1,)) + j with self.assertRaisesRegex(NotImplementedError, 'scalar prefetch'): block_spec_lib.get_stateful_input_fusion_values(f) @parameterized.parameters(True, False) def test_scalar_prefetch_mixed_operands(self, use_literal): i = jnp.array(3.0, dtype=jnp.float32) def f(): a = 10.0 if use_literal else i return jax.jit(lambda a, b: a + jnp.broadcast_to(b, (128, 128)))( a, i * 2.0 ) block_spec = pl.BlockSpec((128, 128), lambda *ij: ij) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=2 )() staged = (jnp.array([3.0], dtype=jnp.float32),) expected = 16.0 if use_literal else 9.0 np.testing.assert_array_equal( kernel_fn((0, 0), staged), jnp.full((128, 128), expected) ) def test_scalar_prefetch_with_closed_over_body_scalar(self): x = jnp.zeros((3, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) scale = jnp.array(5.0, dtype=np.float32) def f(values): (x_,) = values return jax.lax.dynamic_slice(x_, (i, 0), (1, 512)) * scale block_spec = pl.BlockSpec((1, 128), lambda i, *_: (0, i)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=1 )((x,)) self.assertLen(kernel_fn.scalar_prefetch, 1) self.assertEqual(kernel_fn.scalar_prefetch[0], 1) staged = (jnp.array([2], dtype=jnp.int32),) self.assertEqual(value_block_specs[0].index_map(0, *staged), (2, 0)) blk = np.ones((1, 128), dtype=np.float32) np.testing.assert_array_equal(kernel_fn((0,), staged, (blk,)), blk * 5.0) def test_scalar_prefetch_wrong_count_raises(self): x = jax.random.normal(jax.random.key(0), (3, 512), dtype=np.float32) i = jnp.array(1, dtype=jnp.int32) def f(): return jax.lax.dynamic_slice(x, (i, 0), (1, 512)) f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec((1, 128), lambda i, *_: (0, i)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=1 )(new_values) with self.assertRaisesRegex(ValueError, 'prefetches 1 scalars but was'): value_block_specs[0].index_map(0) with self.assertRaisesRegex(ValueError, 'prefetches 1 scalars but was'): kernel_fn((0,), (), (np.ones((1, 128), np.float32),)) def test_non_scalar_prefetch_raises(self): x = jax.random.normal(jax.random.key(0), (3, 512), dtype=np.float32) i = jnp.array([1], dtype=jnp.int32) def f(): return jax.lax.dynamic_slice(x, (i[0], 0), (1, 512)) f2, new_values, _ = block_spec_lib.get_fusion_values(f) block_spec = pl.BlockSpec((1, 128), lambda i, *_: (0, i)) with self.assertRaisesRegex(ValueError, 'Only scalars can be prefetched'): block_spec_lib.pull_block_spec(f2, block_spec, grid_len=1)(new_values) def test_dynamic_update_slice_pull(self): operand = jax.random.normal( jax.random.key(0), (3, 4, 512, 512), dtype=np.float32 ) update = jax.random.normal( jax.random.key(1), (1, 1, 256, 256), dtype=np.float32 ) i, j, k, l = jnp.int32(1), jnp.int32(2), jnp.int32(128), jnp.int32(0) def f(): return jax.lax.dynamic_update_slice(operand, update, (i, j, k, l)) f2, new_values, scalar_prefetch = ( block_spec_lib.get_stateful_input_fusion_values(f)) self.assertLen(new_values, 2) self.assertLen(scalar_prefetch, 4) block_spec = pl.BlockSpec( (1, 1, 128, 128), lambda a, b, c, *_: (a, b, c, 0) ) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch = jax.tree.map(lambda x: x[None], scalar_prefetch) self.assertLen(value_block_specs, 2) op_bs, up_bs = value_block_specs self.assertEqual(op_bs, pl.no_block_spec) self.assertEqual(up_bs.block_shape, (1, 1, 128, 128)) self.assertEqual(up_bs.index_map(2, 2, 2, *scalar_prefetch), (1, 0, 1, 0)) op_tile = np.zeros((1, 1, 128, 128), dtype=np.float32) up_tile = np.ones((1, 1, 128, 128), dtype=np.float32) out = kernel_fn((2, 2, 2), scalar_prefetch, (op_tile, up_tile)) np.testing.assert_array_equal(out, up_tile) def test_dynamic_update_slice_pull_clamping(self): operand = jax.random.normal(jax.random.key(0), (512, 512), dtype=np.float32) update = jax.random.normal(jax.random.key(1), (256, 256), dtype=np.float32) i, j = jnp.int32(1000), jnp.int32(1000) def f(): return jax.lax.dynamic_update_slice(operand, update, (i, j)) # Explicitly confirm standard JAX execution shifts (clamps) out-of-bounds # start indices. On an array of shape (512, 512) with update of shape (256, # 256), start index 1000 is clamped to 256. expected = f() np.testing.assert_array_equal(expected[256:, 256:], update) np.testing.assert_array_equal(expected[:256, :256], operand[:256, :256]) f2, new_values, scalar_prefetch = ( block_spec_lib.get_stateful_input_fusion_values(f)) block_spec = pl.BlockSpec((128, 128), lambda a, b, *_: (a, b)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch = jax.tree.map(lambda x: x[None], scalar_prefetch) op_bs, up_bs = value_block_specs self.assertEqual(op_bs, pl.no_block_spec) # For output block (2, 2), corresponding update block is: # (2 - 2, 2 - 2) = (0, 0). self.assertEqual(up_bs.index_map(2, 2, *scalar_prefetch), (0, 0)) # Verify numerical equivalence of executing the Pallas kernel on tile (2, 2) op_tile = operand[256:384, 256:384] up_tile = update[0:128, 0:128] out_tile = kernel_fn((2, 2), scalar_prefetch, (op_tile, up_tile)) np.testing.assert_array_equal(out_tile, expected[256:384, 256:384]) def test_dynamic_update_slice_non_multiple_block_size(self): operand = jnp.zeros((1, 12, 128), dtype=np.float32) update = jnp.ones((1, 12, 128), dtype=np.float32) i, j, k = jnp.int32(0), jnp.int32(0), jnp.int32(0) def f(): return jax.lax.dynamic_update_slice(operand, update, (i, j, k)) f2, new_values, scalar_prefetch = ( block_spec_lib.get_stateful_input_fusion_values(f) ) block_spec = pl.BlockSpec((1, 8, 128), lambda a, b, *_: (a, b, 0)) _, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch_mem = jax.tree.map(lambda x: x[None], scalar_prefetch) _, up_bs = value_block_specs self.assertEqual(up_bs.block_shape, (1, 8, 128)) self.assertEqual(up_bs.index_map(0, 1, *scalar_prefetch_mem), (0, 1, 0)) out_bs = block_spec_lib.push_block_spec( f2, (pl.no_block_spec, up_bs), scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), grid_len=2, )(new_values) self.assertEqual(out_bs.block_shape, (1, 8, 128)) self.assertEqual(out_bs.index_map(0, 1, *scalar_prefetch_mem), (0, 1, 0)) def test_dynamic_update_slice_pull_unaligned(self): from jax.experimental import checkify operand = jnp.zeros((512, 512), dtype=np.float32) update_unaligned = jnp.ones((150, 150), dtype=np.float32) def f_unaligned_shape(): return jax.lax.dynamic_update_slice( operand, update_unaligned, (jnp.int32(0), jnp.int32(0)) ) f2, new_values, scalar_prefetch = ( block_spec_lib.get_stateful_input_fusion_values(f_unaligned_shape)) block_spec = pl.BlockSpec((128, 128), lambda a, b, *_: (a, b)) _, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch_mem = jax.tree.map(lambda x: x[None], scalar_prefetch) with self.assertRaisesRegex(ValueError, 'not a multiple of block_size'): value_block_specs[1].index_map(0, 0, *scalar_prefetch_mem) update = jnp.ones((256, 256), dtype=np.float32) def f_unaligned_idx(): return jax.lax.dynamic_update_slice( operand, update, (jnp.int32(10), jnp.int32(0)) ) f2, new_values, scalar_prefetch = ( block_spec_lib.get_stateful_input_fusion_values(f_unaligned_idx)) _, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch_mem = jax.tree.map(lambda x: x[None], scalar_prefetch) err, _ = checkify.checkify( lambda: value_block_specs[1].index_map(0, 0, *scalar_prefetch_mem) )() with self.assertRaises(Exception): err.throw() def test_dynamic_update_slice_unsupported_params(self): with self.assertRaisesRegex(NotImplementedError, 'not supported yet'): block_spec_lib._dynamic_update_slice_usage_rule( None, {block_spec_lib.Usage.REGULAR}, allow_negative_indices=True ) with self.assertRaisesRegex(NotImplementedError, 'not supported yet'): block_spec_lib._dynamic_update_slice_eval_rule( None, None, None, allow_negative_indices=True ) with self.assertRaisesRegex(NotImplementedError, 'not supported yet'): block_spec_lib._dynamic_update_slice_rule( None, None, allow_negative_indices=True ) @parameterized.parameters(128, pl.Blocked(128)) def test_concatenate_spanning_blocks(self, block_size): x = jax.random.normal(jax.random.key(0), (256, 256), dtype=np.float32) y = jax.random.normal(jax.random.key(1), (256, 256), dtype=np.float32) def f(x, y): return jnp.concatenate([x, y], axis=0) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, x, y ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((block_size, 128), lambda i, j, *_: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler( 0 ), )(new_values, x, y) ) self.assertEmpty(value_block_specs) self.assertLen(in_block_specs, 2) x_block_spec, y_block_spec = in_block_specs self.assertEqual(x_block_spec.block_shape, (block_size, 128)) self.assertEqual(y_block_spec.block_shape, (block_size, 128)) # Indices in this concatenate should be clamped depending on if the # block is OOB. self.assertEqual( x_block_spec.index_map(0, 0, scalar_prefetch_values), (0, 0) ) self.assertEqual( x_block_spec.index_map(1, 0, scalar_prefetch_values), (1, 0) ) self.assertEqual( x_block_spec.index_map(2, 0, scalar_prefetch_values), (1, 0) ) self.assertEqual( x_block_spec.index_map(3, 0, scalar_prefetch_values), (1, 0) ) self.assertEqual( y_block_spec.index_map(0, 0, scalar_prefetch_values), (0, 0) ) self.assertEqual( y_block_spec.index_map(1, 0, scalar_prefetch_values), (0, 0) ) self.assertEqual( y_block_spec.index_map(2, 0, scalar_prefetch_values), (0, 0) ) self.assertEqual( y_block_spec.index_map(3, 0, scalar_prefetch_values), (1, 0) ) x = jax.random.normal(jax.random.key(0), (128, 128), dtype=np.float32) y = jax.random.normal(jax.random.key(1), (128, 128), dtype=np.float32) # Evaluating should just select the valid block. np.testing.assert_array_equal( kernel_fn((0, 0), scalar_prefetch_values, (), x, y), x ) np.testing.assert_array_equal( kernel_fn((1, 0), scalar_prefetch_values, (), x, y), x ) np.testing.assert_array_equal( kernel_fn((2, 0), scalar_prefetch_values, (), x, y), y ) np.testing.assert_array_equal( kernel_fn((3, 0), scalar_prefetch_values, (), x, y), y ) @parameterized.parameters(512, pl.Blocked(512)) def test_concatenate_full_blocks(self, block_size): x = jax.random.normal(jax.random.key(0), (256, 256), dtype=np.float32) y = jax.random.normal(jax.random.key(1), (256, 256), dtype=np.float32) def f(x, y): return jnp.concatenate([x, y], axis=0) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, x, y ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((block_size, 128), lambda i, j, *_: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler( 0 ), )(new_values, x, y) ) self.assertEmpty(value_block_specs) self.assertLen(in_block_specs, 2) x_block_spec, y_block_spec = in_block_specs self.assertEqual(x_block_spec.block_shape, (256, 128)) self.assertEqual(y_block_spec.block_shape, (256, 128)) # Indices in this concatenate always just fetch the whole block on that # dimension. self.assertEqual( x_block_spec.index_map(0, 0, scalar_prefetch_values), (0, 0) ) self.assertEqual( x_block_spec.index_map(0, 1, scalar_prefetch_values), (0, 1) ) self.assertEqual( y_block_spec.index_map(0, 0, scalar_prefetch_values), (0, 0) ) self.assertEqual( y_block_spec.index_map(0, 1, scalar_prefetch_values), (0, 1) ) x = jax.random.normal(jax.random.key(0), (256, 128), dtype=np.float32) y = jax.random.normal(jax.random.key(1), (256, 128), dtype=np.float32) xy = jnp.concatenate([x, y], axis=0) # Evaluating should just select the valid block. np.testing.assert_array_equal( kernel_fn((0, 0), scalar_prefetch_values, (), x, y), xy ) np.testing.assert_array_equal( kernel_fn((1, 0), scalar_prefetch_values, (), x, y), xy ) @parameterized.product( squeezed=[None, pl.Squeezed()], in_shapes=[((2, 256), (3, 256)), ((1, 256),)], ) def test_concatenate_squeezed_concat_dim(self, squeezed, in_shapes): xs = [ jnp.arange(np.prod(s), dtype=jnp.float32).reshape(s) for s in in_shapes ] block_spec = pl.BlockSpec((squeezed, 128), lambda i, j: (i, j)) kernel_fn, in_block_specs, _ = block_spec_lib.pull_block_spec( lambda *args: lax.concatenate(args, dimension=0), block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(*xs) for bs in in_block_specs: self.assertEqual(bs.block_shape, (squeezed, 128)) idx = len(xs) - 1 row = sum(s[0] for s in in_shapes[:idx]) self.assertEqual(in_block_specs[idx].index_map(row, 1), (0, 1)) x_blocks = [x[0, 128:] for x in xs] np.testing.assert_array_equal( kernel_fn((row, 1), (), *x_blocks), x_blocks[idx] ) def test_transpose_minor(self): x = jax.random.normal(jax.random.key(0), (512, 256), dtype=np.float32) def f(): return jax.lax.transpose(x, (1, 0)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (128, 128)) self.assertEqual( x_block_spec.index_map(0, 1, 2, *scalar_prefetch_values), (1, 0) ) x = jax.random.normal(jax.random.key(0), (128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (x,)), x.T ) def test_transpose_major(self): x = jax.random.normal( jax.random.key(0), (2, 3, 4, 512, 256), dtype=np.float32 ) def f(): return jax.lax.transpose(x, (1, 2, 0, 3, 4)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec( (None, 4, 2, 128, 128), lambda i, j, k, l, m: (i, j, k, l, m) ) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=5, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (2, None, 4, 128, 128)) self.assertEqual( x_block_spec.index_map(0, 1, 2, 3, 4, *scalar_prefetch_values), (2, 0, 1, 3, 4), ) with self.assertRaisesRegex( NotImplementedError, 'Cannot permute last two dimensions' ): block_spec_lib.pull_block_spec( lambda x: jax.lax.transpose(x, (0, 1, 3, 2, 4)), block_spec )(x) x = jax.random.normal(jax.random.key(0), (2, 4, 128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0, 0), scalar_prefetch_values, (x,)), x.swapaxes(0, 1), ) @parameterized.parameters((None, 128), (pl.Squeezed(), pl.Blocked(128))) def test_iota(self, squeezed, blocked): def f(): return jax.lax.broadcasted_iota(jnp.int32, (2, 2, 512, 512), 2) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec( (squeezed, squeezed, blocked, blocked), lambda i, j, k, l: (i, j, k, l) ) kernel_fn, ((),), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) x = jax.lax.broadcasted_iota(jnp.int32, (128, 128), 0) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((1, 1, 0, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 1), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 1, 0), scalar_prefetch_values, ()), x + 128 ) np.testing.assert_array_equal( kernel_fn((0, 0, 3, 0), scalar_prefetch_values, ()), x + 128 * 3 ) @parameterized.parameters(None, pl.Squeezed()) def test_broadcast_size_one_dim_to_squeezed_dim(self, squeezed): x = jnp.arange(512, dtype=jnp.float32).reshape((1, 512)) block_spec = pl.BlockSpec((squeezed, 128), lambda i, j: (i, j)) kernel_fn, (x_block_spec,), _ = block_spec_lib.pull_block_spec( lambda x: lax.broadcast_in_dim(x, (4, 512), (0, 1)), block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(x) self.assertEqual(x_block_spec.block_shape, (squeezed, 128)) self.assertEqual(x_block_spec.index_map(3, 2), (0, 2)) np.testing.assert_array_equal( kernel_fn((3, 2), (), x[0, 256:384]), x[0, 256:384] ) def test_broadcast_scalar(self): def f(): return jnp.full((1, 1, 512, 512), fill_value=1.2345, dtype=jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec( (None, None, 128, 128), lambda i, j, k, l: (i, j, k, l) ) kernel_fn, ((),), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) x = jnp.full((128, 128), fill_value=1.2345, dtype=jnp.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((1, 1, 0, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 1), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 1, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 3, 0), scalar_prefetch_values, ()), x ) def test_broadcast_scalar_with_prefetch(self): a = jnp.array(1.2345) def f(): return jnp.full((1, 1, 512, 512), fill_value=a, dtype=jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertEmpty(new_values) self.assertLen(scalar_prefetch_values, 1) block_spec = pl.BlockSpec( (None, None, 128, 128), lambda i, j, k, l, *_: (i, j, k, l) ) kernel_fn, ((),), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) scalar_prefetch_values = jax.tree.map( lambda x: x[None], scalar_prefetch_values ) x = jnp.full((128, 128), fill_value=a, dtype=jnp.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((1, 1, 0, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 1), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 1, 0), scalar_prefetch_values, ()), x ) np.testing.assert_array_equal( kernel_fn((0, 0, 3, 0), scalar_prefetch_values, ()), x ) @parameterized.parameters( (False, False), (False, True), (True, False), (True, True) ) def test_broadcast_array(self, bcast0, bcast1): x = jnp.ones((1 if bcast0 else 512, 1 if bcast1 else 512)) def f(): return jax.lax.broadcast_in_dim(x, (2, 2, 512, 512), (2, 3)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_shape = (None, 1, 128, 128) block_spec = pl.BlockSpec(block_shape, lambda i, j, k, l: (i, j, k, l)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_index_map = value_block_specs[0].index_map self.assertEqual( x_index_map(0, 0, 1, 2), (0 if bcast0 else 1, 0 if bcast1 else 2) ) self.assertEqual( x_index_map(1, 2, 3, 3), (0 if bcast0 else 3, 0 if bcast1 else 3) ) block_shape = (1 if bcast0 else 128, 1 if bcast1 else 128) self.assertEqual(block_shape, value_block_specs[0].block_shape) x = jnp.full(block_shape, fill_value=1.2345, dtype=jnp.float32) y = jax.lax.broadcast_in_dim(x, (1, 128, 128), (1, 2)) np.testing.assert_array_equal(kernel_fn((0, 0, 0, 0), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((1, 1, 0, 0), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((0, 0, 0, 1), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((0, 0, 1, 0), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((0, 0, 3, 0), (), (x,)), y) @parameterized.parameters(0, 1, 2, 3) def test_broadcast_1d_array(self, bcast_dim): full_shape = (2, 2, 512, 512) x = jnp.ones((full_shape[bcast_dim],)) def f(): return jax.lax.broadcast_in_dim(x, full_shape, (bcast_dim,)) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_shape = (None, 1, 128, 128) block_spec = pl.BlockSpec(block_shape, lambda i, j, k, l: (i, j, k, l)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_index_map = value_block_specs[0].index_map self.assertEqual(x_index_map(0, 0, 1, 2), ((0, 0, 1, 2)[bcast_dim],)) self.assertEqual(x_index_map(1, 2, 3, 3), ((1, 2, 3, 3)[bcast_dim],)) if block_shape[bcast_dim] is None: x = jnp.ones(()) y = jax.lax.broadcast_in_dim(x, (1, 128, 128), ()) else: x = jnp.arange(block_shape[bcast_dim] or 1, dtype=jnp.float32) y = jax.lax.broadcast_in_dim(x, (1, 128, 128), (bcast_dim - 1,)) np.testing.assert_array_equal(kernel_fn((0, 0, 0, 0), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((1, 1, 0, 0), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((0, 0, 0, 1), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((0, 0, 1, 0), (), (x,)), y) np.testing.assert_array_equal(kernel_fn((0, 0, 3, 0), (), (x,)), y) def test_element_indexing(self): x = np.zeros((512, 512), dtype=np.float32) def f(): return x f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) # Block spec with an offset on the first dimension block_spec = pl.BlockSpec( (pl.Element(128, (0, 16)), 128), lambda i, j, k: (128 * i + 16, j) ) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) self.assertEmpty(scalar_prefetch_values) self.assertEqual( value_block_specs[0].block_shape, (pl.Element(128, (0, 16)), 128) ) self.assertEqual(value_block_specs[0].index_map(0, 1, 2), (16, 1)) self.assertEqual(value_block_specs[0].index_map(1, 1, 2), (128 + 16, 1)) x_block = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn( (0, 0, 0), scalar_prefetch_values, (np.ones((128, 128), dtype=np.float32),), ), x_block, ) @parameterized.parameters( # Merge two dimensions. ((8, 8, 128), (1, 2, 128), (1, 1, 2, 128), (0, 2, 3, 5)), ((2, 32, 128), (2, 4, 128), (2, 1, 4, 128), (2, 1, 1, 5)), ((2, 4, 1024), (2, 1, 128), (2, 1, 1, 128), (2, 3, 5, 0)), ((2, 4, 1024), (2, None, 128), (2, 1, 1, 128), (2, 3, 5, 0)), ((2, 4, 1024), (2, pl.Squeezed(), 128), (2, 1, 1, 128), (2, 3, 5, 0)), # Merge three dimensions. ((64, 128), (4, 128), (1, 1, 4, 128), (0, 1, 0, 3)), ((2, 4096), (1, 64), (1, 1, 1, 64), (2, 0, 1, 1)), # Merge two pairs of dimensions. ((8, 1024), (1, 256), (1, 1, 2, 128), (0, 2, 3, 0)), # Merge three dims and expand in trailing dim. ((64, 128, 1), (4, 128, 1), (1, 1, 4, 128), (0, 1, 0, 3)), # Test propagating pl.BoundedSlice. ( (8, 8, 128), (2, pl.BoundedSlice(2), 16), (1, 2, pl.BoundedSlice(2), 16), (1, 0, 3, 5), ), ) def test_reshape( self, shape, block_shape, expected_x_block_shape, expected_x_index ): f = lambda x: x.reshape(shape) in_type = jax.ShapeDtypeStruct((2, 4, 8, 128), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec(block_shape, lambda *pids: pids) kernel_fn, (value_block_specs, x_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=len(shape), scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) pids = (2, 3, 5)[: len(shape)] self.assertEmpty(value_block_specs) self.assertEqual(x_block_spec.block_shape, expected_x_block_shape) self.assertEqual(x_block_spec.index_map(*pids), expected_x_index) concrete_block_shape = block_spec_lib._ref_block_shape(block_shape) concrete_expected_block_shape = block_spec_lib._ref_block_shape( expected_x_block_shape ) x = jnp.arange( np.prod(concrete_block_shape), dtype=jnp.float32, ) x = x.reshape(concrete_expected_block_shape) y = kernel_fn((0, 1, 2), scalar_prefetch_values, (), x) np.testing.assert_array_equal(y, x.reshape(concrete_block_shape)) def test_unsupported_reshape_merge_error_message(self): def f(x): return x.reshape((2, 12)) in_type = jax.ShapeDtypeStruct((2, 3, 4), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) block_spec = pl.BlockSpec((1, 6), lambda i, j: (i, j)) expected_msg = ( r'Cannot pull BlockSpec with block_shape=\(1, 6\) across reshape merge' r' \(2, 3, 4\) -> \(2, 12\): Merging input dimensions \(3, 4\) into' r' output dimension 12 with block size 6 produces a non-rectangular' r' tile in the input buffer\. In Pallas, blocks must be rectangular' r' tiles, but a 1D block of size 6 cuts across inner dimension 4' r' \(remaining size 6 is not divisible by 4\)\.' r' Hint: Make the block size divisible by 4\.' ) with self.assertRaisesRegex(NotImplementedError, expected_msg): block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) def test_basic_reshape_sublanes_to_lanes(self): def f(x): return x.reshape((512, 4096)) in_type = jax.ShapeDtypeStruct((512, 32, 128), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((256, 1024), lambda i, j, k: (i, k)) kernel_fn, (value_block_specs, x_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertEmpty(value_block_specs) self.assertEqual(x_block_spec.block_shape, (256, 8, 128)) self.assertEqual(x_block_spec.index_map(0, 1, 2), (0, 2, 0)) self.assertEqual(x_block_spec.index_map(3, 2, 1), (3, 1, 0)) x = jnp.arange((256 * 1024), dtype=jnp.float32).reshape((256, 8, 128)) y = kernel_fn((0, 1, 2), scalar_prefetch_values, (), x) np.testing.assert_array_equal(y, x.reshape((256, 1024))) @parameterized.parameters( ( (4096,), (32, 128), pl.BlockSpec((8, 128), lambda i, j: (j, 0)), (1024,), ), ( (512, 4096), (512, 32, 128), pl.BlockSpec((256, 8, 128), lambda i, j, k: (i, k, 0)), (256, 1024), ), ( (512, 4096), (512, 32, 128), pl.BlockSpec((256, 8, 128), lambda i, j, k: (i, k, j)), (256, 1024), ), ( (512, 4096), (512, 16, 256), pl.BlockSpec((256, 4, 256), lambda i, j, k: (i, k, 0)), (256, 1024), ), ( (4, 512, 4096), (4, 512, 32, 128), pl.BlockSpec((2, 256, 8, 128), lambda b, i, j, k: (b, i, k, 0)), (2, 256, 1024), ), ) def test_basic_reshape_lanes_to_sublanes( self, in_shape, out_shape, out_block_spec, expected_in_block_shape ): in_type = jax.ShapeDtypeStruct(in_shape, jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( lambda x: x.reshape(out_shape), in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) kernel_fn, (value_block_specs, x_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, out_block_spec, grid_len=len(out_shape), scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertEmpty(value_block_specs) self.assertEqual(x_block_spec.block_shape, expected_in_block_shape) pids = (*range(1, len(out_shape) - 1), 0, 2) expected_idx = (*range(1, len(out_shape) - 1), 2) self.assertEqual(x_block_spec.index_map(*pids), expected_idx) self.assertEqual(jax.jit(x_block_spec.index_map)(*pids), expected_idx) x = jnp.arange(np.prod(expected_in_block_shape), dtype=jnp.float32).reshape( expected_in_block_shape ) y = kernel_fn(pids, scalar_prefetch_values, (), x) np.testing.assert_array_equal(y, x.reshape(out_block_spec.block_shape)) def test_reshape_lanes_to_sublanes_errors(self): in_type = jax.ShapeDtypeStruct((512, 4096), jnp.float32) f2, new_values, _ = block_spec_lib.get_fusion_values( lambda x: x.reshape((512, 32, 128)), in_type ) block_spec = pl.BlockSpec((256, 4, 256), lambda i, j, k: (i, j, k)) with self.assertRaises(NotImplementedError): _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) block_spec = pl.BlockSpec((256, 8, 128), lambda i, j, k: (i, k, 1)) _, (_, x_block_spec), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) with self.assertRaisesRegex( NotImplementedError, 'Must select entire block on last dimension for reshape', ): x_block_spec.index_map(0, 1, 2) def test_basic_swap(self): value = jnp.arange((512 * 1024), dtype=jnp.int32).reshape((512, 1024)) * 2 x = jnp.zeros((256, 512), dtype=jnp.int32) def outer(refs): ref, y_ref = refs def f(values, x): ref, = values return ref.swap(x) in_type = jax.ShapeDtypeStruct((512, 1024), jnp.int32) scalar_prefetch_values = () block_spec = pl.BlockSpec((256, 512), lambda i, j, k: (i, k)) kernel_fn, (value_block_specs, x_block_spec), _ = ( block_spec_lib.pull_block_spec( f, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )((ref,), in_type) ) self.assertLen(value_block_specs, 1) self.assertEqual(x_block_spec.index_map(0, 1, 2), (0, 2)) self.assertEqual(x_block_spec.index_map(3, 2, 1), (3, 1)) y_ref[...] = kernel_fn((0, 1, 1), scalar_prefetch_values, (ref,), x) y = jnp.zeros((256, 512), jnp.int32) _, y = pl.run_state(outer)((value, y)) np.testing.assert_array_equal(y, value[:256, 512:1024]) @parameterized.parameters((None, 512), (pl.Squeezed(), pl.Blocked(512))) def test_swap_with_squeezed_block_dim(self, squeezed, blocked): x = jnp.arange(512, dtype=jnp.int32) def outer(ref): kernel_fn, _, _ = block_spec_lib.pull_block_spec( ref.swap, pl.BlockSpec((squeezed, blocked), lambda i, j: (i, j)), grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(jax.ShapeDtypeStruct((4, 1024), jnp.int32)) kernel_fn((2, 1), (), x) out = pl.run_state(outer)(jnp.zeros((4, 1024), jnp.int32)) np.testing.assert_array_equal(out[2, 512:], x) def test_basic_get(self): value = jnp.arange((512 * 1024), dtype=jnp.int32).reshape((512, 1024)) * 2 def outer(refs): ref, y_ref = refs def f(): return ref.get() block_spec = pl.BlockSpec((256, 512), lambda i, j, k: (i, k)) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )() y_ref[...] = kernel_fn((0, 1, 1), ()) y = jnp.zeros((256, 512), jnp.int32) _, y = pl.run_state(outer)((value, y)) np.testing.assert_array_equal(y, value[:256, 512:1024]) def test_get_with_squeezed_block_spec(self): value = ( jnp.arange((4 * 512 * 1024), dtype=jnp.int32).reshape((4, 512, 1024)) * 2 ) def outer(refs): ref, y_ref = refs def f(): return ref.get() block_spec = pl.BlockSpec( (pl.Squeezed(), 256, 512), lambda i, j, k: (j, i, k) ) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )() y_ref[...] = kernel_fn((0, 3, 1), ()) y = jnp.zeros((256, 512), jnp.int32) _, y = pl.run_state(outer)((value, y)) np.testing.assert_array_equal(y, value[3, :256, 512:1024]) def test_get_with_squeezed_indexer(self): value = ( jnp.arange((4 * 512 * 1024), dtype=jnp.int32).reshape((4, 512, 1024)) * 2 ) def outer(refs): ref, y_ref = refs def f(): return ref[3] block_spec = pl.BlockSpec((256, 512), lambda i, j, k: (i, k)) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )() y_ref[...] = kernel_fn((0, 2, 1), ()) y = jnp.zeros((256, 512), jnp.int32) _, y = pl.run_state(outer)((value, y)) np.testing.assert_array_equal(y, value[3, :256, 512:1024]) def test_get_with_squeezed_indexer_index_map(self): ref_aval = state.AbstractRef(jax.core.ShapedArray((4, 512, 1024), jnp.int32)) def f(r): return r[3] block_spec = pl.BlockSpec((256, 512), lambda i, j: (i, j)) kernel_fn, (ref_bs,), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(ref_aval) self.assertEqual(tuple(ref_bs.block_shape), (pl.Squeezed(), 256, 512)) self.assertEqual(ref_bs.index_map(0, 1), (0, 0, 1)) def test_get_on_sliced_ref(self): value = jnp.arange((640 * 1024), dtype=jnp.int32).reshape((640, 1024)) * 2 def outer(refs): ref, y_ref = refs i = jnp.int32(128) sliced_ref = ref.at[pl.ds(i, 512), :] def f(): return sliced_ref.get() block_spec = pl.BlockSpec((128, 512), lambda i, j: (i, j)) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )() y_ref[...] = kernel_fn((0, 1), ()) y = jnp.zeros((128, 512), jnp.int32) _, y = pl.run_state(outer)((value, y)) np.testing.assert_array_equal(y, value[128:256, 512:1024]) def test_get_on_sliced_ref_misaligned_raises(self): value = jnp.arange((512 * 1024), dtype=jnp.int32).reshape((512, 1024)) * 2 def outer(refs): ref, y_ref = refs i = jnp.int32(5) sliced_ref = ref.at[pl.ds(i, 512), :] def f(): return sliced_ref.get() block_spec = pl.BlockSpec((128, 512), lambda i, j: (i, j)) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )() err, _ = checkify.checkify(lambda: kernel_fn((0, 1), ()))() err.throw() y = jnp.zeros((128, 512), jnp.int32) with self.assertRaisesRegex( checkify.JaxRuntimeError, "slice_start=5 is not a multiple of block_size=128", ): pl.run_state(outer)((value, y)) def test_swap_on_sliced_ref(self): value = jnp.zeros((640, 1024), dtype=jnp.int32) update = jnp.arange(128 * 512, dtype=jnp.int32).reshape((128, 512)) val_full = jax.core.ShapedArray((512, 1024), jnp.int32) def outer(refs): ref, _ = refs i = jnp.int32(128) sliced_ref = ref.at[pl.ds(i, 512), :] def f(val): return sliced_ref.swap(val) block_spec = pl.BlockSpec((128, 512), lambda i, j: (i, j)) kernel_fn, (in_bs,), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(val_full) _ = kernel_fn((0, 1), (), update) value_out, _ = pl.run_state(outer)((value, value)) np.testing.assert_array_equal(value_out[128:256, 512:1024], update) def test_swap_with_squeezed_indexer(self): value = jnp.zeros((2, 640, 1024), dtype=jnp.int32) update = jnp.arange(128 * 512, dtype=jnp.int32).reshape((128, 512)) val_full = jax.core.ShapedArray((512, 1024), jnp.int32) def outer(refs): ref, _ = refs slot = jnp.int32(1) i = jnp.int32(128) sliced_ref = ref.at[slot, pl.ds(i, 512), :] def f(val): return sliced_ref.swap(val) block_spec = pl.BlockSpec((128, 512), lambda i, j: (i, j)) kernel_fn, (in_bs,), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(val_full) del in_bs _ = kernel_fn((0, 1), (), update) value_out, _ = pl.run_state(outer)((value, value)) np.testing.assert_array_equal(value_out[1, 128:256, 512:1024], update) np.testing.assert_array_equal(value_out[0], jnp.zeros((640, 1024), jnp.int32)) @parameterized.parameters(None, pl.Squeezed()) def test_swap_with_squeezed_block_spec(self, squeezed_dim): ref = jax.new_ref(jnp.zeros((4, 8), dtype=jnp.int32)) update = jnp.arange(4, dtype=jnp.int32) @jax.jit def run(ref): block_spec = pl.BlockSpec((squeezed_dim, 4), lambda *ij: ij) kernel_fn, _, _ = block_spec_lib.pull_block_spec( lambda val: ref.at[1:4].swap(val), block_spec, grid_len=2 )(jax.ShapeDtypeStruct((3, 8), jnp.int32)) kernel_fn((2, 1), (), update) run(ref) np.testing.assert_array_equal(ref[3, 4:8], update) def test_get_key_ref(self): """Tests get_p on a key ref (block_shape=None, MemorySpace.KEY). This exercises the _get_pull_rule early return when block_shape is None, and the _get_eval_rule short-circuit when block_shape is None. """ key = jax.random.key(0, impl='threefry2x32') @jax.jit def outer(y_ref): key_ref = jax.new_ref(key, memory_space=pl.MemorySpace.KEY) def f(): k = key_ref.get() return jax.random.uniform(k, (512, 512), dtype=jnp.float32) block_spec = pl.BlockSpec((512, 256), lambda i: (0, i)) kernel_fn, (), _ = block_spec_lib.pull_block_spec( f, block_spec, grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )() for i in range(2): y_ref[:, i * 256 : (i + 1) * 256] = kernel_fn((i,), ()) @jax.jit def gen(idx): k = key for i in idx: k = jax.random.fold_in(k, i) return jax.random.uniform(k, (512, 256), dtype=jnp.float32) y_ref = jax.new_ref(jnp.zeros((512, 512), dtype=jnp.float32)) outer(y_ref) for i in range(2): block = y_ref[:, i * 256 : (i + 1) * 256] expected_block = gen((0, i)) np.testing.assert_array_equal(block, expected_block) def test_random_noise(self): key = jax.random.key(0, impl='threefry2x32') def f(key): return jax.random.uniform(key, (512, 512), dtype=jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, key ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 256), lambda i, j: (i, j)) kernel_fn, (value_block_specs, key_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, key) ) self.assertEmpty(value_block_specs) self.assertEqual(key_block_spec.memory_space, pl.MemorySpace.KEY) self.assertIsNone(key_block_spec.block_shape) @jax.jit def gen(idx): k = key for i in idx: k = jax.random.fold_in(k, i) return jax.random.uniform(k, (128, 256), dtype=jnp.float32) for i in range(4): for j in range(2): out = kernel_fn((i, j), scalar_prefetch_values, (), key) out_ref = gen((i, j)) np.testing.assert_array_equal(out, out_ref) def test_random_fold_in(self): key = jax.random.key(0, impl='threefry2x32') msg = jnp.array(42, dtype=jnp.int32) def f(msg): folded_key = jax.random.fold_in(key, msg) return jax.random.uniform(folded_key, (512, 512), dtype=jnp.float32) in_type = jax.ShapeDtypeStruct((), jnp.int32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) block_spec = pl.BlockSpec((128, 256), lambda i, j: (i, j)) kernel_fn, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertLen(value_block_specs, 1) self.assertEqual(value_block_specs[0].memory_space, pl.MemorySpace.KEY) self.assertIsNone(value_block_specs[0].block_shape) self.assertEqual(in_block_spec, pl.no_block_spec) @jax.jit def gen(idx): k = jax.random.fold_in(key, msg) # When `uniform` gets split into blocks, it produces this loop of # `fold_in`s. for i in idx: k = jax.random.fold_in(k, i) return jax.random.uniform(k, (128, 256), dtype=jnp.float32) for i in range(4): for j in range(2): out = kernel_fn( (i, j), scalar_prefetch_values, tuple(v for v in new_values), msg ) out_ref = gen((i, j)) np.testing.assert_array_equal(out, out_ref) def test_reduce_sum(self): x = jnp.arange(1024 * 256, dtype=jnp.float32).reshape((1024, 256)) def f(): return x.sum(axis=1) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128,), lambda i: (i,)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) y = x[128:256] out = kernel_fn((1,), scalar_prefetch_values, (y,)) np.testing.assert_array_equal(out, y.sum(axis=1)) @parameterized.parameters((512, 128), (128, 512), (128, 128), (256, 128)) def test_tile(self, *value_shape): x = jax.random.normal(jax.random.key(0), value_shape, dtype=np.float32) full_shape = (512, 512) block_shape = (128, 256) reps = tuple(full_dim // dim for full_dim, dim in zip(full_shape, value_shape)) def f(): return jnp.tile(x, reps) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec(block_shape, lambda i, j: (i, j)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] # Expect to load a slice of the input on the non-tiled dimensions # and the whole input on the tiled dimensions. expected_value_block_shape = tuple( min(value_dim, block_dim) for value_dim, block_dim in zip(value_shape, block_shape) ) self.assertEqual(x_block_spec.block_shape, expected_value_block_shape) for idxs in itertools.product(*( range(full_dim // block_dim) for full_dim, block_dim in zip(full_shape, block_shape) )): expected_block_indices = tuple( idx % (value_dim // block_dim) if value_dim >= block_dim else 0 for idx, value_dim, block_dim in zip(idxs, value_shape, block_shape) ) self.assertEqual( x_block_spec.index_map(*idxs, *scalar_prefetch_values), expected_block_indices, ) block_reps = tuple( block_dim // value_block_dim for block_dim, value_block_dim in zip(block_shape, expected_value_block_shape) ) x_block = jax.random.normal( jax.random.key(0), expected_value_block_shape, dtype=np.float32 ) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (x_block,)), jnp.tile(x_block, block_reps), ) @parameterized.parameters((pl.Squeezed(), 128), (None, pl.Blocked(128))) def test_tile_with_squeezed(self, squeezed, blocked): x = jax.random.normal(jax.random.key(0), (1, 64, 256), dtype=np.float32) block_shape = (squeezed, blocked, blocked) reps = (4, 4, 2) # logical shape after tiling: (4, 256, 512) def f(): return jnp.tile(x, reps) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec(block_shape, lambda i, j, k: (i, j, k)) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) x_block_spec = value_block_specs[0] self.assertEqual(x_block_spec.block_shape, (squeezed, 64, 128)) for i, j, k in itertools.product(range(4), range(2), range(4)): self.assertEqual(x_block_spec.index_map(i, j, k), (0, 0, k % 2)) x_block = jax.random.normal(jax.random.key(0), (64, 128), dtype=np.float32) out = kernel_fn((0, 0, 0), (), (x_block,)) np.testing.assert_array_equal(out, jnp.tile(x_block, (2, 1))) def test_pipeline_mode(self): def f(x): return jnp.tanh(x) + x in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) pmode = pl.Buffered(2) block_spec = pl.BlockSpec( (128, 128), lambda i, j, k: (i, j), pipeline_mode=pmode) _, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertEmpty(value_block_specs) self.assertEqual(in_block_spec.block_shape, (128, 128)) self.assertEqual(in_block_spec.pipeline_mode, pmode) def test_block_index_transform_slice_symbolic_equivalence(self): in_type = (jax.ShapeDtypeStruct((2, 512, 512), jnp.float32),) def fn(x): # test symbolic equivalence of block index transforms, the below will # not create BlockIndexTransforms that are equal as objects y = x[0] z = x[0] return lax.add(y, z) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( fn, *in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, *in_type) ) self.assertEmpty(value_block_specs) x_block_spec_1 = in_block_specs[0] self.assertEqual(x_block_spec_1.block_shape, (None, 128, 128)) self.assertEqual(x_block_spec_1.index_map(0, 1, 2), (0, 0, 1)) x = np.ones((2, 128, 128), dtype=np.float32) # kernel_fn does no slicing/squeezing, that is left to blockspec machinery x_block = np.ones((128, 128), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, (), x_block), fn(x), ) def test_block_index_transform_slice_symbolic_non_equivalence(self): in_type = (jax.ShapeDtypeStruct((2, 512, 512), jnp.float32),) def fn(x): # test symbolic non-equivalence of block index transforms, the below will # create different BlockIndexTransforms y = x[0] z = x[1] return lax.add(y, z) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( fn, *in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) with self.assertRaisesRegex(ValueError, 'cannot uniquely pull block specs'): kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, *in_type) ) del kernel_fn, value_block_specs, in_block_specs def test_dot_general_pull(self): lhs_type = jax.ShapeDtypeStruct((512, 256), jnp.float32) rhs_type = jax.ShapeDtypeStruct((256, 128), jnp.float32) def f(x, y): return jax.lax.dot_general(x, y, (((1,), (0,)), ((), ()))) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, lhs_type, rhs_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 64), lambda i, j, k: (i, j)) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, lhs_type, rhs_type) ) self.assertEmpty(value_block_specs) self.assertLen(in_block_specs, 2) lhs_block_spec, rhs_block_spec = in_block_specs self.assertEqual(lhs_block_spec.block_shape, (128, 256)) self.assertEqual(lhs_block_spec.index_map(0, 1, 2), (0, 0)) self.assertEqual(rhs_block_spec.block_shape, (256, 64)) self.assertEqual(rhs_block_spec.index_map(0, 1, 2), (0, 1)) x = np.ones((128, 256), dtype=np.float32) y = np.ones((256, 64), dtype=np.float32) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, new_values, x, y), jax.lax.dot_general(x, y, (((1,), (0,)), ((), ()))), ) @parameterized.parameters(True, False) def test_clamp_pull(self, scalar_limits): in_shape = (256, 128) in_type = jax.ShapeDtypeStruct((256, 128), jnp.float32) lim_shape = () if scalar_limits else in_shape lim_type = jax.ShapeDtypeStruct(lim_shape, jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( jax.lax.clamp, lim_type, in_type, lim_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((128, 128), lambda i, j, k: (i, j)) scalar_prefetch_handler = block_spec_lib.make_scalar_prefetch_handler() kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=3, scalar_prefetch_handler=scalar_prefetch_handler, )(new_values, lim_type, in_type, lim_type) ) self.assertEmpty(value_block_specs) min_block_spec, x_block_spec, max_block_spec = in_block_specs self.assertEqual(x_block_spec.block_shape, (128, 128)) self.assertEqual(x_block_spec.index_map(0, 1, 2), (0, 1)) if scalar_limits: self.assertEqual(min_block_spec, pl.no_block_spec) self.assertEqual(max_block_spec, pl.no_block_spec) else: self.assertEqual(min_block_spec.block_shape, (128, 128)) self.assertEqual(min_block_spec.index_map(0, 1, 2), (0, 1)) self.assertEqual(max_block_spec.block_shape, (128, 128)) self.assertEqual(max_block_spec.index_map(0, 1, 2), (0, 1)) x = jax.random.normal(jax.random.key(0), in_shape, dtype=np.float32) l = jnp.zeros(lim_shape) u = jnp.ones(lim_shape) np.testing.assert_array_equal( kernel_fn((0, 0, 0), scalar_prefetch_values, new_values, l, x, u), jax.lax.clamp(l, x, u), ) def test_split_pull(self): sizes = (8, 8, 16) x = jnp.arange(32 * 256, dtype=jnp.float32).reshape((32, 256)) def f(x): splits = jax.lax.split(x, axis=0, sizes=sizes) return splits[1] f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, x ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec((8, 128), lambda i: (0, i)) scalar_prefetch_handler = block_spec_lib.make_scalar_prefetch_handler() kernel_fn, (value_block_specs, in_block_spec), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=1, scalar_prefetch_handler=scalar_prefetch_handler, )(new_values, x) ) self.assertEmpty(value_block_specs) self.assertEqual(in_block_spec.block_shape, (pl.Blocked(32), 128)) self.assertEqual(in_block_spec.index_map(1), (0, 1)) y = x[:, 128:256] out = kernel_fn((1,), scalar_prefetch_values, new_values, y) np.testing.assert_array_equal(out, jax.lax.split(y, axis=0, sizes=sizes)[1]) @parameterized.parameters(2, pl.Blocked(2)) def test_stack_pull(self, block_size): x = jnp.arange(2 * 256, dtype=jnp.float32).reshape((2, 256)) kernel_fn, in_block_specs, _ = block_spec_lib.pull_block_spec( lambda *xs: lax.stack(xs), pl.BlockSpec((block_size, 128), lambda i: (0, i)), grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(*x) for bs in in_block_specs: self.assertEqual(bs.block_shape, (128,)) self.assertEqual(bs.index_map(1), (1,)) np.testing.assert_array_equal(kernel_fn((1,), (), *x[:, 128:]), x[:, 128:]) def test_pull_dag_with_none_and_squeezed_block_dims(self): x_type = jax.ShapeDtypeStruct((4, 512, 512), jnp.float32) idx_map = lambda i, j: (1, i, j) out_specs = ( pl.BlockSpec((None, 128, 128), idx_map), pl.BlockSpec((pl.Squeezed(), 128, 128), idx_map), ) kernel_fn, (x_block_spec,), _ = block_spec_lib.pull_block_spec( lambda x: (x + 1, x * 2), out_specs, grid_len=2, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(x_type) self.assertEqual(x_block_spec.index_map(2, 3), (1, 2, 3)) x = jnp.ones((128, 128), dtype=jnp.float32) out1, out2 = kernel_fn((2, 3), (), x) np.testing.assert_array_equal(out1, x + 1) np.testing.assert_array_equal(out2, x * 2) class PullBlockSpecHOPTest(jtu.JaxTestCase): def setUp(self): super().setUp() if config.enable_x64.value: self.skipTest('x64 not supported') def test_jit(self): def f(x): return jax.jit(jnp.sin)(x) in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, j, k, l) ) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertEmpty(value_block_specs) x_block_spec = in_block_specs[0] self.assertEqual(x_block_spec.index_map(0, 0, 1, 2, ()), (0, 0, 1, 2)) self.assertEqual(x_block_spec.index_map(1, 2, 3, 3, ()), (1, 2, 3, 3)) x = jax.random.normal(jax.random.key(0), (1, 128, 128), dtype=np.float32) sin_x = jnp.sin(x) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0), scalar_prefetch_values, (), x), sin_x ) def test_custom_jvp(self): def f(x): return jax.nn.relu(x) in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, j, k, l) ) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertEmpty(value_block_specs) x_block_spec = in_block_specs[0] self.assertEqual(x_block_spec.index_map(0, 0, 1, 2, ()), (0, 0, 1, 2)) self.assertEqual(x_block_spec.index_map(1, 2, 3, 3, ()), (1, 2, 3, 3)) x = jax.random.normal(jax.random.key(0), (1, 128, 128), dtype=np.float32) relu_x = jax.nn.relu(x) np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0), scalar_prefetch_values, (), x), relu_x ) def test_custom_vjp(self): @jax.custom_vjp def act(x): return jax.nn.relu(x) * x def act_fwd(x): return jax.nn.relu(x) * x, (x,) def act_bwd(res, dy): (x,) = res return (dy * x * 2.34,) act.defvjp(act_fwd, act_bwd) def f(x): return act(x) in_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values( f, in_type ) self.assertEmpty(new_values) self.assertEmpty(scalar_prefetch_values) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, j, k, l) ) kernel_fn, (value_block_specs, *in_block_specs), _ = ( block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values, in_type) ) self.assertEmpty(value_block_specs) x_block_spec = in_block_specs[0] self.assertEqual(x_block_spec.index_map(0, 0, 1, 2, ()), (0, 0, 1, 2)) self.assertEqual(x_block_spec.index_map(1, 2, 3, 3, ()), (1, 2, 3, 3)) x = jax.random.normal(jax.random.key(0), (1, 128, 128), dtype=np.float32) relu_x_x = jax.nn.relu(x) * x np.testing.assert_array_equal( kernel_fn((0, 0, 0, 0), scalar_prefetch_values, (), x), relu_x_x ) def test_pull_block_spec_handles_closed_over_constants(self): x = jnp.ones((2, 512, 512)) i = jnp.array(1) def f(): return x[i] f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(f) self.assertLen(new_values, 1) self.assertLen(scalar_prefetch_values, 1) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, j, k, l) ) kernel_fn, (value_block_specs,), _ = block_spec_lib.pull_block_spec( f2, block_spec, grid_len=4, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values) self.assertLen(value_block_specs, 1) scalar_prefetch_values = jax.tree.map( lambda x: x[None], scalar_prefetch_values ) fn = lambda x: kernel_fn((0, 0, 0, 0), scalar_prefetch_values, x) new_values_type = (jax.ShapeDtypeStruct((1, 128, 128), jnp.float32),) # Try pulling again # This should not raise an error. _ = block_spec_lib.pull_block_spec( fn, block_spec, grid_len=1, scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), )(new_values_type) class PushBlockSpecTest(parameterized.TestCase): def setUp(self): super().setUp() if config.enable_x64.value: self.skipTest('x64 not supported') def test_binop(self): def f(x): return x + jnp.ones_like(x) block_spec = pl.BlockSpec((128, 128), lambda i, j: (i, j)) x_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, block_spec.block_shape) def f(x, y): return x + y x_block_spec = pl.BlockSpec((128, 128), lambda i, j: (i, j)) y_block_spec = pl.BlockSpec((128, 1), lambda i, j: (i, 0)) x_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) y_type = jax.ShapeDtypeStruct((512, 1), jnp.float32) with self.assertRaisesRegex( ValueError, 'Cannot propagate block spec through RHS broadcast.' ): block_spec_lib.push_block_spec(f, pl.no_block_spec, y_block_spec)( x_type, y_type ) out_block_spec = block_spec_lib.push_block_spec( f, x_block_spec, pl.no_block_spec )(x_type, y_type) self.assertIs(x_block_spec, out_block_spec) x_block_spec = pl.BlockSpec((1, 128), lambda i, j: (0, j)) y_block_spec = pl.BlockSpec((128, 128), lambda i, j: (i, j)) x_type = jax.ShapeDtypeStruct((1, 512), jnp.float32) y_type = jax.ShapeDtypeStruct((512, 512), jnp.float32) with self.assertRaisesRegex( ValueError, 'Cannot propagate block spec through LHS broadcast.' ): block_spec_lib.push_block_spec(f, x_block_spec, pl.no_block_spec)( x_type, y_type ) out_block_spec = block_spec_lib.push_block_spec( f, pl.no_block_spec, y_block_spec )(x_type, y_type) self.assertIs(out_block_spec, y_block_spec) def test_jit(self): def f(x): return jax.jit(jnp.sin)(x) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, l, k, j) ) x_type = jax.ShapeDtypeStruct((1, 1, 512, 512), jnp.float32) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, block_spec.block_shape) def test_custom_jvp(self): def f(x): return jax.nn.relu(x) x_type = jax.ShapeDtypeStruct((1, 1, 512, 512), jnp.float32) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, l, k, j) ) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, block_spec.block_shape) @parameterized.parameters( ((4096,), (32, 128), (1024,), (8, 128)), ((512, 4096), (512, 32, 128), (256, 1024), (256, 8, 128)), ((512, 4096), (512, 16, 256), (256, 1024), (256, 4, 256)), ((4, 512, 4096), (4, 512, 32, 128), (2, 256, 1024), (2, 256, 8, 128)), ) def test_push_reshape_lanes_to_sublanes( self, in_shape, out_shape, in_block_shape, expected_block_shape ): x_type = jax.ShapeDtypeStruct(in_shape, jnp.float32) block_spec = pl.BlockSpec(in_block_shape, lambda *pids: pids) out_block_spec = block_spec_lib.push_block_spec( lambda x: x.reshape(out_shape), block_spec )(x_type) self.assertEqual(out_block_spec.block_shape, expected_block_shape) pids = tuple(range(1, len(in_shape) + 1)) self.assertEqual(out_block_spec.index_map(*pids), (*pids, 0)) def test_custom_vjp(self): @jax.custom_vjp def act(x): return jax.nn.relu(x) * x def act_fwd(x): return jax.nn.relu(x) * x, (x,) def act_bwd(res, dy): (x,) = res return (dy * x * 2.34,) act.defvjp(act_fwd, act_bwd) def f(x): return act(x) x_type = jax.ShapeDtypeStruct((1, 1, 512, 512), jnp.float32) block_spec = pl.BlockSpec( (None, 1, 128, 128), lambda i, j, k, l, _: (i, l, k, j) ) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, block_spec.block_shape) def test_reduce_sum_push(self): def f(x): return x.sum(axis=0) x_type = jax.ShapeDtypeStruct((256, 512), jnp.float32) block_spec = pl.BlockSpec((256, 256), lambda i, j: (i, j)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, (256,)) self.assertEqual(out_block_spec.index_map(2, 3), (3,)) def f(x): return x.sum(axis=1) x_type = jax.ShapeDtypeStruct((128, 512), jnp.float32) block_spec = pl.BlockSpec((64, 512), lambda i, j: (i, j)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, (64,)) self.assertEqual(out_block_spec.index_map(2, 3), (2,)) def test_broadcast_in_dim_push(self): def f(x): return jnp.broadcast_to(x, (128, 512)) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((128,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, (128, 128)) self.assertEqual(out_block_spec.index_map(3), (0, 3)) def f(x): return jnp.broadcast_to(x, (128, 512)) x_type = jax.ShapeDtypeStruct((1, 512), jnp.float32) block_spec = pl.BlockSpec((1, 128), lambda i, j: (i, j)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, (128, 128)) self.assertEqual(out_block_spec.index_map(0, 3), (0, 3)) def f(x): x = jnp.expand_dims(x, axis=1) return jnp.broadcast_to(x, (128, 512)) x_type = jax.ShapeDtypeStruct((128,), jnp.float32) block_spec = pl.BlockSpec((64,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, (64, 512)) self.assertEqual(out_block_spec.index_map(1), (1, 0)) def f(x): x = jnp.expand_dims(x, axis=0) return jnp.broadcast_to(x, (128, 512)) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((256,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec)(x_type) self.assertEqual(out_block_spec.block_shape, (128, 256)) self.assertEqual(out_block_spec.index_map(1), (0, 1)) def test_concatenate_push(self): def f(x1, x2): return jnp.concatenate((x1, x2), axis=0) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((128,), lambda i: (i,)) with self.assertRaisesRegex( NotImplementedError, 'concatenate not supported yet' ): block_spec_lib.push_block_spec(f, block_spec, block_spec)(x_type, x_type) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((512,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec, block_spec)( x_type, x_type ) self.assertEqual(out_block_spec.block_shape, (1024,)) self.assertEqual(out_block_spec.index_map(0), (0,)) def f(x1, x2): return jnp.stack([x1, x2], axis=0) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((128,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec, block_spec)( x_type, x_type ) self.assertEqual(out_block_spec.block_shape, (2, 128)) self.assertEqual(out_block_spec.index_map(3), (0, 3)) def f(x1, x2): return jnp.stack([x1, x2], axis=1) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((128,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec(f, block_spec, block_spec)( x_type, x_type ) self.assertEqual(out_block_spec.block_shape, (128, 2)) self.assertEqual(out_block_spec.index_map(3), (3, 0)) def test_clamp_push(self): def f(x1, x2, x3): return jax.lax.clamp(x1, x2, x3) x_type = jax.ShapeDtypeStruct((512,), jnp.float32) block_spec = pl.BlockSpec((128,), lambda i: (i,)) out_block_spec = block_spec_lib.push_block_spec( f, block_spec, block_spec, block_spec )(x_type, x_type, x_type) self.assertEqual(out_block_spec.block_shape, (128,)) self.assertEqual(out_block_spec.index_map(0), (0,)) def test_split_push(self): sizes = (2, 8, 2) def f(x): return jax.lax.split(x, axis=0, sizes=sizes) x_type = jax.ShapeDtypeStruct((12, 512), jnp.float32) block_spec = pl.BlockSpec((12, 128), lambda i: (0, i)) out_block_specs = block_spec_lib.push_block_spec(f, block_spec)(x_type) for size, out_block_spec in zip(sizes, out_block_specs): self.assertEqual(out_block_spec.block_shape, (size, pl.Blocked(128))) self.assertEqual(out_block_spec.index_map(1), (0, 1)) def test_push_pull_block_spec_on_custom_vjp_traced(self): with config.custom_vjp3(True): @jax.custom_vjp def custom_act(x): return jax.nn.relu(x) def custom_act_fwd(x): return custom_act(x), (x,) def custom_act_bwd(res, g): (x,) = res return (g * (x > 0),) custom_act.defvjp(custom_act_fwd, custom_act_bwd) def f(x): return custom_act(x) x_struct = jax.ShapeDtypeStruct((256, 256), jnp.float32) block_spec = pl.BlockSpec((128, 128), lambda i, j: (i, j)) out_spec = block_spec_lib.push_block_spec(f, block_spec)(x_struct) self.assertEqual(out_spec.block_shape, (128, 128)) _, (in_spec,), _ = block_spec_lib.pull_block_spec(f, block_spec)(x_struct) self.assertEqual(in_spec.block_shape, (128, 128)) def test_dynamic_update_slice_push(self): operand = jax.random.normal(jax.random.key(0), (512, 512), dtype=np.float32) update = jax.random.normal(jax.random.key(1), (256, 256), dtype=np.float32) i, j = jnp.int32(128), jnp.int32(128) def f(): return jax.lax.dynamic_update_slice(operand, update, (i, j)) f2, new_values, _ = block_spec_lib.get_stateful_input_fusion_values(f) update_bs = pl.BlockSpec((128, 128), lambda g, *_: (0, 0)) out_bs = block_spec_lib.push_block_spec( f2, (pl.no_block_spec, update_bs), scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(), grid_len=1, )(new_values) self.assertEqual(out_bs.block_shape, (128, 128)) scalar_prefetch = ( jnp.array([128], dtype=jnp.int32), jnp.array([128], dtype=jnp.int32), ) # update block (0, 0) pushed to operand starting at offset 128 // 128 = 1 -> (1, 1). self.assertEqual(out_bs.index_map(0, *scalar_prefetch), (1, 1)) def test_dynamic_update_slice_push_errors(self): bs = pl.BlockSpec((128, 128), lambda g: (0, 0)) with self.assertRaisesRegex(ValueError, 'simultaneously'): block_spec_lib._dynamic_update_slice_push_rule( None, bs, bs, pl.no_block_spec, pl.no_block_spec ) with self.assertRaisesRegex(ValueError, 'cannot have block specs'): block_spec_lib._dynamic_update_slice_push_rule( None, pl.no_block_spec, bs, bs, pl.no_block_spec ) @parameterized.parameters( ((128, 128), (128, 128)), ((None, 128), (pl.Squeezed(), 128)), ) def test_binop_push_converging_identical_specs(self, bs1_shape, bs2_shape): x_struct = jax.ShapeDtypeStruct((256, 256), jnp.float32) y_struct = jax.ShapeDtypeStruct((256, 256), jnp.float32) def f(x, y): return x + y idx_map = lambda i, j: (i, j) bs1 = pl.BlockSpec(bs1_shape, idx_map) bs2 = pl.BlockSpec(bs2_shape, idx_map) self.assertIsNot(bs1, bs2) out_spec = block_spec_lib.push_block_spec(f, bs1, bs2)( x_struct, y_struct ) self.assertEqual(out_spec.block_shape, bs1_shape) def test_transpose_push(self): x_type = jax.ShapeDtypeStruct((2, 3, 512, 256), jnp.float32) block_spec = pl.BlockSpec((None, 3, 128, 256), lambda i, j, k, l: (i, j, k, l)) out_spec = block_spec_lib.push_block_spec( lambda x: jax.lax.transpose(x, (1, 0, 3, 2)), block_spec )(x_type) self.assertEqual(out_spec.block_shape, (3, None, 256, 128)) self.assertEqual(out_spec.index_map(0, 1, 2, 3), (1, 0, 3, 2)) with self.assertRaisesRegex( NotImplementedError, 'Cannot permute last two dimensions' ): block_spec_lib.push_block_spec( lambda x: jax.lax.transpose(x, (0, 2, 1, 3)), block_spec )(x_type) def test_no_block_spec_for_output_raises(self): x_type = jax.ShapeDtypeStruct((128, 128), jnp.float32) y_type = jax.ShapeDtypeStruct((64,), jnp.float32) block_spec = pl.BlockSpec((64, 64), lambda i, j: (i, j)) with self.assertRaisesRegex(ValueError, 'No block spec found for output'): block_spec_lib.push_block_spec(lambda x: x, pl.no_block_spec)(x_type) with self.assertRaisesRegex(ValueError, 'No block spec found for output'): block_spec_lib.push_block_spec( lambda x, y: (x + 1, y + 1), block_spec, pl.no_block_spec )(x_type, y_type) if __name__ == '__main__': absltest.main(testLoader=jtu.JaxTestLoader())