Repository: google/jax
Pull request #22084 (2024-06-25 10:29:09)

[pallas] Add limited support for shape polymorphism for TPU
[pallas] Add limited support for shape polymorphism for TPU

The main change is to pass the `result_shapes` to the
hlo.CustomCallOp when the output shapes contain dimension
variables. Everything else is already handled by the
support for dynamic bounds sizes for TPU.

Note that this CL only adds limited support for shape
polymorphism: only on TPU, and only when the block
sizes are static.



Patch:
diff --git a/jax/_src/pallas/core.py b/jax/_src/pallas/core.py
index 006ff50be..85a2eda28 100644
--- a/jax/_src/pallas/core.py
+++ b/jax/_src/pallas/core.py
@@ -301,6 +301,15 @@ def _tile_ref(ref: state.AbstractRef, block_shape: tuple[int, ...] | None
   return ref.update(inner_aval=ref.inner_aval.update(shape=shape))
 
 
+def _check_static_ref_shape(ref: state.AbstractRef) -> state.AbstractRef:
+  shape = ref.shape
+  if not jax_core.is_constant_shape(shape):
+    # TODO(necula): thread the tree labels so that we can localize the error
+    raise ValueError("shape polymorphism for Pallas does not support "
+                     f"dynamically-shaped blocks. Found block_shape: {shape}")
+  return ref
+
+
 def _get_ref_avals(grid, in_avals, in_specs, out_avals, out_specs):
   def _get_memory_space(spec):
     if spec is no_block_spec:
@@ -318,13 +327,13 @@ def _get_ref_avals(grid, in_avals, in_specs, out_avals, out_specs):
     in_specs = [None] * len(in_avals)
     out_specs = [None] * len(out_avals)
   tiled_in_ref_avals = [
-      aval if in_spec is no_block_spec
-      else _tile_ref(aval, in_spec.block_shape)
+      _check_static_ref_shape(aval if in_spec is no_block_spec
+                              else _tile_ref(aval, in_spec.block_shape))
       for aval, in_spec in zip(in_ref_avals, in_specs)
   ]
   tiled_out_ref_avals = [
-      aval if out_spec is no_block_spec
-      else _tile_ref(aval, out_spec.block_shape)
+      _check_static_ref_shape(aval if out_spec is no_block_spec
+                              else _tile_ref(aval, out_spec.block_shape))
       for aval, out_spec in zip(out_ref_avals, out_specs)
   ]
   return in_specs, tiled_in_ref_avals, out_specs, tiled_out_ref_avals
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index 3ebbfdb51..76d913bcb 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -71,7 +71,7 @@ TPUMemorySpace = tpu_core.TPUMemorySpace
 VMEM = tpu_core.TPUMemorySpace.VMEM
 SMEM = tpu_core.TPUMemorySpace.SMEM
 
-# The value interpreter as a dynamic dimension by MLIR.
+# The value interpreted as a dynamic dimension by MLIR.
 MLIR_DYNAMIC = -9223372036854775808
 
 partial = functools.partial
diff --git a/jax/_src/tpu_custom_call.py b/jax/_src/tpu_custom_call.py
index a6455837b..d92c0686a 100644
--- a/jax/_src/tpu_custom_call.py
+++ b/jax/_src/tpu_custom_call.py
@@ -214,7 +214,13 @@ def _tpu_custom_call_lowering(
     raise NotImplementedError(
         "Replica lowering for Mosaic kernels not implemented."
     )
-  extra_attributes = {}
+  if all(core.is_constant_shape(aval_out.shape) for aval_out in ctx.avals_out):
+    result_shapes = None
+  else:
+    result_shapes = [
+        mlir.shape_tensor(mlir.eval_dynamic_shape(ctx, aval_out.shape))
+        for aval_out in ctx.avals_out]
+  extra_attributes = None
   # Add kernel_name and kernel_metadata as attributes to the custom call op.
   # This is because we do not want to pollute the backend_config with this
   # information.
@@ -229,7 +235,9 @@ def _tpu_custom_call_lowering(
       operand_output_aliases=dict(input_output_aliases),
       operand_layouts=_avals_to_layouts(ctx.avals_in),
       result_layouts=_avals_to_layouts(ctx.avals_out),
+      result_shapes=result_shapes,
       extra_attributes=extra_attributes)
+
   return call.results
 
 


Test patch:
diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
index b56d16026..0fb7ceab9 100644
--- a/tests/pallas/BUILD
+++ b/tests/pallas/BUILD
@@ -364,3 +364,30 @@ jax_test(
         "//jax:pallas_tpu",  # build_cleaner: keep
     ],
 )
+
+jax_test(
+    name = "pallas_shape_poly_test",
+    srcs = ["pallas_shape_poly_test.py"],
+    config_tags_overrides = {
+        "gpu_a100_x32": {
+            "ondemand": False,  # Include in presubmit.
+        },
+    },
+    disable_configs = [
+        "gpu_x32",
+        "gpu_h100",
+        "gpu_p100",
+        "gpu_p100_x32",
+        "gpu_pjrt_c_api",
+    ],
+    enable_configs = [
+        "gpu_a100_x32",
+    ],
+    tags = [],
+    deps = [
+        "//jax:pallas",
+        "//jax:pallas_gpu",  # build_cleaner: keep
+        "//jax:pallas_tpu",  # build_cleaner: keep
+        "//jax/experimental/export",
+    ],
+)
diff --git a/tests/pallas/pallas_shape_poly_test.py b/tests/pallas/pallas_shape_poly_test.py
new file mode 100644
index 000000000..f6ac2d0a9
--- /dev/null
+++ b/tests/pallas/pallas_shape_poly_test.py
@@ -0,0 +1,220 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ruff: noqa: F401
+
+import functools
+import math
+import os
+import sys
+
+os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.5"
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import jax
+from jax._src import config
+from jax._src import test_util as jtu
+from jax._src.pallas.pallas_call import _trace_to_jaxpr
+from jax._src import tpu_custom_call  # For configuration values
+import jax.numpy as jnp
+from jax.experimental import pallas as pl
+from jax import export
+import numpy as np
+
+
+config.update("jax_traceback_filtering", "off")
+config.parse_flags_with_absl()
+
+
+# TODO(necula): support an activation
+def matmul_kernel(x_ref, y_ref, o_ref):
+  # x shape: (m, l), y shape (l, n), o shape: (m ,n)
+  block_m, block_l = x_ref.shape
+  block_l2, block_n = y_ref.shape
+  assert block_l2 == block_l
+  assert o_ref.shape == (block_m, block_n)
+  @pl.when(pl.program_id(axis=2) == 0)
+  def _():
+    o_ref[...] = jnp.zeros_like(o_ref)
+
+  o_ref[...] += x_ref[...] @ y_ref[...]
+
+
+@functools.partial(jax.jit, static_argnames=['block_shape'])
+def matmul(
+    x: jax.Array,
+    y: jax.Array,
+    *,
+    block_shape=(128, 128, 128)
+):
+  m, l = x.shape
+  l2, n = y.shape
+  assert l2 == l
+  block_m, block_n, block_l = block_shape
+  assert l % block_l == 0, f"{l=}, {block_l=}"
+  assert m % block_m == 0, f"{m=}, {block_m=}"
+  assert n % block_n == 0, f"{n=}, {block_n=}"
+  grid = (m // block_m, n // block_n, l // block_l)
+  fused_matmul = pl.pallas_call(
+      functools.partial(matmul_kernel),
+      out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
+      in_specs=[
+          pl.BlockSpec(lambda i, j, k: (i, k), (block_m, block_l)),
+          pl.BlockSpec(lambda i, j, k: (k, j), (block_l, block_n)),
+      ],
+      out_specs=pl.BlockSpec(lambda i, j, k: (i, j), (block_m, block_n)),
+      grid=grid,
+      interpret=jtu.test_device_matches(["cpu"]),
+  )
+  return fused_matmul(x, y)
+
+
+class ShapePolyTest(jtu.JaxTestCase,
+                    parameterized.TestCase):
+
+  def setUp(self):
+    if jax.config.x64_enabled:
+      self.skipTest("Only works in 32-bit")
+    if (jtu.test_device_matches(["cuda"]) and
+        not jtu.is_cuda_compute_capability_at_least("8.0")):
+      self.skipTest("Only works on GPU with capability >= sm80")
+    if sys.platform == "win32":
+      self.skipTest("Only works on non-Windows platforms")
+    super().setUp()
+    _trace_to_jaxpr.cache_clear()
+
+  def test_copy(self):
+    # The blocks are static, but the input and the grid are of polymorphic
+    # dimensions.
+    block_shape = (8, 128)
+    def f(x, *, eager=False):  # x: i32[w, h]
+      def copy_kernel(x_ref, o_ref):
+        o_ref[...] = x_ref[...]
+      # Use both pl.cdiv and // for specifying the grid
+      grid = (pl.cdiv(x.shape[0], block_shape[0]),
+              (x.shape[1] + 1) // block_shape[1])
+      return pl.pallas_call(
+          copy_kernel,
+          out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
+          in_specs=[pl.BlockSpec(lambda i, j: (i, j), block_shape)],
+          out_specs=pl.BlockSpec(lambda i, j: (i, j), block_shape),
+          grid=grid,
+          interpret=eager and jtu.test_device_matches(["cpu"]))(x)
+
+    shape1 = (128, 256)
+    x1 = jnp.arange(math.prod(shape1), dtype=np.int32).reshape(shape1)
+    res = f(x1, eager=True)
+    self.assertAllClose(res, x1)
+
+    w, h = export.symbolic_shape("w, h")
+    exp = export.export(
+        jax.jit(f),
+        platforms=["tpu"])(jax.ShapeDtypeStruct((w, h), jnp.int32))
+
+    if jtu.test_device_matches(["tpu"]):
+      res_exp_1 = exp.call(x1)
+      self.assertAllClose(res_exp_1, x1)
+
+      shape2 = block_shape
+      x2 = jnp.arange(math.prod(shape2), dtype=np.int32).reshape(shape2)
+      res_exp_2 = exp.call(x2)
+      self.assertAllClose(res_exp_2, x2)
+
+    # TODO(necula): support shape polymorphism for GPU
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        "dynamic grid bounds not supported in the Triton backend"):
+      export.export(
+          jax.jit(f),
+          platforms=["cuda"])(jax.ShapeDtypeStruct((w, h), jnp.int32))
+
+  def test_block_sizes_must_be_static_no_grid(self):
+    def f(x, *, eager=False):  # x: f32[w, h]
+      def copy_one(x_ref, o_ref):
+        o_ref[...] = x_ref[...]
+      return pl.pallas_call(
+          copy_one,
+          out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
+          interpret=eager and jtu.test_device_matches(["cpu"]))(x)
+    shape1 = (128, 256)
+    x1 = jnp.arange(math.prod(shape1), dtype=np.int32).reshape(shape1)
+    res = f(x1, eager=True)
+    self.assertAllClose(res, x1)
+
+    w, h = export.symbolic_shape("w, h")
+    with self.assertRaisesRegex(
+        ValueError,
+        "shape polymorphism for Pallas does not support dynamically-shaped blocks"):
+      export.export(
+          jax.jit(f),
+          platforms=["tpu"])(jax.ShapeDtypeStruct((w, h), jnp.int32))
+
+  def test_block_sizes_must_be_static(self):
+    def f(x, *, eager=False):  # x: f32[w, h]
+      def copy_one(x_ref, o_ref):
+        o_ref[...] = x_ref[...]
+      grid = (2, 2)
+      block_shape = (x.shape[0] // grid[0], x.shape[1] // grid[1])
+      return pl.pallas_call(
+          copy_one,
+          out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
+          in_specs=[pl.BlockSpec(lambda i, j: (i, j), block_shape)],
+          out_specs=pl.BlockSpec(lambda i, j: (i, j), block_shape),
+          grid=grid,
+          interpret=eager and jtu.test_device_matches(["cpu"]))(x)
+    shape1 = (128, 256)
+    x1 = jnp.arange(math.prod(shape1), dtype=np.int32).reshape(shape1)
+    res = f(x1, eager=True)
+    self.assertAllClose(res, x1)
+
+    w, h = export.symbolic_shape("w, h")
+    with self.assertRaisesRegex(
+        ValueError,
+        "shape polymorphism for Pallas does not support dynamically-shaped blocks"):
+
+      export.export(
+          jax.jit(f),
+          platforms=["tpu"])(jax.ShapeDtypeStruct((w, h), jnp.int32))
+
+  @jtu.run_on_devices("tpu")
+  def test_matmul(self):
+    x_shape = (1024, 256)
+    y_shape = (256, 2048)
+
+    key = jax.random.key(42)
+    key1, key2 = jax.random.split(key, 2)
+    x = jax.random.normal(key1, x_shape, dtype=np.float32)
+    y = jax.random.normal(key2, y_shape, dtype=np.float32)
+
+    res = matmul(x, y)
+    self.assertAllClose(res, x @ y, atol=1e-4)
+
+    m, n, l = export.symbolic_shape("m, n, l",
+                                    constraints=["mod(m, 128) == 0",
+                                                 "mod(n, 128) == 0",
+                                                 "mod(l, 128) == 0"])
+    exp = export.export(
+        matmul,
+        platforms=["tpu"])(
+            jax.ShapeDtypeStruct((m, l), jnp.float32),
+            jax.ShapeDtypeStruct((l, n), jnp.float32))
+    if jtu.test_device_matches(["tpu"]):
+      res_exp = exp.call(x, y)
+      self.assertAllClose(res_exp, x @ y, atol=1e-4)
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())