Expose API to control whether to fuse input computation with Pallas kernel on per input basis.

PiperOrigin-RevId: 617975104

Author: jax authors

jax/_src/pallas/mosaic/pallas_call_registration.py

@@ -102,6 +102,7 @@ def _lower_fun(*args):
         cost_estimate=mosaic_params.get("cost_estimate", None),
         vmem_limit_bytes=mosaic_params.get("vmem_limit_bytes", None),
         flags=mosaic_params.get("flags", None),
+        allow_input_fusion=mosaic_params.get("allow_input_fusion", None),
         input_output_aliases=input_output_aliases,
     )(
         *dynamic_grid_args,

tests/pallas/pallas_call_tpu_test.py

@@ -15,6 +15,7 @@
 """Test TPU-specific extensions to pallas_call."""
 
 import functools
+import re
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -1358,6 +1359,32 @@ def kernel(x_ref, y_ref):
         compiler_params=dict(mosaic=dict(vmem_limit_bytes=int(2**18))),
     )(x)
 
+  def test_allow_input_fusion(self):
+    shape = (3, 128, 128)
+
+    def kernel(x_ref, y_ref):
+      y_ref[...] = x_ref[...]
+
+    def f(x, y):
+      z = jax.numpy.add(x, y)
+      return pl.pallas_call(
+          kernel,
+          grid=(3,),
+          in_specs=[pl.BlockSpec(lambda i: (i, 0, 0), (1, 128, 128))],
+          out_specs=pl.BlockSpec(lambda i: (i, 0, 0), (1, 128, 128)),
+          out_shape=x,
+          compiler_params=dict(mosaic=dict(allow_input_fusion=[True])),
+      )(z)
+
+    x = jnp.arange(np.prod(shape), dtype=np.float32).reshape(shape)
+    y = jnp.arange(np.prod(shape), dtype=np.float32).reshape(shape)
+
+    out = f(x, y)
+    expected = x + y
+    np.testing.assert_array_equal(out, expected)
+    compiled = jax.jit(f).lower(x, y).compile().as_text()
+    assert re.search(r'fusion.*kind=kCustom.*fused_computation', compiled)
+
 
 class PallasCallUnblockedIndexingTest(PallasTPUTest):